FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat: Retrieve historical features from BigQuery without entity_df by singhhimanshu0811 · Pull Request #6569 · feast-dev/feast · GitHub

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

Filter by extension

Filter by extension .py  (3) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
163 changes: 131 additions & 32 deletions sdk/python/feast/infra/offline_stores/bigquery.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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

Back | FazBrowse Home | New Git URL