From 4be191173a4780b8fb80c516793b47af9d9cc011 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Sun, 1 Nov 2026 11:57:44 -0800 Subject: [PATCH] vector: keep ListVectors pages and segments flat as an index grows; read only keys in GetVectors (#12300) * vector: scan only the fragments a ListVectors page covers A page pushed its offset into the scan, but the scan still loaded every fragment first, reading each one's deletion file: about 105 S3 GETs per page at 50k rows against 25 at 10k. The page's fragments now come from the live row counts in the manifest, so a page reads only its own. Tokens and row order are unchanged; a manifest without the counts keeps the old scan. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * s3vectors: list a ListVectors segment in the engine The gateway paged through the whole index for every segment and kept one row in segmentCount, so a segment cost segmentCount engine pages and 16 segments listed slower than one. The engine now takes segment_count and segment_index: it reads keys across the stride, keeps the segment's rows and takes their data and metadata alone, one call per page. A segment is still the rows whose position modulo segmentCount is segmentIndex, and its token still the position to resume from. An engine that predates the fields does not set segmented, and the gateway then pages through it as before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: read only the key column across the table in GetVectors The key lookup scanned every column. On S3, lance reads a column during the scan when a row of it is under 1000 bytes, so a 128-dimension vector was read for every row of the table to return 100 of them, and data and metadata were read even when the request did not ask for them. The scan now projects the requested columns and materializes everything but the key late, for the matching rows only. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * s3vectors: start the legacy segment fallback from the page already read With an engine that ignores segment_count, the gateway dropped the page it had just read and fetched it again. The fallback now consumes that page first, so a one-row segment page still costs one engine call. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: check metadata on segmented ListVectors pages The segment test listed data only. Its rows now carry metadata on two in three keys, and every segmented row, read through take_rows, must equal the same row in a plain listing. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: lease the version while a segment takes its data A segmented page with data or metadata reads keys first and then takes the other columns from the same version, so maintenance could reclaim its files in between. That path now holds a lease like GetVectors. A keys-only page is one scan and stays unleased, as plain pages are. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: leave GetVectors materialization to lance Forcing every column late made a 100-key GetVectors with data slower at 10k x 128 in the perf cluster: 54.7 ms against 46.3 ms. The take issued a read per row, where lance's heuristic reads a narrow vector column in one pass per fragment. The projection keeps the gain on its own: with neither data nor metadata requested, the call went from 38.6 ms to 25.3 ms and from 150 to 89 S3 GETs. At 50k x 384 the vector is already late under the heuristic, and a full lookup is unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Opus 5.5 --- seaweed-vector/Cargo.lock | 1 + seaweed-vector/crates/core/src/main.rs | 2 +- seaweed-vector/crates/core/src/query.rs | 25 +- seaweed-vector/crates/server/Cargo.toml | 3 + seaweed-vector/crates/server/src/service.rs | 310 ++++++++++++++++++-- weed/pb/vector.proto | 6 + weed/pb/vector_pb/vector.pb.go | 46 ++- weed/s3api/s3vectors/handlers.go | 83 +++--- weed/s3api/s3vectors/handlers_test.go | 26 +- weed/s3api/s3vectors/list_bounds_test.go | 81 +++++ 10 files changed, 509 insertions(+), 74 deletions(-) diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index 1751228a87..df09ba67fe 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6496,6 +6496,7 @@ dependencies = [ "reqwest 0.12.28", "serde", "serde_json", + "tempfile", "tokio", "tokio-stream", "tonic", diff --git a/seaweed-vector/crates/core/src/main.rs b/seaweed-vector/crates/core/src/main.rs index 2c3d745848..cc293f3067 100644 --- a/seaweed-vector/crates/core/src/main.rs +++ b/seaweed-vector/crates/core/src/main.rs @@ -191,7 +191,7 @@ async fn main() -> Result<()> { let dataset = open_dataset(&client, &id, &fallback).await?; let vector = VectorTable::resolve(&dataset, vector_column.as_deref(), key_column.as_deref())?; - for batch in query::get(&dataset, &vector, &keys).await? { + for batch in query::get(&dataset, &vector, &keys, None).await? { println!("{}", arrow_cast::pretty::pretty_format_batches(&[batch])?); } } diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs index 63ffd72a9c..ea6cf15e80 100644 --- a/seaweed-vector/crates/core/src/query.rs +++ b/seaweed-vector/crates/core/src/query.rs @@ -63,11 +63,12 @@ pub async fn ann(dataset: &Dataset, table: &VectorTable, query: &Query) -> Resul run(scanner(dataset, table, query, true).await?, table).await } -/// Fetch rows by key. +/// Fetch rows by key, projected to `columns` (all of them when None). pub async fn get( dataset: &Dataset, table: &VectorTable, keys: &[String], + columns: Option<&[String]>, ) -> Result> { let Some(key_column) = &table.key_column else { bail!("the table has no key column to look up"); @@ -84,6 +85,9 @@ pub async fn get( scanner .filter(&format!("`{}` in ({list})", key_column.replace('`', "``")))? .with_row_id(); + if let Some(columns) = columns { + scanner.project(columns)?; + } scanner .try_into_stream() .await? @@ -425,11 +429,24 @@ mod tests { let dataset = write_dataset(&dir, &vectors, &["a", "b", "c"]).await; let table = VectorTable::resolve(&dataset, None, None).unwrap(); - let batches = get(&dataset, &table, &["k0".to_string(), "k2".to_string()]) - .await - .unwrap(); + let keys = ["k0".to_string(), "k2".to_string()]; + let batches = get(&dataset, &table, &keys, None).await.unwrap(); let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(rows, 2); + assert!(batches.iter().all(|b| b.column_by_name("vector").is_some())); + + let projected = get(&dataset, &table, &keys, Some(&["key".to_string()])) + .await + .unwrap(); + let mut got: Vec = Vec::new(); + for batch in &projected { + assert!(batch.column_by_name("vector").is_none()); + let column = batch.column_by_name("key").unwrap(); + let column = column.as_any().downcast_ref::().unwrap(); + got.extend(column.iter().flatten().map(String::from)); + } + got.sort(); + assert_eq!(got, keys); } #[tokio::test] diff --git a/seaweed-vector/crates/server/Cargo.toml b/seaweed-vector/crates/server/Cargo.toml index 2727b46f26..5f938f66a4 100644 --- a/seaweed-vector/crates/server/Cargo.toml +++ b/seaweed-vector/crates/server/Cargo.toml @@ -36,6 +36,9 @@ tokio-stream.workspace = true tracing.workspace = true tracing-subscriber.workspace = true +[dev-dependencies] +tempfile = "3" + [build-dependencies] tonic-prost-build = "0.14" protoc-bin-vendored = "3" diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index 4287130b6e..fb8ab22264 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -1,5 +1,7 @@ use anyhow::{Result, bail}; -use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, StringArray}; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, StringArray, UInt64Array, +}; use arrow_schema::{DataType, Field, Schema}; use futures::TryStreamExt; use lance::Dataset; @@ -8,6 +10,7 @@ use lance::dataset::WriteParams; use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; use lance::io::ObjectStoreParams; use lance::io::StorageOptionsAccessor; +use lance::table::format::Fragment; use lance_linalg::distance::DistanceType; use prost_types::{ListValue, Struct, Value, value::Kind}; use std::collections::{BTreeMap, HashMap}; @@ -52,6 +55,7 @@ fn open_status(err: anyhow::Error) -> Status { const KEY_COLUMN: &str = "key"; const METADATA_COLUMN: &str = "metadata"; const VECTOR_COLUMN: &str = "vector"; +const ROW_ID_COLUMN: &str = "_rowid"; fn schema_for( dimension: usize, @@ -213,6 +217,114 @@ fn key_list(keys: &[String]) -> String { ) } +/// The columns a GetVectors or ListVectors response carries. +fn projection(return_data: bool, return_metadata: bool) -> Vec { + let mut columns = vec![KEY_COLUMN.to_string()]; + if return_data { + columns.push(VECTOR_COLUMN.to_string()); + } + if return_metadata { + columns.push(METADATA_COLUMN.to_string()); + } + columns +} + +/// The fragments holding live rows `offset..offset + rows` in scan order, and +/// where `offset` falls in the first of them. A scan reads every listed +/// fragment's deletion file, so a page names only its own. None when the +/// manifest lacks a fragment's live row count. +fn page_fragments(dataset: &Dataset, offset: u64, rows: u64) -> Option<(Vec, u64)> { + dataset.manifest().writer_version.as_ref()?; + let (mut start, mut skip, mut out) = (0u64, 0u64, Vec::new()); + for fragment in dataset.fragments().iter() { + if start >= offset.saturating_add(rows) { + break; + } + let live = fragment.num_rows()? as u64; + if start + live > offset { + if out.is_empty() { + skip = offset - start; + } + out.push(fragment.clone()); + } + start += live; + } + Some((out, skip)) +} + +/// Live rows `offset..offset + rows` in scan order, projected to `columns`. +async fn list_page( + dataset: &Dataset, + columns: &[String], + offset: u64, + rows: u64, +) -> Result> { + let mut scanner = dataset.scan(); + scanner.project(columns)?; + let skip = match page_fragments(dataset, offset, rows) { + Some((fragments, skip)) => { + scanner.with_fragments(fragments); + skip + } + None => offset, + }; + scanner.limit(Some(rows as i64), Some(skip as i64))?; + Ok(scanner.try_into_stream().await?.try_collect().await?) +} + +/// The AWS ListVectors segmentCount limit. +const MAX_SEGMENTS: u64 = 16; + +/// Up to `page` rows whose position modulo `count` is `index`, from position +/// `offset`, and the position to resume from. Keys are read across the +/// stride; the other columns only for the rows kept. +async fn list_segment( + dataset: &Dataset, + columns: &[String], + offset: u64, + page: u64, + count: u64, + index: u64, +) -> Result<(Vec, Option)> { + let first = offset.saturating_add((index + count - offset % count) % count); + let stride = vec![KEY_COLUMN.to_string(), ROW_ID_COLUMN.to_string()]; + let batches = list_page(dataset, &stride, first, page * count).await?; + let (mut keys, mut row_ids, mut scanned) = (Vec::new(), Vec::new(), 0u64); + for batch in &batches { + let (Some(batch_keys), Some(ids)) = ( + batch + .column_by_name(KEY_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()), + batch + .column_by_name(ROW_ID_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()), + ) else { + bail!("the scan did not return keys and row ids"); + }; + for i in 0..batch.num_rows() { + if (scanned + i as u64) % count == 0 { + keys.push(batch_keys.value(i).to_string()); + row_ids.push(ids.value(i)); + } + } + scanned += batch.num_rows() as u64; + } + let next = (row_ids.len() as u64 == page && scanned > (page - 1) * count + 1) + .then(|| first.saturating_add((page - 1) * count + 1)); + if row_ids.is_empty() { + return Ok((Vec::new(), next)); + } + if columns.len() == 1 { + let keys: ArrayRef = Arc::new(StringArray::from(keys)); + return Ok(( + vec![RecordBatch::try_from_iter([(KEY_COLUMN, keys)])?], + next, + )); + } + let projection = ProjectionRequest::from_columns(columns, dataset.schema()); + Ok((vec![dataset.take_rows(&row_ids, projection).await?], next)) +} + fn rows_as_vectors( batches: &[RecordBatch], return_data: bool, @@ -400,7 +512,8 @@ impl VectorService for VectorServer { .await .map_err(internal)?; let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; - let batches = query::get(&dataset, &resolved, &req.keys) + let columns = projection(req.return_data, req.return_metadata); + let batches = query::get(&dataset, &resolved, &req.keys, Some(&columns)) .await .map_err(internal)?; Ok(Response::new(pb::GetVectorsResponse { @@ -436,6 +549,12 @@ impl VectorService for VectorServer { let table = req .table .ok_or_else(|| Status::invalid_argument("missing table"))?; + let segments = req.segment_count as u64; + if segments > MAX_SEGMENTS || (segments > 0 && req.segment_index as u64 >= segments) { + return Err(Status::invalid_argument(format!( + "segment_index must be below segment_count, which is at most {MAX_SEGMENTS}" + ))); + } let dataset = self.open(&table).await.map_err(open_status)?; let offset: u64 = if req.next_token.is_empty() { 0 @@ -444,13 +563,7 @@ impl VectorService for VectorServer { .parse() .map_err(|_| Status::invalid_argument("invalid next_token"))? }; - let mut columns = vec![KEY_COLUMN.to_string()]; - if req.return_data { - columns.push(VECTOR_COLUMN.to_string()); - } - if req.return_metadata { - columns.push(METADATA_COLUMN.to_string()); - } + let columns = projection(req.return_data, req.return_metadata); // An unset max_results arrives as 0; page it like S3 Vectors does. const DEFAULT_PAGE: u32 = 500; const MAX_PAGE: u32 = 1000; @@ -458,17 +571,33 @@ impl VectorService for VectorServer { 0 => DEFAULT_PAGE, n => n.min(MAX_PAGE), }; - let mut scanner = dataset.scan(); - scanner - .project(&columns) - .map_err(internal)? - .limit(Some(page as i64), Some(offset as i64)) - .map_err(internal)?; - let batches: Vec = scanner - .try_into_stream() + if segments > 1 { + // Data and metadata come from a second read of this version, so + // lease it as get_vectors does. + let _lease = if columns.len() > 1 { + weed_vector_core::maintenance::try_pin(&dataset) + .await + .map_err(internal)? + } else { + None + }; + let (vectors, next) = list_segment( + &dataset, + &columns, + offset, + page as u64, + segments, + req.segment_index as u64, + ) .await - .map_err(internal)? - .try_collect() + .map_err(internal)?; + return Ok(Response::new(pb::ListVectorsResponse { + vectors: rows_as_vectors(&vectors, req.return_data, req.return_metadata), + next_token: next.map(|n| n.to_string()).unwrap_or_default(), + segmented: true, + })); + } + let batches = list_page(&dataset, &columns, offset, page as u64) .await .map_err(internal)?; let rows = batches.iter().map(|b| b.num_rows()).sum::(); @@ -481,6 +610,7 @@ impl VectorService for VectorServer { Ok(Response::new(pb::ListVectorsResponse { vectors: rows_as_vectors(&batches, req.return_data, req.return_metadata), next_token, + segmented: segments > 0, })) } @@ -651,6 +781,148 @@ mod tests { assert!(batch_of(&vectors, 4, schema).is_err()); } + fn vectors(keys: std::ops::Range) -> Vec { + keys.map(|i| pb::Vector { + key: format!("k{i:03}"), + data: vec![i as f32, 1.0], + metadata: (i % 3 != 0) + .then(|| parse_struct(&format!(r#"{{"n":{i},"tag":"t{}"}}"#, i % 2)).unwrap()), + }) + .collect() + } + + /// One fragment per range, then the given keys deleted. + async fn fragmented( + dir: &tempfile::TempDir, + ranges: &[std::ops::Range], + deleted: &[String], + ) -> Dataset { + let uri = dir.path().to_str().unwrap(); + let schema = schema_for(2, pb::DistanceMetric::L2, &[]); + let mut dataset: Option = None; + for range in ranges { + let batch = batch_of(&vectors(range.clone()), 2, schema.clone()).unwrap(); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + dataset = Some(match dataset { + None => Dataset::write(reader, uri, None).await.unwrap(), + Some(mut d) => { + d.append(reader, None).await.unwrap(); + d + } + }); + } + let mut dataset = dataset.unwrap(); + if !deleted.is_empty() { + dataset.delete(&key_list(deleted)).await.unwrap(); + } + dataset + } + + fn keys_of(batches: &[RecordBatch]) -> Vec { + rows_as_vectors(batches, false, false) + .into_iter() + .map(|v| v.key) + .collect() + } + + /// Every page names only its fragments yet matches the same window of + /// the full scan, across deletions and a fully deleted fragment. + #[tokio::test] + async fn list_page_matches_a_full_scan() { + let dir = tempfile::tempdir().unwrap(); + let deleted: Vec = ["k001", "k004", "k005", "k006", "k007", "k010"] + .map(String::from) + .to_vec(); + let dataset = fragmented(&dir, &[0..4, 4..8, 8..12, 12..13], &deleted).await; + let columns = vec![KEY_COLUMN.to_string()]; + let all = keys_of( + &dataset + .scan() + .project(&columns) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(), + ); + assert_eq!(all.len(), 7); + let total = all.len() as u64; + assert!(page_fragments(&dataset, 0, total).is_some()); + for offset in 0..=total + 1 { + for rows in [1, 2, 3, 5, total + 2] { + let got = keys_of(&list_page(&dataset, &columns, offset, rows).await.unwrap()); + let end = (offset + rows).min(total) as usize; + let want = &all[(offset.min(total) as usize)..end]; + assert_eq!(got, want, "offset {offset} rows {rows}"); + } + } + } + + /// Following a segment's tokens yields exactly the rows at its positions, + /// in order, and the segments together cover every row once. + #[tokio::test] + async fn list_segment_strides_the_listing() { + let dir = tempfile::tempdir().unwrap(); + let deleted: Vec = ["k002", "k009"].map(String::from).to_vec(); + let dataset = fragmented(&dir, &[0..5, 5..12, 12..13, 13..20], &deleted).await; + let all = keys_of( + &list_page(&dataset, &[KEY_COLUMN.to_string()], 0, 100) + .await + .unwrap(), + ); + let with_data = projection(true, true); + let rows = rows_as_vectors( + &list_page(&dataset, &with_data, 0, 100).await.unwrap(), + true, + true, + ); + assert!(rows.iter().any(|v| v.metadata.is_some())); + for count in [2u64, 3, 4, 16] { + for page in [1u64, 2, 3, 7] { + let mut seen = Vec::new(); + for index in 0..count { + let (mut got, mut offset) = (Vec::new(), Some(0)); + while let Some(from) = offset { + let (batches, next) = + list_segment(&dataset, &with_data, from, page, count, index) + .await + .unwrap(); + for v in rows_as_vectors(&batches, true, true) { + let want = rows.iter().find(|r| r.key == v.key).unwrap(); + assert_eq!(&v, want); + got.push(v.key); + } + assert!(next.is_none() || next > Some(from)); + offset = next; + } + let want: Vec = all + .iter() + .skip(index as usize) + .step_by(count as usize) + .cloned() + .collect(); + assert_eq!(got, want, "count {count} index {index} page {page}"); + seen.extend(got); + } + seen.sort(); + assert_eq!(seen.len(), all.len()); + let mut sorted = all.clone(); + sorted.sort(); + assert_eq!(seen, sorted); + } + } + // Keys alone come from the stride scan, and a page ending on the + // last row carries no token. + let (batches, next) = list_segment(&dataset, &[KEY_COLUMN.to_string()], 0, 9, 2, 1) + .await + .unwrap(); + let odd: Vec = all.iter().skip(1).step_by(2).cloned().collect(); + assert_eq!(keys_of(&batches), odd); + assert_eq!(next, None); + } + /// A catalog that answers one describe with `location`. async fn catalog_describing(location: &str) -> String { use tokio::io::{AsyncReadExt, AsyncWriteExt}; diff --git a/weed/pb/vector.proto b/weed/pb/vector.proto index 723aefed90..8bdaf3e086 100644 --- a/weed/pb/vector.proto +++ b/weed/pb/vector.proto @@ -93,11 +93,17 @@ message ListVectorsRequest { string next_token = 3; bool return_data = 4; bool return_metadata = 5; + // A nonzero segment_count lists only the rows whose position modulo + // segment_count is segment_index; next_token stays a row position. + uint32 segment_count = 6; + uint32 segment_index = 7; } message ListVectorsResponse { repeated Vector vectors = 1; string next_token = 2; + // Set when segment_count was applied; an older engine ignores it. + bool segmented = 3; } message QueryVectorsRequest { diff --git a/weed/pb/vector_pb/vector.pb.go b/weed/pb/vector_pb/vector.pb.go index 83b69c70d0..330a3484d2 100644 --- a/weed/pb/vector_pb/vector.pb.go +++ b/weed/pb/vector_pb/vector.pb.go @@ -693,8 +693,12 @@ type ListVectorsRequest struct { NextToken string `protobuf:"bytes,3,opt,name=next_token,json=nextToken,proto3" json:"next_token,omitempty"` ReturnData bool `protobuf:"varint,4,opt,name=return_data,json=returnData,proto3" json:"return_data,omitempty"` ReturnMetadata bool `protobuf:"varint,5,opt,name=return_metadata,json=returnMetadata,proto3" json:"return_metadata,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // A nonzero segment_count lists only the rows whose position modulo + // segment_count is segment_index; next_token stays a row position. + SegmentCount uint32 `protobuf:"varint,6,opt,name=segment_count,json=segmentCount,proto3" json:"segment_count,omitempty"` + SegmentIndex uint32 `protobuf:"varint,7,opt,name=segment_index,json=segmentIndex,proto3" json:"segment_index,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *ListVectorsRequest) Reset() { @@ -762,10 +766,26 @@ func (x *ListVectorsRequest) GetReturnMetadata() bool { return false } +func (x *ListVectorsRequest) GetSegmentCount() uint32 { + if x != nil { + return x.SegmentCount + } + return 0 +} + +func (x *ListVectorsRequest) GetSegmentIndex() uint32 { + if x != nil { + return x.SegmentIndex + } + return 0 +} + type ListVectorsResponse struct { - state protoimpl.MessageState `protogen:"open.v1"` - Vectors []*Vector `protobuf:"bytes,1,rep,name=vectors,proto3" json:"vectors,omitempty"` - NextToken string `protobuf:"bytes,2,opt,name=next_token,json=nextToken,proto3" json:"next_token,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Vectors []*Vector `protobuf:"bytes,1,rep,name=vectors,proto3" json:"vectors,omitempty"` + NextToken string `protobuf:"bytes,2,opt,name=next_token,json=nextToken,proto3" json:"next_token,omitempty"` + // Set when segment_count was applied; an older engine ignores it. + Segmented bool `protobuf:"varint,3,opt,name=segmented,proto3" json:"segmented,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -814,6 +834,13 @@ func (x *ListVectorsResponse) GetNextToken() string { return "" } +func (x *ListVectorsResponse) GetSegmented() bool { + if x != nil { + return x.Segmented + } + return false +} + type QueryVectorsRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Table *Table `protobuf:"bytes,1,opt,name=table,proto3" json:"table,omitempty"` @@ -1069,7 +1096,7 @@ const file_vector_proto_rawDesc = "" + "\x05table\x18\x01 \x01(\v2\x10.vector_pb.TableR\x05table\x12\x12\n" + "\x04keys\x18\x02 \x03(\tR\x04keys\"1\n" + "\x15DeleteVectorsResponse\x12\x18\n" + - "\aversion\x18\x01 \x01(\x04R\aversion\"\xc6\x01\n" + + "\aversion\x18\x01 \x01(\x04R\aversion\"\x90\x02\n" + "\x12ListVectorsRequest\x12&\n" + "\x05table\x18\x01 \x01(\v2\x10.vector_pb.TableR\x05table\x12\x1f\n" + "\vmax_results\x18\x02 \x01(\rR\n" + @@ -1078,11 +1105,14 @@ const file_vector_proto_rawDesc = "" + "next_token\x18\x03 \x01(\tR\tnextToken\x12\x1f\n" + "\vreturn_data\x18\x04 \x01(\bR\n" + "returnData\x12'\n" + - "\x0freturn_metadata\x18\x05 \x01(\bR\x0ereturnMetadata\"a\n" + + "\x0freturn_metadata\x18\x05 \x01(\bR\x0ereturnMetadata\x12#\n" + + "\rsegment_count\x18\x06 \x01(\rR\fsegmentCount\x12#\n" + + "\rsegment_index\x18\a \x01(\rR\fsegmentIndex\"\x7f\n" + "\x13ListVectorsResponse\x12+\n" + "\avectors\x18\x01 \x03(\v2\x11.vector_pb.VectorR\avectors\x12\x1d\n" + "\n" + - "next_token\x18\x02 \x01(\tR\tnextToken\"\xa0\x02\n" + + "next_token\x18\x02 \x01(\tR\tnextToken\x12\x1c\n" + + "\tsegmented\x18\x03 \x01(\bR\tsegmented\"\xa0\x02\n" + "\x13QueryVectorsRequest\x12&\n" + "\x05table\x18\x01 \x01(\v2\x10.vector_pb.TableR\x05table\x12!\n" + "\fquery_vector\x18\x02 \x03(\x02R\vqueryVector\x12\x13\n" + diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index 983de36233..b02b17b8f3 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -1382,21 +1382,38 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { maxResults = *req.MaxResults } // AWS exposes a static split: segmentIndex/segmentCount divides the full - // listing. The engine service does not know AWS segments, so page through - // the index here and keep only the caller's share; a call scans about - // segmentCount × maxResults rows. The engine's next_token is a row offset, - // so the caller's token is the offset to resume scanning from. - if req.SegmentCount != nil { - var offset uint64 - if req.NextToken != "" { - var err error - if offset, err = strconv.ParseUint(req.NextToken, 10, 64); err != nil { - writeError(w, http.StatusBadRequest, "ValidationException", "invalid nextToken") - return - } + // listing. A segment is the rows whose position modulo segmentCount is + // segmentIndex, and its token is the row position to resume from. + var offset uint64 + if req.SegmentCount != nil && req.NextToken != "" { + var err error + if offset, err = strconv.ParseUint(req.NextToken, 10, 64); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", "invalid nextToken") + return } - out, nextToken, err := s.listSegmentVectors(r.Context(), tableID(arns.bucket, arns.index), - offset, maxResults, uint64(*req.SegmentCount), uint64(*req.SegmentIndex), req.ReturnData, req.ReturnMetadata) + } + listReq := &vector_pb.ListVectorsRequest{ + Table: tableID(arns.bucket, arns.index), + MaxResults: uint32(maxResults), + NextToken: req.NextToken, + ReturnData: req.ReturnData, + ReturnMetadata: req.ReturnMetadata, + } + if req.SegmentCount != nil { + listReq.SegmentCount = uint32(*req.SegmentCount) + listReq.SegmentIndex = uint32(*req.SegmentIndex) + } + resp, err := s.vectors.ListVectors(r.Context(), listReq) + if err != nil { + vectorErr(w, err) + return + } + // An engine that predates segment_count lists every row; page through + // it here, starting with the page already read, and keep only the + // caller's share. + if req.SegmentCount != nil && !resp.Segmented { + out, nextToken, err := s.listSegmentVectors(r.Context(), listReq.Table, + offset, maxResults, uint64(*req.SegmentCount), uint64(*req.SegmentIndex), req.ReturnData, req.ReturnMetadata, resp) if err != nil { vectorErr(w, err) return @@ -1407,17 +1424,6 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { }) return } - resp, err := s.vectors.ListVectors(r.Context(), &vector_pb.ListVectorsRequest{ - Table: tableID(arns.bucket, arns.index), - MaxResults: uint32(maxResults), - NextToken: req.NextToken, - ReturnData: req.ReturnData, - ReturnMetadata: req.ReturnMetadata, - }) - if err != nil { - vectorErr(w, err) - return - } writeJSON(w, http.StatusOK, &ListVectorsResponse{ NextToken: resp.NextToken, Vectors: toOutputVectors(resp.Vectors, req.ReturnData, req.ReturnMetadata), @@ -1461,8 +1467,8 @@ func listVectorsBounds(req *ListVectorsRequest) error { // listSegmentVectors scans the index from offset, keeping the rows that fall // in the caller's segment until limit are collected or the index runs out. // The engine's next_token is a row offset, so the returned token is the offset -// scanning resumes from. -func (s *Server) listSegmentVectors(ctx context.Context, table *vector_pb.Table, offset uint64, maxResults int, segmentCount, segmentIndex uint64, returnData, returnMetadata bool) ([]*vector_pb.Vector, string, error) { +// scanning resumes from. A non-nil first is the engine page at offset. +func (s *Server) listSegmentVectors(ctx context.Context, table *vector_pb.Table, offset uint64, maxResults int, segmentCount, segmentIndex uint64, returnData, returnMetadata bool, first *vector_pb.ListVectorsResponse) ([]*vector_pb.Vector, string, error) { limit := maxResults if limit <= 0 { limit = vectorsDefaultPage @@ -1473,15 +1479,20 @@ func (s *Server) listSegmentVectors(ctx context.Context, table *vector_pb.Table, nextToken := "" token := strconv.FormatUint(offset, 10) for { - page, err := s.vectors.ListVectors(ctx, &vector_pb.ListVectorsRequest{ - Table: table, - MaxResults: vectorsMaxPage, - NextToken: token, - ReturnData: returnData, - ReturnMetadata: returnMetadata, - }) - if err != nil { - return nil, "", err + page := first + first = nil + if page == nil { + var err error + page, err = s.vectors.ListVectors(ctx, &vector_pb.ListVectorsRequest{ + Table: table, + MaxResults: vectorsMaxPage, + NextToken: token, + ReturnData: returnData, + ReturnMetadata: returnMetadata, + }) + if err != nil { + return nil, "", err + } } base := offset offset += uint64(len(page.Vectors)) diff --git a/weed/s3api/s3vectors/handlers_test.go b/weed/s3api/s3vectors/handlers_test.go index 85bd4fd5c7..d7ceccf683 100644 --- a/weed/s3api/s3vectors/handlers_test.go +++ b/weed/s3api/s3vectors/handlers_test.go @@ -24,6 +24,8 @@ import ( // and records the mutating calls a handler test asserts on. type fakeVectors struct { vector_pb.VectorServiceClient + // legacy ignores segment_count, like an engine that predates it. + legacy bool rows []*vector_pb.Vector created []*vector_pb.CreateIndexRequest puts []*vector_pb.PutVectorsRequest @@ -58,6 +60,18 @@ func (f *fakeVectors) ListVectors(ctx context.Context, req *vector_pb.ListVector if offset > len(f.rows) { offset = len(f.rows) } + if req.SegmentCount > 0 && !f.legacy { + count, index := int(req.SegmentCount), int(req.SegmentIndex) + resp := &vector_pb.ListVectorsResponse{Segmented: true} + pos := offset + (index+count-offset%count)%count + for ; pos < len(f.rows) && len(resp.Vectors) < page; pos += count { + resp.Vectors = append(resp.Vectors, f.rows[pos]) + } + if last := pos - count; len(resp.Vectors) == page && last+1 < len(f.rows) { + resp.NextToken = strconv.Itoa(last + 1) + } + return resp, nil + } end := offset + page if end > len(f.rows) { end = len(f.rows) @@ -103,7 +117,7 @@ func TestListSegmentVectorsCoversEveryRowOnce(t *testing.T) { } else if err != nil { t.Fatalf("bad token %q: %v", token, err) } - out, next, err := s.listSegmentVectors(context.Background(), nil, offset, 7, 3, segmentIndex, false, false) + out, next, err := s.listSegmentVectors(context.Background(), nil, offset, 7, 3, segmentIndex, false, false, nil) if err != nil { t.Fatal(err) } @@ -131,7 +145,7 @@ func TestListSegmentVectorsCoversEveryRowOnce(t *testing.T) { func TestListSegmentVectorsStopsInsideEnginePage(t *testing.T) { f := newFakeVectors(100) s := &Server{vectors: f} - out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 5, 2, 0, false, false) + out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 5, 2, 0, false, false, nil) if err != nil { t.Fatal(err) } @@ -143,7 +157,7 @@ func TestListSegmentVectorsStopsInsideEnginePage(t *testing.T) { t.Fatalf("token %q, want 9", next) } offset, _ := strconv.ParseUint(next, 10, 64) - out, next, err = s.listSegmentVectors(context.Background(), nil, offset, 5, 2, 0, false, false) + out, next, err = s.listSegmentVectors(context.Background(), nil, offset, 5, 2, 0, false, false, nil) if err != nil { t.Fatal(err) } @@ -159,7 +173,7 @@ func TestListSegmentVectorsStopsInsideEnginePage(t *testing.T) { func TestListSegmentVectorsEndsOnLastRow(t *testing.T) { f := newFakeVectors(10) s := &Server{vectors: f} - out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 10, 1, 0, false, false) + out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 10, 1, 0, false, false, nil) if err != nil { t.Fatal(err) } @@ -175,14 +189,14 @@ func TestListSegmentVectorsEndsOnLastRow(t *testing.T) { func TestListSegmentVectorsPageLimits(t *testing.T) { f := newFakeVectors(1500) s := &Server{vectors: f} - out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 0, 1, 0, false, false) + out, next, err := s.listSegmentVectors(context.Background(), nil, 0, 0, 1, 0, false, false, nil) if err != nil { t.Fatal(err) } if len(out) != vectorsDefaultPage || next != strconv.Itoa(vectorsDefaultPage) { t.Fatalf("default page: %d rows token %q", len(out), next) } - out, next, err = s.listSegmentVectors(context.Background(), nil, 0, 5000, 1, 0, false, false) + out, next, err = s.listSegmentVectors(context.Background(), nil, 0, 5000, 1, 0, false, false, nil) if err != nil { t.Fatal(err) } diff --git a/weed/s3api/s3vectors/list_bounds_test.go b/weed/s3api/s3vectors/list_bounds_test.go index cdd9e8e3b3..d94907a3cb 100644 --- a/weed/s3api/s3vectors/list_bounds_test.go +++ b/weed/s3api/s3vectors/list_bounds_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "net/http" + "strconv" "strings" "testing" @@ -99,3 +100,83 @@ func TestListVectorsBoundsAccepted(t *testing.T) { } } } + +// A segmented listing comes from the engine in one call per page; an engine +// that ignores segment_count falls back to paging here. Both cover every row +// once, in the same segments. +func TestListVectorsSegmentsFromTheEngine(t *testing.T) { + fs := s3tablestest.Start(t) + modern := &countingVectors{fakeVectors: newFakeVectors(2503)} + legacy := &countingVectors{fakeVectors: newFakeVectors(2503)} + legacy.legacy = true + segments := map[bool][]string{} + for _, engine := range []*countingVectors{modern, legacy} { + s := NewServer(s3tables.NewManagerClient(fs.Client), allowAllAuth{}, engine) + bucket := "vec-modern" + if engine.legacy { + bucket = "vec-legacy" + } + if w := postJSON(t, s.handleCreateVectorBucket, &CreateVectorBucketRequest{VectorBucketName: bucket}); w.Code != http.StatusOK { + t.Fatalf("create bucket: %d %s", w.Code, w.Body.String()) + } + if w := postJSON(t, s.handleCreateIndex, &CreateIndexRequest{ + VectorBucketName: bucket, IndexName: "idx", DataType: "float32", Dimension: 3, DistanceMetric: "euclidean", + }); w.Code != http.StatusOK { + t.Fatalf("create index: %d %s", w.Code, w.Body.String()) + } + engine.lists = 0 + seen, pages := map[string]int{}, 0 + for index := 0; index < 3; index++ { + token := "" + for { + fields := `,"segmentCount":3,"segmentIndex":` + strconv.Itoa(index) + `,"maxResults":100` + if token != "" { + fields += `,"nextToken":"` + token + `"` + } + body := json.RawMessage(`{"vectorBucketName":"` + bucket + `","indexName":"idx"` + fields + `}`) + w := postJSON(t, s.handleListVectors, body) + if w.Code != http.StatusOK { + t.Fatalf("legacy=%v: %d %s", engine.legacy, w.Code, w.Body.String()) + } + var out ListVectorsResponse + if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil { + t.Fatal(err) + } + pages++ + for _, v := range out.Vectors { + seen[v.Key]++ + segments[engine.legacy] = append(segments[engine.legacy], strconv.Itoa(index)+"/"+v.Key) + } + if out.NextToken == "" { + break + } + token = out.NextToken + } + } + if len(seen) != len(engine.rows) { + t.Fatalf("legacy=%v: saw %d distinct rows, want %d", engine.legacy, len(seen), len(engine.rows)) + } + for key, n := range seen { + if n != 1 { + t.Fatalf("legacy=%v: row %s seen %d times", engine.legacy, key, n) + } + } + if !engine.legacy && engine.lists != pages { + t.Fatalf("modern engine took %d calls for %d pages", engine.lists, pages) + } + // The fallback starts from the page the engine already returned. + engine.lists = 0 + body := json.RawMessage(`{"vectorBucketName":"` + bucket + `","indexName":"idx","segmentCount":2,"segmentIndex":0,"maxResults":1}`) + w := postJSON(t, s.handleListVectors, body) + var out ListVectorsResponse + if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil || len(out.Vectors) != 1 || out.Vectors[0].Key != "k0" || out.NextToken != "1" { + t.Fatalf("legacy=%v: one-row segment page %d %s", engine.legacy, w.Code, w.Body.String()) + } + if engine.lists != 1 { + t.Fatalf("legacy=%v: one-row segment page took %d engine calls", engine.legacy, engine.lists) + } + } + if strings.Join(segments[false], ",") != strings.Join(segments[true], ",") { + t.Fatal("legacy fallback and engine segments disagree") + } +}