This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new f80e556 Support high-dimensional DiskANN vectors (#114)
f80e556 is described below
commit f80e5568dc476f7e0bd56c73c88a0b46bf8018ca
Author: jerry <[email protected]>
AuthorDate: Fri Oct 9 16:11:50 2026 +0800
Support high-dimensional DiskANN vectors (#114)
---
core/STORAGE_FORMAT.md | 2 +-
core/src/diskann.rs | 3 -
core/src/diskann_io.rs | 7 +-
core/src/diskann_search.rs | 228 +++++++++++++++++++++++++++++++++++++--------
core/src/index.rs | 29 +++++-
docs/diskann.html | 4 +-
6 files changed, 219 insertions(+), 54 deletions(-)
diff --git a/core/STORAGE_FORMAT.md b/core/STORAGE_FORMAT.md
index 27ec901..2ef8098 100644
--- a/core/STORAGE_FORMAT.md
+++ b/core/STORAGE_FORMAT.md
@@ -103,7 +103,7 @@ a v1 writer may emit requires version 2; it is not a
compatible v1 extension.
The following header invariants are part of v1 and are enforced symmetrically
by
the writer and reader:
-- `1 <= dimension <= 1024`;
+- `1 <= dimension <= u32::MAX`, and `dimension * element_size <= u32::MAX`;
- `metric` is `0`, `1`, or `2`;
- `1 <= vector_count <= u32::MAX` and `entry_node < vector_count`;
- `1 <= pq_m <= dimension`, and `pq_bits` is 4 or 8;
diff --git a/core/src/diskann.rs b/core/src/diskann.rs
index b668b59..2406c93 100644
--- a/core/src/diskann.rs
+++ b/core/src/diskann.rs
@@ -112,9 +112,6 @@ pub(crate) fn validate_diskann_format_configuration(
if dimension == 0 {
return Err(invalid_input("DiskANN dimension must be greater than 0"));
}
- if dimension > 1024 {
- return Err(invalid_input("DiskANN v1 dimension must be at most 1024"));
- }
if pq_m == 0 {
return Err(invalid_input("DiskANN pq.m must be greater than 0"));
}
diff --git a/core/src/diskann_io.rs b/core/src/diskann_io.rs
index b22a2e8..1ccf68b 100644
--- a/core/src/diskann_io.rs
+++ b/core/src/diskann_io.rs
@@ -2002,11 +2002,6 @@ impl DiskAnnHeader {
let vector_record_size = dimension_u32
.checked_mul(build.raw_vector_encoding.element_size() as u32)
.ok_or_else(|| invalid_input("DiskANN vector record size overflows
u32"))?;
- if vector_record_size == 0 || vector_record_size > DISKANN_PAGE_SIZE {
- return Err(invalid_input(
- "DiskANN raw vector does not fit in a logical page",
- ));
- }
if build.storage_layout == DiskAnnStorageLayout::Interleaved {
let maximum_adjacency_bytes = max_degree
.checked_mul(size_of::<u32>() as u32)
@@ -2190,7 +2185,7 @@ impl DiskAnnHeader {
}
fn validate_layout(&self) -> io::Result<()> {
- if self.dimension == 0 || self.dimension > 1024 {
+ if self.dimension == 0 {
return Err(invalid_data("invalid DiskANN dimension"));
}
if self.vector_count == 0 || self.vector_count > u32::MAX as u64 {
diff --git a/core/src/diskann_search.rs b/core/src/diskann_search.rs
index 959e8f0..7ed6aa8 100644
--- a/core/src/diskann_search.rs
+++ b/core/src/diskann_search.rs
@@ -21,8 +21,8 @@ use crate::diskann_io::{
DISKANN_PAGE_SIZE,
};
use crate::distance::{
- fvec_distance, fvec_l2sqr, pq_distance_four_codes, pq_distance_from_table,
preprocess_vectors,
- MetricType,
+ fvec_distance, fvec_inner_product, fvec_l2sqr, fvec_norm_l2sqr,
pq_distance_four_codes,
+ pq_distance_from_table, preprocess_vectors, MetricType,
};
use crate::index_io_util::decode_roaring_filter;
use crate::io::{ReadRequest, SeekRead};
@@ -3453,27 +3453,48 @@ fn metric_distance_from_f16_le_bytes(
"DiskANN f16 raw-vector record has invalid length",
));
}
- if query.len() > 1024 {
- return Err(invalid_data(
- "DiskANN f16 raw-vector dimension exceeds the v1 limit",
- ));
- }
let mut decoded = [0.0f32; 1024];
- for (slot, component) in decoded[..query.len()]
- .iter_mut()
- .zip(bytes.chunks_exact(size_of::<u16>()))
- {
- let value = half::f16::from_bits(u16::from_le_bytes(
- component
- .try_into()
- .expect("validated two-byte raw-vector component"),
- ));
- if !value.is_finite() {
- return Err(invalid_data("DiskANN raw vectors must be finite"));
+ let fits_stack = query.len() <= decoded.len();
+ let mut dot = 0.0;
+ let mut query_norm = 0.0;
+ let mut vector_norm = 0.0;
+ for (query_chunk, bytes_chunk) in
query.chunks(1024).zip(bytes.chunks(2048)) {
+ let values = &mut decoded[..query_chunk.len()];
+ for (slot, component) in values
+ .iter_mut()
+ .zip(bytes_chunk.chunks_exact(size_of::<u16>()))
+ {
+ let value = half::f16::from_bits(u16::from_le_bytes(
+ component
+ .try_into()
+ .expect("validated two-byte raw-vector component"),
+ ));
+ if !value.is_finite() {
+ return Err(invalid_data("DiskANN raw vectors must be finite"));
+ }
+ *slot = value.to_f32();
+ }
+ if fits_stack {
+ return Ok(fvec_distance(query, values, metric));
+ }
+ dot += fvec_inner_product(query_chunk, values);
+ if metric == MetricType::Cosine {
+ query_norm += fvec_norm_l2sqr(query_chunk);
+ vector_norm += fvec_norm_l2sqr(values);
}
- *slot = value.to_f32();
}
- Ok(fvec_distance(query, &decoded[..query.len()], metric))
+ Ok(match metric {
+ MetricType::InnerProduct => -dot,
+ MetricType::Cosine => {
+ let denominator = query_norm.sqrt() * vector_norm.sqrt();
+ if denominator > 0.0 {
+ 1.0 - dot / denominator
+ } else {
+ 1.0
+ }
+ }
+ MetricType::L2 => unreachable!("L2 uses the selected raw-vector
kernel"),
+ })
}
fn l2_distance_from_f16_le_bytes(query: &[f32], bytes: &[u8]) ->
io::Result<f32> {
@@ -3486,11 +3507,6 @@ fn l2_distance_from_f16_le_bytes(query: &[f32], bytes:
&[u8]) -> io::Result<f32>
"DiskANN f16 raw-vector record has invalid length",
));
}
- if query.len() > 1024 {
- return Err(invalid_data(
- "DiskANN f16 raw-vector dimension exceeds the v1 limit",
- ));
- }
#[cfg(all(target_endian = "little", target_arch = "aarch64"))]
if query.len() >= 4 {
// SAFETY: AArch64 guarantees NEON. The kernel uses unaligned loads,
@@ -3499,23 +3515,27 @@ fn l2_distance_from_f16_le_bytes(query: &[f32], bytes:
&[u8]) -> io::Result<f32>
return unsafe { l2_distance_from_f16_le_bytes_neon(query, bytes) };
}
let mut bits = [0u16; 1024];
- for (slot, component) in bits[..query.len()]
- .iter_mut()
- .zip(bytes.chunks_exact(size_of::<u16>()))
- {
- *slot = u16::from_le_bytes(
- component
- .try_into()
- .expect("validated two-byte raw-vector component"),
- );
- }
- let values = bits[..query.len()].reinterpret_cast::<half::f16>();
- if values.iter().any(|value| !value.is_finite()) {
- return Err(invalid_data("DiskANN raw vectors must be finite"));
- }
let mut decoded = [0.0f32; 1024];
- values.convert_to_f32_slice(&mut decoded[..query.len()]);
- Ok(fvec_l2sqr(query, &decoded[..query.len()]))
+ let mut distance = 0.0;
+ for (query_chunk, bytes_chunk) in
query.chunks(1024).zip(bytes.chunks(2048)) {
+ for (slot, component) in bits[..query_chunk.len()]
+ .iter_mut()
+ .zip(bytes_chunk.chunks_exact(size_of::<u16>()))
+ {
+ *slot = u16::from_le_bytes(
+ component
+ .try_into()
+ .expect("validated two-byte raw-vector component"),
+ );
+ }
+ let values = bits[..query_chunk.len()].reinterpret_cast::<half::f16>();
+ if values.iter().any(|value| !value.is_finite()) {
+ return Err(invalid_data("DiskANN raw vectors must be finite"));
+ }
+ values.convert_to_f32_slice(&mut decoded[..query_chunk.len()]);
+ distance += fvec_l2sqr(query_chunk, &decoded[..query_chunk.len()]);
+ }
+ Ok(distance)
}
#[cfg(all(target_endian = "little", target_arch = "aarch64"))]
@@ -4018,6 +4038,63 @@ mod tests {
);
}
+ #[test]
+ fn diskann_f16_high_dimension_distances_support_all_metrics() {
+ let query = vec![1.0f32; 2560];
+ let bytes = half::f16::ONE.to_bits().to_le_bytes().repeat(query.len());
+ for (metric, expected) in [
+ (MetricType::L2, 0.0),
+ (MetricType::InnerProduct, -(query.len() as f32)),
+ (MetricType::Cosine, 0.0),
+ ] {
+ assert_eq!(
+ raw_vector_distance(
+ &query,
+ &bytes,
+ DiskAnnRawVectorEncoding::F16,
+ metric,
+ RawVectorDistanceKernel::Scalar,
+ )
+ .unwrap(),
+ expected
+ );
+ }
+
+ let query = (0..2560)
+ .map(|index| (index % 17) as f32 - 8.0)
+ .collect::<Vec<_>>();
+ let vector = (0..2560)
+ .map(|index| ((index % 13) as f32 - 6.0) / 4.0)
+ .collect::<Vec<_>>();
+ let bytes = vector
+ .iter()
+ .flat_map(|&value|
half::f16::from_f32(value).to_bits().to_le_bytes())
+ .collect::<Vec<_>>();
+ for metric in [MetricType::InnerProduct, MetricType::Cosine] {
+ let expected = fvec_distance(&query, &vector, metric);
+ let actual = metric_distance_from_f16_le_bytes(&query, &bytes,
metric).unwrap();
+ assert!((actual - expected).abs() <= expected.abs().max(1.0) *
1.0e-5);
+ }
+ assert_eq!(
+ metric_distance_from_f16_le_bytes(
+ &query,
+ &vec![0; query.len() * 2],
+ MetricType::Cosine,
+ )
+ .unwrap(),
+ 1.0
+ );
+ let mut non_finite = bytes;
+ non_finite[2048 * 2..2048 * 2 + 2]
+ .copy_from_slice(&half::f16::INFINITY.to_bits().to_le_bytes());
+ assert_eq!(
+ metric_distance_from_f16_le_bytes(&query, &non_finite,
MetricType::InnerProduct)
+ .unwrap_err()
+ .kind(),
+ io::ErrorKind::InvalidData
+ );
+ }
+
#[test]
fn diskann_compact_f16_roundtrips_dense_vector_records() {
let dimension = 8;
@@ -4055,6 +4132,75 @@ mod tests {
assert_eq!(distances, vec![0.0]);
}
+ #[test]
+ fn diskann_compact_roundtrips_vectors_larger_than_a_page() {
+ let dimension = 2560;
+ let count = 16;
+ let data = (0..count)
+ .flat_map(|row| (0..dimension).map(move |column| (row * 7 + column
% 13) as f32))
+ .collect::<Vec<_>>();
+ let ids = (1000..1000 + count as i64).collect::<Vec<_>>();
+ for encoding in [DiskAnnRawVectorEncoding::F32,
DiskAnnRawVectorEncoding::F16] {
+ let mut index = DiskAnnIndex::with_pq_bits(
+ dimension,
+ MetricType::L2,
+ 16,
+ 4,
+ DiskAnnBuildParams {
+ max_degree: 4,
+ build_search_list_size: 8,
+ raw_vector_encoding: encoding,
+ ..DiskAnnBuildParams::default()
+ },
+ );
+ index.train(&data, count).unwrap();
+ index.add(&data, &ids);
+ let mut bytes = Vec::new();
+ write_diskann_index(&index, &mut PosWriter::new(&mut
bytes)).unwrap();
+ let header =
crate::diskann_io::DiskAnnHeader::decode(&bytes[..256]).unwrap();
+ assert!(header.vector_record_size >
crate::diskann_io::DISKANN_PAGE_SIZE);
+ let mut reader =
DiskAnnIndexReader::open(Cursor::new(bytes)).unwrap();
+ let (result_ids, distances) = reader.search(&data[..dimension], 1,
16).unwrap();
+ assert_eq!(result_ids, vec![ids[0]]);
+ assert_eq!(distances, vec![0.0]);
+ let (batch_ids, batch_distances) =
+ reader.search_batch(&data[..dimension * 2], 1, 16).unwrap();
+ assert_eq!(batch_ids, ids[..2]);
+ assert_eq!(batch_distances, vec![0.0; 2]);
+ }
+ }
+
+ #[test]
+ fn diskann_compact_f32_high_dimension_preserves_inner_product_distance() {
+ let dimension = 2560;
+ let count = 16;
+ let data = (0..count)
+ .flat_map(|row| (0..dimension).map(move |column| f32::from(column
== row)))
+ .collect::<Vec<_>>();
+ let ids = (1000..1000 + count as i64).collect::<Vec<_>>();
+ let mut index = DiskAnnIndex::with_pq_bits(
+ dimension,
+ MetricType::InnerProduct,
+ 16,
+ 4,
+ DiskAnnBuildParams {
+ max_degree: 4,
+ build_search_list_size: 8,
+ raw_vector_encoding: DiskAnnRawVectorEncoding::F32,
+ ..DiskAnnBuildParams::default()
+ },
+ );
+ index.train(&data, count).unwrap();
+ index.add(&data, &ids);
+ let mut bytes = Vec::new();
+ write_diskann_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+ let mut reader = DiskAnnIndexReader::open(Cursor::new(bytes)).unwrap();
+
+ let (result_ids, distances) = reader.search(&data[..dimension], 1,
count).unwrap();
+ assert_eq!(result_ids, vec![ids[0]]);
+ assert_eq!(distances, vec![-1.0]);
+ }
+
#[test]
fn
diskann_interleaved_layout_searches_and_batches_without_a_vector_section() {
let dimension = 8;
diff --git a/core/src/index.rs b/core/src/index.rs
index 7640e3c..8e0a792 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -4276,7 +4276,6 @@ mod tests {
#[test]
fn diskann_config_rejects_invalid_shape_and_build_parameters() {
for (extra, expected) in [
- (vec![("dimension", "1025")], "at most 1024"),
(vec![("pq.m", "0")], "pq.m must be greater than 0"),
(vec![("pq.m", "129")], "must not exceed dimension"),
(vec![("pq.bits", "6")], "pq.bits must be 4 or 8"),
@@ -4334,6 +4333,34 @@ mod tests {
}
}
+ #[test]
+ fn diskann_config_accepts_high_dimension_compact_layout() {
+ let config = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "diskann"),
+ ("dimension", "2560"),
+ ("metric", "inner_product"),
+ ("pq.bits", "8"),
+ ("pq.code-ratio", "0.0625"),
+ ("diskann.storage-layout", "compact"),
+ ("diskann.raw-vector-encoding", "f32"),
+ ]))
+ .unwrap();
+
+ let VectorIndexConfig::DiskAnn {
+ pq_m,
+ pq_bits,
+ build,
+ ..
+ } = config
+ else {
+ panic!("expected DiskANN config");
+ };
+ assert_eq!(pq_m, 640);
+ assert_eq!(pq_bits, 8);
+ assert_eq!(build.storage_layout, DiskAnnStorageLayout::Compact);
+ assert_eq!(build.raw_vector_encoding, DiskAnnRawVectorEncoding::F32);
+ }
+
#[test]
fn diskann_trainer_uses_configured_dimension() {
let trainer = VectorIndexTrainer::new(VectorIndexConfig::DiskAnn {
diff --git a/docs/diskann.html b/docs/diskann.html
index cf7ef12..4a1dbe7 100644
--- a/docs/diskann.html
+++ b/docs/diskann.html
@@ -153,7 +153,7 @@ ids, distances = reader.search(
<section class="article-section" id="build-parameters">
<h2>Build parameters</h2>
<div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Default /
validation</th><th>What it controls</th><th>When
increased</th></tr></thead><tbody>
- <tr><td><code>dimension</code></td><td>Inferred by Java/Python
one-shot training; required by streaming APIs;
<code>1..=1024</code></td><td>Raw vector width and PQ codebook
size.</td><td>Raw-vector payload and rerank work grow linearly.</td></tr>
+ <tr><td><code>dimension</code></td><td>Inferred by Java/Python
one-shot training; required by streaming APIs; positive integer</td><td>Raw
vector width and PQ codebook size. Interleaved records must still fit in a 4
KiB page.</td><td>Raw-vector payload and rerank work grow linearly.</td></tr>
<tr><td><code>metric</code></td><td>Required: <code>l2</code>,
<code>inner_product</code>, or <code>cosine</code></td><td>Distance used by
graph build, PQ navigation, and exact reranking. Inner-product scores are
negative dot products; cosine scores are <code>1 -
cosine</code>.</td><td>Semantic, not inferred. Cosine vectors and queries are
normalized internally; zero vectors retain distance 1.</td></tr>
<tr><td><code>target-recall</code></td><td>Optional;
<code>0..=1</code></td><td>Selects <code>fast_build</code> at ≤ 0.85,
<code>high_recall</code> at ≥ 0.97, and <code>balanced</code> between them when
no preset is pinned.</td><td>This is a deterministic starting policy, not a
measured guarantee; validate on held-out queries.</td></tr>
<tr><td><code>max-bytes-per-vector</code></td><td>Optional
positive persisted-size budget</td><td>Guides PQ width and may select 4-bit PQ
and F16 rerank records. Configuration is rejected before training when the
estimated row bytes plus amortized fixed data exceed the budget.</td><td>A
larger allowance preserves more PQ/raw precision. The check is conservative;
compression, headers, and alignment mean it remains a per-vector sizing bound
rather than an exact final-file-size [...]
@@ -254,7 +254,7 @@ ids, distances = reader.search(
<section class="article-section" id="boundaries">
<h2>Current implementation boundaries</h2>
<ul>
- <li>v1 supports squared L2, negative-dot-product inner product,
and <code>1 - cosine</code>, with dimensions up to 1024 and 4/8-bit PQ.</li>
+ <li>v1 supports squared L2, negative-dot-product inner product,
and <code>1 - cosine</code>, with positive dimensions and 4/8-bit PQ.
Interleaved records remain limited by the 4 KiB page.</li>
<li>The index is static. FreshDiskANN-style insert, delete,
consolidation, and graph repair are not implemented.</li>
<li>Serialization builds the complete graph in memory. Raw vectors
are stored as explicit F32 or F16 scalar records; there is no PQ-only result
mode.</li>
<li>Recall is sensitive to both PQ capacity and graph/search
budgets. On the public SIFT1M and GIST1M runs, the default ratio resolved to
<code>m=32</code> and <code>m=240</code>. DiskANN Recall@10 was 0.995 and 0.973
at <code>l_search=200</code>, while GIST rose to 0.991 at
<code>l_search=500</code>; use the public-corpus procedure rather than
synthetic vectors when validating a new ratio.</li>
