Ray node for joining entity dataframes with feature data.
| 102 | |
| 103 | |
| 104 | class RayJoinNode(DAGNode): |
| 105 | """ |
| 106 | Ray node for joining entity dataframes with feature data. |
| 107 | """ |
| 108 | |
| 109 | def __init__( |
| 110 | self, |
| 111 | name: str, |
| 112 | column_info, |
| 113 | config: RayComputeEngineConfig, |
| 114 | is_historical_retrieval: bool = False, |
| 115 | ): |
| 116 | super().__init__(name) |
| 117 | self.column_info = column_info |
| 118 | self.config = config |
| 119 | self.is_historical_retrieval = is_historical_retrieval |
| 120 | |
| 121 | def execute(self, context: ExecutionContext) -> DAGValue: |
| 122 | """Execute the join operation.""" |
| 123 | input_value = self.get_single_input_value(context) |
| 124 | input_value.assert_format(DAGFormat.RAY) |
| 125 | feature_dataset: Dataset = input_value.data |
| 126 | |
| 127 | # If this is not a historical retrieval, just return the feature data |
| 128 | if not self.is_historical_retrieval or context.entity_df is None: |
| 129 | return DAGValue( |
| 130 | data=feature_dataset, |
| 131 | format=DAGFormat.RAY, |
| 132 | metadata={"joined": False}, |
| 133 | ) |
| 134 | |
| 135 | entity_df = context.entity_df |
| 136 | if isinstance(entity_df, pd.DataFrame): |
| 137 | ray_wrapper = get_ray_wrapper() |
| 138 | entity_dataset = ray_wrapper.from_pandas(entity_df) |
| 139 | else: |
| 140 | entity_dataset = entity_df |
| 141 | |
| 142 | join_keys = self.column_info.join_keys |
| 143 | timestamp_col = self.column_info.timestamp_column |
| 144 | requested_feats = getattr(self.column_info, "feature_cols", []) |
| 145 | |
| 146 | # Check if the feature dataset contains aggregated features (from aggregation node) |
| 147 | # If so, we don't need point-in-time join logic - just simple join on entity keys |
| 148 | is_aggregated = ( |
| 149 | input_value.metadata.get("aggregated", False) |
| 150 | if input_value.metadata |
| 151 | else False |
| 152 | ) |
| 153 | |
| 154 | feature_size = feature_dataset.size_bytes() |
| 155 | |
| 156 | if is_aggregated: |
| 157 | # For aggregated features, do simple join on entity keys |
| 158 | feature_df = feature_dataset.to_pandas() |
| 159 | feature_ref = ray.put(feature_df) |
| 160 | |
| 161 | @safe_batch_processor |
no outgoing calls