Skip to content

Commit bff1499

Browse files
cutoutsycutoutsy
committed
refactor: simplify batch write logic and extend to spark/ray engines
Signed-off-by: cutoutsy <cutoutsy@gmail.com>
1 parent 654b5e4 commit bff1499

5 files changed

Lines changed: 114 additions & 42 deletions

File tree

‎sdk/python/feast/infra/compute_engines/local/nodes.py‎

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -377,29 +377,22 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue:
377377
batch_size = (
378378
context.repo_config.materialization_config.online_write_batch_size
379379
)
380-
if batch_size is None:
381-
# Default: write all rows in a single batch (backward compatible)
380+
# Single batch if None (backward compatible), otherwise use configured batch_size
381+
batches = (
382+
[input_table]
383+
if batch_size is None
384+
else input_table.to_batches(max_chunksize=batch_size)
385+
)
386+
for batch in batches:
382387
rows_to_write = _convert_arrow_to_proto(
383-
input_table, self.feature_view, join_key_to_value_type
388+
batch, self.feature_view, join_key_to_value_type
384389
)
385390
online_store.online_write_batch(
386391
config=context.repo_config,
387392
table=self.feature_view,
388393
data=rows_to_write,
389394
progress=lambda x: None,
390395
)
391-
else:
392-
# Batched writes when batch_size is configured
393-
for batch in input_table.to_batches(max_chunksize=batch_size):
394-
rows_to_write = _convert_arrow_to_proto(
395-
batch, self.feature_view, join_key_to_value_type
396-
)
397-
online_store.online_write_batch(
398-
config=context.repo_config,
399-
table=self.feature_view,
400-
data=rows_to_write,
401-
progress=lambda x: None,
402-
)
403396

404397
if self.feature_view.offline:
405398
offline_store = context.offline_store

‎sdk/python/feast/infra/compute_engines/ray/utils.py‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,19 +45,32 @@ def write_to_online_store(
4545
for entity in feature_view.entity_columns
4646
}
4747

48-
rows_to_write = _convert_arrow_to_proto(
49-
arrow_table, feature_view, join_key_to_value_type
48+
batch_size = repo_config.materialization_config.online_write_batch_size
49+
# Single batch if None (backward compatible), otherwise use configured batch_size
50+
batches = (
51+
[arrow_table]
52+
if batch_size is None
53+
else arrow_table.to_batches(max_chunksize=batch_size)
5054
)
5155

52-
if rows_to_write:
53-
online_store.online_write_batch(
54-
config=repo_config,
55-
table=feature_view,
56-
data=rows_to_write,
57-
progress=lambda x: None,
56+
total_rows = 0
57+
for batch in batches:
58+
rows_to_write = _convert_arrow_to_proto(
59+
batch, feature_view, join_key_to_value_type
5860
)
61+
62+
if rows_to_write:
63+
online_store.online_write_batch(
64+
config=repo_config,
65+
table=feature_view,
66+
data=rows_to_write,
67+
progress=lambda x: None,
68+
)
69+
total_rows += len(rows_to_write)
70+
71+
if total_rows > 0:
5972
logger.debug(
60-
f"Successfully wrote {len(rows_to_write)} rows to online store for {feature_view.name}"
73+
f"Successfully wrote {total_rows} rows to online store for {feature_view.name}"
6174
)
6275
else:
6376
logger.warning(f"No rows to write for {feature_view.name}")

‎sdk/python/feast/infra/compute_engines/spark/utils.py‎

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -47,16 +47,24 @@ def map_in_arrow(
4747
for entity in feature_view.entity_columns
4848
}
4949

50-
rows_to_write = _convert_arrow_to_proto(
51-
table, feature_view, join_key_to_value_type
52-
)
53-
54-
online_store.online_write_batch(
55-
config=repo_config,
56-
table=feature_view,
57-
data=rows_to_write,
58-
progress=lambda x: None,
50+
batch_size = repo_config.materialization_config.online_write_batch_size
51+
# Single batch if None (backward compatible), otherwise use configured batch_size
52+
sub_batches = (
53+
[table]
54+
if batch_size is None
55+
else table.to_batches(max_chunksize=batch_size)
5956
)
57+
for sub_batch in sub_batches:
58+
rows_to_write = _convert_arrow_to_proto(
59+
sub_batch, feature_view, join_key_to_value_type
60+
)
61+
62+
online_store.online_write_batch(
63+
config=repo_config,
64+
table=feature_view,
65+
data=rows_to_write,
66+
progress=lambda x: None,
67+
)
6068
if mode == "offline":
6169
offline_store.offline_write_batch(
6270
config=repo_config,
@@ -95,15 +103,23 @@ def map_in_pandas(iterator, serialized_artifacts: SerializedArtifacts):
95103
for entity in feature_view.entity_columns
96104
}
97105

98-
rows_to_write = _convert_arrow_to_proto(
99-
table, feature_view, join_key_to_value_type
100-
)
101-
online_store.online_write_batch(
102-
repo_config,
103-
feature_view,
104-
rows_to_write,
105-
lambda x: None,
106+
batch_size = repo_config.materialization_config.online_write_batch_size
107+
# Single batch if None (backward compatible), otherwise use configured batch_size
108+
sub_batches = (
109+
[table]
110+
if batch_size is None
111+
else table.to_batches(max_chunksize=batch_size)
106112
)
113+
for sub_batch in sub_batches:
114+
rows_to_write = _convert_arrow_to_proto(
115+
sub_batch, feature_view, join_key_to_value_type
116+
)
117+
online_store.online_write_batch(
118+
repo_config,
119+
feature_view,
120+
rows_to_write,
121+
lambda x: None,
122+
)
107123

108124
yield pd.DataFrame(
109125
[pd.Series(range(1, 2))]

‎sdk/python/feast/repo_config.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,8 @@ class MaterializationConfig(BaseModel):
217217
online_write_batch_size: Optional[int] = Field(default=None, gt=0)
218218
""" int: Number of rows to write to online store per batch during materialization.
219219
If None (default), all rows are written in a single batch for backward compatibility.
220-
Set to a positive integer (e.g., 10000) to enable batched writes. """
220+
Set to a positive integer (e.g., 10000) to enable batched writes.
221+
Supported compute engines: local, spark, ray. """
221222

222223

223224
class OpenLineageConfig(FeastBaseModel):

‎sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,3 +217,52 @@ def test_local_output_node():
217217
node.inputs[0].name = "source"
218218
result = node.execute(context)
219219
assert result.num_rows == 4
220+
221+
222+
def test_local_output_node_online_write_default_batch():
223+
"""Test that online_write_batch is called once when batch_size is None (default)."""
224+
# Create a feature view with online=True
225+
feature_view = MagicMock()
226+
feature_view.online = True
227+
feature_view.offline = False
228+
feature_view.entity_columns = []
229+
230+
# Create context with default materialization config (batch_size=None)
231+
context = create_context(
232+
node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(sample_df))}
233+
)
234+
235+
node = LocalOutputNode("output", feature_view)
236+
node.add_input(MagicMock())
237+
node.inputs[0].name = "source"
238+
239+
node.execute(context)
240+
241+
# Verify online_write_batch was called exactly once (all rows in single batch)
242+
assert context.online_store.online_write_batch.call_count == 1
243+
244+
245+
def test_local_output_node_online_write_batched():
246+
"""Test that online_write_batch is called multiple times when batch_size is configured."""
247+
# Create a feature view with online=True
248+
feature_view = MagicMock()
249+
feature_view.online = True
250+
feature_view.offline = False
251+
feature_view.entity_columns = []
252+
253+
# Create context with batch_size=2 (sample_df has 4 rows, so expect 2 batches)
254+
context = create_context(
255+
node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(sample_df))}
256+
)
257+
context.repo_config.materialization_config = MaterializationConfig(
258+
online_write_batch_size=2
259+
)
260+
261+
node = LocalOutputNode("output", feature_view)
262+
node.add_input(MagicMock())
263+
node.inputs[0].name = "source"
264+
265+
node.execute(context)
266+
267+
# Verify online_write_batch was called twice (4 rows / batch_size 2 = 2 batches)
268+
assert context.online_store.online_write_batch.call_count == 2

0 commit comments

Comments
 (0)