Skip to content

Commit 7903e2e

Browse files
fix: Stop the Ray engine from hiding materialization errors (#7026)
* fix: Fail Ray materialization when the online store write fails When the online store refused a write, the Ray engine only logged the error and still reported the job as SUCCEEDED. Users kept getting old feature values and had no sign that anything went wrong. Now the error is raised again after it is logged, and the write step is no longer wrapped in safe_batch_processor, which hid every error. A failed write now fails the job, the same as with the Local engine. Fixes #7024 Signed-off-by: Chaitany Patel <[email protected]> * fix: Keep timestamps when Ray aggregates a BatchFeatureView The Ray aggregation step kept only the entity keys and the aggregated values. It dropped the event timestamp and created timestamp columns. The online store write needs these columns, so writing an aggregated BatchFeatureView to the online store always failed. The previous commit made this failure visible. Now the aggregation keeps the latest event timestamp and created timestamp for each entity. Signed-off-by: Chaitany Patel <[email protected]> * fix: Run the Ray UDF before aggregation in historical retrieval In materialization, the Ray engine ran the UDF on raw rows before aggregation. In historical retrieval, it ran the UDF after aggregation. So a UDF could only work in one of the two paths. In the other path it raised an error, and safe_batch_processor hid that error, so users got values that were never transformed. Now historical retrieval also runs the UDF on raw rows before aggregation. This matches the compute engine docs and the Local and Spark engines. The transformation step is no longer wrapped in safe_batch_processor, so a UDF error now fails the job. A UDF that used aggregated column names (for example sum_conv_rate) must now use the raw column names (for example conv_rate). Signed-off-by: Chaitany Patel <[email protected]> --------- Signed-off-by: Chaitany Patel <[email protected]>
1 parent 68c5d54 commit 7903e2e

6 files changed

Lines changed: 123 additions & 22 deletions

File tree

‎docs/reference/compute-engine/ray.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ The Ray compute engine provides:
3535
The Ray compute engine follows Feast's DAG-based architecture:
3636

3737
```
38-
EntityDF → RayReadNode → RayJoinNode → RayFilterNode → RayAggregationNode → RayTransformationNode → Output
38+
EntityDF → RayReadNode → RayTransformationNode → RayJoinNode → RayFilterNode → RayAggregationNode → Output
3939
```
4040

4141
### Core Components

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

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ def build_aggregation_node(self, view, input_node: DAGNode) -> DAGNode:
8888
config=self.config,
8989
enable_tiling=enable_tiling,
9090
hop_size=tiling_hop_size,
91+
created_timestamp_col=getattr(
92+
view.batch_source, "created_timestamp_column", None
93+
),
9194
)
9295
node.add_input(input_node)
9396

@@ -333,19 +336,19 @@ def _build_aggregation_optimized_plan(self) -> ExecutionPlan:
333336
# 1. Read source data
334337
last_node = self.build_source_node(self.feature_view)
335338

336-
# 2. Apply filters (TTL, custom filters) BEFORE aggregation
339+
# 2. Apply transformations to the raw source data, as materialization does
340+
if self._should_transform(self.feature_view):
341+
last_node = self.build_transformation_node(self.feature_view, [last_node])
342+
343+
# 3. Apply filters (TTL, custom filters) BEFORE aggregation
337344
last_node = self.build_filter_node(self.feature_view, last_node)
338345

339-
# 3. Aggregate across all historical records
346+
# 4. Aggregate across all historical records
340347
last_node = self.build_aggregation_node(self.feature_view, last_node)
341348

342-
# 4. Join with entity_df to get aggregated features for each entity
349+
# 5. Join with entity_df to get aggregated features for each entity
343350
last_node = self.build_join_node(self.feature_view, [last_node])
344351

345-
# 5. Apply transformations to aggregated features
346-
if self._should_transform(self.feature_view):
347-
last_node = self.build_transformation_node(self.feature_view, [last_node])
348-
349352
# 6. Output
350353
last_node = self.build_output_nodes(self.feature_view, last_node)
351354

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from feast.infra.compute_engines.dag.value import DAGValue
2121
from feast.infra.compute_engines.ray.config import RayComputeEngineConfig
2222
from feast.infra.compute_engines.ray.utils import (
23+
_is_empty_batch,
2324
safe_batch_processor,
2425
write_to_online_store,
2526
)
@@ -329,11 +330,13 @@ def __init__(
329330
config: RayComputeEngineConfig,
330331
enable_tiling: bool = False,
331332
hop_size: Optional[timedelta] = None,
333+
created_timestamp_col: Optional[str] = None,
332334
):
333335
super().__init__(name)
334336
self.aggregations = aggregations
335337
self.group_by_keys = group_by_keys
336338
self.timestamp_col = timestamp_col
339+
self.created_timestamp_col = created_timestamp_col
337340
self.config = config
338341
self.enable_tiling = enable_tiling
339342
self.hop_size = hop_size
@@ -443,6 +446,21 @@ def _execute_tiled_aggregation(self, dataset: Dataset) -> DAGValue:
443446
final_pdf, self.group_by_keys, self.timestamp_col
444447
)
445448

449+
# Tiling drops the created timestamp; add back the latest per group.
450+
if (
451+
self.created_timestamp_col
452+
and self.created_timestamp_col in input_pdf.columns
453+
and not final_pdf.empty
454+
):
455+
latest_created = (
456+
input_pdf.groupby(self.group_by_keys)[self.created_timestamp_col]
457+
.max()
458+
.reset_index()
459+
)
460+
final_pdf = final_pdf.merge(
461+
latest_created, on=self.group_by_keys, how="left"
462+
)
463+
446464
aggregated_dataset = ray_wrapper.from_pandas(final_pdf)
447465

448466
return DAGValue(
@@ -543,6 +561,10 @@ def _fallback_pandas_aggregation(self, dataset: Dataset, agg_dict: dict) -> Data
543561

544562
# Combine aggregation results
545563
if agg_results:
564+
# Keep the latest timestamps per group; the online write needs them.
565+
for ts_col in (self.timestamp_col, self.created_timestamp_col):
566+
if ts_col and ts_col in df.columns:
567+
agg_results.append(grouped[ts_col].max())
546568
result_df = pd.concat(agg_results, axis=1)
547569

548570
# Reset index to make group keys regular columns
@@ -708,11 +730,14 @@ def execute(self, context: ExecutionContext) -> DAGValue:
708730
elif callable(self.transformation):
709731
transformation_serialized = dill.dumps(self.transformation)
710732

711-
@safe_batch_processor
733+
# Not wrapped in safe_batch_processor: a failed UDF must fail the job.
712734
def apply_transformation_with_serialized_udf(
713735
batch: pd.DataFrame,
714736
) -> pd.DataFrame:
715737
"""Apply the transformation using pre-serialized UDF."""
738+
if _is_empty_batch(batch):
739+
return batch
740+
716741
if transformation_serialized:
717742
transformation_func = dill.loads(transformation_serialized)
718743
transformed_batch = transformation_func(batch)
@@ -861,9 +886,12 @@ def execute(self, context: ExecutionContext) -> DAGValue:
861886
feature_view=self.feature_view, repo_config=context.repo_config
862887
)
863888

864-
@safe_batch_processor
889+
# Not wrapped in safe_batch_processor: a failed write must fail the job.
865890
def write_batch_with_serialized_artifacts(batch: pd.DataFrame) -> pd.DataFrame:
866891
"""Write each batch using pre-serialized artifacts."""
892+
if _is_empty_batch(batch):
893+
return batch
894+
867895
(
868896
feature_view,
869897
online_store,

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ def write_to_online_store(
8181

8282
except Exception as e:
8383
logger.error(f"Failed to write to online store for {feature_view.name}: {e}")
84+
raise
8485

8586

8687
# Ray Data batch type: pandas DataFrame, numpy dict, or pyarrow Table
@@ -178,6 +179,7 @@ def _write_batch(batch: pa.Table) -> pa.Table:
178179
logger.error(
179180
f"Distributed write to online store failed for {feature_view.name}: {e}"
180181
)
182+
raise
181183

182184

183185
def safe_batch_processor(

‎sdk/python/tests/component/ray/test_compute.py‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,10 @@ def test_ray_compute_engine_get_historical_features(
4141
fs = ray_environment.feature_store
4242
registry = fs.registry
4343

44+
# The udf runs on raw source rows, before aggregation.
4445
def transform_feature(df: pd.DataFrame) -> pd.DataFrame:
45-
df["sum_conv_rate"] = df["sum_conv_rate"] * 2
46-
df["avg_acc_rate"] = df["avg_acc_rate"] * 2
46+
df["conv_rate"] = df["conv_rate"] * 2
47+
df["acc_rate"] = df["acc_rate"] * 2
4748
return df
4849

4950
driver_stats_fv = BatchFeatureView(
@@ -102,11 +103,6 @@ def test_ray_compute_engine_materialize(ray_environment, feature_dataset):
102103
fs = ray_environment.feature_store
103104
registry = fs.registry
104105

105-
def transform_feature(df: pd.DataFrame) -> pd.DataFrame:
106-
df["sum_conv_rate"] = df["sum_conv_rate"] * 2
107-
df["avg_acc_rate"] = df["avg_acc_rate"] * 2
108-
return df
109-
110106
driver_stats_fv = BatchFeatureView(
111107
name="driver_hourly_stats",
112108
entities=[driver],
@@ -115,13 +111,10 @@ def transform_feature(df: pd.DataFrame) -> pd.DataFrame:
115111
Aggregation(column="conv_rate", function="sum"),
116112
Aggregation(column="acc_rate", function="avg"),
117113
],
118-
udf=transform_feature,
119-
udf_string="transform_feature",
120114
ttl=timedelta(days=3),
121115
schema=[
122-
Field(name="conv_rate", dtype=Float32),
123-
Field(name="acc_rate", dtype=Float32),
124-
Field(name="avg_daily_trips", dtype=Int64),
116+
Field(name="sum_conv_rate", dtype=Float32),
117+
Field(name="avg_acc_rate", dtype=Float32),
125118
Field(name="driver_id", dtype=Int32),
126119
],
127120
online=True,
@@ -153,6 +146,17 @@ def tqdm_builder(length):
153146
assert len(ray_materialize_jobs) == 1
154147
assert ray_materialize_jobs[0].status() == MaterializationJobStatus.SUCCEEDED
155148

149+
# driver 1001 has one row in the window: conv_rate 0.8, acc_rate 0.5
150+
online = fs.get_online_features(
151+
features=[
152+
"driver_hourly_stats:sum_conv_rate",
153+
"driver_hourly_stats:avg_acc_rate",
154+
],
155+
entity_rows=[{"driver_id": 1001}],
156+
).to_dict()
157+
assert online["sum_conv_rate"][0] == pytest.approx(0.8)
158+
assert online["avg_acc_rate"][0] == pytest.approx(0.5)
159+
156160

157161
@pytest.mark.integration
158162
@pytest.mark.xdist_group(name="ray")

‎sdk/python/tests/component/ray/test_nodes.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import pytest
77
import ray
88

9+
from feast import Entity, FeatureView, Field, FileSource, RepoConfig
910
from feast.aggregation import Aggregation
1011
from feast.infra.compute_engines.dag.context import ColumnInfo
1112
from feast.infra.compute_engines.dag.model import DAGFormat
@@ -19,13 +20,16 @@
1920
RayJoinNode,
2021
RayReadNode,
2122
RayTransformationNode,
23+
RayWriteNode,
2224
)
25+
from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
2326
from feast.infra.ray_initializer import (
2427
RayConfigManager,
2528
RayExecutionMode,
2629
ensure_ray_initialized,
2730
get_ray_wrapper,
2831
)
32+
from feast.types import Float32
2933

3034

3135
class DummyInputNode(DAGNode):
@@ -195,6 +199,34 @@ def test_ray_aggregation_node(
195199
assert "avg_acc_rate" in result_df.columns
196200

197201

202+
def test_ray_tiled_aggregation_node_keeps_created_timestamp(
203+
ray_session, ray_config, mock_context, sample_data
204+
):
205+
"""Tiled aggregation must keep the latest created timestamp per entity."""
206+
input_value = DAGValue(data=ray.data.from_pandas(sample_data), format=DAGFormat.RAY)
207+
node = RayAggregationNode(
208+
name="aggregation",
209+
aggregations=[
210+
Aggregation(
211+
column="conv_rate", function="sum", time_window=timedelta(days=1)
212+
)
213+
],
214+
group_by_keys=["driver_id"],
215+
timestamp_col="event_timestamp",
216+
config=ray_config,
217+
enable_tiling=True,
218+
hop_size=timedelta(hours=1),
219+
created_timestamp_col="created",
220+
)
221+
node.add_input(DummyInputNode("input_node", input_value))
222+
mock_context.node_outputs = {"input_node": input_value}
223+
224+
result_df = node.execute(mock_context).data.to_pandas().set_index("driver_id")
225+
226+
expected = sample_data.groupby("driver_id")["created"].max()
227+
assert (result_df["created"] == expected.loc[result_df.index]).all()
228+
229+
198230
def test_ray_join_node(ray_session, ray_config, mock_context, sample_data, column_info):
199231
"""Test RayJoinNode functionality."""
200232
entity_data = pd.DataFrame(
@@ -439,6 +471,38 @@ def test_ray_dedup_node_materialization_no_join_keys(
439471
)
440472

441473

474+
def test_ray_write_node_raises_when_online_write_fails(
475+
ray_session, mock_context, sample_data, tmp_path
476+
):
477+
"""A failed online store write must fail the write node, not be swallowed."""
478+
driver = Entity(name="driver", join_keys=["driver_id"])
479+
feature_view = FeatureView(
480+
name="driver_stats",
481+
entities=[driver],
482+
schema=[Field(name="conv_rate", dtype=Float32)],
483+
source=FileSource(
484+
path=str(tmp_path / "source.parquet"), timestamp_field="event_timestamp"
485+
),
486+
online=True,
487+
offline=False,
488+
)
489+
# The online table was never created, so every write raises "no such table".
490+
mock_context.repo_config = RepoConfig(
491+
project="test_project",
492+
registry=str(tmp_path / "registry.db"),
493+
provider="local",
494+
entity_key_serialization_version=3,
495+
online_store=SqliteOnlineStoreConfig(path=str(tmp_path / "online.db")),
496+
)
497+
input_value = DAGValue(data=ray.data.from_pandas(sample_data), format=DAGFormat.RAY)
498+
node = RayWriteNode(name="write", feature_view=feature_view)
499+
node.add_input(DummyInputNode("input_node", input_value))
500+
mock_context.node_outputs = {"input_node": input_value}
501+
502+
with pytest.raises(Exception, match="no such table"):
503+
node.execute(mock_context)
504+
505+
442506
def test_ray_config_validation():
443507
"""Test Ray configuration validation."""
444508
# Test valid configuration

0 commit comments

Comments
 (0)