diff --git a/Cargo.lock b/Cargo.lock index abb8cc27..748a12fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3188,9 +3188,9 @@ dependencies = [ [[package]] name = "scylla-cdc" -version = "0.6.1" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1997e84abfbf7b4312271fe168b9a6d7019d108397346f0f79cd743b711ee80" +checksum = "b4067bc310296f66cff0ca564d32aa2a22a8e8a6b1263a9a195ec01897f87095" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 144b5293..58caee5a 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/vector-store/src/db.rs b/crates/vector-store/src/db.rs index ea92e0dd..f49da00e 100644 --- a/crates/vector-store/src/db.rs +++ b/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, }, @@ -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 { 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::().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, ) -> impl Future + 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 } diff --git a/crates/vector-store/src/db_cdc.rs b/crates/vector-store/src/db_cdc.rs index b3e072c2..c7768f1a 100644 --- a/crates/vector-store/src/db_cdc.rs +++ b/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, - target_column: ColumnName, + backend: DbIndexBackend, tx: mpsc::Sender<(DbEmbedding, Option)>, 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 type" - )) - }) - .collect::>>() - }) + .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, }))) diff --git a/crates/vector-store/src/db_index.rs b/crates/vector-store/src/db_index.rs index 9d18cfb7..15e74b3c 100644 --- a/crates/vector-store/src/db_index.rs +++ b/crates/vector-store/src/db_index.rs @@ -8,13 +8,15 @@ use crate::ColumnName; use crate::Config; use crate::DbEmbedding; use crate::IndexMetadata; -use crate::KeyspaceName; +use crate::KeyspaceIdentifier; use crate::Percentage; use crate::Progress; -use crate::TableName; +use crate::TableIdentifier; use crate::Timestamp; +use crate::Vector; use crate::db_cdc; use crate::db_cdc::CdcReaderConfig; +use crate::db_index_backend; use crate::internals::Internals; use crate::invariant_key::InvariantKey; use crate::node_state::Event; @@ -36,6 +38,7 @@ use scylla::routing::Token; use scylla::statement::prepared::PreparedStatement; use scylla::value::CqlValue; use scylla::value::Row; +use scylla_cdc::CqlIdentifier; use std::collections::HashMap; use std::iter; use std::num::NonZeroUsize; @@ -145,6 +148,18 @@ pub(crate) async fn new( let (tx_index, mut rx_index) = mpsc::channel(CHANNEL_SIZE); let (tx_embeddings, rx_embeddings) = mpsc::channel(CHANNEL_SIZE); + // Wait for initial session to create statements. + let mut statements_session_rx = session_rx.clone(); + while statements_session_rx.borrow().is_none() { + if statements_session_rx.changed().await.is_err() { + return Err(anyhow::anyhow!( + "Session sender dropped before initialization" + )); + } + } + + let statements = Arc::new(Statements::new(statements_session_rx, metadata.clone()).await?); + // Create wide-framed CDC actor let cdc_wide = db_cdc::new( config_rx.clone(), @@ -174,18 +189,6 @@ pub(crate) async fn new( cdc_error_notify.notify_one(); }); - // Wait for initial session to create statements - let mut statements_session_rx = session_rx.clone(); - while statements_session_rx.borrow().is_none() { - if statements_session_rx.changed().await.is_err() { - return Err(anyhow::anyhow!( - "Session sender dropped before initialization" - )); - } - } - - let statements = Arc::new(Statements::new(statements_session_rx, metadata.clone()).await?); - // Spawn main task for full scan and message processing tokio::spawn( async move { @@ -316,30 +319,37 @@ impl Statements { }) .collect(), ); - - let st_partition_key_list = table.partition_key.iter().join(", "); - let st_primary_key_list = primary_key_columns.iter().join(", "); + let st_partition_key_list = table + .partition_key + .iter() + .map(|c| CqlIdentifier::new(c.as_str())) + .join(", "); + let st_primary_key_list = primary_key_columns + .iter() + .map(|c| CqlIdentifier::new(c.as_ref())) + .join(", "); + let keyspace_identifier = KeyspaceIdentifier::from(&metadata.keyspace_name); + let table_identifier = TableIdentifier::from(&metadata.table_name); + let query = db_index_backend::range_scan_query( + &keyspace_identifier, + &table_identifier, + &metadata.target_column, + &st_primary_key_list, + &st_partition_key_list, + ); + let st_range_scan = session + .prepare(query) + .await + .context("range_scan_query")? + .pipe(|mut stmt| { + stmt.set_is_idempotent(true); + stmt + }); Ok(Self { primary_key_columns, - table_columns, - - st_range_scan: session - .prepare(Self::range_scan_query( - &metadata.keyspace_name, - &metadata.table_name, - &st_primary_key_list, - &st_partition_key_list, - &metadata.target_column, - )) - .await - .context("range_scan_query")? - .pipe(|mut stmt| { - stmt.set_is_idempotent(true); - stmt - }), - + st_range_scan, session_rx, }) } @@ -352,25 +362,6 @@ impl Statements { self.table_columns.clone() } - fn range_scan_query( - keyspace: &KeyspaceName, - table: &TableName, - st_primary_key_list: &str, - st_partition_key_list: &str, - embedding: &ColumnName, - ) -> String { - format!( - " - SELECT {st_primary_key_list}, {embedding}, writetime({embedding}) - FROM {keyspace}.{table} - WHERE - token({st_partition_key_list}) >= ? - AND token({st_partition_key_list}) <= ? - BYPASS CACHE - " - ) - } - async fn preform_range_scan(&self, begin: Token, end: Token) -> RangeScanResult { let mut range_scan = self.range_scan_stream(begin, end).await; let mut retry_timeout = START_RETRY_TIMEOUT; @@ -570,24 +561,16 @@ impl Statements { }; let timestamp = Timestamp::UNIX_EPOCH + Duration::from_micros(timestamp as u64); - let Some(CqlValue::Vector(embedding)) = row.columns.pop().unwrap() else { - debug!("range_scan_stream: bad type of an embedding"); + let Some(vector_value) = row.columns.pop().unwrap() else { + debug!("range_scan_stream: missing vector column"); return None; }; - let Ok(embedding) = embedding - .into_iter() - .map(|value| { - let CqlValue::Float(value) = value else { - bail!("range_scan_stream: bad type of an embedding element"); - }; - Ok(value) - }) - .collect::>>() + let Ok(vector) = Vector::try_from(vector_value) .inspect_err(|err| debug!("range_scan_stream: {err}")) else { return None; }; - let embedding = Some(embedding.into()); + let vector = Some(vector); let Ok(primary_key) = row .columns @@ -606,7 +589,7 @@ impl Statements { Some(DbEmbedding { primary_key, - embedding, + embedding: vector, timestamp, }) }) diff --git a/crates/vector-store/src/db_index_backend.rs b/crates/vector-store/src/db_index_backend.rs new file mode 100644 index 00000000..be769d83 --- /dev/null +++ b/crates/vector-store/src/db_index_backend.rs @@ -0,0 +1,368 @@ +/* + * Copyright 2026-present ScyllaDB + * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0 + */ + +use crate::ColumnName; +use crate::CqlLiteral; +use crate::Dimensions; +use crate::IndexMetadata; +use crate::IndexName; +use crate::KeyspaceIdentifier; +use crate::KeyspaceName; +use crate::TableIdentifier; +use crate::TableName; +use crate::Vector; +use crate::vector; +use futures::TryStreamExt; +use regex::Regex; +use scylla::client::session::Session; +use scylla::statement::prepared::PreparedStatement; +use scylla::value::CqlValue; +use scylla_cdc::CqlIdentifier; +use std::collections::BTreeMap; +use std::num::NonZeroUsize; + +pub(crate) struct IndexLocation { + pub keyspace: KeyspaceName, + pub table: TableName, + pub index: IndexName, +} + +pub(crate) enum DbIndexBackend { + Cql { target_column: ColumnName }, + Alternator { target_column: ColumnName }, +} + +impl From<&IndexMetadata> for DbIndexBackend { + fn from(metadata: &IndexMetadata) -> Self { + let target_column = metadata.target_column.clone(); + if metadata.keyspace_name.is_alternator() { + Self::Alternator { target_column } + } else { + Self::Cql { target_column } + } + } +} + +impl DbIndexBackend { + pub fn vector_column_name(&self) -> &str { + match self { + Self::Cql { target_column } => target_column.as_ref(), + Self::Alternator { .. } => ":attrs", + } + } + + pub fn extract_vector(&self, value: CqlValue) -> anyhow::Result> { + match self { + Self::Cql { .. } => Vector::try_from(value).map(Some), + Self::Alternator { target_column } => vector::AlternatorAttrs { + attrs: value, + target_column: target_column.as_ref(), + } + .try_into(), + } + } +} + +/// Builds the CQL range scan query appropriate for the given keyspace. +/// +/// For CQL-native tables, selects the vector column directly. +/// For Alternator tables, selects from the `:attrs` map column. +pub(crate) fn range_scan_query( + keyspace: &KeyspaceIdentifier, + table: &TableIdentifier, + target_column: &ColumnName, + primary_key_list: &str, + partition_key_list: &str, +) -> String { + if keyspace.is_alternator() { + let attributes = CqlIdentifier::new(":attrs"); + let vector = CqlLiteral::new(target_column.as_ref()); + format!( + " + SELECT {primary_key_list}, {attributes}[{vector}], writetime({attributes}[{vector}]) + FROM {keyspace}.{table} + WHERE + token({partition_key_list}) >= ? + AND token({partition_key_list}) <= ? + BYPASS CACHE + " + ) + } else { + let vector = CqlIdentifier::new(target_column.as_ref()); + format!( + " + SELECT {primary_key_list}, {vector}, writetime({vector}) + FROM {keyspace}.{table} + WHERE + token({partition_key_list}) >= ? + AND token({partition_key_list}) <= ? + BYPASS CACHE + " + ) + } +} + +/// Retrieves the vector dimensions for the given index, dispatching to the +/// appropriate strategy based on whether the keyspace is Alternator- or CQL-backed. +pub(crate) async fn get_dimensions( + target_column: &ColumnName, + session: &Session, + st_get_index_target_type: &PreparedStatement, + re_get_index_target_type: &Regex, + st_get_index_options: &PreparedStatement, + location: IndexLocation, +) -> anyhow::Result> { + if location.keyspace.is_alternator() { + get_dimensions_from_index_options(session, st_get_index_options, location).await + } else { + get_dimensions_from_column_type( + target_column, + session, + st_get_index_target_type, + re_get_index_target_type, + location, + ) + .await + } +} + +/// Retrieves the vector dimensions for a CQL-native table by parsing the column type. +async fn get_dimensions_from_column_type( + target_column: &ColumnName, + session: &Session, + st_get_index_target_type: &PreparedStatement, + re_get_index_target_type: &Regex, + location: IndexLocation, +) -> anyhow::Result> { + let column_type = session + .execute_iter( + st_get_index_target_type.clone(), + (location.keyspace, location.table, target_column.clone()), + ) + .await? + .rows_stream::<(String,)>()? + .try_next() + .await?; + let dimensions = column_type + .and_then(|(typ,)| { + re_get_index_target_type + .captures(&typ) + .and_then(|captures| captures["dimensions"].parse::().ok()) + }) + .and_then(|dimensions| NonZeroUsize::new(dimensions).map(|dimensions| dimensions.into())); + Ok(dimensions) +} + +/// Retrieves the vector dimensions for an Alternator table from the index options. +/// +/// In Alternator, the schema has no native `VECTOR` type, so the dimension +/// is stored in the index option `"dimensions"`. +async fn get_dimensions_from_index_options( + session: &Session, + st_get_index_options: &PreparedStatement, + location: IndexLocation, +) -> anyhow::Result> { + let index_options = session + .execute_iter( + st_get_index_options.clone(), + (location.keyspace, location.table, location.index), + ) + .await? + .rows_stream::<(BTreeMap,)>()? + .try_next() + .await?; + let dimensions = index_options + .and_then(|(mut options,)| { + options + .remove("dimensions") + .and_then(|s| s.parse::().ok()) + }) + .and_then(|dimensions| NonZeroUsize::new(dimensions).map(|dimensions| dimensions.into())); + Ok(dimensions) +} + +#[cfg(test)] +mod tests { + use super::*; + use itertools::Itertools; + + #[test] + fn range_scan_query_quotes_lowercase_identifiers() { + let query = range_scan_query( + &KeyspaceIdentifier::from("ks"), + &TableIdentifier::from("tbl"), + &ColumnName::from("embedding"), + &CqlIdentifier::new("id").to_string(), + &CqlIdentifier::new("id").to_string(), + ); + assert!(query.contains(r#""embedding""#)); + assert!(query.contains(r#"FROM "ks"."tbl""#)); + assert!(query.contains(r#"token("id")"#)); + } + + #[test] + fn range_scan_query_quotes_mixed_case_identifiers() { + let pk_list = [ + CqlIdentifier::new("UserId"), + CqlIdentifier::new("CreatedAt"), + ] + .iter() + .join(", "); + let query = range_scan_query( + &KeyspaceIdentifier::from("MyKeyspace"), + &TableIdentifier::from("MyTable"), + &ColumnName::from("EmbeddingCol"), + &pk_list, + &CqlIdentifier::new("UserId").to_string(), + ); + assert!( + query.contains(r#""EmbeddingCol""#), + "mixed-case embedding column must be quoted" + ); + assert!( + query.contains(r#"FROM "MyKeyspace"."MyTable""#), + "mixed-case keyspace/table must be quoted" + ); + assert!( + query.contains(r#""UserId", "CreatedAt""#), + "mixed-case primary key columns must be quoted" + ); + } + + #[test] + fn range_scan_query_quotes_uppercase_identifiers() { + let query = range_scan_query( + &KeyspaceIdentifier::from("UPPER_KS"), + &TableIdentifier::from("UPPER_TBL"), + &ColumnName::from("VEC"), + &CqlIdentifier::new("ID").to_string(), + &CqlIdentifier::new("ID").to_string(), + ); + assert!( + query.contains(r#""VEC""#), + "uppercase embedding column must be quoted" + ); + assert!( + query.contains(r#"FROM "UPPER_KS"."UPPER_TBL""#), + "uppercase keyspace/table must be quoted" + ); + } + + #[test] + fn range_scan_query_quotes_special_character_identifiers() { + let pk_list = [CqlIdentifier::new(":pk"), CqlIdentifier::new(":sk")] + .iter() + .join(", "); + let query = range_scan_query( + &KeyspaceIdentifier::from("my-app"), + &TableIdentifier::from("my-table:v1"), + &ColumnName::from("my-vector"), + &pk_list, + &CqlIdentifier::new(":pk").to_string(), + ); + assert!( + query.contains(r#""my-vector""#), + "hyphenated embedding column must be quoted" + ); + assert!( + query.contains(r#"FROM "my-app"."my-table:v1""#), + "special-character keyspace/table must be quoted" + ); + assert!( + query.contains(r#"token(":pk")"#), + "special-character partition key must be quoted" + ); + } + + #[test] + fn alternator_range_scan_query_basic() { + let pk_list = [CqlIdentifier::new(":pk"), CqlIdentifier::new(":sk")] + .iter() + .join(", "); + let query = range_scan_query( + &KeyspaceIdentifier::from("alternator_my-app"), + &TableIdentifier::from("my-table"), + &ColumnName::from("v"), + &pk_list, + &CqlIdentifier::new(":pk").to_string(), + ); + assert!( + query.contains(r#"":attrs"['v']"#), + "attribute name must be single-quoted inside :attrs map access: {query}" + ); + assert!( + query.contains(r#"writetime(":attrs"['v'])"#), + "writetime must wrap the same :attrs map access: {query}" + ); + assert!( + query.contains(r#"FROM "alternator_my-app"."my-table""#), + "keyspace and table must be double-quoted: {query}" + ); + assert!( + query.contains(r#"token(":pk")"#), + "partition key must be double-quoted: {query}" + ); + } + + #[test] + fn alternator_range_scan_query_special_attribute_name() { + let pk_list = CqlIdentifier::new(":pk").to_string(); + let query = range_scan_query( + &KeyspaceIdentifier::from("alternator_ks"), + &TableIdentifier::from("tbl"), + &ColumnName::from("my-vector:v1"), + &pk_list, + &pk_list, + ); + assert!( + query.contains(r#"":attrs"['my-vector:v1']"#), + "special characters in attribute name must appear verbatim inside single quotes: {query}" + ); + assert!( + query.contains(r#"writetime(":attrs"['my-vector:v1'])"#), + "writetime must use the same single-quoted attribute access: {query}" + ); + } + + #[test] + fn alternator_range_scan_query_mixed_case_attribute() { + let pk_list = CqlIdentifier::new("pk").to_string(); + let query = range_scan_query( + &KeyspaceIdentifier::from("alternator_Ks"), + &TableIdentifier::from("Tbl"), + &ColumnName::from("EmbeddingCol"), + &pk_list, + &pk_list, + ); + assert!( + query.contains(r#"":attrs"['EmbeddingCol']"#), + "mixed-case attribute name must be preserved as-is inside single quotes: {query}" + ); + assert!( + query.contains(r#"FROM "alternator_Ks"."Tbl""#), + "mixed-case keyspace/table must be double-quoted: {query}" + ); + } + + #[test] + fn alternator_range_scan_query_attribute_with_quotes() { + let pk_list = CqlIdentifier::new(":pk").to_string(); + let query = range_scan_query( + &KeyspaceIdentifier::from("alternator_ks"), + &TableIdentifier::from("tbl"), + &ColumnName::from("it's a \"test\""), + &pk_list, + &pk_list, + ); + assert!( + query.contains(r#"":attrs"['it''s a "test"']"#), + "single quotes in attribute name must be escaped by doubling: {query}" + ); + assert!( + query.contains(r#"writetime(":attrs"['it''s a "test"'])"#), + "writetime must use the same escaped attribute access: {query}" + ); + } +} diff --git a/crates/vector-store/src/index/opensearch.rs b/crates/vector-store/src/index/opensearch.rs index 4884d06f..5a863c51 100644 --- a/crates/vector-store/src/index/opensearch.rs +++ b/crates/vector-store/src/index/opensearch.rs @@ -323,7 +323,7 @@ async fn add( &primary_id.as_ref().to_string(), )) .body(json!({ - "vector": embeddings.0, + "vector": embeddings.as_slice(), })) .send() .await @@ -376,7 +376,7 @@ async fn ann( "query": { "knn": { "vector": { - "vector": embedding.0, + "vector": embedding.as_slice(), "k": limit.0, } } diff --git a/crates/vector-store/src/index/usearch.rs b/crates/vector-store/src/index/usearch.rs index 802bae14..feeeedfc 100644 --- a/crates/vector-store/src/index/usearch.rs +++ b/crates/vector-store/src/index/usearch.rs @@ -195,10 +195,10 @@ impl UsearchIndex for ThreadedUsearchIndex { fn add(&self, primary_id: PrimaryId, vector: &Vector) -> anyhow::Result<()> { if self.quantization == ScalarKind::B1 { - let vector = f32_to_b1x8(&vector.0); + let vector = f32_to_b1x8(vector.as_slice()); return Ok(self.inner.add(primary_id.into(), &vector)?); } - Ok(self.inner.add(primary_id.into(), &vector.0)?) + Ok(self.inner.add(primary_id.into(), vector.as_slice())?) } fn remove(&self, primary_id: PrimaryId) -> anyhow::Result<()> { @@ -211,10 +211,10 @@ impl UsearchIndex for ThreadedUsearchIndex { limit: Limit, ) -> anyhow::Result>> { let matches = if self.quantization == ScalarKind::B1 { - let vector = f32_to_b1x8(&vector.0); + let vector = f32_to_b1x8(vector.as_slice()); self.inner.search(&vector, limit.0.get())? } else { - self.inner.search(&vector.0, limit.0.get())? + self.inner.search(vector.as_slice(), limit.0.get())? }; Ok(matches .keys @@ -233,12 +233,14 @@ impl UsearchIndex for ThreadedUsearchIndex { filter: impl Fn(PrimaryId) -> bool, ) -> anyhow::Result>> { let matches = if self.quantization == ScalarKind::B1 { - let vector = f32_to_b1x8(&vector.0); + let vector = f32_to_b1x8(vector.as_slice()); self.inner .filtered_search(&vector, limit.0.get(), |row_id| filter(row_id.into()))? } else { self.inner - .filtered_search(&vector.0, limit.0.get(), |row_id| filter(row_id.into()))? + .filtered_search(vector.as_slice(), limit.0.get(), |row_id| { + filter(row_id.into()) + })? }; Ok(matches .keys diff --git a/crates/vector-store/src/index/validator.rs b/crates/vector-store/src/index/validator.rs index 3f4f38dd..39c79690 100644 --- a/crates/vector-store/src/index/validator.rs +++ b/crates/vector-store/src/index/validator.rs @@ -10,7 +10,7 @@ pub enum Error { } pub fn embedding_dimensions(embedding: &Vector, dimensions: Dimensions) -> anyhow::Result<()> { - let Some(embedding_len) = std::num::NonZeroUsize::new(embedding.0.len()) else { + let Some(embedding_len) = std::num::NonZeroUsize::new(embedding.len()) else { bail!(Error::WrongEmbeddingDimension { expected: dimensions.0.get(), actual: 0, @@ -36,7 +36,7 @@ mod tests { #[test] fn validate_embedding_empty() { - let embedding = Vector(vec![]); + let embedding = Vector::from(vec![]); let dimensions = dims(3); let result = embedding_dimensions(&embedding, dimensions); @@ -52,7 +52,7 @@ mod tests { #[test] fn validate_embedding_too_short() { - let embedding = Vector(vec![0.1, 0.2]); + let embedding = Vector::from(vec![0.1, 0.2]); let dimensions = dims(3); let result = embedding_dimensions(&embedding, dimensions); @@ -68,7 +68,7 @@ mod tests { #[test] fn validate_embedding_too_long() { - let embedding = Vector(vec![0.1, 0.2, 0.3, 0.4]); + let embedding = Vector::from(vec![0.1, 0.2, 0.3, 0.4]); let dimensions = dims(3); let result = embedding_dimensions(&embedding, dimensions); assert!(matches!( @@ -82,7 +82,7 @@ mod tests { #[test] fn validate_embedding_ok() { - let embedding = Vector(vec![0.1, 0.2, 0.3]); + let embedding = Vector::from(vec![0.1, 0.2, 0.3]); let dimensions = dims(3); let result = embedding_dimensions(&embedding, dimensions); diff --git a/crates/vector-store/src/lib.rs b/crates/vector-store/src/lib.rs index 735ef8bf..4b69d3ca 100644 --- a/crates/vector-store/src/lib.rs +++ b/crates/vector-store/src/lib.rs @@ -7,6 +7,7 @@ mod config_manager; pub mod db; mod db_cdc; pub mod db_index; +mod db_index_backend; mod distance; mod engine; pub mod httproutes; @@ -26,6 +27,7 @@ mod primary_key; mod similarity; mod table; mod timestamp; +mod vector; pub use crate::config_manager::ConfigManager; pub use crate::config_manager::load_config; @@ -52,6 +54,7 @@ use scylla::serialize::value::SerializeValue; use scylla::serialize::writers::CellWriter; use scylla::serialize::writers::WrittenCellProof; use scylla::value::CqlValue; +use scylla_cdc::CqlIdentifier; use std::borrow::Cow; use std::collections::HashMap; use std::hash::Hash; @@ -77,6 +80,80 @@ use utoipa::openapi::Schema; use utoipa::openapi::SchemaFormat; use utoipa::openapi::schema::Type; use uuid::Uuid; +pub use vector::Vector; + +/// A CQL string literal that is always properly single-quoted when formatted +/// for use in CQL statements. +/// +/// The inner value stores the already-quoted form of the string. +/// The [`Display`](std::fmt::Display) implementation outputs it in single quotes +/// with embedded single-quote characters escaped by doubling them +/// (`'` -> `''`), following the CQL grammar for string constants. +pub(crate) struct CqlLiteral { + quoted: String, +} + +impl CqlLiteral { + /// Creates a new `CqlLiteral`, preserving the value exactly as given. + /// + /// The value will be single-quoted when formatted, with any embedded + /// single quotes escaped by doubling. + pub(crate) fn new(value: impl AsRef) -> Self { + let quoted = format!("'{}'", value.as_ref().replace('\'', "''")); + Self { quoted } + } +} + +impl std::fmt::Display for CqlLiteral { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.quoted) + } +} + +pub(crate) struct KeyspaceIdentifier { + cql_identifier: CqlIdentifier, + is_alternator: bool, +} + +impl> From for KeyspaceIdentifier { + fn from(value: T) -> Self { + let value = value.as_ref(); + Self { + cql_identifier: CqlIdentifier::new(value), + is_alternator: value.starts_with("alternator_"), + } + } +} + +impl KeyspaceIdentifier { + pub(crate) fn is_alternator(&self) -> bool { + self.is_alternator + } +} + +impl std::fmt::Display for KeyspaceIdentifier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.cql_identifier.fmt(f) + } +} + +pub(crate) struct TableIdentifier { + cql_identifier: CqlIdentifier, +} + +impl> From for TableIdentifier { + fn from(value: T) -> Self { + Self { + cql_identifier: CqlIdentifier::new(value.as_ref()), + } + } +} + +impl std::fmt::Display for TableIdentifier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.cql_identifier.fmt(f) + } +} #[derive(Clone, Debug)] pub struct Config { @@ -154,6 +231,14 @@ pub struct Credentials { /// A keyspace name in a db. pub struct KeyspaceName(String); +impl KeyspaceName { + /// Returns true if this keyspace is backed by Alternator (DynamoDB-compatible API). + /// Alternator keyspaces are prefixed with `alternator_`. + fn is_alternator(&self) -> bool { + self.0.starts_with("alternator_") + } +} + impl SerializeValue for KeyspaceName { fn serialize<'b>( &self, @@ -406,25 +491,6 @@ impl FromStr for Quantization { } } -#[derive( - Clone, - Debug, - PartialEq, - serde::Serialize, - serde::Deserialize, - derive_more::AsRef, - derive_more::From, - utoipa::ToSchema, -)] -/// The vector to use for the Approximate Nearest Neighbor search. The format of data must match the data_type of the index. -pub struct Vector(Vec); - -impl Vector { - pub fn dim(&self) -> Option { - NonZeroUsize::new(self.0.len()).map(Dimensions) - } -} - #[derive( Clone, Copy, diff --git a/crates/vector-store/src/monitor_indexes.rs b/crates/vector-store/src/monitor_indexes.rs index 98f5346b..28e0a476 100644 --- a/crates/vector-store/src/monitor_indexes.rs +++ b/crates/vector-store/src/monitor_indexes.rs @@ -137,6 +137,7 @@ async fn get_indexes(db: &Sender) -> anyhow::Result> idx.keyspace.clone(), idx.table.clone(), idx.target_column.clone(), + idx.index.clone(), ) .await .inspect_err(|err| warn!("unable to get index target dimensions: {err}"))? @@ -470,7 +471,7 @@ mod tests { mock_db .expect_get_index_target_type() - .returning(move |_, _, _, tx| { + .returning(move |_, _, _, _, tx| { async move { // Return dimensions for all indexes tx.send(Ok(Some(NonZeroUsize::new(3).unwrap().into()))) @@ -598,7 +599,7 @@ mod tests { mock_db .expect_get_index_target_type() - .returning(move |_, _, _, tx| { + .returning(move |_, _, _, _, tx| { async move { tx.send(Ok(Some(NonZeroUsize::new(3).unwrap().into()))) .unwrap(); diff --git a/crates/vector-store/src/vector.rs b/crates/vector-store/src/vector.rs new file mode 100644 index 00000000..8ac1054f --- /dev/null +++ b/crates/vector-store/src/vector.rs @@ -0,0 +1,268 @@ +/* + * Copyright 2026-present ScyllaDB + * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0 + */ + +use crate::Dimensions; +use anyhow::anyhow; +use anyhow::bail; +use scylla::value::CqlValue; +use std::num::NonZeroUsize; + +#[derive( + Clone, + Debug, + PartialEq, + serde::Serialize, + serde::Deserialize, + derive_more::AsRef, + derive_more::From, + utoipa::ToSchema, +)] +/// The vector to use for the Approximate Nearest Neighbor search. The format of data must match the data_type of the index. +pub struct Vector(Vec); + +impl Vector { + pub fn as_slice(&self) -> &[f32] { + &self.0 + } + + pub fn is_empty(&self) -> bool { + self.as_slice().is_empty() + } + + pub fn len(&self) -> usize { + self.as_slice().len() + } + + pub fn dim(&self) -> Option { + NonZeroUsize::new(self.len()).map(Dimensions) + } +} + +/// Converts a [`CqlValue`] into a [`Vector`]. +/// +/// Supports two representations: +/// - `CqlValue::Vector` — native CQL `VECTOR` type (used by CQL-native tables). +/// - `CqlValue::Blob` — DynamoDB JSON serialized as bytes (used by Alternator). +impl TryFrom for Vector { + type Error = anyhow::Error; + + fn try_from(value: CqlValue) -> anyhow::Result { + let floats = match value { + CqlValue::Vector(values) => values + .into_iter() + .map(|v| { + let CqlValue::Float(f) = v else { + bail!("bad type of embedding element: expected float, got {v:?}"); + }; + Ok(f) + }) + .collect(), + CqlValue::Blob(bytes) => parse_dynamodb_vector_json(&bytes), + other => Err(anyhow!( + "unsupported CQL type for embedding column: {other:?}" + )), + }?; + Ok(Self(floats)) + } +} + +/// Alternator type tag for the DynamoDB List type (`L`), which is how vector embeddings are serialised. +/// Alternator prefixes each attribute value in the `:attrs` map column with a 1-byte type discriminator. +/// The List type uses the tag value `0x04` (named `NOT_SUPPORTED_YET`) +const ALTERNATOR_TYPE_NOT_SUPPORTED_YET: u8 = 4; + +/// Parses a DynamoDB-style JSON vector stored as raw bytes. +/// +/// Handles two representations: +/// - Plain JSON: `{"L": [{"N": "123.4"}, {"N": "234.5"}, ...]}` +/// - Alternator-prefixed: a 1-byte type tag (`0x04`) followed by the JSON above. +/// This prefix is used by Alternator for attribute values in the `:attrs` map. +fn parse_dynamodb_vector_json(bytes: &[u8]) -> anyhow::Result> { + let bytes = match bytes.first() { + Some(&ALTERNATOR_TYPE_NOT_SUPPORTED_YET) => &bytes[1..], + _ => bytes, + }; + + #[derive(serde::Deserialize)] + struct DynamoDbList { + #[serde(rename = "L")] + l: Vec, + } + + #[derive(serde::Deserialize)] + struct DynamoDbNumber { + #[serde(rename = "N")] + n: String, + } + + let list: DynamoDbList = serde_json::from_slice(bytes)?; + list.l + .into_iter() + .map(|item| { + item.n + .parse::() + .map_err(|e| anyhow!("invalid value in DynamoDB vector element: {e}")) + }) + .collect() +} + +pub(crate) struct AlternatorAttrs<'a> { + pub attrs: CqlValue, + pub target_column: &'a str, +} + +/// Extracts a vector from the Alternator `:attrs` map column. +/// +/// In Alternator, non-key attributes are stored in a `map` column named `:attrs`. +/// Each entry's key is the attribute name and the value is a serialised attribute prefixed with a 1-byte type tag. +impl TryFrom> for Option { + type Error = anyhow::Error; + + fn try_from(input: AlternatorAttrs<'_>) -> anyhow::Result { + let AlternatorAttrs { + attrs, + target_column, + } = input; + let CqlValue::Map(entries) = attrs else { + bail!("expected Map for :attrs column, got {attrs:?}"); + }; + + let target = target_column.as_bytes(); + + entries + .into_iter() + .find_map(|(key, value)| { + let matches = match &key { + CqlValue::Blob(b) => b.as_slice() == target, + CqlValue::Text(s) => s.as_bytes() == target, + _ => false, + }; + matches.then_some(value) + }) + .map(Vector::try_from) + .transpose() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extract_from_cql_vector() { + let value = CqlValue::Vector(vec![ + CqlValue::Float(1.0), + CqlValue::Float(2.5), + CqlValue::Float(3.0), + ]); + let result = Vector::try_from(value).unwrap(); + assert_eq!(result, Vector::from(vec![1.0, 2.5, 3.0])); + } + + #[test] + fn extract_from_dynamodb_json_blob() { + let json = r#"{"L": [{"N": "123.4"}, {"N": "234.5"}, {"N": "345.6"}]}"#; + let value = CqlValue::Blob(json.as_bytes().to_vec()); + let result = Vector::try_from(value).unwrap(); + assert_eq!(result, Vector::from(vec![123.4, 234.5, 345.6])); + } + + #[test] + fn extract_from_dynamodb_json_empty_list() { + let json = r#"{"L": []}"#; + let value = CqlValue::Blob(json.as_bytes().to_vec()); + let result = Vector::try_from(value).unwrap(); + assert_eq!(result, Vector::from(vec![])); + } + + #[test] + fn extract_from_dynamodb_json_invalid_number() { + let json = r#"{"L": [{"N": "not_a_number"}]}"#; + let value = CqlValue::Blob(json.as_bytes().to_vec()); + assert!(Vector::try_from(value).is_err()); + } + + #[test] + fn extract_from_unsupported_type() { + let value = CqlValue::Int(42); + assert!(Vector::try_from(value).is_err()); + } + + #[test] + fn extract_from_cql_vector_wrong_element_type() { + let value = CqlValue::Vector(vec![CqlValue::Int(1)]); + assert!(Vector::try_from(value).is_err()); + } + + /// Helper: prepend the Alternator `NOT_SUPPORTED_YET` tag (0x04) to a + /// DynamoDB JSON string, mirroring how Alternator serialises List values. + fn alternator_blob(json: &str) -> Vec { + let mut v = vec![ALTERNATOR_TYPE_NOT_SUPPORTED_YET]; + v.extend_from_slice(json.as_bytes()); + v + } + + #[test] + fn extract_from_attrs_map_with_blob_keys() { + let json = r#"{"L": [{"N": "1.0"}, {"N": "2.0"}]}"#; + let attrs = CqlValue::Map(vec![ + ( + CqlValue::Blob(b"other".to_vec()), + CqlValue::Blob(alternator_blob(r#"{"S": "ignored"}"#)), + ), + ( + CqlValue::Blob(b"v".to_vec()), + CqlValue::Blob(alternator_blob(json)), + ), + ]); + let result = Option::::try_from(AlternatorAttrs { + attrs, + target_column: "v", + }) + .unwrap(); + assert_eq!(result, Some(Vector::from(vec![1.0, 2.0]))); + } + + #[test] + fn extract_from_attrs_map_with_text_keys() { + let json = r#"{"L": [{"N": "3.0"}]}"#; + let attrs = CqlValue::Map(vec![( + CqlValue::Text("v".to_string()), + CqlValue::Blob(alternator_blob(json)), + )]); + let result = Option::::try_from(AlternatorAttrs { + attrs, + target_column: "v", + }) + .unwrap(); + assert_eq!(result, Some(Vector::from(vec![3.0]))); + } + + #[test] + fn extract_from_attrs_map_missing_target() { + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"other".to_vec()), + CqlValue::Blob(b"data".to_vec()), + )]); + let result = Option::::try_from(AlternatorAttrs { + attrs, + target_column: "v", + }) + .unwrap(); + assert_eq!(result, None); + } + + #[test] + fn extract_from_attrs_non_map() { + let attrs = CqlValue::Int(42); + assert!( + Option::::try_from(AlternatorAttrs { + attrs, + target_column: "v" + }) + .is_err() + ); + } +} diff --git a/crates/vector-store/tests/integration/db_basic.rs b/crates/vector-store/tests/integration/db_basic.rs index ac51c0b8..89c40973 100644 --- a/crates/vector-store/tests/integration/db_basic.rs +++ b/crates/vector-store/tests/integration/db_basic.rs @@ -295,6 +295,7 @@ fn process_db(db: &DbBasic, msg: Db, node_state: Sender) { table, target_column, tx, + .. } => tx .send(Ok(db .0