Skip to content

Commit a952524

Browse files
blaketastic2tchughesiv
blaketastic2
authored and
tchughesiv
committed
fix: Allow get_historical_features with only On Demand Feature View (feast-dev#5256)
Signed-off-by: Blake <blaketastic2@gmail.com>
1 parent 9baea81 commit a952524

3 files changed

Lines changed: 118 additions & 54 deletions

File tree

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ ci = [
142142
"pytest-mock==1.10.4",
143143
"pytest-env",
144144
"Sphinx>4.0.0,<7",
145+
"sqlglot[rs]>=26.12.1",
145146
"testcontainers==4.9.0",
146147
"python-keycloak==4.2.2",
147148
"pre-commit<3.3.2",

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,8 +484,10 @@ def _get_entity_schema(
484484
{% else %}
485485
{{ left_table_query_string }}
486486
{% endif %}
487-
),
488-
487+
)
488+
{% if featureviews | length > 0 %}
489+
,
490+
{% endif %}
489491
{% for featureview in featureviews %}
490492
491493
"{{ featureview.name }}__entity_dataframe" AS (

‎sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py‎

Lines changed: 113 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,10 @@
33
from unittest.mock import MagicMock, patch
44

55
import pandas as pd
6+
import sqlglot
67

78
from feast.entity import Entity
8-
from feast.feature_view import FeatureView, Field
9+
from feast.feature_view import FeatureView, FeatureViewProjection, Field
910
from feast.infra.offline_stores.contrib.postgres_offline_store.postgres import (
1011
PostgreSQLOfflineStore,
1112
PostgreSQLOfflineStoreConfig,
@@ -14,8 +15,9 @@
1415
PostgreSQLSource,
1516
)
1617
from feast.infra.offline_stores.offline_store import RetrievalJob
18+
from feast.on_demand_feature_view import OnDemandFeatureView
1719
from feast.repo_config import RepoConfig
18-
from feast.types import Float32
20+
from feast.types import Float32, ValueType
1921

2022
logging.basicConfig(level=logging.DEBUG)
2123
logger = logging.getLogger(__name__)
@@ -30,15 +32,7 @@ def test_pull_latest_from_table_with_nested_timestamp_or_query(mock_get_conn):
3032
project="test_project",
3133
registry="test_registry",
3234
provider="local",
33-
offline_store=PostgreSQLOfflineStoreConfig(
34-
type="postgres",
35-
host="localhost",
36-
port=5432,
37-
database="test_db",
38-
db_schema="public",
39-
user="test_user",
40-
password="test_password",
41-
),
35+
offline_store=_mock_offline_store_config(),
4236
)
4337

4438
test_data_source = PostgreSQLSource(
@@ -100,15 +94,7 @@ def test_pull_latest_from_table_without_nested_timestamp_or_query(mock_get_conn)
10094
project="test_project",
10195
registry="test_registry",
10296
provider="local",
103-
offline_store=PostgreSQLOfflineStoreConfig(
104-
type="postgres",
105-
host="localhost",
106-
port=5432,
107-
database="test_db",
108-
db_schema="public",
109-
user="test_user",
110-
password="test_password",
111-
),
97+
offline_store=_mock_offline_store_config(),
11298
)
11399

114100
test_data_source = PostgreSQLSource(
@@ -167,15 +153,7 @@ def test_pull_all_from_table_or_query(mock_get_conn):
167153
project="test_project",
168154
registry="test_registry",
169155
provider="local",
170-
offline_store=PostgreSQLOfflineStoreConfig(
171-
type="postgres",
172-
host="localhost",
173-
port=5432,
174-
database="test_db",
175-
db_schema="public",
176-
user="test_user",
177-
password="test_password",
178-
),
156+
offline_store=_mock_offline_store_config(),
179157
)
180158

181159
test_data_source = PostgreSQLSource(
@@ -239,15 +217,7 @@ def test_get_historical_features_entity_select_modes(
239217
project="test_project",
240218
registry="test_registry",
241219
provider="local",
242-
offline_store=PostgreSQLOfflineStoreConfig(
243-
type="postgres",
244-
host="localhost",
245-
port=5432,
246-
database="test_db",
247-
db_schema="public",
248-
user="test_user",
249-
password="test_password",
250-
),
220+
offline_store=_mock_offline_store_config(),
251221
)
252222

253223
test_data_source = PostgreSQLSource(
@@ -259,13 +229,7 @@ def test_get_historical_features_entity_select_modes(
259229

260230
test_feature_view = FeatureView(
261231
name="test_feature_view",
262-
entities=[
263-
Entity(
264-
name="driver_id",
265-
join_keys=["driver_id"],
266-
description="Driver ID",
267-
)
268-
],
232+
entities=_mock_entity(),
269233
schema=[
270234
Field(name="feature1", dtype=Float32),
271235
],
@@ -300,6 +264,9 @@ def test_get_historical_features_entity_select_modes(
300264
all the logic as the field to GROUP BY the data
301265
*/""")
302266

267+
sqlglot.parse(actual_query)
268+
assert True
269+
303270

304271
@patch("feast.infra.offline_stores.contrib.postgres_offline_store.postgres._get_conn")
305272
@patch(
@@ -345,13 +312,7 @@ def test_get_historical_features_entity_select_modes_embed_query(
345312

346313
test_feature_view = FeatureView(
347314
name="test_feature_view",
348-
entities=[
349-
Entity(
350-
name="driver_id",
351-
join_keys=["driver_id"],
352-
description="Driver ID",
353-
)
354-
],
315+
entities=_mock_entity(),
355316
schema=[
356317
Field(name="feature1", dtype=Float32),
357318
],
@@ -388,3 +349,103 @@ def test_get_historical_features_entity_select_modes_embed_query(
388349
assert actual_query.startswith("""WITH
389350
390351
entity_query AS (""")
352+
353+
# Verify the SQL is valid by parsing it
354+
sqlglot.parse(actual_query) # This will raise ParseError if SQL is invalid
355+
assert True # If we get here, the SQL is valid
356+
357+
358+
@patch("feast.infra.offline_stores.contrib.postgres_offline_store.postgres._get_conn")
359+
@patch(
360+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres.df_to_postgres_table"
361+
)
362+
@patch(
363+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres.get_query_schema"
364+
)
365+
@patch("feast.on_demand_feature_view.OnDemandFeatureView.get_requested_odfvs")
366+
def test_get_historical_features_no_feature_view(
367+
mock_get_requested_odfvs,
368+
mock_get_query_schema,
369+
mock_df_to_postgres_table,
370+
mock_get_conn,
371+
):
372+
mock_conn = MagicMock()
373+
mock_get_conn.return_value.__enter__.return_value = mock_conn
374+
375+
# Create a mock OnDemandFeatureView
376+
mock_odfv = MagicMock(spec=OnDemandFeatureView)
377+
mock_odfv.name = "test_odfv"
378+
mock_odfv.features = [Field(name="feature1", dtype=Float32)]
379+
mock_odfv.projection = FeatureViewProjection(
380+
name="test_odfv",
381+
name_alias="test_odfv",
382+
features=[Field(name="feature1", dtype=Float32)],
383+
desired_features=[],
384+
)
385+
mock_get_requested_odfvs.return_value = [mock_odfv]
386+
387+
test_repo_config = RepoConfig(
388+
project="test_project",
389+
registry="test_registry",
390+
provider="local",
391+
offline_store=_mock_offline_store_config(),
392+
)
393+
394+
test_data_source = PostgreSQLSource(
395+
name="test_batch_source",
396+
description="test_batch_source",
397+
table="offline_store_database_name.offline_store_table_name",
398+
timestamp_field="event_published_datetime_utc",
399+
)
400+
401+
test_feature_view = FeatureView(
402+
name="test_feature_view",
403+
entities=_mock_entity(),
404+
schema=[
405+
Field(name="feature1", dtype=Float32),
406+
],
407+
source=test_data_source,
408+
)
409+
410+
mock_registry = MagicMock()
411+
mock_registry.get_on_demand_feature_view.return_value = test_feature_view
412+
mock_registry.list_on_demand_feature_views.return_value = [mock_odfv]
413+
414+
entity_df = pd.DataFrame(
415+
{"event_timestamp": [datetime(2021, 1, 1)], "driver_id": [1]}
416+
)
417+
418+
retrieval_job = PostgreSQLOfflineStore.get_historical_features(
419+
config=test_repo_config,
420+
feature_views=[],
421+
feature_refs=["test_odfv:feature1"],
422+
entity_df=entity_df,
423+
registry=mock_registry,
424+
project="test_project",
425+
)
426+
427+
sqlglot.parse(retrieval_job.to_sql().strip(), dialect="postgres")
428+
assert True
429+
430+
431+
def _mock_offline_store_config():
432+
return PostgreSQLOfflineStoreConfig(
433+
type="postgres",
434+
host="localhost",
435+
port=5432,
436+
database="test_db",
437+
db_schema="public",
438+
user="test_user",
439+
password="test_password",
440+
)
441+
442+
443+
def _mock_entity():
444+
return [
445+
Entity(
446+
name="driver_id",
447+
join_keys=["driver_id"],
448+
description="Driver ID",
449+
value_type=ValueType.INT64,
450+
)
451+
]

0 commit comments

Comments
 (0)