Visitar URL original
feat: Add ScyllaDB online store with vector search by zseta · Pull Request #6508 · feast-dev/feast · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
lint fix
Signed-off-by: Attila Toth <attila.toth@scylladb.com>
  • Loading branch information
zseta committed Jun 18, 2026
commit de1a61b22ffa7fda661549910f427dde1c7e1faf
Original file line number Diff line number Diff line change
Expand Up @@ -93,8 +93,7 @@
DROP_VECTOR_TABLE_CQL = "DROP TABLE IF EXISTS {fqvectable};"

INSERT_VEC_CQL = (
"INSERT INTO {fqvectable} (entity_key, vector_value, event_ts)"
" VALUES (?, ?, ?);"
"INSERT INTO {fqvectable} (entity_key, vector_value, event_ts) VALUES (?, ?, ?);"
)

CREATE_VECTOR_INDEX_CQL = (
Expand Down Expand Up @@ -391,7 +390,7 @@ def _create_table(
logger.info("Creating table %s.", fqtable)
session.execute(CREATE_TABLE_CQL.format(fqtable=fqtable))

for feat_name, dim, sim_func in (vec_features or []):
for feat_name, dim, sim_func in vec_features or []:
fqvectable = self._fq_vec_table_name(
self._keyspace, project, table, feat_name
)
Expand All @@ -417,7 +416,7 @@ def _drop_table(
vec_features: Optional[List[Tuple[str, int, str]]] = None,
) -> None:
session = self._get_session(config)
for feat_name, _dim, _sim in (vec_features or []):
for feat_name, _dim, _sim in vec_features or []:
fqvectable = self._fq_vec_table_name(
self._keyspace, project, table, feat_name
)
Expand Down Expand Up @@ -466,7 +465,9 @@ def teardown(
assert isinstance(online_store_config, ScyllaDBOnlineStoreConfig)

for table in tables:
vec_features = _get_vector_features(table, online_store_config.vector_similarity_function)
vec_features = _get_vector_features(
table, online_store_config.vector_similarity_function
)
self._drop_table(config, project, table, vec_features=vec_features)

# ------------------------------------------------------------------
Expand Down Expand Up @@ -542,10 +543,14 @@ def _main_rows() -> Iterable[Tuple]:

# --- vector table rows (one per entity per vector feature) ---
for feat_name, _dim, _sim in vec_features:
fqvectable = self._fq_vec_table_name(self._keyspace, project, table, feat_name)
fqvectable = self._fq_vec_table_name(
self._keyspace, project, table, feat_name
)
vec_insert_cql = INSERT_VEC_CQL.format(fqvectable=fqvectable)
if vec_insert_cql not in self._prepared_statements:
self._prepared_statements[vec_insert_cql] = session.prepare(vec_insert_cql)
self._prepared_statements[vec_insert_cql] = session.prepare(
vec_insert_cql
)
vec_insert_stmt = self._prepared_statements[vec_insert_cql]

def _vec_rows(
Expand Down Expand Up @@ -715,7 +720,9 @@ def retrieve_online_documents_v2(
)
sim_expr = sim_expr_template

fqvectable = self._fq_vec_table_name(self._keyspace, project, table, vec_feature)
fqvectable = self._fq_vec_table_name(
self._keyspace, project, table, vec_feature
)
ann_cql = ANN_SELECT_CQL.format(
sim_func_call=sim_expr,
fqvectable=fqvectable,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,9 @@ def _make_entity_key(val: str) -> EntityKeyProto:
)


def _make_feature_view(name: str, with_vector: bool = False, ttl: timedelta = None) -> FeatureView:
def _make_feature_view(
name: str, with_vector: bool = False, ttl: timedelta = None
) -> FeatureView:
source = FileSource(path="dummy.parquet", timestamp_field="event_timestamp")
schema: List[Field] = [Field(name="score", dtype=Array(Float32))]
if with_vector:
Expand Down Expand Up @@ -181,14 +183,15 @@ def test_missing_key_returns_none(docker_config):


def test_multiple_features_roundtrip(docker_config):
"""Multiple features of different types all round-trip with the correct value.
"""
"""Multiple features of different types all round-trip with the correct value."""
store = ScyllaDBOnlineStore()
cfg = docker_config
source = FileSource(path="dummy.parquet", timestamp_field="event_timestamp")
fv = FeatureView(
name="test_multi_features",
entities=[Entity(name="item_id", join_keys=["item_id"], value_type=ValueType.STRING)],
entities=[
Entity(name="item_id", join_keys=["item_id"], value_type=ValueType.STRING)
],
schema=[
Field(name="score", dtype=Array(Float32)),
Field(name="priority", dtype=Int64),
Expand Down Expand Up @@ -242,7 +245,11 @@ def test_multiple_entities(docker_config):
batch = [
(
ek,
{"score": ValueProto(float_list_val=FloatList(val=[float(i) * 0.1 + 0.1]))},
{
"score": ValueProto(
float_list_val=FloatList(val=[float(i) * 0.1 + 0.1])
)
},
_utc_now(),
None,
)
Expand Down Expand Up @@ -272,13 +279,27 @@ def test_overwrite_uses_latest_value(docker_config):
store.online_write_batch(
cfg,
fv,
[(ek, {"score": ValueProto(float_list_val=FloatList(val=[0.1]))}, _utc_now(), None)],
[
(
ek,
{"score": ValueProto(float_list_val=FloatList(val=[0.1]))},
_utc_now(),
None,
)
],
None,
)
store.online_write_batch(
cfg,
fv,
[(ek, {"score": ValueProto(float_list_val=FloatList(val=[0.9]))}, _utc_now(), None)],
[
(
ek,
{"score": ValueProto(float_list_val=FloatList(val=[0.9]))},
_utc_now(),
None,
)
],
None,
)
results = store.online_read(cfg, fv, [ek])
Expand Down Expand Up @@ -308,7 +329,14 @@ def test_event_timestamp_returned(docker_config):
store.online_write_batch(
cfg,
fv,
[(ek, {"score": ValueProto(float_list_val=FloatList(val=[0.5]))}, write_ts, None)],
[
(
ek,
{"score": ValueProto(float_list_val=FloatList(val=[0.5]))},
write_ts,
None,
)
],
None,
)
results = store.online_read(cfg, fv, [ek])
Expand Down Expand Up @@ -455,8 +483,7 @@ def test_ttl_expiry(docker_config):
assert len(results) == 1
ts_after, feats_after = results[0]
assert feats_after is None, (
"Row should have expired and return None, "
f"but got features: {feats_after}"
f"Row should have expired and return None, but got features: {feats_after}"
)
finally:
store.teardown(cfg, [fv], [])
Loading