Skip to content

Commit 01224d0

Browse files
YHC66haoxu0
YHC66
and
haoxu0
authored
fix: Apply point-in-time filter correctly with the Polars backend (#6992)
LocalFilterNode filtered rows with pandas-style boolean indexing (df[mask]). polars treats a boolean Series passed to DataFrame.__getitem__ as a column selection, so with the Polars backend every row was kept (including feature rows newer than the entity timestamp) and columns were dropped instead; with a TTL the node then failed with ColumnNotFoundError. Add DataFrameBackend.filter_by_mask and use it for the entity timestamp and TTL filters. Signed-off-by: Yihang Chen <[email protected]> Co-authored-by: Hao Xu <[email protected]>
1 parent c0545f7 commit 01224d0

5 files changed

Lines changed: 68 additions & 2 deletions

File tree

‎sdk/python/feast/infra/compute_engines/backends/base.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ class DataFrameBackend(ABC):
3838
filter(df: Any, expr: str) -> Any
3939
Apply a filter expression (string-based) to the DataFrame.
4040
41+
filter_by_mask(df: Any, mask: Any) -> Any
42+
Keep the rows of the DataFrame for which the boolean mask is true.
43+
4144
to_timedelta_value(delta: timedelta) -> Any
4245
Convert a Python timedelta object to a backend-compatible value
4346
that can be subtracted from a timestamp column.
@@ -65,6 +68,9 @@ def groupby_agg(self, df, group_keys, agg_ops): ...
6568
@abstractmethod
6669
def filter(self, df, expr): ...
6770

71+
@abstractmethod
72+
def filter_by_mask(self, df, mask): ...
73+
6874
@abstractmethod
6975
def to_arrow(self, df): ...
7076

‎sdk/python/feast/infra/compute_engines/backends/pandas_backend.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ def groupby_agg(self, df, group_keys, agg_ops):
3131
def filter(self, df, expr):
3232
return df.query(expr)
3333

34+
def filter_by_mask(self, df, mask):
35+
return df[mask]
36+
3437
def to_arrow(self, df):
3538
return pa.Table.from_pandas(df)
3639

‎sdk/python/feast/infra/compute_engines/backends/polars_backend.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ def groupby_agg(self, df: pl.DataFrame, group_keys, agg_ops) -> pl.DataFrame:
3131
def filter(self, df: pl.DataFrame, expr: str) -> pl.DataFrame:
3232
return df.filter(pl.sql_expr(expr))
3333

34+
def filter_by_mask(self, df: pl.DataFrame, mask) -> pl.DataFrame:
35+
return df.filter(mask)
36+
3437
def to_timedelta_value(self, delta: timedelta):
3538
return pl.duration(milliseconds=delta.total_seconds() * 1000)
3639

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,14 +138,18 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue:
138138

139139
if ENTITY_TS_ALIAS in self.backend.columns(df):
140140
# filter where feature.ts <= entity.event_timestamp
141-
df = df[df[timestamp_column] <= df[ENTITY_TS_ALIAS]]
141+
df = self.backend.filter_by_mask(
142+
df, df[timestamp_column] <= df[ENTITY_TS_ALIAS]
143+
)
142144

143145
# TTL: feature.ts >= entity.event_timestamp - ttl
144146
if self.ttl:
145147
lower_bound = df[ENTITY_TS_ALIAS] - self.backend.to_timedelta_value(
146148
self.ttl
147149
)
148-
df = df[df[timestamp_column] >= lower_bound]
150+
df = self.backend.filter_by_mask(
151+
df, df[timestamp_column] >= lower_bound
152+
)
149153

150154
# Optional user-defined filter expression (e.g., "value > 0")
151155
if self.filter_expr:

‎sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
LocalOutputNode,
1717
LocalTransformationNode,
1818
)
19+
from feast.infra.compute_engines.utils import ENTITY_TS_ALIAS
1920
from feast.repo_config import MaterializationConfig
2021

2122
backend = PandasBackend()
@@ -78,6 +79,55 @@ def test_local_filter_node():
7879
assert result.data.num_rows == 3
7980

8081

82+
@pytest.mark.parametrize("backend_name", ["pandas", "polars"])
83+
@pytest.mark.parametrize(
84+
"ttl, expected_values",
85+
[(None, [20, 30, 40]), (timedelta(minutes=2), [20, 30])],
86+
)
87+
def test_local_filter_node_point_in_time(backend_name, ttl, expected_values):
88+
if backend_name == "polars":
89+
pytest.importorskip("polars")
90+
from feast.infra.compute_engines.backends.factory import BackendFactory
91+
92+
# Rows as they look after the join with the entity dataframe: one row is
93+
# newer than the entity timestamp and one is older than the 2 minute TTL.
94+
joined_df = pd.DataFrame(
95+
{
96+
"entity_id": [1, 1, 2, 2],
97+
"value": [10, 20, 30, 40],
98+
"event_timestamp": [
99+
now + timedelta(minutes=1),
100+
now - timedelta(minutes=1),
101+
now,
102+
now - timedelta(minutes=5),
103+
],
104+
ENTITY_TS_ALIAS: [now] * 4,
105+
}
106+
)
107+
context = create_context(
108+
node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(joined_df))}
109+
)
110+
111+
filter_node = LocalFilterNode(
112+
name="filter",
113+
backend=BackendFactory.from_name(backend_name),
114+
ttl=ttl,
115+
column_info=ColumnInfo(
116+
join_keys=["entity_id"],
117+
feature_cols=["value"],
118+
ts_col="event_timestamp",
119+
created_ts_col=None,
120+
),
121+
)
122+
filter_node.add_input(MagicMock())
123+
filter_node.inputs[0].name = "source"
124+
125+
result = filter_node.execute(context)
126+
assert isinstance(result, ArrowTableValue)
127+
assert set(joined_df.columns) <= set(result.data.column_names)
128+
assert result.data.column("value").to_pylist() == expected_values
129+
130+
81131
def test_local_aggregation_node():
82132
context = create_context(
83133
node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(sample_df))}

0 commit comments

Comments
 (0)