think i got everything working now

franciscojavierarceo · franciscojavierarceo · commit a85854e23d75 · 2024-09-21T13:53:39.000-04:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/inference.py b/sdk/python/feast/inference.py
@@ -204,122 +204,131 @@ def _infer_features_and_entities(
         run_inference_for_features: Whether to run inference for features.
         config: The config for the current feature store.
     """
-    entity_columns: list[Field] = []
     if isinstance(fv, OnDemandFeatureView):
-        columns_to_exclude = set()
-        for (
-            source_feature_view_name,
-            source_feature_view,
-        ) in fv.source_feature_view_projections.items():
-            columns_to_exclude.add(source_feature_view.timestamp_field)
-            columns_to_exclude.add(source_feature_view.created_timestamp_column)
+        return _infer_on_demand_features_and_entities(
+            fv, join_keys, run_inference_for_features, config
+        )
 
-            batch_source = getattr(source_feature_view, "batch_source")
-            batch_field_mapping = getattr(batch_source or None, "field_mapping")
-            if batch_field_mapping:
-                for (
-                    original_col,
-                    mapped_col,
-                ) in batch_field_mapping.items():
-                    if mapped_col in columns_to_exclude:
-                        columns_to_exclude.remove(mapped_col)
-                        columns_to_exclude.add(original_col)
+    columns_to_exclude = {
+        fv.batch_source.timestamp_field,
+        fv.batch_source.created_timestamp_column,
+    }
+    for original_col, mapped_col in fv.batch_source.field_mapping.items():
+        if mapped_col in columns_to_exclude:
+            columns_to_exclude.remove(mapped_col)
+            columns_to_exclude.add(original_col)
 
-                table_column_names_and_types = (
-                    batch_source.get_table_column_names_and_types(config)
-                )
+    table_column_names_and_types = fv.batch_source.get_table_column_names_and_types(
+        config
+    )
 
-            for col_name, col_datatype in table_column_names_and_types:
-                if col_name in columns_to_exclude:
-                    continue
-                elif col_name in join_keys:
-                    field = Field(
-                        name=col_name,
-                        dtype=from_value_type(
-                            batch_source.source_datatype_to_feast_value_type()(
-                                col_datatype
-                            )
-                        ),
-                    )
-                    if field.name not in [
-                        entity_column.name
-                        for entity_column in entity_columns
-                        if hasattr(entity_column, "name")
-                    ]:
-                        entity_columns.append(field)
-                elif not re.match(
-                    "^__|__$", col_name
-                ):  # double underscores often signal an internal-use column
-                    if run_inference_for_features:
-                        feature_name = (
-                            batch_field_mapping[col_name]
-                            if col_name in batch_field_mapping
-                            else col_name
-                        )
-                        field = Field(
-                            name=feature_name,
-                            dtype=from_value_type(
-                                batch_source.source_datatype_to_feast_value_type()(
-                                    col_datatype
-                                )
-                            ),
+    for col_name, col_datatype in table_column_names_and_types:
+        if col_name in columns_to_exclude:
+            continue
+        elif col_name in join_keys:
+            field = Field(
+                name=col_name,
+                dtype=from_value_type(
+                    fv.batch_source.source_datatype_to_feast_value_type()(col_datatype)
+                ),
+            )
+            if field.name not in [
+                entity_column.name for entity_column in fv.entity_columns
+            ]:
+                fv.entity_columns.append(field)
+        elif not re.match(
+            "^__|__$", col_name
+        ):  # double underscores often signal an internal-use column
+            if run_inference_for_features:
+                feature_name = (
+                    fv.batch_source.field_mapping[col_name]
+                    if col_name in fv.batch_source.field_mapping
+                    else col_name
+                )
+                field = Field(
+                    name=feature_name,
+                    dtype=from_value_type(
+                        fv.batch_source.source_datatype_to_feast_value_type()(
+                            col_datatype
                         )
-                        if field.name not in [
-                            feature.name for feature in source_feature_view.features
-                        ]:
-                            source_feature_view.features.append(field)
+                    ),
+                )
+                if field.name not in [feature.name for feature in fv.features]:
+                    fv.features.append(field)
 
-    else:
-        columns_to_exclude = {
-            fv.batch_source.timestamp_field,
-            fv.batch_source.created_timestamp_column,
-        }
-        for original_col, mapped_col in fv.batch_source.field_mapping.items():
-            if mapped_col in columns_to_exclude:
-                columns_to_exclude.remove(mapped_col)
-                columns_to_exclude.add(original_col)
 
-        table_column_names_and_types = fv.batch_source.get_table_column_names_and_types(
-            config
-        )
+def _infer_on_demand_features_and_entities(
+    fv: OnDemandFeatureView,
+    join_keys: Set[Optional[str]],
+    run_inference_for_features,
+    config,
+) -> None:
+    """
+    Updates the specific feature in place with inferred features and entities.
+    Args:
+        fv: The feature view on which to run inference.
+        join_keys: The set of join keys for the feature view's entities.
+        run_inference_for_features: Whether to run inference for features.
+        config: The config for the current feature store.
+    """
+    entity_columns: list[Field] = []
+    columns_to_exclude = set()
+    for (
+        source_feature_view_name,
+        source_feature_view,
+    ) in fv.source_feature_view_projections.items():
+        columns_to_exclude.add(source_feature_view.timestamp_field)
+        columns_to_exclude.add(source_feature_view.created_timestamp_column)
+
+        batch_source = getattr(source_feature_view, "batch_source")
+        batch_field_mapping = getattr(batch_source or None, "field_mapping")
+        if batch_field_mapping:
+            for (
+                original_col,
+                mapped_col,
+            ) in batch_field_mapping.items():
+                if mapped_col in columns_to_exclude:
+                    columns_to_exclude.remove(mapped_col)
+                    columns_to_exclude.add(original_col)
 
+            table_column_names_and_types = (
+                batch_source.get_table_column_names_and_types(config)
+            )
         for col_name, col_datatype in table_column_names_and_types:
             if col_name in columns_to_exclude:
                 continue
             elif col_name in join_keys:
                 field = Field(
                     name=col_name,
                     dtype=from_value_type(
-                        fv.batch_source.source_datatype_to_feast_value_type()(
-                            col_datatype
-                        )
+                        batch_source.source_datatype_to_feast_value_type()(col_datatype)
                     ),
                 )
                 if field.name not in [
                     entity_column.name
-                    if not isinstance(entity_column, str)
-                    else entity_column
                     for entity_column in entity_columns
+                    if hasattr(entity_column, "name")
                 ]:
                     entity_columns.append(field)
             elif not re.match(
                 "^__|__$", col_name
             ):  # double underscores often signal an internal-use column
                 if run_inference_for_features:
                     feature_name = (
-                        fv.batch_source.field_mapping[col_name]
-                        if col_name in fv.batch_source.field_mapping
+                        batch_field_mapping[col_name]
+                        if col_name in batch_field_mapping
                         else col_name
                     )
                     field = Field(
                         name=feature_name,
                         dtype=from_value_type(
-                            fv.batch_source.source_datatype_to_feast_value_type()(
+                            batch_source.source_datatype_to_feast_value_type()(
                                 col_datatype
                             )
                         ),
                     )
-                    if field.name not in [feature.name for feature in fv.features]:
-                        fv.features.append(field)
-
+                    if field.name not in [
+                        feature.name for feature in source_feature_view.features
+                    ]:
+                        source_feature_view.features.append(field)
     fv.entity_columns = entity_columns
