Skip to content

Commit 39e4148

Browse files
felixwang9817felixwang9817
committed
Fix SFV tests to use actual data to allow for inference
Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent b7416eb commit 39e4148

1 file changed

Lines changed: 122 additions & 114 deletions

File tree

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

Lines changed: 122 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -2,139 +2,147 @@
22

33
import pytest
44

5-
from feast import Entity, Field, FileSource
65
from feast.aggregation import Aggregation
76
from feast.data_format import AvroFormat
87
from feast.data_source import KafkaSource
8+
from feast.entity import Entity
9+
from feast.field import Field
910
from feast.stream_feature_view import stream_feature_view
1011
from feast.types import Float32
12+
from tests.utils.cli_utils import CliRunner, get_example_repo
13+
from tests.utils.data_source_utils import prep_file_source
1114

1215

1316
@pytest.mark.integration
14-
def test_apply_stream_feature_view(environment) -> None:
17+
def test_apply_stream_feature_view(simple_dataset_1) -> None:
1518
"""
1619
Test apply of StreamFeatureView.
1720
"""
18-
fs = environment.feature_store
19-
20-
# Create Feature Views
21-
entity = Entity(name="driver_entity", join_keys=["test_key"])
22-
23-
stream_source = KafkaSource(
24-
name="kafka",
25-
timestamp_field="event_timestamp",
26-
bootstrap_servers="",
27-
message_format=AvroFormat(""),
28-
topic="topic",
29-
batch_source=FileSource(path="test_path", timestamp_field="event_timestamp"),
30-
watermark=timedelta(days=1),
31-
)
32-
33-
@stream_feature_view(
34-
entities=[entity],
35-
ttl=timedelta(days=30),
36-
owner="test@example.com",
37-
online=True,
38-
schema=[Field(name="dummy_field", dtype=Float32)],
39-
description="desc",
40-
aggregations=[
41-
Aggregation(
42-
column="dummy_field", function="max", time_window=timedelta(days=1),
43-
),
44-
Aggregation(
45-
column="dummy_field2", function="count", time_window=timedelta(days=24),
46-
),
47-
],
48-
timestamp_field="event_timestamp",
49-
mode="spark",
50-
source=stream_source,
51-
tags={},
52-
)
53-
def simple_sfv(df):
54-
return df
55-
56-
fs.apply([entity, simple_sfv])
57-
stream_feature_views = fs.list_stream_feature_views()
58-
assert len(stream_feature_views) == 1
59-
assert stream_feature_views[0] == simple_sfv
60-
61-
entities = fs.list_entities()
62-
assert len(entities) == 1
63-
assert entities[0] == entity
64-
65-
features = fs.get_online_features(
66-
features=["simple_sfv:dummy_field"], entity_rows=[{"test_key": 1001}],
67-
).to_dict(include_event_timestamps=True)
68-
69-
assert "test_key" in features
70-
assert features["test_key"] == [1001]
71-
assert "dummy_field" in features
72-
assert features["dummy_field"] == [None]
21+
runner = CliRunner()
22+
with runner.local_repo(
23+
get_example_repo("example_feature_repo_1.py"), "bigquery"
24+
) as fs, prep_file_source(
25+
df=simple_dataset_1, timestamp_field="ts_1"
26+
) as file_source:
27+
entity = Entity(name="driver_entity", join_keys=["test_key"])
28+
29+
stream_source = KafkaSource(
30+
name="kafka",
31+
timestamp_field="event_timestamp",
32+
bootstrap_servers="",
33+
message_format=AvroFormat(""),
34+
topic="topic",
35+
batch_source=file_source,
36+
watermark=timedelta(days=1),
37+
)
38+
39+
@stream_feature_view(
40+
entities=[entity],
41+
ttl=timedelta(days=30),
42+
owner="test@example.com",
43+
online=True,
44+
schema=[Field(name="dummy_field", dtype=Float32)],
45+
description="desc",
46+
aggregations=[
47+
Aggregation(
48+
column="dummy_field", function="max", time_window=timedelta(days=1),
49+
),
50+
Aggregation(
51+
column="dummy_field2",
52+
function="count",
53+
time_window=timedelta(days=24),
54+
),
55+
],
56+
timestamp_field="event_timestamp",
57+
mode="spark",
58+
source=stream_source,
59+
tags={},
60+
)
61+
def simple_sfv(df):
62+
return df
63+
64+
fs.apply([entity, simple_sfv])
65+
66+
stream_feature_views = fs.list_stream_feature_views()
67+
assert len(stream_feature_views) == 1
68+
assert stream_feature_views[0] == simple_sfv
69+
70+
features = fs.get_online_features(
71+
features=["simple_sfv:dummy_field"], entity_rows=[{"test_key": 1001}],
72+
).to_dict(include_event_timestamps=True)
73+
74+
assert "test_key" in features
75+
assert features["test_key"] == [1001]
76+
assert "dummy_field" in features
77+
assert features["dummy_field"] == [None]
7378

7479

7580
@pytest.mark.integration
76-
def test_stream_feature_view_udf(environment) -> None:
81+
def test_stream_feature_view_udf(simple_dataset_1) -> None:
7782
"""
7883
Test apply of StreamFeatureView udfs are serialized correctly and usable.
7984
"""
80-
fs = environment.feature_store
81-
82-
# Create Feature Views
83-
entity = Entity(name="driver_entity", join_keys=["test_key"])
84-
85-
stream_source = KafkaSource(
86-
name="kafka",
87-
timestamp_field="event_timestamp",
88-
bootstrap_servers="",
89-
message_format=AvroFormat(""),
90-
topic="topic",
91-
batch_source=FileSource(path="test_path", timestamp_field="event_timestamp"),
92-
watermark=timedelta(days=1),
93-
)
94-
95-
@stream_feature_view(
96-
entities=[entity],
97-
ttl=timedelta(days=30),
98-
owner="test@example.com",
99-
online=True,
100-
schema=[Field(name="dummy_field", dtype=Float32)],
101-
description="desc",
102-
aggregations=[
103-
Aggregation(
104-
column="dummy_field", function="max", time_window=timedelta(days=1),
105-
),
106-
Aggregation(
107-
column="dummy_field2", function="count", time_window=timedelta(days=24),
108-
),
109-
],
110-
timestamp_field="event_timestamp",
111-
mode="spark",
112-
source=stream_source,
113-
tags={},
114-
)
115-
def pandas_view(pandas_df):
116-
import pandas as pd
117-
118-
assert type(pandas_df) == pd.DataFrame
119-
df = pandas_df.transform(lambda x: x + 10, axis=1)
120-
df.insert(2, "C", [20.2, 230.0, 34.0], True)
121-
return df
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+
bootstrap_servers="",
97+
message_format=AvroFormat(""),
98+
topic="topic",
99+
batch_source=file_source,
100+
watermark=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
122132

123-
import pandas as pd
124-
125-
df = pd.DataFrame({"A": [1, 2, 3], "B": [10, 20, 30]})
133+
import pandas as pd
126134

127-
fs.apply([entity, pandas_view])
128-
stream_feature_views = fs.list_stream_feature_views()
129-
assert len(stream_feature_views) == 1
130-
assert stream_feature_views[0].name == "pandas_view"
131-
assert stream_feature_views[0] == pandas_view
135+
fs.apply([entity, pandas_view])
132136

133-
sfv = stream_feature_views[0]
137+
stream_feature_views = fs.list_stream_feature_views()
138+
assert len(stream_feature_views) == 1
139+
assert stream_feature_views[0] == pandas_view
134140

135-
new_df = sfv.udf(df)
141+
sfv = stream_feature_views[0]
136142

137-
expected_df = pd.DataFrame(
138-
{"A": [11, 12, 13], "B": [20, 30, 40], "C": [20.2, 230.0, 34.0]}
139-
)
140-
assert new_df.equals(expected_df)
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)