From 35dca232a09a12edeceb6a898407b0defb90d07c Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Mon, 20 Jul 2026 20:31:31 +0530 Subject: [PATCH] fix: prefer UDF source over dill; Postgres SELECT *; fresh online registry reads Rehydrate BatchFeatureView/ODFV callables from body_text (strip leading decorators) before dill.loads to avoid Spark driver exit 139 and cross-Python serve failures. Treat empty Postgres feature_name_columns as SELECT *. Disable registry cache on the online request path so FeatureView state gates see AVAILABLE_ONLINE immediately after materialize. Co-authored-by: Cursor --- sdk/python/feast/feature_server.py | 4 +- .../compute_engines/spark/feature_builder.py | 5 +- .../infra/compute_engines/spark/nodes.py | 25 +++- .../postgres_offline_store/postgres.py | 56 +++++++-- .../transformation/pandas_transformation.py | 13 +- .../transformation/python_transformation.py | 13 +- .../transformation/ray_transformation.py | 14 ++- .../feast/transformation/udf_rehydrate.py | 115 ++++++++++++++++++ sdk/python/feast/utils.py | 7 +- .../postgres_offline_store/test_postgres.py | 63 ++++++++++ .../unit/transformation/test_udf_rehydrate.py | 65 ++++++++++ 11 files changed, 350 insertions(+), 30 deletions(-) create mode 100644 sdk/python/feast/transformation/udf_rehydrate.py create mode 100644 sdk/python/tests/unit/transformation/test_udf_rehydrate.py diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py index bba91130db3..e97665f3c61 100644 --- a/sdk/python/feast/feature_server.py +++ b/sdk/python/feast/feature_server.py @@ -225,7 +225,7 @@ async def _get_features( ): if request.feature_service: feature_service = await run_in_threadpool( - store.get_feature_service, request.feature_service, allow_cache=True + store.get_feature_service, request.feature_service, allow_cache=False ) assert_permissions( resource=feature_service, actions=[AuthzedAction.READ_ONLINE] @@ -237,7 +237,7 @@ async def _get_features( store.registry, store.project, request.features, - allow_cache=True, + allow_cache=False, hide_dummy_entity=False, ) for feature_view in all_feature_views: diff --git a/sdk/python/feast/infra/compute_engines/spark/feature_builder.py b/sdk/python/feast/infra/compute_engines/spark/feature_builder.py index 94f29220513..1766f7c29a6 100644 --- a/sdk/python/feast/infra/compute_engines/spark/feature_builder.py +++ b/sdk/python/feast/infra/compute_engines/spark/feature_builder.py @@ -109,7 +109,10 @@ def build_dedup_node(self, view, input_node): def build_transformation_node(self, view, input_nodes): udf_name = view.feature_transformation.name udf = view.feature_transformation.udf - node = SparkTransformationNode(udf_name, udf, inputs=input_nodes) + udf_string = getattr(view.feature_transformation, "udf_string", "") or "" + node = SparkTransformationNode( + udf_name, udf, inputs=input_nodes, udf_string=udf_string + ) self.nodes.append(node) return node diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py index 92964b72bc9..ca76d97cf3c 100644 --- a/sdk/python/feast/infra/compute_engines/spark/nodes.py +++ b/sdk/python/feast/infra/compute_engines/spark/nodes.py @@ -38,6 +38,7 @@ create_offline_store_retrieval_job, infer_entity_timestamp_column, ) +from feast.transformation.udf_rehydrate import resolve_udf from feast.infra.offline_stores.contrib.spark_offline_store.spark import ( SparkRetrievalJob, _get_entity_schema, @@ -602,9 +603,29 @@ def execute(self, context: ExecutionContext) -> DAGValue: class SparkTransformationNode(DAGNode): - def __init__(self, name: str, udf: Callable, inputs: List[DAGNode]): + def __init__( + self, + name: str, + udf: Callable, + inputs: List[DAGNode], + udf_string: str = "", + ): super().__init__(name, inputs) self.udf = udf + self.udf_string = udf_string or "" + + def _resolve_udf(self) -> Callable: + """Prefer source reconstruction over dill callables. + + Dill-deserialized functions that call DataFrame.withColumn / __getitem__ + can segfault (exit 139) on Spark 4.0.1. Re-executing ``udf_string`` + yields a healthy callable. + """ + return resolve_udf( + udf_string=self.udf_string, + fallback_udf=self.udf, + preferred_name=self.name, + ) def execute(self, context: ExecutionContext) -> DAGValue: input_values = self.get_input_values(context) @@ -613,7 +634,7 @@ def execute(self, context: ExecutionContext) -> DAGValue: input_dfs: List[DataFrame] = [val.data for val in input_values] - transformed_df = self.udf(*input_dfs) + transformed_df = self._resolve_udf()(*input_dfs) return DAGValue( data=transformed_df, format=DAGFormat.SPARK, metadata={"transformed": True} diff --git a/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py b/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py index d2bbeb90133..0bd09bc8656 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py +++ b/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py @@ -99,14 +99,39 @@ def pull_latest_from_table_or_query( if created_timestamp_column: timestamps.append(created_timestamp_column) timestamp_desc_string = " DESC, ".join(_append_alias(timestamps, "a")) + " DESC" - a_field_string = ", ".join( - _append_alias(join_key_columns + feature_name_columns + timestamps, "a") - ) - b_field_string = ", ".join( - _append_alias(join_key_columns + feature_name_columns + timestamps, "b") - ) - - query = f""" + # Empty feature_name_columns means "all source columns". BatchFeatureView + # python/pandas/ray transforms signal this via get_column_info. Selecting + # only join keys + timestamps would starve the UDF of input features. + if not feature_name_columns: + distinct_on = ", ".join(f'a."{c}"' for c in join_key_columns) or ( + f'a."{timestamp_field}"' + ) + order_by_parts = [f'a."{c}"' for c in join_key_columns] + [ + f'a."{timestamp_field}" DESC' + ] + if created_timestamp_column: + order_by_parts.append(f'a."{created_timestamp_column}" DESC') + order_by = ", ".join(order_by_parts) + query = f""" + SELECT DISTINCT ON ({distinct_on}) + a.* + {f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""} + FROM {from_expression} a + WHERE a."{timestamp_field}" BETWEEN '{start_date}'::timestamptz AND '{end_date}'::timestamptz + ORDER BY {order_by} + """ + else: + a_field_string = ", ".join( + _append_alias( + join_key_columns + feature_name_columns + timestamps, "a" + ) + ) + b_field_string = ", ".join( + _append_alias( + join_key_columns + feature_name_columns + timestamps, "b" + ) + ) + query = f""" SELECT {b_field_string} {f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""} @@ -271,12 +296,17 @@ def pull_all_from_table_or_query( timestamp_fields = [timestamp_field] if created_timestamp_column: timestamp_fields.append(created_timestamp_column) - field_string = ", ".join( - _append_alias( - join_key_columns + feature_name_columns + timestamp_fields, - "paftoq_alias", + # Empty feature_name_columns => SELECT * (BatchFeatureView python mode). + # Default materialization uses pull_all (pull_latest_features=False). + if not feature_name_columns: + field_string = "paftoq_alias.*" + else: + field_string = ", ".join( + _append_alias( + join_key_columns + feature_name_columns + timestamp_fields, + "paftoq_alias", + ) ) - ) timestamp_filter = get_timestamp_filter_sql( start_date, diff --git a/sdk/python/feast/transformation/pandas_transformation.py b/sdk/python/feast/transformation/pandas_transformation.py index 6e073c30100..c0ecb69b5d8 100644 --- a/sdk/python/feast/transformation/pandas_transformation.py +++ b/sdk/python/feast/transformation/pandas_transformation.py @@ -1,7 +1,6 @@ import inspect from typing import Any, Callable, Optional, cast, get_type_hints -import dill import pandas as pd import pyarrow @@ -146,7 +145,15 @@ def __eq__(self, other): @classmethod def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto): + from feast.transformation.udf_rehydrate import resolve_udf + + udf_string = user_defined_function_proto.body_text or "" + udf = resolve_udf( + udf_string=udf_string, + body=user_defined_function_proto.body or None, + preferred_name=user_defined_function_proto.name or None, + ) return PandasTransformation( - udf=dill.loads(user_defined_function_proto.body), - udf_string=user_defined_function_proto.body_text, + udf=udf, + udf_string=udf_string, ) diff --git a/sdk/python/feast/transformation/python_transformation.py b/sdk/python/feast/transformation/python_transformation.py index 68e9eee95f6..2153b93fbdb 100644 --- a/sdk/python/feast/transformation/python_transformation.py +++ b/sdk/python/feast/transformation/python_transformation.py @@ -1,7 +1,6 @@ from types import FunctionType from typing import Any, Dict, Optional, cast -import dill import pyarrow from feast.field import Field, from_value_type @@ -164,7 +163,15 @@ def __reduce__(self): @classmethod def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto): + from feast.transformation.udf_rehydrate import resolve_udf + + udf_string = user_defined_function_proto.body_text or "" + udf = resolve_udf( + udf_string=udf_string, + body=user_defined_function_proto.body or None, + preferred_name=user_defined_function_proto.name or None, + ) return PythonTransformation( - udf=dill.loads(user_defined_function_proto.body), - udf_string=user_defined_function_proto.body_text, + udf=udf, + udf_string=udf_string, ) diff --git a/sdk/python/feast/transformation/ray_transformation.py b/sdk/python/feast/transformation/ray_transformation.py index b592ce8b0d7..3c20a891a19 100644 --- a/sdk/python/feast/transformation/ray_transformation.py +++ b/sdk/python/feast/transformation/ray_transformation.py @@ -1,8 +1,6 @@ import inspect from typing import Any, Callable, Optional, cast, get_type_hints -import dill - from feast.field import Field, from_value_type from feast.protos.feast.core.Transformation_pb2 import ( UserDefinedFunctionV2 as UserDefinedFunctionProto, @@ -284,9 +282,17 @@ def __eq__(self, other): @classmethod def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto): + from feast.transformation.udf_rehydrate import resolve_udf + + udf_string = user_defined_function_proto.body_text or "" + udf = resolve_udf( + udf_string=udf_string, + body=user_defined_function_proto.body or None, + preferred_name=user_defined_function_proto.name or None, + ) return RayTransformation( - udf=dill.loads(user_defined_function_proto.body), - udf_string=user_defined_function_proto.body_text, + udf=udf, + udf_string=udf_string, name=user_defined_function_proto.name if user_defined_function_proto.name else None, diff --git a/sdk/python/feast/transformation/udf_rehydrate.py b/sdk/python/feast/transformation/udf_rehydrate.py new file mode 100644 index 00000000000..72e4de387b8 --- /dev/null +++ b/sdk/python/feast/transformation/udf_rehydrate.py @@ -0,0 +1,115 @@ +"""Rehydrate trusted UDF callables from registry source text. + +Feast persists both dill bytes (``body``) and source (``body_text`` / ``udf_string``). +Preferring source avoids: + +* Spark driver segfaults (exit 139) when dill-restored callables touch DataFrames +* Cross-Python-version serve failures when apply and feature-server differ +""" + +from __future__ import annotations + +import logging +import re +from typing import Callable, Optional + +import dill + +logger = logging.getLogger(__name__) + +# Leading decorators (e.g. @on_demand_feature_view(...)) are often included in +# dill.source.getsource() output. They are not needed to recover the callable and +# usually reference names unavailable at rehydrate time. +_DECORATOR_BLOCK = re.compile( + r"^(?:\s*@[\s\S]*?\n)+(?=\s*(?:async\s+)?def\s+)", + re.MULTILINE, +) + + +def _strip_leading_decorators(udf_string: str) -> str: + stripped = _DECORATOR_BLOCK.sub("", udf_string.lstrip(), count=1) + return stripped if stripped.strip() else udf_string + + +def _exec_namespace() -> dict: + """Minimal globals so typical UDF source can exec without full feast apply ctx.""" + ns: dict = {"__name__": "feast_udf_rehydrate"} + try: + import pandas as pd + + ns["pd"] = pd + ns["pandas"] = pd + except ImportError: + pass + try: + import numpy as np + + ns["np"] = np + ns["numpy"] = np + except ImportError: + pass + return ns + + +def rehydrate_udf_from_source( + udf_string: str, + *, + preferred_name: Optional[str] = None, +) -> Optional[Callable]: + """Exec ``udf_string`` and return the resulting callable, or None. + + ``udf_string`` is treated as trusted registry content written by ``feast apply``. + Returns None on failure so callers can fall back to dill. + """ + if not (udf_string or "").strip(): + return None + + source = _strip_leading_decorators(udf_string) + ns = _exec_namespace() + try: + exec(source, ns) # noqa: S102 — trusted registry source + except Exception as e: + logger.debug("udf source rehydrate failed: %s", e) + return None + + if preferred_name and preferred_name in ns and callable(ns[preferred_name]): + return ns[preferred_name] + + for value in ns.values(): + if not callable(value): + continue + name = getattr(value, "__name__", None) + if name in (None, "", "__build_class__"): + continue + # Skip imported modules / classes we seeded + if name in ("DataFrame",): + continue + return value + + return None + + +def resolve_udf( + *, + udf_string: str = "", + body: Optional[bytes] = None, + fallback_udf: Optional[Callable] = None, + preferred_name: Optional[str] = None, +) -> Callable: + """Resolve a UDF: source first, then ``fallback_udf``, then dill ``body``.""" + rehydrated = rehydrate_udf_from_source( + udf_string, preferred_name=preferred_name + ) + if rehydrated is not None: + return rehydrated + + if fallback_udf is not None: + return fallback_udf + + if body: + return dill.loads(body) + + raise ValueError( + "Cannot resolve UDF: empty udf_string/body_text and no dill body or " + "fallback callable. Re-run feast apply so body_text is persisted." + ) diff --git a/sdk/python/feast/utils.py b/sdk/python/feast/utils.py index 4c10c1903e5..72ab13b6292 100644 --- a/sdk/python/feast/utils.py +++ b/sdk/python/feast/utils.py @@ -1315,7 +1315,10 @@ def _get_online_request_context( ): from feast.feature_view import DUMMY_ENTITY_NAME - _feature_refs = _get_features(registry, project, features, allow_cache=True) + # Online serve must see fresh FeatureViewState (AVAILABLE_ONLINE vs + # MATERIALIZING/GENERATED). Cached registry reads cause brief or stuck 500s + # after remote materialize completes. + _feature_refs = _get_features(registry, project, features, allow_cache=False) ( requested_feature_views, @@ -1324,7 +1327,7 @@ def _get_online_request_context( registry=registry, project=project, features=features, - allow_cache=True, + allow_cache=False, hide_dummy_entity=False, ) diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py b/sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py index 3eb01a0e6d0..3fb4c84c60c 100644 --- a/sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py @@ -194,6 +194,69 @@ def test_pull_all_from_table_or_query(mock_get_conn): assert actual_query == expected_query +@patch("feast.infra.offline_stores.contrib.postgres_offline_store.postgres._get_conn") +def test_pull_all_empty_feature_cols_selects_star(mock_get_conn): + """Empty feature_name_columns means all source columns (BFV transform modes).""" + mock_conn = MagicMock() + mock_get_conn.return_value.__enter__.return_value = mock_conn + + test_repo_config = RepoConfig( + project="test_project", + registry="test_registry", + provider="local", + offline_store=_mock_offline_store_config(), + ) + test_data_source = PostgreSQLSource( + name="test_batch_source", + table="offline_store_database_name.offline_store_table_name", + timestamp_field="event_timestamp", + ) + retrieval_job = PostgreSQLOfflineStore.pull_all_from_table_or_query( + config=test_repo_config, + data_source=test_data_source, + join_key_columns=["entity_id"], + feature_name_columns=[], + timestamp_field="event_timestamp", + start_date=datetime(2021, 1, 1, tzinfo=timezone.utc), + end_date=datetime(2021, 1, 2, tzinfo=timezone.utc), + ) + sql = retrieval_job.to_sql() + assert "paftoq_alias.*" in sql + assert 'paftoq_alias."entity_id"' not in sql + + +@patch("feast.infra.offline_stores.contrib.postgres_offline_store.postgres._get_conn") +def test_pull_latest_empty_feature_cols_selects_star(mock_get_conn): + mock_conn = MagicMock() + mock_get_conn.return_value.__enter__.return_value = mock_conn + + test_repo_config = RepoConfig( + project="test_project", + registry="test_registry", + provider="local", + offline_store=_mock_offline_store_config(), + ) + test_data_source = PostgreSQLSource( + name="test_batch_source", + table="offline_store_database_name.offline_store_table_name", + timestamp_field="event_timestamp", + ) + retrieval_job = PostgreSQLOfflineStore.pull_latest_from_table_or_query( + config=test_repo_config, + data_source=test_data_source, + join_key_columns=["entity_id"], + feature_name_columns=[], + timestamp_field="event_timestamp", + created_timestamp_column=None, + start_date=datetime(2021, 1, 1, tzinfo=timezone.utc), + end_date=datetime(2021, 1, 2, tzinfo=timezone.utc), + ) + sql = retrieval_job.to_sql() + assert "DISTINCT ON" in sql + assert "a.*" in sql + assert 'ORDER BY a."entity_id", a."event_timestamp" DESC' in sql + + @patch("feast.infra.offline_stores.contrib.postgres_offline_store.postgres._get_conn") @patch( "feast.infra.offline_stores.contrib.postgres_offline_store.postgres.df_to_postgres_table" diff --git a/sdk/python/tests/unit/transformation/test_udf_rehydrate.py b/sdk/python/tests/unit/transformation/test_udf_rehydrate.py new file mode 100644 index 00000000000..c8838b3d1e6 --- /dev/null +++ b/sdk/python/tests/unit/transformation/test_udf_rehydrate.py @@ -0,0 +1,65 @@ +import dill +import pandas as pd + +from feast.protos.feast.core.Transformation_pb2 import ( + UserDefinedFunctionV2 as UserDefinedFunctionProto, +) +from feast.transformation.pandas_transformation import PandasTransformation +from feast.transformation.udf_rehydrate import rehydrate_udf_from_source, resolve_udf + + +def _sample_udf(df: pd.DataFrame) -> pd.DataFrame: + out = pd.DataFrame() + out["doubled"] = df["x"] * 2 + return out + + +_SAMPLE_SRC = ''' +def _sample_udf(df): + import pandas as pd + out = pd.DataFrame() + out["doubled"] = df["x"] * 2 + return out +''' + + +def test_rehydrate_udf_from_source_prefers_named_function(): + fn = rehydrate_udf_from_source(_SAMPLE_SRC, preferred_name="_sample_udf") + assert fn is not None + result = fn(pd.DataFrame({"x": [1, 2]})) + assert list(result["doubled"]) == [2, 4] + + +def test_resolve_udf_prefers_source_over_garbage_dill_body(): + # Garbage dill body would fail dill.loads; source must win. + udf = resolve_udf( + udf_string=_SAMPLE_SRC, + body=b"not-valid-dill", + preferred_name="_sample_udf", + ) + result = udf(pd.DataFrame({"x": [3]})) + assert result["doubled"].iloc[0] == 6 + + +def test_resolve_udf_falls_back_to_dill_when_source_empty(): + body = dill.dumps(_sample_udf, recurse=True) + udf = resolve_udf(udf_string="", body=body) + result = udf(pd.DataFrame({"x": [4]})) + assert result["doubled"].iloc[0] == 8 + + +def test_rehydrate_strips_on_demand_decorator(): + src = '''@on_demand_feature_view( + sources=[feature_view_1], + schema=[Field(name="metric_sum", dtype=Float64)], +) +def metric_sum_odfv(inputs): + import pandas as pd + df = pd.DataFrame() + df["metric_sum"] = inputs["metric_a"] + inputs["metric_b"] + return df +''' + fn = rehydrate_udf_from_source(src, preferred_name="metric_sum_odfv") + assert fn is not None + result = fn(pd.DataFrame({"metric_a": [1.0], "metric_b": [2.0]})) + assert result["metric_sum"].iloc[0] == 3.0