Execute the join operation.
(self, context: ExecutionContext)
| 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 |
| 162 | def join_with_aggregated_features(batch: pd.DataFrame) -> pd.DataFrame: |
| 163 | features = ray.get(feature_ref) |
| 164 | if join_keys: |
| 165 | result = pd.merge( |
| 166 | batch, |
| 167 | features, |
| 168 | on=join_keys, |
| 169 | how="left", |
| 170 | suffixes=("", "_feature"), |
| 171 | ) |
| 172 | else: |
| 173 | result = batch.copy() |
| 174 | return result |
| 175 | |
| 176 | joined_dataset = entity_dataset.map_batches( |
| 177 | join_with_aggregated_features, |
| 178 | batch_format="pandas", |