Skip to content

Commit eea45c9

Browse files
kevjumbakevjumba
committed
add back in
Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
1 parent d0c679c commit eea45c9

2 files changed

Lines changed: 143 additions & 143 deletions

File tree

‎sdk/python/tests/integration/online_store/test_universal_online.py‎

Lines changed: 74 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -441,80 +441,80 @@ def test_online_retrieval_with_event_timestamps(
441441
)
442442

443443

444-
# @pytest.mark.integration
445-
# @pytest.mark.universal_online_stores
446-
# # @pytest.mark.goserver Disabling because the go fs tests are flaking in CI. TODO(achals): uncomment after fixed.
447-
# @pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: str(v))
448-
# def test_stream_feature_view_online_retrieval(
449-
# environment, universal_data_sources, feature_server_endpoint, full_feature_names
450-
# ):
451-
# """
452-
# Tests materialization and online retrieval for stream feature views.
453-
454-
# This test is separate from test_online_retrieval since combining feature views and
455-
# stream feature views into a single test resulted in test flakiness. This is tech
456-
# debt that should be resolved soon.
457-
# """
458-
# # Set up feature store.
459-
# fs = environment.feature_store
460-
# entities, datasets, data_sources = universal_data_sources
461-
# feature_views = construct_universal_feature_views(data_sources)
462-
# pushable_feature_view = feature_views.pushed_locations
463-
# fs.apply([location(), pushable_feature_view])
464-
465-
# # Materialize.
466-
# fs.materialize(
467-
# environment.start_date - timedelta(days=1),
468-
# environment.end_date + timedelta(days=1),
469-
# )
470-
471-
# # Get online features by randomly sampling 10 entities that exist in the batch source.
472-
# sample_locations = datasets.location_df.sample(10)["location_id"]
473-
# entity_rows = [
474-
# {"location_id": sample_location} for sample_location in sample_locations
475-
# ]
476-
477-
# feature_refs = [
478-
# "pushable_location_stats:temperature",
479-
# ]
480-
# unprefixed_feature_refs = [f.rsplit(":", 1)[-1] for f in feature_refs if ":" in f]
481-
482-
# online_features_dict = get_online_features_dict(
483-
# environment=environment,
484-
# endpoint=feature_server_endpoint,
485-
# features=feature_refs,
486-
# entity_rows=entity_rows,
487-
# full_feature_names=full_feature_names,
488-
# )
489-
490-
# # Check that the response has the expected set of keys.
491-
# keys = set(online_features_dict.keys())
492-
# expected_keys = set(
493-
# f.replace(":", "__") if full_feature_names else f.split(":")[-1]
494-
# for f in feature_refs
495-
# ) | {"location_id"}
496-
# assert (
497-
# keys == expected_keys
498-
# ), f"Response keys are different from expected: {keys - expected_keys} (extra) and {expected_keys - keys} (missing)"
499-
500-
# # Check that the feature values match.
501-
# tc = unittest.TestCase()
502-
# for i, entity_row in enumerate(entity_rows):
503-
# df_features = get_latest_feature_values_from_location_df(
504-
# entity_row, datasets.location_df
505-
# )
506-
507-
# assert df_features["location_id"] == online_features_dict["location_id"][i]
508-
# for unprefixed_feature_ref in unprefixed_feature_refs:
509-
# tc.assertAlmostEqual(
510-
# df_features[unprefixed_feature_ref],
511-
# online_features_dict[
512-
# response_feature_name(
513-
# unprefixed_feature_ref, feature_refs, full_feature_names
514-
# )
515-
# ][i],
516-
# delta=0.0001,
517-
# )
444+
@pytest.mark.integration
445+
@pytest.mark.universal_online_stores
446+
# @pytest.mark.goserver Disabling because the go fs tests are flaking in CI. TODO(achals): uncomment after fixed.
447+
@pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: str(v))
448+
def test_stream_feature_view_online_retrieval(
449+
environment, universal_data_sources, feature_server_endpoint, full_feature_names
450+
):
451+
"""
452+
Tests materialization and online retrieval for stream feature views.
453+
454+
This test is separate from test_online_retrieval since combining feature views and
455+
stream feature views into a single test resulted in test flakiness. This is tech
456+
debt that should be resolved soon.
457+
"""
458+
# Set up feature store.
459+
fs = environment.feature_store
460+
entities, datasets, data_sources = universal_data_sources
461+
feature_views = construct_universal_feature_views(data_sources)
462+
pushable_feature_view = feature_views.pushed_locations
463+
fs.apply([location(), pushable_feature_view])
464+
465+
# Materialize.
466+
fs.materialize(
467+
environment.start_date - timedelta(days=1),
468+
environment.end_date + timedelta(days=1),
469+
)
470+
471+
# Get online features by randomly sampling 10 entities that exist in the batch source.
472+
sample_locations = datasets.location_df.sample(10)["location_id"]
473+
entity_rows = [
474+
{"location_id": sample_location} for sample_location in sample_locations
475+
]
476+
477+
feature_refs = [
478+
"pushable_location_stats:temperature",
479+
]
480+
unprefixed_feature_refs = [f.rsplit(":", 1)[-1] for f in feature_refs if ":" in f]
481+
482+
online_features_dict = get_online_features_dict(
483+
environment=environment,
484+
endpoint=feature_server_endpoint,
485+
features=feature_refs,
486+
entity_rows=entity_rows,
487+
full_feature_names=full_feature_names,
488+
)
489+
490+
# Check that the response has the expected set of keys.
491+
keys = set(online_features_dict.keys())
492+
expected_keys = set(
493+
f.replace(":", "__") if full_feature_names else f.split(":")[-1]
494+
for f in feature_refs
495+
) | {"location_id"}
496+
assert (
497+
keys == expected_keys
498+
), f"Response keys are different from expected: {keys - expected_keys} (extra) and {expected_keys - keys} (missing)"
499+
500+
# Check that the feature values match.
501+
tc = unittest.TestCase()
502+
for i, entity_row in enumerate(entity_rows):
503+
df_features = get_latest_feature_values_from_location_df(
504+
entity_row, datasets.location_df
505+
)
506+
507+
assert df_features["location_id"] == online_features_dict["location_id"][i]
508+
for unprefixed_feature_ref in unprefixed_feature_refs:
509+
tc.assertAlmostEqual(
510+
df_features[unprefixed_feature_ref],
511+
online_features_dict[
512+
response_feature_name(
513+
unprefixed_feature_ref, feature_refs, full_feature_names
514+
)
515+
][i],
516+
delta=0.0001,
517+
)
518518

519519

520520
@pytest.mark.integration

‎sdk/python/tests/integration/registration/test_stream_feature_view_apply.py‎

Lines changed: 69 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -77,72 +77,72 @@ def simple_sfv(df):
7777
assert features["dummy_field"] == [None]
7878

7979

80-
# @pytest.mark.integration
81-
# def test_stream_feature_view_udf(simple_dataset_1) -> None:
82-
# """
83-
# Test apply of StreamFeatureView udfs are serialized correctly and usable.
84-
# """
85-
# runner = CliRunner()
86-
# with runner.local_repo(
87-
# get_example_repo("example_feature_repo_1.py"), "bigquery"
88-
# ) as fs, prep_file_source(
89-
# df=simple_dataset_1, timestamp_field="ts_1"
90-
# ) as file_source:
91-
# entity = Entity(name="driver_entity", join_keys=["test_key"])
92-
93-
# stream_source = KafkaSource(
94-
# name="kafka",
95-
# timestamp_field="event_timestamp",
96-
# kafka_bootstrap_servers="",
97-
# message_format=AvroFormat(""),
98-
# topic="topic",
99-
# batch_source=file_source,
100-
# watermark_delay_threshold=timedelta(days=1),
101-
# )
102-
103-
# @stream_feature_view(
104-
# entities=[entity],
105-
# ttl=timedelta(days=30),
106-
# owner="test@example.com",
107-
# online=True,
108-
# schema=[Field(name="dummy_field", dtype=Float32)],
109-
# description="desc",
110-
# aggregations=[
111-
# Aggregation(
112-
# column="dummy_field", function="max", time_window=timedelta(days=1),
113-
# ),
114-
# Aggregation(
115-
# column="dummy_field2",
116-
# function="count",
117-
# time_window=timedelta(days=24),
118-
# ),
119-
# ],
120-
# timestamp_field="event_timestamp",
121-
# mode="spark",
122-
# source=stream_source,
123-
# tags={},
124-
# )
125-
# def pandas_view(pandas_df):
126-
# import pandas as pd
127-
128-
# assert type(pandas_df) == pd.DataFrame
129-
# df = pandas_df.transform(lambda x: x + 10, axis=1)
130-
# df.insert(2, "C", [20.2, 230.0, 34.0], True)
131-
# return df
132-
133-
# import pandas as pd
134-
135-
# fs.apply([entity, pandas_view])
136-
137-
# stream_feature_views = fs.list_stream_feature_views()
138-
# assert len(stream_feature_views) == 1
139-
# assert stream_feature_views[0] == pandas_view
140-
141-
# sfv = stream_feature_views[0]
142-
143-
# df = pd.DataFrame({"A": [1, 2, 3], "B": [10, 20, 30]})
144-
# new_df = sfv.udf(df)
145-
# expected_df = pd.DataFrame(
146-
# {"A": [11, 12, 13], "B": [20, 30, 40], "C": [20.2, 230.0, 34.0]}
147-
# )
148-
# assert new_df.equals(expected_df)
80+
@pytest.mark.integration
81+
def test_stream_feature_view_udf(simple_dataset_1) -> None:
82+
"""
83+
Test apply of StreamFeatureView udfs are serialized correctly and usable.
84+
"""
85+
runner = CliRunner()
86+
with runner.local_repo(
87+
get_example_repo("example_feature_repo_1.py"), "bigquery"
88+
) as fs, prep_file_source(
89+
df=simple_dataset_1, timestamp_field="ts_1"
90+
) as file_source:
91+
entity = Entity(name="driver_entity", join_keys=["test_key"])
92+
93+
stream_source = KafkaSource(
94+
name="kafka",
95+
timestamp_field="event_timestamp",
96+
kafka_bootstrap_servers="",
97+
message_format=AvroFormat(""),
98+
topic="topic",
99+
batch_source=file_source,
100+
watermark_delay_threshold=timedelta(days=1),
101+
)
102+
103+
@stream_feature_view(
104+
entities=[entity],
105+
ttl=timedelta(days=30),
106+
owner="test@example.com",
107+
online=True,
108+
schema=[Field(name="dummy_field", dtype=Float32)],
109+
description="desc",
110+
aggregations=[
111+
Aggregation(
112+
column="dummy_field", function="max", time_window=timedelta(days=1),
113+
),
114+
Aggregation(
115+
column="dummy_field2",
116+
function="count",
117+
time_window=timedelta(days=24),
118+
),
119+
],
120+
timestamp_field="event_timestamp",
121+
mode="spark",
122+
source=stream_source,
123+
tags={},
124+
)
125+
def pandas_view(pandas_df):
126+
import pandas as pd
127+
128+
assert type(pandas_df) == pd.DataFrame
129+
df = pandas_df.transform(lambda x: x + 10, axis=1)
130+
df.insert(2, "C", [20.2, 230.0, 34.0], True)
131+
return df
132+
133+
import pandas as pd
134+
135+
fs.apply([entity, pandas_view])
136+
137+
stream_feature_views = fs.list_stream_feature_views()
138+
assert len(stream_feature_views) == 1
139+
assert stream_feature_views[0] == pandas_view
140+
141+
sfv = stream_feature_views[0]
142+
143+
df = pd.DataFrame({"A": [1, 2, 3], "B": [10, 20, 30]})
144+
new_df = sfv.udf(df)
145+
expected_df = pd.DataFrame(
146+
{"A": [11, 12, 13], "B": [20, 30, 40], "C": [20.2, 230.0, 34.0]}
147+
)
148+
assert new_df.equals(expected_df)

0 commit comments

Comments
 (0)