diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 583b76285..45e920c91 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -67,6 +67,125 @@ impl ParquetFormatReader { } } +pub(crate) async fn has_beneficial_offset_index( + reader: Box, + file_size: u64, + column_name: &str, + sample_ranges: &[RowRange], + scan_ranges: &[RowRange], +) -> crate::Result { + let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); + let mut reader = ArrowFileReader::new(file_size, reader.into()); + let metadata = reader.get_metadata(Some(&options)).await?; + Ok(metadata_has_beneficial_offset_index( + &metadata, + column_name, + sample_ranges, + scan_ranges, + )) +} + +fn metadata_has_beneficial_offset_index( + metadata: &ParquetMetaData, + column_name: &str, + sample_ranges: &[RowRange], + scan_ranges: &[RowRange], +) -> bool { + 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::>(); + let Some(offset_index) = metadata.offset_index() else { + return false; + }; + if columns.is_empty() || offset_index.len() != metadata.row_groups().len() { + return false; + } + let mut sample_selection = build_row_ranges_selection(metadata.row_groups(), sample_ranges); + let mut scan_selection = build_row_ranges_selection(metadata.row_groups(), scan_ranges); + let mut checked = false; + let mut sample_bytes = 0u64; + let mut scan_bytes = 0u64; + 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 false; + }; + let sample_row_group_selection = sample_selection.split_off(row_count); + let scan_row_group_selection = scan_selection.split_off(row_count); + let sample_selected = sample_row_group_selection.selects_any(); + let scan_selected = scan_row_group_selection.selects_any(); + checked |= sample_selected; + if !sample_selected && !scan_selected { + continue; + } + let mut sample_byte_ranges = Vec::new(); + let mut scan_byte_ranges = Vec::new(); + for index in &columns { + let column = row_group.column(*index); + let Some(page_locations) = indexes.get(*index).map(|index| index.page_locations()) + else { + return false; + }; + let Some(first_page) = page_locations.first() else { + return false; + }; + let Ok(column_start) = u64::try_from( + column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()), + ) else { + return false; + }; + let Ok(first_page_offset) = u64::try_from(first_page.offset) else { + return false; + }; + for (selection, selected_ranges) in [ + (&sample_row_group_selection, &mut sample_byte_ranges), + (&scan_row_group_selection, &mut scan_byte_ranges), + ] { + if !selection.selects_any() { + continue; + } + if column_start < first_page_offset { + selected_ranges.push(column_start..first_page_offset); + } + selected_ranges.extend(selection.scan_ranges(page_locations)); + } + } + for (selected_ranges, total_bytes) in [ + (sample_byte_ranges, &mut sample_bytes), + (scan_byte_ranges, &mut scan_bytes), + ] { + let Some(group_bytes) = merge_byte_ranges(&selected_ranges, RANGE_COALESCE_BYTES) + .into_iter() + .try_fold(0u64, |total, range| { + total.checked_add(range.end.checked_sub(range.start)?) + }) + else { + return false; + }; + let Some(total) = total_bytes.checked_add(group_bytes) else { + return false; + }; + *total_bytes = total; + } + } + // Sparse training is followed by a full scan, so require it to skip at + // least half of the projected bytes instead of accepting marginal savings. + checked && sample_bytes <= scan_bytes / 2 +} + enum ParquetRowGroupMessage { Batch(RecordBatch), Error(Error), @@ -487,12 +606,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 @@ -500,39 +619,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::>(); + let projected_bytes = selected_row_groups + .iter() + .map(|(_, _, projected_bytes)| *projected_bytes) .collect::>(); 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; }; @@ -557,6 +687,7 @@ impl FormatFileReader for ParquetFormatReader { row_group_mask, row_group_index, batch_size, + selection, permit, batch_tx, )); @@ -660,6 +791,7 @@ async fn read_row_group( projection: ProjectionMask, row_group_index: usize, batch_size: Option, + selection: Option, _permit: ParquetReadPermit, sender: mpsc::Sender, ) { @@ -669,6 +801,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); } @@ -2248,8 +2383,8 @@ fn split_ranges_for_concurrency(merged: Vec>, concurrency: usize) -> mod tests { use super::build_parquet_row_filter; use super::{ - forward_row_group_batches, FilePredicates, ParquetFormatReader, ParquetFormatWriter, - ParquetRowGroupMessage, + forward_row_group_batches, metadata_has_beneficial_offset_index, FilePredicates, + ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage, }; use super::{ AsyncArrowWriter, Bytes, PageIndexPolicy, ParquetMetaDataReader, Predicate, @@ -2832,6 +2967,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(ParquetReadBudget::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::>() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect::>() + }) + .collect::>(); + let expected = ranges + .iter() + .flat_map(|range| range.from() as i32..=range.to() as i32) + .collect::>(); + + 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 = ParquetReadBudget::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; @@ -2921,7 +3163,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(ParquetReadBudget::new(8, 256 * 1024 * 1024).unwrap()); budget.enable_diagnostics(); let file_size = data.len() as u64; @@ -2945,7 +3187,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] @@ -3384,11 +3626,13 @@ mod tests { row_group_rows: usize, total_rows: i32, statistics: EnabledStatistics, + offset_index_disabled: bool, ) -> Vec { 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(); @@ -3404,7 +3648,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()); @@ -3432,19 +3676,90 @@ mod tests { assert_eq!(selection.row_count(), 10); } + #[tokio::test] + async fn test_sparse_row_selection_requires_offset_index_and_page_savings() { + let bytes = write_multi_row_group_parquet(10, 30, EnabledStatistics::Chunk, false).await; + let metadata = load_metadata_with_page_index(&bytes, true); + assert!( + !metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 19)], + &[RowRange::new(0, 29)], + ), + "reading two of three row groups is not sufficiently sparse" + ); + + let bytes = write_multi_page_parquet(10, 80).await; + let metadata = load_metadata_with_page_index(&bytes, true); + assert!( + !metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(20, 38)], + &[RowRange::new(20, 39)], + ), + "near-full reads within one shard must not use the whole file as the baseline" + ); + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 69)], + &[RowRange::new(0, 79)], + )); + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 0), RowRange::new(79, 79)], + &[RowRange::new(0, 79)], + )); + + let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; + let metadata = load_metadata_with_page_index(&bytes, true); + + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 19)], + &[RowRange::new(0, 19)], + )); + let sparse_ranges = [RowRange::new(0, 0)]; + assert!(metadata_has_beneficial_offset_index( + &metadata, + "value", + &sparse_ranges, + &[RowRange::new(0, 19)], + )); + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "missing", + &sparse_ranges, + &[RowRange::new(0, 19)], + )); + let bytes_without_index = + write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; + let metadata_without_index = load_metadata_with_page_index(&bytes_without_index, true); + assert!(!metadata_has_beneficial_offset_index( + &metadata_without_index, + "value", + &sparse_ranges, + &[RowRange::new(0, 19)], + )); + } + #[tokio::test] async fn test_row_group_selection_in_fails_open_on_unusable_stats() { 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 @@ -3513,6 +3828,26 @@ mod tests { struct TrackingFileRead { data: Bytes, ranges: Arc>>>, + resident_bytes: Arc, + peak_resident_bytes: Arc, + } + + struct TrackedReadBuffer { + data: Box<[u8]>, + resident_bytes: Arc, + } + + 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 { @@ -3520,6 +3855,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)), } } @@ -3541,7 +3878,135 @@ mod tests { impl crate::io::FileRead for TrackingFileRead { async fn read(&self, range: std::ops::Range) -> crate::Result { 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::>(); + for cancel in [false, true] { + let tracker = TrackingFileRead::new(data.clone()); + let budget = Arc::new(ParquetReadBudget::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::()); + } + 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/arrow/parquet_read_budget.rs b/crates/paimon/src/arrow/parquet_read_budget.rs index bef9c6ebd..a11111759 100644 --- a/crates/paimon/src/arrow/parquet_read_budget.rs +++ b/crates/paimon/src/arrow/parquet_read_budget.rs @@ -108,6 +108,10 @@ impl ParquetReadBudget { self.parallelism } + pub(crate) fn max_inflight_bytes(&self) -> u64 { + self.max_inflight_bytes + } + pub(crate) fn enable_diagnostics(&self) { self.diagnostics.enabled.store(true, Ordering::Relaxed); } 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 cf7e910e5..dcda47217 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 { + 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, +) -> Result { DataSplitBuilder::new() .with_snapshot(shard.snapshot_id) .with_partition(shard.partition.clone()) @@ -30,10 +40,7 @@ pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result( index_column: &str, dimension: usize, expected_row_id: &mut i64, +) -> Result> { + 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> { + 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> { let vector_index = batch .schema() @@ -163,21 +237,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)?; diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index 3b4a1fb84..f4d35e126 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -17,8 +17,14 @@ use crate::spec::{CoreOptions, DataField, ManifestEntry}; use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; -use crate::table::RowRange; -use crate::Result; +use crate::table::{merge_row_ranges, RowRange}; +use crate::{Error, Result}; + +use super::validation::checked_row_count; + +// Keep samples short enough to avoid storage-order bias; fall back before range I/O explodes. +const MAX_IVF_TRAINING_RANGE_ROWS: usize = 128; +const MAX_IVF_TRAINING_RANGES: usize = 4_096; pub(crate) type VindexIndexShard = VectorIndexShard; @@ -45,3 +51,78 @@ pub(super) fn plan_vindex_shards( "vindex", ) } + +pub(super) fn plan_ivf_training_ranges( + shard: &VindexIndexShard, + training_rows: usize, +) -> Result>> { + let shard_rows = usize::try_from(checked_row_count( + shard.row_range_start, + shard.row_range_end, + )?) + .map_err(|error| Error::DataInvalid { + message: "vindex shard row count does not fit usize".to_string(), + source: Some(Box::new(error)), + })?; + if training_rows == 0 || training_rows > shard_rows { + return Err(Error::DataInvalid { + message: format!( + "Invalid IVF training row count: {training_rows}; shard contains {shard_rows} rows" + ), + source: None, + }); + } + let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS); + if range_count == 1 || range_count > MAX_IVF_TRAINING_RANGES { + return Ok(None); + } + let seed = ivf_training_seed(shard); + let mut cursor = shard.row_range_start; + let mut ranges = Vec::with_capacity(range_count); + + for range_index in 0..range_count { + let stratum_length = + shard_rows / range_count + usize::from(range_index < shard_rows % range_count); + let length = + training_rows / range_count + usize::from(range_index < training_rows % range_count); + debug_assert!(length <= stratum_length); + let available_offsets = stratum_length - length + 1; + let offset = mix_seed(seed ^ range_index as u64) as usize % available_offsets; + let start = checked_add_offset(cursor, offset, "training range")?; + let end = checked_add_offset(start, length - 1, "training range")?; + ranges.push(RowRange::new(start, end)); + cursor = checked_add_offset(cursor, stratum_length, "training stratum")?; + } + + Ok(Some(merge_row_ranges(ranges))) +} + +pub(super) fn ivf_training_seed(shard: &VindexIndexShard) -> u64 { + let mut seed = mix_seed( + (shard.snapshot_id as u64) + ^ (shard.row_range_start as u64).rotate_left(21) + ^ (shard.row_range_end as u64).rotate_left(42) + ^ (shard.source_bucket as u64).rotate_left(11), + ); + for byte in &shard.partition_bytes { + seed = mix_seed(seed ^ u64::from(*byte)); + } + seed +} + +fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { + let offset = i64::try_from(offset).map_err(|error| Error::DataInvalid { + message: format!("vindex {name} offset does not fit i64"), + source: Some(Box::new(error)), + })?; + value.checked_add(offset).ok_or_else(|| Error::DataInvalid { + message: format!("vindex {name} offset overflows i64"), + source: None, + }) +} + +fn mix_seed(mut value: u64) -> u64 { + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) +} 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 a4636b578..5c11add03 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::validate_vector_batch; -use super::planning::{plan_vindex_shards, VindexIndexShard}; +use super::extraction::{validate_vector_batch, validate_vector_batch_ranges}; +use super::planning::{plan_ivf_training_ranges, plan_vindex_shards, VindexIndexShard}; use super::validation::{ checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, find_index_field, validate_vector_field, @@ -143,6 +143,130 @@ fn test_planner_splits_single_file_across_shards() { ); } +#[test] +fn test_ivf_training_ranges_are_bounded_and_exact() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(100), 1_000))], + 1_000, + ) + .unwrap() + .remove(0); + + for training_rows in [129, 200, 899] { + let ranges = plan_ivf_training_ranges(&shard, training_rows) + .unwrap() + .expect("sparse ranges expected"); + + assert!(ranges.len() <= 64); + assert_eq!( + ranges.iter().map(RowRange::count).sum::(), + training_rows as i64 + ); + assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); + assert!(ranges.first().unwrap().from() >= shard.row_range_start); + assert!(ranges.last().unwrap().to() <= shard.row_range_end); + assert_eq!( + ranges, + plan_ivf_training_ranges(&shard, training_rows) + .unwrap() + .expect("sparse ranges expected") + ); + } + + let ranges = plan_ivf_training_ranges(&shard, 200) + .unwrap() + .expect("sparse ranges expected"); + let mut other_snapshot = shard.clone(); + other_snapshot.snapshot_id += 1; + assert_ne!( + ranges, + plan_ivf_training_ranges(&other_snapshot, 200) + .unwrap() + .expect("sparse ranges expected") + ); +} + +#[test] +fn test_ivf_training_ranges_fall_back_for_a_single_range_sample() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(100), 1_000))], + 1_000, + ) + .unwrap() + .remove(0); + + for rows in [1, 127, 128] { + assert!(plan_ivf_training_ranges(&shard, rows).unwrap().is_none()); + } + let ranges = plan_ivf_training_ranges(&shard, 129).unwrap().unwrap(); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 129); + assert!(ranges + .iter() + .all(|range| range.from() >= 100 && range.to() < 1_100)); + assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); +} + +#[test] +fn test_ivf_training_ranges_keep_segments_short() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 1_000_000))], + 1_000_000, + ) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 65_536) + .unwrap() + .expect("sparse ranges expected"); + + assert_eq!(ranges.len(), 512); + assert!(ranges.iter().all(|range| range.count() <= 128)); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 65_536); +} + +#[test] +fn test_ivf_training_ranges_scale_to_keep_segments_short() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 2_000_000))], + 2_000_000, + ) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 262_144) + .unwrap() + .expect("sparse ranges expected"); + + assert_eq!(ranges.len(), 2_048); + assert!(ranges.iter().all(|range| range.count() <= 128)); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 262_144); +} + +#[test] +fn test_ivf_training_ranges_fall_back_above_range_limit() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 2_000_000))], + 2_000_000, + ) + .unwrap() + .remove(0); + + assert!(plan_ivf_training_ranges(&shard, 524_289).unwrap().is_none()); +} + +#[test] +fn test_ivf_training_ranges_cover_full_shard_without_empty_sentinel() { + let shard = plan(vec![manifest_entry(data_file("a", Some(0), 1_000))], 1_000) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 1_000) + .unwrap() + .expect("full range expected"); + + assert_eq!(ranges, vec![RowRange::new(0, 999)]); +} + #[test] fn test_planner_rejects_missing_first_row_id() { let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) @@ -252,6 +376,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_sparse_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_sparse_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)]); @@ -559,26 +747,79 @@ async fn vindex_second_build_without_new_data_is_noop() { #[tokio::test] async fn vindex_incremental_build_indexes_only_new_rows() { let table_path = "memory:/test_vindex_incremental"; - let table = vindex_e2e_table(table_path, "10"); + let mut options = vindex_e2e_options("2048"); + options.insert("ivf-flat.dimension".to_string(), "4096".to_string()); + options.insert("ivf-flat.nlist".to_string(), "3".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; // Build #1 over the initial batch via a real end-to-end build. write_vectors( &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + (0..1024).collect(), + (0..1024) + .map(|id| { + (0..4096) + .map(|component| (id * 4096 + component) as f32) + .collect() + }) + .collect(), ) .await; - let first_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .with_options(HashMap::from([( - "ivf-flat.train.sample-ratio".to_string(), - "0.9".to_string(), - )])) - .execute() + let build_options = + HashMap::from([("ivf-flat.train.sample-ratio".to_string(), "0.2".to_string())]); + let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); + let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) .await .unwrap(); + let core_options = CoreOptions::new(table.schema().options()); + let shards = plan_vindex_shards( + table_path, + table.schema().partition_keys(), + table.schema().fields(), + &core_options, + snapshot.id(), + entries, + core_options.global_index_row_count_per_shard().unwrap(), + &[], + ) + .unwrap(); + assert_eq!(shards.len(), 1); + let options = crate::vindex::VindexVectorIndexOptions::new( + table.schema().options(), + &build_options, + IVF_FLAT_IDENTIFIER, + find_index_field(&table, "embedding").unwrap(), + ) + .unwrap(); + let training_rows = paimon_vindex_core::autotune::default_training_vector_count( + checked_training_vector_count( + (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, + options.train_sample_ratio, + ) + .unwrap(), + options.config.nlist(), + ) + .unwrap(); + let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + builder + .with_index_column("embedding") + .with_options(build_options); + assert!(builder + .sparse_training_ranges(&shards[0], "embedding", training_rows) + .await + .unwrap() + .is_some()); + let first_built = builder.execute().await.unwrap(); assert!(first_built > 0, "first build must index the initial rows"); // First appended row-id, derived from the data manifest (never hard-coded). @@ -596,18 +837,62 @@ async fn vindex_incremental_build_indexes_only_new_rows() { // Append a second batch (new row-ids [n..]). write_vectors( &table, - vec![4, 5, 6], - vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], + vec![1024, 1025, 1026], + (1024..1027) + .map(|id| { + (0..4096) + .map(|component| (id * 4096 + component) as f32) + .collect() + }) + .collect(), ) .await; - // End-to-end: build #2 must SUCCEED and index the appended rows. - let second_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() + let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) .await .unwrap(); + let shards = plan_vindex_shards( + table_path, + table.schema().partition_keys(), + table.schema().fields(), + &core_options, + snapshot.id(), + entries, + core_options.global_index_row_count_per_shard().unwrap(), + &indexed_coverage, + ) + .unwrap(); + assert_eq!(shards.len(), 1); + let options = crate::vindex::VindexVectorIndexOptions::new( + table.schema().options(), + &HashMap::new(), + IVF_FLAT_IDENTIFIER, + find_index_field(&table, "embedding").unwrap(), + ) + .unwrap(); + let training_rows = paimon_vindex_core::autotune::default_training_vector_count( + checked_training_vector_count( + (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, + options.train_sample_ratio, + ) + .unwrap(), + options.config.nlist(), + ) + .unwrap(); + let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + builder.with_index_column("embedding"); + assert!(builder + .sparse_training_ranges(&shards[0], "embedding", training_rows) + .await + .unwrap() + .is_none()); + // End-to-end: build #2 must SUCCEED and index the appended rows. + let second_built = builder.execute().await.unwrap(); assert!(second_built > 0, "appended rows must be indexed"); let all_files = latest_vindex_index_files(&table).await; @@ -643,6 +928,113 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } +#[tokio::test] +async fn vindex_sparse_probe_errors_fall_back() { + let table_path = "memory:/test_vindex_probe_fallback"; + let table = vindex_e2e_table(table_path, "1024"); + let mut shard = plan( + vec![manifest_entry(data_file("broken.parquet", Some(0), 1024))], + 1024, + ) + .unwrap() + .remove(0); + shard.bucket_path = table_path.to_string(); + let builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + assert!(plan_ivf_training_ranges(&shard, 192).unwrap().is_some()); + + // A failed open and an unreadable footer both disable only the optimization. + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .unwrap() + .is_none()); + let path = shard.files[0].data_file_path(&shard.bucket_path); + table + .file_io() + .new_output(&path) + .unwrap() + .write(vec![0; 128].into()) + .await + .unwrap(); + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .unwrap() + .is_none()); + + // Invalid source metadata is not an optional probe failure. + shard.files[0].file_size = -1; + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .is_err()); +} + +#[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(|id| (900..1000).contains(id)), + "{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 3221dc4c1..26f7039dd 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -41,9 +41,12 @@ 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) capability_check: Duration, + pub(super) sample_read: Duration, pub(super) train_finish: Duration, pub(super) raw_temp_reread: Duration, pub(super) index_add: Duration, + pub(super) full_scan_add: Duration, pub(super) serialize_upload: Duration, pub(super) rows: usize, pub(super) training_rows_seen: usize, @@ -58,17 +61,25 @@ pub(super) struct VectorIndexBuildTiming { impl VectorIndexBuildTiming { pub(super) fn log(self, index_type: &str, commit: Duration) { let total = self.total_without_commit.saturating_add(commit); - let accounted = self - .source_batch_wait - .saturating_add(self.raw_temp_write) - .saturating_add(self.train_finish) - .saturating_add(self.raw_temp_reread) - .saturating_add(self.index_add) + let build = self + .capability_check + .saturating_add(if self.raw_temp_bytes != 0 { + self.source_batch_wait + .saturating_add(self.raw_temp_write) + .saturating_add(self.train_finish) + .saturating_add(self.raw_temp_reread) + .saturating_add(self.index_add) + } else { + self.sample_read + .saturating_add(self.train_finish) + .saturating_add(self.full_scan_add) + }); + let accounted = build .saturating_add(self.serialize_upload) .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_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms=0.000 full_scan_add_ms=0.000 pipeline_blocked_ms=0.000 producer_blocked_ms=0.000 consumer_add_ms=0.000 data_file_count={} data_file_read_concurrency=1 peak_ready_batches=0 total_ms={:.3} unattributed_ms={:.3}", + "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_write_ms={:.3} capability_check_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} data_file_count={} data_file_read_concurrency=1 total_ms={:.3} unattributed_ms={:.3}", index_type, self.file_name, self.rows, @@ -89,11 +100,14 @@ impl VectorIndexBuildTiming { self.parquet_projected_bytes_total, self.parquet_peak_inflight_row_groups, self.raw_temp_write.as_secs_f64() * 1000.0, + self.capability_check.as_secs_f64() * 1000.0, self.train_finish.as_secs_f64() * 1000.0, self.raw_temp_reread.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, + self.sample_read.as_secs_f64() * 1000.0, + self.full_scan_add.as_secs_f64() * 1000.0, self.data_file_count, total.as_secs_f64() * 1000.0, unattributed.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 803c08887..1dcfb705b 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -15,21 +15,26 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::{data_split_for_shard, validate_vector_batch}; -use super::planning::VindexIndexShard; +use super::extraction::{ + data_split_for_shard, data_split_for_shard_ranges, validate_vector_batch, + validate_vector_batch_ranges, +}; +use super::planning::{ivf_training_seed, plan_ivf_training_ranges, VindexIndexShard}; use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, }; use super::VindexIndexBuildBuilder; +use crate::arrow::format::parquet::has_beneficial_offset_index; 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::table::RowRange; +use crate::vindex::{VindexVectorIndexOptions, DISKANN_IDENTIFIER}; use crate::{Error, Result}; use arrow_buffer::MutableBuffer; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use paimon_vindex_core::autotune::default_training_vector_count; use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::PosWriter; @@ -41,12 +46,130 @@ use tokio_util::io::SyncIoBridge; const INDEX_DIR: &str = "index"; const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024; +const READ_ADD_QUEUE_CAPACITY: usize = 2; pub(super) struct BuiltIndexFile { pub(super) meta: IndexFileMeta, pub(super) timing: Option, } impl<'a> VindexIndexBuildBuilder<'a> { + /// Plans sparse training reads, falling back when the optional file probe fails. + pub(super) async fn sparse_training_ranges( + &self, + shard: &VindexIndexShard, + index_column: &str, + training_rows_retained: usize, + ) -> Result>> { + 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 mut sparse_ranges = + if training_rows_retained > 0 && training_rows_retained < row_count_usize { + plan_ivf_training_ranges(shard, training_rows_retained)? + } else { + None + }; + if let Some(ranges) = sparse_ranges.as_ref() { + let mut checks = Vec::new(); + let mut usable = true; + for file in &shard.files { + 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 { + usable = false; + break; + }; + let local_ranges = ranges + .iter() + .filter_map(|range| { + let from = range.from().max(file_start); + let to = range.to().min(file_end); + (from <= to).then(|| RowRange::new(from - file_start, to - file_start)) + }) + .collect::>(); + if local_ranges.is_empty() { + continue; + } + let scan_ranges = vec![RowRange::new( + shard.row_range_start.max(file_start) - file_start, + shard.row_range_end.min(file_end) - file_start, + )]; + let path = file.data_file_path(&shard.bucket_path); + if !path.to_ascii_lowercase().ends_with(".parquet") { + usable = false; + break; + } + 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)), + })?; + checks.push((path, file_size, local_ranges, scan_ranges)); + } + let found_vector_file = !checks.is_empty(); + if usable && found_vector_file { + let concurrency = self + .table + .schema() + .core_options() + .parquet_row_group_parallelism()?; + let file_io = self.table.file_io(); + let mut checks = futures::stream::iter(checks) + .map(|(path, file_size, local_ranges, scan_ranges)| async move { + let input = file_io.new_input(&path)?; + let reader = Box::new(input.reader().await?); + has_beneficial_offset_index( + reader, + file_size, + index_column, + &local_ranges, + &scan_ranges, + ) + .await + }) + .buffer_unordered(concurrency); + while let Some(result) = checks.next().await { + match result { + Ok(true) => {} + Ok(false) => { + usable = false; + break; + } + Err(error) => { + log::warn!( + "vindex sparse training probe failed for column '{}' in shard [{}, {}]; falling back to a full scan: {}", + index_column, + shard.row_range_start, + shard.row_range_end, + error, + ); + return Ok(None); + } + } + } + } + if !usable || !found_vector_file { + log::warn!( + "vindex sparse training read is unavailable for column '{}' in shard [{}, {}]; falling back to a full scan (beneficial_offset_indexes={usable}, found_vector_file={found_vector_file})", + index_column, + shard.row_range_start, + shard.row_range_end, + ); + sparse_ranges = None; + } + } + Ok(sparse_ranges) + } + pub(super) async fn build_index_file( &self, shard: &VindexIndexShard, @@ -87,245 +210,460 @@ impl<'a> VindexIndexBuildBuilder<'a> { let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; let training_vector_count = checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let training_buffer_rows = - (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); - let training_buffer_floats = training_buffer_rows - .checked_mul(dimension_usize) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; + let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { + 0 + } else { + // This only gates sparse reads; errors must preserve the full-scan fallback. + default_training_vector_count(training_vector_count, options.config.nlist()) + .unwrap_or(0) + }; + let capability_start = timing_enabled.then(Instant::now); + let sparse_ranges = self + .sparse_training_ranges(shard, index_column, training_rows_retained) + .await?; + let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); + if let Some(budget) = parquet_read_budget.as_ref() { + eprintln!( + "event=paimon_vector_index_build_plan snapshot_id={} row_range_start={} row_range_end={} source_bucket={} sparse={} training_seed={} training_range_count={} parquet_row_group_parallelism={} parquet_max_inflight_bytes={}", + shard.snapshot_id, + shard.row_range_start, + shard.row_range_end, + shard.source_bucket, + sparse_ranges.is_some(), + ivf_training_seed(shard), + sparse_ranges.as_ref().map_or(0, Vec::len), + budget.parallelism(), + budget.max_inflight_bytes(), + ); + } 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 raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { - message: format!("Failed to create temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - let mut raw_file = tokio::fs::File::from_std(raw_file); - let split = data_split_for_shard(shard)?; - 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.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let mut expected_row_id = shard.row_range_start; - let mut rows_seen = 0usize; + let mut sample_read = Duration::ZERO; + let mut full_scan_add = Duration::ZERO; + let raw_temp_reread; + let index_add; + let train_finish; let mut bytes_written = 0usize; - let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(training_buffer_floats); - - loop { - let source_start = timing_enabled.then(Instant::now); - let batch = batches.try_next().await?; - if let Some(source_start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); - } - let Some(batch) = batch else { break }; - batch_count += 1; - let vectors = - validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; - let batch_end = - rows_seen - .checked_add(vectors.row_count) - .ok_or_else(|| Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), - source: None, - })?; + let training_rows_seen; - if training_vector_count == row_count_usize { + let writer = if let Some(ranges) = sparse_ranges { + let sample_start = timing_enabled.then(Instant::now); + let split = data_split_for_shard_ranges(shard, ranges.clone())?; + 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.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let mut range_index = 0usize; + let mut expected_row_id = ranges[0].from(); + let mut rows_seen = 0usize; + while let Some(batch) = batches.try_next().await? { + let vectors = validate_vector_batch_ranges( + &batch, + index_column, + dimension_usize, + &ranges, + &mut range_index, + &mut expected_row_id, + )?; + rows_seen = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training row count overflows usize".to_string(), + source: None, + })?; trainer .add_training_vectors_mut(vectors.values, vectors.row_count) .map_err(|e| Error::DataInvalid { message: format!("Failed to add vindex training vectors: {e}"), source: Some(Box::new(e)), })?; - } else { - while next_training_sample < training_vector_count { - let sample_row = checked_training_sample_index( - next_training_sample, - row_count_usize, - training_vector_count, - )?; - if sample_row >= batch_end { - break; - } - let start = (sample_row - rows_seen) * dimension_usize; - training_buffer - .extend_from_slice(&vectors.values[start..start + dimension_usize]); - next_training_sample += 1; - if training_buffer.len() == training_buffer_floats { - trainer - .add_training_vectors_mut( - &training_buffer, - training_buffer.len() / dimension_usize, - ) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - training_buffer.clear(); - } - } } + if rows_seen != training_rows_retained || range_index != ranges.len() { + return Err(Error::DataInvalid { + message: format!( + "vindex sparse training data mismatch: rows={rows_seen}/{training_rows_retained}, ranges={range_index}/{}", + ranges.len() + ), + source: None, + }); + } + training_rows_seen = rows_seen; + sample_read = sample_start.map_or(Duration::ZERO, |start| start.elapsed()); - let raw_write_start = timing_enabled.then(Instant::now); - raw_file - .write_all(vectors.bytes) + let train_start = timing_enabled.then(Instant::now); + let training = tokio::task::spawn_blocking(move || trainer.finish()) .await .map_err(|e| Error::UnexpectedError { - message: format!("Failed to spill vindex vectors: {e}"), + 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)), })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let writer = VectorIndexWriter::new(training); + + let split = data_split_for_shard(shard)?; + 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.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let expected_end = + shard + .row_range_end + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range end overflows i64".to_string(), + source: None, + })?; + let index_column = index_column.to_string(); + let row_range_start = shard.row_range_start; + let (sender, mut receiver) = + tokio::sync::mpsc::channel::(READ_ADD_QUEUE_CAPACITY); + let full_scan_start = timing_enabled.then(Instant::now); + let consumer = tokio::task::spawn_blocking(move || -> Result<_> { + let mut writer = writer; + let mut expected_row_id = row_range_start; + let mut rows_added = 0usize; + let mut batches_added = 0usize; + let mut ids = Vec::new(); + while let Some(batch) = receiver.blocking_recv() { + let vectors = validate_vector_batch( + &batch, + &index_column, + dimension_usize, + &mut expected_row_id, + )?; + let batch_end = rows_added.checked_add(vectors.row_count).ok_or_else(|| { + Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + } + })?; + ids.clear(); + for row in rows_added..batch_end { + ids.push(i64::try_from(row).map_err(|e| Error::DataInvalid { + message: "vindex row id does not fit i64".to_string(), + source: Some(Box::new(e)), + })?); + } + 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 = batch_end; + batches_added += 1; + } + Ok((writer, rows_added, expected_row_id, batches_added)) + }); + + let mut producer_error = None; + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await; + if let Some(start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(start.elapsed()); + } + let batch = match batch { + Ok(Some(batch)) => batch, + Ok(None) => break, + Err(error) => { + producer_error = Some(error); + break; + } + }; + if sender.send(batch).await.is_err() { + break; + } + } + drop(sender); + let consumer_result = consumer.await; + let (writer, rows_added, next_row_id, batches_added) = + consumer_result.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })??; + if let Some(error) = producer_error { + return Err(error); } - bytes_written = bytes_written - .checked_add(vectors.bytes.len()) + if rows_added != row_count_usize || next_row_id != expected_end { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_added}/{row_count_usize}, next_row_id={next_row_id}/{expected_end}" + ), + source: None, + }); + } + batch_count = batches_added; + full_scan_add = full_scan_start.map_or(Duration::ZERO, |start| start.elapsed()); + raw_temp_reread = Duration::ZERO; + index_add = Duration::ZERO; + writer + } else { + let training_buffer_rows = + (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); + let training_buffer_floats = training_buffer_rows + .checked_mul(dimension_usize) .ok_or_else(|| Error::DataInvalid { - message: "vindex spilled byte count overflows usize".to_string(), + message: "vindex training buffer length overflows usize".to_string(), source: None, })?; - rows_seen = batch_end; - } + let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { + message: format!("Failed to create temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let mut raw_file = tokio::fs::File::from_std(raw_file); + let split = data_split_for_shard(shard)?; + 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.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let mut expected_row_id = shard.row_range_start; + let mut rows_seen = 0usize; + let mut next_training_sample = 0usize; + let mut training_buffer = Vec::with_capacity(training_buffer_floats); - if !training_buffer.is_empty() { - trainer - .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension_usize) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } - if rows_seen != row_count_usize - || expected_row_id - != shard + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await?; + if let Some(source_start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); + } + let Some(batch) = batch else { break }; + batch_count += 1; + let vectors = validate_vector_batch( + &batch, + index_column, + dimension_usize, + &mut expected_row_id, + )?; + let batch_end = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + })?; + + if training_vector_count == row_count_usize { + trainer + .add_training_vectors_mut(vectors.values, vectors.row_count) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } else { + while next_training_sample < training_vector_count { + let sample_row = checked_training_sample_index( + next_training_sample, + row_count_usize, + training_vector_count, + )?; + if sample_row >= batch_end { + break; + } + let start = (sample_row - rows_seen) * dimension_usize; + training_buffer + .extend_from_slice(&vectors.values[start..start + dimension_usize]); + next_training_sample += 1; + if training_buffer.len() == training_buffer_floats { + trainer + .add_training_vectors_mut( + &training_buffer, + training_buffer.len() / dimension_usize, + ) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + training_buffer.clear(); + } + } + } + + let raw_write_start = timing_enabled.then(Instant::now); + raw_file + .write_all(vectors.bytes) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to spill vindex vectors: {e}"), + source: Some(Box::new(e)), + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + bytes_written = + bytes_written + .checked_add(vectors.bytes.len()) + .ok_or_else(|| Error::DataInvalid { + message: "vindex spilled byte count overflows usize".to_string(), + source: None, + })?; + rows_seen = batch_end; + } + + if !training_buffer.is_empty() { + trainer + .add_training_vectors_mut( + &training_buffer, + training_buffer.len() / dimension_usize, + ) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } + let expected_end = + shard .row_range_end .checked_add(1) .ok_or_else(|| Error::DataInvalid { message: "vindex row range end overflows i64".to_string(), source: None, - })? - || (training_vector_count != row_count_usize - && next_training_sample != training_vector_count) - || bytes_written != expected_bytes - { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" - ), - source: None, - }); - } - let raw_write_start = timing_enabled.then(Instant::now); - raw_file.flush().await.map_err(|e| Error::UnexpectedError { - message: format!("Failed to flush temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - let raw_file_len = raw_file - .metadata() - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to inspect temporary vindex vector file: {e}"), + })?; + if rows_seen != row_count_usize + || expected_row_id != expected_end + || (training_vector_count != row_count_usize + && next_training_sample != training_vector_count) + || bytes_written != expected_bytes + { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" + ), + source: None, + }); + } + training_rows_seen = training_vector_count; + let raw_write_start = timing_enabled.then(Instant::now); + raw_file.flush().await.map_err(|e| Error::UnexpectedError { + message: format!("Failed to flush temporary vindex vector file: {e}"), source: Some(Box::new(e)), - })? - .len(); - if raw_file_len != expected_bytes as u64 { - return Err(Error::DataInvalid { - message: format!( - "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" - ), - source: None, - }); - } - let raw_file = raw_file.into_std().await; - // Diagnostics only: never fail the build for a timing log field. - let training_rows_retained = if timing_enabled { - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) - } else { - 0 - }; + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + let raw_file_len = raw_file + .metadata() + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to inspect temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })? + .len(); + if raw_file_len != expected_bytes as u64 { + return Err(Error::DataInvalid { + message: format!( + "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" + ), + source: None, + }); + } + let raw_file = raw_file.into_std().await; - let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( - move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { - let train_start = timing_enabled.then(Instant::now); - let training = trainer.finish()?; - let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let mut writer = VectorIndexWriter::new(training); - let mut raw_temp_reread = Duration::ZERO; - let mut index_add = Duration::ZERO; - let mut raw_file = raw_file; - let reread_start = timing_enabled.then(Instant::now); - raw_file.seek(SeekFrom::Start(0))?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - let batch_rows = training_buffer_rows.min(row_count_usize); - let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; - let mut buffer = MutableBuffer::new(batch_bytes); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let result = tokio::task::spawn_blocking( + move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { + let train_start = timing_enabled.then(Instant::now); + let training = trainer.finish()?; + let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let mut writer = VectorIndexWriter::new(training); + let mut raw_temp_reread = Duration::ZERO; + let mut index_add = Duration::ZERO; + let mut raw_file = raw_file; let reread_start = timing_enabled.then(Instant::now); - raw_file.read_exact(buffer.as_slice_mut())?; + raw_file.seek(SeekFrom::Start(0))?; if let Some(start) = reread_start { raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); } - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); + let batch_rows = training_buffer_rows.min(row_count_usize); + let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; + let mut buffer = MutableBuffer::new(batch_bytes); + let mut ids = Vec::with_capacity(batch_rows); + let mut rows_added = 0usize; + while rows_added < row_count_usize { + let rows = batch_rows.min(row_count_usize - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let reread_start = timing_enabled.then(Instant::now); + raw_file.read_exact(buffer.as_slice_mut())?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + ids.clear(); + for row in rows_added..rows_added + rows { + ids.push(i64::try_from(row).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "vindex row id does not fit i64", + ) + })?); + } + let add_start = timing_enabled.then(Instant::now); + writer.add_vectors(&ids, buffer.typed_data::(), rows)?; + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + rows_added += rows; } - let add_start = timing_enabled.then(Instant::now); - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - if let Some(start) = add_start { - index_add = index_add.saturating_add(start.elapsed()); + let mut trailing = [0u8; 1]; + let reread_start = timing_enabled.then(Instant::now); + if raw_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "temporary vindex vector file contains trailing bytes", + )); } - rows_added += rows; - } - let mut trailing = [0u8; 1]; - let reread_start = timing_enabled.then(Instant::now); - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - Ok((writer, train_finish, raw_temp_reread, index_add)) - }, - ) - .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 or add vectors to vindex index: {e}"), - source: Some(Box::new(e)), - })?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + Ok((writer, train_finish, raw_temp_reread, index_add)) + }, + ) + .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 or add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; + train_finish = result.1; + raw_temp_reread = result.2; + index_add = result.3; + result.0 + }; let serialize_upload_start = timing_enabled.then(Instant::now); self.table @@ -425,12 +763,15 @@ impl<'a> VindexIndexBuildBuilder<'a> { parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, raw_temp_write, + capability_check, + sample_read, train_finish, raw_temp_reread, index_add, + full_scan_add, serialize_upload, rows: row_count_usize, - training_rows_seen: training_vector_count, + training_rows_seen, training_rows_retained, batch_count, raw_temp_bytes: bytes_written,