Skip to content

Commit abfac01

Browse files
franciscojavierarceofranciscojavierarceo
authored
feat: Adding support for dictionary writes to online store (#4156)
* feat: Adding support for dictionary writes to online store Signed-off-by: Francisco Javier Arceo <[email protected]> * Simple approach Signed-off-by: Francisco Javier Arceo <[email protected]> * lint Signed-off-by: Francisco Javier Arceo <[email protected]> * adding error if both are missing Signed-off-by: Francisco Javier Arceo <[email protected]> * rename dict to input_dict Signed-off-by: Francisco Javier Arceo <[email protected]> * updated input arg to test Signed-off-by: Francisco Javier Arceo <[email protected]> * Renaming function argument Signed-off-by: Francisco Javier Arceo <[email protected]> * updated docstring Signed-off-by: Francisco Javier Arceo <[email protected]> * updated type signature Signed-off-by: Francisco Javier Arceo <[email protected]> * updated type to be more explicit Signed-off-by: Francisco Javier Arceo <[email protected]> --------- Signed-off-by: Francisco Javier Arceo <[email protected]>
1 parent c91dd69 commit abfac01

3 files changed

Lines changed: 162 additions & 1 deletion

File tree

‎sdk/python/feast/errors.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,3 +420,10 @@ def __init__(self, push_source_name: str):
420420
class ReadOnlyRegistryException(Exception):
421421
def __init__(self):
422422
super().__init__("Registry implementation is read-only.")
423+
424+
425+
class DataFrameSerializationError(Exception):
426+
def __init__(self, input_dict: dict):
427+
super().__init__(
428+
f"Failed to serialize the provided dictionary into a pandas DataFrame: {input_dict.keys()}"
429+
)

‎sdk/python/feast/feature_store.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
from feast.dqm.errors import ValidationFailed
5555
from feast.entity import Entity
5656
from feast.errors import (
57+
DataFrameSerializationError,
5758
DataSourceRepeatNamesException,
5859
EntityNotFoundException,
5960
FeatureNameCollisionError,
@@ -1406,7 +1407,8 @@ def push(
14061407
def write_to_online_store(
14071408
self,
14081409
feature_view_name: str,
1409-
df: pd.DataFrame,
1410+
df: Optional[pd.DataFrame] = None,
1411+
inputs: Optional[Union[Dict[str, List[Any]], pd.DataFrame]] = None,
14101412
allow_registry_cache: bool = True,
14111413
):
14121414
"""
@@ -1415,6 +1417,7 @@ def write_to_online_store(
14151417
Args:
14161418
feature_view_name: The feature view to which the dataframe corresponds.
14171419
df: The dataframe to be persisted.
1420+
inputs: Optional the dictionary object to be written
14181421
allow_registry_cache (optional): Whether to allow retrieving feature views from a cached registry.
14191422
"""
14201423
# TODO: restrict this to work with online StreamFeatureViews and validate the FeatureView type
@@ -1426,6 +1429,18 @@ def write_to_online_store(
14261429
feature_view = self.get_feature_view(
14271430
feature_view_name, allow_registry_cache=allow_registry_cache
14281431
)
1432+
if df is not None and inputs is not None:
1433+
raise ValueError("Both df and inputs cannot be provided at the same time.")
1434+
if df is None and inputs is not None:
1435+
if isinstance(inputs, dict):
1436+
try:
1437+
df = pd.DataFrame(inputs)
1438+
except Exception as _:
1439+
raise DataFrameSerializationError(inputs)
1440+
elif isinstance(inputs, pd.DataFrame):
1441+
pass
1442+
else:
1443+
raise ValueError("inputs must be a dictionary or a pandas DataFrame.")
14291444
provider = self._get_provider()
14301445
provider.ingest_df(feature_view, df)
14311446

‎sdk/python/tests/unit/online_store/test_online_writes.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2022 The Feast Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
import tempfile
17+
import unittest
18+
from datetime import datetime, timedelta
19+
from typing import Any, Dict
20+
21+
from feast import Entity, FeatureStore, FeatureView, FileSource, RepoConfig
22+
from feast.driver_test_data import create_driver_hourly_stats_df
23+
from feast.field import Field
24+
from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
25+
from feast.on_demand_feature_view import on_demand_feature_view
26+
from feast.types import Float32, Float64, Int64
27+
28+
29+
class TestOnlineWrites(unittest.TestCase):
30+
def setUp(self):
31+
with tempfile.TemporaryDirectory() as data_dir:
32+
self.store = FeatureStore(
33+
config=RepoConfig(
34+
project="test_write_to_online_store",
35+
registry=os.path.join(data_dir, "registry.db"),
36+
provider="local",
37+
entity_key_serialization_version=2,
38+
online_store=SqliteOnlineStoreConfig(
39+
path=os.path.join(data_dir, "online.db")
40+
),
41+
)
42+
)
43+
44+
# Generate test data.
45+
end_date = datetime.now().replace(microsecond=0, second=0, minute=0)
46+
start_date = end_date - timedelta(days=15)
47+
48+
driver_entities = [1001, 1002, 1003, 1004, 1005]
49+
driver_df = create_driver_hourly_stats_df(
50+
driver_entities, start_date, end_date
51+
)
52+
driver_stats_path = os.path.join(data_dir, "driver_stats.parquet")
53+
driver_df.to_parquet(
54+
path=driver_stats_path, allow_truncated_timestamps=True
55+
)
56+
57+
driver = Entity(name="driver", join_keys=["driver_id"])
58+
59+
driver_stats_source = FileSource(
60+
name="driver_hourly_stats_source",
61+
path=driver_stats_path,
62+
timestamp_field="event_timestamp",
63+
created_timestamp_column="created",
64+
)
65+
66+
driver_stats_fv = FeatureView(
67+
name="driver_hourly_stats",
68+
entities=[driver],
69+
ttl=timedelta(days=0),
70+
schema=[
71+
Field(name="conv_rate", dtype=Float32),
72+
Field(name="acc_rate", dtype=Float32),
73+
Field(name="avg_daily_trips", dtype=Int64),
74+
],
75+
online=True,
76+
source=driver_stats_source,
77+
)
78+
79+
@on_demand_feature_view(
80+
sources=[driver_stats_fv[["conv_rate", "acc_rate"]]],
81+
schema=[Field(name="conv_rate_plus_acc", dtype=Float64)],
82+
mode="python",
83+
)
84+
def test_view(inputs: Dict[str, Any]) -> Dict[str, Any]:
85+
output: Dict[str, Any] = {
86+
"conv_rate_plus_acc": [
87+
conv_rate + acc_rate
88+
for conv_rate, acc_rate in zip(
89+
inputs["conv_rate"], inputs["acc_rate"]
90+
)
91+
]
92+
}
93+
return output
94+
95+
self.store.apply(
96+
[
97+
driver,
98+
driver_stats_source,
99+
driver_stats_fv,
100+
test_view,
101+
]
102+
)
103+
self.store.write_to_online_store(
104+
feature_view_name="driver_hourly_stats", df=driver_df
105+
)
106+
# This will give the intuitive structure of the data as:
107+
# {"driver_id": [..], "conv_rate": [..], "acc_rate": [..], "avg_daily_trips": [..]}
108+
driver_dict = driver_df.to_dict(orient="list")
109+
self.store.write_to_online_store(
110+
feature_view_name="driver_hourly_stats",
111+
inputs=driver_dict,
112+
)
113+
114+
def test_online_retrieval(self):
115+
entity_rows = [
116+
{
117+
"driver_id": 1001,
118+
}
119+
]
120+
121+
online_python_response = self.store.get_online_features(
122+
entity_rows=entity_rows,
123+
features=[
124+
"driver_hourly_stats:conv_rate",
125+
"driver_hourly_stats:acc_rate",
126+
"test_view:conv_rate_plus_acc",
127+
],
128+
).to_dict()
129+
130+
assert len(online_python_response) == 4
131+
assert all(
132+
key in online_python_response.keys()
133+
for key in [
134+
"driver_id",
135+
"acc_rate",
136+
"conv_rate",
137+
"conv_rate_plus_acc",
138+
]
139+
)

0 commit comments

Comments
 (0)