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") + } +}