Test RayJoinNode functionality.
(ray_session, ray_config, mock_context, sample_data, column_info)
| 196 | |
| 197 | |
| 198 | def test_ray_join_node(ray_session, ray_config, mock_context, sample_data, column_info): |
| 199 | """Test RayJoinNode functionality.""" |
| 200 | entity_data = pd.DataFrame( |
| 201 | [ |
| 202 | {"driver_id": 1001, "event_timestamp": datetime.now()}, |
| 203 | {"driver_id": 1002, "event_timestamp": datetime.now()}, |
| 204 | ] |
| 205 | ) |
| 206 | feature_dataset = ray.data.from_pandas(sample_data) |
| 207 | feature_value = DAGValue(data=feature_dataset, format=DAGFormat.RAY) |
| 208 | dummy_node = DummyInputNode("feature_node", feature_value) |
| 209 | node = RayJoinNode( |
| 210 | name="join", |
| 211 | column_info=column_info, |
| 212 | config=ray_config, |
| 213 | ) |
| 214 | node.add_input(dummy_node) |
| 215 | mock_context.node_outputs = {"feature_node": feature_value} |
| 216 | mock_context.entity_df = entity_data |
| 217 | result = node.execute(mock_context) |
| 218 | assert isinstance(result, DAGValue) |
| 219 | assert result.format == DAGFormat.RAY |
| 220 | result_df = result.data.to_pandas() |
| 221 | assert len(result_df) >= 2 |
| 222 | assert "driver_id" in result_df.columns |
| 223 | |
| 224 | |
| 225 | def test_ray_transformation_node( |
nothing calls this directly
no test coverage detected