|
6 | 6 | import pytest |
7 | 7 | import ray |
8 | 8 |
|
| 9 | +from feast import Entity, FeatureView, Field, FileSource, RepoConfig |
9 | 10 | from feast.aggregation import Aggregation |
10 | 11 | from feast.infra.compute_engines.dag.context import ColumnInfo |
11 | 12 | from feast.infra.compute_engines.dag.model import DAGFormat |
|
19 | 20 | RayJoinNode, |
20 | 21 | RayReadNode, |
21 | 22 | RayTransformationNode, |
| 23 | + RayWriteNode, |
22 | 24 | ) |
| 25 | +from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig |
23 | 26 | from feast.infra.ray_initializer import ( |
24 | 27 | RayConfigManager, |
25 | 28 | RayExecutionMode, |
26 | 29 | ensure_ray_initialized, |
27 | 30 | get_ray_wrapper, |
28 | 31 | ) |
| 32 | +from feast.types import Float32 |
29 | 33 |
|
30 | 34 |
|
31 | 35 | class DummyInputNode(DAGNode): |
@@ -195,6 +199,34 @@ def test_ray_aggregation_node( |
195 | 199 | assert "avg_acc_rate" in result_df.columns |
196 | 200 |
|
197 | 201 |
|
| 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 | + |
198 | 230 | def test_ray_join_node(ray_session, ray_config, mock_context, sample_data, column_info): |
199 | 231 | """Test RayJoinNode functionality.""" |
200 | 232 | entity_data = pd.DataFrame( |
@@ -439,6 +471,38 @@ def test_ray_dedup_node_materialization_no_join_keys( |
439 | 471 | ) |
440 | 472 |
|
441 | 473 |
|
| 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 | + |
442 | 506 | def test_ray_config_validation(): |
443 | 507 | """Test Ray configuration validation.""" |
444 | 508 | # Test valid configuration |
|
0 commit comments