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-rust.git
The following commit(s) were added to refs/heads/main by this push:
new a48c01b1 perf(vindex): optimize IVF index build reads (#800)
a48c01b1 is described below
commit a48c01b1523da53335c245614701f1e7f9c316c4
Author: jerry <[email protected]>
AuthorDate: Mon Sep 21 10:25:15 2026 +0800
perf(vindex): optimize IVF index build reads (#800)
---
crates/paimon/src/arrow/format/parquet.rs | 545 ++++++-
.../paimon/src/table/vindex_index_build_builder.rs | 1 +
.../table/vindex_index_build_builder/extraction.rs | 108 +-
.../table/vindex_index_build_builder/pipeline.rs | 1670 ++++++++++++++++++++
.../src/table/vindex_index_build_builder/tests.rs | 379 ++++-
.../src/table/vindex_index_build_builder/timing.rs | 8 +-
.../src/table/vindex_index_build_builder/writer.rs | 149 +-
crates/paimon/src/vindex/mod.rs | 74 +
docs/src/sql.md | 1 +
9 files changed, 2839 insertions(+), 96 deletions(-)
diff --git a/crates/paimon/src/arrow/format/parquet.rs
b/crates/paimon/src/arrow/format/parquet.rs
index bc7403a3..34cd8335 100644
--- a/crates/paimon/src/arrow/format/parquet.rs
+++ b/crates/paimon/src/arrow/format/parquet.rs
@@ -98,6 +98,141 @@ impl ParquetFormatReader {
}
}
+#[derive(Debug, Clone)]
+pub(crate) struct ParquetGranule {
+ pub(crate) first_row: i64,
+ pub(crate) row_count: i64,
+ pub(crate) byte_ranges: Vec<Range<u64>>,
+}
+
+pub(crate) async fn parquet_granules(
+ reader: Box<dyn FileRead>,
+ file_size: u64,
+ column_name: &str,
+ page_index_enabled: bool,
+) -> crate::Result<(Vec<ParquetGranule>, bool)> {
+ let mut options = ArrowReaderOptions::new();
+ if page_index_enabled {
+ options = options.with_offset_index_policy(PageIndexPolicy::Optional);
+ }
+ let mut reader = ArrowFileReader::new(file_size, reader.into());
+ let metadata = reader.get_metadata(Some(&options)).await?;
+ let columns = metadata
+ .file_metadata()
+ .schema_descr()
+ .columns()
+ .iter()
+ .enumerate()
+ .filter_map(|(index, column)| {
+ column
+ .path()
+ .parts()
+ .first()
+ .is_some_and(|part| part == column_name)
+ .then_some(index)
+ })
+ .collect::<Vec<_>>();
+ if columns.is_empty() {
+ return Err(Error::DataInvalid {
+ message: format!("Parquet column '{column_name}' not found"),
+ source: None,
+ });
+ }
+ let Some(offset_index) = metadata
+ .offset_index()
+ .filter(|index| index.len() == metadata.row_groups().len())
+ else {
+ return Ok((row_group_granules(&metadata, &columns), false));
+ };
+
+ let mut granules = Vec::new();
+ let mut cursor = 0i64;
+ for (row_group, indexes) in metadata.row_groups().iter().zip(offset_index)
{
+ let Ok(row_count) = usize::try_from(row_group.num_rows()) else {
+ return Ok((row_group_granules(&metadata, &columns), false));
+ };
+ let Some(leaf_pages) = columns
+ .iter()
+ .map(|index| indexes.get(*index).map(|index|
index.page_locations()))
+ .collect::<Option<Vec<_>>>()
+ else {
+ return Ok((row_group_granules(&metadata, &columns), false));
+ };
+ let pages = leaf_pages[0];
+ if pages.is_empty()
+ || leaf_pages.iter().any(|other| {
+ !page_boundaries_valid(other, row_count)
+ || other.len() != pages.len()
+ || other
+ .iter()
+ .zip(pages)
+ .any(|(a, b)| a.first_row_index != b.first_row_index)
+ })
+ {
+ return Ok((row_group_granules(&metadata, &columns), false));
+ }
+ for (page_index, page) in pages.iter().enumerate() {
+ let next = pages
+ .get(page_index + 1)
+ .map_or(row_group.num_rows(), |page| page.first_row_index);
+ let mut byte_ranges = Vec::with_capacity(columns.len() * 2);
+ for (column_index, pages) in columns.iter().zip(&leaf_pages) {
+ let column = row_group.column(*column_index);
+ let (column_start, _) = column.byte_range();
+ if let Some(first) = pages.first() {
+ let Ok(first_page) = u64::try_from(first.offset) else {
+ return Ok((row_group_granules(&metadata, &columns),
false));
+ };
+ if column_start < first_page {
+ byte_ranges.push(column_start..first_page);
+ }
+ }
+ let page = &pages[page_index];
+ let (Ok(start), Ok(length)) = (
+ u64::try_from(page.offset),
+ u64::try_from(page.compressed_page_size),
+ ) else {
+ return Ok((row_group_granules(&metadata, &columns),
false));
+ };
+ let Some(end) = start.checked_add(length).filter(|end| *end <=
file_size) else {
+ return Ok((row_group_granules(&metadata, &columns),
false));
+ };
+ byte_ranges.push(start..end);
+ }
+ granules.push(ParquetGranule {
+ first_row: cursor + page.first_row_index,
+ row_count: next - page.first_row_index,
+ byte_ranges,
+ });
+ }
+ cursor += row_group.num_rows();
+ }
+ Ok((granules, true))
+}
+
+fn row_group_granules(metadata: &ParquetMetaData, columns: &[usize]) ->
Vec<ParquetGranule> {
+ let mut cursor = 0i64;
+ metadata
+ .row_groups()
+ .iter()
+ .map(|row_group| {
+ let granule = ParquetGranule {
+ first_row: cursor,
+ row_count: row_group.num_rows(),
+ byte_ranges: columns
+ .iter()
+ .map(|index| {
+ let (start, length) =
row_group.column(*index).byte_range();
+ start..start + length
+ })
+ .collect(),
+ };
+ cursor += row_group.num_rows();
+ granule
+ })
+ .collect()
+}
+
enum ParquetRowGroupMessage {
Batch(RecordBatch),
Error(Error),
@@ -631,12 +766,12 @@ impl FormatFileReader for ParquetFormatReader {
// predicate-free path and run a bounded number concurrently.
//
// Row-group receivers are consumed in order and buffer one batch each,
- // preserving positional `_ROW_ID`, sort order, and batch
backpressure. Reads
- // with predicates or an explicit row selection retain the original
- // single-stream path until their selections are split per row group.
- let read_budget = self.read_budget.as_ref().filter(|_| {
- preds.is_empty() && row_filter_factory.is_none() &&
row_selection.is_none()
- });
+ // preserving positional `_ROW_ID`, sort order, and batch backpressure.
+ // Reads with predicates retain the original single-stream path.
+ let read_budget = self
+ .read_budget
+ .as_ref()
+ .filter(|_| preds.is_empty() && row_filter_factory.is_none());
let row_group_parallelism = read_budget
.map(|budget| {
budget
@@ -644,39 +779,50 @@ impl FormatFileReader for ParquetFormatReader {
.min(batch_stream_builder.metadata().num_row_groups())
})
.unwrap_or(1);
- let projected_bytes = self
+ let selected_row_groups = self
.read_budget
.as_ref()
.filter(|budget| row_group_parallelism > 1 ||
budget.diagnostics_enabled())
.map(|budget| {
- let mut diagnostic_selection = combined_selection;
- let projected_bytes = batch_stream_builder
+ let mut row_group_selection = combined_selection;
+ let selected_row_groups = batch_stream_builder
.metadata()
.row_groups()
.iter()
- .filter(|row_group| {
- diagnostic_selection.as_mut().is_none_or(|selection| {
- selection
- .split_off(row_group.num_rows() as usize)
- .selects_any()
- })
+ .enumerate()
+ .filter_map(|(row_group_index, row_group)| {
+ let selection = row_group_selection
+ .as_mut()
+ .map(|selection|
selection.split_off(row_group.num_rows() as usize));
+ if selection
+ .as_ref()
+ .is_some_and(|selection| !selection.selects_any())
+ {
+ return None;
+ }
+ let projected_bytes =
projected_row_group_bytes(row_group, &mask);
+ Some((row_group_index, selection, projected_bytes))
})
- .map(|row_group| projected_row_group_bytes(row_group,
&mask))
+ .collect::<Vec<_>>();
+ let projected_bytes = selected_row_groups
+ .iter()
+ .map(|(_, _, projected_bytes)| *projected_bytes)
.collect::<Vec<_>>();
budget.record_projected_row_groups(&projected_bytes);
- projected_bytes
+ selected_row_groups
});
if row_group_parallelism > 1 {
- let row_group_count =
batch_stream_builder.metadata().num_row_groups();
+ let selected_row_groups =
+ selected_row_groups.expect("parallel row-group reads need a
selection plan");
+ let row_group_count = selected_row_groups.len();
let reader_metadata = ArrowReaderMetadata::try_new(
batch_stream_builder.metadata().clone(),
ArrowReaderOptions::new(),
)?;
- let projected_bytes = projected_bytes.expect("parallel row-group
reads need sizes");
let read_budget = Arc::clone(read_budget.expect("checked above"));
let (row_group_tx, mut row_group_rx) =
mpsc::channel(row_group_parallelism);
tokio::spawn(async move {
- for (row_group_index, projected_bytes) in
projected_bytes.into_iter().enumerate() {
+ for (row_group_index, selection, projected_bytes) in
selected_row_groups {
let Ok(slot) = row_group_tx.reserve().await else {
return;
};
@@ -701,6 +847,7 @@ impl FormatFileReader for ParquetFormatReader {
row_group_mask,
row_group_index,
batch_size,
+ selection,
permit,
batch_tx,
));
@@ -830,6 +977,7 @@ async fn read_row_group(
projection: ProjectionMask,
row_group_index: usize,
batch_size: Option<usize>,
+ selection: Option<RowSelection>,
_permit: ReadPermit,
sender: mpsc::Sender<ParquetRowGroupMessage>,
) {
@@ -839,6 +987,9 @@ async fn read_row_group(
)
.with_projection(projection)
.with_row_groups(vec![row_group_index]);
+ if let Some(selection) = selection {
+ builder = builder.with_row_selection(selection);
+ }
if let Some(size) = batch_size {
builder = builder.with_batch_size(size);
}
@@ -2359,6 +2510,13 @@ fn merge_byte_ranges(ranges: &[Range<u64>], coalesce:
u64) -> Vec<Range<u64>> {
merged
}
+pub(crate) fn coalesced_parquet_range_bytes(ranges: &[Range<u64>]) -> u64 {
+ merge_byte_ranges(ranges, RANGE_COALESCE_BYTES)
+ .into_iter()
+ .map(|range| range.end - range.start)
+ .sum()
+}
+
/// Split merged ranges into fixed-size batches to utilize concurrency,
/// Each merged range is divided into chunks of `expected_size`,
/// with the last chunk taking whatever remains.
@@ -2418,8 +2576,8 @@ fn split_ranges_for_concurrency(merged: Vec<Range<u64>>,
concurrency: usize) ->
mod tests {
use super::build_parquet_row_filter;
use super::{
- forward_row_group_batches, parse_compression, supported_compressions,
FilePredicates,
- ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage,
+ forward_row_group_batches, parquet_granules, parse_compression,
supported_compressions,
+ FilePredicates, ParquetFormatReader, ParquetFormatWriter,
ParquetRowGroupMessage,
};
use super::{
AsyncArrowWriter, Bytes, PageIndexPolicy, ParquetMetaDataReader,
Predicate,
@@ -3004,6 +3162,113 @@ mod tests {
);
}
+ #[tokio::test]
+ async fn test_sparse_row_groups_preserve_selection_order_and_budget() {
+ let data = write_multi_row_group_parquet(64, 384,
EnabledStatistics::Chunk, false).await;
+ let in_flight = Arc::new(AtomicUsize::new(0));
+ let max_in_flight = Arc::new(AtomicUsize::new(0));
+ let file_reader = ConcurrentTrackingFileRead {
+ data: Bytes::from(data),
+ in_flight,
+ max_in_flight,
+ };
+ let file_size = file_reader.data.len() as u64;
+ let ranges = vec![
+ RowRange::new(60, 68),
+ RowRange::new(130, 135),
+ RowRange::new(258, 263),
+ RowRange::new(380, 383),
+ ];
+ let budget = Arc::new(ReadBudget::new(2, 256 * 1024 * 1024).unwrap());
+ budget.enable_diagnostics();
+
+ let batches =
ParquetFormatReader::with_read_budget(Arc::clone(&budget))
+ .read_batch_stream(
+ Box::new(file_reader),
+ file_size,
+ &[int_field("id")],
+ None,
+ Some(32),
+ Some(ranges.clone()),
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let actual = batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .iter()
+ .copied()
+ .collect::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
+ let expected = ranges
+ .iter()
+ .flat_map(|range| range.from() as i32..=range.to() as i32)
+ .collect::<Vec<_>>();
+
+ assert_eq!(actual, expected);
+ let diagnostics = budget.diagnostics();
+ assert_eq!(diagnostics.row_group_count, 5);
+ assert_eq!(diagnostics.peak_inflight, 2);
+ assert_eq!(diagnostics.current_inflight, 0);
+ }
+
+ #[tokio::test]
+ async fn test_sparse_page_budget_charges_full_projected_row_group() {
+ const MIB: i64 = 1024 * 1024;
+
+ let bytes = write_multi_page_parquet(10, 80).await;
+ let metadata = load_metadata_with_page_index(&bytes, true);
+ let offset_index = &metadata.offset_index().unwrap()[0];
+ let page_locations = offset_index[0].page_locations();
+ let compressed_bytes = page_locations
+ .iter()
+ .map(|page| i64::from(page.compressed_page_size))
+ .sum();
+ let mut row_group = metadata.row_groups()[0].clone();
+ let column = row_group
+ .column(0)
+ .clone()
+ .into_builder()
+ .set_total_compressed_size(compressed_bytes)
+ .set_total_uncompressed_size(308 * MIB)
+ .set_data_page_offset(page_locations[0].offset)
+ .set_dictionary_page_offset(None)
+ .build()
+ .unwrap();
+ row_group.columns_mut()[0] = column;
+
+ let projection =
super::ProjectionMask::roots(row_group.schema_descr(), [0]);
+ let projected_bytes = super::projected_row_group_bytes(&row_group,
&projection);
+
+ assert_eq!(projected_bytes, 308 * MIB as u64);
+ let budget = ReadBudget::new(8, 256 * MIB as u64).unwrap();
+ budget.enable_diagnostics();
+ let permit = budget.acquire(projected_bytes).await.unwrap();
+ assert!(
+ tokio::time::timeout(Duration::from_millis(20), budget.acquire(1))
+ .await
+ .is_err()
+ );
+ drop(permit);
+ let permit = tokio::time::timeout(Duration::from_secs(1),
budget.acquire(projected_bytes))
+ .await
+ .unwrap()
+ .unwrap();
+ assert_eq!(budget.diagnostics().peak_inflight, 1);
+ drop(permit);
+ assert_eq!(budget.diagnostics().current_inflight, 0);
+ }
+
#[tokio::test]
async fn test_parquet_read_budget_is_shared_across_readers() {
const ROWS: i32 = 256;
@@ -3093,7 +3358,7 @@ mod tests {
#[tokio::test]
async fn test_parquet_diagnostics_include_reads_with_row_selection() {
- let data = write_multi_row_group_parquet(32, 64,
EnabledStatistics::Chunk).await;
+ let data = write_multi_row_group_parquet(32, 64,
EnabledStatistics::Chunk, false).await;
let budget = Arc::new(ReadBudget::new(8, 256 * 1024 * 1024).unwrap());
budget.enable_diagnostics();
let file_size = data.len() as u64;
@@ -3117,7 +3382,7 @@ mod tests {
let diagnostics = budget.diagnostics();
assert_eq!(diagnostics.row_group_count, 1);
assert!(diagnostics.projected_bytes_total > 0);
- assert_eq!(diagnostics.peak_inflight, 0);
+ assert_eq!(diagnostics.peak_inflight, 1);
}
#[tokio::test]
@@ -3706,11 +3971,13 @@ mod tests {
row_group_rows: usize,
total_rows: i32,
statistics: EnabledStatistics,
+ offset_index_disabled: bool,
) -> Vec<u8> {
let schema = writer_arrow_schema();
let props = parquet::file::properties::WriterProperties::builder()
.set_max_row_group_row_count(Some(row_group_rows))
.set_statistics_enabled(statistics)
+ .set_offset_index_disabled(offset_index_disabled)
.build();
let mut buf = Vec::new();
let mut writer = AsyncArrowWriter::try_new(&mut buf, schema.clone(),
Some(props)).unwrap();
@@ -3726,7 +3993,7 @@ mod tests {
#[tokio::test]
async fn test_row_group_selection_in_uses_min_max_without_page_index() {
- let bytes = write_multi_row_group_parquet(10, 20,
EnabledStatistics::Chunk).await;
+ let bytes = write_multi_row_group_parquet(10, 20,
EnabledStatistics::Chunk, false).await;
let metadata = load_metadata_with_page_index(&bytes, false);
assert_eq!(metadata.row_groups().len(), 2);
assert!(metadata.column_index().is_none());
@@ -3759,14 +4026,14 @@ mod tests {
let fields = vec![int_field("id"), int_field("value")];
let predicates = vec![id_leaf(PredicateOperator::In,
vec![Datum::Int(100)])];
- let bytes = write_multi_row_group_parquet(10, 10,
EnabledStatistics::None).await;
+ let bytes = write_multi_row_group_parquet(10, 10,
EnabledStatistics::None, false).await;
let metadata = load_metadata_with_page_index(&bytes, false);
let selection =
super::build_predicate_row_selection(metadata.row_groups(),
&predicates, &fields)
.unwrap();
assert!(selection.is_none(), "missing stats must fail open");
- let bytes = write_multi_row_group_parquet(10, 10,
EnabledStatistics::Chunk).await;
+ let bytes = write_multi_row_group_parquet(10, 10,
EnabledStatistics::Chunk, false).await;
let metadata = load_metadata_with_page_index(&bytes, false);
let mut damaged_row_group = metadata.row_groups()[0].clone();
let damaged_id_column = damaged_row_group
@@ -3831,6 +4098,78 @@ mod tests {
buf
}
+ #[tokio::test]
+ async fn
test_parquet_granules_prefers_pages_and_falls_back_to_row_groups() {
+ let bytes = Bytes::from(write_multi_page_parquet(10, 80).await);
+ let (granules, page_level) = parquet_granules(
+ Box::new(TrackingFileRead::new(bytes.clone())),
+ bytes.len() as u64,
+ "value",
+ true,
+ )
+ .await
+ .unwrap();
+ assert!(page_level);
+ assert!(granules.len() > 1);
+ assert_eq!(
+ granules
+ .iter()
+ .map(|granule| granule.row_count)
+ .sum::<i64>(),
+ 80
+ );
+ assert!(granules
+ .iter()
+ .all(|granule| !granule.byte_ranges.is_empty()));
+
+ let bytes = Bytes::from(
+ write_multi_row_group_parquet(10, 30, EnabledStatistics::Chunk,
true).await,
+ );
+ let (granules, page_level) = parquet_granules(
+ Box::new(TrackingFileRead::new(bytes.clone())),
+ bytes.len() as u64,
+ "value",
+ true,
+ )
+ .await
+ .unwrap();
+ assert!(!page_level);
+ assert_eq!(granules.len(), 3);
+ assert_eq!(
+ granules
+ .iter()
+ .map(|granule| granule.row_count)
+ .sum::<i64>(),
+ 30
+ );
+ }
+
+ #[tokio::test]
+ async fn test_parquet_granules_uses_row_groups_when_page_index_disabled() {
+ let bytes = Bytes::from(write_multi_page_parquet(10, 80).await);
+ let metadata = load_metadata_with_page_index(&bytes, true);
+ assert!(metadata.offset_index().is_some());
+ let (granules, page_level) = parquet_granules(
+ Box::new(TrackingFileRead::new(bytes.clone())),
+ bytes.len() as u64,
+ "value",
+ false,
+ )
+ .await
+ .unwrap();
+
+ assert!(!page_level);
+ assert_eq!(granules.len(), 1);
+ assert_eq!(granules[0].first_row, 0);
+ assert_eq!(granules[0].row_count, 80);
+ // Charge the complete projected column chunk, not individual selected
pages.
+ let (start, length) = metadata.row_group(0).column(1).byte_range();
+ assert_eq!(
+ granules[0].byte_ranges,
+ std::iter::once(start..start + length).collect::<Vec<_>>()
+ );
+ }
+
async fn write_page_pruning_io_parquet() -> Vec<u8> {
const ROWS: i32 = 1024;
let schema = Arc::new(ArrowSchema::new(vec![
@@ -3928,6 +4267,26 @@ mod tests {
struct TrackingFileRead {
data: Bytes,
ranges: Arc<std::sync::Mutex<Vec<std::ops::Range<u64>>>>,
+ resident_bytes: Arc<AtomicUsize>,
+ peak_resident_bytes: Arc<AtomicUsize>,
+ }
+
+ struct TrackedReadBuffer {
+ data: Box<[u8]>,
+ resident_bytes: Arc<AtomicUsize>,
+ }
+
+ impl AsRef<[u8]> for TrackedReadBuffer {
+ fn as_ref(&self) -> &[u8] {
+ &self.data
+ }
+ }
+
+ impl Drop for TrackedReadBuffer {
+ fn drop(&mut self) {
+ self.resident_bytes
+ .fetch_sub(self.data.len(), AtomicOrdering::SeqCst);
+ }
}
impl TrackingFileRead {
@@ -3935,6 +4294,8 @@ mod tests {
Self {
data,
ranges: Arc::new(std::sync::Mutex::new(Vec::new())),
+ resident_bytes: Arc::new(AtomicUsize::new(0)),
+ peak_resident_bytes: Arc::new(AtomicUsize::new(0)),
}
}
@@ -3956,7 +4317,135 @@ mod tests {
impl crate::io::FileRead for TrackingFileRead {
async fn read(&self, range: std::ops::Range<u64>) ->
crate::Result<Bytes> {
self.ranges.lock().unwrap().push(range.clone());
- Ok(self.data.slice(range.start as usize..range.end as usize))
+ // Count each source allocation until its last slice is dropped,
not slice lengths.
+ let data = self.data[range.start as usize..range.end as usize]
+ .to_vec()
+ .into_boxed_slice();
+ let current = self
+ .resident_bytes
+ .fetch_add(data.len(), AtomicOrdering::SeqCst)
+ + data.len();
+ self.peak_resident_bytes
+ .fetch_max(current, AtomicOrdering::SeqCst);
+ Ok(Bytes::from_owner(TrackedReadBuffer {
+ data,
+ resident_bytes: Arc::clone(&self.resident_bytes),
+ }))
+ }
+ }
+
+ #[tokio::test]
+ async fn test_sparse_read_buffer_owners_and_cancellation() {
+ use crate::io::FileRead;
+ use rand::{RngCore, SeedableRng};
+
+ const MIB: usize = 1024 * 1024;
+ const GROUP_ROWS: usize = 2 * MIB;
+ const PAGE_ROWS: usize = 64 * 1024;
+ const BUDGET: usize = 20 * MIB;
+
+ let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+ "id",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let props = parquet::file::properties::WriterProperties::builder()
+ .set_max_row_group_row_count(Some(GROUP_ROWS))
+ .set_data_page_size_limit(usize::MAX)
+ .set_data_page_row_count_limit(PAGE_ROWS)
+ .set_write_batch_size(PAGE_ROWS)
+ .set_dictionary_enabled(false)
+
.set_compression(parquet::basic::Compression::ZSTD(Default::default()))
+ .build();
+ let mut data = Vec::new();
+ let mut writer =
+ AsyncArrowWriter::try_new(&mut data, Arc::clone(&schema),
Some(props)).unwrap();
+ let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+ for _ in 0..4 {
+ let values =
Int32Array::from_iter_values((0..GROUP_ROWS).map(|row| {
+ if (row / PAGE_ROWS).is_multiple_of(2) {
+ 0
+ } else {
+ rng.next_u32() as i32
+ }
+ }));
+ writer
+ .write(&RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(values)]).unwrap())
+ .await
+ .unwrap();
+ }
+ writer.close().await.unwrap();
+ let metadata = load_metadata_with_page_index(&data, true);
+ assert_eq!(metadata.num_row_groups(), 4);
+ let pages = metadata.offset_index().unwrap()[0][0].page_locations();
+ assert!(
+ pages[1].compressed_page_size > pages[0].compressed_page_size *
100,
+ "adjacent equal-row-count pages must have very different
compression ratios"
+ );
+ let projection = super::ProjectionMask::all();
+ let projected =
super::projected_row_group_bytes(&metadata.row_groups()[0], &projection);
+ assert!(projected > 8 * MIB as u64 && projected < 9 * MIB as u64);
+
+ let data = Bytes::from(data);
+ let tracker = TrackingFileRead::new(data.clone());
+ let buffer = tracker.read(0..1024).await.unwrap();
+ let slice = buffer.slice(0..1);
+ drop(buffer);
+ assert_eq!(tracker.resident_bytes.load(AtomicOrdering::SeqCst), 1024);
+ drop(slice);
+ assert_eq!(tracker.resident_bytes.load(AtomicOrdering::SeqCst), 0);
+
+ let ranges = (0..4)
+ .flat_map(|group| {
+ (1..GROUP_ROWS / PAGE_ROWS).step_by(2).map(move |page| {
+ let start = (group * GROUP_ROWS + page * PAGE_ROWS) as i64;
+ RowRange::new(start, start + 255)
+ })
+ })
+ .collect::<Vec<_>>();
+ for cancel in [false, true] {
+ let tracker = TrackingFileRead::new(data.clone());
+ let budget = Arc::new(ReadBudget::new(8, BUDGET as u64).unwrap());
+ budget.enable_diagnostics();
+ let mut stream =
ParquetFormatReader::with_read_budget(Arc::clone(&budget))
+ .read_batch_stream(
+ Box::new(tracker.clone()),
+ data.len() as u64,
+ &[int_field("id")],
+ None,
+ Some(128),
+ Some(ranges.clone()),
+ )
+ .await
+ .unwrap();
+ let mut rows =
stream.try_next().await.unwrap().unwrap().num_rows();
+ if !cancel {
+ while let Some(batch) = stream.try_next().await.unwrap() {
+ rows += batch.num_rows();
+ }
+ assert_eq!(rows as i64,
ranges.iter().map(RowRange::count).sum::<i64>());
+ }
+ drop(stream);
+ tokio::time::timeout(Duration::from_secs(2), async {
+ while budget.diagnostics().current_inflight != 0
+ || tracker.resident_bytes.load(AtomicOrdering::SeqCst) != 0
+ {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("all buffer owners and row-group permits must be
released");
+ let peak =
tracker.peak_resident_bytes.load(AtomicOrdering::SeqCst);
+ assert!(
+ peak > MIB && peak <= BUDGET,
+ "resident source buffers: {peak}"
+ );
+ assert_eq!(budget.diagnostics().peak_inflight, 2);
+ let _permit =
+ tokio::time::timeout(Duration::from_secs(1),
budget.acquire(BUDGET as u64))
+ .await
+ .unwrap()
+ .unwrap();
}
}
diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs
b/crates/paimon/src/table/vindex_index_build_builder.rs
index 9e06ff9a..6b255713 100644
--- a/crates/paimon/src/table/vindex_index_build_builder.rs
+++ b/crates/paimon/src/table/vindex_index_build_builder.rs
@@ -16,6 +16,7 @@
// under the License.
mod extraction;
+mod pipeline;
mod planning;
mod timing;
mod validation;
diff --git a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs
b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs
index cf7e910e..37a8bb51 100644
--- a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs
+++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs
@@ -23,6 +23,16 @@ use crate::{Error, Result};
use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array,
ListArray, RecordBatch};
pub(super) fn data_split_for_shard(shard: &VindexIndexShard) ->
Result<DataSplit> {
+ data_split_for_shard_ranges(
+ shard,
+ vec![RowRange::new(shard.row_range_start, shard.row_range_end)],
+ )
+}
+
+pub(super) fn data_split_for_shard_ranges(
+ shard: &VindexIndexShard,
+ row_ranges: Vec<RowRange>,
+) -> Result<DataSplit> {
DataSplitBuilder::new()
.with_snapshot(shard.snapshot_id)
.with_partition(shard.partition.clone())
@@ -30,24 +40,97 @@ pub(super) fn data_split_for_shard(shard:
&VindexIndexShard) -> Result<DataSplit
.with_bucket_path(shard.bucket_path.clone())
.with_total_buckets(shard.total_buckets)
.with_data_files(shard.files.clone())
- .with_row_ranges(vec![RowRange::new(
- shard.row_range_start,
- shard.row_range_end,
- )])
+ .with_row_ranges(row_ranges)
.build()
}
pub(super) struct ValidatedVectorBatch<'a> {
pub(super) values: &'a [f32],
pub(super) bytes: &'a [u8],
+ pub(super) row_ids: &'a [i64],
pub(super) row_count: usize,
}
+pub(super) fn extract_vector_batch<'a>(
+ batch: &'a RecordBatch,
+ index_column: &str,
+ dimension: usize,
+) -> Result<ValidatedVectorBatch<'a>> {
+ validate_vector_batch_with(batch, index_column, dimension, |_| Ok(()))
+}
+
pub(super) fn validate_vector_batch<'a>(
batch: &'a RecordBatch,
index_column: &str,
dimension: usize,
expected_row_id: &mut i64,
+) -> Result<ValidatedVectorBatch<'a>> {
+ validate_vector_batch_with(batch, index_column, dimension, |row_id| {
+ if row_id != *expected_row_id {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "vindex vector extraction expected _ROW_ID {}, got {}",
+ expected_row_id, row_id
+ ),
+ source: None,
+ });
+ }
+ *expected_row_id = expected_row_id
+ .checked_add(1)
+ .ok_or_else(|| Error::DataInvalid {
+ message: "vindex expected row id overflows i64".to_string(),
+ source: None,
+ })?;
+ Ok(())
+ })
+}
+
+pub(super) fn validate_vector_batch_ranges<'a>(
+ batch: &'a RecordBatch,
+ index_column: &str,
+ dimension: usize,
+ ranges: &[RowRange],
+ range_index: &mut usize,
+ expected_row_id: &mut i64,
+) -> Result<ValidatedVectorBatch<'a>> {
+ validate_vector_batch_with(batch, index_column, dimension, |row_id| {
+ let range = ranges.get(*range_index).ok_or_else(|| Error::DataInvalid {
+ message: format!("vindex vector extraction got unexpected _ROW_ID
{row_id}"),
+ source: None,
+ })?;
+ if row_id != *expected_row_id {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "vindex vector extraction expected _ROW_ID {}, got {}",
+ expected_row_id, row_id
+ ),
+ source: None,
+ });
+ }
+ if row_id == range.to() {
+ *range_index += 1;
+ *expected_row_id = match ranges.get(*range_index) {
+ Some(next) => next.from(),
+ None => row_id.checked_add(1).ok_or_else(|| Error::DataInvalid
{
+ message: "vindex expected row id overflows
i64".to_string(),
+ source: None,
+ })?,
+ };
+ } else {
+ *expected_row_id = row_id.checked_add(1).ok_or_else(||
Error::DataInvalid {
+ message: "vindex expected row id overflows i64".to_string(),
+ source: None,
+ })?;
+ }
+ Ok(())
+ })
+}
+
+fn validate_vector_batch_with<'a>(
+ batch: &'a RecordBatch,
+ index_column: &str,
+ dimension: usize,
+ mut validate_row_id: impl FnMut(i64) -> Result<()>,
) -> Result<ValidatedVectorBatch<'a>> {
let vector_index = batch
.schema()
@@ -163,21 +246,7 @@ pub(super) fn validate_vector_batch<'a>(
});
}
for row_id in row_ids.values() {
- if *row_id != *expected_row_id {
- return Err(Error::DataInvalid {
- message: format!(
- "vindex vector extraction expected _ROW_ID {}, got {}",
- expected_row_id, row_id
- ),
- source: None,
- });
- }
- *expected_row_id = expected_row_id
- .checked_add(1)
- .ok_or_else(|| Error::DataInvalid {
- message: "vindex expected row id overflows i64".to_string(),
- source: None,
- })?;
+ validate_row_id(*row_id)?;
}
let byte_start = checked_vector_bytes(start, 1)?;
@@ -185,6 +254,7 @@ pub(super) fn validate_vector_batch<'a>(
Ok(ValidatedVectorBatch {
values: &values.values()[start..end],
bytes: &values.values().inner().as_slice()[byte_start..byte_end],
+ row_ids: row_ids.values(),
row_count: batch.num_rows(),
})
}
diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs
b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs
new file mode 100644
index 00000000..ebd2b8b3
--- /dev/null
+++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs
@@ -0,0 +1,1670 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::extraction::{
+ data_split_for_shard_ranges, extract_vector_batch,
validate_vector_batch_ranges,
+};
+use super::planning::VindexIndexShard;
+use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming};
+use super::validation::{
+ checked_row_count, checked_training_sample_index,
checked_training_vector_count,
+ checked_vector_bytes,
+};
+use super::writer::BuiltIndexFile;
+use super::VindexIndexBuildBuilder;
+use crate::arrow::format::parquet::{
+ coalesced_parquet_range_bytes, parquet_granules, ParquetGranule,
+};
+use crate::spec::ROW_ID_FIELD_NAME;
+use crate::table::data_file_reader::DataFileReadTiming;
+use crate::table::table_read::configured_parquet_read_budget;
+use crate::table::{merge_row_ranges, ArrowRecordBatchStream, RowRange};
+use crate::vindex::VindexVectorIndexOptions;
+use crate::{Error, Result};
+use arrow_array::RecordBatch;
+use arrow_buffer::MutableBuffer;
+use futures::{StreamExt, TryStreamExt};
+use paimon_vindex_core::autotune::default_training_vector_count;
+use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexTraining,
VectorIndexWriter};
+use std::collections::{HashMap, HashSet};
+use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
+use std::ops::Range;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tokio::sync::{mpsc, oneshot};
+use tokio::task::JoinHandle;
+
+const MIN_STRATA: usize = 256;
+const ROWS_PER_STRATUM: usize = 128;
+const FIRST_BYTES_NUMERATOR: u64 = 3;
+const FIRST_BYTES_DENOMINATOR: u64 = 10;
+const QUEUE_CAPACITY: usize = 2;
+const BUFFER_BYTES: usize = 8 * 1024 * 1024;
+const REPLAY_TARGET_BYTES: usize = 32 * 1024 * 1024;
+
+#[derive(Clone, Debug)]
+struct Granule {
+ range: RowRange,
+ file_index: usize,
+ byte_ranges: Vec<Range<u64>>,
+}
+
+#[derive(Debug)]
+pub(super) struct GranulePlan {
+ pub(super) first: Vec<RowRange>,
+ pub(super) rest: Vec<RowRange>,
+ first_rows: usize,
+ // None preserves full-spill's global stride sampling when first covers
the shard.
+ training: Option<Vec<i64>>,
+}
+
+struct SpillRecord {
+ ids: Vec<i64>,
+ bytes: Vec<u8>,
+}
+
+enum AddItem {
+ Batch(RecordBatch, Vec<i64>),
+ Spilled(Vec<i64>, MutableBuffer),
+}
+
+type SpillTask = oneshot::Receiver<std::io::Result<(std::fs::File, u64,
Duration)>>;
+type ConsumerTask = oneshot::Receiver<Result<(VectorIndexWriter, usize, usize,
Duration)>>;
+type TrainingTask = JoinHandle<std::io::Result<(VectorIndexTraining,
Duration)>>;
+type ReplayTask = JoinHandle<std::io::Result<(usize, Duration)>>;
+
+fn spawn_channel_worker<T, F>(name: &'static str, worker: F) ->
Result<oneshot::Receiver<T>>
+where
+ T: Send + 'static,
+ F: FnOnce() -> T + Send + 'static,
+{
+ let (sender, receiver) = oneshot::channel();
+ std::thread::Builder::new()
+ .name(name.to_string())
+ .spawn(move || {
+ let _ = sender.send(worker());
+ })
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to spawn {name} worker: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ Ok(receiver)
+}
+
+struct SpillWriter {
+ sender: mpsc::Sender<SpillRecord>,
+ task: SpillTask,
+}
+
+impl SpillWriter {
+ async fn finish(self) -> Result<(std::fs::File, u64, Duration)> {
+ drop(self.sender);
+ join_spill(self.task).await
+ }
+}
+
+struct LivePipeline {
+ sender: mpsc::Sender<AddItem>,
+ consumer: ConsumerTask,
+ replay: ReplayTask,
+ spill_bytes: u64,
+ spill_write: Duration,
+}
+
+fn spawn_spill_writer(timing_enabled: bool) -> Result<SpillWriter> {
+ let file = tempfile::tempfile().map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to create temporary vindex vector file: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ let (sender, mut receiver) = mpsc::channel::<SpillRecord>(QUEUE_CAPACITY);
+ let task = spawn_channel_worker("paimon-vindex-spill", move || ->
std::io::Result<_> {
+ let mut writer = BufWriter::with_capacity(BUFFER_BYTES, file);
+ let mut spill_bytes = 0u64;
+ let mut spill_write = Duration::ZERO;
+ while let Some(record) = receiver.blocking_recv() {
+ let write_start = timing_enabled.then(Instant::now);
+ let count = record.ids.len() as u64;
+ writer.write_all(&count.to_le_bytes())?;
+ for id in record.ids {
+ writer.write_all(&id.to_le_bytes())?;
+ }
+ writer.write_all(&record.bytes)?;
+ spill_bytes = spill_bytes
+ .saturating_add(8)
+ .saturating_add(count.saturating_mul(8))
+ .saturating_add(record.bytes.len() as u64);
+ if let Some(start) = write_start {
+ spill_write = spill_write.saturating_add(start.elapsed());
+ }
+ }
+ let write_start = timing_enabled.then(Instant::now);
+ writer.flush()?;
+ let file = writer.into_inner().map_err(|e| e.into_error())?;
+ if let Some(start) = write_start {
+ spill_write = spill_write.saturating_add(start.elapsed());
+ }
+ Ok((file, spill_bytes, spill_write))
+ })?;
+ Ok(SpillWriter { sender, task })
+}
+
+async fn join_spill(task: SpillTask) -> Result<(std::fs::File, u64, Duration)>
{
+ task.await
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("vindex spill task failed: {e}"),
+ source: None,
+ })?
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to spill vindex vectors: {e}"),
+ source: Some(Box::new(e)),
+ })
+}
+
+fn spawn_add_consumer(
+ writer: VectorIndexWriter,
+ mut receiver: mpsc::Receiver<AddItem>,
+ index_column: String,
+ dimension: usize,
+ timing_enabled: bool,
+) -> Result<ConsumerTask> {
+ spawn_channel_worker("paimon-vindex-add", move || -> Result<_> {
+ let mut writer = writer;
+ let mut rows_added = 0usize;
+ let mut replay_rows = 0usize;
+ let mut index_add = Duration::ZERO;
+ while let Some(item) = receiver.blocking_recv() {
+ let add_start = timing_enabled.then(Instant::now);
+ match item {
+ AddItem::Batch(batch, ids) => {
+ let vectors = extract_vector_batch(&batch, &index_column,
dimension)?;
+ if ids.len() != vectors.row_count {
+ return Err(Error::DataInvalid {
+ message: "vindex add batch id count
mismatch".to_string(),
+ source: None,
+ });
+ }
+ writer
+ .add_vectors(&ids, vectors.values, vectors.row_count)
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to add vectors to vindex
index: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ rows_added += vectors.row_count;
+ }
+ AddItem::Spilled(ids, buffer) => {
+ let values = buffer.typed_data::<f32>();
+ if values.len() != ids.len() * dimension {
+ return Err(Error::DataInvalid {
+ message: "vindex spilled vector length
mismatch".to_string(),
+ source: None,
+ });
+ }
+ writer.add_vectors(&ids, values, ids.len()).map_err(|e| {
+ Error::UnexpectedError {
+ message: format!("Failed to add spilled vectors to
vindex index: {e}"),
+ source: Some(Box::new(e)),
+ }
+ })?;
+ rows_added += ids.len();
+ replay_rows += ids.len();
+ }
+ }
+ if let Some(start) = add_start {
+ index_add = index_add.saturating_add(start.elapsed());
+ }
+ }
+ Ok((writer, rows_added, replay_rows, index_add))
+ })
+}
+
+async fn join_consumer(task: ConsumerTask) -> Result<(VectorIndexWriter,
usize, usize, Duration)> {
+ task.await.map_err(|e| Error::UnexpectedError {
+ message: format!("vindex add task failed: {e}"),
+ source: None,
+ })?
+}
+
+async fn join_training(task: TrainingTask) -> Result<(VectorIndexTraining,
Duration)> {
+ task.await
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("vindex training task failed: {e}"),
+ source: None,
+ })?
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to train vindex index: {e}"),
+ source: Some(Box::new(e)),
+ })
+}
+
+async fn start_live_pipeline(
+ training: TrainingTask,
+ spill: SpillWriter,
+ index_column: String,
+ dimension: usize,
+ timing_enabled: bool,
+) -> Result<(LivePipeline, Duration)> {
+ let trained = join_training(training).await;
+ let spilled = spill.finish().await;
+ let (trained, train_finish) = trained?;
+ let (file, spill_bytes, spill_write) = spilled?;
+ let (sender, receiver) = mpsc::channel(QUEUE_CAPACITY);
+ let consumer = spawn_add_consumer(
+ VectorIndexWriter::new(trained),
+ receiver,
+ index_column,
+ dimension,
+ timing_enabled,
+ )?;
+ let replay = spawn_replay(file, sender.clone(), dimension, timing_enabled);
+ Ok((
+ LivePipeline {
+ sender,
+ consumer,
+ replay,
+ spill_bytes,
+ spill_write,
+ },
+ train_finish,
+ ))
+}
+
+async fn finish_live_pipeline(
+ pipeline: LivePipeline,
+) -> (
+ Result<(VectorIndexWriter, usize, usize, Duration)>,
+ Result<(usize, Duration)>,
+ u64,
+ Duration,
+) {
+ let replay = pipeline
+ .replay
+ .await
+ .map_err(|e| Error::UnexpectedError {
+ message: format!("vindex replay task failed: {e}"),
+ source: None,
+ })
+ .and_then(|result| {
+ result.map_err(|e| Error::UnexpectedError {
+ message: format!("Failed to replay spilled vindex vectors:
{e}"),
+ source: Some(Box::new(e)),
+ })
+ });
+ drop(pipeline.sender);
+ let consumer = join_consumer(pipeline.consumer).await;
+ (consumer, replay, pipeline.spill_bytes, pipeline.spill_write)
+}
+
+fn spawn_replay(
+ mut file: std::fs::File,
+ sender: mpsc::Sender<AddItem>,
+ dimension: usize,
+ timing_enabled: bool,
+) -> ReplayTask {
+ tokio::task::spawn_blocking(move || -> std::io::Result<(usize, Duration)> {
+ let mut spill_read = Duration::ZERO;
+ let read_start = timing_enabled.then(Instant::now);
+ file.seek(SeekFrom::Start(0))?;
+ if let Some(start) = read_start {
+ spill_read = spill_read.saturating_add(start.elapsed());
+ }
+ let mut reader = BufReader::with_capacity(BUFFER_BYTES, file);
+ let mut rows = 0usize;
+ let mut ids = Vec::new();
+ let mut vectors = MutableBuffer::new(REPLAY_TARGET_BYTES);
+ loop {
+ let read_start = timing_enabled.then(Instant::now);
+ let mut header = [0u8; 8];
+ let read = reader.read(&mut header)?;
+ if read == 0 {
+ if let Some(start) = read_start {
+ spill_read = spill_read.saturating_add(start.elapsed());
+ }
+ break;
+ }
+ reader.read_exact(&mut header[read..])?;
+ let count =
usize::try_from(u64::from_le_bytes(header)).map_err(|_| {
+ std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid
spill row count")
+ })?;
+ let id_bytes_len = count.checked_mul(8).ok_or_else(|| {
+ std::io::Error::new(std::io::ErrorKind::InvalidData, "spill id
length overflow")
+ })?;
+ let mut id_bytes = vec![0u8; id_bytes_len];
+ reader.read_exact(&mut id_bytes)?;
+ let (id_chunks, remainder) = id_bytes.as_chunks::<8>();
+ debug_assert!(remainder.is_empty());
+ ids.extend(id_chunks.iter().map(|bytes|
i64::from_le_bytes(*bytes)));
+ let vector_bytes = count
+ .checked_mul(dimension)
+ .and_then(|value| value.checked_mul(4))
+ .ok_or_else(|| {
+ std::io::Error::new(
+ std::io::ErrorKind::InvalidData,
+ "spill vector length overflow",
+ )
+ })?;
+ let offset = vectors.len();
+ vectors.resize(offset + vector_bytes, 0);
+ reader.read_exact(&mut vectors.as_slice_mut()[offset..])?;
+ if let Some(start) = read_start {
+ spill_read = spill_read.saturating_add(start.elapsed());
+ }
+ rows += count;
+ if vectors.len() >= REPLAY_TARGET_BYTES {
+ let item = AddItem::Spilled(
+ std::mem::take(&mut ids),
+ std::mem::replace(&mut vectors,
MutableBuffer::new(REPLAY_TARGET_BYTES)),
+ );
+ if sender.blocking_send(item).is_err() {
+ return Ok((rows, spill_read));
+ }
+ }
+ }
+ if !ids.is_empty() {
+ let _ = sender.blocking_send(AddItem::Spilled(ids, vectors));
+ }
+ Ok((rows, spill_read))
+ })
+}
+
+fn splitmix64(state: &mut u64) -> u64 {
+ *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+ let mut z = *state;
+ z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+ z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+ z ^ (z >> 31)
+}
+
+fn shard_seed(shard: &VindexIndexShard) -> u64 {
+ let mut seed = 0u64;
+ let mut absorb = |value: u64| {
+ let mut state = seed ^ value;
+ seed = splitmix64(&mut state);
+ };
+ absorb(shard.snapshot_id as u64);
+ absorb(shard.source_bucket as u64);
+ absorb(shard.row_range_start as u64);
+ absorb(shard.row_range_end as u64);
+ absorb(shard.partition_bytes.len() as u64);
+ for chunk in shard.partition_bytes.chunks(8) {
+ let mut word = [0u8; 8];
+ word[..chunk.len()].copy_from_slice(chunk);
+ absorb(u64::from_le_bytes(word));
+ }
+ seed
+}
+
+fn random_below(state: &mut u64, bound: usize) -> usize {
+ let bound = bound as u64;
+ let threshold = bound.wrapping_neg() % bound;
+ loop {
+ let value = splitmix64(state);
+ if value >= threshold {
+ return (value % bound) as usize;
+ }
+ }
+}
+
+fn pick_indices(total: usize, count: usize, state: &mut u64) -> impl
Iterator<Item = usize> + '_ {
+ (0..count).map(move |index| {
+ let start = (index as u128 * total as u128 / count as u128) as usize;
+ let end = ((index + 1) as u128 * total as u128 / count as u128) as
usize;
+ start + random_below(state, end - start)
+ })
+}
+
+// Boundaries depend only on the layout, never on which granules were sampled.
+// Keep original indices: a unit can cross files, but physical reads cannot.
+fn sampling_units(granules: &[Granule]) -> Vec<Range<usize>> {
+ let mut units: Vec<Range<usize>> = Vec::new();
+ let mut start = 0;
+ for end in 1..=granules.len() {
+ let rows = granules[end - 1].range.to() - granules[start].range.from()
+ 1;
+ if rows as usize >= 2 * ROWS_PER_STRATUM {
+ units.push(start..end);
+ start = end;
+ }
+ }
+ if start < granules.len() {
+ if let Some(last) = units.last_mut() {
+ last.end = granules.len();
+ } else {
+ units.push(0..granules.len());
+ }
+ }
+ units
+}
+
+fn unit_training_offsets(rows: usize, quota: usize, offset: usize) -> impl
Iterator<Item = usize> {
+ (0..quota).map(move |index| {
+ ((index as u128 * rows as u128 + offset as u128) / quota as u128) as
usize
+ })
+}
+
+fn granule_bytes(granules: &[Granule], selected: Option<&HashSet<usize>>) ->
u64 {
+ let mut files = HashMap::<usize, Vec<Range<u64>>>::new();
+ for (index, granule) in granules.iter().enumerate() {
+ if selected.is_none_or(|selected| selected.contains(&index)) {
+ files
+ .entry(granule.file_index)
+ .or_default()
+ .extend(granule.byte_ranges.iter().cloned());
+ }
+ }
+ files
+ .values()
+ .map(|ranges| coalesced_parquet_range_bytes(ranges))
+ .sum()
+}
+
+fn select_first(
+ granules: &[Granule],
+ eligible: usize,
+ retained: usize,
+ mut seed: u64,
+) -> Result<GranulePlan> {
+ // Callers have validated that granules partition one nonempty shard.
+ let range = RowRange::new(
+ granules[0].range.from(),
+ granules.last().unwrap().range.to(),
+ );
+ let rows = checked_row_count(range.from(), range.to())? as usize;
+ let whole_shard = || GranulePlan {
+ first: vec![range.clone()],
+ rest: Vec::new(),
+ first_rows: rows,
+ training: None,
+ };
+ if rows < 2 * ROWS_PER_STRATUM {
+ return Ok(whole_shard());
+ }
+ let units = sampling_units(granules);
+ let strata = retained
+ .div_ceil(ROWS_PER_STRATUM)
+ .max(MIN_STRATA)
+ .min(rows);
+ let candidates = strata
+ .checked_mul(ROWS_PER_STRATUM)
+ .ok_or_else(|| Error::DataInvalid {
+ message: "vindex training candidate count overflows
usize".to_string(),
+ source: None,
+ })?;
+ let mut quotas = vec![0usize; units.len()];
+ for row in pick_indices(rows, strata, &mut seed) {
+ let row_id = range.from() + row as i64;
+ let unit = units.partition_point(|unit| granules[unit.end -
1].range.to() < row_id);
+ quotas[unit] += ROWS_PER_STRATUM;
+ }
+ let mut selected = HashSet::new();
+ for (unit, quota) in units.iter().zip("as) {
+ let unit_rows = granules[unit.end - 1].range.to() -
granules[unit.start].range.from() + 1;
+ if *quota > unit_rows as usize {
+ return Ok(whole_shard());
+ }
+ if *quota > 0 {
+ selected.extend(unit.clone());
+ }
+ }
+ let first_bytes = granule_bytes(granules, Some(&selected));
+ let total_bytes = granule_bytes(granules, None);
+ if selected.len() == granules.len()
+ || selected.len() < MIN_STRATA.min(granules.len())
+ || total_bytes == 0
+ || first_bytes.saturating_mul(FIRST_BYTES_DENOMINATOR)
+ > total_bytes.saturating_mul(FIRST_BYTES_NUMERATOR)
+ {
+ return Ok(whole_shard());
+ }
+
+ let mut training = Vec::with_capacity(candidates);
+ for (unit, quota) in units.iter().zip(quotas) {
+ if quota == 0 {
+ continue;
+ }
+ let start = granules[unit.start].range.from();
+ let unit_rows = (granules[unit.end - 1].range.to() - start + 1) as
usize;
+ let offset = random_below(&mut seed, unit_rows);
+ training
+ .extend(unit_training_offsets(unit_rows, quota, offset).map(|row|
start + row as i64));
+ }
+ debug_assert_eq!(training.len(), candidates);
+ if eligible < candidates {
+ for index in 0..eligible {
+ let other = index + random_below(&mut seed, candidates - index);
+ training.swap(index, other);
+ }
+ training.truncate(eligible);
+ training.sort_unstable();
+ }
+ let mut first = Vec::with_capacity(selected.len());
+ let mut rest = Vec::with_capacity(granules.len() - selected.len());
+ let mut first_rows = 0;
+ for (index, granule) in granules.iter().enumerate() {
+ if selected.contains(&index) {
+ first_rows += granule.range.count() as usize;
+ first.push(granule.range.clone());
+ } else {
+ rest.push(granule.range.clone());
+ }
+ }
+ Ok(GranulePlan {
+ first: merge_row_ranges(first),
+ rest: merge_row_ranges(rest),
+ first_rows,
+ training: Some(training),
+ })
+}
+
+fn append_shard_granules(
+ granules: &mut Vec<Granule>,
+ file_index: usize,
+ file_start: i64,
+ shard_range: &RowRange,
+ file_granules: Vec<ParquetGranule>,
+) -> Result<()> {
+ for granule in file_granules {
+ let from = file_start
+ .checked_add(granule.first_row)
+ .ok_or_else(|| Error::DataInvalid {
+ message: "vindex granule row id overflows i64".to_string(),
+ source: None,
+ })?;
+ let to = from
+ .checked_add(granule.row_count - 1)
+ .ok_or_else(|| Error::DataInvalid {
+ message: "vindex granule row range overflows i64".to_string(),
+ source: None,
+ })?;
+ if let Some(range) = shard_range.intersect_inclusive(from, to) {
+ granules.push(Granule {
+ range,
+ file_index,
+ byte_ranges: granule.byte_ranges,
+ });
+ }
+ }
+ Ok(())
+}
+
+fn granules_partition_shard(granules: &[Granule], shard_range: &RowRange) ->
bool {
+ granules
+ .first()
+ .is_some_and(|granule| granule.range.from() == shard_range.from())
+ && granules
+ .last()
+ .is_some_and(|granule| granule.range.to() == shard_range.to())
+ && granules
+ .windows(2)
+ .all(|pair| pair[0].range.to().checked_add(1) ==
Some(pair[1].range.from()))
+}
+
+fn local_ids(row_ids: &[i64], start: i64, row_count: usize) ->
Result<Vec<i64>> {
+ let end = start
+ .checked_add(i64::try_from(row_count).map_err(|e| Error::DataInvalid {
+ message: "vindex row count does not fit i64".to_string(),
+ source: Some(Box::new(e)),
+ })?)
+ .ok_or_else(|| Error::DataInvalid {
+ message: "vindex row range overflows i64".to_string(),
+ source: None,
+ })?;
+ row_ids
+ .iter()
+ .map(|row_id| {
+ if *row_id < start || *row_id >= end {
+ Err(Error::DataInvalid {
+ message: format!("vindex row id {row_id} is outside shard
[{start}, {end})"),
+ source: None,
+ })
+ } else {
+ Ok(*row_id - start)
+ }
+ })
+ .collect()
+}
+
+impl<'a> VindexIndexBuildBuilder<'a> {
+ fn open_vector_stream(
+ &self,
+ shard: &VindexIndexShard,
+ ranges: Vec<RowRange>,
+ index_column: &str,
+ read_timing: Option<&Arc<DataFileReadTiming>>,
+ parquet_read_budget: Option<&Arc<crate::arrow::ReadBudget>>,
+ ) -> Result<ArrowRecordBatchStream> {
+ let split = data_split_for_shard_ranges(shard, ranges)?;
+ let mut read_builder = self.table.new_read_builder();
+ read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?;
+ let read = read_builder.new_read()?;
+ let read = match read_timing {
+ Some(timing) =>
read.with_data_file_read_timing(Arc::clone(timing)),
+ None => read,
+ };
+ let read = match parquet_read_budget {
+ Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)),
+ None => read,
+ };
+ read.to_arrow(&[split])
+ }
+
+ pub(super) async fn plan_granules(
+ &self,
+ shard: &VindexIndexShard,
+ index_column: &str,
+ eligible: usize,
+ retained: usize,
+ ) -> Result<GranulePlan> {
+ let shard_range = RowRange::new(shard.row_range_start,
shard.row_range_end);
+ let mut granules = Vec::new();
+ let mut use_whole_shard = false;
+ let mut parquet_files = Vec::new();
+
+ for (file_index, file) in shard.files.iter().enumerate() {
+ if file
+ .write_cols
+ .as_ref()
+ .is_some_and(|columns| !columns.iter().any(|column| column ==
index_column))
+ {
+ continue;
+ }
+ let Some((file_start, file_end)) = file.row_id_range() else {
+ use_whole_shard = true;
+ break;
+ };
+ let Some(range) = shard_range.intersect_inclusive(file_start,
file_end) else {
+ continue;
+ };
+ let file_size = u64::try_from(file.file_size).map_err(|e|
Error::DataInvalid {
+ message: format!(
+ "Invalid data file size for '{}': {}",
+ file.file_name, file.file_size
+ ),
+ source: Some(Box::new(e)),
+ })?;
+ let path = file.data_file_path(&shard.bucket_path);
+ if path.to_ascii_lowercase().ends_with(".parquet") {
+ parquet_files.push((file_index, path, file_size, file_start,
file_end));
+ } else {
+ granules.push(Granule {
+ range,
+ file_index,
+ byte_ranges: std::iter::once(0..file_size).collect(),
+ });
+ }
+ }
+
+ if !use_whole_shard && !parquet_files.is_empty() {
+ let page_index_enabled = self
+ .table
+ .schema()
+ .core_options()
+ .parquet_filter_column_index_enabled()?;
+ let concurrency = self
+ .table
+ .schema()
+ .core_options()
+ .parquet_row_group_parallelism()?
+ .max(1);
+ let file_io = self.table.file_io();
+ let mut results = futures::stream::iter(parquet_files)
+ .map(
+ |(file_index, path, file_size, file_start, file_end)|
async move {
+ let input = file_io.new_input(&path)?;
+ let reader = Box::new(input.reader().await?);
+ let (granules, _) =
+ parquet_granules(reader, file_size, index_column,
page_index_enabled)
+ .await?;
+ Ok::<_, Error>((file_index, file_start, file_end,
granules))
+ },
+ )
+ .buffer_unordered(concurrency);
+ while let Some(result) = results.next().await {
+ match result {
+ Ok((file_index, file_start, file_end, file_granules)) => {
+ let covered = file_granules
+ .iter()
+ .map(|granule| granule.row_count)
+ .sum::<i64>();
+ if covered != file_end - file_start + 1 {
+ use_whole_shard = true;
+ break;
+ }
+ append_shard_granules(
+ &mut granules,
+ file_index,
+ file_start,
+ &shard_range,
+ file_granules,
+ )?;
+ }
+ Err(error) => {
+ log::warn!(
+ "vindex granule metadata read failed; using the
whole shard as the first granule: {error}"
+ );
+ use_whole_shard = true;
+ break;
+ }
+ }
+ }
+ }
+
+ granules.sort_by_key(|granule| granule.range.from());
+ if !use_whole_shard {
+ // Data Evolution may retain multiple physical providers for the
same logical rows.
+ // Let the reader choose the provider instead of counting those
rows twice here.
+ use_whole_shard = !granules_partition_shard(&granules,
&shard_range);
+ }
+ if use_whole_shard || granules.is_empty() {
+ granules = vec![Granule {
+ range: shard_range,
+ file_index: 0,
+ byte_ranges: std::iter::once(
+ 0..shard
+ .files
+ .iter()
+ .map(|file| file.file_size.max(0) as u64)
+ .sum(),
+ )
+ .collect(),
+ }];
+ }
+
+ select_first(&granules, eligible, retained, shard_seed(shard))
+ }
+
+ #[allow(clippy::too_many_arguments)]
+ pub(super) async fn build_index_file_granule(
+ &self,
+ shard: &VindexIndexShard,
+ index_column: &str,
+ dimension: i32,
+ index_field_id: i32,
+ options: &VindexVectorIndexOptions,
+ index_meta: Vec<u8>,
+ ) -> Result<BuiltIndexFile> {
+ let timing_enabled = vector_index_build_timing_enabled();
+ let total_start = timing_enabled.then(Instant::now);
+ let mut source_batch_wait = Duration::ZERO;
+ let mut batch_count = 0usize;
+ let read_timing = timing_enabled.then(||
Arc::new(DataFileReadTiming::default()));
+ let parquet_read_budget = if timing_enabled {
+ let budget = configured_parquet_read_budget(self.table)?;
+ budget.enable_diagnostics();
+ Some(budget)
+ } else {
+ None
+ };
+ let row_count = checked_row_count(shard.row_range_start,
shard.row_range_end)?;
+ let row_count_usize = usize::try_from(row_count).map_err(|e|
Error::DataInvalid {
+ message: format!("Invalid vindex row count: {row_count}"),
+ source: Some(Box::new(e)),
+ })?;
+ let dimension = usize::try_from(dimension).map_err(|e|
Error::DataInvalid {
+ message: format!("Invalid vindex dimension: {dimension}"),
+ source: Some(Box::new(e)),
+ })?;
+ if dimension == 0 {
+ return Err(Error::DataInvalid {
+ message: "vindex vector dimension must be
positive".to_string(),
+ source: None,
+ });
+ }
+ checked_vector_bytes(row_count_usize, dimension)?;
+ let eligible = checked_training_vector_count(row_count_usize,
options.train_sample_ratio)?;
+ let retained =
+ default_training_vector_count(eligible,
options.config.nlist()).unwrap_or(eligible);
+ let plan = self
+ .plan_granules(shard, index_column, eligible, retained)
+ .await?;
+
+ let mut trainer =
+ VectorIndexTrainer::new(options.config.clone()).map_err(|e|
Error::DataInvalid {
+ message: format!("Failed to initialize vindex trainer: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ let mut spill = Some(spawn_spill_writer(timing_enabled)?);
+ let training_rows = plan.training.as_ref().map_or(eligible, Vec::len);
+ let training_buffer_rows = (BUFFER_BYTES / checked_vector_bytes(1,
dimension)?).max(1);
+ let training_buffer_floats = training_buffer_rows * dimension;
+ let mut training_buffer = Vec::with_capacity(training_buffer_floats);
+ let mut next_training_sample = 0usize;
+ let mut first_rows = 0usize;
+ let mut range_index = 0usize;
+ let first_result: Result<()> = async {
+ let mut stream = self.open_vector_stream(
+ shard,
+ plan.first.clone(),
+ index_column,
+ read_timing.as_ref(),
+ parquet_read_budget.as_ref(),
+ )?;
+ let mut expected_row_id = plan.first[0].from();
+ loop {
+ let source_start = timing_enabled.then(Instant::now);
+ let batch = stream.try_next().await?;
+ if let Some(start) = source_start {
+ source_batch_wait =
source_batch_wait.saturating_add(start.elapsed());
+ }
+ let Some(batch) = batch else { break };
+ batch_count += 1;
+ let vectors = validate_vector_batch_ranges(
+ &batch,
+ index_column,
+ dimension,
+ &plan.first,
+ &mut range_index,
+ &mut expected_row_id,
+ )?;
+ let batch_end =
first_rows.checked_add(vectors.row_count).ok_or_else(|| {
+ Error::DataInvalid {
+ message: "vindex first-batch row count overflows
usize".to_string(),
+ source: None,
+ }
+ })?;
+ while next_training_sample < training_rows {
+ let row = if let Some(training) = &plan.training {
+ let sample = training[next_training_sample];
+ match vectors.row_ids.binary_search(&sample) {
+ Ok(row) => row,
+ Err(row) if row == vectors.row_count => break,
+ Err(_) => return Err(Error::DataInvalid {
+ message: format!("Missing vindex training row
{sample}"),
+ source: None,
+ }),
+ }
+ } else {
+ let sample = checked_training_sample_index(
+ next_training_sample,
+ plan.first_rows,
+ training_rows,
+ )?;
+ if sample >= batch_end {
+ break;
+ }
+ sample - first_rows
+ };
+ let offset = row * dimension;
+
training_buffer.extend_from_slice(&vectors.values[offset..offset + dimension]);
+ next_training_sample += 1;
+ if training_buffer.len() == training_buffer_floats {
+ trainer
+ .add_training_vectors_mut(
+ &training_buffer,
+ training_buffer.len() / dimension,
+ )
+ .map_err(|e| Error::DataInvalid {
+ message: format!("Failed to add vindex
training vectors: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ training_buffer.clear();
+ }
+ }
+ let record = SpillRecord {
+ ids: local_ids(vectors.row_ids, shard.row_range_start,
row_count_usize)?,
+ bytes: vectors.bytes.to_vec(),
+ };
+ first_rows = batch_end;
+ if spill.as_ref().unwrap().sender.send(record).await.is_err() {
+ return Err(Error::UnexpectedError {
+ message: "vindex spill writer stopped
unexpectedly".to_string(),
+ source: None,
+ });
+ }
+ }
+ if !training_buffer.is_empty() {
+ trainer
+ .add_training_vectors_mut(&training_buffer,
training_buffer.len() / dimension)
+ .map_err(|e| Error::DataInvalid {
+ message: format!("Failed to add vindex training
vectors: {e}"),
+ source: Some(Box::new(e)),
+ })?;
+ }
+ if first_rows != plan.first_rows
+ || range_index != plan.first.len()
+ || next_training_sample != training_rows
+ {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "vindex first-batch mismatch: rows={first_rows}/{},
ranges={range_index}/{}, training={next_training_sample}/{training_rows}",
+ plan.first_rows,
+ plan.first.len()
+ ),
+ source: None,
+ });
+ }
+ Ok(())
+ }
+ .await;
+ if let Err(error) = first_result {
+ return match spill.take().unwrap().finish().await {
+ Ok(_) => Err(error),
+ Err(spill_error) => Err(spill_error),
+ };
+ }
+
+ let mut training: Option<TrainingTask> =
Some(tokio::task::spawn_blocking(
+ move || -> std::io::Result<_> {
+ let start = timing_enabled.then(Instant::now);
+ let training = trainer.finish()?;
+ Ok((
+ training,
+ start.map_or(Duration::ZERO, |start| start.elapsed()),
+ ))
+ },
+ ));
+ let mut live: Option<LivePipeline> = None;
+ let mut train_finish = Duration::ZERO;
+ let index_column = index_column.to_string();
+
+ macro_rules! go_live {
+ () => {{
+ let (pipeline, duration) = start_live_pipeline(
+ training.take().expect("training task"),
+ spill.take().expect("spill writer"),
+ index_column.clone(),
+ dimension,
+ timing_enabled,
+ )
+ .await?;
+ train_finish = duration;
+ live = Some(pipeline);
+ }};
+ }
+
+ let mut rest_rows = 0usize;
+ let producer_result: Result<()> = async {
+ if !plan.rest.is_empty() {
+ let mut stream = self.open_vector_stream(
+ shard,
+ plan.rest.clone(),
+ &index_column,
+ read_timing.as_ref(),
+ parquet_read_budget.as_ref(),
+ )?;
+ let mut range_index = 0usize;
+ let mut expected_row_id = plan.rest[0].from();
+ loop {
+ let source_start = timing_enabled.then(Instant::now);
+ let batch = stream.try_next().await?;
+ if let Some(start) = source_start {
+ source_batch_wait =
source_batch_wait.saturating_add(start.elapsed());
+ }
+ let Some(batch) = batch else { break };
+ batch_count += 1;
+ if live.is_none() && training.as_ref().is_some_and(|task|
task.is_finished()) {
+ go_live!();
+ }
+ let vectors = validate_vector_batch_ranges(
+ &batch,
+ &index_column,
+ dimension,
+ &plan.rest,
+ &mut range_index,
+ &mut expected_row_id,
+ )?;
+ let ids = local_ids(vectors.row_ids,
shard.row_range_start, row_count_usize)?;
+ rest_rows =
rest_rows.checked_add(vectors.row_count).ok_or_else(|| {
+ Error::DataInvalid {
+ message: "vindex remaining row count overflows
usize".to_string(),
+ source: None,
+ }
+ })?;
+ if let Some(pipeline) = &live {
+ if pipeline
+ .sender
+ .send(AddItem::Batch(batch, ids))
+ .await
+ .is_err()
+ {
+ return Err(Error::UnexpectedError {
+ message: "vindex add consumer stopped
unexpectedly".to_string(),
+ source: None,
+ });
+ }
+ } else {
+ let record = SpillRecord {
+ ids,
+ bytes: vectors.bytes.to_vec(),
+ };
+ if
spill.as_ref().unwrap().sender.send(record).await.is_err() {
+ return Err(Error::UnexpectedError {
+ message: "vindex spill writer stopped
unexpectedly".to_string(),
+ source: None,
+ });
+ }
+ }
+ }
+ if rest_rows != row_count_usize - plan.first_rows
+ || range_index != plan.rest.len()
+ {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "vindex remaining-batch mismatch:
rows={rest_rows}/{}, ranges={range_index}/{}",
+ row_count_usize - plan.first_rows,
+ plan.rest.len()
+ ),
+ source: None,
+ });
+ }
+ }
+ Ok(())
+ }
+ .await;
+
+ if let Err(producer_error) = producer_result {
+ if let Some(pipeline) = live.take() {
+ let (consumer, replay, _, _) =
finish_live_pipeline(pipeline).await;
+ consumer?;
+ replay?;
+ } else {
+ let spill_result = match spill.take() {
+ Some(spill) => spill.finish().await.map(|_| ()),
+ None => Ok(()),
+ };
+ let training_result = match training.take() {
+ Some(training) => join_training(training).await.map(|_|
()),
+ None => Ok(()),
+ };
+ spill_result?;
+ training_result?;
+ }
+ return Err(producer_error);
+ }
+ if live.is_none() {
+ go_live!();
+ }
+
+ let (consumer, replay, granule_spill_bytes, granule_spill_write) =
+ finish_live_pipeline(live.unwrap()).await;
+ let (writer, rows_added, consumer_replay_rows, index_add) = consumer?;
+ let (replay_rows, granule_spill_read) = replay?;
+ if rows_added != row_count_usize || replay_rows !=
consumer_replay_rows {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "vindex pipelined add mismatch:
rows={rows_added}/{row_count_usize},
replay={consumer_replay_rows}/{replay_rows}"
+ ),
+ source: None,
+ });
+ }
+
+ let serialize_upload_start = timing_enabled.then(Instant::now);
+ let meta = self
+ .finish_index_file(writer, shard, index_field_id, index_meta,
row_count)
+ .await?;
+ let serialize_upload =
+ serialize_upload_start.map_or(Duration::ZERO, |start|
start.elapsed());
+ let (oss_read, parquet_decode) = read_timing
+ .as_ref()
+ .map_or((Duration::ZERO, Duration::ZERO), |timing| {
+ (timing.file_read(), timing.parquet_decode())
+ });
+ let (file_schema_open, first_batch_wait, remaining_batch_wait) =
read_timing
+ .as_ref()
+ .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing|
{
+ timing.file_waits()
+ });
+ let parquet_diagnostics = parquet_read_budget
+ .as_ref()
+ .map_or_else(Default::default, |budget| budget.diagnostics());
+ let timing = total_start.map(|start| VectorIndexBuildTiming {
+ total_without_commit: start.elapsed(),
+ source_batch_wait,
+ oss_read,
+ parquet_decode,
+ file_schema_open,
+ first_batch_wait,
+ remaining_batch_wait,
+ parquet_row_group_count: parquet_diagnostics.row_group_count,
+ parquet_projected_bytes_min:
parquet_diagnostics.projected_bytes_min,
+ parquet_projected_bytes_max:
parquet_diagnostics.projected_bytes_max,
+ parquet_projected_bytes_total:
parquet_diagnostics.projected_bytes_total,
+ parquet_peak_inflight_row_groups:
parquet_diagnostics.peak_inflight,
+ raw_temp_write: Duration::ZERO,
+ granule_spill_write,
+ train_finish,
+ raw_temp_reread: Duration::ZERO,
+ granule_spill_read,
+ index_add,
+ serialize_upload,
+ rows: row_count_usize,
+ training_rows_seen: training_rows,
+ training_rows_retained: retained,
+ batch_count,
+ raw_temp_bytes: 0,
+ granule_spill_bytes,
+ index_bytes: meta.file_size as u64,
+ data_file_count: shard.files.len(),
+ file_name: meta.file_name.clone(),
+ });
+ Ok(BuiltIndexFile { meta, timing })
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn granules_with_rows(rows: impl IntoIterator<Item = usize>) ->
Vec<Granule> {
+ let mut start = 0;
+ rows.into_iter()
+ .enumerate()
+ .map(|(file_index, rows)| {
+ let range = RowRange::new(start, start + rows as i64 - 1);
+ start += rows as i64;
+ Granule {
+ range,
+ file_index,
+ byte_ranges: std::iter::once(0..rows as u64).collect(),
+ }
+ })
+ .collect()
+ }
+
+ #[test]
+ fn a_single_large_granule_is_not_excluded_from_training() {
+ for large_index in [0, 2_048, 4_095] {
+ let granules = granules_with_rows((0..4_096).map(|index| {
+ if index == large_index {
+ 524_288
+ } else {
+ 128
+ }
+ }));
+ for seed in 0..32 {
+ let plan = select_first(&granules, 65_536, 65_536,
seed).unwrap();
+ // Row-proportional bytes force whole-shard when the large
granule is hit.
+ assert!(plan.training.is_none(), "index {large_index}, seed
{seed}");
+ assert_eq!(plan.first_rows, 4_095 * 128 + 524_288);
+ assert!(plan.rest.is_empty());
+ }
+ }
+ }
+
+ #[test]
+ fn static_units_merge_short_pages_and_the_trailing_remainder() {
+ for page_rows in [64, 85, 128, 170, 256, 341] {
+ let granules = granules_with_rows(vec![page_rows; 101]);
+ let units = sampling_units(&granules);
+ assert_eq!(units.first().unwrap().start, 0);
+ assert_eq!(units.last().unwrap().end, granules.len());
+ assert!(units.windows(2).all(|pair| pair[0].end == pair[1].start));
+ assert!(units.iter().all(|unit| unit.len() * page_rows >= 256));
+ }
+ // A row-group tail can join the next file's first page; retain both
indices.
+ let granules = granules_with_rows([341, 40, 341, 40]);
+ assert_eq!(sampling_units(&granules), vec![0..1, 1..4]);
+ assert_ne!(granules[1].file_index, granules[2].file_index);
+
+ // 170-row pages need two pages per unit, doubling first-pass rows
here.
+ let granules = granules_with_rows(vec![170; 16_384]);
+ for seed in 0..3 {
+ let plan = select_first(&granules, 65_536, 65_536, seed).unwrap();
+ assert_eq!(plan.first_rows, 512 * 340);
+ assert_eq!(plan.training.unwrap().len(), 65_536);
+ }
+ }
+
+ #[test]
+ fn systematic_unit_rows_have_equal_marginal_probability() {
+ for (rows, quota) in [(341, 128), (256, 256), (129, 128)] {
+ let mut counts = vec![0; rows];
+ for offset in 0..rows {
+ let selected = unit_training_offsets(rows, quota,
offset).collect::<Vec<_>>();
+ assert_eq!(selected.len(), quota);
+ assert!(selected.windows(2).all(|pair| pair[0] < pair[1]));
+ for row in selected {
+ assert!(row < rows);
+ counts[row] += 1;
+ }
+ }
+ assert!(counts.iter().all(|count| *count == quota));
+ }
+ }
+
+ #[test]
+ fn training_shrink_preserves_exact_counts_and_shard_offsets() {
+ let mut granules = granules_with_rows(vec![256; 4_096]);
+ for granule in &mut granules {
+ granule.range = RowRange::new(granule.range.from() + 17,
granule.range.to() + 17);
+ }
+ for eligible in [1, 100, 127, 128, 129, 32_767, 32_768] {
+ for seed in 0..3 {
+ let plan = select_first(&granules, eligible, eligible,
seed).unwrap();
+ let training = plan.training.unwrap();
+ assert_eq!(training.len(), eligible);
+ assert!(training.windows(2).all(|pair| pair[0] < pair[1]));
+ assert!(training
+ .iter()
+ .all(|row| (17..17 + 4_096 * 256).contains(row)));
+ assert!(training.iter().all(|row| plan
+ .first
+ .iter()
+ .any(|range| { range.from() <= *row && *row <= range.to()
})));
+ }
+ }
+ // Three or more hits can exceed a 256-row unit: do not truncate
quotas.
+ let granules = granules_with_rows(vec![256; 64]);
+ let plan = select_first(&granules, 128, 128, 0).unwrap();
+ assert!(plan.training.is_none());
+ assert_eq!(plan.first_rows, 16_384);
+ }
+
+ #[test]
+ fn large_granule_training_is_row_weighted_when_the_byte_gate_allows_it() {
+ let mut granules =
+ granules_with_rows((0..4_096).map(|index| if index == 2_048 {
524_288 } else { 128 }));
+ for granule in &mut granules {
+ granule.byte_ranges = std::iter::once(0..1).collect();
+ }
+ let large = &granules[2_048].range;
+ for seed in 0..32 {
+ let training = select_first(&granules, 65_536, 65_536, seed)
+ .unwrap()
+ .training
+ .unwrap();
+ let large_rows = training
+ .iter()
+ .filter(|row| large.from() <= **row && **row <= large.to())
+ .count();
+ assert!(
+ (32_512..=33_024).contains(&large_rows),
+ "seed {seed}: {large_rows}"
+ );
+ }
+ }
+
+ #[test]
+ fn minimum_granule_count_can_force_whole_shard_without_the_byte_gate() {
+ let mut granules =
+ granules_with_rows((0..4_096).map(|index| if index == 0 {
1_048_576 } else { 256 }));
+ for granule in &mut granules {
+ granule.byte_ranges = std::iter::once(0..1).collect();
+ }
+ // All units are single granules. Prove quota and byte gates pass,
while
+ // many row strata hit the same large granule and the count gate fails.
+ let rows = granules.last().unwrap().range.to() as usize + 1;
+ let mut quotas = vec![0; granules.len()];
+ for row in pick_indices(rows, MIN_STRATA, &mut 7) {
+ let index = granules.partition_point(|granule| granule.range.to()
< row as i64);
+ quotas[index] += ROWS_PER_STRATUM;
+ }
+ let selected = quotas
+ .iter()
+ .enumerate()
+ .filter_map(|(index, quota)| (*quota > 0).then_some(index))
+ .collect::<HashSet<_>>();
+ assert!(quotas
+ .iter()
+ .zip(&granules)
+ .all(|(quota, granule)| { *quota <= granule.range.count() as usize
}));
+ assert!(selected.len() < MIN_STRATA);
+ assert!(
+ granule_bytes(&granules, Some(&selected)) * FIRST_BYTES_DENOMINATOR
+ <= granule_bytes(&granules, None) * FIRST_BYTES_NUMERATOR
+ );
+ assert!(select_first(&granules, 32_768, 32_768, 7)
+ .unwrap()
+ .training
+ .is_none());
+ }
+
+ #[test]
+ fn stratified_first_and_rest_cover_every_granule_once() {
+ let granules = granules_with_rows(vec![256; 4_096]);
+ let plan = select_first(&granules, 65_536, 32_768, 7).unwrap();
+ let training = plan.training.as_ref().unwrap();
+ assert_eq!(training.len(), 32_768);
+ assert!(training.windows(2).all(|pair| pair[0] < pair[1]));
+ assert!(training.iter().all(|row| plan
+ .first
+ .iter()
+ .any(|range| range.from() <= *row && *row <= range.to())));
+ let mut all = plan.first.clone();
+ all.extend(plan.rest.clone());
+ assert_eq!(
+ merge_row_ranges(all),
+ vec![RowRange::new(0, 4_096 * 256 - 1)]
+ );
+ assert_eq!(
+ plan.first_rows
+ + plan
+ .rest
+ .iter()
+ .map(|range| range.count() as usize)
+ .sum::<usize>(),
+ 4_096 * 256
+ );
+ for stratum in 0..256 {
+ let (start, end) = (stratum * 4_096, (stratum + 1) * 4_096);
+ assert_eq!(
+ training
+ .iter()
+ .filter(|row| start <= **row && **row < end)
+ .count(),
+ 128,
+ "stratum {stratum}"
+ );
+ }
+ assert_eq!(
+ plan.training,
+ select_first(&granules, 65_536, 32_768, 7).unwrap().training
+ );
+ assert_ne!(
+ plan.training,
+ select_first(&granules, 65_536, 32_768, 8).unwrap().training
+ );
+ }
+
+ #[test]
+ fn shard_seed_follows_the_whole_shard_identity() {
+ let shard = VindexIndexShard {
+ partition: crate::spec::BinaryRow::new(0),
+ partition_bytes: vec![1, 2, 3, 4, 5, 6, 7, 8, 9],
+ files: Vec::new(),
+ row_range_start: 0,
+ row_range_end: 999,
+ snapshot_id: 3,
+ source_bucket: 0,
+ total_buckets: 2,
+ bucket_path: "memory:/t/bucket-0".to_string(),
+ };
+ let picks = |shard: &VindexIndexShard| {
+ pick_indices(4_096, 512, &mut
shard_seed(shard)).collect::<Vec<_>>()
+ };
+ assert_eq!(picks(&shard), picks(&shard.clone()));
+
+ let mut others = vec![shard.clone(); 5];
+ others[0].partition_bytes = vec![1, 2, 3, 4, 5, 6, 7, 8, 10];
+ others[1].partition_bytes = Vec::new();
+ others[2].source_bucket = 1;
+ others[3].snapshot_id = 4;
+ others[4].row_range_end = 1_999;
+ for other in &others {
+ assert_ne!(picks(&shard), picks(other), "{other:?}");
+ }
+ }
+
+ #[test]
+ fn first_picks_do_not_alias_with_periodic_granule_order() {
+ let granules = granules_with_rows(vec![256; 4_096]);
+ for seed in 0..32u64 {
+ let training = select_first(&granules, 65_536, 65_536, seed)
+ .unwrap()
+ .training
+ .unwrap();
+ for period in [2usize, 4, 8, 16] {
+ let mut hits = vec![0usize; period];
+ for row in &training {
+ hits[*row as usize / 256 % period] += 1;
+ }
+ let expected = training.len() / period;
+ assert!(
+ hits.iter()
+ .all(|count| *count * 4 > expected && *count <
expected * 2),
+ "seed {seed} period {period}: {hits:?}"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn unequal_granules_are_sampled_in_proportion() {
+ let granules =
+ granules_with_rows((0..4_096).map(|index| if index % 8 == 4 { 16 }
else { 256 }));
+ let short_share = 16.0 / (16.0 + 7.0 * 256.0);
+ for seed in 0..32u64 {
+ let training = select_first(&granules, 65_536, 65_536, seed)
+ .unwrap()
+ .training
+ .unwrap();
+ assert_eq!(training.len(), 65_536);
+ let short_rows = training
+ .iter()
+ .filter(|row| granules.partition_point(|g| g.range.to() <
**row) % 8 == 4)
+ .count();
+ let share = short_rows as f64 / training.len() as f64;
+ assert!(
+ share > short_share / 2.0 && share < short_share * 2.0,
+ "seed {seed}: {share}"
+ );
+ }
+ }
+
+ #[test]
+ fn short_granules_use_static_units_instead_of_filling_from_the_head() {
+ let granules = granules_with_rows(vec![64; 4_096]);
+ let plan = select_first(&granules, 32_768, 32_768, 7).unwrap();
+ assert_eq!(plan.first_rows, 65_536);
+ let training = plan.training.unwrap();
+ for quarter in 0..4 {
+ let picked = training
+ .iter()
+ .filter(|row| **row / 65_536 == quarter)
+ .count();
+ assert_eq!(picked, 8_192, "quarter {quarter}");
+ }
+ }
+
+ #[test]
+ fn coalesced_page_holes_force_all_first() {
+ let mut granules = granules_with_rows(vec![256; 4_096]);
+ assert!(select_first(&granules, 65_536, 65_536, 7)
+ .unwrap()
+ .training
+ .is_some());
+ for (index, granule) in granules.iter_mut().enumerate() {
+ granule.file_index = 0;
+ granule.byte_ranges =
+ std::iter::once(index as u64 * 512..index as u64 * 512 +
256).collect();
+ }
+ // Same rows, but coalescing reads the holes too: the byte gate alone
changes the path.
+ assert!(select_first(&granules, 65_536, 65_536, 7)
+ .unwrap()
+ .training
+ .is_none());
+ }
+
+ #[test]
+ fn coarse_granules_read_all_first() {
+ let granules = (0..16)
+ .map(|row| Granule {
+ range: RowRange::new(row, row),
+ file_index: 0,
+ byte_ranges: std::iter::once(row as u64..row as u64 +
1).collect(),
+ })
+ .collect::<Vec<_>>();
+ assert!(select_first(&granules, 8, 8, 7).unwrap().training.is_none());
+ }
+
+ #[test]
+ fn near_full_shard_in_shared_file_reads_all_first() {
+ let file_granules = (0..1_200)
+ .map(|row| ParquetGranule {
+ first_row: row * 256,
+ row_count: 256,
+ byte_ranges: std::iter::once(
+ row as u64 * 2 * 1024 * 1024..row as u64 * 2 * 1024 * 1024
+ 1,
+ )
+ .collect(),
+ })
+ .collect();
+ let shard_range = RowRange::new(400 * 256, 700 * 256 - 1);
+ let mut granules = Vec::new();
+
+ append_shard_granules(&mut granules, 0, 0, &shard_range,
file_granules).unwrap();
+
+ assert_eq!(granules.len(), 300);
+ assert!(select_first(&granules, 32_768, 32_768, 7)
+ .unwrap()
+ .training
+ .is_none());
+ }
+
+ #[test]
+ fn granule_partition_requires_exact_contiguous_coverage() {
+ for (ranges, shard, expected) in [
+ (vec![], (0, 9), false),
+ (vec![(0, 9)], (0, 9), true),
+ (vec![(0, 4), (5, 9)], (0, 9), true),
+ (vec![(0, 4), (6, 9)], (0, 9), false),
+ (vec![(0, 5), (5, 9)], (0, 9), false),
+ (vec![(1, 9)], (0, 9), false),
+ (vec![(0, 8)], (0, 9), false),
+ (vec![(0, 9), (0, 9)], (0, 9), false),
+ (vec![(0, 2), (6, 7), (3, 5), (8, 9)], (0, 9), false),
+ (vec![(i64::MAX, i64::MAX)], (i64::MAX, i64::MAX), true),
+ (
+ vec![(i64::MAX - 1, i64::MAX - 1), (i64::MAX, i64::MAX)],
+ (i64::MAX - 1, i64::MAX),
+ true,
+ ),
+ (
+ vec![(i64::MAX, i64::MAX), (i64::MAX, i64::MAX)],
+ (i64::MAX, i64::MAX),
+ false,
+ ),
+ ] {
+ let granules = ranges
+ .iter()
+ .map(|&(from, to)| Granule {
+ range: RowRange::new(from, to),
+ file_index: 0,
+ byte_ranges: vec![],
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(
+ granules_partition_shard(&granules, &RowRange::new(shard.0,
shard.1)),
+ expected,
+ "ranges={ranges:?}, shard={shard:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn overlapping_data_evolution_providers_do_not_partition_shard() {
+ let shard_range = RowRange::new(0, 99);
+ let mut granules = Vec::new();
+ for file_index in 0..2 {
+ append_shard_granules(
+ &mut granules,
+ file_index,
+ 0,
+ &shard_range,
+ vec![ParquetGranule {
+ first_row: 0,
+ row_count: 100,
+ byte_ranges: std::iter::once(0..1).collect(),
+ }],
+ )
+ .unwrap();
+ }
+ granules.sort_by_key(|granule| granule.range.from());
+
+ assert!(!granules_partition_shard(&granules, &shard_range));
+ granules.truncate(1);
+ assert!(granules_partition_shard(&granules, &shard_range));
+ }
+
+ #[test]
+ fn live_pipeline_completes_with_one_blocking_thread() {
+ let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .max_blocking_threads(1)
+ .enable_all()
+ .build()
+ .unwrap();
+
+ runtime.block_on(async {
+ tokio::time::timeout(Duration::from_secs(5), async {
+ let vectors = vec![0.0f32, 0.0, 1.0, 1.0, 2.0, 2.0];
+ let ids = vec![0, 1, 2];
+ let config =
+
paimon_vindex_core::index::VectorIndexConfig::from_options(&HashMap::from([
+ ("index.type".to_string(), "ivf_flat".to_string()),
+ ("dimension".to_string(), "2".to_string()),
+ ("nlist".to_string(), "1".to_string()),
+ ("metric".to_string(), "l2".to_string()),
+ ]))
+ .unwrap();
+ let mut trainer = VectorIndexTrainer::new(config).unwrap();
+ trainer
+ .add_training_vectors_mut(&vectors, ids.len())
+ .unwrap();
+
+ let spill = spawn_spill_writer(false).unwrap();
+ for _ in 0..2 {
+ spill
+ .sender
+ .send(SpillRecord {
+ ids: Vec::new(),
+ bytes: Vec::new(),
+ })
+ .await
+ .unwrap();
+ }
+ spill
+ .sender
+ .send(SpillRecord {
+ ids,
+ bytes: vectors
+ .iter()
+ .flat_map(|value| value.to_ne_bytes())
+ .collect(),
+ })
+ .await
+ .unwrap();
+ let training = tokio::task::spawn_blocking(move || ->
std::io::Result<_> {
+ Ok((trainer.finish()?, Duration::ZERO))
+ });
+
+ let (pipeline, _) =
+ start_live_pipeline(training, spill,
"embedding".to_string(), 2, false)
+ .await
+ .unwrap();
+ let (consumer, replay, _, _) =
finish_live_pipeline(pipeline).await;
+ let (_, rows_added, consumer_replay_rows, _) =
consumer.unwrap();
+ let (replay_rows, _) = replay.unwrap();
+ assert_eq!((rows_added, consumer_replay_rows, replay_rows),
(3, 3, 3));
+ })
+ .await
+ .expect("vindex live pipeline deadlocked on the shared blocking
pool");
+ });
+ }
+
+ #[tokio::test]
+ async fn spill_reports_written_bytes() {
+ let spill = spawn_spill_writer(true).unwrap();
+ spill
+ .sender
+ .send(SpillRecord {
+ ids: vec![1, 2],
+ bytes: vec![0; 16],
+ })
+ .await
+ .unwrap();
+
+ let (file, bytes, _) = spill.finish().await.unwrap();
+ assert_eq!(bytes, 40);
+
+ let (sender, mut receiver) = mpsc::channel(1);
+ let (rows, _) = spawn_replay(file, sender, 2,
true).await.unwrap().unwrap();
+ assert_eq!(rows, 2);
+ assert!(matches!(
+ receiver.recv().await,
+ Some(AddItem::Spilled(ids, _)) if ids == vec![1, 2]
+ ));
+ }
+}
diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs
b/crates/paimon/src/table/vindex_index_build_builder/tests.rs
index a4636b57..fb5c1499 100644
--- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs
+++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use super::extraction::validate_vector_batch;
+use super::extraction::{validate_vector_batch, validate_vector_batch_ranges};
use super::planning::{plan_vindex_shards, VindexIndexShard};
use super::validation::{
checked_training_sample_index, checked_training_vector_count,
checked_vector_bytes,
@@ -252,6 +252,70 @@ fn test_extract_vectors_accepts_list_float32_and_row_ids()
{
assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]);
}
+#[test]
+fn test_ranged_vector_validation_accepts_gaps_across_batches() {
+ let ranges = vec![RowRange::new(10, 11), RowRange::new(15, 16)];
+ let batches = [
+ vector_batch(
+ vec![
+ Some(vec![Some(1.0), Some(2.0)]),
+ Some(vec![Some(3.0), Some(4.0)]),
+ ],
+ vec![Some(10), Some(11)],
+ ),
+ vector_batch(
+ vec![
+ Some(vec![Some(5.0), Some(6.0)]),
+ Some(vec![Some(7.0), Some(8.0)]),
+ ],
+ vec![Some(15), Some(16)],
+ ),
+ ];
+ let mut range_index = 0;
+ let mut expected_row_id = ranges[0].from();
+
+ for batch in &batches {
+ validate_vector_batch_ranges(
+ batch,
+ "embedding",
+ 2,
+ &ranges,
+ &mut range_index,
+ &mut expected_row_id,
+ )
+ .unwrap();
+ }
+
+ assert_eq!(range_index, ranges.len());
+ assert_eq!(expected_row_id, 17);
+}
+
+#[test]
+fn test_ranged_vector_validation_rejects_bad_row_ids() {
+ let ranges = vec![RowRange::new(10, 11), RowRange::new(15, 16)];
+ for row_ids in [
+ vec![Some(10), Some(10)],
+ vec![Some(10), Some(15)],
+ vec![Some(9), Some(10)],
+ vec![Some(10), Some(11), Some(16), Some(15)],
+ vec![Some(10), Some(11), Some(15), Some(16), Some(17)],
+ ] {
+ let rows = row_ids.len();
+ let batch = vector_batch(vec![Some(vec![Some(1.0), Some(2.0)]); rows],
row_ids);
+ let mut range_index = 0;
+ let mut expected_row_id = ranges[0].from();
+ assert!(validate_vector_batch_ranges(
+ &batch,
+ "embedding",
+ 2,
+ &ranges,
+ &mut range_index,
+ &mut expected_row_id,
+ )
+ .is_err());
+ }
+}
+
#[test]
fn test_extract_vectors_rejects_dimension_mismatch() {
let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]);
@@ -605,6 +669,10 @@ async fn vindex_incremental_build_indexes_only_new_rows() {
let second_built = table
.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER)
.with_index_column("embedding")
+ .with_options(HashMap::from([(
+ "vindex.build.granule.enabled".to_string(),
+ "false".to_string(),
+ )]))
.execute()
.await
.unwrap();
@@ -643,6 +711,315 @@ async fn vindex_incremental_build_indexes_only_new_rows()
{
}
}
+#[tokio::test]
+async fn vindex_small_training_sample_preserves_tail_cluster_recall() {
+ let table_path = "memory:/test_vindex_small_sample_recall";
+ let mut options = table_options("1000");
+ for (key, value) in [
+ ("ivf-sq.dimension", "1"),
+ ("ivf-sq.nlist", "1"),
+ ("ivf-sq.metric", "l2"),
+ ("ivf-sq.train.sample-ratio", "0.1"),
+ ] {
+ options.insert(key.to_string(), value.to_string());
+ }
+ let table = test_table_with_io(
+ FileIOBuilder::new("memory").build().unwrap(),
+ table_path,
+ vindex_schema_builder(options).build().unwrap(),
+ );
+ setup_dirs(table.file_io(), table_path).await;
+ write_vectors(
+ &table,
+ (0..1000).collect(),
+ (0..1000)
+ .map(|id| {
+ vec![if id < 450 {
+ 0.0
+ } else if id < 900 {
+ 1.0
+ } else {
+ 100.0
+ }]
+ })
+ .collect(),
+ )
+ .await;
+ assert_eq!(
+ table
+ .new_vindex_index_build_builder(crate::vindex::IVF_SQ_IDENTIFIER)
+ .with_index_column("embedding")
+ .execute()
+ .await
+ .unwrap(),
+ 1
+ );
+
+ let result = table
+ .new_vector_search_builder()
+ .with_vector_column("embedding")
+ .with_query_vector(vec![100.0])
+ .with_limit(10)
+ .with_options(HashMap::from([(
+ "ivf-sq.nprobe".to_string(),
+ "1".to_string(),
+ )]))
+ .execute()
+ .await
+ .unwrap();
+ let row_ids = &result.row_ids().unwrap().row_ids;
+ assert_eq!(row_ids.len(), 10);
+ // Equal-distance IDs need not have a stable order; all hits must be in
the tail cluster.
+ assert!(
+ row_ids.iter().all(|row_id| (900..1000).contains(row_id)),
+ "{result:?}"
+ );
+}
+
+/// Every eighth data file holds a cluster the other files do not. With 4,096
one-page
+/// files the granule planner reads 512 of them first, one per stratum of
eight, so a
+/// fixed position in the stratum would train without ever seeing that cluster.
+#[tokio::test]
+async fn vindex_granule_training_sees_file_periodic_cluster() {
+ const FILES: usize = 4_096;
+ const ROWS_PER_FILE: usize = 256;
+ let table_path = "memory:/test_vindex_granule_periodic_cluster";
+ let mut options = table_options("2000000");
+ for (key, value) in [
+ ("ivf-sq.dimension", "1"),
+ ("ivf-sq.nlist", "1"),
+ ("ivf-sq.metric", "l2"),
+ ("target-file-size", "1b"),
+ ] {
+ options.insert(key.to_string(), value.to_string());
+ }
+ let table = test_table_with_io(
+ FileIOBuilder::new("memory").build().unwrap(),
+ table_path,
+ vindex_schema_builder(options).build().unwrap(),
+ );
+ setup_dirs(table.file_io(), table_path).await;
+ let mut table_write = TableWrite::new(&table,
"test-user".to_string()).unwrap();
+ for file in 0..FILES {
+ let value = if file % 8 == 7 {
+ 100.0
+ } else {
+ (file % 2) as f32
+ };
+ let first_id = (file * ROWS_PER_FILE) as i32;
+ table_write
+ .write_arrow_batch(&build_vector_batch(
+ (first_id..first_id + ROWS_PER_FILE as i32).collect(),
+ vec![vec![value]; ROWS_PER_FILE],
+ ))
+ .await
+ .unwrap();
+ }
+ let messages = table_write.prepare_commit().await.unwrap();
+ TableCommit::new(table.clone(), "test-user".to_string())
+ .commit(messages)
+ .await
+ .unwrap();
+
+ assert_eq!(
+ table
+ .new_vindex_index_build_builder(crate::vindex::IVF_SQ_IDENTIFIER)
+ .with_index_column("embedding")
+ .execute()
+ .await
+ .unwrap(),
+ 1
+ );
+
+ let result = table
+ .new_vector_search_builder()
+ .with_vector_column("embedding")
+ .with_query_vector(vec![100.0])
+ .with_limit(10)
+ .with_options(HashMap::from([(
+ "ivf-sq.nprobe".to_string(),
+ "1".to_string(),
+ )]))
+ .execute()
+ .await
+ .unwrap();
+ let row_ids = &result.row_ids().unwrap().row_ids;
+ assert_eq!(row_ids.len(), 10);
+ // Equal-distance IDs need not have a stable order; all hits must be in
the periodic cluster.
+ assert!(
+ row_ids
+ .iter()
+ .all(|row_id| (*row_id as usize / ROWS_PER_FILE) % 8 == 7),
+ "{result:?}"
+ );
+}
+
+#[tokio::test]
+async fn vindex_granule_training_sees_one_oversized_row_group() {
+ use crate::arrow::format::parquet::parquet_granules;
+ use parquet::arrow::AsyncArrowWriter;
+ use parquet::basic::Compression;
+ use parquet::file::properties::WriterProperties;
+
+ const FILES: usize = 4_096;
+ const SMALL_ROWS: usize = 128;
+ const LARGE_ROWS: usize = 524_288;
+ let total_rows = (FILES - 1) * SMALL_ROWS + LARGE_ROWS;
+ // Reuse physical file contents; logical row IDs are assigned by the
commit.
+ let mut contents = Vec::new();
+ for (rows, value) in [(SMALL_ROWS, 0.0), (SMALL_ROWS, 1.0), (LARGE_ROWS,
100.0)] {
+ let batch = build_vector_batch((0..rows as i32).collect(),
vec![vec![value]; rows]);
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(Some(rows))
+ .set_offset_index_disabled(true)
+ .set_dictionary_enabled(false)
+ .set_compression(Compression::UNCOMPRESSED)
+ .build();
+ let mut bytes = Vec::new();
+ let mut writer =
+ AsyncArrowWriter::try_new(&mut bytes, batch.schema(),
Some(props)).unwrap();
+ writer.write(&batch).await.unwrap();
+ writer.close().await.unwrap();
+ contents.push(bytes::Bytes::from(bytes));
+ }
+
+ // The control uses full spill. Three real snapshot identities exercise the
+ // production seed derivation, without adding a seed override to the
builder.
+ for (granule_enabled, snapshots) in [(false, 1), (true, 1), (true, 2),
(true, 3)] {
+ let table_path =
format!("memory:/oversized_granule_{granule_enabled}_{snapshots}");
+ let mut options = table_options("2000000");
+ for (key, value) in [
+ ("ivf-sq.dimension", "1"),
+ ("ivf-sq.nlist", "1"),
+ ("ivf-sq.metric", "l2"),
+ ] {
+ options.insert(key.to_string(), value.to_string());
+ }
+ let table = test_table_with_io(
+ FileIOBuilder::new("memory").build().unwrap(),
+ &table_path,
+ vindex_schema_builder(options).build().unwrap(),
+ );
+ setup_dirs(table.file_io(), &table_path).await;
+ let mut files = Vec::new();
+ for index in 0..FILES {
+ let large = index == FILES - 1;
+ let bytes = &contents[if large { 2 } else { index % 2 }];
+ let name = format!("data-{index:04}.parquet");
+ table
+ .file_io()
+ .new_output(&format!("{table_path}/bucket-0/{name}"))
+ .unwrap()
+ .write(bytes.clone())
+ .await
+ .unwrap();
+ let mut file = data_file(
+ &name,
+ None,
+ if large { LARGE_ROWS } else { SMALL_ROWS } as i64,
+ );
+ file.file_size = bytes.len() as i64;
+ file.file_source = Some(0); // APPEND: the commit assigns row IDs.
+ files.push(file);
+ }
+ for group in files.chunks(FILES.div_ceil(snapshots)) {
+ TableCommit::new(table.clone(), "test-user".to_string())
+ .commit(vec![CommitMessage::new(
+ BinaryRow::new(0).to_serialized_bytes(),
+ 0,
+ group.to_vec(),
+ )])
+ .await
+ .unwrap();
+ }
+ let snapshot = SnapshotManager::new(table.file_io().clone(),
table_path.clone())
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .unwrap();
+ assert_eq!(snapshot.id(), snapshots as i64);
+ let entries = table
+ .new_read_builder()
+ .new_scan()
+ .with_scan_all_files()
+ .plan_manifest_entries(&snapshot)
+ .await
+ .unwrap();
+ let large_file = entries
+ .iter()
+ .find(|entry| entry.file().row_count == LARGE_ROWS as i64)
+ .unwrap()
+ .file();
+ let (large_start, large_end) = large_file.row_id_range().unwrap();
+ let input = table
+ .file_io()
+
.new_input(&large_file.data_file_path(&format!("{table_path}/bucket-0")))
+ .unwrap();
+ let (physical_granules, has_offset_index) = parquet_granules(
+ Box::new(input.reader().await.unwrap()),
+ large_file.file_size as u64,
+ "embedding",
+ true,
+ )
+ .await
+ .unwrap();
+ assert!(!has_offset_index);
+ assert_eq!(physical_granules.len(), 1);
+ assert_eq!(physical_granules[0].row_count, LARGE_ROWS as i64);
+
+ let mut builder =
table.new_vindex_index_build_builder(crate::vindex::IVF_SQ_IDENTIFIER);
+ builder
+ .with_index_column("embedding")
+ .with_options(HashMap::from([(
+ "vindex.build.granule.enabled".to_string(),
+ granule_enabled.to_string(),
+ )]));
+ if granule_enabled {
+ let shards = plan_vindex_shards(
+ table.location(),
+ table.schema().partition_keys(),
+ table.schema().fields(),
+ &CoreOptions::new(table.schema().options()),
+ snapshot.id(),
+ entries,
+ 2_000_000,
+ &[],
+ )
+ .unwrap();
+ assert_eq!(shards.len(), 1);
+ let plan = builder
+ .plan_granules(&shards[0], "embedding", total_rows, 65_536)
+ .await
+ .unwrap();
+ // The independently trained cluster occupies half the logical rows
+ // and projected bytes, so selecting it must trigger the byte gate.
+ assert!(plan.rest.is_empty(), "snapshot {}", snapshot.id());
+ assert_eq!(plan.first, vec![RowRange::new(0, total_rows as i64 -
1)]);
+ }
+ assert_eq!(builder.execute().await.unwrap(), 1);
+ let result = table
+ .new_vector_search_builder()
+ .with_vector_column("embedding")
+ .with_query_vector(vec![100.0])
+ .with_limit(10)
+ .with_options(HashMap::from([(
+ "ivf-sq.nprobe".to_string(),
+ "1".to_string(),
+ )]))
+ .execute()
+ .await
+ .unwrap();
+ let row_ids = &result.row_ids().unwrap().row_ids;
+ assert_eq!(row_ids.len(), 10);
+ assert!(
+ row_ids
+ .iter()
+ .all(|row| (large_start as u64..=large_end as
u64).contains(row)),
+ "granule={granule_enabled}, snapshot={snapshots}: {result:?}"
+ );
+ }
+}
+
#[tokio::test]
async fn vindex_build_cleans_written_shards_when_later_shard_fails() {
let table_path = "memory:/test_vindex_abort_written_shard";
diff --git a/crates/paimon/src/table/vindex_index_build_builder/timing.rs
b/crates/paimon/src/table/vindex_index_build_builder/timing.rs
index 3221dc4c..ae43568f 100644
--- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs
+++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs
@@ -41,8 +41,10 @@ pub(super) struct VectorIndexBuildTiming {
pub(super) parquet_projected_bytes_total: u64,
pub(super) parquet_peak_inflight_row_groups: usize,
pub(super) raw_temp_write: Duration,
+ pub(super) granule_spill_write: Duration,
pub(super) train_finish: Duration,
pub(super) raw_temp_reread: Duration,
+ pub(super) granule_spill_read: Duration,
pub(super) index_add: Duration,
pub(super) serialize_upload: Duration,
pub(super) rows: usize,
@@ -50,6 +52,7 @@ pub(super) struct VectorIndexBuildTiming {
pub(super) training_rows_retained: usize,
pub(super) batch_count: usize,
pub(super) raw_temp_bytes: usize,
+ pub(super) granule_spill_bytes: u64,
pub(super) index_bytes: u64,
pub(super) data_file_count: usize,
pub(super) file_name: String,
@@ -68,7 +71,7 @@ impl VectorIndexBuildTiming {
.saturating_add(commit);
let unattributed = total.saturating_sub(accounted);
eprintln!(
- "event=paimon_vector_index_build index_type={} file={} rows={}
training_rows_seen={} training_rows_retained={} batch_count={}
raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3}
parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3}
remaining_batch_wait_ms={:.3} parquet_row_group_count={}
parquet_projected_bytes_min={} parquet_projected_bytes_max={}
parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={}
raw_temp_wri [...]
+ "event=paimon_vector_index_build index_type={} file={} rows={}
training_rows_seen={} training_rows_retained={} batch_count={}
raw_temp_bytes={} granule_spill_bytes={} index_bytes={}
source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3}
file_schema_open_ms={:.3} first_batch_wait_ms={:.3}
remaining_batch_wait_ms={:.3} parquet_row_group_count={}
parquet_projected_bytes_min={} parquet_projected_bytes_max={}
parquet_projected_bytes_total={} parquet_peak_inflight_row [...]
index_type,
self.file_name,
self.rows,
@@ -76,6 +79,7 @@ impl VectorIndexBuildTiming {
self.training_rows_retained,
self.batch_count,
self.raw_temp_bytes,
+ self.granule_spill_bytes,
self.index_bytes,
self.source_batch_wait.as_secs_f64() * 1000.0,
self.oss_read.as_secs_f64() * 1000.0,
@@ -89,8 +93,10 @@ impl VectorIndexBuildTiming {
self.parquet_projected_bytes_total,
self.parquet_peak_inflight_row_groups,
self.raw_temp_write.as_secs_f64() * 1000.0,
+ self.granule_spill_write.as_secs_f64() * 1000.0,
self.train_finish.as_secs_f64() * 1000.0,
self.raw_temp_reread.as_secs_f64() * 1000.0,
+ self.granule_spill_read.as_secs_f64() * 1000.0,
self.index_add.as_secs_f64() * 1000.0,
self.serialize_upload.as_secs_f64() * 1000.0,
commit.as_secs_f64() * 1000.0,
diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs
b/crates/paimon/src/table/vindex_index_build_builder/writer.rs
index 803c0888..8c7fc126 100644
--- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs
+++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs
@@ -26,7 +26,7 @@ use super::VindexIndexBuildBuilder;
use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME};
use crate::table::data_file_reader::DataFileReadTiming;
use crate::table::table_read::configured_parquet_read_budget;
-use crate::vindex::VindexVectorIndexOptions;
+use crate::vindex::{VindexVectorIndexOptions, DISKANN_IDENTIFIER};
use crate::{Error, Result};
use arrow_buffer::MutableBuffer;
use futures::TryStreamExt;
@@ -41,6 +41,7 @@ use tokio_util::io::SyncIoBridge;
const INDEX_DIR: &str = "index";
const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024;
+
pub(super) struct BuiltIndexFile {
pub(super) meta: IndexFileMeta,
pub(super) timing: Option<VectorIndexBuildTiming>,
@@ -55,6 +56,44 @@ impl<'a> VindexIndexBuildBuilder<'a> {
index_field_id: i32,
options: &VindexVectorIndexOptions,
index_meta: Vec<u8>,
+ ) -> Result<BuiltIndexFile> {
+ let use_granule = self.index_type != DISKANN_IDENTIFIER &&
options.granule_build_enabled;
+ log::info!(
+ "vindex build strategy: index_type={}, strategy={}",
+ self.index_type,
+ if use_granule { "granule" } else { "full-spill" }
+ );
+ if use_granule {
+ return self
+ .build_index_file_granule(
+ shard,
+ index_column,
+ dimension,
+ index_field_id,
+ options,
+ index_meta,
+ )
+ .await;
+ }
+ self.build_full_spill_index_file(
+ shard,
+ index_column,
+ dimension,
+ index_field_id,
+ options,
+ index_meta,
+ )
+ .await
+ }
+
+ async fn build_full_spill_index_file(
+ &self,
+ shard: &VindexIndexShard,
+ index_column: &str,
+ dimension: i32,
+ index_field_id: i32,
+ options: &VindexVectorIndexOptions,
+ index_meta: Vec<u8>,
) -> Result<BuiltIndexFile> {
let timing_enabled = vector_index_build_timing_enabled();
let total_start = timing_enabled.then(Instant::now);
@@ -328,6 +367,65 @@ impl<'a> VindexIndexBuildBuilder<'a> {
})?;
let serialize_upload_start = timing_enabled.then(Instant::now);
+ let meta = self
+ .finish_index_file(writer, shard, index_field_id, index_meta,
row_count)
+ .await?;
+ let serialize_upload =
+ serialize_upload_start.map_or(Duration::ZERO, |start|
start.elapsed());
+ let (oss_read, parquet_decode) = read_timing
+ .as_ref()
+ .map_or((Duration::ZERO, Duration::ZERO), |timing| {
+ (timing.file_read(), timing.parquet_decode())
+ });
+ let (file_schema_open, first_batch_wait, remaining_batch_wait) =
read_timing
+ .as_ref()
+ .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing|
{
+ timing.file_waits()
+ });
+ let parquet_diagnostics = parquet_read_budget
+ .as_ref()
+ .map_or_else(Default::default, |budget| budget.diagnostics());
+ let timing = total_start.map(|start| VectorIndexBuildTiming {
+ total_without_commit: start.elapsed(),
+ source_batch_wait,
+ oss_read,
+ parquet_decode,
+ file_schema_open,
+ first_batch_wait,
+ remaining_batch_wait,
+ parquet_row_group_count: parquet_diagnostics.row_group_count,
+ parquet_projected_bytes_min:
parquet_diagnostics.projected_bytes_min,
+ parquet_projected_bytes_max:
parquet_diagnostics.projected_bytes_max,
+ parquet_projected_bytes_total:
parquet_diagnostics.projected_bytes_total,
+ parquet_peak_inflight_row_groups:
parquet_diagnostics.peak_inflight,
+ raw_temp_write,
+ granule_spill_write: Duration::ZERO,
+ train_finish,
+ raw_temp_reread,
+ granule_spill_read: Duration::ZERO,
+ index_add,
+ serialize_upload,
+ rows: row_count_usize,
+ training_rows_seen: training_vector_count,
+ training_rows_retained,
+ batch_count,
+ raw_temp_bytes: bytes_written,
+ granule_spill_bytes: 0,
+ index_bytes: meta.file_size as u64,
+ data_file_count: shard.files.len(),
+ file_name: meta.file_name.clone(),
+ });
+ Ok(BuiltIndexFile { meta, timing })
+ }
+
+ pub(super) async fn finish_index_file(
+ &self,
+ writer: VectorIndexWriter,
+ shard: &VindexIndexShard,
+ index_field_id: i32,
+ index_meta: Vec<u8>,
+ row_count: i64,
+ ) -> Result<IndexFileMeta> {
self.table
.file_io()
.mkdirs(&format!(
@@ -377,11 +475,9 @@ impl<'a> VindexIndexBuildBuilder<'a> {
return Err(error);
}
};
- let serialize_upload =
- serialize_upload_start.map_or(Duration::ZERO, |start|
start.elapsed());
- let meta = IndexFileMeta {
+ Ok(IndexFileMeta {
index_type: self.index_type.clone(),
- file_name: file_name.clone(),
+ file_name,
file_size: checked_i64(
status.size,
"Index file is too large for Rust IndexFileMeta",
@@ -397,47 +493,6 @@ impl<'a> VindexIndexBuildBuilder<'a> {
source_meta: None,
index_meta: Some(index_meta),
}),
- };
- let (oss_read, parquet_decode) = read_timing
- .as_ref()
- .map_or((Duration::ZERO, Duration::ZERO), |timing| {
- (timing.file_read(), timing.parquet_decode())
- });
- let (file_schema_open, first_batch_wait, remaining_batch_wait) =
read_timing
- .as_ref()
- .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing|
{
- timing.file_waits()
- });
- let parquet_diagnostics = parquet_read_budget
- .as_ref()
- .map_or_else(Default::default, |budget| budget.diagnostics());
- let timing = total_start.map(|start| VectorIndexBuildTiming {
- total_without_commit: start.elapsed(),
- source_batch_wait,
- oss_read,
- parquet_decode,
- file_schema_open,
- first_batch_wait,
- remaining_batch_wait,
- parquet_row_group_count: parquet_diagnostics.row_group_count,
- parquet_projected_bytes_min:
parquet_diagnostics.projected_bytes_min,
- parquet_projected_bytes_max:
parquet_diagnostics.projected_bytes_max,
- parquet_projected_bytes_total:
parquet_diagnostics.projected_bytes_total,
- parquet_peak_inflight_row_groups:
parquet_diagnostics.peak_inflight,
- raw_temp_write,
- train_finish,
- raw_temp_reread,
- index_add,
- serialize_upload,
- rows: row_count_usize,
- training_rows_seen: training_vector_count,
- training_rows_retained,
- batch_count,
- raw_temp_bytes: bytes_written,
- index_bytes: status.size,
- data_file_count: shard.files.len(),
- file_name,
- });
- Ok(BuiltIndexFile { meta, timing })
+ })
}
}
diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs
index d3c807f0..90d6db1f 100644
--- a/crates/paimon/src/vindex/mod.rs
+++ b/crates/paimon/src/vindex/mod.rs
@@ -40,6 +40,7 @@ const DEFAULT_NLIST: &str = "256";
const DEFAULT_PQ_M: &str = "16";
const DEFAULT_PQ_USE_OPQ: &str = "false";
const DEFAULT_TRAIN_SAMPLE_RATIO: f64 = 1.0;
+const VINDEX_BUILD_GRANULE_ENABLED: &str = "vindex.build.granule.enabled";
const VECTOR_SEARCH_TIMING_ENV: &str = "PAIMON_LOG_VECTOR_SEARCH_TIMING";
const DISKANN_OPTION_KEYS: &[(&str, &str)] = &[
("deployment-profile", "deployment-profile"),
@@ -114,6 +115,7 @@ pub(crate) struct VindexVectorIndexOptions {
pub config: VectorIndexConfig,
pub native_options: HashMap<String, String>,
pub train_sample_ratio: f64,
+ pub granule_build_enabled: bool,
}
impl VindexVectorIndexOptions {
@@ -247,10 +249,12 @@ impl VindexVectorIndexOptions {
})?;
let train_sample_ratio =
resolve_train_sample_ratio(table_options, user_options,
index_type, field.name())?;
+ let granule_build_enabled =
resolve_granule_build_enabled(table_options, user_options)?;
Ok(Self {
config,
native_options,
train_sample_ratio,
+ granule_build_enabled,
})
}
@@ -309,6 +313,9 @@ fn is_supported_user_option_key(key: &str, index_type:
&str, field_name: &str) -
if key == "index.type" {
return true;
}
+ if key == VINDEX_BUILD_GRANULE_ENABLED {
+ return index_type != DISKANN_IDENTIFIER;
+ }
if is_allowed_native_key(key, index_type) {
return true;
}
@@ -326,6 +333,25 @@ fn is_supported_user_option_key(key: &str, index_type:
&str, field_name: &str) -
false
}
+fn resolve_granule_build_enabled(
+ table_options: &HashMap<String, String>,
+ user_options: &HashMap<String, String>,
+) -> crate::Result<bool> {
+ match user_options
+ .get(VINDEX_BUILD_GRANULE_ENABLED)
+ .or_else(|| table_options.get(VINDEX_BUILD_GRANULE_ENABLED))
+ {
+ Some(value) => value
+ .parse::<bool>()
+ .map_err(|_| crate::Error::ConfigInvalid {
+ message: format!(
+ "Invalid vindex option
{VINDEX_BUILD_GRANULE_ENABLED}='{value}'; expected true or false"
+ ),
+ }),
+ None => Ok(true),
+ }
+}
+
fn is_allowed_native_key(key: &str, index_type: &str) -> bool {
match key {
"dimension" | "metric" => true,
@@ -926,6 +952,53 @@ mod tests {
assert_eq!(diskann.train_sample_ratio, 0.5);
}
+ #[test]
+ fn test_vindex_options_granule_build_enabled() {
+ let field = array_float_field();
+ let defaults = VindexVectorIndexOptions::new(
+ &HashMap::new(),
+ &HashMap::new(),
+ IVF_FLAT_IDENTIFIER,
+ &field,
+ )
+ .unwrap();
+ assert!(defaults.granule_build_enabled);
+
+ let table_options = HashMap::from([(
+ VINDEX_BUILD_GRANULE_ENABLED.to_string(),
+ "false".to_string(),
+ )]);
+ let disabled = VindexVectorIndexOptions::new(
+ &table_options,
+ &HashMap::new(),
+ IVF_FLAT_IDENTIFIER,
+ &field,
+ )
+ .unwrap();
+ assert!(!disabled.granule_build_enabled);
+
+ let user_options =
+ HashMap::from([(VINDEX_BUILD_GRANULE_ENABLED.to_string(),
"true".to_string())]);
+ let overridden = VindexVectorIndexOptions::new(
+ &table_options,
+ &user_options,
+ IVF_FLAT_IDENTIFIER,
+ &field,
+ )
+ .unwrap();
+ assert!(overridden.granule_build_enabled);
+
+ let invalid =
+ HashMap::from([(VINDEX_BUILD_GRANULE_ENABLED.to_string(),
"yes".to_string())]);
+ assert!(VindexVectorIndexOptions::new(
+ &HashMap::new(),
+ &invalid,
+ IVF_FLAT_IDENTIFIER,
+ &field,
+ )
+ .is_err());
+ }
+
#[test]
fn test_vindex_options_reject_invalid_train_sample_ratio() {
for value in ["0", "-0.1", "1.1", "NaN", "inf", "not-a-number"] {
@@ -1049,6 +1122,7 @@ mod tests {
),
(IVF_FLAT_IDENTIFIER, "diskann.max-degree"),
(DISKANN_IDENTIFIER, "diskann.nlist"),
+ (DISKANN_IDENTIFIER, VINDEX_BUILD_GRANULE_ENABLED),
(DISKANN_IDENTIFIER, "diskann.ivf.coarse-assignment"),
(
DISKANN_IDENTIFIER,
diff --git a/docs/src/sql.md b/docs/src/sql.md
index 7598ee53..7848f91c 100644
--- a/docs/src/sql.md
+++ b/docs/src/sql.md
@@ -1264,6 +1264,7 @@ Supported vindex options:
| `<index-type>.ivf.coarse-assignment` | `auto` | all IVF types | Build-time
list assignment mode: `auto` uses Vamana for large centroid matrices; `exact`
always uses exact assignment. |
| `<index-type>.ivf.train.max-points-per-centroid` | `256` | all IVF types |
Positive limit on coarse K-means training data: at most `nlist × value`
vectors. |
| `<index-type>.train.sample-ratio` or `fields.<field>.train.sample-ratio` |
`1.0` | all vindex types | Fraction of shard rows selected evenly for training.
Must be in `(0, 1]`; all rows are still added to the index. The field-specific
option takes precedence. |
+| `vindex.build.granule.enabled` | `true` | all IVF types | Use the granule
build pipeline. Set to `false` to use the full-spill build path for future
index builds. |
| `<index-type>.pq.m` | `16` | `ivf-pq` | Number of product-quantization
sub-vectors. The dimension must be divisible by this value. |
| `<index-type>.pq.use-opq` | `false` | `ivf-pq` | Whether to enable OPQ
before PQ encoding. |
| `ivf-pq.ivf.pq-encoding` | `auto` | `ivf-pq` | Build-time PQ encoding mode:
`auto` selects an accelerated backend when supported; `canonical` uses the
canonical encoder. |
