Skip to content

Commit 5fd7af7

Browse files
aniketpalucursoragent
aniketpalu
and
cursoragent
authored
fix: UDF/ODFV source rehydrate (+ Postgres / online cache) (#6655)
* 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. RayTransformation.from_proto left unchanged (follow-up). Signed-off-by: Aniket Paluskar <[email protected]> * Lint Signed-off-by: Aniket Paluskar <[email protected]> * fix: drop global allow_cache=False on online path Post-materialize MATERIALIZING gate lag is handled by client retry / registry TTL refresh, not by bypassing the registry cache on every online request. Addresses review feedback on #6655. Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]> * fix: replace UDF decorator strip regex with linear scan Avoid CodeQL ReDoS finding on nested @/newline regex when rehydrating body_text. Add multiline and adversarial @ spam unit tests. Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]> * fix: cast resolve_udf result to FunctionType for mypy Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]> * fix: compare UDF transformations by udf_string when present Source-first from_proto rebuilds callables whose bytecode differs from the live repo function; requiring co_code equality made no-op feast apply rewrite ODFVs and broke universal CLI integration tests. Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]> * lint Signed-off-by: Aniket Paluskar <[email protected]> * fix: align Ray __eq__ with source identity; cache Spark UDF resolve RayTransformation now compares by udf_string when present (same contract as Pandas/Python) so a future source-first from_proto will not break no-op apply. SparkTransformationNode caches resolve_udf so strip+exec runs once per node lifetime. Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]> --------- Signed-off-by: Aniket Paluskar <[email protected]> Co-authored-by: Cursor <[email protected]>
1 parent f771ea4 commit 5fd7af7

12 files changed

Lines changed: 513 additions & 42 deletions

File tree

‎.secrets.baseline‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎sdk/python/feast/infra/compute_engines/spark/feature_builder.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,10 @@ def build_dedup_node(self, view, input_node):
109109
def build_transformation_node(self, view, input_nodes):
110110
udf_name = view.feature_transformation.name
111111
udf = view.feature_transformation.udf
112-
node = SparkTransformationNode(udf_name, udf, inputs=input_nodes)
112+
udf_string = getattr(view.feature_transformation, "udf_string", "") or ""
113+
node = SparkTransformationNode(
114+
udf_name, udf, inputs=input_nodes, udf_string=udf_string
115+
)
113116
self.nodes.append(node)
114117
return node
115118

‎sdk/python/feast/infra/compute_engines/spark/nodes.py‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import (
4646
SparkSource,
4747
)
48+
from feast.transformation.udf_rehydrate import resolve_udf
4849

4950
logger = logging.getLogger(__name__)
5051

@@ -602,9 +603,32 @@ def execute(self, context: ExecutionContext) -> DAGValue:
602603

603604

604605
class SparkTransformationNode(DAGNode):
605-
def __init__(self, name: str, udf: Callable, inputs: List[DAGNode]):
606+
def __init__(
607+
self,
608+
name: str,
609+
udf: Callable,
610+
inputs: List[DAGNode],
611+
udf_string: str = "",
612+
):
606613
super().__init__(name, inputs)
607614
self.udf = udf
615+
self.udf_string = udf_string or ""
616+
self._resolved_udf: Optional[Callable] = None
617+
618+
def _resolve_udf(self) -> Callable:
619+
"""Prefer source reconstruction over dill callables.
620+
621+
Dill-deserialized functions that call DataFrame.withColumn / __getitem__
622+
can segfault (exit 139) on Spark 4.0.1. Re-executing ``udf_string``
623+
yields a healthy callable. Result is cached for the lifetime of the node.
624+
"""
625+
if self._resolved_udf is None:
626+
self._resolved_udf = resolve_udf(
627+
udf_string=self.udf_string,
628+
fallback_udf=self.udf,
629+
preferred_name=self.name,
630+
)
631+
return self._resolved_udf
608632

609633
def execute(self, context: ExecutionContext) -> DAGValue:
610634
input_values = self.get_input_values(context)
@@ -613,7 +637,7 @@ def execute(self, context: ExecutionContext) -> DAGValue:
613637

614638
input_dfs: List[DataFrame] = [val.data for val in input_values]
615639

616-
transformed_df = self.udf(*input_dfs)
640+
transformed_df = self._resolve_udf()(*input_dfs)
617641

618642
return DAGValue(
619643
data=transformed_df, format=DAGFormat.SPARK, metadata={"transformed": True}

‎sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -101,14 +101,35 @@ def pull_latest_from_table_or_query(
101101
if created_timestamp_column:
102102
timestamps.append(created_timestamp_column)
103103
timestamp_desc_string = " DESC, ".join(_append_alias(timestamps, "a")) + " DESC"
104-
a_field_string = ", ".join(
105-
_append_alias(join_key_columns + feature_name_columns + timestamps, "a")
106-
)
107-
b_field_string = ", ".join(
108-
_append_alias(join_key_columns + feature_name_columns + timestamps, "b")
109-
)
110-
111-
query = f"""
104+
# Empty feature_name_columns means "all source columns". BatchFeatureView
105+
# python/pandas/ray transforms signal this via get_column_info. Selecting
106+
# only join keys + timestamps would starve the UDF of input features.
107+
if not feature_name_columns:
108+
distinct_on = ", ".join(f'a."{c}"' for c in join_key_columns) or (
109+
f'a."{timestamp_field}"'
110+
)
111+
order_by_parts = [f'a."{c}"' for c in join_key_columns] + [
112+
f'a."{timestamp_field}" DESC'
113+
]
114+
if created_timestamp_column:
115+
order_by_parts.append(f'a."{created_timestamp_column}" DESC')
116+
order_by = ", ".join(order_by_parts)
117+
query = f"""
118+
SELECT DISTINCT ON ({distinct_on})
119+
a.*
120+
{f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""}
121+
FROM {from_expression} a
122+
WHERE a."{timestamp_field}" BETWEEN '{start_date}'::timestamptz AND '{end_date}'::timestamptz
123+
ORDER BY {order_by}
124+
"""
125+
else:
126+
a_field_string = ", ".join(
127+
_append_alias(join_key_columns + feature_name_columns + timestamps, "a")
128+
)
129+
b_field_string = ", ".join(
130+
_append_alias(join_key_columns + feature_name_columns + timestamps, "b")
131+
)
132+
query = f"""
112133
SELECT
113134
{b_field_string}
114135
{f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""}
@@ -275,12 +296,17 @@ def pull_all_from_table_or_query(
275296
timestamp_fields = [timestamp_field]
276297
if created_timestamp_column:
277298
timestamp_fields.append(created_timestamp_column)
278-
field_string = ", ".join(
279-
_append_alias(
280-
join_key_columns + feature_name_columns + timestamp_fields,
281-
"paftoq_alias",
299+
# Empty feature_name_columns => SELECT * (BatchFeatureView python mode).
300+
# Default materialization uses pull_all (pull_latest_features=False).
301+
if not feature_name_columns:
302+
field_string = "paftoq_alias.*"
303+
else:
304+
field_string = ", ".join(
305+
_append_alias(
306+
join_key_columns + feature_name_columns + timestamp_fields,
307+
"paftoq_alias",
308+
)
282309
)
283-
)
284310

285311
timestamp_filter = get_timestamp_filter_sql(
286312
start_date,

‎sdk/python/feast/transformation/pandas_transformation.py‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import inspect
22
from typing import Any, Callable, Optional, cast, get_type_hints
33

4-
import dill
54
import pandas as pd
65
import pyarrow
76

@@ -134,17 +133,28 @@ def __eq__(self, other):
134133
if not isinstance(other, PandasTransformation):
135134
return False
136135

137-
if (
138-
self.udf_string != other.udf_string
139-
or self.udf.__code__.co_code != other.udf.__code__.co_code
140-
):
141-
return False
136+
# udf_string is the canonical diff identity. Source-first from_proto
137+
# rebuilds a new callable (strip+exec) whose bytecode differs from the
138+
# live repo function even when the source is unchanged — do not require
139+
# co_code equality when both sides have source text.
140+
left = self.udf_string or ""
141+
right = other.udf_string or ""
142+
if left and right:
143+
return left == right
142144

143-
return True
145+
return self.udf.__code__.co_code == other.udf.__code__.co_code
144146

145147
@classmethod
146148
def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto):
149+
from feast.transformation.udf_rehydrate import resolve_udf
150+
151+
udf_string = user_defined_function_proto.body_text or ""
152+
udf = resolve_udf(
153+
udf_string=udf_string,
154+
body=user_defined_function_proto.body or None,
155+
preferred_name=user_defined_function_proto.name or None,
156+
)
147157
return PandasTransformation(
148-
udf=dill.loads(user_defined_function_proto.body),
149-
udf_string=user_defined_function_proto.body_text,
158+
udf=udf,
159+
udf_string=udf_string,
150160
)

‎sdk/python/feast/transformation/python_transformation.py‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from types import FunctionType
22
from typing import Any, Dict, Optional, cast
33

4-
import dill
54
import pyarrow
65

76
from feast.field import Field, from_value_type
@@ -145,13 +144,16 @@ def __eq__(self, other):
145144
if not isinstance(other, PythonTransformation):
146145
return False
147146

148-
if (
149-
self.udf_string != other.udf_string
150-
or self.udf.__code__.co_code != other.udf.__code__.co_code
151-
):
152-
return False
147+
# udf_string is the canonical diff identity. Source-first from_proto
148+
# rebuilds a new callable (strip+exec) whose bytecode differs from the
149+
# live repo function even when the source is unchanged — do not require
150+
# co_code equality when both sides have source text.
151+
left = self.udf_string or ""
152+
right = other.udf_string or ""
153+
if left and right:
154+
return left == right
153155

154-
return True
156+
return self.udf.__code__.co_code == other.udf.__code__.co_code
155157

156158
def __reduce__(self):
157159
"""Support for pickle/dill serialization."""
@@ -162,7 +164,15 @@ def __reduce__(self):
162164

163165
@classmethod
164166
def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto):
167+
from feast.transformation.udf_rehydrate import resolve_udf
168+
169+
udf_string = user_defined_function_proto.body_text or ""
170+
udf = resolve_udf(
171+
udf_string=udf_string,
172+
body=user_defined_function_proto.body or None,
173+
preferred_name=user_defined_function_proto.name or None,
174+
)
165175
return PythonTransformation(
166-
udf=dill.loads(user_defined_function_proto.body),
167-
udf_string=user_defined_function_proto.body_text,
176+
udf=cast(FunctionType, udf),
177+
udf_string=udf_string,
168178
)

‎sdk/python/feast/transformation/ray_transformation.py‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -272,13 +272,14 @@ def __eq__(self, other):
272272
if not isinstance(other, RayTransformation):
273273
return False
274274

275-
if (
276-
self.udf_string != other.udf_string
277-
or self.udf.__code__.co_code != other.udf.__code__.co_code
278-
):
279-
return False
280-
281-
return True
275+
# Match Pandas/Python: udf_string is the canonical diff identity so a
276+
# future source-first from_proto does not break no-op feast apply.
277+
left = self.udf_string or ""
278+
right = other.udf_string or ""
279+
if left and right:
280+
return left == right
281+
282+
return self.udf.__code__.co_code == other.udf.__code__.co_code
282283

283284
@classmethod
284285
def from_proto(cls, user_defined_function_proto: UserDefinedFunctionProto):

0 commit comments

Comments
 (0)