Skip to content

Commit 7a4be44

Browse files
kevjumbakevjumba
committed
Fix
Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
1 parent 8049390 commit 7a4be44

2 files changed

Lines changed: 238 additions & 5 deletions

File tree

‎sdk/python/tests/integration/offline_store/test_offline_push.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
@pytest.mark.universal_online_stores
1515
def test_writing_incorrect_order_fails(environment, universal_data_sources):
1616
# TODO(kevjumba) handle incorrect order later, for now schema must be in the order that the filesource is in
17+
"""This test tests if we have incorrect order when writing to offline store.
18+
Specifically, event_timestamp should be the first column to adhere with the filesource column order.
19+
"""
1720
store = environment.feature_store
1821
_, _, data_sources = universal_data_sources
1922
driver_stats = FeatureView(
@@ -43,7 +46,7 @@ def test_writing_incorrect_order_fails(environment, universal_data_sources):
4346
assert df["conv_rate"].isnull().all()
4447
assert df["avg_daily_trips"].isnull().all()
4548

46-
expected_df = pd.DataFrame.from_dict(
49+
df = pd.DataFrame.from_dict(
4750
{
4851
"driver_id": [1001, 1002],
4952
"event_timestamp": [ts - timedelta(hours=3), ts],
@@ -54,14 +57,17 @@ def test_writing_incorrect_order_fails(environment, universal_data_sources):
5457
)
5558
with pytest.raises(ValueError):
5659
store.write_to_offline_store(
57-
driver_stats.name, expected_df, allow_registry_cache=False
60+
driver_stats.name, df, allow_registry_cache=False
5861
)
5962

6063

6164
@pytest.mark.integration
6265
@pytest.mark.universal_online_stores
6366
def test_writing_incorrect_schema_fails(environment, universal_data_sources):
6467
# TODO(kevjumba) handle incorrect order later, for now schema must be in the order that the filesource is in
68+
"""This test tests if we have incorrect attribute when writing to offline store.
69+
Specifically, `incorrect_attribute` is an inccorect column to adhere with the filesource column order.
70+
"""
6571
store = environment.feature_store
6672
_, _, data_sources = universal_data_sources
6773
driver_stats = FeatureView(
@@ -91,18 +97,18 @@ def test_writing_incorrect_schema_fails(environment, universal_data_sources):
9197
assert df["conv_rate"].isnull().all()
9298
assert df["avg_daily_trips"].isnull().all()
9399

94-
expected_df = pd.DataFrame.from_dict(
100+
df = pd.DataFrame.from_dict(
95101
{
96102
"event_timestamp": [ts - timedelta(hours=3), ts],
97103
"driver_id": [1001, 1002],
98104
"conv_rate": [random.random(), random.random()],
99-
"incorrect_schema": [random.randint(0, 10), random.randint(0, 10)],
105+
"incorrect_attribute": [random.randint(0, 10), random.randint(0, 10)],
100106
"created": [ts, ts],
101107
},
102108
)
103109
with pytest.raises(ValueError):
104110
store.write_to_offline_store(
105-
driver_stats.name, expected_df, allow_registry_cache=False
111+
driver_stats.name, df, allow_registry_cache=False
106112
)
107113

108114

@@ -143,6 +149,7 @@ def test_writing_consecutively_to_offline_store(environment, universal_data_sour
143149
assert df["conv_rate"].isnull().all()
144150
assert df["avg_daily_trips"].isnull().all()
145151

152+
# This dataframe has its columns ordered exactly as it is in the parquet file generated by driver_test_data.py.
146153
first_df = pd.DataFrame.from_dict(
147154
{
148155
"event_timestamp": [ts - timedelta(hours=4), ts - timedelta(hours=3)],

‎sdk/python/tests/integration/offline_store/test_offline_write.py‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,226 @@
1+
import random
2+
from datetime import datetime, timedelta
3+
4+
import numpy as np
5+
import pandas as pd
6+
import pytest
7+
8+
from feast import FeatureView, Field
9+
from feast.types import Float32, Int32
10+
from tests.integration.feature_repos.universal.entities import driver
11+
12+
13+
@pytest.mark.integration
14+
@pytest.mark.universal_online_stores
15+
def test_writing_incorrect_order_fails(environment, universal_data_sources):
16+
# TODO(kevjumba) handle incorrect order later, for now schema must be in the order that the filesource is in
17+
store = environment.feature_store
18+
_, _, data_sources = universal_data_sources
19+
driver_stats = FeatureView(
20+
name="driver_stats",
21+
entities=["driver"],
22+
schema=[
23+
Field(name="avg_daily_trips", dtype=Int32),
24+
Field(name="conv_rate", dtype=Float32),
25+
],
26+
source=data_sources.driver,
27+
)
28+
29+
now = datetime.utcnow()
30+
ts = pd.Timestamp(now).round("ms")
31+
32+
entity_df = pd.DataFrame.from_dict(
33+
{"driver_id": [1001, 1002], "event_timestamp": [ts - timedelta(hours=3), ts]}
34+
)
35+
36+
store.apply([driver(), driver_stats])
37+
df = store.get_historical_features(
38+
entity_df=entity_df,
39+
features=["driver_stats:conv_rate", "driver_stats:avg_daily_trips"],
40+
full_feature_names=False,
41+
).to_df()
42+
43+
assert df["conv_rate"].isnull().all()
44+
assert df["avg_daily_trips"].isnull().all()
45+
46+
expected_df = pd.DataFrame.from_dict(
47+
{
48+
"driver_id": [1001, 1002],
49+
"event_timestamp": [ts - timedelta(hours=3), ts],
50+
"conv_rate": [random.random(), random.random()],
51+
"avg_daily_trips": [random.randint(0, 10), random.randint(0, 10)],
52+
"created": [ts, ts],
53+
},
54+
)
55+
with pytest.raises(ValueError):
56+
store.write_to_offline_store(
57+
driver_stats.name, expected_df, allow_registry_cache=False
58+
)
59+
60+
61+
@pytest.mark.integration
62+
@pytest.mark.universal_online_stores
63+
def test_writing_incorrect_schema_fails(environment, universal_data_sources):
64+
# TODO(kevjumba) handle incorrect order later, for now schema must be in the order that the filesource is in
65+
store = environment.feature_store
66+
_, _, data_sources = universal_data_sources
67+
driver_stats = FeatureView(
68+
name="driver_stats",
69+
entities=["driver"],
70+
schema=[
71+
Field(name="avg_daily_trips", dtype=Int32),
72+
Field(name="conv_rate", dtype=Float32),
73+
],
74+
source=data_sources.driver,
75+
)
76+
77+
now = datetime.utcnow()
78+
ts = pd.Timestamp(now).round("ms")
79+
80+
entity_df = pd.DataFrame.from_dict(
81+
{"driver_id": [1001, 1002], "event_timestamp": [ts - timedelta(hours=3), ts]}
82+
)
83+
84+
store.apply([driver(), driver_stats])
85+
df = store.get_historical_features(
86+
entity_df=entity_df,
87+
features=["driver_stats:conv_rate", "driver_stats:avg_daily_trips"],
88+
full_feature_names=False,
89+
).to_df()
90+
91+
assert df["conv_rate"].isnull().all()
92+
assert df["avg_daily_trips"].isnull().all()
93+
94+
expected_df = pd.DataFrame.from_dict(
95+
{
96+
"event_timestamp": [ts - timedelta(hours=3), ts],
97+
"driver_id": [1001, 1002],
98+
"conv_rate": [random.random(), random.random()],
99+
"incorrect_schema": [random.randint(0, 10), random.randint(0, 10)],
100+
"created": [ts, ts],
101+
},
102+
)
103+
with pytest.raises(ValueError):
104+
store.write_to_offline_store(
105+
driver_stats.name, expected_df, allow_registry_cache=False
106+
)
107+
108+
109+
@pytest.mark.integration
110+
@pytest.mark.universal_online_stores
111+
def test_writing_consecutively_to_offline_store(environment, universal_data_sources):
112+
store = environment.feature_store
113+
_, _, data_sources = universal_data_sources
114+
driver_stats = FeatureView(
115+
name="driver_stats",
116+
entities=["driver"],
117+
schema=[
118+
Field(name="avg_daily_trips", dtype=Int32),
119+
Field(name="conv_rate", dtype=Float32),
120+
Field(name="acc_rate", dtype=Float32),
121+
],
122+
source=data_sources.driver,
123+
ttl=timedelta(minutes=10),
124+
)
125+
126+
now = datetime.utcnow()
127+
ts = pd.Timestamp(now, unit="ns")
128+
129+
entity_df = pd.DataFrame.from_dict(
130+
{
131+
"driver_id": [1001, 1001],
132+
"event_timestamp": [ts - timedelta(hours=4), ts - timedelta(hours=3)],
133+
}
134+
)
135+
136+
store.apply([driver(), driver_stats])
137+
df = store.get_historical_features(
138+
entity_df=entity_df,
139+
features=["driver_stats:conv_rate", "driver_stats:avg_daily_trips"],
140+
full_feature_names=False,
141+
).to_df()
142+
143+
assert df["conv_rate"].isnull().all()
144+
assert df["avg_daily_trips"].isnull().all()
145+
146+
first_df = pd.DataFrame.from_dict(
147+
{
148+
"event_timestamp": [ts - timedelta(hours=4), ts - timedelta(hours=3)],
149+
"driver_id": [1001, 1001],
150+
"conv_rate": [random.random(), random.random()],
151+
"acc_rate": [random.random(), random.random()],
152+
"avg_daily_trips": [random.randint(0, 10), random.randint(0, 10)],
153+
"created": [ts, ts],
154+
},
155+
)
156+
store.write_to_offline_store(
157+
driver_stats.name, first_df, allow_registry_cache=False
158+
)
159+
160+
after_write_df = store.get_historical_features(
161+
entity_df=entity_df,
162+
features=["driver_stats:conv_rate", "driver_stats:avg_daily_trips"],
163+
full_feature_names=False,
164+
).to_df()
165+
166+
assert len(after_write_df) == len(first_df)
167+
assert np.where(
168+
after_write_df["conv_rate"].reset_index(drop=True)
169+
== first_df["conv_rate"].reset_index(drop=True)
170+
)
171+
assert np.where(
172+
after_write_df["avg_daily_trips"].reset_index(drop=True)
173+
== first_df["avg_daily_trips"].reset_index(drop=True)
174+
)
175+
176+
second_df = pd.DataFrame.from_dict(
177+
{
178+
"event_timestamp": [ts - timedelta(hours=1), ts],
179+
"driver_id": [1001, 1001],
180+
"conv_rate": [random.random(), random.random()],
181+
"acc_rate": [random.random(), random.random()],
182+
"avg_daily_trips": [random.randint(0, 10), random.randint(0, 10)],
183+
"created": [ts, ts],
184+
},
185+
)
186+
187+
store.write_to_offline_store(
188+
driver_stats.name, second_df, allow_registry_cache=False
189+
)
190+
191+
entity_df = pd.DataFrame.from_dict(
192+
{
193+
"driver_id": [1001, 1001, 1001, 1001],
194+
"event_timestamp": [
195+
ts - timedelta(hours=4),
196+
ts - timedelta(hours=3),
197+
ts - timedelta(hours=1),
198+
ts,
199+
],
200+
}
201+
)
202+
203+
after_write_df = store.get_historical_features(
204+
entity_df=entity_df,
205+
features=[
206+
"driver_stats:conv_rate",
207+
"driver_stats:acc_rate",
208+
"driver_stats:avg_daily_trips",
209+
],
210+
full_feature_names=False,
211+
).to_df()
212+
213+
expected_df = pd.concat([first_df, second_df])
214+
assert len(after_write_df) == len(expected_df)
215+
assert np.where(
216+
after_write_df["conv_rate"].reset_index(drop=True)
217+
== expected_df["conv_rate"].reset_index(drop=True)
218+
)
219+
assert np.where(
220+
after_write_df["acc_rate"].reset_index(drop=True)
221+
== expected_df["acc_rate"].reset_index(drop=True)
222+
)
223+
assert np.where(
224+
after_write_df["avg_daily_trips"].reset_index(drop=True)
225+
== expected_df["avg_daily_trips"].reset_index(drop=True)
226+
)

0 commit comments

Comments
 (0)