Skip to content

Commit c39130f

Browse files
author
Himanshu Singh
committed
linter errors
1 parent ca988b9 commit c39130f

3 files changed

Lines changed: 63 additions & 29 deletions

File tree

sdk/python/feast/infra/offline_stores/bigquery.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -369,6 +369,7 @@ def query_generator() -> Iterator[str]:
369369
expected_join_keys = offline_utils.get_expected_join_keys(
370370
project, feature_views, registry
371371
)
372+
assert entity_schema is not None
372373
offline_utils.assert_expected_columns_in_entity_df(
373374
entity_schema, expected_join_keys, event_timestamp_col
374375
)

sdk/python/feast/infra/offline_stores/offline_utils.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -431,8 +431,6 @@ def format_date(val: Union[str, datetime]) -> str:
431431
return " AND ".join(filters) if filters else ""
432432

433433

434-
435-
436434
def gather_all_entities(fv_query_contexts: List[FeatureViewQueryContext]):
437435
all_entities: List[str] = []
438436
for ctx in fv_query_contexts:

sdk/python/tests/unit/infra/offline_stores/test_bigquery_non_entity_mode.py

Lines changed: 62 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,12 @@
1818
BigQueryOfflineStoreConfig,
1919
BigQueryRetrievalJob,
2020
_bq_create_entity_union_table,
21-
2221
)
23-
from feast.infra.offline_stores.offline_utils import gather_all_entities
2422
from feast.infra.offline_stores.bigquery_source import BigQuerySource
25-
from feast.infra.offline_stores.offline_utils import FeatureViewQueryContext
23+
from feast.infra.offline_stores.offline_utils import (
24+
FeatureViewQueryContext,
25+
gather_all_entities,
26+
)
2627
from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
2728
from feast.repo_config import RepoConfig
2829

@@ -145,9 +146,7 @@ def _make_client(self) -> MagicMock:
145146

146147
@patch("feast.infra.offline_stores.bigquery.block_until_done")
147148
@patch("feast.infra.offline_stores.bigquery._utc_now")
148-
def test_single_view_creates_table_with_correct_sql(
149-
self, mock_utc_now, mock_block
150-
):
149+
def test_single_view_creates_table_with_correct_sql(self, mock_utc_now, mock_block):
151150
mock_utc_now.return_value = END
152151
client = self._make_client()
153152
fv = _make_feature_view_mock("fv1", ["customer_id"])
@@ -271,7 +270,9 @@ def test_timestamp_filter_uses_start_and_end(self, mock_utc_now, mock_block):
271270

272271
sql = client.query.call_args[0][0]
273272
assert "2023-01-01T00:00:00" in sql # start_date
274-
assert "2024-01-01T00:00:00" in sql # end_date (appears twice: WHERE + entity_ts value)
273+
assert (
274+
"2024-01-01T00:00:00" in sql
275+
) # end_date (appears twice: WHERE + entity_ts value)
275276
assert "BETWEEN TIMESTAMP(" in sql
276277

277278
@patch("feast.infra.offline_stores.bigquery.block_until_done")
@@ -319,9 +320,9 @@ def test_sets_table_expiry(self, mock_utc_now, mock_block):
319320
client.update_table.assert_called_once()
320321
updated_table = client.update_table.call_args[0][0]
321322
# expiry should be 30 minutes after _utc_now()
322-
assert updated_table.expires == datetime(2024, 1, 1, tzinfo=timezone.utc) + timedelta(
323-
minutes=30
324-
)
323+
assert updated_table.expires == datetime(
324+
2024, 1, 1, tzinfo=timezone.utc
325+
) + timedelta(minutes=30)
325326

326327

327328
# ---------------------------------------------------------------------------
@@ -338,7 +339,9 @@ def repo_config():
338339
def mock_bq_client():
339340
client = MagicMock()
340341
client.project = "my-project"
341-
client.query.return_value = MagicMock(state="DONE", exception=lambda timeout=None: None)
342+
client.query.return_value = MagicMock(
343+
state="DONE", exception=lambda timeout=None: None
344+
)
342345
client.get_table.return_value = MagicMock()
343346
return client
344347

@@ -348,8 +351,12 @@ class TestGetHistoricalFeaturesNonEntityMode:
348351
@patch("feast.infra.offline_stores.bigquery._utc_now")
349352
@patch("feast.infra.offline_stores.bigquery._bq_create_entity_union_table")
350353
@patch("feast.infra.offline_stores.bigquery._upload_entity_df")
351-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
352-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
354+
@patch(
355+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
356+
)
357+
@patch(
358+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
359+
)
353360
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
354361
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
355362
def test_non_entity_mode_calls_union_table_not_upload(
@@ -403,8 +410,12 @@ def test_non_entity_mode_calls_union_table_not_upload(
403410
@patch("feast.infra.offline_stores.bigquery._utc_now")
404411
@patch("feast.infra.offline_stores.bigquery._bq_create_entity_union_table")
405412
@patch("feast.infra.offline_stores.bigquery._upload_entity_df")
406-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
407-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
413+
@patch(
414+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
415+
)
416+
@patch(
417+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
418+
)
408419
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
409420
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
410421
def test_non_entity_mode_uses_entity_ts_as_timestamp_col(
@@ -455,8 +466,12 @@ def test_non_entity_mode_uses_entity_ts_as_timestamp_col(
455466
@patch("feast.infra.offline_stores.bigquery._utc_now")
456467
@patch("feast.infra.offline_stores.bigquery._bq_create_entity_union_table")
457468
@patch("feast.infra.offline_stores.bigquery._upload_entity_df")
458-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
459-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
469+
@patch(
470+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
471+
)
472+
@patch(
473+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
474+
)
460475
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
461476
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
462477
def test_non_entity_mode_passes_start_and_end_to_union_table(
@@ -507,8 +522,12 @@ def test_non_entity_mode_passes_start_and_end_to_union_table(
507522
@patch("feast.infra.offline_stores.bigquery.block_until_done")
508523
@patch("feast.infra.offline_stores.bigquery._utc_now")
509524
@patch("feast.infra.offline_stores.bigquery._bq_create_entity_union_table")
510-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
511-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
525+
@patch(
526+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
527+
)
528+
@patch(
529+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
530+
)
512531
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
513532
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
514533
def test_non_entity_mode_returns_retrieval_job(
@@ -553,8 +572,12 @@ def test_non_entity_mode_returns_retrieval_job(
553572
@patch("feast.infra.offline_stores.bigquery.block_until_done")
554573
@patch("feast.infra.offline_stores.bigquery._utc_now")
555574
@patch("feast.infra.offline_stores.bigquery._bq_create_entity_union_table")
556-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
557-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
575+
@patch(
576+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
577+
)
578+
@patch(
579+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
580+
)
558581
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
559582
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
560583
def test_non_entity_mode_metadata_excludes_timestamp_col_from_keys(
@@ -610,9 +633,15 @@ class TestGetHistoricalFeaturesEntityDfMode:
610633
@patch("feast.infra.offline_stores.bigquery._get_entity_df_event_timestamp_range")
611634
@patch("feast.infra.offline_stores.bigquery._get_entity_schema")
612635
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_expected_join_keys")
613-
@patch("feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df")
614-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
615-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
636+
@patch(
637+
"feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df"
638+
)
639+
@patch(
640+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
641+
)
642+
@patch(
643+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
644+
)
616645
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
617646
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
618647
def test_entity_df_mode_calls_upload_not_union_table(
@@ -671,9 +700,15 @@ def test_entity_df_mode_calls_upload_not_union_table(
671700
@patch("feast.infra.offline_stores.bigquery._get_entity_df_event_timestamp_range")
672701
@patch("feast.infra.offline_stores.bigquery._get_entity_schema")
673702
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_expected_join_keys")
674-
@patch("feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df")
675-
@patch("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context")
676-
@patch("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query")
703+
@patch(
704+
"feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df"
705+
)
706+
@patch(
707+
"feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
708+
)
709+
@patch(
710+
"feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
711+
)
677712
@patch("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity")
678713
@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
679714
def test_entity_df_sql_string_mode_works(

0 commit comments

Comments
 (0)