Skip to content

Commit 563d9cf

Browse files
MattDelacMattDelac
authored
New SQL template for BQ historical retrieval (#1602)
Signed-off-by: Matt Delacour <matt.delacour@shopify.com>
1 parent c3ee2cc commit 563d9cf

2 files changed

Lines changed: 119 additions & 94 deletions

File tree

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

Lines changed: 109 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,10 @@ def _get_bigquery_client():
394394
# * Create temporary tables instead of keeping all tables in memory
395395

396396
SINGLE_FEATURE_VIEW_POINT_IN_TIME_JOIN = """
397+
/*
398+
Compute a deterministic hash for the `left_table_query_string` that will be used throughout
399+
all the logic as the field to GROUP BY the data
400+
*/
397401
WITH entity_dataframe AS (
398402
SELECT
399403
*,
@@ -405,110 +409,128 @@ def _get_bigquery_client():
405409
) AS entity_row_unique_id
406410
FROM {{ left_table_query_string }}
407411
),
412+
408413
{% for featureview in featureviews %}
414+
409415
/*
410416
This query template performs the point-in-time correctness join for a single feature set table
411417
to the provided entity table.
412-
1. Concatenate the timestamp and entities from the feature set table with the entity dataset.
413-
Feature values are joined to this table later for improved efficiency.
414-
featureview_timestamp is equal to null in rows from the entity dataset.
415-
*/
416-
{{ featureview.name }}__union_features AS (
417-
SELECT
418-
-- unique identifier for each row in the entity dataset.
419-
entity_row_unique_id,
420-
-- event_timestamp contains the timestamps to join onto
421-
{{entity_df_event_timestamp_col}} AS event_timestamp,
422-
-- the feature_timestamp, i.e. the latest occurrence of the requested feature relative to the entity_dataset timestamp
423-
NULL as {{ featureview.name }}_feature_timestamp,
424-
-- created timestamp of the feature at the corresponding feature_timestamp
425-
{{ 'NULL as created_timestamp,' if featureview.created_timestamp_column else '' }}
426-
-- select only entities belonging to this feature set
427-
{{ featureview.entities | join(', ')}},
428-
-- boolean for filtering the dataset later
429-
true AS is_entity_table
430-
FROM entity_dataframe
431-
UNION ALL
432-
SELECT
433-
NULL as entity_row_unique_id,
434-
{{ featureview.event_timestamp_column }} as event_timestamp,
435-
{{ featureview.event_timestamp_column }} as {{ featureview.name }}_feature_timestamp,
436-
{{ featureview.created_timestamp_column ~ ' as created_timestamp,' if featureview.created_timestamp_column else '' }}
437-
{{ featureview.entity_selections | join(', ')}},
438-
false AS is_entity_table
439-
FROM {{ featureview.table_subquery }} WHERE {{ featureview.event_timestamp_column }} <= '{{ max_timestamp }}'
440-
{% if featureview.ttl == 0 %}{% else %}AND {{ featureview.event_timestamp_column }} >= Timestamp_sub(TIMESTAMP '{{ min_timestamp }}', interval {{ featureview.ttl }} second){% endif %}
418+
419+
1. We first join the current feature_view to the entity dataframe that has been passed.
420+
This JOIN has the following logic:
421+
- For each row of the entity dataframe, only keep the rows where the `event_timestamp_column`
422+
is less than the one provided in the entity dataframe
423+
- If there a TTL for the current feature_view, also keep the rows where the `event_timestamp_column`
424+
is higher the the one provided minus the TTL
425+
- For each row, Join on the entity key and retrieve the `entity_row_unique_id` that has been
426+
computed previously
427+
428+
The output of this CTE will contain all the necessary information and already filtered out most
429+
of the data that is not relevant.
430+
*/
431+
432+
{{ featureview.name }}__subquery AS (
433+
SELECT
434+
{{ featureview.event_timestamp_column }} as event_timestamp,
435+
{{ featureview.created_timestamp_column ~ ' as created_timestamp,' if featureview.created_timestamp_column else '' }}
436+
{{ featureview.entity_selections | join(', ')}},
437+
{% for feature in featureview.features %}
438+
{{ feature }} as {{ featureview.name }}__{{ feature }}{% if loop.last %}{% else %}, {% endif %}
439+
{% endfor %}
440+
FROM {{ featureview.table_subquery }}
441441
),
442+
443+
{{ featureview.name }}__base AS (
444+
SELECT
445+
subquery.*,
446+
entity_dataframe.{{entity_df_event_timestamp_col}} AS entity_timestamp,
447+
entity_dataframe.entity_row_unique_id
448+
FROM {{ featureview.name }}__subquery AS subquery
449+
INNER JOIN entity_dataframe
450+
ON TRUE
451+
AND subquery.event_timestamp <= entity_dataframe.{{entity_df_event_timestamp_col}}
452+
453+
{% if featureview.ttl == 0 %}{% else %}
454+
AND subquery.event_timestamp >= Timestamp_sub(entity_dataframe.{{entity_df_event_timestamp_col}}, interval {{ featureview.ttl }} second)
455+
{% endif %}
456+
457+
{% for entity in featureview.entities %}
458+
AND subquery.{{ entity }} = entity_dataframe.{{ entity }}
459+
{% endfor %}
460+
),
461+
442462
/*
443-
2. Window the data in the unioned dataset, partitioning by entity and ordering by event_timestamp, as
444-
well as is_entity_table.
445-
Within each window, back-fill the feature_timestamp - as a result of this, the null feature_timestamps
446-
in the rows from the entity table should now contain the latest timestamps relative to the row's
447-
event_timestamp.
448-
For rows where event_timestamp(provided datetime) - feature_timestamp > max age, set the
449-
feature_timestamp to null.
450-
*/
451-
{{ featureview.name }}__joined AS (
452-
SELECT
453-
entity_row_unique_id,
454-
event_timestamp,
455-
{{ featureview.entities | join(', ')}},
456-
{% for feature in featureview.features %}
457-
IF(event_timestamp >= {{ featureview.name }}_feature_timestamp {% if featureview.ttl == 0 %}{% else %}AND Timestamp_sub(event_timestamp, interval {{ featureview.ttl }} second) < {{ featureview.name }}_feature_timestamp{% endif %}, {{ featureview.name }}__{{ feature }}, NULL) as {{ featureview.name }}__{{ feature }}{% if loop.last %}{% else %}, {% endif %}
458-
{% endfor %}
459-
FROM (
460-
SELECT
461-
entity_row_unique_id,
462-
event_timestamp,
463-
{{ featureview.entities | join(', ')}},
464-
{{ 'FIRST_VALUE(created_timestamp IGNORE NULLS) over w AS created_timestamp,' if featureview.created_timestamp_column else '' }}
465-
FIRST_VALUE({{ featureview.name }}_feature_timestamp IGNORE NULLS) over w AS {{ featureview.name }}_feature_timestamp,
466-
is_entity_table
467-
FROM {{ featureview.name }}__union_features
468-
WINDOW w AS (PARTITION BY {{ featureview.entities | join(', ') }} ORDER BY event_timestamp DESC, is_entity_table DESC{{', created_timestamp DESC' if featureview.created_timestamp_column else ''}} ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)
469-
)
463+
2. If the `created_timestamp_column` has been set, we need to
464+
deduplicate the data first. This is done by calculating the
465+
`MAX(created_at_timestamp)` for each event_timestamp.
466+
We then join the data on the next CTE
467+
*/
468+
{% if featureview.created_timestamp_column %}
469+
{{ featureview.name }}__dedup AS (
470+
SELECT
471+
entity_row_unique_id,
472+
event_timestamp,
473+
MAX(created_timestamp) as created_timestamp,
474+
FROM {{ featureview.name }}__base
475+
GROUP BY entity_row_unique_id, event_timestamp
476+
),
477+
{% endif %}
478+
470479
/*
471-
3. Select only the rows from the entity table, and join the features from the original feature set table
472-
to the dataset using the entity values, feature_timestamp, and created_timestamps.
473-
*/
474-
LEFT JOIN (
475-
SELECT
476-
{{ featureview.event_timestamp_column }} as {{ featureview.name }}_feature_timestamp,
477-
{{ featureview.created_timestamp_column ~ ' as created_timestamp,' if featureview.created_timestamp_column else '' }}
478-
{{ featureview.entity_selections | join(', ')}},
479-
{% for feature in featureview.features %}
480-
{{ feature }} as {{ featureview.name }}__{{ feature }}{% if loop.last %}{% else %}, {% endif %}
481-
{% endfor %}
482-
FROM {{ featureview.table_subquery }} WHERE {{ featureview.event_timestamp_column }} <= '{{ max_timestamp }}'
483-
{% if featureview.ttl == 0 %}{% else %}AND {{ featureview.event_timestamp_column }} >= Timestamp_sub(TIMESTAMP '{{ min_timestamp }}', interval {{ featureview.ttl }} second){% endif %}
484-
) USING ({{ featureview.name }}_feature_timestamp,{{ ' created_timestamp,' if featureview.created_timestamp_column else '' }} {{ featureview.entities | join(', ')}})
485-
WHERE is_entity_table
480+
3. The data has been filtered during the first CTE "*__base"
481+
Thus we only need to compute the latest timestamp of each feature.
482+
*/
483+
{{ featureview.name }}__latest AS (
484+
SELECT
485+
entity_row_unique_id,
486+
MAX(event_timestamp) AS event_timestamp
487+
{% if featureview.created_timestamp_column %}
488+
,ANY_VALUE(created_timestamp) AS created_timestamp
489+
{% endif %}
490+
491+
FROM {{ featureview.name }}__base
492+
{% if featureview.created_timestamp_column %}
493+
INNER JOIN {{ featureview.name }}__dedup
494+
USING (entity_row_unique_id, event_timestamp, created_timestamp)
495+
{% endif %}
496+
497+
GROUP BY entity_row_unique_id
486498
),
499+
487500
/*
488-
4. Finally, deduplicate the rows by selecting the first occurrence of each entity table entity_row_unique_id.
489-
*/
490-
{{ featureview.name }}__deduped AS (SELECT
491-
k.*
492-
FROM (
493-
SELECT ARRAY_AGG(row LIMIT 1)[OFFSET(0)] k
494-
FROM {{ featureview.name }}__joined row
495-
GROUP BY entity_row_unique_id
496-
)){% if loop.last %}{% else %}, {% endif %}
501+
4. Once we know the latest value of each feature for a given timestamp,
502+
we can join again the data back to the original "base" dataset
503+
*/
504+
{{ featureview.name }}__cleaned AS (
505+
SELECT base.*
506+
FROM {{ featureview.name }}__base as base
507+
INNER JOIN {{ featureview.name }}__latest
508+
USING(
509+
entity_row_unique_id,
510+
event_timestamp
511+
{% if featureview.created_timestamp_column %}
512+
,created_timestamp
513+
{% endif %}
514+
)
515+
){% if loop.last %}{% else %}, {% endif %}
516+
497517
498518
{% endfor %}
499519
/*
500520
Joins the outputs of multiple time travel joins to a single table.
521+
The entity_dataframe dataset being our source of truth here.
501522
*/
502-
SELECT edf.{{entity_df_event_timestamp_col}} as {{entity_df_event_timestamp_col}}, * EXCEPT (entity_row_unique_id, {{entity_df_event_timestamp_col}}) FROM entity_dataframe edf
523+
524+
SELECT * EXCEPT (entity_row_unique_id)
525+
FROM entity_dataframe
503526
{% for featureview in featureviews %}
504527
LEFT JOIN (
505528
SELECT
506-
entity_row_unique_id,
507-
{% for feature in featureview.features %}
508-
{{ featureview.name }}__{{ feature }}{% if loop.last %}{% else %}, {% endif %}
509-
{% endfor %}
510-
FROM {{ featureview.name }}__deduped
529+
entity_row_unique_id,
530+
{% for feature in featureview.features %}
531+
{{ featureview.name }}__{{ feature }},
532+
{% endfor %}
533+
FROM {{ featureview.name }}__cleaned
511534
) USING (entity_row_unique_id)
512535
{% endfor %}
513-
ORDER BY {{entity_df_event_timestamp_col}}
514536
"""

‎sdk/python/tests/test_historical_retrieval.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,6 @@ def test_historical_features_from_bigquery_sources(
334334
start_date,
335335
) = generate_entities(start_date, infer_event_timestamp_col)
336336

337-
# bigquery_dataset = "test_hist_retrieval_static"
338337
bigquery_dataset = (
339338
f"test_hist_retrieval_{int(time.time_ns())}_{random.randint(1000, 9999)}"
340339
)
@@ -452,13 +451,16 @@ def test_historical_features_from_bigquery_sources(
452451
)
453452
)
454453

454+
assert sorted(expected_df.columns) == sorted(
455+
actual_df_from_sql_entities.columns
456+
)
455457
assert_frame_equal(
456458
expected_df.sort_values(
457459
by=[event_timestamp, "order_id", "driver_id", "customer_id"]
458460
).reset_index(drop=True),
459-
actual_df_from_sql_entities.sort_values(
460-
by=[event_timestamp, "order_id", "driver_id", "customer_id"]
461-
).reset_index(drop=True),
461+
actual_df_from_sql_entities[expected_df.columns]
462+
.sort_values(by=[event_timestamp, "order_id", "driver_id", "customer_id"])
463+
.reset_index(drop=True),
462464
check_dtype=False,
463465
)
464466

@@ -532,12 +534,13 @@ def test_historical_features_from_bigquery_sources(
532534
)
533535
)
534536

537+
assert sorted(expected_df.columns) == sorted(actual_df_from_df_entities.columns)
535538
assert_frame_equal(
536539
expected_df.sort_values(
537540
by=[event_timestamp, "order_id", "driver_id", "customer_id"]
538541
).reset_index(drop=True),
539-
actual_df_from_df_entities.sort_values(
540-
by=[event_timestamp, "order_id", "driver_id", "customer_id"]
541-
).reset_index(drop=True),
542+
actual_df_from_df_entities[expected_df.columns]
543+
.sort_values(by=[event_timestamp, "order_id", "driver_id", "customer_id"])
544+
.reset_index(drop=True),
542545
check_dtype=False,
543546
)

0 commit comments

Comments
 (0)