@@ -1048,12 +1048,11 @@ def to_spark_df(self) -> pyspark.sql.DataFrame: |
1048 | 1048 | """ |
1049 | 1049 | spark_df = self._to_spark_df_internal() |
1050 | 1050 | if self.on_demand_feature_views: |
1051 | | - return _apply_odfv_transformations( |
1052 | | - spark_df=spark_df, |
| 1051 | + return SparkOnDemandFeatureViewProcessor( |
1053 | 1052 | on_demand_feature_views=self.on_demand_feature_views, |
1054 | 1053 | full_feature_names=self.full_feature_names, |
1055 | 1054 | metadata=self.metadata, |
1056 | | - ) |
| 1055 | + ).transform(spark_df) |
1057 | 1056 | return spark_df |
1058 | 1057 |
|
1059 | 1058 | def _to_spark_df_internal(self) -> pyspark.sql.DataFrame: |
@@ -1265,140 +1264,155 @@ def metadata(self) -> Optional[RetrievalMetadata]: |
1265 | 1264 | return self._metadata |
1266 | 1265 |
|
1267 | 1266 |
|
1268 | | -def _apply_odfv_transformations( |
1269 | | - spark_df: pyspark.sql.DataFrame, |
1270 | | - on_demand_feature_views: List[OnDemandFeatureView], |
1271 | | - full_feature_names: bool, |
1272 | | - metadata: Optional[RetrievalMetadata], |
1273 | | -) -> pyspark.sql.DataFrame: |
| 1267 | +class SparkOnDemandFeatureViewProcessor: |
1274 | 1268 | """ |
1275 | | - Append the requested on demand features to ``spark_df``, computed on the executors. |
| 1269 | + Appends the requested on demand features to a Spark DataFrame on the executors. |
1276 | 1270 |
|
1277 | 1271 | The requested ODFV features are added under the names ``RetrievalJob.to_arrow`` |
1278 | 1272 | gives them, or every declared feature of an ODFV when none of its features is |
1279 | 1273 | listed. Each executor rebuilds the ODFVs from their protos and applies them to |
1280 | 1274 | every Arrow batch, casting the results to the declared feature types where Feast |
1281 | 1275 | has an Arrow type for them. |
1282 | 1276 | """ |
1283 | | - from pyspark.sql.types import StructField, StructType |
1284 | | - |
1285 | | - from feast.infra.compute_engines.spark.nodes import from_feast_to_spark_type |
1286 | | - |
1287 | | - odfv_feature_refs: Dict[str, set[str]] = {} |
1288 | | - if metadata and metadata.features: |
1289 | | - for feature_ref in metadata.features: |
1290 | | - if ":" in feature_ref: |
1291 | | - view_name, feature_name = feature_ref.split(":", 1) |
1292 | | - for odfv in on_demand_feature_views: |
1293 | | - if ( |
1294 | | - odfv.name == view_name |
1295 | | - or odfv.projection.name_to_use() == view_name |
1296 | | - ): |
1297 | | - if view_name not in odfv_feature_refs: |
1298 | | - odfv_feature_refs[view_name] = set() |
1299 | | - expected_col_name = ( |
1300 | | - f"{odfv.projection.name_to_use()}__{feature_name}" |
1301 | | - if full_feature_names |
1302 | | - else feature_name |
1303 | | - ) |
1304 | | - odfv_feature_refs[view_name].add(expected_col_name) |
1305 | | - |
1306 | | - output_features: Dict[str, Field] = {} |
1307 | | - for odfv in on_demand_feature_views: |
1308 | | - requested_features_for_odfv = ( |
1309 | | - odfv_feature_refs.get(odfv.name) |
1310 | | - if odfv.name in odfv_feature_refs |
1311 | | - else odfv_feature_refs.get(odfv.projection.name_to_use()) |
1312 | | - ) |
1313 | | - for feature in odfv.features: |
1314 | | - col = ( |
1315 | | - f"{odfv.projection.name_to_use()}__{feature.name}" |
1316 | | - if full_feature_names |
1317 | | - else feature.name |
1318 | | - ) |
1319 | | - if ( |
1320 | | - requested_features_for_odfv is None |
1321 | | - or col in requested_features_for_odfv |
1322 | | - ): |
1323 | | - output_features[col] = feature |
1324 | | - |
1325 | | - spark_fields = list(spark_df.schema.fields) |
1326 | | - # Feast-declared Arrow types to cast the ODFV outputs to. Types without one |
1327 | | - # (e.g. Bytes) keep the type Arrow inferred. |
1328 | | - output_pa_types: Dict[str, pyarrow.DataType] = {} |
1329 | | - for name, feature in output_features.items(): |
1330 | | - spark_type = from_feast_to_spark_type(feature.dtype) |
1331 | | - if spark_type is None: |
1332 | | - raise ValueError( |
1333 | | - f"Could not resolve Spark type for on demand feature '{name}' " |
1334 | | - f"(dtype={feature.dtype})." |
1335 | | - ) |
1336 | | - spark_fields.append(StructField(name, spark_type, nullable=True)) |
1337 | | - try: |
1338 | | - pa_type = from_feast_to_pyarrow_type(feature.dtype) |
1339 | | - except (ValueError, KeyError, TypeError): |
1340 | | - continue |
1341 | | - output_pa_types[name] = pa_type |
1342 | 1277 |
|
1343 | | - # The ODFV proto does not carry its projection (alias, version tag), so ship the |
1344 | | - # projection proto alongside and re-apply it with with_projection on the executors. |
1345 | | - odfv_protos = [ |
1346 | | - ( |
1347 | | - odfv.to_proto().SerializeToString(), |
1348 | | - odfv.projection.to_proto().SerializeToString(), |
| 1278 | + def __init__( |
| 1279 | + self, |
| 1280 | + on_demand_feature_views: List[OnDemandFeatureView], |
| 1281 | + full_feature_names: bool, |
| 1282 | + metadata: Optional[RetrievalMetadata], |
| 1283 | + ): |
| 1284 | + self.on_demand_feature_views = on_demand_feature_views |
| 1285 | + self.full_feature_names = full_feature_names |
| 1286 | + self.metadata = metadata |
| 1287 | + |
| 1288 | + def transform(self, spark_df: pyspark.sql.DataFrame) -> pyspark.sql.DataFrame: |
| 1289 | + from pyspark.sql.types import StructField, StructType |
| 1290 | + |
| 1291 | + from feast.infra.compute_engines.spark.nodes import from_feast_to_spark_type |
| 1292 | + |
| 1293 | + output_features = self._get_requested_features() |
| 1294 | + |
| 1295 | + spark_fields = list(spark_df.schema.fields) |
| 1296 | + # Feast-declared Arrow types to cast the ODFV outputs to. Types without one |
| 1297 | + # (e.g. Bytes) keep the type Arrow inferred. |
| 1298 | + output_pa_types: Dict[str, pyarrow.DataType] = {} |
| 1299 | + for name, feature in output_features.items(): |
| 1300 | + spark_type = from_feast_to_spark_type(feature.dtype) |
| 1301 | + if spark_type is None: |
| 1302 | + raise ValueError( |
| 1303 | + f"Could not resolve Spark type for on demand feature '{name}' " |
| 1304 | + f"(dtype={feature.dtype})." |
| 1305 | + ) |
| 1306 | + spark_fields.append(StructField(name, spark_type, nullable=True)) |
| 1307 | + try: |
| 1308 | + pa_type = from_feast_to_pyarrow_type(feature.dtype) |
| 1309 | + except (ValueError, KeyError, TypeError): |
| 1310 | + continue |
| 1311 | + output_pa_types[name] = pa_type |
| 1312 | + |
| 1313 | + # The ODFV proto does not carry its projection (alias, version tag), so ship |
| 1314 | + # the projection proto alongside and re-apply it with with_projection on the |
| 1315 | + # executors. |
| 1316 | + odfv_protos = [ |
| 1317 | + ( |
| 1318 | + odfv.to_proto().SerializeToString(), |
| 1319 | + odfv.projection.to_proto().SerializeToString(), |
| 1320 | + ) |
| 1321 | + for odfv in self.on_demand_feature_views |
| 1322 | + ] |
| 1323 | + # Keep self out of the closure, so the ODFVs reach the executors only as protos. |
| 1324 | + full_feature_names = self.full_feature_names |
| 1325 | + output_columns = list(output_features) |
| 1326 | + return spark_df.mapInArrow( |
| 1327 | + lambda x: SparkOnDemandFeatureViewProcessor._map_in_arrow( |
| 1328 | + x, |
| 1329 | + odfv_protos, |
| 1330 | + full_feature_names=full_feature_names, |
| 1331 | + output_columns=output_columns, |
| 1332 | + output_pa_types=output_pa_types, |
| 1333 | + ), |
| 1334 | + StructType(spark_fields), |
1349 | 1335 | ) |
1350 | | - for odfv in on_demand_feature_views |
1351 | | - ] |
1352 | | - return spark_df.mapInArrow( |
1353 | | - lambda x: _map_in_arrow_odfv( |
1354 | | - x, |
1355 | | - odfv_protos, |
1356 | | - full_feature_names=full_feature_names, |
1357 | | - output_columns=list(output_features), |
1358 | | - output_pa_types=output_pa_types, |
1359 | | - ), |
1360 | | - StructType(spark_fields), |
1361 | | - ) |
1362 | | - |
1363 | 1336 |
|
1364 | | -def _map_in_arrow_odfv( |
1365 | | - iterator: Iterable[pyarrow.RecordBatch], |
1366 | | - odfv_protos: List[Tuple[bytes, bytes]], |
1367 | | - full_feature_names: bool, |
1368 | | - output_columns: List[str], |
1369 | | - output_pa_types: Dict[str, pyarrow.DataType], |
1370 | | -) -> Iterator[pyarrow.RecordBatch]: |
1371 | | - for batch in iterator: |
1372 | | - features_table = pyarrow.Table.from_batches([batch]) |
1373 | | - |
1374 | | - odfvs = [] |
1375 | | - for serialized_odfv, serialized_projection in odfv_protos: |
1376 | | - odfv_proto = OnDemandFeatureViewProto.FromString(serialized_odfv) |
1377 | | - projection_proto = FeatureViewProjectionProto.FromString( |
1378 | | - serialized_projection |
| 1337 | + def _get_requested_features(self) -> Dict[str, Field]: |
| 1338 | + odfv_feature_refs: Dict[str, set[str]] = {} |
| 1339 | + if self.metadata and self.metadata.features: |
| 1340 | + for feature_ref in self.metadata.features: |
| 1341 | + if ":" in feature_ref: |
| 1342 | + view_name, feature_name = feature_ref.split(":", 1) |
| 1343 | + for odfv in self.on_demand_feature_views: |
| 1344 | + if ( |
| 1345 | + odfv.name == view_name |
| 1346 | + or odfv.projection.name_to_use() == view_name |
| 1347 | + ): |
| 1348 | + if view_name not in odfv_feature_refs: |
| 1349 | + odfv_feature_refs[view_name] = set() |
| 1350 | + expected_col_name = ( |
| 1351 | + f"{odfv.projection.name_to_use()}__{feature_name}" |
| 1352 | + if self.full_feature_names |
| 1353 | + else feature_name |
| 1354 | + ) |
| 1355 | + odfv_feature_refs[view_name].add(expected_col_name) |
| 1356 | + |
| 1357 | + output_features: Dict[str, Field] = {} |
| 1358 | + for odfv in self.on_demand_feature_views: |
| 1359 | + requested_features_for_odfv = ( |
| 1360 | + odfv_feature_refs.get(odfv.name) |
| 1361 | + if odfv.name in odfv_feature_refs |
| 1362 | + else odfv_feature_refs.get(odfv.projection.name_to_use()) |
1379 | 1363 | ) |
1380 | | - odfvs.append( |
1381 | | - OnDemandFeatureView.from_proto(odfv_proto).with_projection( |
1382 | | - FeatureViewProjection.from_proto(projection_proto) |
| 1364 | + for feature in odfv.features: |
| 1365 | + col = ( |
| 1366 | + f"{odfv.projection.name_to_use()}__{feature.name}" |
| 1367 | + if self.full_feature_names |
| 1368 | + else feature.name |
| 1369 | + ) |
| 1370 | + if ( |
| 1371 | + requested_features_for_odfv is None |
| 1372 | + or col in requested_features_for_odfv |
| 1373 | + ): |
| 1374 | + output_features[col] = feature |
| 1375 | + return output_features |
| 1376 | + |
| 1377 | + @staticmethod |
| 1378 | + def _map_in_arrow( |
| 1379 | + iterator: Iterable[pyarrow.RecordBatch], |
| 1380 | + odfv_protos: List[Tuple[bytes, bytes]], |
| 1381 | + full_feature_names: bool, |
| 1382 | + output_columns: List[str], |
| 1383 | + output_pa_types: Dict[str, pyarrow.DataType], |
| 1384 | + ) -> Iterator[pyarrow.RecordBatch]: |
| 1385 | + for batch in iterator: |
| 1386 | + features_table = pyarrow.Table.from_batches([batch]) |
| 1387 | + |
| 1388 | + odfvs = [] |
| 1389 | + for serialized_odfv, serialized_projection in odfv_protos: |
| 1390 | + odfv_proto = OnDemandFeatureViewProto.FromString(serialized_odfv) |
| 1391 | + projection_proto = FeatureViewProjectionProto.FromString( |
| 1392 | + serialized_projection |
| 1393 | + ) |
| 1394 | + odfvs.append( |
| 1395 | + OnDemandFeatureView.from_proto(odfv_proto).with_projection( |
| 1396 | + FeatureViewProjection.from_proto(projection_proto) |
| 1397 | + ) |
| 1398 | + ) |
| 1399 | + |
| 1400 | + for odfv in odfvs: |
| 1401 | + transformed_arrow = odfv.transform_arrow( |
| 1402 | + features_table, full_feature_names |
1383 | 1403 | ) |
1384 | | - ) |
1385 | 1404 |
|
1386 | | - for odfv in odfvs: |
1387 | | - transformed_arrow = odfv.transform_arrow(features_table, full_feature_names) |
1388 | | - |
1389 | | - for col in transformed_arrow.column_names: |
1390 | | - if col.startswith("__index"): |
1391 | | - continue |
1392 | | - if col in output_columns: |
1393 | | - column = transformed_arrow[col] |
1394 | | - if col in output_pa_types: |
1395 | | - column = column.cast(output_pa_types[col]) |
1396 | | - features_table = features_table.append_column(col, column) |
1397 | | - |
1398 | | - # mapInArrow requires the declared column order. |
1399 | | - yield from features_table.select( |
1400 | | - batch.schema.names + output_columns |
1401 | | - ).to_batches() |
| 1405 | + for col in transformed_arrow.column_names: |
| 1406 | + if col in output_columns: |
| 1407 | + column = transformed_arrow[col] |
| 1408 | + if col in output_pa_types: |
| 1409 | + column = column.cast(output_pa_types[col]) |
| 1410 | + features_table = features_table.append_column(col, column) |
| 1411 | + |
| 1412 | + # mapInArrow requires the declared column order. |
| 1413 | + yield from features_table.select( |
| 1414 | + batch.schema.names + output_columns |
| 1415 | + ).to_batches() |
1402 | 1416 |
|
1403 | 1417 |
|
1404 | 1418 | def _register_iceberg_source_as_temp_view( |
|