Support Alternator in Vector Store

Cargo.toml

@@ -47,7 +47,7 @@ rcgen = "0.14.5"
 regex = "1.11.1"
 reqwest = { version = "0.12.15", default-features = false, features = ["json", "rustls-tls"] }
 scylla = { version = "1.5.0", features = ["time-03", "rustls-023", "metrics"] }
-scylla-cdc = "0.6.1"
+scylla-cdc = "0.6.3"
 scylla-proxy = "0.0.5"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"

crates/vector-store/src/db.rs

@@ -23,6 +23,7 @@ use crate::SpaceType;
 use crate::TableName;
 use crate::db_index;
 use crate::db_index::DbIndex;
+use crate::db_index_backend;
 use crate::internals::Internals;
 use crate::internals::InternalsExt;
 use crate::node_state::Event;
@@ -46,7 +47,6 @@ use scylla::statement::prepared::PreparedStatement;
 use scylla::value::CqlTimeuuid;
 use secrecy::ExposeSecret;
 use std::collections::BTreeMap;
-use std::num::NonZeroUsize;
 use std::sync::Arc;
 use std::time::Duration;
 use tap::Pipe;
@@ -110,6 +110,7 @@ pub enum Db {
         keyspace: KeyspaceName,
         table: TableName,
         target_column: ColumnName,
+        index: IndexName,
         tx: oneshot::Sender<GetIndexTargetTypeR>,
     },
 
@@ -152,6 +153,7 @@ pub(crate) trait DbExt {
         keyspace: KeyspaceName,
         table: TableName,
         target_column: ColumnName,
+        index: IndexName,
     ) -> GetIndexTargetTypeR;
 
     async fn get_index_params(
@@ -205,12 +207,14 @@ impl DbExt for mpsc::Sender<Db> {
         keyspace: KeyspaceName,
         table: TableName,
         target_column: ColumnName,
+        index: IndexName,
     ) -> GetIndexTargetTypeR {
         let (tx, rx) = oneshot::channel();
         self.send(Db::GetIndexTargetType {
             keyspace,
             table,
             target_column,
+            index,
             tx,
         })
         .await?;
@@ -437,11 +441,12 @@ async fn process(
             keyspace,
             table,
             target_column,
+            index,
             tx,
         } => tx
             .send(
                 statements
-                    .get_index_target_type(keyspace, table, target_column)
+                    .get_index_target_type(keyspace, table, target_column, index)
                     .await,
             )
             .unwrap_or_else(|_| trace!("process: Db::GetIndexTargetType: unable to send response")),
@@ -756,29 +761,27 @@ impl Statements {
         keyspace: KeyspaceName,
         table: TableName,
         target_column: ColumnName,
+        index: IndexName,
     ) -> GetIndexTargetTypeR {
         let session = self
             .session_rx
             .borrow()
             .clone()
             .ok_or_else(|| anyhow::anyhow!("No active session"))?;
-        Ok(session
-            .execute_iter(
-                self.st_get_index_target_type.clone(),
-                (keyspace, table, target_column),
-            )
-            .await?
-            .rows_stream::<(String,)>()?
-            .try_next()
-            .await?
-            .and_then(|(typ,)| {
-                self.re_get_index_target_type
-                    .captures(&typ)
-                    .and_then(|captures| captures["dimensions"].parse::<usize>().ok())
-            })
-            .and_then(|dimensions| {
-                NonZeroUsize::new(dimensions).map(|dimensions| dimensions.into())
-            }))
+
+        db_index_backend::get_dimensions(
+            &target_column,
+            &session,
+            &self.st_get_index_target_type,
+            &self.re_get_index_target_type,
+            &self.st_get_index_options,
+            db_index_backend::IndexLocation {
+                keyspace,
+                table,
+                index,
+            },
+        )
+        .await
     }
 
     const ST_GET_INDEX_OPTIONS: &str = "
@@ -1017,6 +1020,7 @@ pub(crate) mod tests {
             keyspace: KeyspaceName,
             table: TableName,
             target_column: ColumnName,
+            index: IndexName,
             tx: oneshot::Sender<GetIndexTargetTypeR>,
         ) -> impl Future<Output = ()> + Send + 'static;
 
@@ -1065,9 +1069,10 @@ pub(crate) mod tests {
                             keyspace,
                             table,
                             target_column,
+                            index,
                             tx,
                         } => {
-                            sim.get_index_target_type(keyspace, table, target_column, tx)
+                            sim.get_index_target_type(keyspace, table, target_column, index, tx)
                                 .await
                         }
 

crates/vector-store/src/db_cdc.rs

@@ -8,6 +8,7 @@ use crate::ColumnName;
 use crate::Config;
 use crate::DbEmbedding;
 use crate::IndexMetadata;
+use crate::db_index_backend::DbIndexBackend;
 use crate::internals::Internals;
 use crate::internals::InternalsExt;
 use ::time::Date;
@@ -21,7 +22,6 @@ use anyhow::bail;
 use async_trait::async_trait;
 use futures::FutureExt;
 use scylla::client::session::Session;
-use scylla::value::CqlValue;
 use scylla_cdc::consumer::CDCRow;
 use scylla_cdc::consumer::Consumer;
 use scylla_cdc::consumer::ConsumerFactory;
@@ -484,7 +484,7 @@ fn spawn_handler_task(
 
 struct CdcConsumerData {
     primary_key_columns: Vec<ColumnName>,
-    target_column: ColumnName,
+    backend: DbIndexBackend,
     tx: mpsc::Sender<(DbEmbedding, Option<AsyncInProgress>)>,
     gregorian_epoch: PrimitiveDateTime,
 }
@@ -499,28 +499,16 @@ impl Consumer for CdcConsumer {
             return Ok(());
         }
 
-        let target_column = self.0.target_column.as_ref();
-        if !row.column_deletable(target_column) {
-            bail!("CDC error: target column {target_column} should be deletable");
+        let source = &self.0.backend;
+        let column = source.vector_column_name();
+        if !row.column_deletable(column) {
+            bail!("CDC error: column {column} should be deletable");
         }
-
         let embedding = row
-            .take_value(target_column)
-            .map(|value| {
-                let CqlValue::Vector(value) = value else {
-                    bail!("CDC error: target column {target_column} should be VECTOR type");
-                };
-                value
-                    .into_iter()
-                    .map(|value| {
-                        value.as_float().ok_or(anyhow!(
-                            "CDC error: target column {target_column} should be VECTOR<float> type"
-                        ))
-                    })
-                    .collect::<anyhow::Result<Vec<_>>>()
-            })
+            .take_value(column)
+            .map(|v| source.extract_vector(v))
             .transpose()?
-            .map(|embedding| embedding.into());
+            .flatten();
 
         let primary_key = self
             .0
@@ -603,9 +591,11 @@ impl CdcConsumerFactory {
             Time::MIDNIGHT,
         );
 
+        let backend = DbIndexBackend::from(metadata);
+
         Ok(Self(Arc::new(CdcConsumerData {
             primary_key_columns,
-            target_column: metadata.target_column.clone(),
+            backend,
             tx,
             gregorian_epoch,
         })))