diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index d4ddd5846..7b3b125e8 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -95,6 +95,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>, +} + +pub(crate) async fn parquet_granules( + reader: Box, + file_size: u64, + column_name: &str, + page_index_enabled: bool, +) -> crate::Result<(Vec, 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::>(); + 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::>>() + 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 { + 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), @@ -562,12 +697,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 @@ -575,39 +710,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; }; @@ -632,6 +778,7 @@ impl FormatFileReader for ParquetFormatReader { row_group_mask, row_group_index, batch_size, + selection, permit, batch_tx, )); @@ -735,6 +882,7 @@ async fn read_row_group( projection: ProjectionMask, row_group_index: usize, batch_size: Option, + selection: Option, _permit: ReadPermit, sender: mpsc::Sender, ) { @@ -744,6 +892,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); } @@ -2264,6 +2415,13 @@ fn merge_byte_ranges(ranges: &[Range], coalesce: u64) -> Vec> { merged } +pub(crate) fn coalesced_parquet_range_bytes(ranges: &[Range]) -> 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. @@ -2323,8 +2481,8 @@ fn split_ranges_for_concurrency(merged: Vec>, 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, @@ -2908,6 +3066,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::>() + .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 = 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; @@ -2997,7 +3262,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; @@ -3021,7 +3286,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] @@ -3460,11 +3725,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(); @@ -3480,7 +3747,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()); @@ -3513,14 +3780,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 @@ -3585,6 +3852,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::(), + 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::(), + 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::>() + ); + } + async fn write_page_pruning_io_parquet() -> Vec { const ROWS: i32 = 1024; let schema = Arc::new(ArrowSchema::new(vec![ @@ -3682,6 +4021,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 { @@ -3689,6 +4048,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)), } } @@ -3710,7 +4071,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(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::()); + } + 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 9e06ff9a8..6b255713f 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 cf7e910e5..37a8bb515 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,24 +40,97 @@ pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result { 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> { + 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> { + 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 +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 000000000..ebd2b8b3d --- /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>, +} + +#[derive(Debug)] +pub(super) struct GranulePlan { + pub(super) first: Vec, + pub(super) rest: Vec, + first_rows: usize, + // None preserves full-spill's global stride sampling when first covers the shard. + training: Option>, +} + +struct SpillRecord { + ids: Vec, + bytes: Vec, +} + +enum AddItem { + Batch(RecordBatch, Vec), + Spilled(Vec, MutableBuffer), +} + +type SpillTask = oneshot::Receiver>; +type ConsumerTask = oneshot::Receiver>; +type TrainingTask = JoinHandle>; +type ReplayTask = JoinHandle>; + +fn spawn_channel_worker(name: &'static str, worker: F) -> Result> +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, + 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, + consumer: ConsumerTask, + replay: ReplayTask, + spill_bytes: u64, + spill_write: Duration, +} + +fn spawn_spill_writer(timing_enabled: bool) -> Result { + 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::(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, + index_column: String, + dimension: usize, + timing_enabled: bool, +) -> Result { + 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::(); + 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, + 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 + '_ { + (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> { + let mut units: Vec> = 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 { + (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>) -> u64 { + let mut files = HashMap::>>::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 { + // 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, + file_index: usize, + file_start: i64, + shard_range: &RowRange, + file_granules: Vec, +) -> 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> { + 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, + index_column: &str, + read_timing: Option<&Arc>, + parquet_read_budget: Option<&Arc>, + ) -> Result { + 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 { + 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::(); + 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, + ) -> Result { + 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 = 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 = 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) -> Vec { + 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::>(); + 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::>(); + 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::(), + 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::>() + }; + 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::>(); + 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::>(); + 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 a4636b578..fb5c14990 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 3221dc4c1..ae43568f2 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_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={} 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_groups={} raw_temp_write_ms={:.3} granule_spill_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} granule_spill_read_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}", 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 803c08887..8c7fc1260 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, @@ -55,6 +56,44 @@ impl<'a> VindexIndexBuildBuilder<'a> { index_field_id: i32, options: &VindexVectorIndexOptions, index_meta: Vec, + ) -> Result { + 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, ) -> Result { 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, + row_count: i64, + ) -> Result { 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 d3c807f06..90d6db1fc 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, 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, + user_options: &HashMap, +) -> crate::Result { + match user_options + .get(VINDEX_BUILD_GRANULE_ENABLED) + .or_else(|| table_options.get(VINDEX_BUILD_GRANULE_ENABLED)) + { + Some(value) => value + .parse::() + .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 28c292bbb..918c999e4 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1247,6 +1247,7 @@ Supported vindex options: | `.ivf.coarse-assignment` | `auto` | all IVF types | Build-time list assignment mode: `auto` uses Vamana for large centroid matrices; `exact` always uses exact assignment. | | `.ivf.train.max-points-per-centroid` | `256` | all IVF types | Positive limit on coarse K-means training data: at most `nlist × value` vectors. | | `.train.sample-ratio` or `fields..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. | | `.pq.m` | `16` | `ivf-pq` | Number of product-quantization sub-vectors. The dimension must be divisible by this value. | | `.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. |