leaves12138 commented on code in PR #62:
URL:
https://github.com/apache/paimon-vector-index/pull/62#discussion_r3650843840
##########
core/src/ivfflat_io.rs:
##########
@@ -229,79 +378,175 @@ impl<R: SeekRead> IVFFlatIndexReader<R> {
return Ok(());
}
- let mut cursor = PreadCursor::new(&mut self.reader,
IVFFLAT_HEADER_SIZE as u64);
- self.quantizer_centroids =
- read_f32_vec(&mut cursor, checked_section_size(self.nlist,
self.d)?)?;
+ let centroid_count = checked_section_size(self.nlist, self.d)?;
+ let centroid_bytes = centroid_count.checked_mul(4).ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "IVF-FLAT centroid bytes overflow",
+ )
+ })?;
+ let table_bytes = self.nlist.checked_mul(16).ok_or_else(|| {
+ io::Error::new(io::ErrorKind::InvalidData, "IVF-FLAT offset table
overflow")
+ })?;
+ let mut metadata = vec![
+ 0u8;
+ centroid_bytes.checked_add(table_bytes).ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "IVF-FLAT metadata size overflow",
+ )
+ })?
+ ];
+ self.reader
+ .pread(&mut [ReadRequest::new(IVFFLAT_HEADER_SIZE as u64, &mut
metadata)])?;
+ self.quantizer_centroids =
bytes_to_f32_vec(&metadata[..centroid_bytes])?;
self.list_offsets = vec![0; self.nlist];
self.list_counts = vec![0; self.nlist];
self.list_id_bytes_lens = vec![0; self.nlist];
+ let mut actual_total = 0i64;
for list_id in 0..self.nlist {
- self.list_offsets[list_id] = read_i64_le(&mut cursor)?;
- let count = read_i32_le(&mut cursor)?;
+ let base = centroid_bytes + list_id * 16;
+ self.list_offsets[list_id] =
+ i64::from_le_bytes(metadata[base..base +
8].try_into().unwrap());
+ let count = i32::from_le_bytes(metadata[base + 8..base +
12].try_into().unwrap());
if count < 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("negative list count {} at list {}", count,
list_id),
));
}
self.list_counts[list_id] = count;
- let id_bytes_len = read_i32_le(&mut cursor)?;
+ actual_total = actual_total.checked_add(count as
i64).ok_or_else(|| {
+ io::Error::new(io::ErrorKind::InvalidData, "IVF-FLAT vector
count overflow")
+ })?;
+ let id_bytes_len =
+ i32::from_le_bytes(metadata[base + 12..base +
16].try_into().unwrap());
if id_bytes_len < 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("negative id_bytes_len {} at list {}",
id_bytes_len, list_id),
));
}
+ if count > 0 && id_bytes_len == 0 {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("missing delta ID bytes for non-empty IVF-FLAT
list {list_id}"),
+ ));
+ }
self.list_id_bytes_lens[list_id] = id_bytes_len;
}
+ if actual_total != self.total_vectors {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!(
+ "IVF-FLAT header vector count {} does not match list total
{actual_total}",
+ self.total_vectors
+ ),
+ ));
+ }
self.loaded = true;
Ok(())
}
pub fn read_inverted_list(&mut self, list_id: usize) ->
io::Result<(Vec<i64>, Vec<f32>)> {
+ let mut lists = self.read_inverted_lists(&[list_id])?;
+ let list = lists.pop().expect("one requested list has one result");
+ let vectors = list.vectors().to_vec();
+ Ok((list.ids, vectors))
+ }
+
+ fn read_inverted_lists(&mut self, list_ids: &[usize]) ->
io::Result<Vec<FlatListData>> {
self.ensure_loaded()?;
- if list_id >= self.nlist {
+ if !self.delta_ids {
return Err(io::Error::new(
- io::ErrorKind::InvalidInput,
- format!("list_id {} out of range (nlist={})", list_id,
self.nlist),
+ io::ErrorKind::InvalidData,
+ "IVF-FLAT reader only supports delta IDs",
));
}
- let count = self.list_counts[list_id] as usize;
- if count == 0 {
- return Ok((Vec::new(), Vec::new()));
- }
-
- let offset = checked_list_offset(self.list_offsets[list_id], list_id)?;
- let vector_bytes = checked_list_bytes(count, self.d * 4)?;
- if self.delta_ids {
+ let mut results = (0..list_ids.len()).map(|_|
None).collect::<Vec<_>>();
+ let mut metas = Vec::new();
+ let mut payloads = Vec::new();
+ for (input_index, &list_id) in list_ids.iter().enumerate() {
+ if list_id >= self.nlist {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!("list_id {} out of range (nlist={})", list_id,
self.nlist),
+ ));
+ }
+ let count = self.list_counts[list_id] as usize;
+ if count == 0 {
+ results[input_index] = Some(FlatListData {
+ list_id,
+ ids: Vec::new(),
+ payload: AlignedFlatPayload::empty(),
+ });
+ continue;
+ }
+ let offset = checked_list_offset(self.list_offsets[list_id],
list_id)?;
+ let vector_bytes = checked_list_bytes(count, self.d * 4)?;
let id_bytes_len = self.list_id_bytes_lens[list_id] as usize;
let payload_len = 12usize
.checked_add(id_bytes_len)
.and_then(|len| len.checked_add(vector_bytes))
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "IVF-FLAT list
payload overflow")
})?;
- let mut payload = vec![0u8; payload_len];
- self.reader
- .pread(&mut [ReadRequest::new(offset, &mut payload)])?;
- let base_id =
i64::from_le_bytes(payload[0..8].try_into().unwrap());
- let encoded_len =
i32::from_le_bytes(payload[8..12].try_into().unwrap());
- if encoded_len < 0 || encoded_len as usize != id_bytes_len {
- return Err(io::Error::new(
- io::ErrorKind::InvalidData,
- "IVF-FLAT id_bytes_len mismatch",
- ));
+ metas.push(FlatListRead {
+ input_index,
+ list_id,
+ count,
+ id_bytes_len,
+ offset,
+ });
+ payloads.push(AlignedFlatPayload::new(
Review Comment:
These payloads are allocated for every selected list before
`pread_batched_slices` starts. `MAX_IVF_BATCH_READ_BYTES` only bounds each
`pread` submission; it does not bound the aggregate live allocation. A single
search passes all probed lists here, and batch search passes the union of all
queries' lists. The same pattern exists in IVF-PQ and IVF-SQ. For example,
IVF-FLAT with 1M 768-dimensional vectors and `nprobe == nlist` can retain
roughly 3 GiB of vector payloads before scanning. This is a regression from the
previous per-list read/scan behavior. Please partition selected lists by
aggregate payload bytes/range count and perform read -> decode/scan -> release
for each chunk, similar to the bounded IVF-RQ path.
##########
core/src/index.rs:
##########
@@ -1218,6 +2412,561 @@ mod tests {
.collect()
}
+ #[test]
+ fn diskann_config_index_type_code_and_name() {
+ let index_type = IndexType::from_code(5).expect("DiskANN index type
code should exist");
+ assert_eq!(index_type.as_str(), "diskann");
+ }
+
+ #[test]
+ fn ivfsq_uses_a_new_type_code_and_retired_type_codes_stay_reserved() {
+ assert_eq!(IndexType::from_code(6), Some(IndexType::IvfSq));
+ assert_eq!(IndexType::IvfSq.as_str(), "ivf_sq");
+ assert_eq!(IndexType::from_code(2), None);
+ assert_eq!(IndexType::from_code(3), None);
+ }
+
+ #[test]
+ fn diskann_config_parses_without_nlist() {
+ let config = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ]))
+ .expect("valid DiskANN options should infer pq.m without nlist");
+
+ assert_eq!(config.index_type(), IndexType::DiskAnn);
+ assert_eq!(config.dimension(), 128);
+ assert_eq!(config.nlist(), 1);
+ let VectorIndexConfig::DiskAnn {
+ pq_m,
+ pq_bits,
+ build,
+ ..
+ } = config
+ else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(pq_m, 32);
+ assert_eq!(pq_bits, 8);
+ assert_eq!(build.memory_budget_bytes, 8 * 1024 * 1024 * 1024);
+ }
+
+ #[test]
+ fn pq_config_parses_code_ratio_and_explicit_m_takes_precedence() {
+ let auto = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_pq"),
+ ("dimension", "128"),
+ ("nlist", "4"),
+ ("metric", "l2"),
+ ("pq.code-ratio", "0.125"),
+ ]))
+ .expect("IVF-PQ should infer pq.m from a relative code budget");
+ let VectorIndexConfig::IvfPq { m, .. } = auto else {
+ panic!("expected IVF-PQ config");
+ };
+ assert_eq!(m, 64);
+
+ let explicit = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ("pq.m", "16"),
+ ("pq.code-ratio", "0.125"),
+ ]))
+ .expect("explicit pq.m should override the code ratio");
+ let VectorIndexConfig::DiskAnn { pq_m, .. } = explicit else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(pq_m, 16);
+ }
+
+ #[test]
+ fn ivf_config_resolves_automatic_nlist_from_expected_count() {
+ let plan = VectorIndexBuildPlan::from_options(&options(&[
+ ("index.type", "ivf_sq"),
+ ("dimension", "100"),
+ ("expected-vector-count", "1183514"),
+ ("nlist", "auto"),
+ ("metric", "cosine"),
+ ]))
+ .unwrap();
+ assert_eq!(plan.expected_vector_count, Some(1_183_514));
+ assert_eq!(plan.config.nlist(), 1024);
+ assert_eq!(plan.config.resolved().nlist, 1024);
+
+ let missing_count = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_flat"),
+ ("dimension", "8"),
+ ("metric", "l2"),
+ ]))
+ .unwrap_err();
+ assert!(missing_count.to_string().contains("expected-vector-count"));
+ }
+
+ #[test]
+ fn capacity_goal_resolves_rq_bits_and_diskann_deployment_layout() {
+ let rq = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_rq"),
+ ("dimension", "100"),
+ ("nlist", "16"),
+ ("metric", "l2"),
+ ("max-bytes-per-vector", "88"),
+ ]))
+ .unwrap();
+ assert_eq!(rq.resolved().rq_bits, Some(4));
+
+ let diskann = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ("deployment-profile", "local_storage"),
+ ("diskann.storage-layout", "auto"),
+ ]))
+ .unwrap();
+ assert_eq!(
+ diskann.resolved().diskann_build.unwrap().storage_layout,
+ DiskAnnStorageLayout::Interleaved
+ );
+ }
+
+ #[test]
+ fn build_time_goal_is_never_silently_guessed_from_hardware() {
+ let values = options(&[
+ ("index.type", "ivf_sq"),
+ ("dimension", "16"),
+ ("nlist", "4"),
+ ("metric", "l2"),
+ ("max-build-seconds", "10"),
+ ]);
+ let plan = VectorIndexBuildPlan::from_options(&values).unwrap();
+ assert_eq!(plan.objective.max_build_seconds, Some(10.0));
+
+ let error = VectorIndexConfig::from_options(&values).unwrap_err();
+ assert!(error
+ .to_string()
+ .contains("requires measured offline calibration"));
+ }
+
+ #[test]
+ fn tagged_search_width_rejects_cross_index_parameters() {
+ assert_eq!(
+ VectorSearchParams::automatic(10)
+ .resolve_ivf_nprobe(1024, 1_000_000, None)
+ .unwrap(),
+ 64
+ );
+ assert!(VectorSearchParams::with_l_search(10, 100)
+ .resolve_ivf_nprobe(1024, 1_000_000, None)
+ .unwrap_err()
+ .to_string()
+ .contains("cannot be used with an IVF"));
+ assert!(VectorSearchParams::new(10, 64)
+ .resolve_diskann_l_search()
+ .unwrap_err()
+ .to_string()
+ .contains("cannot be used with a DiskANN"));
+ }
+
+ #[test]
+ fn automatic_filtered_search_expands_until_results_are_filled() {
+ let mut observed = Vec::new();
+ let result = progressive_ivf_search(
+ VectorSearchParams::automatic(2),
+ 16,
+ 2,
+ 1,
+ 2,
+ 10,
+ |nprobe| {
+ observed.push(nprobe);
+ if nprobe < 8 {
+ Ok((vec![7, -1], vec![1.0, f32::MAX]))
+ } else {
+ Ok((vec![7, 8], vec![1.0, 2.0]))
+ }
+ },
+ )
+ .unwrap();
+ assert_eq!(observed, vec![2, 4, 8]);
+ assert_eq!(result.0, vec![7, 8]);
+ }
+
+ #[test]
+ fn pq_config_constructors_use_the_default_relative_budget() {
+ let ivf = VectorIndexConfig::ivf_pq(128, 4, MetricType::L2,
false).unwrap();
+ let VectorIndexConfig::IvfPq { m, .. } = ivf else {
+ panic!("expected IVF-PQ config");
+ };
+ assert_eq!(m, 32);
+
+ let diskann =
+ VectorIndexConfig::disk_ann(960, MetricType::L2, 8,
DiskAnnBuildParams::default())
+ .unwrap();
+ let VectorIndexConfig::DiskAnn { pq_m, .. } = diskann else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(pq_m, 240);
+ }
+
+ #[test]
+ fn diskann_config_parses_explicit_build_parameters() {
+ let config = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ("pq.m", "16"),
+ ("pq.bits", "4"),
+ ("diskann.max-degree", "32"),
+ ("diskann.build-search-list-size", "64"),
+ ("diskann.alpha", "1.4"),
+ ("diskann.seed", "7"),
+ ("diskann.memory-budget-bytes", "123456"),
+ ("diskann.storage-layout", "interleaved"),
+ ("diskann.raw-vector-encoding", "f16"),
+ ("diskann.build-distance", "full_precision"),
+ ]))
+ .expect("explicit DiskANN build parameters should parse");
+
+ let VectorIndexConfig::DiskAnn {
+ pq_m,
+ pq_bits,
+ build,
+ ..
+ } = config
+ else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(pq_m, 16);
+ assert_eq!(pq_bits, 4);
+ assert_eq!(build.max_degree, 32);
+ assert_eq!(build.build_search_list_size, 64);
+ assert_eq!(build.alpha, 1.4);
+ assert_eq!(build.seed, 7);
+ assert_eq!(build.memory_budget_bytes, 123456);
+ assert_eq!(build.storage_layout, DiskAnnStorageLayout::Interleaved);
+ assert_eq!(build.raw_vector_encoding, DiskAnnRawVectorEncoding::F16);
+ assert_eq!(build.build_distance, DiskAnnBuildDistance::FullPrecision);
+ }
+
+ #[test]
+ fn diskann_build_width_default_follows_explicit_degree() {
+ let diskann = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ("diskann.max-degree", "128"),
+ ]))
+ .expect("omitted Lbuild should follow an explicit DiskANN degree");
+ let VectorIndexConfig::DiskAnn { build, .. } = diskann else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(build.max_degree, 128);
+ assert_eq!(build.build_search_list_size, 128);
+ }
+
+ #[test]
+ fn explicit_diskann_build_search_width_is_preserved() {
+ let diskann = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", "l2"),
+ ("diskann.max-degree", "128"),
+ ("diskann.build-search-list-size", "256"),
+ ]))
+ .expect("explicit Lbuild should take precedence over its automatic
default");
+ let VectorIndexConfig::DiskAnn { build, .. } = diskann else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(build.build_search_list_size, 256);
+ }
+
+ #[test]
+ fn diskann_config_accepts_all_supported_metrics() {
+ for (name, expected) in [
+ ("l2", MetricType::L2),
+ ("inner_product", MetricType::InnerProduct),
+ ("cosine", MetricType::Cosine),
+ ] {
+ let config = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "128"),
+ ("metric", name),
+ ("pq.m", "16"),
+ ]))
+ .expect("DiskANN should accept every public metric");
+ let VectorIndexConfig::DiskAnn { metric, .. } = config else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(metric, expected);
+ }
+ }
+
+ fn assert_diskann_metric_roundtrip(
+ metric: MetricType,
+ data: &[f32],
+ ids: &[i64],
+ queries: &[f32],
+ expected_ids: &[i64],
+ expected_distances: &[f32],
+ ) {
+ let dimension = 2;
+ let count = ids.len();
+ let config = VectorIndexConfig::DiskAnn {
+ dimension,
+ metric,
+ pq_m: 1,
+ pq_bits: 4,
+ build: DiskAnnBuildParams {
+ max_degree: 8,
+ build_search_list_size: 16,
+ raw_vector_encoding: DiskAnnRawVectorEncoding::F32,
+ build_distance: DiskAnnBuildDistance::ProductQuantized,
+ ..DiskAnnBuildParams::default()
+ },
+ };
+ let mut writer = build_writer(config, data, count);
+ writer.add_vectors(ids, data, count).unwrap();
+ let mut bytes = Vec::new();
+ writer.write(&mut PosWriter::new(&mut bytes)).unwrap();
+
+ let mut reader = VectorIndexReader::open(Cursor::new(bytes)).unwrap();
+ assert_eq!(reader.metadata().metric, metric);
+ let params = VectorSearchParams::with_l_search(1, count);
Review Comment:
This sets `l_search` to the full vector count, so the new IP/cosine tests
validate format propagation and exact reranking but do not exercise approximate
graph/PQ candidate quality. Please add a deterministic recall regression with
`l_search << N` (ideally covering both PQ-guided/full-precision build and
batch/filter paths). I ran an additional 2K x 32 random-data check and saw
about 99.5%--99.6% recall@10 at `l_search=50`, so this is a coverage request
rather than a currently observed correctness failure.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
