forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathstores.py
More file actions
98 lines (80 loc) · 3.09 KB
/
Copy pathstores.py
File metadata and controls
98 lines (80 loc) · 3.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
from feast.types import FeatureRow_pb2 as FeatureRowProto
from feast.core import FeatureSet_pb2 as FeatureSetProto
import sqlite3
from typing import Dict, List
from feast.entity import Entity
from feast.value_type import ValueType
from feast.feature_set import FeatureSet, Feature
from feast.types import (
FeatureRow_pb2 as FeatureRowProto,
Field_pb2 as FieldProto,
Value_pb2 as ValueProto,
)
from google.protobuf.timestamp_pb2 import Timestamp
class Database:
pass
class SQLiteDatabase(Database):
def __init__(self):
self._conn = sqlite3.connect(":memory:")
self._c = self._conn.cursor()
def register_feature_set(self, feature_set: FeatureSetProto.FeatureSetSpec):
query = build_sqlite_create_table_query(feature_set)
print(query)
self._c.execute(query)
self._c.execute("SELECT name FROM sqlite_master WHERE type='table';")
available_table = self._c.fetchall()
print(available_table)
def upsert_feature_row(
self,
feature_set: FeatureSetProto.FeatureSetSpec,
feature_row: FeatureRowProto.FeatureRow,
):
values = (feature_row.event_timestamp,)
for entity in list(feature_set.entities):
values = values + (get_feature_row_value_by_name(feature_row, entity.name),)
values = values + (feature_row.SerializeToString(),)
self._c.execute(build_sqlite_insert_feature_row_query(feature_set), values)
def build_sqlite_create_table_query(feature_set: FeatureSetProto.FeatureSetSpec):
query = (
"""
CREATE TABLE IF NOT EXISTS {} (
{}
PRIMARY KEY ({})
);
"""
).format(
get_table_name(feature_set),
" ".join([column + " text NOT NULL," for column in get_columns(feature_set)]),
", ".join(
get_columns(feature_set)[1:]
), # exclude event_timestamp column for online stores
)
# Hyphens become three underscores
query = query.replace("-", "___")
return query
def build_sqlite_insert_feature_row_query(feature_set: FeatureSetProto.FeatureSetSpec):
return """
INSERT OR REPLACE INTO {} ({})
VALUES(?,?,?,?,?,?)
""".format(
get_table_name(feature_set), ",".join(get_columns(feature_set))
)
def get_columns(feature_set: FeatureSetProto.FeatureSetSpec) -> List[str]:
return (
["event_timestamp"]
+ [field.name for field in list(feature_set.entities)]
+ ["value"]
)
def get_feature_row_value_by_name(feature_row, name):
values = [field.value for field in list(feature_row.fields) if field.name == name]
if len(values) != 1:
raise Exception(
"Invalid number of features with name {} in feature row {}".format(
name, feature_row.name
)
)
return values[0]
def get_table_name(feature_set: FeatureSetProto.FeatureSetSpec) -> str:
if not feature_set.name and not feature_set.version:
raise ValueError("Feature set name or version is missing")
return (feature_set.name + "_" + str(feature_set.version)).replace("-", "___")