|
| 1 | +from typing import List |
| 2 | + |
| 3 | +import pyarrow as pa |
| 4 | +import pyarrow.parquet as pq |
| 5 | +import s3fs |
| 6 | +from bytewax import Dataflow, cluster_main # type: ignore |
| 7 | +from bytewax.inputs import AdvanceTo, Emit, ManualInputConfig, distribute |
| 8 | +from bytewax.parse import proc_env |
| 9 | +from tqdm import tqdm |
| 10 | + |
| 11 | +from feast import FeatureStore, FeatureView, RepoConfig |
| 12 | +from feast.utils import _convert_arrow_to_proto, _run_pyarrow_field_mapping |
| 13 | + |
| 14 | + |
| 15 | +class BytewaxMaterializationDataflow: |
| 16 | + def __init__( |
| 17 | + self, |
| 18 | + config: RepoConfig, |
| 19 | + feature_view: FeatureView, |
| 20 | + paths: List[str], |
| 21 | + ): |
| 22 | + self.config = config |
| 23 | + self.feature_store = FeatureStore(config=config) |
| 24 | + |
| 25 | + self.feature_view = feature_view |
| 26 | + self.paths = paths |
| 27 | + |
| 28 | + self._run_dataflow() |
| 29 | + |
| 30 | + def process_path(self, path): |
| 31 | + fs = s3fs.S3FileSystem() |
| 32 | + dataset = pq.ParquetDataset(path, filesystem=fs, use_legacy_dataset=False) |
| 33 | + batches = [] |
| 34 | + for fragment in dataset.fragments: |
| 35 | + for batch in fragment.to_table().to_batches(): |
| 36 | + batches.append(batch) |
| 37 | + |
| 38 | + return batches |
| 39 | + |
| 40 | + def input_builder(self, worker_index, worker_count, resume_epoch): |
| 41 | + worker_paths = distribute(self.paths, worker_index, worker_count) |
| 42 | + epoch = 0 |
| 43 | + for path in worker_paths: |
| 44 | + yield AdvanceTo(epoch) |
| 45 | + yield Emit(path) |
| 46 | + epoch += 1 |
| 47 | + |
| 48 | + return |
| 49 | + |
| 50 | + def output_builder(self, worker_index, worker_count): |
| 51 | + def output_fn(epoch_batch): |
| 52 | + _, batch = epoch_batch |
| 53 | + |
| 54 | + table = pa.Table.from_batches([batch]) |
| 55 | + |
| 56 | + if self.feature_view.batch_source.field_mapping is not None: |
| 57 | + table = _run_pyarrow_field_mapping( |
| 58 | + table, self.feature_view.batch_source.field_mapping |
| 59 | + ) |
| 60 | + |
| 61 | + join_key_to_value_type = { |
| 62 | + entity.name: entity.dtype.to_value_type() |
| 63 | + for entity in self.feature_view.entity_columns |
| 64 | + } |
| 65 | + |
| 66 | + rows_to_write = _convert_arrow_to_proto( |
| 67 | + table, self.feature_view, join_key_to_value_type |
| 68 | + ) |
| 69 | + provider = self.feature_store._get_provider() |
| 70 | + with tqdm(total=len(rows_to_write)) as progress: |
| 71 | + provider.online_write_batch( |
| 72 | + config=self.config, |
| 73 | + table=self.feature_view, |
| 74 | + data=rows_to_write, |
| 75 | + progress=progress.update, |
| 76 | + ) |
| 77 | + |
| 78 | + return output_fn |
| 79 | + |
| 80 | + def _run_dataflow(self): |
| 81 | + flow = Dataflow() |
| 82 | + flow.flat_map(self.process_path) |
| 83 | + flow.capture() |
| 84 | + cluster_main( |
| 85 | + flow, |
| 86 | + ManualInputConfig(self.input_builder), |
| 87 | + self.output_builder, |
| 88 | + **proc_env(), |
| 89 | + ) |
0 commit comments