Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 131 additions & 32 deletions sdk/python/feast/infra/offline_stores/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,13 @@
ContextManager,
Dict,
Iterator,
KeysView,
List,
Literal,
Optional,
Tuple,
Union,
cast,
)

import numpy as np
Expand Down Expand Up @@ -57,7 +59,7 @@
from feast.on_demand_feature_view import OnDemandFeatureView
from feast.repo_config import FeastConfigBaseModel, RepoConfig
from feast.saved_dataset import SavedDatasetStorage
from feast.utils import _utc_now, get_user_agent
from feast.utils import _utc_now, compute_non_entity_date_range, get_user_agent

from .bigquery_source import (
BigQueryLoggingDestination,
Expand Down Expand Up @@ -267,10 +269,11 @@ def get_historical_features(
config: RepoConfig,
feature_views: List[FeatureView],
feature_refs: List[str],
entity_df: Union[pd.DataFrame, str],
entity_df: Optional[Union[pd.DataFrame, str]],
registry: BaseRegistry,
project: str,
full_feature_names: bool = False,
**kwargs: Any,
) -> RetrievalJob:
# TODO: Add entity_df validation in order to fail before interacting with BigQuery
assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)
Expand All @@ -297,60 +300,100 @@ def get_historical_features(
config.offline_store.table_create_disposition,
)

entity_schema = _get_entity_schema(
client=client,
entity_df=entity_df,
)
# Non-entity mode: create a left temporary table from entity keys - any entity key having an event in the time window

entity_df_event_timestamp_col = (
offline_utils.infer_event_timestamp_from_entity_df(entity_schema)
)
non_entity_mode = entity_df is None

entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(
entity_df,
entity_df_event_timestamp_col,
client,
)
if non_entity_mode:
start_date, end_date = compute_non_entity_date_range(
feature_views,
start_date=kwargs.get("start_date"),
end_date=kwargs.get("end_date"),
)
entity_df_event_timestamp_range = (start_date, end_date)

@contextlib.contextmanager
def query_generator() -> Iterator[str]:
_upload_entity_df(
# Pre-compute query contexts to collect entity column names per feature view.
fv_query_contexts_pre = offline_utils.get_feature_view_query_context(
Comment thread
singhhimanshu0811 marked this conversation as resolved.
feature_refs,
feature_views,
registry,
project,
entity_df_event_timestamp_range,
)
all_entities = offline_utils.gather_all_entities(fv_query_contexts_pre)
event_timestamp_col = "entity_ts"
entity_schema_keys: KeysView[str] = cast(
KeysView[str],
{k: None for k in (all_entities + [event_timestamp_col])}.keys(),
)

entity_schema = None
else:
entity_schema = _get_entity_schema(
client=client,
table_name=table_reference,
entity_df=entity_df,
)

expected_join_keys = offline_utils.get_expected_join_keys(
project, feature_views, registry
event_timestamp_col = offline_utils.infer_event_timestamp_from_entity_df(
entity_schema
)

offline_utils.assert_expected_columns_in_entity_df(
entity_schema, expected_join_keys, entity_df_event_timestamp_col
entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(
entity_df,
event_timestamp_col,
client,
)

# Build a query context containing all information required to template the BigQuery SQL query
query_context = offline_utils.get_feature_view_query_context(
entity_schema_keys = entity_schema.keys()
all_entities = []
start_date = entity_df_event_timestamp_range[0]
end_date = entity_df_event_timestamp_range[1]
fv_query_contexts_pre = offline_utils.get_feature_view_query_context(
feature_refs,
feature_views,
registry,
project,
entity_df_event_timestamp_range,
)

@contextlib.contextmanager
def query_generator() -> Iterator[str]:
if non_entity_mode:
_bq_create_entity_union_table(
client=client,
table_name=table_reference,
feature_views=feature_views,
fv_query_contexts=fv_query_contexts_pre,
start_date=start_date,
end_date=end_date,
all_entities=all_entities,
event_timestamp_col=event_timestamp_col,
)
else:
_upload_entity_df(
client=client,
table_name=table_reference,
entity_df=entity_df,
)
expected_join_keys = offline_utils.get_expected_join_keys(
project, feature_views, registry
)
assert entity_schema is not None
offline_utils.assert_expected_columns_in_entity_df(
entity_schema, expected_join_keys, event_timestamp_col
)

# Generate the BigQuery SQL query from the query context
query = offline_utils.build_point_in_time_query(
query_context,
feature_view_query_contexts=fv_query_contexts_pre,
left_table_query_string=table_reference,
entity_df_event_timestamp_col=entity_df_event_timestamp_col,
entity_df_columns=entity_schema.keys(),
entity_df_event_timestamp_col=event_timestamp_col,
entity_df_columns=entity_schema_keys,
query_template=MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN,
full_feature_names=full_feature_names,
)

try:
yield query
finally:
# Asynchronously clean up the uploaded Bigquery table, which will expire
# Asynchronously clean up the uploaded BigQuery table, which will expire
# if cleanup fails
client.delete_table(table=table_reference, not_found_ok=True)

Expand All @@ -364,7 +407,7 @@ def query_generator() -> Iterator[str]:
),
metadata=RetrievalMetadata(
features=feature_refs,
keys=list(entity_schema.keys() - {entity_df_event_timestamp_col}),
keys=list(set(entity_schema_keys) - {event_timestamp_col}),
min_event_timestamp=entity_df_event_timestamp_range[0],
max_event_timestamp=entity_df_event_timestamp_range[1],
),
Expand Down Expand Up @@ -567,6 +610,62 @@ def clear_monitoring_baseline(
)


def _bq_create_entity_union_table(
client: "Client",
table_name: str,
feature_views: List[FeatureView],
fv_query_contexts: List[offline_utils.FeatureViewQueryContext],
start_date: datetime,
end_date: datetime,
all_entities: List[str],
event_timestamp_col: str,
) -> None:
"""
Creates a BigQuery temp table containing the UNION DISTINCT of entity keys observed
across all feature views in [start_date, end_date], plus a stable as-of timestamp
column set to end_date. Used as the left table for PIT joins in non-entity mode.
"""
start_str = start_date.strftime("%Y-%m-%dT%H:%M:%S")
end_str = end_date.strftime("%Y-%m-%dT%H:%M:%S")

per_view_selects: List[str] = []
for fv, ctx in zip(feature_views, fv_query_contexts):
assert isinstance(fv.batch_source, BigQuerySource)
from_expression = fv.batch_source.get_table_query_string()
timestamp_field = ctx.timestamp_field

ctx_entities_set = set(ctx.entities)
select_entities: List[str] = []
for col in all_entities:
if col in ctx_entities_set:
select_entities.append(f"`{col}`")
else:
select_entities.append(f"CAST(NULL AS STRING) AS `{col}`")

per_view_selects.append(
f"SELECT DISTINCT {', '.join(select_entities)} "
f"FROM {from_expression} "
f"WHERE `{timestamp_field}` BETWEEN TIMESTAMP('{start_str}') AND TIMESTAMP('{end_str}')"
)

union_query = "\nUNION DISTINCT\n".join(per_view_selects)
entity_cols = (
", ".join(f"`{e}`" for e in all_entities) if all_entities else "TRUE AS _dummy"
)

create_sql = (
f"CREATE TABLE `{table_name}` AS "
f"SELECT {entity_cols}, TIMESTAMP('{end_str}') AS `{event_timestamp_col}` "
f"FROM ({union_query}) AS _entity_union"
)

block_until_done(client, client.query(create_sql))

table = client.get_table(table_name)
table.expires = _utc_now() + timedelta(minutes=30)
client.update_table(table, ["expires"])


# ------------------------------------------------------------------ #
# BigQuery monitoring metrics (native)
# ------------------------------------------------------------------ #
Expand Down
9 changes: 9 additions & 0 deletions sdk/python/feast/infra/offline_stores/offline_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -429,3 +429,12 @@ def format_date(val: Union[str, datetime]) -> str:
filters.append(f"{dp_field} <= '{format_date(end_date)}'")

return " AND ".join(filters) if filters else ""


def gather_all_entities(fv_query_contexts: List[FeatureViewQueryContext]):
all_entities: List[str] = []
for ctx in fv_query_contexts:
for e in ctx.entities:
if e not in all_entities:
all_entities.append(e)
return all_entities
Loading
Loading