diff --git a/README.md b/README.md index dfcfeca2..e0258a24 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ service supports also `.env` files. | `VECTOR_STORE_CDC_FINE_SAFETY_INTERVAL` | Fine-grained CDC reader's safety interval for low-latency updates (ie. `100ms`) | `100ms` | | `VECTOR_STORE_CDC_FINE_SLEEP_INTERVAL` | Fine-grained CDC reader's sleep interval for low-latency updates (ie. `500ms`) | `500ms` | | `VECTOR_STORE_MONITOR_INDEXES_INTERVAL` | How often to poll Scylla for schema changes (new/removed vector indexes). The value is in human readable format (ie. `100ms`) | `1s` | +| `VECTOR_STORE_INDEX_STATUS_UPDATE_INTERVAL` | How often to sync index status (e.g., BOOTSTRAPPING->SERVING) into the engine's cached state. The value is in human readable format (ie. `100ms`) | `1s` | | `VECTOR_STORE_USEARCH_SIMULATOR` | Enable simulator for USearch. Provides human readable delays for simulated operations (`search:add-remove:reserve`). | | | `VECTOR_STORE_ALTER_INDEX_SIMULATOR` | Enable simulator for missing `ALTER INDEX`. When enable indexes aren't deleted and their version is not checked. | `false` | diff --git a/api/openapi.json b/api/openapi.json index 2e520f9e..40b0aa44 100644 --- a/api/openapi.json +++ b/api/openapi.json @@ -407,6 +407,13 @@ "limit": { "$ref": "#/components/schemas/Limit" }, + "return_columns": { + "type": "array", + "items": { + "$ref": "#/components/schemas/ColumnName" + }, + "description": "Filtering-column names whose stored values should be returned alongside\nthe primary keys. Empty (the default) means return no column values." + }, "vector": { "$ref": "#/components/schemas/Vector" } @@ -420,6 +427,18 @@ "similarity_scores" ], "properties": { + "column_values": { + "type": "object", + "description": "Per-column stored values for the filtering columns requested in\n`return_columns`. Each entry maps a column name to a Vec of\n`Option` — one entry per returned nearest neighbour, in the\nsame order as `similarity_scores`. `None` means the value was not\npresent for that row (e.g. the attribute did not exist when the row\nwas indexed). Absent when `return_columns` was empty.", + "additionalProperties": { + "type": "array", + "items": {} + }, + "propertyNames": { + "type": "string", + "description": "Name of the column in a db table." + } + }, "distances": { "type": "array", "items": { diff --git a/crates/httpapi/src/lib.rs b/crates/httpapi/src/lib.rs index 48ca4e45..725c9d99 100644 --- a/crates/httpapi/src/lib.rs +++ b/crates/httpapi/src/lib.rs @@ -270,6 +270,10 @@ pub struct PostIndexAnnRequest { pub filter: Option, #[serde(default)] pub limit: Limit, + /// Filtering-column names whose stored values should be returned alongside + /// the primary keys. Empty (the default) means return no column values. + #[serde(default)] + pub return_columns: Vec, } #[derive(serde::Deserialize, serde::Serialize, utoipa::ToSchema)] @@ -277,6 +281,14 @@ pub struct PostIndexAnnResponse { pub primary_keys: HashMap>, pub distances: Vec, pub similarity_scores: Vec, + /// Per-column stored values for the filtering columns requested in + /// `return_columns`. Each entry maps a column name to a Vec of + /// `Option` — one entry per returned nearest neighbour, in the + /// same order as `similarity_scores`. `None` means the value was not + /// present for that row (e.g. the attribute did not exist when the row + /// was indexed). Absent when `return_columns` was empty. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub column_values: HashMap>>, } #[derive( diff --git a/crates/httpclient/src/lib.rs b/crates/httpclient/src/lib.rs index a4316024..29698df5 100644 --- a/crates/httpclient/src/lib.rs +++ b/crates/httpclient/src/lib.rs @@ -88,6 +88,7 @@ impl HttpClient { vector, filter, limit, + return_columns: vec![], }; self.post_ann_data(keyspace_name, index_name, &request) .await diff --git a/crates/testclient/src/lib.rs b/crates/testclient/src/lib.rs index 344d2c78..b7e98137 100644 --- a/crates/testclient/src/lib.rs +++ b/crates/testclient/src/lib.rs @@ -59,6 +59,7 @@ impl TestClient { vector, filter, limit, + return_columns: vec![], }) .await } diff --git a/crates/validator/src/cdc.rs b/crates/validator/src/cdc.rs index 90978010..006faa40 100644 --- a/crates/validator/src/cdc.rs +++ b/crates/validator/src/cdc.rs @@ -51,6 +51,11 @@ pub(crate) async fn new() -> TestCase { timeout, cql_per_row_ttl_expires_from_index, ) + .with_test( + "cql_explicit_vector_column_deletion_removes_from_index", + timeout, + cql_explicit_vector_column_deletion_removes_from_index, + ) } #[framed] @@ -659,3 +664,94 @@ async fn cql_per_row_ttl_expires_from_index(actors: TestActors) { info!("finished"); } + +/// Regression test: explicitly deleting the vector column (`DELETE v FROM t WHERE pk = ?`) +/// must remove the item from the vector index, not leave it as "unchanged". +/// +/// Previously the CQL CDC path only treated full row/partition deletions as +/// removals. An explicit column tombstone produced `embedding = None` +/// ("unchanged") instead of `Some(None)` ("delete from index"), so the stale +/// vector remained searchable. +#[framed] +async fn cql_explicit_vector_column_deletion_removes_from_index(actors: TestActors) { + info!("started"); + + let (session, clients) = prepare_connection_single_vs(&actors).await; + let client = clients.first().unwrap(); + let keyspace = create_keyspace(&session).await; + let table = create_table(&session, "pk INT PRIMARY KEY, v VECTOR", None).await; + + session + .query_unpaged( + format!("INSERT INTO {table} (pk, v) VALUES (1, [1.0, 0.0, 0.0])"), + (), + ) + .await + .expect("failed to insert data"); + session + .query_unpaged( + format!("INSERT INTO {table} (pk, v) VALUES (2, [0.0, 1.0, 0.0])"), + (), + ) + .await + .expect("failed to insert data"); + + let index = create_index(CreateIndexQuery::new(&session, &clients, &table, "v")).await; + + let status = wait_for_index(client, &index).await; + assert_eq!( + status.count, 2, + "Index should have 2 vectors after full scan" + ); + + // Explicitly delete only the vector column of pk=1 (not the whole row). + // This is a column tombstone, not a row deletion — the CQL CDC path + // must recognise it via is_value_deleted() and remove the vector from the index. + session + .query_unpaged(format!("DELETE v FROM {table} WHERE pk = 1"), ()) + .await + .expect("failed to delete vector column"); + + // Wait for the index count to drop to 1. + wait_for( + || async { + let status = client.index_status(&index.keyspace, &index.index).await; + matches!(status, Ok(s) if s.count == 1) + }, + "Waiting for index count to drop to 1 after explicit column deletion", + FINE_GRAINED_CDC_MAX_LATENCY, + ) + .await; + + // ANN query should only return pk=2. + let result = wait_for_value( + || async { + let result = get_opt_query_results( + format!("SELECT pk FROM {table} ORDER BY v ANN OF [0.0, 1.0, 0.0] LIMIT 10"), + &session, + ) + .await; + result.filter(|r| r.rows_num() == 1) + }, + "Waiting for ANN query to return only pk=2", + FINE_GRAINED_CDC_MAX_LATENCY, + ) + .await; + let pks: Vec = result + .rows::<(i32,)>() + .expect("failed to get rows") + .map(|r| r.expect("row error").0) + .collect(); + assert_eq!( + pks, + vec![2], + "Only pk=2 should remain after explicit vector column deletion" + ); + + session + .query_unpaged(format!("DROP KEYSPACE {keyspace}"), ()) + .await + .expect("failed to drop keyspace"); + + info!("finished"); +} diff --git a/crates/validator/src/filtering.rs b/crates/validator/src/filtering.rs index eb13c644..289809cb 100644 --- a/crates/validator/src/filtering.rs +++ b/crates/validator/src/filtering.rs @@ -100,6 +100,11 @@ pub(crate) async fn new() -> TestCase { timeout, local_ann_with_timestamp_gte_filter, ) + .with_test( + "ann_filter_by_non_pk_filtering_column", + timeout, + ann_filter_by_non_pk_filtering_column, + ) } /// Test ANN search filtered by partition key equality. @@ -1044,6 +1049,87 @@ async fn local_ann_with_timestamp_gte_filter(actors: TestActors) { info!("finished"); } +/// Test ANN search filtered by a non-primary-key CQL filtering column. +/// +/// The index is created with `category` as a filtering column. Rows have two +/// category values (0 and 1). Querying with `WHERE category = 0` must return +/// only the rows whose category is 0. +#[framed] +async fn ann_filter_by_non_pk_filtering_column(actors: TestActors) { + info!("started"); + + let (session, clients) = prepare_connection(&actors).await; + + let keyspace = create_keyspace(&session).await; + let table = create_table( + &session, + "pk INT PRIMARY KEY, category INT, v VECTOR", + None, + ) + .await; + + // Insert 6 rows: pk 0..5, alternating category 0 and 1. + for pk in 0..6_i32 { + session + .query_unpaged( + format!("INSERT INTO {table} (pk, category, v) VALUES (?, ?, ?)"), + (pk, pk % 2, &vec![pk as f32, 0.0, 0.0]), + ) + .await + .expect("failed to insert data"); + } + + let index = create_index( + CreateIndexQuery::new(&session, &clients, &table, "v").filter_columns(["category"]), + ) + .await; + + for client in &clients { + let index_status = wait_for_index(client, &index).await; + assert_eq!(index_status.count, 6, "Expected 6 vectors to be indexed"); + } + + let result = wait_for_value( + || async { + let result = get_opt_query_results( + format!( + "SELECT pk, category FROM {table} \ + WHERE category = 0 \ + ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 10 \ + ALLOW FILTERING" + ), + &session, + ) + .await; + result.filter(|r| r.rows_num() == 3) + }, + "Waiting for category=0 filtered ANN query to return 3 rows", + DEFAULT_OPERATION_TIMEOUT, + ) + .await; + + let rows: Vec<(i32, i32)> = result + .rows::<(i32, i32)>() + .expect("failed to get rows") + .map(|row| row.expect("failed to get row")) + .collect(); + + assert_eq!(rows.len(), 3, "Expected exactly 3 rows with category=0"); + for (pk, category) in &rows { + assert_eq!( + *category, 0, + "Expected all rows to have category=0, got pk={pk} category={category}" + ); + } + + session + .query_unpaged(format!("DROP KEYSPACE {keyspace}"), ()) + .await + .expect("failed to drop a keyspace"); + + info!("finished"); +} + #[framed] async fn ann_filter_by_clustering_key_only_requires_allow_filtering(actors: TestActors) { info!("started"); diff --git a/crates/vector-store/benches/pipeline.rs b/crates/vector-store/benches/pipeline.rs index 041ce512..a4ffc35b 100644 --- a/crates/vector-store/benches/pipeline.rs +++ b/crates/vector-store/benches/pipeline.rs @@ -265,8 +265,9 @@ fn scan_fn_mpsc( .send(( DbEmbedding { primary_key, - embedding, + embedding: Some(embedding), timestamp, + column_values: Default::default(), }, in_progress, )) @@ -1139,7 +1140,7 @@ where }) }) .collect_vec(); - stream::iter(tasks.into_iter()) + stream::iter(tasks) .then(|task| async move { task.await.unwrap() }) .fold(Duration::ZERO, |acc, x| async move { acc + x }) .await diff --git a/crates/vector-store/src/config_manager.rs b/crates/vector-store/src/config_manager.rs index b772bff9..0c355035 100644 --- a/crates/vector-store/src/config_manager.rs +++ b/crates/vector-store/src/config_manager.rs @@ -491,6 +491,12 @@ pub async fn load_config(env: impl Fn(&str) -> anyhow::Result) -> anyhow .transpose()? .map(|v| v.into()); + config.engine_status_update_interval = env("VECTOR_STORE_INDEX_STATUS_UPDATE_INTERVAL") + .ok() + .map(|v| v.parse::()) + .transpose()? + .map(|v| v.into()); + config.cql_uri_translation_map = env("VECTOR_STORE_CQL_URI_TRANSLATION_MAP") .ok() .map(|v| serde_json::from_str(&v)) diff --git a/crates/vector-store/src/db.rs b/crates/vector-store/src/db.rs index 90215edd..f419af1a 100644 --- a/crates/vector-store/src/db.rs +++ b/crates/vector-store/src/db.rs @@ -42,7 +42,6 @@ use rustls_pki_types::pem::PemObject; use scylla::client::session::Session; use scylla::client::session::TlsContext; use scylla::client::session_builder::SessionBuilder; -use scylla::cluster::metadata::ColumnType; use scylla::cluster::metadata::Table; use scylla::statement::prepared::PreparedStatement; use scylla::value::CqlTimeuuid; @@ -751,7 +750,7 @@ impl Statements { anyhow!("table {table_name} does not exist").context(InvalidMetadata) })?; Ok(options.remove("target").and_then(|target| { - from_target_option(table, target) + from_target_option(table, &KeyspaceName::from(keyspace_name.as_str()), target) .map( |(index_type, target_column, filtering_columns)| DbCustomIndex { keyspace: keyspace_name.into(), @@ -1015,6 +1014,7 @@ fn convert_legacy_target_option( fn from_target_option( table: &Table, + keyspace_name: &KeyspaceName, value: String, ) -> anyhow::Result<(DbIndexType, ColumnName, Vec)> { let Some(target) = parse_target_option(table, &value)? else { @@ -1023,13 +1023,7 @@ fn from_target_option( }; let validate_target_type = |target_name: &str| -> anyhow::Result<()> { - let column = table.columns.get(target_name).ok_or_else(|| { - anyhow!("invalid target option: column {target_name} does not exist in a table") - })?; - if !matches!(column.typ, ColumnType::Vector { .. }) { - bail!("invalid target option: column {target_name} is not a vector column in a table"); - } - Ok(()) + db_index_backend::validate_target_type(table, keyspace_name, target_name) }; validate_target_type(&target.target_column)?; diff --git a/crates/vector-store/src/db_cdc.rs b/crates/vector-store/src/db_cdc.rs index 28df1ede..a9604210 100644 --- a/crates/vector-store/src/db_cdc.rs +++ b/crates/vector-store/src/db_cdc.rs @@ -26,6 +26,7 @@ use scylla::client::session::Session; use scylla_cdc::consumer::CDCRow; use scylla_cdc::consumer::Consumer; use scylla_cdc::consumer::ConsumerFactory; +use scylla_cdc::consumer::OperationType; use scylla_cdc::log_reader::CDCLogReaderBuilder; use std::sync::Arc; use std::time::Duration; @@ -507,15 +508,28 @@ impl Consumer for CdcConsumer { } 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 vector_col = source.vector_column_name(); + if !row.column_deletable(vector_col) { + bail!("CDC error: column {vector_col} should be deletable"); } - let embedding = row - .take_value(column) - .map(|v| source.extract_vector(v)) - .transpose()? - .flatten(); + + // Is this a row/partition deletion operation? + let is_deletion_op = matches!( + row.operation, + OperationType::RowDelete + | OperationType::PartitionDelete + | OperationType::RowRangeDelInclLeft + | OperationType::RowRangeDelExclLeft + | OperationType::RowRangeDelInclRight + | OperationType::RowRangeDelExclRight + ); + + let (column_values, embedding) = source.get_column_values( + &mut row, + is_deletion_op, + &self.0.primary_key_columns, + vector_col, + )?; let primary_key = self .0 @@ -546,6 +560,14 @@ impl Consumer for CdcConsumer { )) .into(); + // All columns changed in a single CDC event share the same write timestamp + // (the timestamp of the originating CQL statement). Attach it to each + // filtering column value so the index can apply last-writer-wins per column. + let column_values = column_values + .into_iter() + .map(|(k, v)| (k, (timestamp, v))) + .collect(); + _ = self .0 .tx @@ -554,6 +576,7 @@ impl Consumer for CdcConsumer { primary_key, embedding, timestamp, + column_values, }, None, )) diff --git a/crates/vector-store/src/db_index.rs b/crates/vector-store/src/db_index.rs index db35eaa4..9bca6fac 100644 --- a/crates/vector-store/src/db_index.rs +++ b/crates/vector-store/src/db_index.rs @@ -294,6 +294,9 @@ struct Statements { partition_key_count: usize, table_columns: GetTableColumnsR, st_range_scan: PreparedStatement, + /// Non-primary-key filtering columns included in the range scan query (in SELECT order). + scan_filtering_columns: Vec, + backend: db_index_backend::DbIndexBackend, } impl Statements { @@ -346,8 +349,17 @@ impl Statements { None } }) - .collect(), + .collect::>(), ); + + // For Alternator tables, non-key attributes live in `:attrs` and are not schema + // columns. Inject them as Blob so that Table::new() can create storage for them. + // Using Blob (rather than Text) preserves the Alternator type-tag byte, which lets + // cql_cmp distinguish S-type strings from N-type numbers and apply the correct + // comparison semantics (e.g. the string "1" must not match a numeric filter < 5). + let backend = db_index_backend::DbIndexBackend::from(&metadata); + let table_columns = backend.enrich_table_columns(table_columns, &metadata.filtering_columns); + let st_partition_key_list = table .partition_key .iter() @@ -359,12 +371,24 @@ impl Statements { .join(", "); let keyspace_identifier = KeyspaceIdentifier::from(&metadata.keyspace_name); let table_identifier = TableIdentifier::from(&metadata.table_name); + + // Determine which filtering columns are non-PK (these will appear as extra + // SELECT columns in the range scan query, after the vector+writetime columns). + let scan_filtering_columns: Vec = metadata + .filtering_columns + .iter() + .filter(|c| !primary_key_columns.contains(c)) + .cloned() + .collect(); + let query = db_index_backend::range_scan_query( &keyspace_identifier, &table_identifier, &metadata.target_column, + &primary_key_columns, &st_primary_key_list, &st_partition_key_list, + &metadata.filtering_columns, ); let st_range_scan = session .prepare(query) @@ -380,6 +404,8 @@ impl Statements { partition_key_count, table_columns, st_range_scan, + scan_filtering_columns, + backend, session_rx, }) } @@ -553,8 +579,9 @@ impl Statements { begin: Token, end: Token, ) -> anyhow::Result> { - // last two columns are embedding and writetime - let columns_len_expected = self.primary_key_columns.len() + 2; + // Columns: pk_cols | vector | writetime(vector) | col1 | writetime(col1) | ... (in order) + let n_extra = self.scan_filtering_columns.len(); + let columns_len_expected = self.primary_key_columns.len() + 2 + 2 * n_extra; // wait for an active session let session = { @@ -570,6 +597,9 @@ impl Statements { } }; + let scan_filtering_columns = self.scan_filtering_columns.clone(); + let backend = self.backend.clone(); + Ok(session .execute_iter(self.st_range_scan.clone(), (begin.value(), end.value())) .await? @@ -585,6 +615,23 @@ impl Statements { return None; } + // Pop extra filtering columns (appended after vector writetime) in reverse + // order. Each column is followed by its own writetime, so we pop the + // writetime first, then the value. + let mut column_values = std::collections::BTreeMap::new(); + for col_name in scan_filtering_columns.iter().rev() { + let Some(CqlValue::BigInt(col_ts_us)) = row.columns.pop().unwrap() else { + debug!("range_scan_stream: bad type of writetime for filtering column"); + return None; + }; + let col_ts = Timestamp::UNIX_EPOCH + Duration::from_micros(col_ts_us as u64); + let raw = row.columns.pop().unwrap(); + // A null/absent column in the base table is stored as a tombstone (`None`) + // so that previously-stored values for this row are cleared. + let scalar = raw.and_then(|v| backend.extract_scan_column_value(v)); + column_values.insert(col_name.clone(), (col_ts, scalar)); + } + let Some(CqlValue::BigInt(timestamp)) = row.columns.pop().unwrap() else { debug!("range_scan_stream: bad type of a writetime"); return None; @@ -619,8 +666,9 @@ impl Statements { Some(DbEmbedding { primary_key, - embedding: vector, + embedding: Some(vector), timestamp, + column_values, }) }) .filter_map(|value| async move { diff --git a/crates/vector-store/src/db_index_backend.rs b/crates/vector-store/src/db_index_backend.rs index be769d83..1b7b2987 100644 --- a/crates/vector-store/src/db_index_backend.rs +++ b/crates/vector-store/src/db_index_backend.rs @@ -14,14 +14,21 @@ use crate::TableIdentifier; use crate::TableName; use crate::Vector; use crate::vector; +use anyhow::bail; use futures::TryStreamExt; use regex::Regex; use scylla::client::session::Session; +use scylla::cluster::metadata::ColumnType; +use scylla::cluster::metadata::NativeType; +use scylla::cluster::metadata::Table; use scylla::statement::prepared::PreparedStatement; use scylla::value::CqlValue; use scylla_cdc::CqlIdentifier; +use scylla_cdc::consumer::CDCRow; use std::collections::BTreeMap; +use std::collections::HashMap; use std::num::NonZeroUsize; +use std::sync::Arc; pub(crate) struct IndexLocation { pub keyspace: KeyspaceName, @@ -29,18 +36,33 @@ pub(crate) struct IndexLocation { pub index: IndexName, } +#[derive(Clone)] pub(crate) enum DbIndexBackend { - Cql { target_column: ColumnName }, - Alternator { target_column: ColumnName }, + Cql { + target_column: ColumnName, + filtering_columns: Arc>, + }, + Alternator { + target_column: ColumnName, + /// Non-primary-key filtering columns, extracted from `:attrs` during scan/CDC. + filtering_columns: Arc>, + }, } impl From<&IndexMetadata> for DbIndexBackend { fn from(metadata: &IndexMetadata) -> Self { let target_column = metadata.target_column.clone(); + let filtering_columns = Arc::clone(&metadata.filtering_columns); if metadata.keyspace_name.is_alternator() { - Self::Alternator { target_column } + Self::Alternator { + target_column, + filtering_columns, + } } else { - Self::Cql { target_column } + Self::Cql { + target_column, + filtering_columns, + } } } } @@ -48,7 +70,7 @@ impl From<&IndexMetadata> for DbIndexBackend { impl DbIndexBackend { pub fn vector_column_name(&self) -> &str { match self { - Self::Cql { target_column } => target_column.as_ref(), + Self::Cql { target_column, .. } => target_column.as_ref(), Self::Alternator { .. } => ":attrs", } } @@ -56,32 +78,244 @@ impl DbIndexBackend { 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 { + Self::Alternator { target_column, .. } => vector::AlternatorAttrs { attrs: value, target_column: target_column.as_ref(), } .try_into(), } } + + /// Extracts filtering-column values and the embedding action from a CDC row delta. + /// + /// Returns `(column_values, embedding)` where: + /// - `column_values`: map from filtering column name to new value (`None` = deleted). + /// - `embedding`: + /// - `None` = vector unchanged (no update needed) + /// - `Some(None)` = vector deleted + /// - `Some(Some(v))` = new vector value + pub(crate) fn get_column_values( + &self, + row: &mut CDCRow<'_>, + is_deletion_op: bool, + primary_key_columns: &[ColumnName], + vector_col: &str, + ) -> anyhow::Result<(BTreeMap>, Option>)> { + let raw_vector_value = row.take_value(vector_col); + + match self { + Self::Alternator { + filtering_columns, + target_column, + .. + } => { + // cdc$deleted_elements_:attrs lists map keys explicitly deleted in this delta. + let deleted_elements: Vec = if row.collection_exists(vector_col) { + row.take_deleted_elements(vector_col) + } else { + vec![] + }; + + // cdc$deleted_:attrs is true when the whole :attrs collection was tombstoned + // (e.g. PutItem replaces the item entirely). + let whole_col_deleted = row.is_value_deleted(vector_col); + + // Returns true if the given attribute name was explicitly deleted from :attrs. + let attr_key_deleted = |attr: &str| -> bool { + deleted_elements.iter().any(|k| match k { + CqlValue::Text(s) | CqlValue::Ascii(s) => s == attr, + CqlValue::Blob(b) => b.as_slice() == attr.as_bytes(), + _ => false, + }) + }; + + let mut col_vals = BTreeMap::new(); + + if is_deletion_op { + // Full row/partition deletion: tombstone every filtering column. + for fc in filtering_columns.iter() { + col_vals.insert(fc.clone(), None); + } + } else { + // Partial update (RowUpdate / RowInsert / PostImage): + // CDC carries a delta — only changed entries appear in :attrs. + // 1. Extract new values from the :attrs delta for filtering columns. + if let Some(ref attrs) = raw_vector_value { + let delta_vals = vector::extract_alternator_scalars(attrs, filtering_columns); + col_vals.extend(delta_vals); + } + // 2. Tombstone filtering columns that were explicitly deleted or whose + // collection was replaced (whole_col_deleted) and are absent from the delta. + for fc in filtering_columns.iter() { + if attr_key_deleted(fc.as_ref()) + || (whole_col_deleted && !col_vals.contains_key(fc)) + { + col_vals.insert(fc.clone(), None); + } + } + } + + // Determine embedding action (tri-state): + // Some(None) = delete vector + // Some(Some(v)) = set vector + // None = vector unchanged, skip + let emb: Option> = + if is_deletion_op || attr_key_deleted(target_column.as_ref()) { + // Row deleted, or vector attribute explicitly removed from the map. + Some(None) + } else { + // Try to extract the vector from the :attrs delta. + let extracted = raw_vector_value + .as_ref() + .map(|v| self.extract_vector(v.clone())) + .transpose()?; + // extracted: Option> + // None => no :attrs delta at all + // Some(None) => :attrs delta present but vector key absent + // Some(Some(v)) => vector present in delta + if whole_col_deleted { + // Collection replaced: if vector not in new delta it was deleted. + Some(extracted.flatten()) + } else { + // Normal partial update: vector absent from delta = unchanged. + extracted.flatten().map(Some) + } + }; + + Ok((col_vals, emb)) + } + Self::Cql { filtering_columns, .. } => { + // For CQL tables a CDC row may be a partial update (only the changed columns + // appear). We can therefore only record values for columns that are present + // in this event; absent columns are left unchanged. + let mut vals = BTreeMap::new(); + for fc in filtering_columns.iter() { + if !primary_key_columns.contains(fc) { + if let Some(v) = row.take_value(fc.as_ref()) { + vals.insert(fc.clone(), Some(v)); + } else if row.is_value_deleted(fc.as_ref()) { + // Explicit column tombstone, e.g. `DELETE fc FROM t WHERE pk = ?`. + vals.insert(fc.clone(), None); + } + } + } + + // For CQL, derive embedding from the raw vector column value. + // Signal index removal (`Some(None)`) for: + // - row/partition deletions (is_deletion_op) + // - explicit column tombstones, e.g. `DELETE v FROM t WHERE pk = ?` + // or `UPDATE t SET v = null WHERE pk = ?`, where the column is + // absent from the CDC row but marked deleted via is_value_deleted. + let emb = if is_deletion_op || row.is_value_deleted(vector_col) { + Some(None) + } else { + raw_vector_value + .map(|v| self.extract_vector(v)) + .transpose()? + .flatten() + .map(Some) + }; + + Ok((vals, emb)) + } + } + } + + /// Converts a raw scan value for a filtering column into its stored form. + /// + /// For Alternator tables the scan returns a raw blob from `:attrs['col']`; + /// pass it through `extract_alternator_scalar` so only scalar S/N attributes + /// are kept (as `CqlValue::Blob` with the type-tag byte preserved) and + /// unsupported types are dropped. + /// For CQL tables the value is already typed; return it directly. + pub(crate) fn extract_scan_column_value(&self, raw: CqlValue) -> Option { + match self { + Self::Alternator { .. } => { + if let CqlValue::Blob(blob) = &raw { + vector::extract_alternator_scalar(blob) + } else { + Some(raw) + } + } + Self::Cql { .. } => Some(raw), + } + } + /// Injects Alternator filtering columns (absent from the CQL schema) into the + /// table-columns map as `NativeType::Blob` so that downstream code can create + /// storage for them. For CQL backends the map is returned unchanged. + pub(crate) fn enrich_table_columns( + &self, + table_columns: Arc>, + filtering_columns: &[ColumnName], + ) -> Arc> { + match self { + Self::Alternator { .. } => { + let mut cols = (*table_columns).clone(); + for fc in filtering_columns { + if !cols.contains_key(fc) { + cols.insert(fc.clone(), NativeType::Blob); + } + } + Arc::new(cols) + } + Self::Cql { .. } => table_columns, + } + } +} + +/// Validates that the target column of an index is a CQL vector column. +/// +/// For Alternator keyspaces the vector attribute is stored inside the `:attrs` +/// map and is not a real CQL column, so the check is skipped for those. +pub(crate) fn validate_target_type( + table: &Table, + keyspace_name: &KeyspaceName, + target_name: &str, +) -> anyhow::Result<()> { + if keyspace_name.is_alternator() { + return Ok(()); + } + let column = table.columns.get(target_name).ok_or_else(|| { + anyhow::anyhow!("invalid target option: column {target_name} does not exist in a table") + })?; + if !matches!(column.typ, ColumnType::Vector { .. }) { + bail!("invalid target option: column {target_name} is not a vector column in a table"); + } + Ok(()) } /// 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. +/// For CQL-native tables, selects the vector column and any filtering columns directly. +/// For Alternator tables, selects from the `:attrs` map column for the vector, +/// and from individual `:attrs['col']` subscripts for each non-primary-key filtering column. pub(crate) fn range_scan_query( keyspace: &KeyspaceIdentifier, table: &TableIdentifier, target_column: &ColumnName, + primary_key_columns: &[ColumnName], primary_key_list: &str, partition_key_list: &str, + filtering_columns: &[ColumnName], ) -> String { if keyspace.is_alternator() { let attributes = CqlIdentifier::new(":attrs"); let vector = CqlLiteral::new(target_column.as_ref()); + // Collect non-primary-key filtering columns to fetch from :attrs. + let extra_cols: Vec<_> = filtering_columns + .iter() + .filter(|c| !primary_key_columns.contains(c)) + .collect(); + let extra_select: String = extra_cols + .iter() + .map(|c| { + let lit = CqlLiteral::new(c.as_ref()); + format!(", {attributes}[{lit}], writetime({attributes}[{lit}])") + }) + .collect(); format!( " - SELECT {primary_key_list}, {attributes}[{vector}], writetime({attributes}[{vector}]) + SELECT {primary_key_list}, {attributes}[{vector}], writetime({attributes}[{vector}]){extra_select} FROM {keyspace}.{table} WHERE token({partition_key_list}) >= ? @@ -91,9 +325,21 @@ pub(crate) fn range_scan_query( ) } else { let vector = CqlIdentifier::new(target_column.as_ref()); + // Collect non-primary-key filtering columns to fetch directly. + let extra_cols: Vec<_> = filtering_columns + .iter() + .filter(|c| !primary_key_columns.contains(c)) + .collect(); + let extra_select: String = extra_cols + .iter() + .map(|c| { + let ident = CqlIdentifier::new(c.as_ref()); + format!(", {ident}, writetime({ident})") + }) + .collect(); format!( " - SELECT {primary_key_list}, {vector}, writetime({vector}) + SELECT {primary_key_list}, {vector}, writetime({vector}){extra_select} FROM {keyspace}.{table} WHERE token({partition_key_list}) >= ? @@ -194,8 +440,10 @@ mod tests { &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""#)); @@ -214,8 +462,10 @@ mod tests { &KeyspaceIdentifier::from("MyKeyspace"), &TableIdentifier::from("MyTable"), &ColumnName::from("EmbeddingCol"), + &[], &pk_list, &CqlIdentifier::new("UserId").to_string(), + &[], ); assert!( query.contains(r#""EmbeddingCol""#), @@ -237,8 +487,10 @@ mod tests { &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""#), @@ -259,8 +511,10 @@ mod tests { &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""#), @@ -285,8 +539,10 @@ mod tests { &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']"#), @@ -313,8 +569,10 @@ mod tests { &KeyspaceIdentifier::from("alternator_ks"), &TableIdentifier::from("tbl"), &ColumnName::from("my-vector:v1"), + &[], &pk_list, &pk_list, + &[], ); assert!( query.contains(r#"":attrs"['my-vector:v1']"#), @@ -333,8 +591,10 @@ mod tests { &KeyspaceIdentifier::from("alternator_Ks"), &TableIdentifier::from("Tbl"), &ColumnName::from("EmbeddingCol"), + &[], &pk_list, &pk_list, + &[], ); assert!( query.contains(r#"":attrs"['EmbeddingCol']"#), @@ -353,8 +613,10 @@ mod tests { &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"']"#), diff --git a/crates/vector-store/src/engine.rs b/crates/vector-store/src/engine.rs index 3d8895db..aca810d2 100644 --- a/crates/vector-store/src/engine.rs +++ b/crates/vector-store/src/engine.rs @@ -121,14 +121,17 @@ pub(crate) async fn new( config_rx.clone(), ) .await?; + let check_interval = config_rx + .borrow() + .engine_status_update_interval + .unwrap_or(Duration::from_secs(1)); let memory_actor = memory::new(config_rx); tokio::spawn( async move { debug!("starting"); - const CHECK_INTERVAL: Duration = Duration::from_secs(1); - let mut interval = time::interval(CHECK_INTERVAL); + let mut interval = time::interval(check_interval); loop { tokio::select! { msg = rx.recv() => { diff --git a/crates/vector-store/src/httproutes.rs b/crates/vector-store/src/httproutes.rs index 02bb36ba..9b5e550d 100644 --- a/crates/vector-store/src/httproutes.rs +++ b/crates/vector-store/src/httproutes.rs @@ -526,61 +526,71 @@ async fn post_index_ann( let index_key = IndexKey::new(&keyspace, &index_name); let (equality_cols, range_cols) = restriction_columns(&request.filter); let allow_filtering = request.filter.as_ref().is_some_and(|f| f.allow_filtering); - let (routed_key, index, primary_key_columns, table_columns) = match state - .indexes - .read() - .unwrap() - .best_index(&index_key, &equality_cols, &range_cols) - { - indexes::BestIndexState::Serving { - key: routed_key, - index, - needs_filtering, - primary_key_columns, - table_columns, - } => { - if matches!(needs_filtering, indexes::NeedsFiltering::Yes(_)) && !allow_filtering { - timer.observe_duration(); - - let msg = format!( - "Index {keyspace}.{index_name} requires ALLOW FILTERING for this query" - ); - debug!("post_index_ann: {msg}"); - return (StatusCode::BAD_REQUEST, msg).into_response(); - } - (routed_key, index, primary_key_columns, table_columns) - } - indexes::BestIndexState::NotServing(progress) => { - timer.observe_duration(); + let (routed_key, index, primary_key_columns, table_columns, index_filtering_columns) = + match state + .indexes + .read() + .unwrap() + .best_index(&index_key, &equality_cols, &range_cols) + { + indexes::BestIndexState::Serving { + key: routed_key, + index, + needs_filtering, + primary_key_columns, + table_columns, + filtering_columns, + } => { + if matches!(needs_filtering, indexes::NeedsFiltering::Yes(_)) + && !allow_filtering + { + timer.observe_duration(); - match progress { - Progress::InProgress(percentage) => { let msg = format!( - "Index {keyspace}.{index_name} is not available yet \ - as it is still being constructed, progress: {:.3}%", - percentage.get() + "Index {keyspace}.{index_name} requires ALLOW FILTERING for this query" ); debug!("post_index_ann: {msg}"); - return (StatusCode::SERVICE_UNAVAILABLE, msg).into_response(); + return (StatusCode::BAD_REQUEST, msg).into_response(); } - Progress::Done => { - let msg = format!( - "Index {keyspace}.{index_name} is not serving, \ + ( + routed_key, + index, + primary_key_columns, + table_columns, + filtering_columns, + ) + } + indexes::BestIndexState::NotServing(progress) => { + timer.observe_duration(); + + match progress { + Progress::InProgress(percentage) => { + let msg = format!( + "Index {keyspace}.{index_name} is not available yet \ + as it is still being constructed, progress: {:.3}%", + percentage.get() + ); + debug!("post_index_ann: {msg}"); + return (StatusCode::SERVICE_UNAVAILABLE, msg).into_response(); + } + Progress::Done => { + let msg = format!( + "Index {keyspace}.{index_name} is not serving, \ but full scan did finish." - ); - debug!("post_index_ann: {msg}"); - return (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response(); + ); + debug!("post_index_ann: {msg}"); + return (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response(); + } } } - } - indexes::BestIndexState::NotFound => { - timer.observe_duration(); + indexes::BestIndexState::NotFound => { + timer.observe_duration(); - let msg = format!("missing index: {keyspace}.{index_name}"); - debug!("post_index_ann: {msg}"); - return (StatusCode::NOT_FOUND, msg).into_response(); - } - }; + let msg = format!("missing index: {keyspace}.{index_name}"); + debug!("post_index_ann: {msg}"); + return (StatusCode::NOT_FOUND, msg).into_response(); + } + }; #[cfg(feature = "slow-test-hooks")] state @@ -592,11 +602,19 @@ async fn post_index_ann( )) .await; + let return_columns: std::sync::Arc<[crate::ColumnName]> = request + .return_columns + .into_iter() + .map(crate::ColumnName::from) + .collect::>() + .into(); let search_result = if let Some(filter) = request.filter { let filter = match try_from_post_index_ann_filter( filter, &primary_key_columns, &table_columns, + &index_filtering_columns, + keyspace.is_alternator(), ) { Ok(filter) => filter, Err(err) => { @@ -610,11 +628,17 @@ async fn post_index_ann( request.vector.into(), filter, request.limit.into(), + std::sync::Arc::clone(&return_columns), ) .await } else { index - .ann(routed_key, request.vector.into(), request.limit.into()) + .ann( + routed_key, + request.vector.into(), + request.limit.into(), + std::sync::Arc::clone(&return_columns), + ) .await }; @@ -630,7 +654,7 @@ async fn post_index_ann( (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response() } }, - Ok((primary_keys, distances)) => { + Ok((primary_keys, distances, column_values_per_row)) => { if primary_keys.len() != distances.len() { let msg = format!( "wrong size of an ann response: \ @@ -682,15 +706,35 @@ async fn post_index_ann( debug!("post_index_ann: {err}"); (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() } - Ok(primary_keys) => ( - StatusCode::OK, - response::Json(httpapi::PostIndexAnnResponse { - primary_keys, - distances: distances.into_iter().map(|d| d.into()).collect(), - similarity_scores, - }), - ) - .into_response(), + Ok(primary_keys) => { + let column_values: HashMap< + httpapi::ColumnName, + Vec>, + > = return_columns + .iter() + .map(|col_name| { + let values = column_values_per_row + .iter() + .map(|row_vals| { + row_vals + .get(col_name) + .and_then(|v| try_to_json(v.clone()).ok()) + }) + .collect(); + (col_name.clone().into(), values) + }) + .collect(); + ( + StatusCode::OK, + response::Json(httpapi::PostIndexAnnResponse { + primary_keys, + distances: distances.into_iter().map(|d| d.into()).collect(), + similarity_scores, + column_values, + }), + ) + .into_response() + } } } } @@ -703,6 +747,8 @@ fn try_from_post_index_ann_filter( json_filter: httpapi::PostIndexAnnFilter, primary_key_columns: &[crate::ColumnName], table_columns: &HashMap, + filtering_columns: &[crate::ColumnName], + is_alternator: bool, ) -> anyhow::Result { let is_same_len = |columns: &[crate::ColumnName], values: &[Value]| -> anyhow::Result<()> { if columns.len() != values.len() { @@ -715,14 +761,34 @@ fn try_from_post_index_ann_filter( Ok(()) }; let from_json = |column: &crate::ColumnName, value: Value| -> anyhow::Result { - if !primary_key_columns.contains(column) { - bail!("Filtering on non primary key columns is not supported"); + if !primary_key_columns.contains(column) && !filtering_columns.contains(column) { + bail!("Column '{column}' is not a primary key column or a filtering column"); }; let Some(native_type) = table_columns.get(column) else { bail!( "Column '{column}' in filter restriction is not part of the table or is not a supported native type", ) }; + // Alternator non-key filtering columns are injected as NativeType::Blob (raw `:attrs` + // bytes with type-tag). A JSON string filter value is an S-type comparison + // (CqlValue::Text) and a JSON number filter value is an N-type comparison + // (CqlValue::Decimal). cql_cmp inspects the tag byte in the stored blob to enforce + // type-correct comparisons, so e.g. the string "1" will not match a numeric filter < 5. + // Gate this on is_alternator so that genuine CQL blob columns keep the normal + // hex-encoded-string → CqlValue::Blob conversion path. + if is_alternator + && !primary_key_columns.contains(column) + && matches!(native_type, NativeType::Blob) + { + return match value { + Value::String(s) => Ok(CqlValue::Text(s)), + Value::Number(_) => try_from_json(value, &NativeType::Decimal), + _ => bail!( + "Column '{column}' is an Alternator attribute column; \ + only string (S-type) and number (N-type) filter values are supported" + ), + }; + } try_from_json(value, native_type) }; Ok(Filter { @@ -911,7 +977,9 @@ fn try_to_json(value: CqlValue) -> anyhow::Result { CqlValue::Decimal(value) => Ok(Value::String(BigDecimal::from(value).to_string())), - _ => unimplemented!(), + CqlValue::Inet(value) => Ok(Value::String(value.to_string())), + + other => bail!("unsupported CqlValue variant for JSON conversion: {other:?}"), } } @@ -1189,6 +1257,8 @@ mod tests { .unwrap(), &primary_key_columns, &table_columns, + &[], + false, ) .unwrap(); assert!(filter.allow_filtering); @@ -1267,7 +1337,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1285,7 +1357,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1301,7 +1375,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1317,7 +1393,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1333,7 +1411,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1351,7 +1431,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &[("pk".into(), NativeType::Int),].into_iter().collect() + &[("pk".into(), NativeType::Int),].into_iter().collect(), + &[], + false ) .is_err() ); @@ -1369,7 +1451,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1387,7 +1471,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1405,7 +1491,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1423,7 +1511,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1441,7 +1531,9 @@ mod tests { ) .unwrap(), &primary_key_columns, - &table_columns + &table_columns, + &[], + false ) .is_err() ); @@ -1817,6 +1909,22 @@ mod tests { .unwrap(), Value::String("-98765432109876543210.123456789".to_string()) ); + + // Inet (IPv4 and IPv6) + assert_eq!( + try_to_json(CqlValue::Inet( + "192.168.1.1".parse::().unwrap() + )) + .unwrap(), + Value::String("192.168.1.1".to_string()) + ); + assert_eq!( + try_to_json(CqlValue::Inet("::1".parse::().unwrap())).unwrap(), + Value::String("::1".to_string()) + ); + + // Unsupported variant returns an error instead of panicking. + assert!(try_to_json(CqlValue::Counter(scylla::value::Counter(0))).is_err()); } #[test] @@ -1858,4 +1966,144 @@ mod tests { httpapi::IndexStatus::Serving ); } + + // Helpers shared by the filtering-column tests below. + fn filtering_table_columns() -> HashMap { + [ + ("pk".into(), NativeType::Int), + ("ck".into(), NativeType::Int), + // A plain CQL text filtering column + ("tag".into(), NativeType::Text), + // A blob column — used for Alternator non-key attributes + ("attr".into(), NativeType::Blob), + ] + .into_iter() + .collect() + } + + #[test] + fn filter_on_non_pk_filtering_column_ok() { + // Filtering on a non-primary-key column that appears in filtering_columns succeeds. + let filter = try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "==", "lhs": "tag", "rhs": "hello" }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &["tag".into()], + false, + ) + .unwrap(); + assert!( + matches!(filter.restrictions.first(), Some(Restriction::Eq { lhs, rhs }) + if *lhs == "tag".into() && *rhs == CqlValue::Text("hello".to_string())) + ); + } + + #[test] + fn filter_on_non_pk_column_not_in_filtering_columns_fails() { + // A non-primary-key column that is NOT listed in filtering_columns must be rejected. + assert!( + try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "==", "lhs": "tag", "rhs": "hello" }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &[], // tag is not listed here + false, + ) + .is_err() + ); + } + + #[test] + fn alternator_blob_column_string_maps_to_text() { + // For Alternator indexes, a JSON string value for a Blob filtering column + // must be converted to CqlValue::Text (S-type comparison). + let filter = try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "==", "lhs": "attr", "rhs": "hello" }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &["attr".into()], + true, // is_alternator + ) + .unwrap(); + assert!( + matches!(filter.restrictions.first(), Some(Restriction::Eq { lhs, rhs }) + if *lhs == "attr".into() && *rhs == CqlValue::Text("hello".to_string())) + ); + } + + #[test] + fn alternator_blob_column_number_maps_to_decimal() { + // For Alternator indexes, a JSON number value for a Blob filtering column + // must be converted to CqlValue::Decimal (N-type comparison). + let filter = try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "<", "lhs": "attr", "rhs": 42 }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &["attr".into()], + true, // is_alternator + ) + .unwrap(); + assert!( + matches!(filter.restrictions.first(), Some(Restriction::Lt { lhs, rhs }) + if *lhs == "attr".into() && matches!(rhs, CqlValue::Decimal(_))) + ); + } + + #[test] + fn alternator_blob_column_unsupported_type_fails() { + // A boolean value for an Alternator blob filtering column must be rejected. + assert!( + try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "==", "lhs": "attr", "rhs": true }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &["attr".into()], + true, // is_alternator + ) + .is_err() + ); + } + + #[test] + fn non_alternator_blob_column_uses_hex_path() { + // For non-Alternator (plain CQL) indexes the Blob column must still accept + // a "0x"-prefixed hex string and produce CqlValue::Blob. + let filter = try_from_post_index_ann_filter( + serde_json::from_str( + r#"{ "restrictions": [{ "type": "==", "lhs": "attr", "rhs": "0xdeadbeef" }], + "allow_filtering": true }"#, + ) + .unwrap(), + &["pk".into(), "ck".into()], + &filtering_table_columns(), + &["attr".into()], + false, // is_alternator = false + ) + .unwrap(); + assert!( + matches!(filter.restrictions.first(), Some(Restriction::Eq { lhs, rhs }) + if *lhs == "attr".into() + && *rhs == CqlValue::Blob(vec![0xde, 0xad, 0xbe, 0xef])) + ); + } } diff --git a/crates/vector-store/src/index/actor.rs b/crates/vector-store/src/index/actor.rs index bc0866f3..8f291c3a 100644 --- a/crates/vector-store/src/index/actor.rs +++ b/crates/vector-store/src/index/actor.rs @@ -3,6 +3,7 @@ * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0 */ use crate::AsyncInProgress; +use crate::ColumnName; use crate::Distance; use crate::Filter; use crate::IndexKey; @@ -11,10 +12,17 @@ use crate::PrimaryKey; use crate::Vector; use crate::table::PartitionId; use crate::table::PrimaryId; +use scylla::value::CqlValue; +use std::collections::BTreeMap; +use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::oneshot; -pub(crate) type AnnR = anyhow::Result<(Vec, Vec)>; +pub(crate) type AnnR = anyhow::Result<( + Vec, + Vec, + Vec>, +)>; pub(crate) type CountR = anyhow::Result; pub enum Index { @@ -36,6 +44,7 @@ pub enum Index { index_key: IndexKey, embedding: Vector, limit: Limit, + return_columns: Arc<[ColumnName]>, tx: oneshot::Sender, }, FilteredAnn { @@ -43,6 +52,7 @@ pub enum Index { embedding: Vector, filter: Filter, limit: Limit, + return_columns: Arc<[ColumnName]>, tx: oneshot::Sender, }, Count { @@ -66,13 +76,20 @@ pub(crate) trait IndexExt { in_progress: Option, ); async fn remove_partition(&self, partition_id: PartitionId); - async fn ann(&self, index_key: IndexKey, embedding: Vector, limit: Limit) -> AnnR; + async fn ann( + &self, + index_key: IndexKey, + embedding: Vector, + limit: Limit, + return_columns: Arc<[ColumnName]>, + ) -> AnnR; async fn filtered_ann( &self, index_key: IndexKey, embedding: Vector, filter: Filter, limit: Limit, + return_columns: Arc<[ColumnName]>, ) -> AnnR; async fn count(&self, index_key: IndexKey) -> CountR; } @@ -120,12 +137,19 @@ impl IndexExt for mpsc::Sender { } #[hotpath::measure] - async fn ann(&self, index_key: IndexKey, embedding: Vector, limit: Limit) -> AnnR { + async fn ann( + &self, + index_key: IndexKey, + embedding: Vector, + limit: Limit, + return_columns: Arc<[ColumnName]>, + ) -> AnnR { let (tx, rx) = oneshot::channel(); self.send(Index::Ann { index_key, embedding, limit, + return_columns, tx, }) .await?; @@ -139,6 +163,7 @@ impl IndexExt for mpsc::Sender { embedding: Vector, filter: Filter, limit: Limit, + return_columns: Arc<[ColumnName]>, ) -> AnnR { let (tx, rx) = oneshot::channel(); self.send(Index::FilteredAnn { @@ -146,6 +171,7 @@ impl IndexExt for mpsc::Sender { embedding, filter, limit, + return_columns, tx, }) .await?; diff --git a/crates/vector-store/src/index/opensearch.rs b/crates/vector-store/src/index/opensearch.rs index 44d968ce..0d160473 100644 --- a/crates/vector-store/src/index/opensearch.rs +++ b/crates/vector-store/src/index/opensearch.rs @@ -3,6 +3,7 @@ * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0 */ +use crate::ColumnName; use crate::Connectivity; use crate::Dimensions; use crate::Distance; @@ -295,11 +296,20 @@ async fn process( Index::Ann { embedding, limit, + return_columns, tx, .. } => { ann( - index_key, tx, embedding, dimensions, limit, space_type, table, client, + index_key, + tx, + embedding, + dimensions, + limit, + space_type, + return_columns, + table, + client, ) .await } @@ -360,6 +370,7 @@ async fn ann( dimensions: Dimensions, limit: Limit, space_type: SpaceType, + return_columns: Arc<[ColumnName]>, table: Arc>, client: Arc, ) { @@ -417,7 +428,7 @@ async fn ann( let index_id = IndexIdGenerator::new().next(true).unwrap(); let partition_id = PartitionId::global(index_id); - let hits = { + let (keys, scores, column_values_vec) = { let table = table.read().unwrap(); hits.unwrap() .iter() @@ -426,12 +437,16 @@ async fn ann( let score = hit["_score"].as_f64().unwrap(); let primary_id = PrimaryId::from(id.parse::().unwrap()); let primary_key = table.primary_key(partition_id, primary_id).unwrap(); - (primary_key, score) + let col_vals = if return_columns.is_empty() { + Default::default() + } else { + table.column_values_for(partition_id, primary_id, &return_columns) + }; + (primary_key, score, col_vals) }) - .collect::>() + .collect::<(Vec<_>, Vec<_>, Vec<_>)>() }; - let (keys, scores): (Vec<_>, Vec<_>) = hits.iter().cloned().unzip(); let distances: anyhow::Result> = scores .iter() .map(|score| Distance::try_from((*score as f32, space_type, Some(dimensions)))) @@ -445,7 +460,7 @@ async fn ann( }; tx_ann - .send(Ok((keys, distances))) + .send(Ok((keys, distances, column_values_vec))) .unwrap_or_else(|_| trace!("ann: unable to send response")); } diff --git a/crates/vector-store/src/index/usearch.rs b/crates/vector-store/src/index/usearch.rs index 0a4154c1..c4cdb4d3 100644 --- a/crates/vector-store/src/index/usearch.rs +++ b/crates/vector-store/src/index/usearch.rs @@ -3,6 +3,7 @@ * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0 */ +use crate::ColumnName; use crate::Config; use crate::Dimensions; use crate::Distance; @@ -30,7 +31,6 @@ use crate::worker; use crate::worker::Worker; use crate::worker::WorkerExt; use anyhow::anyhow; -use itertools::Itertools; use std::collections::BTreeMap; use std::collections::BTreeSet; use std::sync::Arc; @@ -774,12 +774,13 @@ where index_key, embedding, limit, + return_columns, tx, } => { let Some((partition_id, _)) = table.read().unwrap().partition_id(&index_key, None) else { warn!("partition id not found for index key {index_key:?} during ann"); - _ = tx.send(Ok((vec![], vec![]))); + _ = tx.send(Ok((vec![], vec![], vec![]))); return None; }; let index_id = partition_id.index_id(); @@ -789,7 +790,7 @@ where .map(|(state, partition)| (state, Arc::clone(partition))) else { warn!("state or partition not found for index key {index_key:?} during ann"); - _ = tx.send(Ok((vec![], vec![]))); + _ = tx.send(Ok((vec![], vec![], vec![]))); return None; }; Some(( @@ -798,6 +799,7 @@ where Index::Ann { embedding, limit, + return_columns, tx, index_key, }, @@ -809,6 +811,7 @@ where embedding, filter, limit, + return_columns, tx, } => { let Some((partition_id, restrictions)) = table @@ -817,7 +820,7 @@ where .partition_id(&index_key, Some(filter.restrictions)) else { warn!("partition id not found for index key {index_key:?} during filtered ann"); - _ = tx.send(Ok((vec![], vec![]))); + _ = tx.send(Ok((vec![], vec![], vec![]))); return None; }; let index_id = partition_id.index_id(); @@ -830,7 +833,7 @@ where "state or partition not found for index key {index_key:?} \ during filtered ann" ); - _ = tx.send(Ok((vec![], vec![]))); + _ = tx.send(Ok((vec![], vec![], vec![]))); return None; }; let msg = if let Some(restrictions) = restrictions { @@ -841,6 +844,7 @@ where restrictions, allow_filtering: filter.allow_filtering, }, + return_columns, tx, index_key, } @@ -848,6 +852,7 @@ where Index::Ann { embedding, limit, + return_columns, tx, index_key, } @@ -961,11 +966,12 @@ fn process( Index::Ann { embedding, limit, + return_columns, tx, .. } => { if let Some(tx) = validate_dimensions(tx, &embedding, dimensions) { - ann(partition, tx, &table, embedding, limit); + ann(partition, tx, &table, embedding, limit, return_columns); } } @@ -973,11 +979,20 @@ fn process( embedding, limit, filter, + return_columns, tx, .. } => { if let Some(tx) = validate_dimensions(tx, &embedding, dimensions) { - filtered_ann(partition, tx, &table, embedding, filter, limit); + filtered_ann( + partition, + tx, + &table, + embedding, + filter, + limit, + return_columns, + ); } } @@ -1063,6 +1078,7 @@ fn ann( table: &Arc>, embedding: Vector, limit: Limit, + return_columns: Arc<[ColumnName]>, ) where I: UsearchIndex + Send + Sync + 'static, { @@ -1074,23 +1090,35 @@ fn ann( .map_err(|err| anyhow!("ann: search failed: {err}")) .and_then(|matches| { let table = table.read().unwrap(); - let (primary_keys, distances) = itertools::process_results( - matches.filter_map_ok(|(primary_id, distance)| { - table - .primary_key(partition.partition_id, primary_id) - .or_else(|| { - debug!( - "not defined primary key for partition_id {partition_id:?} \ - and primary_id {primary_id:?}", - partition_id = partition.partition_id, - ); - None - }) - .map(|primary_key| (primary_key, distance)) - }), - |it| it.unzip(), - )?; - Ok((primary_keys, distances)) + let mut primary_keys = Vec::new(); + let mut distances = Vec::new(); + let mut column_values_vec = Vec::new(); + for result in matches { + let (primary_id, distance) = result?; + let Some(primary_key) = + table.primary_key(partition.partition_id, primary_id) + else { + debug!( + "not defined primary key for partition_id {partition_id:?} \ + and primary_id {primary_id:?}", + partition_id = partition.partition_id, + ); + continue; + }; + let col_vals = if return_columns.is_empty() { + BTreeMap::new() + } else { + table.column_values_for( + partition.partition_id, + primary_id, + &return_columns, + ) + }; + primary_keys.push(primary_key); + distances.push(distance); + column_values_vec.push(col_vals); + } + Ok((primary_keys, distances, column_values_vec)) }), ) .unwrap_or_else(|_| trace!("ann: unable to send response")); @@ -1104,6 +1132,7 @@ fn filtered_ann( embedding: Vector, filter: Filter, limit: Limit, + return_columns: Arc<[ColumnName]>, ) where I: UsearchIndex + Send + Sync + 'static, { @@ -1123,23 +1152,35 @@ fn filtered_ann( .map_err(|err| anyhow!("ann: search failed: {err}")) .and_then(|matches| { let table = table.read().unwrap(); - let (primary_keys, distances) = itertools::process_results( - matches.filter_map_ok(|(primary_id, distance)| { - table - .primary_key(partition.partition_id, primary_id) - .or_else(|| { - debug!( - "not defined primary key for partition_id {partition_id:?} \ - and primary_id {primary_id:?}", - partition_id = partition.partition_id, - ); - None - }) - .map(|primary_key| (primary_key, distance)) - }), - |it| it.unzip(), - )?; - Ok((primary_keys, distances)) + let mut primary_keys = Vec::new(); + let mut distances = Vec::new(); + let mut column_values_vec = Vec::new(); + for result in matches { + let (primary_id, distance) = result?; + let Some(primary_key) = + table.primary_key(partition.partition_id, primary_id) + else { + debug!( + "not defined primary key for partition_id {partition_id:?} \ + and primary_id {primary_id:?}", + partition_id = partition.partition_id, + ); + continue; + }; + let col_vals = if return_columns.is_empty() { + BTreeMap::new() + } else { + table.column_values_for( + partition.partition_id, + primary_id, + &return_columns, + ) + }; + primary_keys.push(primary_key); + distances.push(distance); + column_values_vec.push(col_vals); + } + Ok((primary_keys, distances, column_values_vec)) }), ) .unwrap_or_else(|_| trace!("ann: unable to send response")); @@ -1261,6 +1302,7 @@ mod tests { index_key.clone(), vec![0.0f32; dimensions.get()].into(), limit, + Arc::new([]), ) .await; } @@ -1334,11 +1376,12 @@ mod tests { .once() .returning(|_, _| Some([CqlValue::Int(2)].into())); - let (primary_keys, distances) = actor + let (primary_keys, distances, _) = actor .ann( index_key.clone(), vec![2.2, -2.2, 2.2].into(), NonZeroUsize::new(1).unwrap().into(), + Arc::new([]), ) .await .unwrap(); @@ -1365,6 +1408,7 @@ mod tests { index_key.clone(), vec![2.2, -2.2, 2.2].into(), NonZeroUsize::new(1).unwrap().into(), + Arc::new([]), ) .await .unwrap() @@ -1397,11 +1441,12 @@ mod tests { .once() .returning(|_, _| Some([CqlValue::Int(2)].into())); - let (primary_keys, distances) = actor + let (primary_keys, distances, _) = actor .ann( index_key, vec![2.2, -2.2, 2.2].into(), NonZeroUsize::new(1).unwrap().into(), + Arc::new([]), ) .await .unwrap(); diff --git a/crates/vector-store/src/indexes.rs b/crates/vector-store/src/indexes.rs index f15d671c..de99442d 100644 --- a/crates/vector-store/src/indexes.rs +++ b/crates/vector-store/src/indexes.rs @@ -213,6 +213,8 @@ pub(crate) enum BestIndexState { primary_key_columns: Arc>, table_columns: Arc>, needs_filtering: NeedsFiltering, + /// Filtering columns for the selected index (includes primary key columns). + filtering_columns: Arc>, }, } @@ -320,6 +322,7 @@ impl Indexes { primary_key_columns: Arc::clone(&routed_entry.primary_key_columns), table_columns: Arc::clone(&routed_entry.table_columns), needs_filtering: needs_filtering.clone(), + filtering_columns: Arc::clone(&routed_entry.filtering_columns), } } None => BestIndexState::NotServing(requested_entry.progress), diff --git a/crates/vector-store/src/lib.rs b/crates/vector-store/src/lib.rs index c8395cb0..22d5818d 100644 --- a/crates/vector-store/src/lib.rs +++ b/crates/vector-store/src/lib.rs @@ -62,6 +62,7 @@ use scylla::serialize::writers::CellWriter; use scylla::serialize::writers::WrittenCellProof; use scylla::value::CqlValue; use scylla_cdc::CqlIdentifier; +use std::collections::BTreeMap; use std::collections::HashMap; use std::hash::Hash; use std::net::SocketAddr; @@ -173,6 +174,7 @@ pub struct Config { pub cdc_fine_safety_interval: Option, pub cdc_fine_sleep_interval: Option, pub monitor_indexes_interval: Option, + pub engine_status_update_interval: Option, pub disable_colors: bool, pub tls_cert_path: Option, pub tls_key_path: Option, @@ -207,6 +209,7 @@ impl Default for Config { cdc_fine_safety_interval: None, cdc_fine_sleep_interval: None, monitor_indexes_interval: None, + engine_status_update_interval: None, } } } @@ -237,7 +240,7 @@ 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 { + pub(crate) fn is_alternator(&self) -> bool { self.0.starts_with("alternator_") } } @@ -608,8 +611,18 @@ impl DbCustomIndex { #[derive(Clone, Debug, PartialEq)] pub struct DbEmbedding { pub primary_key: PrimaryKey, - pub embedding: Option, + /// `None` = vector unchanged (skip update), `Some(None)` = vector deleted, + /// `Some(Some(v))` = vector set to `v`. + pub embedding: Option>, pub timestamp: Timestamp, + /// Values for filtering columns, keyed by column name. + /// For Alternator tables, these are extracted from the `:attrs` map. + /// Each value carries the write timestamp of that column so that the index + /// layer can apply last-writer-wins per column independently of the vector + /// column's timestamp. `None` represents a tombstone: the column was + /// explicitly absent or deleted in this write, which clears any + /// previously-stored value for this row. + pub column_values: BTreeMap)>, } #[derive(Clone, derive_more::From)] diff --git a/crates/vector-store/src/monitor_items.rs b/crates/vector-store/src/monitor_items.rs index bfe254be..a85bb22e 100644 --- a/crates/vector-store/src/monitor_items.rs +++ b/crates/vector-store/src/monitor_items.rs @@ -174,8 +174,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: Some(vec![1.].into()), + embedding: Some(Some(vec![1.].into())), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() @@ -209,8 +210,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: Some(vec![1.].into()), + embedding: Some(Some(vec![1.].into())), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; let (tx_progress, _rx_progress) = mpsc::channel(1); table @@ -269,8 +271,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: Some(vec![1.].into()), + embedding: Some(Some(vec![1.].into())), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() @@ -325,8 +328,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: Some(vec![1.].into()), + embedding: Some(Some(vec![1.].into())), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() @@ -398,8 +402,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: Some(vec![1.].into()), + embedding: Some(Some(vec![1.].into())), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() @@ -496,8 +501,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: None, + embedding: Some(None), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() @@ -548,8 +554,9 @@ mod tests { let embedding = DbEmbedding { primary_key: [CqlValue::Int(1)].into(), - embedding: None, + embedding: Some(None), timestamp: Timestamp::from_unix_timestamp(10), + column_values: Default::default(), }; table .write() diff --git a/crates/vector-store/src/table.rs b/crates/vector-store/src/table.rs index 2b6bcd48..25f0ec22 100644 --- a/crates/vector-store/src/table.rs +++ b/crates/vector-store/src/table.rs @@ -12,6 +12,9 @@ use crate::Restriction; use crate::Timestamp; use crate::Vector; use crate::primary_key::normalize; +use crate::vector::ALTERNATOR_TYPE_JSON; +use crate::vector::ALTERNATOR_TYPE_N; +use crate::vector::ALTERNATOR_TYPE_S; use anyhow::anyhow; use anyhow::bail; use bigdecimal::BigDecimal; @@ -325,6 +328,12 @@ enum TValue { Some(Timestamp, T), } impl TValue { + fn timestamp(&self) -> Timestamp { + match self { + Self::None(t) | Self::Some(t, _) => *t, + } + } + fn get(&self) -> Option<&T> { match self { Self::None(_) => None, @@ -412,6 +421,30 @@ impl Column { } } + fn stored_timestamp(&self, primary_id: PrimaryId) -> Option { + match self { + Self::Ascii(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::BigInt(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Blob(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Boolean(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Date(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Decimal(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Double(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Float(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Inet(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Int(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::SmallInt(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Text(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Time(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Timestamp(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Timeuuid(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::TinyInt(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Uuid(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::Varint(vec) => vec.get(primary_id).map(|v| v.timestamp()), + Self::PrimaryKey(_) => None, + } + } + #[allow(dead_code)] fn insert_cqlvalue( &mut self, @@ -532,6 +565,37 @@ impl Column { } } + /// Records a tombstone for this column at the given primary id and timestamp, clearing any + /// previously-stored value. Analogous to [`Self::insert_cqlvalue`] but writes + /// [`TValue::None`] so the column appears absent to the filter evaluator. + fn insert_tombstone( + &mut self, + primary_id: PrimaryId, + timestamp: Timestamp, + ) -> anyhow::Result<()> { + match self { + Self::Ascii(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::BigInt(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Blob(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Boolean(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Date(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Decimal(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Double(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Float(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Inet(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Int(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::SmallInt(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Text(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Time(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Timestamp(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Timeuuid(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::TinyInt(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Uuid(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::Varint(vec) => vec.update(primary_id, TValue::None(timestamp)), + Self::PrimaryKey(_) => bail!("Cannot insert tombstone into PrimaryKey column"), + } + } + fn get( &self, primary_id: PrimaryId, @@ -1084,6 +1148,7 @@ impl TableAdd for Table { let primary_key = db_embedding.primary_key; let vector = db_embedding.embedding; + let column_values = db_embedding.column_values; let normalized_key = self.normalize_primary_key(&primary_key); let row_map = &mut self.primary_ids; @@ -1115,47 +1180,77 @@ impl TableAdd for Table { missing partition id for index_id {index_id:?} and primary_id {primary_id:?}" ) })?; - if let Some(vector) = &vector { - if vector_already_exists { - operations.push(Operation::RemoveBeforeAddVector { - primary_id, - partition_id, - }); - } + match &vector { + Some(Some(vector)) => { + if vector_already_exists { + operations.push(Operation::RemoveBeforeAddVector { + primary_id, + partition_id, + }); + } - let primary_id = primary_id.next_epoch(); - let timestamp = db_embedding.timestamp; - operations.push(Operation::AddVector { - primary_id, - partition_id, - vector: vector.clone(), - is_update: vector_already_exists, - }); - index - .vector_timestamps - .update_epoch_timestamp(primary_id, timestamp)?; - } else { - let epoch = primary_id.epoch().next(); - index - .vector_timestamps - .update(primary_id, ETValue::None(epoch, *timestamp))?; - if vector_already_exists { - operations.push(Operation::RemoveVector { + let primary_id = primary_id.next_epoch(); + let timestamp = db_embedding.timestamp; + operations.push(Operation::AddVector { primary_id, partition_id, + vector: vector.clone(), + is_update: vector_already_exists, }); - if index.data.remove_row(primary_id) { - operations.push(Operation::RemovePartition { partition_id }); + index + .vector_timestamps + .update_epoch_timestamp(primary_id, timestamp)?; + } + Some(None) => { + // Explicit deletion: clear the vector. + let epoch = primary_id.epoch().next(); + index + .vector_timestamps + .update(primary_id, ETValue::None(epoch, *timestamp))?; + if vector_already_exists { + operations.push(Operation::RemoveVector { + primary_id, + partition_id, + }); + if index.data.remove_row(primary_id) { + operations.push(Operation::RemovePartition { partition_id }); + } } } + None => { + // Vector unchanged; skip the vector update entirely. + } } Ok(()) })?; + // Update filtering column values for this row. + // Apply last-writer-wins per column using the column's own timestamp, + // which may differ from the vector's timestamp. + // A `None` value is a tombstone that clears a previously-stored value. + for (col_name, (col_ts, val)) in column_values { + if let Some(col) = self.columns.get_mut(&col_name) { + if col + .stored_timestamp(primary_id) + .is_some_and(|t| t >= col_ts) + { + continue; + } + let result = match val { + Some(val) => { + col.insert_cqlvalue(primary_id, col_ts, val) + } + None => col.insert_tombstone(primary_id, col_ts), + }; + if let Err(err) = result { + warn!("Failed to update filtering column '{col_name}': {err}"); + } + } + } Ok(operations) } Entry::Vacant(entry) => { - if let Some(vector) = &vector { + if let Some(Some(vector)) = &vector { let primary_id = self.free_primary_ids.take_id()?; entry.insert(primary_id); self.primary_keys @@ -1188,9 +1283,25 @@ impl TableAdd for Table { Ok(()) }, )?; - } else { - warn!("Added row with no vector, skipping vector addition"); + // Store filtering column values for this new row. + // A `None` value is a tombstone that records the column as absent at + // this timestamp, preventing stale out-of-order CDC updates from + // setting it to a value that was already cleared. + for (col_name, (col_ts, val)) in column_values { + if let Some(col) = self.columns.get_mut(&col_name) { + let result = match val { + Some(val) => { + col.insert_cqlvalue(primary_id, col_ts, val) + } + None => col.insert_tombstone(primary_id, col_ts), + }; + if let Err(err) = result { + warn!("Failed to store filtering column '{col_name}': {err}"); + } + } + } } + // None or Some(None): no vector to add (skip or deletion of non-existent row). Ok(operations) } } @@ -1216,6 +1327,18 @@ pub(crate) trait TableSearch { primary_id: PrimaryId, restriction: &Restriction, ) -> bool; + + /// Return the stored values for the given filtering columns for a single + /// row identified by `primary_id`. Columns that have no stored value + /// (e.g. the attribute was absent when the row was indexed) are omitted + /// from the returned map. Returns an empty map when `columns` is empty + /// or the `primary_id` is invalid. + fn column_values_for( + &self, + partition_id: PartitionId, + primary_id: PrimaryId, + columns: &[ColumnName], + ) -> std::collections::BTreeMap; } impl TableSearch for Table { @@ -1351,6 +1474,26 @@ impl TableSearch for Table { } } } + + fn column_values_for( + &self, + partition_id: PartitionId, + primary_id: PrimaryId, + columns: &[ColumnName], + ) -> std::collections::BTreeMap { + if !self.is_valid_primary_id(partition_id, primary_id) { + return std::collections::BTreeMap::new(); + } + columns + .iter() + .filter_map(|col_name| { + self.columns + .get(col_name) + .and_then(|col| col.get(primary_id, &self.primary_keys)) + .map(|value| (col_name.clone(), value)) + }) + .collect() + } } /// Construct a partition key from the given restrictions. @@ -1392,6 +1535,54 @@ fn partition_key_from_restrictions( }) } +/// Compare an Alternator `:attrs` blob (S-type) against a Text filter restriction. +/// Returns None if the blob is not an S-type attribute (type mismatch). +fn alternator_blob_cmp_text(blob: &[u8], b: &str) -> Option { + let (&tag, rest) = blob.split_first()?; + match tag { + ALTERNATOR_TYPE_S => { + let a = std::str::from_utf8(rest).ok()?; + Some(a.cmp(b)) + } + ALTERNATOR_TYPE_JSON => { + let obj: serde_json::Value = serde_json::from_slice(rest).ok()?; + let s = obj.as_object()?.get("S")?.as_str()?; + Some(s.cmp(b)) + } + _ => None, // N-type or unknown vs Text filter → type mismatch + } +} + +/// Compare an Alternator `:attrs` blob (N-type) against a Decimal filter restriction. +/// Returns None if the blob is not an N-type attribute (type mismatch). +fn alternator_blob_cmp_decimal(blob: &[u8], b: &CqlDecimal) -> Option { + let (&tag, rest) = blob.split_first()?; + let (b_bytes, b_scale) = b.as_signed_be_bytes_slice_and_exponent(); + let b_bd = BigDecimal::from(CqlDecimalBorrowed::from_signed_be_bytes_slice_and_exponent( + b_bytes, b_scale, + )); + let a_bd: BigDecimal = match tag { + ALTERNATOR_TYPE_N => { + if rest.len() < 4 { + return None; + } + let scale = i32::from_be_bytes([rest[0], rest[1], rest[2], rest[3]]); + let varint_bytes = &rest[4..]; + BigDecimal::from(CqlDecimalBorrowed::from_signed_be_bytes_slice_and_exponent( + varint_bytes, + scale, + )) + } + ALTERNATOR_TYPE_JSON => { + let obj: serde_json::Value = serde_json::from_slice(rest).ok()?; + let n_str = obj.as_object()?.get("N")?.as_str()?; + n_str.parse().ok()? + } + _ => return None, // S-type or unknown vs Decimal filter → type mismatch + }; + Some(a_bd.cmp(&b_bd)) +} + /// Compare two CqlValues, returning an Ordering if they are comparable. /// Only Numeric, Text, Date, Time, and Timestamp types support comparison operators. fn cql_cmp(lhs: &CqlValue, rhs: &CqlValue) -> Option { @@ -1425,6 +1616,21 @@ fn cql_cmp(lhs: &CqlValue, rhs: &CqlValue) -> Option { ); Some(a_bd.cmp(&b_bd)) } + // Cross-type: Blob (raw Alternator `:attrs` value with type-tag byte) vs Text or Decimal. + // The tag byte determines the DynamoDB attribute type: + // ALTERNATOR_TYPE_S (0x00): S-type, payload = raw UTF-8 bytes + // ALTERNATOR_TYPE_N (0x03): N-type, payload = CQL decimal (4-byte BE scale + BE varint) + // ALTERNATOR_TYPE_JSON (0x04): JSON-wrapped {"S":...} or {"N":...} + // A type mismatch (S-type blob vs Decimal restriction, or N-type blob vs Text restriction) + // returns None, so that e.g. the string "1" does not match a numeric filter < 5. + (CqlValue::Blob(blob), CqlValue::Text(b)) => alternator_blob_cmp_text(blob, b), + (CqlValue::Text(a), CqlValue::Blob(blob)) => { + alternator_blob_cmp_text(blob, a).map(Ordering::reverse) + } + (CqlValue::Blob(blob), CqlValue::Decimal(b)) => alternator_blob_cmp_decimal(blob, b), + (CqlValue::Decimal(a), CqlValue::Blob(blob)) => { + alternator_blob_cmp_decimal(blob, a).map(Ordering::reverse) + } // Text types (CqlValue::Text(a), CqlValue::Text(b)) => Some(a.cmp(b)), (CqlValue::Ascii(a), CqlValue::Ascii(b)) => Some(a.cmp(b)), @@ -1519,8 +1725,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(1)].into(), - embedding: Some(vec![0.1, 0.2, 0.3].into()), + embedding: Some(Some(vec![0.1, 0.2, 0.3].into())), timestamp: Timestamp::from_unix_timestamp(100), + column_values: Default::default(), }, ) .unwrap(); @@ -1544,8 +1751,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(2)].into(), - embedding: Some(vec![0.2, 0.2, 0.3].into()), + embedding: Some(Some(vec![0.2, 0.2, 0.3].into())), timestamp: Timestamp::from_unix_timestamp(100), + column_values: Default::default(), }, ) .unwrap(); @@ -1571,8 +1779,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(3)].into(), - embedding: Some(vec![0.3, 0.2, 0.3].into()), + embedding: Some(Some(vec![0.3, 0.2, 0.3].into())), timestamp: Timestamp::from_unix_timestamp(100), + column_values: Default::default(), }, ) .unwrap(); @@ -1638,8 +1847,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(2)].into(), - embedding: Some(vec![0.2, 0.2, 0.3].into()), + embedding: Some(Some(vec![0.2, 0.2, 0.3].into())), timestamp: Timestamp::from_unix_timestamp(50), + column_values: Default::default(), }, ) .unwrap(); @@ -1651,8 +1861,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(2)].into(), - embedding: Some(vec![0.5, 0.5, 0.3].into()), + embedding: Some(Some(vec![0.5, 0.5, 0.3].into())), timestamp: Timestamp::from_unix_timestamp(150), + column_values: Default::default(), }, ) .unwrap(); @@ -1693,8 +1904,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(1)].into(), - embedding: None, + embedding: Some(None), timestamp: Timestamp::from_unix_timestamp(200), + column_values: Default::default(), }, ) .unwrap(); @@ -1716,8 +1928,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(2)].into(), - embedding: None, + embedding: Some(None), timestamp: Timestamp::from_unix_timestamp(200), + column_values: Default::default(), }, ) .unwrap(); @@ -1739,8 +1952,9 @@ mod tests { &index_key, DbEmbedding { primary_key: [CqlValue::Int(1), CqlValue::Int(3)].into(), - embedding: None, + embedding: Some(None), timestamp: Timestamp::from_unix_timestamp(200), + column_values: Default::default(), }, ) .unwrap(); @@ -2174,4 +2388,246 @@ mod tests { )) ); } + + // Helpers shared by the Alternator blob comparison tests below. + fn s_blob(s: &str) -> CqlValue { + let mut b = vec![ALTERNATOR_TYPE_S]; + b.extend_from_slice(s.as_bytes()); + CqlValue::Blob(b) + } + + // N-type blob: [ALTERNATOR_TYPE_N] + 4-byte BE scale + signed BE varint. + // value = unscaled_bigint * 10^(-scale) + fn n_blob(scale: i32, unscaled: &[u8]) -> CqlValue { + let mut b = vec![ALTERNATOR_TYPE_N]; + b.extend_from_slice(&scale.to_be_bytes()); + b.extend_from_slice(unscaled); + CqlValue::Blob(b) + } + + fn json_blob(json: &str) -> CqlValue { + let mut b = vec![ALTERNATOR_TYPE_JSON]; + b.extend_from_slice(json.as_bytes()); + CqlValue::Blob(b) + } + + fn decimal(s: &str) -> CqlValue { + CqlValue::Decimal(CqlDecimal::try_from(s.parse::().unwrap()).unwrap()) + } + + #[test] + fn alternator_blob_text_s_type() { + assert_eq!( + cql_cmp(&s_blob("apple"), &CqlValue::Text("apple".into())), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&s_blob("apple"), &CqlValue::Text("banana".into())), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&s_blob("banana"), &CqlValue::Text("apple".into())), + Some(Ordering::Greater) + ); + // empty S-type payload is a valid empty string + assert_eq!( + cql_cmp(&s_blob(""), &CqlValue::Text("".into())), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&s_blob(""), &CqlValue::Text("a".into())), + Some(Ordering::Less) + ); + } + + #[test] + fn alternator_blob_text_s_type_reversed() { + // (Text, Blob) path should be the mirror image + assert_eq!( + cql_cmp(&CqlValue::Text("apple".into()), &s_blob("apple")), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&CqlValue::Text("apple".into()), &s_blob("banana")), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&CqlValue::Text("banana".into()), &s_blob("apple")), + Some(Ordering::Greater) + ); + } + + #[test] + fn alternator_blob_text_json_s_type() { + assert_eq!( + cql_cmp( + &json_blob(r#"{"S":"hello"}"#), + &CqlValue::Text("hello".into()) + ), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&json_blob(r#"{"S":"aaa"}"#), &CqlValue::Text("bbb".into())), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&json_blob(r#"{"S":"bbb"}"#), &CqlValue::Text("aaa".into())), + Some(Ordering::Greater) + ); + } + + #[test] + fn alternator_blob_text_type_mismatch() { + // N-type blob vs Text -> type mismatch -> None + assert_eq!( + cql_cmp(&n_blob(0, &[0x05]), &CqlValue::Text("5".into())), + None + ); + // JSON N-type blob vs Text -> None (no "S" key in the object) + assert_eq!( + cql_cmp(&json_blob(r#"{"N":"5"}"#), &CqlValue::Text("5".into())), + None + ); + // Unknown tag byte -> None + assert_eq!( + cql_cmp( + &CqlValue::Blob(vec![0x99, b'x']), + &CqlValue::Text("x".into()) + ), + None + ); + } + + #[test] + fn alternator_blob_text_edge_cases() { + // Empty blob (no tag byte) -> None + assert_eq!( + cql_cmp(&CqlValue::Blob(vec![]), &CqlValue::Text("x".into())), + None + ); + // Invalid UTF-8 in S-type payload -> None + assert_eq!( + cql_cmp( + &CqlValue::Blob(vec![ALTERNATOR_TYPE_S, 0xFF, 0xFE]), + &CqlValue::Text("x".into()) + ), + None + ); + // Malformed JSON -> None + assert_eq!( + cql_cmp(&json_blob("not json"), &CqlValue::Text("x".into())), + None + ); + // JSON with neither "S" nor correct key -> None + assert_eq!( + cql_cmp( + &json_blob(r#"{"X":"hello"}"#), + &CqlValue::Text("hello".into()) + ), + None + ); + } + + #[test] + fn alternator_blob_decimal_n_type() { + // 5 = 5 * 10^0 : scale=0, unscaled=[0x05] + assert_eq!( + cql_cmp(&n_blob(0, &[0x05]), &decimal("5")), + Some(Ordering::Equal) + ); + // 1.5 = 15 * 10^-1 : scale=1, unscaled=[0x0F] + assert_eq!( + cql_cmp(&n_blob(1, &[0x0F]), &decimal("1.5")), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&n_blob(0, &[0x01]), &decimal("2")), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&n_blob(0, &[0x03]), &decimal("2")), + Some(Ordering::Greater) + ); + // negative: -1 = 0xFF in signed BE, scale=0 + assert_eq!( + cql_cmp(&n_blob(0, &[0xFF]), &decimal("0")), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&n_blob(0, &[0xFF]), &decimal("-1")), + Some(Ordering::Equal) + ); + } + + #[test] + fn alternator_blob_decimal_n_type_reversed() { + // (Decimal, Blob) path + assert_eq!( + cql_cmp(&decimal("1.5"), &n_blob(1, &[0x0F])), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&decimal("2"), &n_blob(1, &[0x0F])), + Some(Ordering::Greater) + ); + assert_eq!( + cql_cmp(&decimal("1"), &n_blob(1, &[0x0F])), + Some(Ordering::Less) + ); + } + + #[test] + fn alternator_blob_decimal_json_n_type() { + assert_eq!( + cql_cmp(&json_blob(r#"{"N":"3.14"}"#), &decimal("3.14")), + Some(Ordering::Equal) + ); + assert_eq!( + cql_cmp(&json_blob(r#"{"N":"1"}"#), &decimal("2")), + Some(Ordering::Less) + ); + assert_eq!( + cql_cmp(&json_blob(r#"{"N":"3"}"#), &decimal("2")), + Some(Ordering::Greater) + ); + } + + #[test] + fn alternator_blob_decimal_type_mismatch() { + // S-type blob vs Decimal -> None + assert_eq!(cql_cmp(&s_blob("1"), &decimal("1")), None); + // JSON S-type blob vs Decimal -> None (no "N" key in the object) + assert_eq!(cql_cmp(&json_blob(r#"{"S":"1"}"#), &decimal("1")), None); + // Unknown tag -> None + assert_eq!( + cql_cmp(&CqlValue::Blob(vec![0x99, 0, 0, 0, 0, 1]), &decimal("1")), + None + ); + } + + #[test] + fn alternator_blob_decimal_edge_cases() { + // Empty blob -> None + assert_eq!(cql_cmp(&CqlValue::Blob(vec![]), &decimal("1")), None); + // N-type but payload too short for 4-byte scale -> None + assert_eq!( + cql_cmp( + &CqlValue::Blob(vec![ALTERNATOR_TYPE_N, 0, 0]), + &decimal("1") + ), + None + ); + // Malformed JSON -> None + assert_eq!(cql_cmp(&json_blob("not json"), &decimal("1")), None); + // JSON with no "N" key -> None + assert_eq!( + cql_cmp(&json_blob(r#"{"X":"3.14"}"#), &decimal("3.14")), + None + ); + // JSON "N" value that is not a valid decimal string -> None + assert_eq!( + cql_cmp(&json_blob(r#"{"N":"not-a-number"}"#), &decimal("1")), + None + ); + } } diff --git a/crates/vector-store/src/vector.rs b/crates/vector-store/src/vector.rs index 8f421a0a..c5c5d1e8 100644 --- a/crates/vector-store/src/vector.rs +++ b/crates/vector-store/src/vector.rs @@ -7,6 +7,7 @@ use crate::Dimensions; use anyhow::anyhow; use anyhow::bail; use scylla::value::CqlValue; +use std::collections::BTreeMap; use std::num::NonZeroUsize; #[derive(Clone, Debug, PartialEq, derive_more::AsRef, derive_more::From)] @@ -59,10 +60,15 @@ impl TryFrom for Vector { } } +/// Alternator type tags, matching the `alternator_type` enum in Scylla's serialization.hh. +/// These values are written to disk and must not be reordered. +pub(crate) const ALTERNATOR_TYPE_S: u8 = 0; // String: raw UTF-8 bytes follow the tag +pub(crate) const ALTERNATOR_TYPE_N: u8 = 3; // Number: CQL decimal (4-byte big-endian scale + big-endian signed varint) follows the tag + /// Alternator type tag for unoptimized JSON encoding. /// Type `0x04` (`NOT_SUPPORTED_YET`) is used for any type that does not have an optimized encoding. /// The payload is an unoptimized JSON value. -const ALTERNATOR_TYPE_JSON: u8 = 4; +pub(crate) const ALTERNATOR_TYPE_JSON: u8 = 4; /// Alternator type tag for the optimized `FLOAT32VECTOR` type. /// The value is serialized as this 1-byte tag followed by sequential 32-bit big-endian floats, @@ -166,6 +172,65 @@ impl TryFrom> for Option { } } +/// Extracts a scalar value from an Alternator-encoded attribute blob, preserving +/// the raw bytes (including the type-tag byte) as a [`CqlValue::Blob`]. +/// +/// The type tag is preserved so that [`cql_cmp`] can later distinguish S-type (string) +/// from N-type (number) attributes and apply type-correct comparisons. Without the tag, +/// a string `"1"` and a number `1` would both be stored as `"1"` and a numeric filter +/// `< 5` would incorrectly match the string. +/// +/// Supported types: S (`0x00`), N (`0x03`), and JSON-wrapped (`0x04`). All other +/// attribute types (B, BOOL, L, M, SS, NS, BS, NULL) return `None`. +pub(crate) fn extract_alternator_scalar(blob: &[u8]) -> Option { + let (&tag, _rest) = blob.split_first()?; + match tag { + ALTERNATOR_TYPE_S | ALTERNATOR_TYPE_N | ALTERNATOR_TYPE_JSON => { + Some(CqlValue::Blob(blob.to_vec())) + } + _ => None, + } +} + +/// Extracts scalar values for the given attribute names from the Alternator `:attrs` map. +/// +/// Each entry whose key matches a requested column name and whose value is a scalar S or N +/// attribute is included in the returned map. Non-scalar or unrecognised attributes are skipped. +/// +/// Filtering columns that are absent from the `:attrs` map are **not** included. Callers that +/// need tombstones for absent columns (e.g. when the collection was fully replaced) must add +/// them separately. +pub(crate) fn extract_alternator_scalars( + attrs: &CqlValue, + columns: &[crate::ColumnName], +) -> BTreeMap> { + let CqlValue::Map(entries) = attrs else { + // `:attrs` is not a map (unexpected); return empty — callers handle tombstones. + return BTreeMap::new(); + }; + let mut result = BTreeMap::new(); + for (key, value) in entries { + let name = match key { + CqlValue::Blob(b) => std::str::from_utf8(b).ok().map(str::to_owned), + CqlValue::Text(s) => Some(s.clone()), + _ => None, + }; + let Some(name) = name else { continue }; + let col = crate::ColumnName::from(name); + if !columns.contains(&col) { + continue; + } + let blob = match value { + CqlValue::Blob(b) => b, + _ => continue, + }; + if let Some(scalar) = extract_alternator_scalar(blob) { + result.insert(col, Some(scalar)); + } + } + result +} + #[cfg(test)] mod tests { use super::*; @@ -350,4 +415,211 @@ mod tests { .unwrap(); assert_eq!(result, Some(Vector::from(vec![1.0, 2.0, 3.0]))); } + + // --- extract_alternator_scalar --- + + fn s_blob(s: &str) -> Vec { + let mut b = vec![ALTERNATOR_TYPE_S]; + b.extend_from_slice(s.as_bytes()); + b + } + + fn n_blob(scale: i32, unscaled: &[u8]) -> Vec { + let mut b = vec![ALTERNATOR_TYPE_N]; + b.extend_from_slice(&scale.to_be_bytes()); + b.extend_from_slice(unscaled); + b + } + + fn json_s_blob(s: &str) -> Vec { + let json = format!(r#"{{"S":"{}"}}"#, s); + let mut b = vec![ALTERNATOR_TYPE_JSON]; + b.extend_from_slice(json.as_bytes()); + b + } + + #[test] + fn extract_scalar_s_type_returns_raw_blob() { + let raw = s_blob("hello"); + let result = extract_alternator_scalar(&raw); + assert_eq!(result, Some(CqlValue::Blob(raw))); + } + + #[test] + fn extract_scalar_n_type_returns_raw_blob() { + // 5 = 5 * 10^0 + let raw = n_blob(0, &[0x05]); + let result = extract_alternator_scalar(&raw); + assert_eq!(result, Some(CqlValue::Blob(raw))); + } + + #[test] + fn extract_scalar_json_type_returns_raw_blob() { + let raw = json_s_blob("world"); + let result = extract_alternator_scalar(&raw); + assert_eq!(result, Some(CqlValue::Blob(raw))); + } + + #[test] + fn extract_scalar_empty_blob_returns_none() { + assert_eq!(extract_alternator_scalar(&[]), None); + } + + #[test] + fn extract_scalar_float32vector_tag_returns_none() { + // FLOAT32VECTOR (0x05) is not a scalar attribute type + let raw = alternator_vector_blob(&[1.0, 2.0]); + assert_eq!(extract_alternator_scalar(&raw), None); + } + + #[test] + fn extract_scalar_unknown_tag_returns_none() { + assert_eq!(extract_alternator_scalar(&[0x99, b'x']), None); + } + + // --- extract_alternator_scalars --- + + #[test] + fn extract_scalars_non_map_returns_empty() { + let result = extract_alternator_scalars(&CqlValue::Int(42), &["col".into()]); + assert!(result.is_empty()); + } + + #[test] + fn extract_scalars_s_type_included() { + let raw = s_blob("hello"); + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"col".to_vec()), + CqlValue::Blob(raw.clone()), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert_eq!(result.get(&"col".into()), Some(&Some(CqlValue::Blob(raw)))); + } + + #[test] + fn extract_scalars_n_type_included() { + let raw = n_blob(0, &[0x07]); + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"num".to_vec()), + CqlValue::Blob(raw.clone()), + )]); + let result = extract_alternator_scalars(&attrs, &["num".into()]); + assert_eq!(result.get(&"num".into()), Some(&Some(CqlValue::Blob(raw)))); + } + + #[test] + fn extract_scalars_json_type_included() { + let raw = json_s_blob("value"); + let attrs = CqlValue::Map(vec![( + CqlValue::Text("col".to_string()), + CqlValue::Blob(raw.clone()), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert_eq!(result.get(&"col".into()), Some(&Some(CqlValue::Blob(raw)))); + } + + #[test] + fn extract_scalars_float32vector_skipped() { + // FLOAT32VECTOR is not a scalar; extract_alternator_scalar returns None -> entry omitted + let raw = alternator_vector_blob(&[1.0, 2.0]); + let attrs = CqlValue::Map(vec![(CqlValue::Blob(b"vec".to_vec()), CqlValue::Blob(raw))]); + let result = extract_alternator_scalars(&attrs, &["vec".into()]); + assert!(result.is_empty()); + } + + #[test] + fn extract_scalars_unrequested_column_skipped() { + let raw = s_blob("hi"); + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"other".to_vec()), + CqlValue::Blob(raw), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert!(result.is_empty()); + } + + #[test] + fn extract_scalars_non_blob_value_skipped() { + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"col".to_vec()), + CqlValue::Text("not a blob".to_string()), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert!(result.is_empty()); + } + + #[test] + fn extract_scalars_non_utf8_blob_key_skipped() { + let raw = s_blob("hi"); + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(vec![0xFF, 0xFE]), // invalid UTF-8 + CqlValue::Blob(raw), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert!(result.is_empty()); + } + + #[test] + fn extract_scalars_text_key_works() { + let raw = s_blob("val"); + let attrs = CqlValue::Map(vec![( + CqlValue::Text("col".to_string()), + CqlValue::Blob(raw.clone()), + )]); + let result = extract_alternator_scalars(&attrs, &["col".into()]); + assert_eq!(result.get(&"col".into()), Some(&Some(CqlValue::Blob(raw)))); + } + + #[test] + fn extract_scalars_mixed_map() { + // S scalar, N scalar, vector (skipped), unrequested column (skipped), non-blob (skipped) + let s_raw = s_blob("hello"); + let n_raw = n_blob(0, &[0x03]); + let attrs = CqlValue::Map(vec![ + ( + CqlValue::Blob(b"s_col".to_vec()), + CqlValue::Blob(s_raw.clone()), + ), + ( + CqlValue::Blob(b"n_col".to_vec()), + CqlValue::Blob(n_raw.clone()), + ), + ( + CqlValue::Blob(b"vec".to_vec()), + CqlValue::Blob(alternator_vector_blob(&[1.0])), + ), + ( + CqlValue::Blob(b"other".to_vec()), + CqlValue::Blob(s_blob("ignored")), + ), + (CqlValue::Blob(b"non_blob".to_vec()), CqlValue::Int(99)), + ]); + let result = + extract_alternator_scalars(&attrs, &["s_col".into(), "n_col".into(), "vec".into()]); + assert_eq!(result.len(), 2); + assert_eq!( + result.get(&"s_col".into()), + Some(&Some(CqlValue::Blob(s_raw))) + ); + assert_eq!( + result.get(&"n_col".into()), + Some(&Some(CqlValue::Blob(n_raw))) + ); + assert!(!result.contains_key(&"other".into())); + assert!(!result.contains_key(&"non_blob".into())); + // vec is requested but FLOAT32VECTOR is not a scalar -> absent + assert!(!result.contains_key(&"vec".into())); + } + + #[test] + fn extract_scalars_absent_column_not_included() { + // Absent columns must NOT appear in the result (callers add tombstones separately) + let attrs = CqlValue::Map(vec![( + CqlValue::Blob(b"present".to_vec()), + CqlValue::Blob(s_blob("hi")), + )]); + let result = extract_alternator_scalars(&attrs, &["present".into(), "absent".into()]); + assert!(result.contains_key(&"present".into())); + assert!(!result.contains_key(&"absent".into())); + } } diff --git a/crates/vector-store/tests/integration/db_basic.rs b/crates/vector-store/tests/integration/db_basic.rs index ecc6026e..63f3e761 100644 --- a/crates/vector-store/tests/integration/db_basic.rs +++ b/crates/vector-store/tests/integration/db_basic.rs @@ -56,8 +56,9 @@ pub(crate) fn scan_fn( .send(( DbEmbedding { primary_key, - embedding, + embedding: Some(embedding), timestamp, + column_values: Default::default(), }, Some(tx_in_progress.clone().into()), )) diff --git a/crates/vector-store/tests/integration/https.rs b/crates/vector-store/tests/integration/https.rs index 7baaa46d..3b537178 100644 --- a/crates/vector-store/tests/integration/https.rs +++ b/crates/vector-store/tests/integration/https.rs @@ -106,6 +106,7 @@ async fn test_https_server_responds() { vector: vec![1.0].into(), filter: None, limit: NonZeroUsize::new(1).unwrap().into(), + return_columns: vec![], }) .send() .await @@ -118,6 +119,7 @@ async fn test_https_server_responds() { vector: vec![1.0].into(), filter: None, limit: NonZeroUsize::new(1).unwrap().into(), + return_columns: vec![], }) .send() .await diff --git a/crates/vector-store/tests/integration/routing.rs b/crates/vector-store/tests/integration/routing.rs index f2e184ed..33311c8e 100644 --- a/crates/vector-store/tests/integration/routing.rs +++ b/crates/vector-store/tests/integration/routing.rs @@ -522,16 +522,13 @@ async fn ann_routes_to_local_index_with_filter_columns_covering_restriction() { allow_filtering: true, }; - // TODO: update this assertion to expect StatusCode::OK once filtering on - // non-primary-key columns is supported end-to-end. Currently the request - // is routed correctly but fails at the filter validation layer. let response = assert_ann_served_by( &client, &covering, post_ann_with_filter(&client, &non_covering, filter), ) .await; - assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(response.status(), StatusCode::OK); } #[tokio::test] @@ -597,16 +594,13 @@ async fn ann_routes_to_global_index_with_filter_columns_covering_restriction() { allow_filtering: true, }; - // TODO: update this assertion to expect StatusCode::OK once filtering on - // non-primary-key columns is supported end-to-end. Currently the request - // is routed correctly but fails at the filter validation layer. let response = assert_ann_served_by( &client, &covering, post_ann_with_filter(&client, &non_covering, filter), ) .await; - assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(response.status(), StatusCode::OK); } #[tokio::test]