Execute the write operation.
(self, context: ExecutionContext)
| 852 | self.config = config |
| 853 | |
| 854 | def execute(self, context: ExecutionContext) -> DAGValue: |
| 855 | """Execute the write operation.""" |
| 856 | input_value = self.get_single_input_value(context) |
| 857 | input_value.assert_format(DAGFormat.RAY) |
| 858 | dataset: Dataset = input_value.data |
| 859 | |
| 860 | serialized_artifacts = SerializedArtifacts.serialize( |
| 861 | feature_view=self.feature_view, repo_config=context.repo_config |
| 862 | ) |
| 863 | |
| 864 | @safe_batch_processor |
| 865 | def write_batch_with_serialized_artifacts(batch: pd.DataFrame) -> pd.DataFrame: |
| 866 | """Write each batch using pre-serialized artifacts.""" |
| 867 | ( |
| 868 | feature_view, |
| 869 | online_store, |
| 870 | offline_store, |
| 871 | repo_config, |
| 872 | ) = serialized_artifacts.unserialize() |
| 873 | |
| 874 | arrow_table = pa.Table.from_pandas(batch) |
| 875 | |
| 876 | # Write to online store if enabled |
| 877 | write_to_online_store( |
| 878 | arrow_table=arrow_table, |
| 879 | feature_view=feature_view, |
| 880 | online_store=online_store, |
| 881 | repo_config=repo_config, |
| 882 | ) |
| 883 | |
| 884 | # Write to offline store if enabled |
| 885 | if getattr(feature_view, "offline", False): |
| 886 | offline_store.offline_write_batch( |
| 887 | config=repo_config, |
| 888 | feature_view=feature_view, |
| 889 | table=arrow_table, |
| 890 | progress=lambda x: None, |
| 891 | ) |
| 892 | |
| 893 | return batch |
| 894 | |
| 895 | # Resolve write concurrency from config. |
| 896 | # write_concurrency takes precedence; falls back to max_workers, then 1. |
| 897 | if self.config is not None and self.config.write_concurrency is not None: |
| 898 | _write_concurrency = self.config.write_concurrency |
| 899 | elif self.config is not None and self.config.max_workers is not None: |
| 900 | _write_concurrency = self.config.max_workers |
| 901 | else: |
| 902 | _write_concurrency = 1 |
| 903 | |
| 904 | written_dataset = dataset.map_batches( |
| 905 | write_batch_with_serialized_artifacts, |
| 906 | batch_format="pandas", |
| 907 | concurrency=_write_concurrency, |
| 908 | ) |
| 909 | written_dataset = written_dataset.materialize() |
| 910 | |
| 911 | return DAGValue( |
nothing calls this directly
no test coverage detected