Update Bytewax to latest version.

whoahbot · whoahbot · commit a8d4ac471cbc · 2022-11-29T16:03:13.000-08:00
Signed-off-by: Dan Herrera &lt;whoahbot@bytewax.io&gt;
diff --git a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_dataflow.py b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_dataflow.py
@@ -3,9 +3,12 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 import s3fs
-from bytewax import Dataflow, cluster_main  # type: ignore
-from bytewax.inputs import AdvanceTo, Emit, ManualInputConfig, distribute
-from bytewax.parse import proc_env
+
+from bytewax.dataflow import Dataflow  # type: ignore
+from bytewax.execution import cluster_main
+from bytewax.inputs import ManualInputConfig, distribute
+from bytewax.outputs import ManualOutputConfig
+
 from tqdm import tqdm
 
 from feast import FeatureStore, FeatureView, RepoConfig
@@ -37,20 +40,15 @@ def process_path(self, path):
 
         return batches
 
-    def input_builder(self, worker_index, worker_count, resume_epoch):
+    def input_builder(self, worker_index, worker_count, _state):
         worker_paths = distribute(self.paths, worker_index, worker_count)
-        epoch = 0
         for path in worker_paths:
-            yield AdvanceTo(epoch)
-            yield Emit(path)
-            epoch += 1
+            yield None, path
 
         return
 
     def output_builder(self, worker_index, worker_count):
-        def output_fn(epoch_batch):
-            _, batch = epoch_batch
-
+        def output_fn(batch):
             table = pa.Table.from_batches([batch])
 
             if self.feature_view.batch_source.field_mapping is not None:
@@ -79,8 +77,7 @@ def output_fn(epoch_batch):
 
     def _run_dataflow(self):
         flow = Dataflow()
+        flow.input("inp", ManualInputConfig(self.input_builder))
         flow.flat_map(self.process_path)
-        flow.capture()
-        cluster_main(
-            flow, ManualInputConfig(self.input_builder), self.output_builder, [], 0
-        )
+        flow.capture(ManualOutputConfig(self.output_builder))
+        cluster_main(flow, [], 0)
diff --git a/setup.py b/setup.py
@@ -93,7 +93,7 @@
 
 AWS_REQUIRED = ["boto3>=1.17.0,<=1.20.23", "docker>=5.0.2", "s3fs>=0.4.0,<=2022.01.0"]
 
-BYTEWAX_REQUIRED = ["bytewax==0.10.0", "docker>=5.0.2", "kubernetes<=20.13.0"]
+BYTEWAX_REQUIRED = ["bytewax==0.13.1", "docker>=5.0.2", "kubernetes<=20.13.0"]
 
 SNOWFLAKE_REQUIRED = [
     "snowflake-connector-python[pandas]>=2.7.3,<3",
