Skip to content

Commit f2b4eb9

Browse files
pushkarmoipushkarmoi
authored
feat: Feast/IKV online store contrib plugin integration (#4068)
1 parent b66baa4 commit f2b4eb9

File tree

6 files changed

+320
-0
lines changed

6 files changed

+320
-0
lines changed

‎sdk/python/feast/cli.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,7 @@ def materialize_incremental_command(ctx: click.Context, end_ts: str, views: List
595595
"cassandra",
596596
"rockset",
597597
"hazelcast",
598+
"ikv",
598599
],
599600
case_sensitive=False,
600601
),

‎sdk/python/feast/infra/online_stores/contrib/ikv_online_store/__init__.py‎

Whitespace-only changes.

‎sdk/python/feast/infra/online_stores/contrib/ikv_online_store/ikv.py‎

Lines changed: 300 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,300 @@
1+
from datetime import datetime
2+
from typing import (
3+
Any,
4+
Callable,
5+
Dict,
6+
Iterator,
7+
List,
8+
Literal,
9+
Optional,
10+
Sequence,
11+
Tuple,
12+
)
13+
14+
from ikvpy.client import IKVReader, IKVWriter
15+
from ikvpy.clientoptions import ClientOptions, ClientOptionsBuilder
16+
from ikvpy.document import IKVDocument, IKVDocumentBuilder
17+
from ikvpy.factory import create_new_reader, create_new_writer
18+
from pydantic import StrictStr
19+
20+
from feast import Entity, FeatureView, utils
21+
from feast.infra.online_stores.helpers import compute_entity_id
22+
from feast.infra.online_stores.online_store import OnlineStore
23+
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
24+
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
25+
from feast.repo_config import FeastConfigBaseModel, RepoConfig
26+
from feast.usage import log_exceptions_and_usage
27+
28+
PRIMARY_KEY_FIELD_NAME: str = "_entity_key"
29+
EVENT_CREATION_TIMESTAMP_FIELD_NAME: str = "_event_timestamp"
30+
CREATION_TIMESTAMP_FIELD_NAME: str = "_created_timestamp"
31+
32+
33+
class IKVOnlineStoreConfig(FeastConfigBaseModel):
34+
"""Online store config for IKV store"""
35+
36+
type: Literal["ikv"] = "ikv"
37+
"""Online store type selector"""
38+
39+
account_id: StrictStr
40+
"""(Required) IKV account id"""
41+
42+
account_passkey: StrictStr
43+
"""(Required) IKV account passkey"""
44+
45+
store_name: StrictStr
46+
"""(Required) IKV store name"""
47+
48+
mount_directory: Optional[StrictStr] = None
49+
"""(Required only for reader) IKV mount point i.e. directory for storing IKV data locally."""
50+
51+
52+
class IKVOnlineStore(OnlineStore):
53+
"""
54+
IKV (inlined.io key value) store implementation of the online store interface.
55+
"""
56+
57+
# lazy initialization
58+
_reader: Optional[IKVReader] = None
59+
_writer: Optional[IKVWriter] = None
60+
61+
@log_exceptions_and_usage(online_store="ikv")
62+
def online_write_batch(
63+
self,
64+
config: RepoConfig,
65+
table: FeatureView,
66+
data: List[
67+
Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
68+
],
69+
progress: Optional[Callable[[int], Any]],
70+
) -> None:
71+
"""
72+
Writes a batch of feature rows to the online store.
73+
74+
If a tz-naive timestamp is passed to this method, it is assumed to be UTC.
75+
76+
Args:
77+
config: The config for the current feature store.
78+
table: Feature view to which these feature rows correspond.
79+
data: A list of quadruplets containing feature data. Each quadruplet contains an entity
80+
key, a dict containing feature values, an event timestamp for the row, and the created
81+
timestamp for the row if it exists.
82+
progress: Function to be called once a batch of rows is written to the online store, used
83+
to show progress.
84+
"""
85+
# update should have been called before
86+
if self._writer is None:
87+
return
88+
89+
for entity_key, features, event_timestamp, _ in data:
90+
entity_id: str = compute_entity_id(
91+
entity_key,
92+
entity_key_serialization_version=config.entity_key_serialization_version,
93+
)
94+
document: IKVDocument = IKVOnlineStore._create_document(
95+
entity_id, table, features, event_timestamp
96+
)
97+
self._writer.upsert_fields(document)
98+
if progress:
99+
progress(1)
100+
101+
@log_exceptions_and_usage(online_store="ikv")
102+
def online_read(
103+
self,
104+
config: RepoConfig,
105+
table: FeatureView,
106+
entity_keys: List[EntityKeyProto],
107+
requested_features: Optional[List[str]] = None,
108+
) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
109+
"""
110+
Reads features values for the given entity keys.
111+
112+
Args:
113+
config: The config for the current feature store.
114+
table: The feature view whose feature values should be read.
115+
entity_keys: The list of entity keys for which feature values should be read.
116+
requested_features: The list of features that should be read.
117+
118+
Returns:
119+
A list of the same length as entity_keys. Each item in the list is a tuple where the first
120+
item is the event timestamp for the row, and the second item is a dict mapping feature names
121+
to values, which are returned in proto format.
122+
"""
123+
if not len(entity_keys):
124+
return []
125+
126+
# create IKV primary keys
127+
primary_keys = [
128+
compute_entity_id(ek, config.entity_key_serialization_version)
129+
for ek in entity_keys
130+
]
131+
132+
# create IKV field names
133+
if requested_features is None:
134+
requested_features = []
135+
136+
field_names: List[Optional[str]] = [None] * (1 + len(requested_features))
137+
field_names[0] = EVENT_CREATION_TIMESTAMP_FIELD_NAME
138+
for i, fn in enumerate(requested_features):
139+
field_names[i + 1] = IKVOnlineStore._create_ikv_field_name(table, fn)
140+
141+
assert self._reader is not None
142+
value_iter = self._reader.multiget_bytes_values(
143+
bytes_primary_keys=[],
144+
str_primary_keys=primary_keys,
145+
field_names=field_names,
146+
)
147+
148+
# decode results
149+
return [
150+
IKVOnlineStore._decode_fields_for_primary_key(
151+
requested_features, value_iter
152+
)
153+
for _ in range(0, len(primary_keys))
154+
]
155+
156+
@staticmethod
157+
def _decode_fields_for_primary_key(
158+
requested_features: List[str], value_iter: Iterator[Optional[bytes]]
159+
) -> Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]:
160+
# decode timestamp
161+
dt: Optional[datetime] = None
162+
dt_bytes = next(value_iter)
163+
if dt_bytes:
164+
dt = datetime.fromisoformat(str(dt_bytes, "utf-8"))
165+
166+
# decode other features
167+
features = {}
168+
for requested_feature in requested_features:
169+
value_proto_bytes: Optional[bytes] = next(value_iter)
170+
if value_proto_bytes:
171+
value_proto = ValueProto()
172+
value_proto.ParseFromString(value_proto_bytes)
173+
features[requested_feature] = value_proto
174+
175+
return dt, features
176+
177+
# called before any read/write requests are issued
178+
@log_exceptions_and_usage(online_store="ikv")
179+
def update(
180+
self,
181+
config: RepoConfig,
182+
tables_to_delete: Sequence[FeatureView],
183+
tables_to_keep: Sequence[FeatureView],
184+
entities_to_delete: Sequence[Entity],
185+
entities_to_keep: Sequence[Entity],
186+
partial: bool,
187+
):
188+
"""
189+
Reconciles cloud resources with the specified set of Feast objects.
190+
191+
Args:
192+
config: The config for the current feature store.
193+
tables_to_delete: Feature views whose corresponding infrastructure should be deleted.
194+
tables_to_keep: Feature views whose corresponding infrastructure should not be deleted, and
195+
may need to be updated.
196+
entities_to_delete: Entities whose corresponding infrastructure should be deleted.
197+
entities_to_keep: Entities whose corresponding infrastructure should not be deleted, and
198+
may need to be updated.
199+
partial: If true, tables_to_delete and tables_to_keep are not exhaustive lists, so
200+
infrastructure corresponding to other feature views should be not be touched.
201+
"""
202+
self._init_clients(config=config)
203+
assert self._writer is not None
204+
205+
# note: we assume tables_to_keep does not overlap with tables_to_delete
206+
207+
for feature_view in tables_to_delete:
208+
# each field in an IKV document is prefixed by the feature-view's name
209+
self._writer.drop_fields_by_name_prefix([feature_view.name])
210+
211+
@log_exceptions_and_usage(online_store="ikv")
212+
def teardown(
213+
self,
214+
config: RepoConfig,
215+
tables: Sequence[FeatureView],
216+
entities: Sequence[Entity],
217+
):
218+
"""
219+
Tears down all cloud resources for the specified set of Feast objects.
220+
221+
Args:
222+
config: The config for the current feature store.
223+
tables: Feature views whose corresponding infrastructure should be deleted.
224+
entities: Entities whose corresponding infrastructure should be deleted.
225+
"""
226+
self._init_clients(config=config)
227+
assert self._writer is not None
228+
229+
# drop fields corresponding to this feature-view
230+
for feature_view in tables:
231+
self._writer.drop_fields_by_name_prefix([feature_view.name])
232+
233+
# shutdown clients
234+
self._writer.shutdown()
235+
self._writer = None
236+
237+
if self._reader is not None:
238+
self._reader.shutdown()
239+
self._reader = None
240+
241+
@staticmethod
242+
def _create_ikv_field_name(feature_view: FeatureView, feature_name: str) -> str:
243+
return "{}_{}".format(feature_view.name, feature_name)
244+
245+
@staticmethod
246+
def _create_document(
247+
entity_id: str,
248+
feature_view: FeatureView,
249+
values: Dict[str, ValueProto],
250+
event_timestamp: datetime,
251+
) -> IKVDocument:
252+
"""Converts feast key-value pairs into an IKV document."""
253+
254+
# initialie builder by inserting primary key and row creation timestamp
255+
event_timestamp_str: str = utils.make_tzaware(event_timestamp).isoformat()
256+
builder = (
257+
IKVDocumentBuilder()
258+
.put_string_field(PRIMARY_KEY_FIELD_NAME, entity_id)
259+
.put_bytes_field(
260+
EVENT_CREATION_TIMESTAMP_FIELD_NAME, event_timestamp_str.encode("utf-8")
261+
)
262+
)
263+
264+
for feature_name, feature_value in values.items():
265+
field_name = IKVOnlineStore._create_ikv_field_name(
266+
feature_view, feature_name
267+
)
268+
builder.put_bytes_field(field_name, feature_value.SerializeToString())
269+
270+
return builder.build()
271+
272+
def _init_clients(self, config: RepoConfig):
273+
"""Initializes (if required) reader/writer ikv clients."""
274+
online_config = config.online_store
275+
assert isinstance(online_config, IKVOnlineStoreConfig)
276+
client_options = IKVOnlineStore._config_to_client_options(online_config)
277+
278+
# initialize writer
279+
if self._writer is None:
280+
self._writer = create_new_writer(client_options)
281+
282+
# initialize reader, iff mount_dir is specified
283+
if self._reader is None:
284+
if online_config.mount_directory and len(online_config.mount_directory) > 0:
285+
self._reader = create_new_reader(client_options)
286+
287+
@staticmethod
288+
def _config_to_client_options(config: IKVOnlineStoreConfig) -> ClientOptions:
289+
"""Utility for IKVOnlineStoreConfig to IKV ClientOptions conversion."""
290+
builder = (
291+
ClientOptionsBuilder()
292+
.with_account_id(config.account_id)
293+
.with_account_passkey(config.account_passkey)
294+
.with_store_name(config.store_name)
295+
)
296+
297+
if config.mount_directory and len(config.mount_directory) > 0:
298+
builder = builder.with_mount_directory(config.mount_directory)
299+
300+
return builder.build()

‎sdk/python/feast/repo_config.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
"mysql": "feast.infra.online_stores.contrib.mysql_online_store.mysql.MySQLOnlineStore",
6464
"rockset": "feast.infra.online_stores.contrib.rockset_online_store.rockset.RocksetOnlineStore",
6565
"hazelcast": "feast.infra.online_stores.contrib.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore",
66+
"ikv": "feast.infra.online_stores.contrib.ikv_online_store.ikv.IKVOnlineStore",
6667
}
6768

6869
OFFLINE_STORE_CLASS_FOR_TYPE = {

‎sdk/python/tests/integration/feature_repos/repo_configuration.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,14 @@
100100
"host": os.getenv("ROCKSET_APISERVER", "api.rs2.usw2.rockset.com"),
101101
}
102102

103+
IKV_CONFIG = {
104+
"type": "ikv",
105+
"account_id": os.getenv("IKV_ACCOUNT_ID", ""),
106+
"account_passkey": os.getenv("IKV_ACCOUNT_PASSKEY", ""),
107+
"store_name": os.getenv("IKV_STORE_NAME", ""),
108+
"mount_directory": os.getenv("IKV_MOUNT_DIR", ""),
109+
}
110+
103111
OFFLINE_STORE_TO_PROVIDER_CONFIG: Dict[str, Tuple[str, Type[DataSourceCreator]]] = {
104112
"file": ("local", FileDataSourceCreator),
105113
"bigquery": ("gcp", BigQueryDataSourceCreator),
@@ -139,6 +147,11 @@
139147
# containerized version of Rockset.
140148
# AVAILABLE_ONLINE_STORES["rockset"] = (ROCKSET_CONFIG, None)
141149

150+
# Uncomment to test using private IKV account. Currently not enabled as
151+
# there is no dedicated IKV instance for CI testing and there is no
152+
# containerized version of IKV.
153+
# AVAILABLE_ONLINE_STORES["ikv"] = (IKV_CONFIG, None)
154+
142155

143156
full_repo_configs_module = os.environ.get(FULL_REPO_CONFIGS_MODULE_ENV_NAME)
144157
if full_repo_configs_module is not None:

‎setup.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,10 @@
130130
"rockset>=1.0.3",
131131
]
132132

133+
IKV_REQUIRED = [
134+
"ikvpy>=0.0.23",
135+
]
136+
133137
HAZELCAST_REQUIRED = [
134138
"hazelcast-python-client>=5.1",
135139
]
@@ -372,6 +376,7 @@ def run(self):
372376
"rockset": ROCKSET_REQUIRED,
373377
"ibis": IBIS_REQUIRED,
374378
"duckdb": DUCKDB_REQUIRED,
379+
"ikv": IKV_REQUIRED
375380
},
376381
include_package_data=True,
377382
license="Apache",

0 commit comments

Comments
 (0)