Skip to content

Commit 2e41d03

Browse files
judahrandjudahrand
authored
Pre compute the timestamp range for feature views (feast-dev#2103)
* Pre compute the timestamp range for feature views This allows BigQuery to take advantage of Partitions when querying historical features. Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Correct comparison typo Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Use `isinstance` rather than `type` Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Convert min and max `event_timestamp` to UTC Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Move `_to_naive_utc` to avoid circular import Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com>
1 parent a710cb9 commit 2e41d03

6 files changed

Lines changed: 148 additions & 29 deletions

File tree

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

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22
from datetime import datetime
33
from pathlib import Path
44

5-
import pytz
6-
75
from feast.feature_view import FeatureView
86
from feast.infra.passthrough_provider import PassthroughProvider
97
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
@@ -24,13 +22,6 @@ def _table_id(project: str, table: FeatureView) -> str:
2422
return f"{project}_{table.name}"
2523

2624

27-
def _to_naive_utc(ts: datetime):
28-
if ts.tzinfo is None:
29-
return ts
30-
else:
31-
return ts.astimezone(pytz.utc).replace(tzinfo=None)
32-
33-
3425
class LocalRegistryStore(RegistryStore):
3526
def __init__(self, registry_config: RegistryConfig, repo_path: Path):
3627
registry_path = Path(registry_config.path)

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

Lines changed: 55 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,16 @@
11
import contextlib
22
import uuid
33
from datetime import date, datetime, timedelta
4-
from typing import Callable, ContextManager, Dict, Iterator, List, Optional, Union
4+
from typing import (
5+
Callable,
6+
ContextManager,
7+
Dict,
8+
Iterator,
9+
List,
10+
Optional,
11+
Tuple,
12+
Union,
13+
)
514

615
import numpy as np
716
import pandas as pd
@@ -156,9 +165,17 @@ def query_generator() -> Iterator[str]:
156165
entity_schema, expected_join_keys, entity_df_event_timestamp_col
157166
)
158167

168+
entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(
169+
entity_df, entity_df_event_timestamp_col, client, table_reference,
170+
)
171+
159172
# Build a query context containing all information required to template the BigQuery SQL query
160173
query_context = offline_utils.get_feature_view_query_context(
161-
feature_refs, feature_views, registry, project,
174+
feature_refs,
175+
feature_views,
176+
registry,
177+
project,
178+
entity_df_event_timestamp_range,
162179
)
163180

164181
# Generate the BigQuery SQL query from the query context
@@ -368,7 +385,7 @@ def _upload_entity_df_and_get_entity_schema(
368385
) -> Dict[str, np.dtype]:
369386
"""Uploads a Pandas entity dataframe into a BigQuery table and returns the resulting table"""
370387

371-
if type(entity_df) is str:
388+
if isinstance(entity_df, str):
372389
job = client.query(f"CREATE TABLE {table_name} AS ({entity_df})")
373390
block_until_done(client, job)
374391

@@ -394,6 +411,39 @@ def _upload_entity_df_and_get_entity_schema(
394411
return entity_schema
395412

396413

414+
def _get_entity_df_event_timestamp_range(
415+
entity_df: Union[pd.DataFrame, str],
416+
entity_df_event_timestamp_col: str,
417+
client: Client,
418+
table_name: str,
419+
) -> Tuple[datetime, datetime]:
420+
if type(entity_df) is str:
421+
job = client.query(
422+
f"SELECT MIN({entity_df_event_timestamp_col}) AS min, MAX({entity_df_event_timestamp_col}) AS max FROM {table_name}"
423+
)
424+
res = next(job.result())
425+
entity_df_event_timestamp_range = (
426+
res.get("min"),
427+
res.get("max"),
428+
)
429+
elif isinstance(entity_df, pd.DataFrame):
430+
entity_df_event_timestamp = entity_df.loc[
431+
:, entity_df_event_timestamp_col
432+
].infer_objects()
433+
if pd.api.types.is_string_dtype(entity_df_event_timestamp):
434+
entity_df_event_timestamp = pd.to_datetime(
435+
entity_df_event_timestamp, utc=True
436+
)
437+
entity_df_event_timestamp_range = (
438+
entity_df_event_timestamp.min(),
439+
entity_df_event_timestamp.max(),
440+
)
441+
else:
442+
raise InvalidEntityType(type(entity_df))
443+
444+
return entity_df_event_timestamp_range
445+
446+
397447
def _get_bigquery_client(project: Optional[str] = None, location: Optional[str] = None):
398448
try:
399449
client = bigquery.Client(project=project, location=location)
@@ -484,9 +534,9 @@ def _get_bigquery_client(project: Optional[str] = None, location: Optional[str]
484534
{{ feature }} as {% if full_feature_names %}{{ featureview.name }}__{{feature}}{% else %}{{ feature }}{% endif %}{% if loop.last %}{% else %}, {% endif %}
485535
{% endfor %}
486536
FROM {{ featureview.table_subquery }}
487-
WHERE {{ featureview.event_timestamp_column }} <= (SELECT MAX(entity_timestamp) FROM entity_dataframe)
537+
WHERE {{ featureview.event_timestamp_column }} <= '{{ featureview.max_event_timestamp }}'
488538
{% if featureview.ttl == 0 %}{% else %}
489-
AND {{ featureview.event_timestamp_column }} >= Timestamp_sub((SELECT MIN(entity_timestamp) FROM entity_dataframe), interval {{ featureview.ttl }} second)
539+
AND {{ featureview.event_timestamp_column }} >= '{{ featureview.min_event_timestamp }}'
490540
{% endif %}
491541
),
492542

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import importlib
22
import uuid
33
from dataclasses import asdict, dataclass
4-
from datetime import timedelta
4+
from datetime import datetime, timedelta
55
from typing import Any, Dict, KeysView, List, Optional, Set, Tuple
66

77
import numpy as np
@@ -20,6 +20,7 @@
2020
from feast.infra.offline_stores.offline_store import OfflineStore
2121
from feast.infra.provider import _get_requested_feature_views_to_features_dict
2222
from feast.registry import Registry
23+
from feast.utils import to_naive_utc
2324

2425
DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL = "event_timestamp"
2526

@@ -90,13 +91,16 @@ class FeatureViewQueryContext:
9091
created_timestamp_column: Optional[str]
9192
table_subquery: str
9293
entity_selections: List[str]
94+
min_event_timestamp: Optional[str]
95+
max_event_timestamp: str
9396

9497

9598
def get_feature_view_query_context(
9699
feature_refs: List[str],
97100
feature_views: List[FeatureView],
98101
registry: Registry,
99102
project: str,
103+
entity_df_timestamp_range: Tuple[datetime, datetime],
100104
) -> List[FeatureViewQueryContext]:
101105
"""Build a query context containing all information required to template a BigQuery and Redshift point-in-time SQL query"""
102106

@@ -130,6 +134,14 @@ def get_feature_view_query_context(
130134
event_timestamp_column = feature_view.input.event_timestamp_column
131135
created_timestamp_column = feature_view.input.created_timestamp_column
132136

137+
min_event_timestamp = None
138+
if feature_view.ttl:
139+
min_event_timestamp = to_naive_utc(
140+
entity_df_timestamp_range[0] - feature_view.ttl
141+
).isoformat()
142+
143+
max_event_timestamp = to_naive_utc(entity_df_timestamp_range[1]).isoformat()
144+
133145
context = FeatureViewQueryContext(
134146
name=feature_view.projection.name_to_use(),
135147
ttl=ttl_seconds,
@@ -144,6 +156,8 @@ def get_feature_view_query_context(
144156
# TODO: Make created column optional and not hardcoded
145157
table_subquery=feature_view.input.get_table_query_string(),
146158
entity_selections=entity_selections,
159+
min_event_timestamp=min_event_timestamp,
160+
max_event_timestamp=max_event_timestamp,
147161
)
148162
query_context.append(context)
149163
return query_context

‎sdk/python/feast/infra/offline_stores/redshift.py‎

Lines changed: 68 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,21 @@
11
import contextlib
22
import uuid
33
from datetime import datetime
4-
from typing import Callable, ContextManager, Dict, Iterator, List, Optional, Union
4+
from typing import (
5+
Callable,
6+
ContextManager,
7+
Dict,
8+
Iterator,
9+
List,
10+
Optional,
11+
Tuple,
12+
Union,
13+
)
514

615
import numpy as np
716
import pandas as pd
817
import pyarrow as pa
18+
from dateutil import parser
919
from pydantic import StrictStr
1020
from pydantic.typing import Literal
1121
from pytz import utc
@@ -145,9 +155,21 @@ def query_generator() -> Iterator[str]:
145155
entity_schema, expected_join_keys, entity_df_event_timestamp_col
146156
)
147157

158+
entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(
159+
entity_df,
160+
entity_df_event_timestamp_col,
161+
redshift_client,
162+
config,
163+
table_name,
164+
)
165+
148166
# Build a query context containing all information required to template the Redshift SQL query
149167
query_context = offline_utils.get_feature_view_query_context(
150-
feature_refs, feature_views, registry, project,
168+
feature_refs,
169+
feature_views,
170+
registry,
171+
project,
172+
entity_df_event_timestamp_range,
151173
)
152174

153175
# Generate the Redshift SQL query from the query context
@@ -357,6 +379,48 @@ def _upload_entity_df_and_get_entity_schema(
357379
raise InvalidEntityType(type(entity_df))
358380

359381

382+
def _get_entity_df_event_timestamp_range(
383+
entity_df: Union[pd.DataFrame, str],
384+
entity_df_event_timestamp_col: str,
385+
redshift_client,
386+
config: RepoConfig,
387+
table_name: str,
388+
) -> Tuple[datetime, datetime]:
389+
if isinstance(entity_df, pd.DataFrame):
390+
entity_df_event_timestamp = entity_df.loc[
391+
:, entity_df_event_timestamp_col
392+
].infer_objects()
393+
if pd.api.types.is_string_dtype(entity_df_event_timestamp):
394+
entity_df_event_timestamp = pd.to_datetime(
395+
entity_df_event_timestamp, utc=True
396+
)
397+
entity_df_event_timestamp_range = (
398+
entity_df_event_timestamp.min(),
399+
entity_df_event_timestamp.max(),
400+
)
401+
elif isinstance(entity_df, str):
402+
# If the entity_df is a string (SQL query), determine range
403+
# from table
404+
statement_id = aws_utils.execute_redshift_statement(
405+
redshift_client,
406+
config.offline_store.cluster_id,
407+
config.offline_store.database,
408+
config.offline_store.user,
409+
f"SELECT MIN({entity_df_event_timestamp_col}) AS min, MAX({entity_df_event_timestamp_col}) AS max FROM {table_name}",
410+
)
411+
res = aws_utils.get_redshift_statement_result(redshift_client, statement_id)[
412+
"Records"
413+
][0]
414+
entity_df_event_timestamp_range = (
415+
parser.parse(res[0]["stringValue"]),
416+
parser.parse(res[1]["stringValue"]),
417+
)
418+
else:
419+
raise InvalidEntityType(type(entity_df))
420+
421+
return entity_df_event_timestamp_range
422+
423+
360424
# This query is based on sdk/python/feast/infra/offline_stores/bigquery.py:MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN
361425
# There are couple of changes from BigQuery:
362426
# 1. Use VARCHAR instead of STRING type
@@ -428,9 +492,9 @@ def _upload_entity_df_and_get_entity_schema(
428492
{{ feature }} as {% if full_feature_names %}{{ featureview.name }}__{{feature}}{% else %}{{ feature }}{% endif %}{% if loop.last %}{% else %}, {% endif %}
429493
{% endfor %}
430494
FROM {{ featureview.table_subquery }}
431-
WHERE {{ featureview.event_timestamp_column }} <= (SELECT MAX(entity_timestamp) FROM entity_dataframe)
495+
WHERE {{ featureview.event_timestamp_column }} <= '{{ featureview.max_event_timestamp }}'
432496
{% if featureview.ttl == 0 %}{% else %}
433-
AND {{ featureview.event_timestamp_column }} >= (SELECT MIN(entity_timestamp) FROM entity_dataframe) - {{ featureview.ttl }} * interval '1' second
497+
AND {{ featureview.event_timestamp_column }} >= '{{ featureview.min_event_timestamp }}'
434498
{% endif %}
435499
),
436500

‎sdk/python/feast/infra/online_stores/sqlite.py‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
from pathlib import Path
1919
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
2020

21-
import pytz
2221
from pydantic import StrictStr
2322
from pydantic.schema import Literal
2423

@@ -33,6 +32,7 @@
3332
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
3433
from feast.repo_config import FeastConfigBaseModel, RepoConfig
3534
from feast.usage import log_exceptions_and_usage, tracing_span
35+
from feast.utils import to_naive_utc
3636

3737

3838
class SqliteOnlineStoreConfig(FeastConfigBaseModel):
@@ -95,9 +95,9 @@ def online_write_batch(
9595
with conn:
9696
for entity_key, values, timestamp, created_ts in data:
9797
entity_key_bin = serialize_entity_key(entity_key)
98-
timestamp = _to_naive_utc(timestamp)
98+
timestamp = to_naive_utc(timestamp)
9999
if created_ts is not None:
100-
created_ts = _to_naive_utc(created_ts)
100+
created_ts = to_naive_utc(created_ts)
101101

102102
for feature_name, val in values.items():
103103
conn.execute(
@@ -222,13 +222,6 @@ def _table_id(project: str, table: FeatureView) -> str:
222222
return f"{project}_{table.name}"
223223

224224

225-
def _to_naive_utc(ts: datetime):
226-
if ts.tzinfo is None:
227-
return ts
228-
else:
229-
return ts.astimezone(pytz.utc).replace(tzinfo=None)
230-
231-
232225
class SqliteTable(InfraObject):
233226
"""
234227
A Sqlite table managed by Feast.

‎sdk/python/feast/utils.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,10 @@ def make_tzaware(t: datetime) -> datetime:
99
return t.replace(tzinfo=utc)
1010
else:
1111
return t
12+
13+
14+
def to_naive_utc(ts: datetime) -> datetime:
15+
if ts.tzinfo is None:
16+
return ts
17+
else:
18+
return ts.astimezone(utc).replace(tzinfo=None)

0 commit comments

Comments
 (0)