feat: make online_write_batch_size configurable in MaterializationConfig

cutoutsy · cutoutsy · commit 654b5e4e4e81 · 2026-04-13T14:15:20.000+08:00
Signed-off-by: cutoutsy &lt;cutoutsy@gmail.com&gt;
diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py
@@ -374,16 +374,32 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue:
                 for entity in self.feature_view.entity_columns
             }
 
-            rows_to_write = _convert_arrow_to_proto(
-                input_table, self.feature_view, join_key_to_value_type
-            )
-
-            online_store.online_write_batch(
-                config=context.repo_config,
-                table=self.feature_view,
-                data=rows_to_write,
-                progress=lambda x: None,
+            batch_size = (
+                context.repo_config.materialization_config.online_write_batch_size
             )
+            if batch_size is None:
+                # Default: write all rows in a single batch (backward compatible)
+                rows_to_write = _convert_arrow_to_proto(
+                    input_table, self.feature_view, join_key_to_value_type
+                )
+                online_store.online_write_batch(
+                    config=context.repo_config,
+                    table=self.feature_view,
+                    data=rows_to_write,
+                    progress=lambda x: None,
+                )
+            else:
+                # Batched writes when batch_size is configured
+                for batch in input_table.to_batches(max_chunksize=batch_size):
+                    rows_to_write = _convert_arrow_to_proto(
+                        batch, self.feature_view, join_key_to_value_type
+                    )
+                    online_store.online_write_batch(
+                        config=context.repo_config,
+                        table=self.feature_view,
+                        data=rows_to_write,
+                        progress=lambda x: None,
+                    )
 
         if self.feature_view.offline:
             offline_store = context.offline_store
diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -214,6 +214,11 @@ class MaterializationConfig(BaseModel):
     """ bool: If true, feature retrieval jobs will only pull the latest feature values for each entity.
         If false, feature retrieval jobs will pull all feature values within the specified time range. """
 
+    online_write_batch_size: Optional[int] = Field(default=None, gt=0)
+    """ int: Number of rows to write to online store per batch during materialization.
+        If None (default), all rows are written in a single batch for backward compatibility.
+        Set to a positive integer (e.g., 10000) to enable batched writes. """
+
 
 class OpenLineageConfig(FeastBaseModel):
     """Configuration for OpenLineage integration.
diff --git a/sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py b/sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py
@@ -15,6 +15,7 @@
     LocalOutputNode,
     LocalTransformationNode,
 )
+from feast.repo_config import MaterializationConfig
 
 backend = PandasBackend()
 now = pd.Timestamp.utcnow()
@@ -37,9 +38,11 @@
 
 def create_context(node_outputs):
     # Setup execution context
+    repo_config = MagicMock()
+    repo_config.materialization_config = MaterializationConfig()
     return ExecutionContext(
         project="test_proj",
-        repo_config=MagicMock(),
+        repo_config=repo_config,
         offline_store=MagicMock(),
         online_store=MagicMock(),
         entity_defs=MagicMock(),
