From f2061fdda203f7566dc528e159c97f9e5141b6c3 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 9 Sep 2026 17:28:43 +0800 Subject: [PATCH 01/24] feat: optimize IVF index build reads --- crates/paimon/src/arrow/format/parquet.rs | 68 +- .../vindex_index_build_builder/extraction.rs | 98 ++- .../vindex_index_build_builder/planning.rs | 82 ++- .../table/vindex_index_build_builder/tests.rs | 61 +- .../vindex_index_build_builder/timing.rs | 30 +- .../vindex_index_build_builder/writer.rs | 681 ++++++++++++------ 6 files changed, 782 insertions(+), 238 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 758fb36e3..1008a4b56 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -67,6 +67,47 @@ impl ParquetFormatReader { } } +pub(crate) async fn has_usable_offset_index( + reader: Box, + file_size: u64, + column_name: &str, +) -> crate::Result { + let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); + let mut reader = ArrowFileReader::new(file_size, reader.into()); + let metadata = reader.get_metadata(Some(&options)).await?; + Ok(metadata_has_usable_offset_index(&metadata, column_name)) +} + +fn metadata_has_usable_offset_index(metadata: &ParquetMetaData, column_name: &str) -> bool { + let columns = metadata + .file_metadata() + .schema_descr() + .columns() + .iter() + .enumerate() + .filter_map(|(index, column)| { + column + .path() + .parts() + .first() + .is_some_and(|part| part == column_name) + .then_some(index) + }) + .collect::>(); + let Some(offset_index) = metadata.offset_index() else { + return false; + }; + !columns.is_empty() + && offset_index.len() == metadata.row_groups().len() + && offset_index.iter().all(|row_group| { + columns.iter().all(|index| { + row_group + .get(*index) + .is_some_and(|index| !index.page_locations().is_empty()) + }) + }) +} + enum ParquetRowGroupMessage { Batch(RecordBatch), Error(Error), @@ -2290,7 +2331,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, + forward_row_group_batches, metadata_has_usable_offset_index, parse_compression, + supported_compressions, FilePredicates, ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage, }; use super::{ @@ -3427,11 +3469,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(); @@ -3447,7 +3491,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()); @@ -3475,19 +3519,35 @@ mod tests { assert_eq!(selection.row_count(), 10); } + #[tokio::test] + async fn test_sparse_row_selection_requires_offset_index_for_projected_column() { + let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; + let metadata = load_metadata_with_page_index(&bytes, true); + + assert!(metadata_has_usable_offset_index(&metadata, "value")); + assert!(!metadata_has_usable_offset_index(&metadata, "missing")); + let bytes_without_index = + write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; + let metadata_without_index = load_metadata_with_page_index(&bytes_without_index, true); + assert!(!metadata_has_usable_offset_index( + &metadata_without_index, + "value" + )); + } + #[tokio::test] async fn test_row_group_selection_in_fails_open_on_unusable_stats() { let fields = vec![int_field("id"), int_field("value")]; let predicates = vec![id_leaf(PredicateOperator::In, vec![Datum::Int(100)])]; - let bytes = write_multi_row_group_parquet(10, 10, EnabledStatistics::None).await; + let bytes = write_multi_row_group_parquet(10, 10, EnabledStatistics::None, false).await; let metadata = load_metadata_with_page_index(&bytes, false); let selection = super::build_predicate_row_selection(metadata.row_groups(), &predicates, &fields) .unwrap(); assert!(selection.is_none(), "missing stats must fail open"); - let bytes = write_multi_row_group_parquet(10, 10, EnabledStatistics::Chunk).await; + let bytes = write_multi_row_group_parquet(10, 10, EnabledStatistics::Chunk, false).await; let metadata = load_metadata_with_page_index(&bytes, false); let mut damaged_row_group = metadata.row_groups()[0].clone(); let damaged_id_column = damaged_row_group 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..ca48a3f4c 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs @@ -23,6 +23,16 @@ use crate::{Error, Result}; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result { + data_split_for_shard_ranges( + shard, + vec![RowRange::new(shard.row_range_start, shard.row_range_end)], + ) +} + +pub(super) fn data_split_for_shard_ranges( + shard: &VindexIndexShard, + row_ranges: Vec, +) -> Result { DataSplitBuilder::new() .with_snapshot(shard.snapshot_id) .with_partition(shard.partition.clone()) @@ -30,10 +40,7 @@ pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result( index_column: &str, dimension: usize, expected_row_id: &mut i64, +) -> Result> { + validate_vector_batch_with(batch, index_column, dimension, |row_id| { + if row_id != *expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "vindex vector extraction expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + *expected_row_id = expected_row_id + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex expected row id overflows i64".to_string(), + source: None, + })?; + Ok(()) + }) +} + +pub(super) fn validate_vector_batch_ranges<'a>( + batch: &'a RecordBatch, + index_column: &str, + dimension: usize, + ranges: &[RowRange], + range_index: &mut usize, + expected_row_id: &mut i64, +) -> Result> { + validate_vector_batch_with(batch, index_column, dimension, |row_id| { + let range = ranges.get(*range_index).ok_or_else(|| Error::DataInvalid { + message: format!("vindex vector extraction got unexpected _ROW_ID {row_id}"), + source: None, + })?; + if row_id != *expected_row_id || row_id > range.to() { + return Err(Error::DataInvalid { + message: format!( + "vindex vector extraction expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + if row_id == range.to() { + *range_index += 1; + *expected_row_id = match ranges.get(*range_index) { + Some(next) => next.from(), + None => row_id.checked_add(1).ok_or_else(|| Error::DataInvalid { + message: "vindex expected row id overflows i64".to_string(), + source: None, + })?, + }; + } else { + *expected_row_id = row_id.checked_add(1).ok_or_else(|| Error::DataInvalid { + message: "vindex expected row id overflows i64".to_string(), + source: None, + })?; + } + Ok(()) + }) +} + +fn validate_vector_batch_with<'a>( + batch: &'a RecordBatch, + index_column: &str, + dimension: usize, + mut validate_row_id: impl FnMut(i64) -> Result<()>, ) -> Result> { let vector_index = batch .schema() @@ -163,21 +237,7 @@ pub(super) fn validate_vector_batch<'a>( }); } for row_id in row_ids.values() { - if *row_id != *expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "vindex vector extraction expected _ROW_ID {}, got {}", - expected_row_id, row_id - ), - source: None, - }); - } - *expected_row_id = expected_row_id - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex expected row id overflows i64".to_string(), - source: None, - })?; + validate_row_id(*row_id)?; } let byte_start = checked_vector_bytes(start, 1)?; diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index 3b4a1fb84..b9ac30d2d 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -18,7 +18,11 @@ use crate::spec::{CoreOptions, DataField, ManifestEntry}; use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; use crate::table::RowRange; -use crate::Result; +use crate::{Error, Result}; + +use super::validation::checked_row_count; + +const MAX_IVF_TRAINING_RANGES: usize = 64; pub(crate) type VindexIndexShard = VectorIndexShard; @@ -45,3 +49,79 @@ pub(super) fn plan_vindex_shards( "vindex", ) } + +pub(super) fn plan_ivf_training_ranges( + shard: &VindexIndexShard, + training_rows: usize, +) -> Result> { + let shard_rows = usize::try_from(checked_row_count( + shard.row_range_start, + shard.row_range_end, + )?) + .map_err(|error| Error::DataInvalid { + message: "vindex shard row count does not fit usize".to_string(), + source: Some(Box::new(error)), + })?; + if training_rows == 0 || training_rows > shard_rows { + return Err(Error::DataInvalid { + message: format!( + "Invalid IVF training row count: {training_rows}; shard contains {shard_rows} rows" + ), + source: None, + }); + } + if training_rows == shard_rows { + return Ok(Vec::new()); + } + + let range_count = training_rows.min(MAX_IVF_TRAINING_RANGES); + let gap_count = range_count + 1; + let skipped_rows = shard_rows - training_rows; + let seed = mix_seed( + (shard.snapshot_id as u64) + ^ (shard.row_range_start as u64).rotate_left(21) + ^ (shard.row_range_end as u64).rotate_left(42), + ); + let range_extra_offset = seed as usize % range_count; + let gap_extra_offset = seed.rotate_left(17) as usize % gap_count; + let mut cursor = shard.row_range_start; + let mut ranges = Vec::with_capacity(range_count); + + for gap_index in 0..range_count { + let gap = skipped_rows / gap_count + + usize::from( + (gap_index + gap_count - gap_extra_offset) % gap_count < skipped_rows % gap_count, + ); + cursor = checked_add_offset(cursor, gap, "training gap")?; + let length = training_rows / range_count + + usize::from( + (gap_index + range_count - range_extra_offset) % range_count + < training_rows % range_count, + ); + let end = checked_add_offset(cursor, length - 1, "training range")?; + ranges.push(RowRange::new(cursor, end)); + cursor = end.checked_add(1).ok_or_else(|| Error::DataInvalid { + message: "vindex training range end overflows i64".to_string(), + source: None, + })?; + } + + Ok(ranges) +} + +fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { + let offset = i64::try_from(offset).map_err(|error| Error::DataInvalid { + message: format!("vindex {name} offset does not fit i64"), + source: Some(Box::new(error)), + })?; + value.checked_add(offset).ok_or_else(|| Error::DataInvalid { + message: format!("vindex {name} offset overflows i64"), + source: None, + }) +} + +fn mix_seed(mut value: u64) -> u64 { + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs index a4636b578..23055cb45 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::validate_vector_batch; -use super::planning::{plan_vindex_shards, VindexIndexShard}; +use super::extraction::{validate_vector_batch, validate_vector_batch_ranges}; +use super::planning::{plan_ivf_training_ranges, plan_vindex_shards, VindexIndexShard}; use super::validation::{ checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, find_index_field, validate_vector_field, @@ -143,6 +143,25 @@ fn test_planner_splits_single_file_across_shards() { ); } +#[test] +fn test_ivf_training_ranges_are_bounded_and_exact() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(100), 1_000))], + 1_000, + ) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 200).unwrap(); + + assert_eq!(ranges.len(), 64); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 200); + assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); + assert!(ranges.first().unwrap().from() >= 100); + assert!(ranges.last().unwrap().to() <= 1_099); + assert_eq!(ranges, plan_ivf_training_ranges(&shard, 200).unwrap()); +} + #[test] fn test_planner_rejects_missing_first_row_id() { let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) @@ -252,6 +271,44 @@ fn test_extract_vectors_accepts_list_float32_and_row_ids() { assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); } +#[test] +fn test_sparse_vector_validation_accepts_gaps_across_batches() { + let ranges = vec![RowRange::new(10, 11), RowRange::new(15, 16)]; + let batches = [ + vector_batch( + vec![ + Some(vec![Some(1.0), Some(2.0)]), + Some(vec![Some(3.0), Some(4.0)]), + ], + vec![Some(10), Some(11)], + ), + vector_batch( + vec![ + Some(vec![Some(5.0), Some(6.0)]), + Some(vec![Some(7.0), Some(8.0)]), + ], + vec![Some(15), Some(16)], + ), + ]; + let mut range_index = 0; + let mut expected_row_id = ranges[0].from(); + + for batch in &batches { + validate_vector_batch_ranges( + batch, + "embedding", + 2, + &ranges, + &mut range_index, + &mut expected_row_id, + ) + .unwrap(); + } + + assert_eq!(range_index, ranges.len()); + assert_eq!(expected_row_id, 17); +} + #[test] fn test_extract_vectors_rejects_dimension_mismatch() { let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); 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..289e1a05e 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -41,9 +41,14 @@ 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) sample_read: Duration, pub(super) train_finish: Duration, pub(super) raw_temp_reread: Duration, pub(super) index_add: Duration, + pub(super) full_scan_add: Duration, + pub(super) pipeline_blocked: Duration, + pub(super) producer_blocked: Duration, + pub(super) consumer_add: Duration, pub(super) serialize_upload: Duration, pub(super) rows: usize, pub(super) training_rows_seen: usize, @@ -58,17 +63,23 @@ pub(super) struct VectorIndexBuildTiming { impl VectorIndexBuildTiming { pub(super) fn log(self, index_type: &str, commit: Duration) { let total = self.total_without_commit.saturating_add(commit); - let accounted = self - .source_batch_wait - .saturating_add(self.raw_temp_write) - .saturating_add(self.train_finish) - .saturating_add(self.raw_temp_reread) - .saturating_add(self.index_add) + let build = if self.raw_temp_bytes != 0 { + self.source_batch_wait + .saturating_add(self.raw_temp_write) + .saturating_add(self.train_finish) + .saturating_add(self.raw_temp_reread) + .saturating_add(self.index_add) + } else { + self.sample_read + .saturating_add(self.train_finish) + .saturating_add(self.full_scan_add) + }; + let accounted = build .saturating_add(self.serialize_upload) .saturating_add(commit); let unattributed = total.saturating_sub(accounted); eprintln!( - "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms=0.000 full_scan_add_ms=0.000 pipeline_blocked_ms=0.000 producer_blocked_ms=0.000 consumer_add_ms=0.000 data_file_count={} data_file_read_concurrency=1 peak_ready_batches=0 total_ms={:.3} unattributed_ms={:.3}", + "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} pipeline_blocked_ms={:.3} producer_blocked_ms={:.3} consumer_add_ms={:.3} 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, @@ -94,6 +105,11 @@ impl VectorIndexBuildTiming { self.index_add.as_secs_f64() * 1000.0, self.serialize_upload.as_secs_f64() * 1000.0, commit.as_secs_f64() * 1000.0, + self.sample_read.as_secs_f64() * 1000.0, + self.full_scan_add.as_secs_f64() * 1000.0, + self.pipeline_blocked.as_secs_f64() * 1000.0, + self.producer_blocked.as_secs_f64() * 1000.0, + self.consumer_add.as_secs_f64() * 1000.0, self.data_file_count, total.as_secs_f64() * 1000.0, unattributed.as_secs_f64() * 1000.0, diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs index 803c08887..3113ef6c2 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -15,18 +15,22 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::{data_split_for_shard, validate_vector_batch}; -use super::planning::VindexIndexShard; +use super::extraction::{ + data_split_for_shard, data_split_for_shard_ranges, validate_vector_batch, + validate_vector_batch_ranges, +}; +use super::planning::{plan_ivf_training_ranges, VindexIndexShard}; use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, }; use super::VindexIndexBuildBuilder; +use crate::arrow::format::parquet::has_usable_offset_index; use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME}; use crate::table::data_file_reader::DataFileReadTiming; use crate::table::table_read::configured_parquet_read_budget; -use crate::vindex::VindexVectorIndexOptions; +use crate::vindex::{VindexVectorIndexOptions, DISKANN_IDENTIFIER}; use crate::{Error, Result}; use arrow_buffer::MutableBuffer; use futures::TryStreamExt; @@ -87,245 +91,507 @@ impl<'a> VindexIndexBuildBuilder<'a> { let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; let training_vector_count = checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let training_buffer_rows = - (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); - let training_buffer_floats = training_buffer_rows - .checked_mul(dimension_usize) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; + let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { + 0 + } else { + default_training_vector_count(training_vector_count, options.config.nlist()).map_err( + |e| Error::DataInvalid { + message: format!("Failed to calculate IVF training vector count: {e}"), + source: Some(Box::new(e)), + }, + )? + }; + let mut sparse_ranges = + if training_rows_retained > 0 && training_rows_retained < row_count_usize { + Some(plan_ivf_training_ranges(shard, training_rows_retained)?) + } else { + None + }; + if sparse_ranges.is_some() { + let mut found_vector_file = false; + for file in shard.files.iter().filter(|file| { + file.write_cols + .as_ref() + .is_none_or(|columns| columns.iter().any(|column| column == index_column)) + }) { + found_vector_file = true; + let path = file.data_file_path(&shard.bucket_path); + if !path.to_ascii_lowercase().ends_with(".parquet") { + sparse_ranges = None; + break; + } + let file_size = u64::try_from(file.file_size).map_err(|e| Error::DataInvalid { + message: format!( + "Invalid data file size for '{}': {}", + file.file_name, file.file_size + ), + source: Some(Box::new(e)), + })?; + let input = self.table.file_io().new_input(&path)?; + if !has_usable_offset_index( + Box::new(input.reader().await?), + file_size, + index_column, + ) + .await? + { + sparse_ranges = None; + break; + } + } + if !found_vector_file { + sparse_ranges = None; + } + } let mut trainer = VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { message: format!("Failed to initialize vindex trainer: {e}"), source: Some(Box::new(e)), })?; - let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { - message: format!("Failed to create temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - let mut raw_file = tokio::fs::File::from_std(raw_file); - let split = data_split_for_shard(shard)?; - let mut read_builder = self.table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let read = match read_timing.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let mut expected_row_id = shard.row_range_start; - let mut rows_seen = 0usize; + let mut sample_read = Duration::ZERO; + let mut full_scan_add = Duration::ZERO; + let mut pipeline_blocked = Duration::ZERO; + let mut producer_blocked = Duration::ZERO; + let mut consumer_add = Duration::ZERO; + let raw_temp_reread; + let index_add; + let train_finish; let mut bytes_written = 0usize; - let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(training_buffer_floats); + let training_rows_seen; - loop { - let source_start = timing_enabled.then(Instant::now); - let batch = batches.try_next().await?; - if let Some(source_start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); - } - let Some(batch) = batch else { break }; - batch_count += 1; - let vectors = - validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; - let batch_end = - rows_seen - .checked_add(vectors.row_count) - .ok_or_else(|| Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), - source: None, - })?; - - if training_vector_count == row_count_usize { + let writer = if let Some(ranges) = sparse_ranges { + let sample_start = timing_enabled.then(Instant::now); + let split = data_split_for_shard_ranges(shard, ranges.clone())?; + let mut read_builder = self.table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let read = match read_timing.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let mut range_index = 0usize; + let mut expected_row_id = ranges[0].from(); + let mut rows_seen = 0usize; + while let Some(batch) = batches.try_next().await? { + let vectors = validate_vector_batch_ranges( + &batch, + index_column, + dimension_usize, + &ranges, + &mut range_index, + &mut expected_row_id, + )?; + rows_seen = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training row count overflows usize".to_string(), + source: None, + })?; trainer .add_training_vectors_mut(vectors.values, vectors.row_count) .map_err(|e| Error::DataInvalid { message: format!("Failed to add vindex training vectors: {e}"), source: Some(Box::new(e)), })?; - } else { - while next_training_sample < training_vector_count { - let sample_row = checked_training_sample_index( - next_training_sample, - row_count_usize, - training_vector_count, - )?; - if sample_row >= batch_end { - break; - } - let start = (sample_row - rows_seen) * dimension_usize; - training_buffer - .extend_from_slice(&vectors.values[start..start + dimension_usize]); - next_training_sample += 1; - if training_buffer.len() == training_buffer_floats { - trainer - .add_training_vectors_mut( - &training_buffer, - training_buffer.len() / dimension_usize, - ) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - training_buffer.clear(); - } - } } + if rows_seen != training_rows_retained || range_index != ranges.len() { + return Err(Error::DataInvalid { + message: format!( + "vindex sparse training data mismatch: rows={rows_seen}/{training_rows_retained}, ranges={range_index}/{}", + ranges.len() + ), + source: None, + }); + } + training_rows_seen = rows_seen; + sample_read = sample_start.map_or(Duration::ZERO, |start| start.elapsed()); - let raw_write_start = timing_enabled.then(Instant::now); - raw_file - .write_all(vectors.bytes) + let train_start = timing_enabled.then(Instant::now); + let training = tokio::task::spawn_blocking(move || trainer.finish()) .await .map_err(|e| Error::UnexpectedError { - message: format!("Failed to spill vindex vectors: {e}"), + message: format!("vindex training task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to train vindex index: {e}"), source: Some(Box::new(e)), })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let writer = VectorIndexWriter::new(training); + + let split = data_split_for_shard(shard)?; + let mut read_builder = self.table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let read = match read_timing.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let expected_end = + shard + .row_range_end + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range end overflows i64".to_string(), + source: None, + })?; + let index_column = index_column.to_string(); + let row_range_start = shard.row_range_start; + let (sender, mut receiver) = tokio::sync::mpsc::channel(2); + let full_scan_start = timing_enabled.then(Instant::now); + let consumer = tokio::task::spawn_blocking(move || -> Result<_> { + let mut writer = writer; + let mut expected_row_id = row_range_start; + let mut rows_added = 0usize; + let mut batches_added = 0usize; + let mut blocked = Duration::ZERO; + let mut add = Duration::ZERO; + let mut ids = Vec::new(); + loop { + let wait_start = timing_enabled.then(Instant::now); + let batch = receiver.blocking_recv(); + if let Some(start) = wait_start { + blocked = blocked.saturating_add(start.elapsed()); + } + let Some(batch) = batch else { break }; + let vectors = validate_vector_batch( + &batch, + &index_column, + dimension_usize, + &mut expected_row_id, + )?; + let batch_end = rows_added.checked_add(vectors.row_count).ok_or_else(|| { + Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + } + })?; + ids.clear(); + for row in rows_added..batch_end { + ids.push(i64::try_from(row).map_err(|e| Error::DataInvalid { + message: "vindex row id does not fit i64".to_string(), + source: Some(Box::new(e)), + })?); + } + let add_start = timing_enabled.then(Instant::now); + 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)), + })?; + if let Some(start) = add_start { + add = add.saturating_add(start.elapsed()); + } + rows_added = batch_end; + batches_added += 1; + } + if rows_added != row_count_usize || expected_row_id != expected_end { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_added}/{row_count_usize}, next_row_id={expected_row_id}/{expected_end}" + ), + source: None, + }); + } + Ok((writer, batches_added, blocked, add)) + }); + + let mut producer_error = None; + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await; + if let Some(start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(start.elapsed()); + } + let batch = match batch { + Ok(Some(batch)) => batch, + Ok(None) => break, + Err(error) => { + producer_error = Some(error); + break; + } + }; + let send_start = timing_enabled.then(Instant::now); + let send_result = sender.send(batch).await; + if let Some(start) = send_start { + producer_blocked = producer_blocked.saturating_add(start.elapsed()); + } + if send_result.is_err() { + break; + } } - bytes_written = bytes_written - .checked_add(vectors.bytes.len()) + drop(sender); + let consumer_result = consumer.await; + if let Some(error) = producer_error { + return Err(error); + } + let (writer, batches_added, blocked, add) = + consumer_result.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })??; + batch_count = batches_added; + pipeline_blocked = blocked; + consumer_add = add; + full_scan_add = full_scan_start.map_or(Duration::ZERO, |start| start.elapsed()); + raw_temp_reread = Duration::ZERO; + index_add = Duration::ZERO; + writer + } else { + let training_buffer_rows = + (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); + let training_buffer_floats = training_buffer_rows + .checked_mul(dimension_usize) .ok_or_else(|| Error::DataInvalid { - message: "vindex spilled byte count overflows usize".to_string(), + message: "vindex training buffer length overflows usize".to_string(), source: None, })?; - rows_seen = batch_end; - } + let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { + message: format!("Failed to create temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let mut raw_file = tokio::fs::File::from_std(raw_file); + let split = data_split_for_shard(shard)?; + let mut read_builder = self.table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let read = match read_timing.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let mut expected_row_id = shard.row_range_start; + let mut rows_seen = 0usize; + let mut next_training_sample = 0usize; + let mut training_buffer = Vec::with_capacity(training_buffer_floats); - if !training_buffer.is_empty() { - trainer - .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension_usize) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } - if rows_seen != row_count_usize - || expected_row_id - != shard + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await?; + if let Some(source_start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); + } + let Some(batch) = batch else { break }; + batch_count += 1; + let vectors = validate_vector_batch( + &batch, + index_column, + dimension_usize, + &mut expected_row_id, + )?; + let batch_end = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + })?; + + if training_vector_count == row_count_usize { + trainer + .add_training_vectors_mut(vectors.values, vectors.row_count) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } else { + while next_training_sample < training_vector_count { + let sample_row = checked_training_sample_index( + next_training_sample, + row_count_usize, + training_vector_count, + )?; + if sample_row >= batch_end { + break; + } + let start = (sample_row - rows_seen) * dimension_usize; + training_buffer + .extend_from_slice(&vectors.values[start..start + dimension_usize]); + next_training_sample += 1; + if training_buffer.len() == training_buffer_floats { + trainer + .add_training_vectors_mut( + &training_buffer, + training_buffer.len() / dimension_usize, + ) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + training_buffer.clear(); + } + } + } + + let raw_write_start = timing_enabled.then(Instant::now); + raw_file + .write_all(vectors.bytes) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to spill vindex vectors: {e}"), + source: Some(Box::new(e)), + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + bytes_written = + bytes_written + .checked_add(vectors.bytes.len()) + .ok_or_else(|| Error::DataInvalid { + message: "vindex spilled byte count overflows usize".to_string(), + source: None, + })?; + rows_seen = batch_end; + } + + if !training_buffer.is_empty() { + trainer + .add_training_vectors_mut( + &training_buffer, + training_buffer.len() / dimension_usize, + ) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } + let expected_end = + shard .row_range_end .checked_add(1) .ok_or_else(|| Error::DataInvalid { message: "vindex row range end overflows i64".to_string(), source: None, - })? - || (training_vector_count != row_count_usize - && next_training_sample != training_vector_count) - || bytes_written != expected_bytes - { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" - ), - source: None, - }); - } - let raw_write_start = timing_enabled.then(Instant::now); - raw_file.flush().await.map_err(|e| Error::UnexpectedError { - message: format!("Failed to flush temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - let raw_file_len = raw_file - .metadata() - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to inspect temporary vindex vector file: {e}"), + })?; + if rows_seen != row_count_usize + || expected_row_id != expected_end + || (training_vector_count != row_count_usize + && next_training_sample != training_vector_count) + || bytes_written != expected_bytes + { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" + ), + source: None, + }); + } + training_rows_seen = training_vector_count; + let raw_write_start = timing_enabled.then(Instant::now); + raw_file.flush().await.map_err(|e| Error::UnexpectedError { + message: format!("Failed to flush temporary vindex vector file: {e}"), source: Some(Box::new(e)), - })? - .len(); - if raw_file_len != expected_bytes as u64 { - return Err(Error::DataInvalid { - message: format!( - "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" - ), - source: None, - }); - } - let raw_file = raw_file.into_std().await; - // Diagnostics only: never fail the build for a timing log field. - let training_rows_retained = if timing_enabled { - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) - } else { - 0 - }; + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + let raw_file_len = raw_file + .metadata() + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to inspect temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })? + .len(); + if raw_file_len != expected_bytes as u64 { + return Err(Error::DataInvalid { + message: format!( + "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" + ), + source: None, + }); + } + let raw_file = raw_file.into_std().await; - let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( - move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { - let train_start = timing_enabled.then(Instant::now); - let training = trainer.finish()?; - let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let mut writer = VectorIndexWriter::new(training); - let mut raw_temp_reread = Duration::ZERO; - let mut index_add = Duration::ZERO; - let mut raw_file = raw_file; - let reread_start = timing_enabled.then(Instant::now); - raw_file.seek(SeekFrom::Start(0))?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - let batch_rows = training_buffer_rows.min(row_count_usize); - let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; - let mut buffer = MutableBuffer::new(batch_bytes); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let result = tokio::task::spawn_blocking( + move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { + let train_start = timing_enabled.then(Instant::now); + let training = trainer.finish()?; + let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let mut writer = VectorIndexWriter::new(training); + let mut raw_temp_reread = Duration::ZERO; + let mut index_add = Duration::ZERO; + let mut raw_file = raw_file; let reread_start = timing_enabled.then(Instant::now); - raw_file.read_exact(buffer.as_slice_mut())?; + raw_file.seek(SeekFrom::Start(0))?; if let Some(start) = reread_start { raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); } - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); + let batch_rows = training_buffer_rows.min(row_count_usize); + let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; + let mut buffer = MutableBuffer::new(batch_bytes); + let mut ids = Vec::with_capacity(batch_rows); + let mut rows_added = 0usize; + while rows_added < row_count_usize { + let rows = batch_rows.min(row_count_usize - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let reread_start = timing_enabled.then(Instant::now); + raw_file.read_exact(buffer.as_slice_mut())?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + ids.clear(); + for row in rows_added..rows_added + rows { + ids.push(i64::try_from(row).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "vindex row id does not fit i64", + ) + })?); + } + let add_start = timing_enabled.then(Instant::now); + writer.add_vectors(&ids, buffer.typed_data::(), rows)?; + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + rows_added += rows; } - let add_start = timing_enabled.then(Instant::now); - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - if let Some(start) = add_start { - index_add = index_add.saturating_add(start.elapsed()); + let mut trailing = [0u8; 1]; + let reread_start = timing_enabled.then(Instant::now); + if raw_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "temporary vindex vector file contains trailing bytes", + )); } - rows_added += rows; - } - let mut trailing = [0u8; 1]; - let reread_start = timing_enabled.then(Instant::now); - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - Ok((writer, train_finish, raw_temp_reread, index_add)) - }, - ) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("vindex training task failed: {e}"), - source: None, - })? - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to train or add vectors to vindex index: {e}"), - source: Some(Box::new(e)), - })?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + Ok((writer, train_finish, raw_temp_reread, index_add)) + }, + ) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex training task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to train or add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; + train_finish = result.1; + raw_temp_reread = result.2; + index_add = result.3; + result.0 + }; let serialize_upload_start = timing_enabled.then(Instant::now); self.table @@ -425,12 +691,17 @@ impl<'a> VindexIndexBuildBuilder<'a> { parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, raw_temp_write, + sample_read, train_finish, raw_temp_reread, index_add, + full_scan_add, + pipeline_blocked, + producer_blocked, + consumer_add, serialize_upload, rows: row_count_usize, - training_rows_seen: training_vector_count, + training_rows_seen, training_rows_retained, batch_count, raw_temp_bytes: bytes_written, From 3bc5ad50ce0513a2f48d4bbaddb8d0a07b91315c Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 12:52:30 +0800 Subject: [PATCH 02/24] perf(vindex): complete IVF build read optimization --- crates/paimon/src/arrow/format/parquet.rs | 55 +++- crates/paimon/src/io/file_io.rs | 3 + .../paimon/src/io/file_io/multipart_test.rs | 296 ++++++++++++++++++ crates/paimon/src/table/data_file_reader.rs | 17 + .../table/vindex_index_build_builder/tests.rs | 118 +++++++ .../vindex_index_build_builder/timing.rs | 44 ++- .../vindex_index_build_builder/writer.rs | 138 ++++++-- 7 files changed, 614 insertions(+), 57 deletions(-) create mode 100644 crates/paimon/src/io/file_io/multipart_test.rs diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 1008a4b56..ec1a85060 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -71,14 +71,23 @@ pub(crate) async fn has_usable_offset_index( reader: Box, file_size: u64, column_name: &str, + row_ranges: &[RowRange], ) -> crate::Result { let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); let mut reader = ArrowFileReader::new(file_size, reader.into()); let metadata = reader.get_metadata(Some(&options)).await?; - Ok(metadata_has_usable_offset_index(&metadata, column_name)) + Ok(metadata_has_usable_offset_index( + &metadata, + column_name, + row_ranges, + )) } -fn metadata_has_usable_offset_index(metadata: &ParquetMetaData, column_name: &str) -> bool { +fn metadata_has_usable_offset_index( + metadata: &ParquetMetaData, + column_name: &str, + row_ranges: &[RowRange], +) -> bool { let columns = metadata .file_metadata() .schema_descr() @@ -97,15 +106,31 @@ fn metadata_has_usable_offset_index(metadata: &ParquetMetaData, column_name: &st let Some(offset_index) = metadata.offset_index() else { return false; }; - !columns.is_empty() - && offset_index.len() == metadata.row_groups().len() - && offset_index.iter().all(|row_group| { - columns.iter().all(|index| { - row_group + if columns.is_empty() || offset_index.len() != metadata.row_groups().len() { + return false; + } + let mut row_start = 0i64; + let mut checked = false; + for (row_group, indexes) in metadata.row_groups().iter().zip(offset_index) { + let Some(row_end) = row_start.checked_add(row_group.num_rows()) else { + return false; + }; + if row_ranges + .iter() + .any(|range| range.from() < row_end && range.to() >= row_start) + { + checked = true; + if !columns.iter().all(|index| { + indexes .get(*index) .is_some_and(|index| !index.page_locations().is_empty()) - }) - }) + }) { + return false; + } + } + row_start = row_end; + } + checked } enum ParquetRowGroupMessage { @@ -3524,14 +3549,20 @@ mod tests { let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; let metadata = load_metadata_with_page_index(&bytes, true); - assert!(metadata_has_usable_offset_index(&metadata, "value")); - assert!(!metadata_has_usable_offset_index(&metadata, "missing")); + let ranges = [RowRange::new(0, 19)]; + assert!(metadata_has_usable_offset_index( + &metadata, "value", &ranges + )); + assert!(!metadata_has_usable_offset_index( + &metadata, "missing", &ranges + )); let bytes_without_index = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; let metadata_without_index = load_metadata_with_page_index(&bytes_without_index, true); assert!(!metadata_has_usable_offset_index( &metadata_without_index, - "value" + "value", + &ranges )); } diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs index 720225ebc..f6c50ea9f 100644 --- a/crates/paimon/src/io/file_io.rs +++ b/crates/paimon/src/io/file_io.rs @@ -34,6 +34,9 @@ use snafu::ResultExt; use tokio_util::compat::FuturesAsyncWriteCompatExt; use url::Url; +#[cfg(test)] +pub(crate) mod multipart_test; + use super::cache::{CachedFileReader, LocalCache}; use super::Storage; diff --git a/crates/paimon/src/io/file_io/multipart_test.rs b/crates/paimon/src/io/file_io/multipart_test.rs new file mode 100644 index 000000000..e6fa73edd --- /dev/null +++ b/crates/paimon/src/io/file_io/multipart_test.rs @@ -0,0 +1,296 @@ +// 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::*; +use crate::io::FileIOBuilder; +use opendal::raw::{ + oio, OpCopier, OpCopy, OpCreateDir, OpList, OpPresign, OpRead, OpRename, OpStat, OpWrite, + RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, +}; +use opendal::{Buffer, Capability, Metadata, OperationContext}; +use std::collections::BTreeMap; +use std::sync::Mutex; +use tokio::sync::Notify; + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub(crate) enum Fault { + #[default] + None, + Part, + Close, +} + +#[derive(Debug, Default)] +pub(crate) struct UploadState { + pub(crate) fault: Fault, + pub(crate) fail_on_index: usize, + pub(crate) index_writes: usize, + pub(crate) concurrency: Vec, + pub(crate) uploads: BTreeMap>, + pub(crate) completed_parts: Vec, + pub(crate) aborts: usize, +} + +/// Exercise OpenDAL's real multipart scheduler, storing completed objects in memory. +#[derive(Clone, Debug)] +pub(crate) struct MultipartProvider { + memory: Operator, + part_size: usize, + pub(crate) state: Arc>, +} + +impl MultipartProvider { + pub(crate) fn new(part_size: usize) -> Self { + Self { + memory: Operator::via_iter(opendal::services::MEMORY_SCHEME, []).unwrap(), + part_size, + state: Arc::default(), + } + } + + pub(crate) fn file_io(&self) -> FileIO { + FileIOBuilder::new("memory") + .build() + .unwrap() + .with_provider(Arc::new(self.clone())) + } +} + +#[async_trait::async_trait] +impl FileIOProvider for MultipartProvider { + async fn create(&self, path: &str) -> crate::Result<(Operator, String)> { + Ok(( + Operator::from_parts(OperationContext::default(), Arc::new(self.clone())), + Url::parse(path) + .unwrap() + .path() + .trim_start_matches('/') + .to_string(), + )) + } +} + +impl Service for MultipartProvider { + type Reader = oio::Reader; + type Writer = oio::Writer; + type Lister = oio::Lister; + type Deleter = oio::Deleter; + type Copier = oio::Copier; + + fn info(&self) -> ServiceInfo { + self.memory.service().info() + } + + fn capability(&self) -> Capability { + Capability { + write_can_multi: true, + write_multi_max_size: Some(self.part_size), + ..self.memory.service().capability() + } + } + + fn write( + &self, + ctx: &OperationContext, + path: &str, + args: OpWrite, + ) -> opendal::Result { + let mut state = self.state.lock().unwrap(); + state.concurrency.push(args.concurrent()); + if !path.ends_with(".index") { + return self.memory.service().write(ctx, path, args); + } + state.index_writes += 1; + let fault = if state.index_writes == state.fail_on_index { + state.fault + } else { + Fault::None + }; + Ok(Box::new(oio::MultipartWriter::new( + ctx.executor().clone(), + Upload { + provider: self.clone(), + path: path.to_string(), + fault, + reorder: args.concurrent() > 1, + second_part: Notify::new(), + }, + args.concurrent(), + ))) + } + + async fn create_dir( + &self, + ctx: &OperationContext, + path: &str, + args: OpCreateDir, + ) -> opendal::Result { + self.memory.service().create_dir(ctx, path, args).await + } + async fn stat( + &self, + ctx: &OperationContext, + path: &str, + args: OpStat, + ) -> opendal::Result { + self.memory.service().stat(ctx, path, args).await + } + fn read( + &self, + ctx: &OperationContext, + path: &str, + args: OpRead, + ) -> opendal::Result { + self.memory.service().read(ctx, path, args) + } + fn delete(&self, ctx: &OperationContext) -> opendal::Result { + self.memory.service().delete(ctx) + } + fn list( + &self, + ctx: &OperationContext, + path: &str, + args: OpList, + ) -> opendal::Result { + self.memory.service().list(ctx, path, args) + } + fn copy( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpCopy, + opts: OpCopier, + ) -> opendal::Result { + self.memory.service().copy(ctx, from, to, args, opts) + } + async fn rename( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpRename, + ) -> opendal::Result { + self.memory.service().rename(ctx, from, to, args).await + } + async fn presign( + &self, + ctx: &OperationContext, + path: &str, + args: OpPresign, + ) -> opendal::Result { + self.memory.service().presign(ctx, path, args).await + } +} + +struct Upload { + provider: MultipartProvider, + path: String, + fault: Fault, + reorder: bool, + second_part: Notify, +} + +fn injected_error() -> opendal::Error { + opendal::Error::new(opendal::ErrorKind::Unexpected, "injected multipart failure") +} + +impl oio::MultipartWrite for Upload { + async fn write_once(&self, size: u64, body: Buffer) -> opendal::Result { + assert_eq!( + self.fault, + Fault::None, + "failure test must exercise multipart" + ); + self.provider.memory.write(&self.path, body).await?; + Ok(Metadata::default().with_content_length(size)) + } + + async fn initiate_part(&self) -> opendal::Result { + self.provider + .state + .lock() + .unwrap() + .uploads + .insert(self.path.clone(), BTreeMap::new()); + Ok(self.path.clone()) + } + + async fn write_part( + &self, + upload_id: &str, + part_number: usize, + size: u64, + body: Buffer, + ) -> opendal::Result { + assert_eq!(size as usize, body.len()); + if self.reorder && part_number == 0 { + self.second_part.notified().await; + } + { + let mut state = self.provider.state.lock().unwrap(); + state.completed_parts.push(part_number); + state + .uploads + .get_mut(upload_id) + .unwrap() + .insert(part_number, body.to_bytes()); + } + if part_number == 1 { + self.second_part.notify_one(); + } + if self.fault == Fault::Part && part_number == 0 { + return Err(injected_error()); + } + Ok(oio::MultipartPart { + part_number, + etag: part_number.to_string(), + checksum: None, + size: Some(size), + }) + } + + async fn complete_part( + &self, + upload_id: &str, + parts: &[oio::MultipartPart], + ) -> opendal::Result { + if self.fault == Fault::Close { + return Err(injected_error()); + } + let mut bytes = Vec::new(); + { + let mut state = self.provider.state.lock().unwrap(); + let uploaded = state.uploads.remove(upload_id).unwrap(); + assert_eq!(parts.len(), uploaded.len()); + for (expected, part) in parts.iter().enumerate() { + assert_eq!(part.part_number, expected); + bytes.extend_from_slice(&uploaded[&part.part_number]); + } + } + let size = bytes.len() as u64; + self.provider.memory.write(&self.path, bytes).await?; + Ok(Metadata::default().with_content_length(size)) + } + + async fn abort_part(&self, upload_id: &str) -> opendal::Result<()> { + let mut state = self.provider.state.lock().unwrap(); + state.aborts += 1; + state.uploads.remove(upload_id); + Ok(()) + } +} diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 7019e9bc4..5f442e8eb 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -44,6 +44,8 @@ use std::time::{Duration, Instant}; #[derive(Debug, Default)] pub(crate) struct DataFileReadTiming { file_read_nanos: AtomicU64, + file_read_bytes: AtomicU64, + file_read_requests: AtomicU64, parquet_decode_nanos: AtomicU64, file_schema_open_nanos: AtomicU64, first_batch_wait_nanos: AtomicU64, @@ -56,6 +58,12 @@ impl DataFileReadTiming { .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); } + fn add_file_io(&self, bytes: usize) { + self.file_read_bytes + .fetch_add(bytes as u64, Ordering::Relaxed); + self.file_read_requests.fetch_add(1, Ordering::Relaxed); + } + fn add_parquet_decode(&self, duration: Duration) { self.parquet_decode_nanos .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); @@ -79,6 +87,13 @@ impl DataFileReadTiming { Duration::from_nanos(self.file_read_nanos.load(Ordering::Relaxed)) } + pub(crate) fn file_io(&self) -> (u64, u64) { + ( + self.file_read_bytes.load(Ordering::Relaxed), + self.file_read_requests.load(Ordering::Relaxed), + ) + } + pub(crate) fn parquet_decode(&self) -> Duration { Duration::from_nanos(self.parquet_decode_nanos.load(Ordering::Relaxed)) } @@ -102,6 +117,8 @@ impl FileRead for TimedFileRead { let start = Instant::now(); let result = self.inner.read(range).await; self.timing.add_file_read(start.elapsed()); + self.timing + .add_file_io(result.as_ref().map_or(0, bytes::Bytes::len)); result } } 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 23055cb45..a5813b762 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -309,6 +309,32 @@ fn test_sparse_vector_validation_accepts_gaps_across_batches() { assert_eq!(expected_row_id, 17); } +#[test] +fn test_sparse_vector_validation_rejects_bad_row_ids() { + let ranges = vec![RowRange::new(10, 11), RowRange::new(15, 16)]; + for row_ids in [ + vec![Some(10), Some(10)], + vec![Some(10), Some(15)], + vec![Some(9), Some(10)], + vec![Some(10), Some(11), Some(16), Some(15)], + vec![Some(10), Some(11), Some(15), Some(16), Some(17)], + ] { + let rows = row_ids.len(); + let batch = vector_batch(vec![Some(vec![Some(1.0), Some(2.0)]); rows], row_ids); + let mut range_index = 0; + let mut expected_row_id = ranges[0].from(); + assert!(validate_vector_batch_ranges( + &batch, + "embedding", + 2, + &ranges, + &mut range_index, + &mut expected_row_id, + ) + .is_err()); + } +} + #[test] fn test_extract_vectors_rejects_dimension_mismatch() { let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); @@ -700,6 +726,98 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } +#[tokio::test] +async fn vindex_upload_failure_preserves_committed_index() { + use crate::io::multipart_test::{Fault, MultipartProvider}; + + for fault in [Fault::Part, Fault::Close] { + let provider = MultipartProvider::new(128); + let table_path = "memory:/test_vindex_upload_failure"; + let table = test_table_with_io( + provider.file_io(), + table_path, + vindex_schema_builder(vindex_e2e_options("3")) + .build() + .unwrap(), + ); + setup_dirs(table.file_io(), table_path).await; + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .unwrap(); + let existing = latest_vindex_index_files(&table).await; + let old_path = format!("{table_path}/{INDEX_DIR}/{}", existing[0].file_name); + let old_bytes = table + .file_io() + .new_input(&old_path) + .unwrap() + .read() + .await + .unwrap(); + let mut search = table.new_vector_search_builder(); + search + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1); + let old_result = search.execute().await.unwrap(); + assert!(!old_result.is_empty()); + + write_vectors(&table, vec![4, 5, 6, 7, 8, 9], vec![vec![-1.0, 0.0]; 6]).await; + let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); + let before = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + { + let mut state = provider.state.lock().unwrap(); + state.fault = fault; + state.fail_on_index = state.index_writes + 2; + state.concurrency.clear(); + } + let error = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .expect_err("injected upload must fail"); + assert!( + error.to_string().contains("injected multipart failure"), + "{error}" + ); + let after = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + assert_eq!(before.id(), after.id()); + assert_eq!(before.index_manifest(), after.index_manifest()); + assert_eq!(latest_vindex_index_files(&table).await, existing); + assert_eq!( + table + .file_io() + .new_input(&old_path) + .unwrap() + .read() + .await + .unwrap(), + old_bytes + ); + assert_eq!(search.execute().await.unwrap(), old_result); + let files = table + .file_io() + .list_status(&format!("{table_path}/{INDEX_DIR}/")) + .await + .unwrap(); + assert_eq!(files.len(), 1); + assert!(table.file_io().exists(&old_path).await.unwrap()); + let state = provider.state.lock().unwrap(); + assert_eq!(state.concurrency, [1, 1]); + assert_eq!(state.uploads.len(), 1); + assert_eq!(state.aborts, 0); + } +} + #[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 289e1a05e..c391ad35f 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -31,6 +31,8 @@ pub(super) struct VectorIndexBuildTiming { pub(super) total_without_commit: Duration, pub(super) source_batch_wait: Duration, pub(super) oss_read: Duration, + pub(super) oss_read_bytes: u64, + pub(super) oss_range_requests: u64, pub(super) parquet_decode: Duration, pub(super) file_schema_open: Duration, pub(super) first_batch_wait: Duration, @@ -41,6 +43,7 @@ pub(super) struct VectorIndexBuildTiming { pub(super) parquet_projected_bytes_total: u64, pub(super) parquet_peak_inflight_row_groups: usize, pub(super) raw_temp_write: Duration, + pub(super) capability_check: Duration, pub(super) sample_read: Duration, pub(super) train_finish: Duration, pub(super) raw_temp_reread: Duration, @@ -48,12 +51,16 @@ pub(super) struct VectorIndexBuildTiming { pub(super) full_scan_add: Duration, pub(super) pipeline_blocked: Duration, pub(super) producer_blocked: Duration, - pub(super) consumer_add: Duration, + pub(super) consumer_validate_add: Duration, pub(super) serialize_upload: Duration, pub(super) rows: usize, pub(super) training_rows_seen: usize, pub(super) training_rows_retained: usize, pub(super) batch_count: usize, + pub(super) batch_bytes_min: usize, + pub(super) batch_bytes_max: usize, + pub(super) batch_bytes_total: usize, + pub(super) peak_ready_batches: usize, pub(super) raw_temp_bytes: usize, pub(super) index_bytes: u64, pub(super) data_file_count: usize, @@ -63,33 +70,40 @@ pub(super) struct VectorIndexBuildTiming { impl VectorIndexBuildTiming { pub(super) fn log(self, index_type: &str, commit: Duration) { let total = self.total_without_commit.saturating_add(commit); - let build = if self.raw_temp_bytes != 0 { - self.source_batch_wait - .saturating_add(self.raw_temp_write) - .saturating_add(self.train_finish) - .saturating_add(self.raw_temp_reread) - .saturating_add(self.index_add) - } else { - self.sample_read - .saturating_add(self.train_finish) - .saturating_add(self.full_scan_add) - }; + let build = self + .capability_check + .saturating_add(if self.raw_temp_bytes != 0 { + self.source_batch_wait + .saturating_add(self.raw_temp_write) + .saturating_add(self.train_finish) + .saturating_add(self.raw_temp_reread) + .saturating_add(self.index_add) + } else { + self.sample_read + .saturating_add(self.train_finish) + .saturating_add(self.full_scan_add) + }); let accounted = build .saturating_add(self.serialize_upload) .saturating_add(commit); let unattributed = total.saturating_sub(accounted); eprintln!( - "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} pipeline_blocked_ms={:.3} producer_blocked_ms={:.3} consumer_add_ms={:.3} 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={} batch_bytes_min={} batch_bytes_max={} batch_bytes_total={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} oss_read_bytes={} oss_range_requests={} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} capability_check_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} pipeline_blocked_ms={:.3} producer_blocked_ms={:.3} consumer_validate_add_ms={:.3} data_file_count={} data_file_read_concurrency=1 peak_ready_batches={} total_ms={:.3} unattributed_ms={:.3}", index_type, self.file_name, self.rows, self.training_rows_seen, self.training_rows_retained, self.batch_count, + self.batch_bytes_min, + self.batch_bytes_max, + self.batch_bytes_total, self.raw_temp_bytes, self.index_bytes, self.source_batch_wait.as_secs_f64() * 1000.0, self.oss_read.as_secs_f64() * 1000.0, + self.oss_read_bytes, + self.oss_range_requests, self.parquet_decode.as_secs_f64() * 1000.0, self.file_schema_open.as_secs_f64() * 1000.0, self.first_batch_wait.as_secs_f64() * 1000.0, @@ -100,6 +114,7 @@ impl VectorIndexBuildTiming { self.parquet_projected_bytes_total, self.parquet_peak_inflight_row_groups, self.raw_temp_write.as_secs_f64() * 1000.0, + self.capability_check.as_secs_f64() * 1000.0, self.train_finish.as_secs_f64() * 1000.0, self.raw_temp_reread.as_secs_f64() * 1000.0, self.index_add.as_secs_f64() * 1000.0, @@ -109,8 +124,9 @@ impl VectorIndexBuildTiming { self.full_scan_add.as_secs_f64() * 1000.0, self.pipeline_blocked.as_secs_f64() * 1000.0, self.producer_blocked.as_secs_f64() * 1000.0, - self.consumer_add.as_secs_f64() * 1000.0, + self.consumer_validate_add.as_secs_f64() * 1000.0, self.data_file_count, + self.peak_ready_batches, total.as_secs_f64() * 1000.0, unattributed.as_secs_f64() * 1000.0, ); diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs index 3113ef6c2..673e36445 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -30,10 +30,11 @@ use crate::arrow::format::parquet::has_usable_offset_index; use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME}; use crate::table::data_file_reader::DataFileReadTiming; use crate::table::table_read::configured_parquet_read_budget; +use crate::table::RowRange; use crate::vindex::{VindexVectorIndexOptions, DISKANN_IDENTIFIER}; use crate::{Error, Result}; use arrow_buffer::MutableBuffer; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use paimon_vindex_core::autotune::default_training_vector_count; use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::PosWriter; @@ -45,6 +46,7 @@ use tokio_util::io::SyncIoBridge; const INDEX_DIR: &str = "index"; const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024; +const READ_ADD_QUEUE_CAPACITY: usize = 2; pub(super) struct BuiltIndexFile { pub(super) meta: IndexFileMeta, pub(super) timing: Option, @@ -107,19 +109,36 @@ impl<'a> VindexIndexBuildBuilder<'a> { } else { None }; - if sparse_ranges.is_some() { - let mut found_vector_file = false; - for file in shard.files.iter().filter(|file| { - file.write_cols - .as_ref() - .is_none_or(|columns| columns.iter().any(|column| column == index_column)) - }) { - found_vector_file = true; + let capability_start = (timing_enabled && sparse_ranges.is_some()).then(Instant::now); + if let Some(ranges) = sparse_ranges.as_ref() { + let mut checks = Vec::new(); + let mut usable = true; + for file in &shard.files { let path = file.data_file_path(&shard.bucket_path); if !path.to_ascii_lowercase().ends_with(".parquet") { - sparse_ranges = None; + usable = false; break; } + let Some((file_start, file_end)) = file.row_id_range() else { + usable = false; + break; + }; + let local_ranges = ranges + .iter() + .filter_map(|range| { + let from = range.from().max(file_start); + let to = range.to().min(file_end); + (from <= to).then(|| RowRange::new(from - file_start, to - file_start)) + }) + .collect::>(); + if local_ranges.is_empty() + || file + .write_cols + .as_ref() + .is_some_and(|columns| !columns.iter().any(|column| column == index_column)) + { + continue; + } let file_size = u64::try_from(file.file_size).map_err(|e| Error::DataInvalid { message: format!( "Invalid data file size for '{}': {}", @@ -127,22 +146,40 @@ impl<'a> VindexIndexBuildBuilder<'a> { ), source: Some(Box::new(e)), })?; - let input = self.table.file_io().new_input(&path)?; - if !has_usable_offset_index( - Box::new(input.reader().await?), - file_size, - index_column, - ) - .await? - { - sparse_ranges = None; - break; + checks.push((path, file_size, local_ranges)); + } + let found_vector_file = !checks.is_empty(); + if usable && found_vector_file { + let concurrency = self + .table + .schema() + .core_options() + .parquet_row_group_parallelism()?; + let file_io = self.table.file_io(); + let mut checks = futures::stream::iter(checks) + .map(|(path, file_size, local_ranges)| async move { + let input = file_io.new_input(&path)?; + has_usable_offset_index( + Box::new(input.reader().await?), + file_size, + index_column, + &local_ranges, + ) + .await + }) + .buffer_unordered(concurrency); + while let Some(file_usable) = checks.try_next().await? { + if !file_usable { + usable = false; + break; + } } } - if !found_vector_file { + if !usable || !found_vector_file { sparse_ranges = None; } } + let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); let mut trainer = VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { @@ -153,7 +190,11 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut full_scan_add = Duration::ZERO; let mut pipeline_blocked = Duration::ZERO; let mut producer_blocked = Duration::ZERO; - let mut consumer_add = Duration::ZERO; + let mut consumer_validate_add = Duration::ZERO; + let mut batch_bytes_min = 0usize; + let mut batch_bytes_max = 0usize; + let mut batch_bytes_total = 0usize; + let mut peak_ready_batches = 0usize; let raw_temp_reread; let index_add; let train_finish; @@ -250,7 +291,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { })?; let index_column = index_column.to_string(); let row_range_start = shard.row_range_start; - let (sender, mut receiver) = tokio::sync::mpsc::channel(2); + let (sender, mut receiver) = + tokio::sync::mpsc::channel::(READ_ADD_QUEUE_CAPACITY); let full_scan_start = timing_enabled.then(Instant::now); let consumer = tokio::task::spawn_blocking(move || -> Result<_> { let mut writer = writer; @@ -258,7 +300,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut rows_added = 0usize; let mut batches_added = 0usize; let mut blocked = Duration::ZERO; - let mut add = Duration::ZERO; + let mut validate_add = Duration::ZERO; + let mut batch_bytes_min = usize::MAX; + let mut batch_bytes_max = 0usize; + let mut batch_bytes_total = 0usize; let mut ids = Vec::new(); loop { let wait_start = timing_enabled.then(Instant::now); @@ -267,6 +312,11 @@ impl<'a> VindexIndexBuildBuilder<'a> { blocked = blocked.saturating_add(start.elapsed()); } let Some(batch) = batch else { break }; + let validate_add_start = timing_enabled.then(Instant::now); + let batch_bytes = batch.get_array_memory_size(); + batch_bytes_min = batch_bytes_min.min(batch_bytes); + batch_bytes_max = batch_bytes_max.max(batch_bytes); + batch_bytes_total = batch_bytes_total.saturating_add(batch_bytes); let vectors = validate_vector_batch( &batch, &index_column, @@ -286,15 +336,14 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: Some(Box::new(e)), })?); } - let add_start = timing_enabled.then(Instant::now); 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)), })?; - if let Some(start) = add_start { - add = add.saturating_add(start.elapsed()); + if let Some(start) = validate_add_start { + validate_add = validate_add.saturating_add(start.elapsed()); } rows_added = batch_end; batches_added += 1; @@ -307,7 +356,19 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: None, }); } - Ok((writer, batches_added, blocked, add)) + Ok(( + writer, + batches_added, + blocked, + validate_add, + if batches_added == 0 { + 0 + } else { + batch_bytes_min + }, + batch_bytes_max, + batch_bytes_total, + )) }); let mut producer_error = None; @@ -333,20 +394,25 @@ impl<'a> VindexIndexBuildBuilder<'a> { if send_result.is_err() { break; } + peak_ready_batches = peak_ready_batches + .max(READ_ADD_QUEUE_CAPACITY.saturating_sub(sender.capacity())); } drop(sender); let consumer_result = consumer.await; if let Some(error) = producer_error { return Err(error); } - let (writer, batches_added, blocked, add) = + let (writer, batches_added, blocked, validate_add, min_bytes, max_bytes, total_bytes) = consumer_result.map_err(|e| Error::UnexpectedError { message: format!("vindex add task failed: {e}"), source: None, })??; batch_count = batches_added; pipeline_blocked = blocked; - consumer_add = add; + consumer_validate_add = validate_add; + batch_bytes_min = min_bytes; + batch_bytes_max = max_bytes; + batch_bytes_total = total_bytes; full_scan_add = full_scan_start.map_or(Duration::ZERO, |start| start.elapsed()); raw_temp_reread = Duration::ZERO; index_add = Duration::ZERO; @@ -669,6 +735,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { .map_or((Duration::ZERO, Duration::ZERO), |timing| { (timing.file_read(), timing.parquet_decode()) }); + let (oss_read_bytes, oss_range_requests) = read_timing + .as_ref() + .map_or((0, 0), |timing| timing.file_io()); let (file_schema_open, first_batch_wait, remaining_batch_wait) = read_timing .as_ref() .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing| { @@ -681,6 +750,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { total_without_commit: start.elapsed(), source_batch_wait, oss_read, + oss_read_bytes, + oss_range_requests, parquet_decode, file_schema_open, first_batch_wait, @@ -691,6 +762,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, raw_temp_write, + capability_check, sample_read, train_finish, raw_temp_reread, @@ -698,12 +770,16 @@ impl<'a> VindexIndexBuildBuilder<'a> { full_scan_add, pipeline_blocked, producer_blocked, - consumer_add, + consumer_validate_add, serialize_upload, rows: row_count_usize, training_rows_seen, training_rows_retained, batch_count, + batch_bytes_min, + batch_bytes_max, + batch_bytes_total, + peak_ready_batches, raw_temp_bytes: bytes_written, index_bytes: status.size, data_file_count: shard.files.len(), From fad54e31c771e6fbfb322f3f70e00654e0e09e6c Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 16:01:49 +0800 Subject: [PATCH 03/24] fix(vindex): stratify IVF training sampling --- .../vindex_index_build_builder/planning.rs | 45 +++++++++---------- .../table/vindex_index_build_builder/tests.rs | 30 ++++++++++--- 2 files changed, 43 insertions(+), 32 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index b9ac30d2d..c5f54b079 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -17,7 +17,7 @@ use crate::spec::{CoreOptions, DataField, ManifestEntry}; use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; -use crate::table::RowRange; +use crate::table::{merge_row_ranges, RowRange}; use crate::{Error, Result}; use super::validation::checked_row_count; @@ -75,38 +75,33 @@ pub(super) fn plan_ivf_training_ranges( } let range_count = training_rows.min(MAX_IVF_TRAINING_RANGES); - let gap_count = range_count + 1; - let skipped_rows = shard_rows - training_rows; - let seed = mix_seed( + let mut seed = mix_seed( (shard.snapshot_id as u64) ^ (shard.row_range_start as u64).rotate_left(21) - ^ (shard.row_range_end as u64).rotate_left(42), + ^ (shard.row_range_end as u64).rotate_left(42) + ^ (shard.source_bucket as u64).rotate_left(11), ); - let range_extra_offset = seed as usize % range_count; - let gap_extra_offset = seed.rotate_left(17) as usize % gap_count; + for byte in &shard.partition_bytes { + seed = mix_seed(seed ^ u64::from(*byte)); + } let mut cursor = shard.row_range_start; let mut ranges = Vec::with_capacity(range_count); - for gap_index in 0..range_count { - let gap = skipped_rows / gap_count - + usize::from( - (gap_index + gap_count - gap_extra_offset) % gap_count < skipped_rows % gap_count, - ); - cursor = checked_add_offset(cursor, gap, "training gap")?; - let length = training_rows / range_count - + usize::from( - (gap_index + range_count - range_extra_offset) % range_count - < training_rows % range_count, - ); - let end = checked_add_offset(cursor, length - 1, "training range")?; - ranges.push(RowRange::new(cursor, end)); - cursor = end.checked_add(1).ok_or_else(|| Error::DataInvalid { - message: "vindex training range end overflows i64".to_string(), - source: None, - })?; + for range_index in 0..range_count { + let stratum_length = + shard_rows / range_count + usize::from(range_index < shard_rows % range_count); + let length = + training_rows / range_count + usize::from(range_index < training_rows % range_count); + debug_assert!(length <= stratum_length); + let available_offsets = stratum_length - length + 1; + let offset = mix_seed(seed ^ range_index as u64) as usize % available_offsets; + let start = checked_add_offset(cursor, offset, "training range")?; + let end = checked_add_offset(start, length - 1, "training range")?; + ranges.push(RowRange::new(start, end)); + cursor = checked_add_offset(cursor, stratum_length, "training stratum")?; } - Ok(ranges) + Ok(merge_row_ranges(ranges)) } fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { 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 a5813b762..0863dad73 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -152,14 +152,30 @@ fn test_ivf_training_ranges_are_bounded_and_exact() { .unwrap() .remove(0); - let ranges = plan_ivf_training_ranges(&shard, 200).unwrap(); + for training_rows in [1, 63, 64, 65, 200, 899] { + let ranges = plan_ivf_training_ranges(&shard, training_rows).unwrap(); + + assert!(ranges.len() <= 64); + assert_eq!( + ranges.iter().map(RowRange::count).sum::(), + training_rows as i64 + ); + assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); + assert!(ranges.first().unwrap().from() >= shard.row_range_start); + assert!(ranges.last().unwrap().to() <= shard.row_range_end); + assert_eq!( + ranges, + plan_ivf_training_ranges(&shard, training_rows).unwrap() + ); + } - assert_eq!(ranges.len(), 64); - assert_eq!(ranges.iter().map(RowRange::count).sum::(), 200); - assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); - assert!(ranges.first().unwrap().from() >= 100); - assert!(ranges.last().unwrap().to() <= 1_099); - assert_eq!(ranges, plan_ivf_training_ranges(&shard, 200).unwrap()); + let ranges = plan_ivf_training_ranges(&shard, 200).unwrap(); + let mut other_snapshot = shard.clone(); + other_snapshot.snapshot_id += 1; + assert_ne!( + ranges, + plan_ivf_training_ranges(&other_snapshot, 200).unwrap() + ); } #[test] From 0ffc3938ea37979aeb508bac39e90d8ed6d9a969 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 16:43:00 +0800 Subject: [PATCH 04/24] fix(vindex): preserve IVF sampling quality --- .../vindex_index_build_builder/planning.rs | 15 ++-- .../table/vindex_index_build_builder/tests.rs | 77 ++++++++++++++++++- .../vindex_index_build_builder/writer.rs | 33 ++++---- 3 files changed, 97 insertions(+), 28 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index c5f54b079..745fc7d80 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -22,7 +22,9 @@ use crate::{Error, Result}; use super::validation::checked_row_count; -const MAX_IVF_TRAINING_RANGES: usize = 64; +// Keep samples short enough to avoid storage-order bias; fall back before range I/O explodes. +const MAX_IVF_TRAINING_RANGE_ROWS: usize = 128; +const MAX_IVF_TRAINING_RANGES: usize = 4_096; pub(crate) type VindexIndexShard = VectorIndexShard; @@ -53,7 +55,7 @@ pub(super) fn plan_vindex_shards( pub(super) fn plan_ivf_training_ranges( shard: &VindexIndexShard, training_rows: usize, -) -> Result> { +) -> Result>> { let shard_rows = usize::try_from(checked_row_count( shard.row_range_start, shard.row_range_end, @@ -70,11 +72,10 @@ pub(super) fn plan_ivf_training_ranges( source: None, }); } - if training_rows == shard_rows { - return Ok(Vec::new()); + let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS); + if range_count > MAX_IVF_TRAINING_RANGES { + return Ok(None); } - - let range_count = training_rows.min(MAX_IVF_TRAINING_RANGES); let mut seed = mix_seed( (shard.snapshot_id as u64) ^ (shard.row_range_start as u64).rotate_left(21) @@ -101,7 +102,7 @@ pub(super) fn plan_ivf_training_ranges( cursor = checked_add_offset(cursor, stratum_length, "training stratum")?; } - Ok(merge_row_ranges(ranges)) + Ok(Some(merge_row_ranges(ranges))) } fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { 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 0863dad73..11f73b552 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -153,7 +153,9 @@ fn test_ivf_training_ranges_are_bounded_and_exact() { .remove(0); for training_rows in [1, 63, 64, 65, 200, 899] { - let ranges = plan_ivf_training_ranges(&shard, training_rows).unwrap(); + let ranges = plan_ivf_training_ranges(&shard, training_rows) + .unwrap() + .expect("sparse ranges expected"); assert!(ranges.len() <= 64); assert_eq!( @@ -165,19 +167,86 @@ fn test_ivf_training_ranges_are_bounded_and_exact() { assert!(ranges.last().unwrap().to() <= shard.row_range_end); assert_eq!( ranges, - plan_ivf_training_ranges(&shard, training_rows).unwrap() + plan_ivf_training_ranges(&shard, training_rows) + .unwrap() + .expect("sparse ranges expected") ); } - let ranges = plan_ivf_training_ranges(&shard, 200).unwrap(); + let ranges = plan_ivf_training_ranges(&shard, 200) + .unwrap() + .expect("sparse ranges expected"); let mut other_snapshot = shard.clone(); other_snapshot.snapshot_id += 1; assert_ne!( ranges, - plan_ivf_training_ranges(&other_snapshot, 200).unwrap() + plan_ivf_training_ranges(&other_snapshot, 200) + .unwrap() + .expect("sparse ranges expected") ); } +#[test] +fn test_ivf_training_ranges_keep_segments_short() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 1_000_000))], + 1_000_000, + ) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 65_536) + .unwrap() + .expect("sparse ranges expected"); + + assert_eq!(ranges.len(), 512); + assert!(ranges.iter().all(|range| range.count() <= 128)); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 65_536); +} + +#[test] +fn test_ivf_training_ranges_scale_to_keep_segments_short() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 2_000_000))], + 2_000_000, + ) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 262_144) + .unwrap() + .expect("sparse ranges expected"); + + assert_eq!(ranges.len(), 2_048); + assert!(ranges.iter().all(|range| range.count() <= 128)); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 262_144); +} + +#[test] +fn test_ivf_training_ranges_fall_back_above_range_limit() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(0), 2_000_000))], + 2_000_000, + ) + .unwrap() + .remove(0); + + assert!(plan_ivf_training_ranges(&shard, 524_289).unwrap().is_none()); +} + +#[test] +fn test_ivf_training_ranges_cover_full_shard_without_empty_sentinel() { + let shard = plan(vec![manifest_entry(data_file("a", Some(0), 1_000))], 1_000) + .unwrap() + .remove(0); + + let ranges = plan_ivf_training_ranges(&shard, 1_000) + .unwrap() + .expect("full range expected"); + + assert_eq!(ranges, vec![RowRange::new(0, 999)]); +} + #[test] fn test_planner_rejects_missing_first_row_id() { let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) 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 673e36445..42577e02c 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -96,16 +96,13 @@ impl<'a> VindexIndexBuildBuilder<'a> { let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { 0 } else { - default_training_vector_count(training_vector_count, options.config.nlist()).map_err( - |e| Error::DataInvalid { - message: format!("Failed to calculate IVF training vector count: {e}"), - source: Some(Box::new(e)), - }, - )? + // This only gates sparse reads; errors must preserve the full-scan fallback. + default_training_vector_count(training_vector_count, options.config.nlist()) + .unwrap_or(0) }; let mut sparse_ranges = if training_rows_retained > 0 && training_rows_retained < row_count_usize { - Some(plan_ivf_training_ranges(shard, training_rows_retained)?) + plan_ivf_training_ranges(shard, training_rows_retained)? } else { None }; @@ -114,10 +111,12 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut checks = Vec::new(); let mut usable = true; for file in &shard.files { - let path = file.data_file_path(&shard.bucket_path); - if !path.to_ascii_lowercase().ends_with(".parquet") { - usable = false; - break; + if file + .write_cols + .as_ref() + .is_some_and(|columns| !columns.iter().any(|column| column == index_column)) + { + continue; } let Some((file_start, file_end)) = file.row_id_range() else { usable = false; @@ -131,14 +130,14 @@ impl<'a> VindexIndexBuildBuilder<'a> { (from <= to).then(|| RowRange::new(from - file_start, to - file_start)) }) .collect::>(); - if local_ranges.is_empty() - || file - .write_cols - .as_ref() - .is_some_and(|columns| !columns.iter().any(|column| column == index_column)) - { + if local_ranges.is_empty() { continue; } + let path = file.data_file_path(&shard.bucket_path); + if !path.to_ascii_lowercase().ends_with(".parquet") { + usable = false; + break; + } let file_size = u64::try_from(file.file_size).map_err(|e| Error::DataInvalid { message: format!( "Invalid data file size for '{}': {}", From af2bcd3079ae506350ae7653eaa50887e8e99e28 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 17:36:40 +0800 Subject: [PATCH 05/24] fix(vindex): harden sparse build pipeline --- .../table/vindex_index_build_builder/tests.rs | 11 +++-- .../vindex_index_build_builder/writer.rs | 41 +++++++++++++------ 2 files changed, 37 insertions(+), 15 deletions(-) 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 11f73b552..cb61491a8 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -733,8 +733,13 @@ async fn vindex_incremental_build_indexes_only_new_rows() { // Build #1 over the initial batch via a real end-to-end build. write_vectors( &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + vec![0, 1, 2, 3], + vec![ + vec![1.0, 0.0], + vec![0.0, 1.0], + vec![1.0, 1.0], + vec![2.0, 1.0], + ], ) .await; let first_built = table @@ -742,7 +747,7 @@ async fn vindex_incremental_build_indexes_only_new_rows() { .with_index_column("embedding") .with_options(HashMap::from([( "ivf-flat.train.sample-ratio".to_string(), - "0.9".to_string(), + "0.5".to_string(), )])) .execute() .await 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 42577e02c..5c4ad87b0 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -175,6 +175,12 @@ impl<'a> VindexIndexBuildBuilder<'a> { } } if !usable || !found_vector_file { + log::warn!( + "vindex sparse training read is unavailable for column '{}' in shard [{}, {}]; falling back to a full scan (usable_offset_indexes={usable}, found_vector_file={found_vector_file})", + index_column, + shard.row_range_start, + shard.row_range_end, + ); sparse_ranges = None; } } @@ -347,16 +353,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { rows_added = batch_end; batches_added += 1; } - if rows_added != row_count_usize || expected_row_id != expected_end { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_added}/{row_count_usize}, next_row_id={expected_row_id}/{expected_end}" - ), - source: None, - }); - } Ok(( writer, + rows_added, + expected_row_id, batches_added, blocked, validate_add, @@ -398,14 +398,31 @@ impl<'a> VindexIndexBuildBuilder<'a> { } drop(sender); let consumer_result = consumer.await; + let ( + writer, + rows_added, + next_row_id, + batches_added, + blocked, + validate_add, + min_bytes, + max_bytes, + total_bytes, + ) = consumer_result.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })??; if let Some(error) = producer_error { return Err(error); } - let (writer, batches_added, blocked, validate_add, min_bytes, max_bytes, total_bytes) = - consumer_result.map_err(|e| Error::UnexpectedError { - message: format!("vindex add task failed: {e}"), + if rows_added != row_count_usize || next_row_id != expected_end { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_added}/{row_count_usize}, next_row_id={next_row_id}/{expected_end}" + ), source: None, - })??; + }); + } batch_count = batches_added; pipeline_blocked = blocked; consumer_validate_add = validate_add; From 999a5302e513559ed2ad6fddb79d172229a823a9 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 10:19:41 +0800 Subject: [PATCH 06/24] perf(vindex): parallelize sparse parquet reads --- crates/paimon/src/arrow/format/parquet.rs | 121 +++++++++++++++--- crates/paimon/src/arrow/read_budget.rs | 4 + crates/paimon/src/table/data_file_reader.rs | 14 +- .../vindex_index_build_builder/planning.rs | 23 ++-- .../table/vindex_index_build_builder/tests.rs | 75 +++++++++++ .../vindex_index_build_builder/timing.rs | 25 ++++ .../vindex_index_build_builder/writer.rs | 50 ++++++-- 7 files changed, 268 insertions(+), 44 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index ec1a85060..bf4ffab76 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -595,12 +595,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 @@ -608,39 +608,53 @@ 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; + } + Some(( + row_group_index, + selection, + projected_row_group_bytes(row_group, &mask), + )) }) - .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; }; @@ -665,6 +679,7 @@ impl FormatFileReader for ParquetFormatReader { row_group_mask, row_group_index, batch_size, + selection, permit, batch_tx, )); @@ -768,6 +783,7 @@ async fn read_row_group( projection: ProjectionMask, row_group_index: usize, batch_size: Option, + selection: Option, _permit: ReadPermit, sender: mpsc::Sender, ) { @@ -777,6 +793,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); } @@ -2942,6 +2961,66 @@ mod tests { ); } + #[tokio::test] + async fn test_sparse_row_groups_preserve_selection_order_and_budget() { + let data = write_multi_row_group_parquet(64, 384, EnabledStatistics::Chunk, false).await; + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let file_reader = ConcurrentTrackingFileRead { + data: Bytes::from(data), + in_flight, + max_in_flight, + }; + let file_size = file_reader.data.len() as u64; + let ranges = vec![ + RowRange::new(60, 68), + RowRange::new(130, 135), + RowRange::new(258, 263), + RowRange::new(380, 383), + ]; + let budget = Arc::new(ParquetReadBudget::new(2, 256 * 1024 * 1024).unwrap()); + budget.enable_diagnostics(); + + let batches = ParquetFormatReader::with_read_budget(Arc::clone(&budget)) + .read_batch_stream( + Box::new(file_reader), + file_size, + &[int_field("id")], + None, + Some(32), + Some(ranges.clone()), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect::>() + }) + .collect::>(); + let expected = ranges + .iter() + .flat_map(|range| range.from() as i32..=range.to() as i32) + .collect::>(); + + assert_eq!(actual, expected); + let diagnostics = budget.diagnostics(); + assert_eq!(diagnostics.row_group_count, 5); + assert_eq!(diagnostics.peak_inflight, 2); + assert_eq!(diagnostics.current_inflight, 0); + } + #[tokio::test] async fn test_parquet_read_budget_is_shared_across_readers() { const ROWS: i32 = 256; @@ -3055,7 +3134,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] diff --git a/crates/paimon/src/arrow/read_budget.rs b/crates/paimon/src/arrow/read_budget.rs index 334ac8e34..42aa4411f 100644 --- a/crates/paimon/src/arrow/read_budget.rs +++ b/crates/paimon/src/arrow/read_budget.rs @@ -140,6 +140,10 @@ impl ReadBudget { self.parallelism } + pub(crate) fn max_inflight_bytes(&self) -> u64 { + self.max_inflight_bytes + } + pub(crate) fn enable_diagnostics(&self) { self.diagnostics.enabled.store(true, Ordering::Relaxed); } diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 5f442e8eb..0e16766d6 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -53,7 +53,7 @@ pub(crate) struct DataFileReadTiming { } impl DataFileReadTiming { - fn add_file_read(&self, duration: Duration) { + pub(crate) fn add_file_read(&self, duration: Duration) { self.file_read_nanos .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); } @@ -94,6 +94,13 @@ impl DataFileReadTiming { ) } + pub(crate) fn wrap_reader(self: &Arc, inner: Box) -> Box { + Box::new(TimedFileRead { + inner, + timing: Arc::clone(self), + }) + } + pub(crate) fn parquet_decode(&self) -> Duration { Duration::from_nanos(self.parquet_decode_nanos.load(Ordering::Relaxed)) } @@ -554,10 +561,7 @@ impl DataFileReader { timing.add_file_read(start.elapsed()); } let file_reader: Box = match read_timing.as_ref() { - Some(timing) => Box::new(TimedFileRead { - inner: Box::new(file_reader), - timing: Arc::clone(timing), - }), + Some(timing) => timing.wrap_reader(Box::new(file_reader)), None => Box::new(file_reader), }; let is_parquet = path_to_read.to_ascii_lowercase().ends_with(".parquet"); diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index 745fc7d80..89c0fcd69 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -76,15 +76,7 @@ pub(super) fn plan_ivf_training_ranges( if range_count > MAX_IVF_TRAINING_RANGES { return Ok(None); } - let mut seed = mix_seed( - (shard.snapshot_id as u64) - ^ (shard.row_range_start as u64).rotate_left(21) - ^ (shard.row_range_end as u64).rotate_left(42) - ^ (shard.source_bucket as u64).rotate_left(11), - ); - for byte in &shard.partition_bytes { - seed = mix_seed(seed ^ u64::from(*byte)); - } + let seed = ivf_training_seed(shard); let mut cursor = shard.row_range_start; let mut ranges = Vec::with_capacity(range_count); @@ -105,6 +97,19 @@ pub(super) fn plan_ivf_training_ranges( Ok(Some(merge_row_ranges(ranges))) } +pub(super) fn ivf_training_seed(shard: &VindexIndexShard) -> u64 { + let mut seed = mix_seed( + (shard.snapshot_id as u64) + ^ (shard.row_range_start as u64).rotate_left(21) + ^ (shard.row_range_end as u64).rotate_left(42) + ^ (shard.source_bucket as u64).rotate_left(11), + ); + for byte in &shard.partition_bytes { + seed = mix_seed(seed ^ u64::from(*byte)); + } + seed +} + fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { let offset = i64::try_from(offset).map_err(|error| Error::DataInvalid { message: format!("vindex {name} offset does not fit i64"), 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 cb61491a8..a256663b5 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -816,6 +816,81 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } +#[test] +fn vindex_build_logs_read_phases() { + // Run the real sparse and fallback builds in a separate process so the + // timing environment variable and stderr capture cannot race other tests. + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "table::vindex_index_build_builder::tests::vindex_incremental_build_indexes_only_new_rows", + "--nocapture", + ]) + .env("PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING", "1") + .output() + .unwrap(); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!(output.status.success(), "{stderr}"); + let events = stderr + .lines() + .filter(|line| line.starts_with("event=paimon_vector_index_build")) + .map(|line| { + line.split_whitespace() + .filter_map(|field| field.split_once('=')) + .collect::>() + }) + .collect::>(); + let plans = events + .iter() + .filter(|event| event["event"] == "paimon_vector_index_build_plan") + .collect::>(); + assert_eq!(plans.len(), 2, "missing build provenance: {stderr}"); + assert_eq!(plans[0]["sparse"], "true"); + assert_eq!(plans[1]["sparse"], "false"); + for plan in &plans { + assert!(plan["snapshot_id"].parse::().unwrap() > 0); + assert!(plan["training_seed"].parse::().is_ok()); + assert_eq!(plan["parquet_row_group_parallelism"], "8"); + assert_eq!(plan["parquet_max_inflight_bytes"], "268435456"); + } + let phases = events + .iter() + .filter(|event| event["event"] == "paimon_vector_index_build_read") + .collect::>(); + assert_eq!( + phases + .iter() + .map(|event| event["phase"]) + .collect::>(), + ["probe", "sample", "full_scan", "full_scan"] + ); + for phase in &phases { + assert_eq!(phase["io_scope"], "file_read_wrapper"); + assert!(phase["read_bytes"].parse::().unwrap() > 0); + assert!(phase["read_calls"].parse::().unwrap() > 0); + assert!(phase["read_ms"].parse::().unwrap() >= 0.0); + } + let totals = events + .iter() + .filter(|event| event["event"] == "paimon_vector_index_build") + .collect::>(); + for (total, reads) in [(totals[0], &phases[1..3]), (totals[1], &phases[3..4])] { + for (total_key, phase_key) in [ + ("oss_read_bytes", "read_bytes"), + ("oss_range_requests", "read_calls"), + ] { + assert_eq!( + total[total_key].parse::().unwrap(), + reads + .iter() + .map(|event| event[phase_key].parse::().unwrap()) + .sum::(), + "phase counters must exclude probe and partition the existing total" + ); + } + } +} + #[tokio::test] async fn vindex_upload_failure_preserves_committed_index() { use crate::io::multipart_test::{Fault, MultipartProvider}; 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 c391ad35f..2181ed5d7 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use super::planning::VindexIndexShard; +use crate::table::data_file_reader::DataFileReadTiming; use std::sync::OnceLock; use std::time::Duration; @@ -27,6 +29,29 @@ pub(super) fn vector_index_build_timing_enabled() -> bool { }) } +// Snapshot only at drained phase boundaries. These are completed FileRead +// calls (plus reader-open time), not OSS wire requests; concurrent read_ms +// is cumulative and must not be added to wall-clock stage durations. +pub(super) fn log_read_phase( + phase: &str, + shard: &VindexIndexShard, + timing: &DataFileReadTiming, + previous: (Duration, u64, u64), +) -> (Duration, u64, u64) { + let read = timing.file_read(); + let (bytes, calls) = timing.file_io(); + eprintln!( + "event=paimon_vector_index_build_read phase={phase} snapshot_id={} row_range_start={} row_range_end={} io_scope=file_read_wrapper read_ms={:.3} read_bytes={} read_calls={}", + shard.snapshot_id, + shard.row_range_start, + shard.row_range_end, + read.saturating_sub(previous.0).as_secs_f64() * 1000.0, + bytes.saturating_sub(previous.1), + calls.saturating_sub(previous.2), + ); + (read, bytes, calls) +} + pub(super) struct VectorIndexBuildTiming { pub(super) total_without_commit: Duration, pub(super) source_batch_wait: Duration, 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 5c4ad87b0..ed7114924 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -19,8 +19,8 @@ use super::extraction::{ data_split_for_shard, data_split_for_shard_ranges, validate_vector_batch, validate_vector_batch_ranges, }; -use super::planning::{plan_ivf_training_ranges, VindexIndexShard}; -use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; +use super::planning::{ivf_training_seed, plan_ivf_training_ranges, VindexIndexShard}; +use super::timing::{log_read_phase, vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, @@ -107,6 +107,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { None }; let capability_start = (timing_enabled && sparse_ranges.is_some()).then(Instant::now); + let capability_read_timing = + capability_start.map(|_| Arc::new(DataFileReadTiming::default())); if let Some(ranges) = sparse_ranges.as_ref() { let mut checks = Vec::new(); let mut usable = true; @@ -155,16 +157,21 @@ impl<'a> VindexIndexBuildBuilder<'a> { .core_options() .parquet_row_group_parallelism()?; let file_io = self.table.file_io(); + let timing = capability_read_timing.as_ref(); let mut checks = futures::stream::iter(checks) .map(|(path, file_size, local_ranges)| async move { let input = file_io.new_input(&path)?; - has_usable_offset_index( - Box::new(input.reader().await?), - file_size, - index_column, - &local_ranges, - ) - .await + let open_start = timing.map(|_| Instant::now()); + let reader = Box::new(input.reader().await?); + if let (Some(timing), Some(start)) = (timing, open_start) { + timing.add_file_read(start.elapsed()); + } + let reader = match timing { + Some(timing) => timing.wrap_reader(reader), + None => reader, + }; + has_usable_offset_index(reader, file_size, index_column, &local_ranges) + .await }) .buffer_unordered(concurrency); while let Some(file_usable) = checks.try_next().await? { @@ -185,6 +192,23 @@ impl<'a> VindexIndexBuildBuilder<'a> { } } let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); + if let Some(timing) = capability_read_timing.as_ref() { + log_read_phase("probe", shard, timing, Default::default()); + } + if let Some(budget) = parquet_read_budget.as_ref() { + eprintln!( + "event=paimon_vector_index_build_plan snapshot_id={} row_range_start={} row_range_end={} source_bucket={} sparse={} training_seed={} training_range_count={} parquet_row_group_parallelism={} parquet_max_inflight_bytes={}", + shard.snapshot_id, + shard.row_range_start, + shard.row_range_end, + shard.source_bucket, + sparse_ranges.is_some(), + ivf_training_seed(shard), + sparse_ranges.as_ref().map_or(0, Vec::len), + budget.parallelism(), + budget.max_inflight_bytes(), + ); + } let mut trainer = VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { @@ -205,6 +229,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { let train_finish; let mut bytes_written = 0usize; let training_rows_seen; + let mut sample_io = Default::default(); let writer = if let Some(ranges) = sparse_ranges { let sample_start = timing_enabled.then(Instant::now); @@ -258,6 +283,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { } training_rows_seen = rows_seen; sample_read = sample_start.map_or(Duration::ZERO, |start| start.elapsed()); + if let Some(timing) = read_timing.as_ref() { + sample_io = log_read_phase("sample", shard, timing, Default::default()); + } let train_start = timing_enabled.then(Instant::now); let training = tokio::task::spawn_blocking(move || trainer.finish()) @@ -675,6 +703,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { result.0 }; + if let Some(timing) = read_timing.as_ref() { + log_read_phase("full_scan", shard, timing, sample_io); + } + let serialize_upload_start = timing_enabled.then(Instant::now); self.table .file_io() From 17db1f87f99294a837da62e073af80f25a19e291 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 13:07:00 +0800 Subject: [PATCH 07/24] perf(parquet): budget sparse reads by selected pages --- crates/paimon/src/arrow/format/parquet.rs | 109 ++++++++++++++++++++-- 1 file changed, 102 insertions(+), 7 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index bf4ffab76..9814c0f95 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -46,6 +46,7 @@ use parquet::file::metadata::{ KeyValue, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData, }; use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics as ParquetStatistics; use std::cmp::Ordering; @@ -614,6 +615,7 @@ impl FormatFileReader for ParquetFormatReader { .filter(|budget| row_group_parallelism > 1 || budget.diagnostics_enabled()) .map(|budget| { let mut row_group_selection = combined_selection; + let offset_index = batch_stream_builder.metadata().offset_index(); let selected_row_groups = batch_stream_builder .metadata() .row_groups() @@ -629,11 +631,15 @@ impl FormatFileReader for ParquetFormatReader { { return None; } - Some(( - row_group_index, - selection, - projected_row_group_bytes(row_group, &mask), - )) + let projected_bytes = projected_row_group_bytes( + row_group, + &mask, + selection.as_ref(), + offset_index + .and_then(|index| index.get(row_group_index)) + .map(Vec::as_slice), + ); + Some((row_group_index, selection, projected_bytes)) }) .collect::>(); let projected_bytes = selected_row_groups @@ -765,13 +771,50 @@ impl FormatFileReader for ParquetFormatReader { } } -fn projected_row_group_bytes(row_group: &RowGroupMetaData, projection: &ProjectionMask) -> u64 { +fn projected_row_group_bytes( + row_group: &RowGroupMetaData, + projection: &ProjectionMask, + selection: Option<&RowSelection>, + offset_index: Option<&[OffsetIndexMetaData]>, +) -> u64 { row_group .columns() .iter() .enumerate() .filter(|(leaf_index, _)| projection.leaf_included(*leaf_index)) - .filter_map(|(_, column)| u64::try_from(column.uncompressed_size()).ok()) + .filter_map(|(leaf_index, column)| { + let uncompressed_bytes = u64::try_from(column.uncompressed_size()).ok()?; + let selected_bytes = selection + .zip(offset_index.and_then(|index| index.get(leaf_index))) + .and_then(|(selection, page_index)| { + let page_locations = page_index.page_locations(); + let column_start = u64::try_from( + column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()), + ) + .ok()?; + let dictionary_bytes = u64::try_from(page_locations.first()?.offset) + .ok()? + .checked_sub(column_start)?; + let selected_compressed_bytes = selection + .scan_ranges(page_locations) + .into_iter() + .try_fold(dictionary_bytes, |total, range| { + total.checked_add(range.end.checked_sub(range.start)?) + })?; + let compressed_bytes = u64::try_from(column.compressed_size()).ok()?; + if compressed_bytes == 0 || selected_compressed_bytes > compressed_bytes { + return None; + } + Some( + ((u128::from(uncompressed_bytes) * u128::from(selected_compressed_bytes)) + .div_ceil(u128::from(compressed_bytes))) + .min(u128::from(u64::MAX)) as u64, + ) + }); + Some(selected_bytes.unwrap_or(uncompressed_bytes)) + }) .fold(0u64, u64::saturating_add) } @@ -3021,6 +3064,58 @@ mod tests { assert_eq!(diagnostics.current_inflight, 0); } + #[tokio::test] + async fn test_sparse_page_budget_allows_large_row_groups_to_overlap() { + 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 selection = RowSelection::from_consecutive_ranges(std::iter::once(0..1), 80); + let projected_bytes = super::projected_row_group_bytes( + &row_group, + &projection, + Some(&selection), + Some(offset_index.as_slice()), + ); + + assert!(projected_bytes < 128 * MIB as u64); + assert_eq!( + super::projected_row_group_bytes( + &row_group, + &projection, + None, + Some(offset_index.as_slice()), + ), + 308 * MIB as u64 + ); + + let budget = ParquetReadBudget::new(8, 256 * MIB as u64).unwrap(); + budget.enable_diagnostics(); + let _first = budget.acquire(projected_bytes).await.unwrap(); + let _second = budget.acquire(projected_bytes).await.unwrap(); + assert_eq!(budget.diagnostics().peak_inflight, 2); + } + #[tokio::test] async fn test_parquet_read_budget_is_shared_across_readers() { const ROWS: i32 = 256; From e7bb7b737c4ffe6731bb99f9c3725ed6643fe163 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 14:56:49 +0800 Subject: [PATCH 08/24] fix(vindex): guard sparse training reads --- crates/paimon/src/arrow/format/parquet.rs | 187 +++++++++--------- .../vindex_index_build_builder/planning.rs | 2 +- .../table/vindex_index_build_builder/tests.rs | 50 +++-- .../vindex_index_build_builder/writer.rs | 6 +- 4 files changed, 131 insertions(+), 114 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 9814c0f95..9871c64a4 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -46,7 +46,6 @@ use parquet::file::metadata::{ KeyValue, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData, }; use parquet::file::page_index::column_index::ColumnIndexMetaData; -use parquet::file::page_index::offset_index::OffsetIndexMetaData; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics as ParquetStatistics; use std::cmp::Ordering; @@ -68,7 +67,7 @@ impl ParquetFormatReader { } } -pub(crate) async fn has_usable_offset_index( +pub(crate) async fn has_beneficial_offset_index( reader: Box, file_size: u64, column_name: &str, @@ -77,14 +76,14 @@ pub(crate) async fn has_usable_offset_index( let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); let mut reader = ArrowFileReader::new(file_size, reader.into()); let metadata = reader.get_metadata(Some(&options)).await?; - Ok(metadata_has_usable_offset_index( + Ok(metadata_has_beneficial_offset_index( &metadata, column_name, row_ranges, )) } -fn metadata_has_usable_offset_index( +fn metadata_has_beneficial_offset_index( metadata: &ParquetMetaData, column_name: &str, row_ranges: &[RowRange], @@ -110,28 +109,69 @@ fn metadata_has_usable_offset_index( if columns.is_empty() || offset_index.len() != metadata.row_groups().len() { return false; } - let mut row_start = 0i64; + let mut selection = build_row_ranges_selection(metadata.row_groups(), row_ranges); let mut checked = false; + let mut full_bytes = 0u64; + let mut selected_bytes = 0u64; for (row_group, indexes) in metadata.row_groups().iter().zip(offset_index) { - let Some(row_end) = row_start.checked_add(row_group.num_rows()) else { + let Ok(row_count) = usize::try_from(row_group.num_rows()) else { return false; }; - if row_ranges - .iter() - .any(|range| range.from() < row_end && range.to() >= row_start) - { + let row_group_selection = selection.split_off(row_count); + let mut selected_ranges = Vec::new(); + for index in &columns { + let column = row_group.column(*index); + let Ok(column_bytes) = u64::try_from(column.compressed_size()) else { + return false; + }; + let Some(total) = full_bytes.checked_add(column_bytes) else { + return false; + }; + full_bytes = total; + + if !row_group_selection.selects_any() { + continue; + } checked = true; - if !columns.iter().all(|index| { - indexes - .get(*index) - .is_some_and(|index| !index.page_locations().is_empty()) - }) { + let Some(page_locations) = indexes.get(*index).map(|index| index.page_locations()) + else { + return false; + }; + let Some(first_page) = page_locations.first() else { return false; + }; + let Ok(column_start) = u64::try_from( + column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()), + ) else { + return false; + }; + let Ok(first_page_offset) = u64::try_from(first_page.offset) else { + return false; + }; + if column_start < first_page_offset { + selected_ranges.push(column_start..first_page_offset); } + selected_ranges.extend(row_group_selection.scan_ranges(page_locations)); + } + if row_group_selection.selects_any() { + let Some(group_selected_bytes) = + merge_byte_ranges(&selected_ranges, RANGE_COALESCE_BYTES) + .into_iter() + .try_fold(0u64, |total, range| { + total.checked_add(range.end.checked_sub(range.start)?) + }) + else { + return false; + }; + let Some(total) = selected_bytes.checked_add(group_selected_bytes) else { + return false; + }; + selected_bytes = total; } - row_start = row_end; } - checked + checked && selected_bytes < full_bytes } enum ParquetRowGroupMessage { @@ -615,7 +655,6 @@ impl FormatFileReader for ParquetFormatReader { .filter(|budget| row_group_parallelism > 1 || budget.diagnostics_enabled()) .map(|budget| { let mut row_group_selection = combined_selection; - let offset_index = batch_stream_builder.metadata().offset_index(); let selected_row_groups = batch_stream_builder .metadata() .row_groups() @@ -631,14 +670,7 @@ impl FormatFileReader for ParquetFormatReader { { return None; } - let projected_bytes = projected_row_group_bytes( - row_group, - &mask, - selection.as_ref(), - offset_index - .and_then(|index| index.get(row_group_index)) - .map(Vec::as_slice), - ); + let projected_bytes = projected_row_group_bytes(row_group, &mask); Some((row_group_index, selection, projected_bytes)) }) .collect::>(); @@ -771,50 +803,13 @@ impl FormatFileReader for ParquetFormatReader { } } -fn projected_row_group_bytes( - row_group: &RowGroupMetaData, - projection: &ProjectionMask, - selection: Option<&RowSelection>, - offset_index: Option<&[OffsetIndexMetaData]>, -) -> u64 { +fn projected_row_group_bytes(row_group: &RowGroupMetaData, projection: &ProjectionMask) -> u64 { row_group .columns() .iter() .enumerate() .filter(|(leaf_index, _)| projection.leaf_included(*leaf_index)) - .filter_map(|(leaf_index, column)| { - let uncompressed_bytes = u64::try_from(column.uncompressed_size()).ok()?; - let selected_bytes = selection - .zip(offset_index.and_then(|index| index.get(leaf_index))) - .and_then(|(selection, page_index)| { - let page_locations = page_index.page_locations(); - let column_start = u64::try_from( - column - .dictionary_page_offset() - .unwrap_or_else(|| column.data_page_offset()), - ) - .ok()?; - let dictionary_bytes = u64::try_from(page_locations.first()?.offset) - .ok()? - .checked_sub(column_start)?; - let selected_compressed_bytes = selection - .scan_ranges(page_locations) - .into_iter() - .try_fold(dictionary_bytes, |total, range| { - total.checked_add(range.end.checked_sub(range.start)?) - })?; - let compressed_bytes = u64::try_from(column.compressed_size()).ok()?; - if compressed_bytes == 0 || selected_compressed_bytes > compressed_bytes { - return None; - } - Some( - ((u128::from(uncompressed_bytes) * u128::from(selected_compressed_bytes)) - .div_ceil(u128::from(compressed_bytes))) - .min(u128::from(u64::MAX)) as u64, - ) - }); - Some(selected_bytes.unwrap_or(uncompressed_bytes)) - }) + .filter_map(|(_, column)| u64::try_from(column.uncompressed_size()).ok()) .fold(0u64, u64::saturating_add) } @@ -2418,7 +2413,7 @@ fn split_ranges_for_concurrency(merged: Vec>, concurrency: usize) -> mod tests { use super::build_parquet_row_filter; use super::{ - forward_row_group_batches, metadata_has_usable_offset_index, parse_compression, + forward_row_group_batches, metadata_has_beneficial_offset_index, parse_compression, supported_compressions, FilePredicates, ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage, }; @@ -3065,7 +3060,7 @@ mod tests { } #[tokio::test] - async fn test_sparse_page_budget_allows_large_row_groups_to_overlap() { + 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; @@ -3090,30 +3085,9 @@ mod tests { row_group.columns_mut()[0] = column; let projection = super::ProjectionMask::roots(row_group.schema_descr(), [0]); - let selection = RowSelection::from_consecutive_ranges(std::iter::once(0..1), 80); - let projected_bytes = super::projected_row_group_bytes( - &row_group, - &projection, - Some(&selection), - Some(offset_index.as_slice()), - ); - - assert!(projected_bytes < 128 * MIB as u64); - assert_eq!( - super::projected_row_group_bytes( - &row_group, - &projection, - None, - Some(offset_index.as_slice()), - ), - 308 * MIB as u64 - ); + let projected_bytes = super::projected_row_group_bytes(&row_group, &projection); - let budget = ParquetReadBudget::new(8, 256 * MIB as u64).unwrap(); - budget.enable_diagnostics(); - let _first = budget.acquire(projected_bytes).await.unwrap(); - let _second = budget.acquire(projected_bytes).await.unwrap(); - assert_eq!(budget.diagnostics().peak_inflight, 2); + assert_eq!(projected_bytes, 308 * MIB as u64); } #[tokio::test] @@ -3719,24 +3693,41 @@ mod tests { } #[tokio::test] - async fn test_sparse_row_selection_requires_offset_index_for_projected_column() { + async fn test_sparse_row_selection_requires_offset_index_and_page_savings() { + let bytes = write_multi_page_parquet(10, 80).await; + let metadata = load_metadata_with_page_index(&bytes, true); + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 0), RowRange::new(79, 79)] + )); + let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; let metadata = load_metadata_with_page_index(&bytes, true); - let ranges = [RowRange::new(0, 19)]; - assert!(metadata_has_usable_offset_index( - &metadata, "value", &ranges + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 19)] + )); + let sparse_ranges = [RowRange::new(0, 0)]; + assert!(metadata_has_beneficial_offset_index( + &metadata, + "value", + &sparse_ranges )); - assert!(!metadata_has_usable_offset_index( - &metadata, "missing", &ranges + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "missing", + &sparse_ranges )); let bytes_without_index = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; let metadata_without_index = load_metadata_with_page_index(&bytes_without_index, true); - assert!(!metadata_has_usable_offset_index( + assert!(!metadata_has_beneficial_offset_index( &metadata_without_index, "value", - &ranges + &sparse_ranges )); } diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index 89c0fcd69..f4d35e126 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -73,7 +73,7 @@ pub(super) fn plan_ivf_training_ranges( }); } let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS); - if range_count > MAX_IVF_TRAINING_RANGES { + if range_count == 1 || range_count > MAX_IVF_TRAINING_RANGES { return Ok(None); } let seed = ivf_training_seed(shard); 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 a256663b5..b10570c0a 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -152,7 +152,7 @@ fn test_ivf_training_ranges_are_bounded_and_exact() { .unwrap() .remove(0); - for training_rows in [1, 63, 64, 65, 200, 899] { + for training_rows in [129, 200, 899] { let ranges = plan_ivf_training_ranges(&shard, training_rows) .unwrap() .expect("sparse ranges expected"); @@ -186,6 +186,18 @@ fn test_ivf_training_ranges_are_bounded_and_exact() { ); } +#[test] +fn test_ivf_training_ranges_fall_back_for_a_single_range_sample() { + let shard = plan( + vec![manifest_entry(data_file("a", Some(100), 1_000))], + 1_000, + ) + .unwrap() + .remove(0); + + assert!(plan_ivf_training_ranges(&shard, 128).unwrap().is_none()); +} + #[test] fn test_ivf_training_ranges_keep_segments_short() { let shard = plan( @@ -727,19 +739,27 @@ async fn vindex_second_build_without_new_data_is_noop() { #[tokio::test] async fn vindex_incremental_build_indexes_only_new_rows() { let table_path = "memory:/test_vindex_incremental"; - let table = vindex_e2e_table(table_path, "10"); + let mut options = vindex_e2e_options("2048"); + options.insert("ivf-flat.dimension".to_string(), "4096".to_string()); + options.insert("ivf-flat.nlist".to_string(), "3".to_string()); + let table = test_table_with_io( + FileIOBuilder::new("memory").build().unwrap(), + table_path, + vindex_schema_builder(options).build().unwrap(), + ); setup_dirs(table.file_io(), table_path).await; // Build #1 over the initial batch via a real end-to-end build. write_vectors( &table, - vec![0, 1, 2, 3], - vec![ - vec![1.0, 0.0], - vec![0.0, 1.0], - vec![1.0, 1.0], - vec![2.0, 1.0], - ], + (0..1024).collect(), + (0..1024) + .map(|id| { + (0..4096) + .map(|component| (id * 4096 + component) as f32) + .collect() + }) + .collect(), ) .await; let first_built = table @@ -747,7 +767,7 @@ async fn vindex_incremental_build_indexes_only_new_rows() { .with_index_column("embedding") .with_options(HashMap::from([( "ivf-flat.train.sample-ratio".to_string(), - "0.5".to_string(), + "0.2".to_string(), )])) .execute() .await @@ -769,8 +789,14 @@ async fn vindex_incremental_build_indexes_only_new_rows() { // Append a second batch (new row-ids [n..]). write_vectors( &table, - vec![4, 5, 6], - vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], + vec![1024, 1025, 1026], + (1024..1027) + .map(|id| { + (0..4096) + .map(|component| (id * 4096 + component) as f32) + .collect() + }) + .collect(), ) .await; 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 ed7114924..3fa53a780 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::validation::{ checked_training_vector_count, checked_vector_bytes, }; use super::VindexIndexBuildBuilder; -use crate::arrow::format::parquet::has_usable_offset_index; +use crate::arrow::format::parquet::has_beneficial_offset_index; use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME}; use crate::table::data_file_reader::DataFileReadTiming; use crate::table::table_read::configured_parquet_read_budget; @@ -170,7 +170,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { Some(timing) => timing.wrap_reader(reader), None => reader, }; - has_usable_offset_index(reader, file_size, index_column, &local_ranges) + has_beneficial_offset_index(reader, file_size, index_column, &local_ranges) .await }) .buffer_unordered(concurrency); @@ -183,7 +183,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { } if !usable || !found_vector_file { log::warn!( - "vindex sparse training read is unavailable for column '{}' in shard [{}, {}]; falling back to a full scan (usable_offset_indexes={usable}, found_vector_file={found_vector_file})", + "vindex sparse training read is unavailable for column '{}' in shard [{}, {}]; falling back to a full scan (beneficial_offset_indexes={usable}, found_vector_file={found_vector_file})", index_column, shard.row_range_start, shard.row_range_end, From 000a1c169e3f967eb09e641cf93da76fc5098d57 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 11:49:57 +0800 Subject: [PATCH 09/24] fix(vindex): require meaningful sparse savings and cover safety regressions --- crates/paimon/src/arrow/format/parquet.rs | 184 +++++++++++++++++- .../table/vindex_index_build_builder/tests.rs | 76 +++++++- 2 files changed, 257 insertions(+), 3 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 9871c64a4..ef506af84 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -171,7 +171,9 @@ fn metadata_has_beneficial_offset_index( selected_bytes = total; } } - checked && selected_bytes < full_bytes + // Sparse training is followed by a full scan, so require it to skip at + // least half of the projected bytes instead of accepting marginal savings. + checked && selected_bytes <= full_bytes / 2 } enum ParquetRowGroupMessage { @@ -3088,6 +3090,22 @@ mod tests { let projected_bytes = super::projected_row_group_bytes(&row_group, &projection); assert_eq!(projected_bytes, 308 * MIB as u64); + let budget = ParquetReadBudget::new(8, 256 * MIB as u64).unwrap(); + budget.enable_diagnostics(); + let permit = budget.acquire(projected_bytes).await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(20), budget.acquire(1)) + .await + .is_err() + ); + drop(permit); + let permit = tokio::time::timeout(Duration::from_secs(1), budget.acquire(projected_bytes)) + .await + .unwrap() + .unwrap(); + assert_eq!(budget.diagnostics().peak_inflight, 1); + drop(permit); + assert_eq!(budget.diagnostics().current_inflight, 0); } #[tokio::test] @@ -3694,8 +3712,20 @@ mod tests { #[tokio::test] async fn test_sparse_row_selection_requires_offset_index_and_page_savings() { + let bytes = write_multi_row_group_parquet(10, 30, EnabledStatistics::Chunk, false).await; + let metadata = load_metadata_with_page_index(&bytes, true); + assert!( + !metadata_has_beneficial_offset_index(&metadata, "value", &[RowRange::new(0, 19)]), + "reading two of three row groups is not sufficiently sparse" + ); + let bytes = write_multi_page_parquet(10, 80).await; let metadata = load_metadata_with_page_index(&bytes, true); + assert!(!metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 69)] + )); assert!(!metadata_has_beneficial_offset_index( &metadata, "value", @@ -3812,6 +3842,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 { @@ -3819,6 +3869,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)), } } @@ -3840,7 +3892,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 % 2 == 0 { + 0 + } else { + rng.next_u32() as i32 + } + })); + writer + .write(&RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values)]).unwrap()) + .await + .unwrap(); + } + writer.close().await.unwrap(); + let metadata = load_metadata_with_page_index(&data, true); + assert_eq!(metadata.num_row_groups(), 4); + let pages = metadata.offset_index().unwrap()[0][0].page_locations(); + assert!( + pages[1].compressed_page_size > pages[0].compressed_page_size * 100, + "adjacent equal-row-count pages must have very different compression ratios" + ); + let projection = super::ProjectionMask::all(); + let projected = super::projected_row_group_bytes(&metadata.row_groups()[0], &projection); + assert!(projected > 8 * MIB as u64 && projected < 9 * MIB as u64); + + let data = Bytes::from(data); + let tracker = TrackingFileRead::new(data.clone()); + let buffer = tracker.read(0..1024).await.unwrap(); + let slice = buffer.slice(0..1); + drop(buffer); + assert_eq!(tracker.resident_bytes.load(AtomicOrdering::SeqCst), 1024); + drop(slice); + assert_eq!(tracker.resident_bytes.load(AtomicOrdering::SeqCst), 0); + + let ranges = (0..4) + .flat_map(|group| { + (1..GROUP_ROWS / PAGE_ROWS).step_by(2).map(move |page| { + let start = (group * GROUP_ROWS + page * PAGE_ROWS) as i64; + RowRange::new(start, start + 255) + }) + }) + .collect::>(); + for cancel in [false, true] { + let tracker = TrackingFileRead::new(data.clone()); + let budget = Arc::new(ParquetReadBudget::new(8, BUDGET as u64).unwrap()); + budget.enable_diagnostics(); + let mut stream = ParquetFormatReader::with_read_budget(Arc::clone(&budget)) + .read_batch_stream( + Box::new(tracker.clone()), + data.len() as u64, + &[int_field("id")], + None, + Some(128), + Some(ranges.clone()), + ) + .await + .unwrap(); + let mut rows = stream.try_next().await.unwrap().unwrap().num_rows(); + if !cancel { + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows(); + } + assert_eq!(rows as i64, ranges.iter().map(RowRange::count).sum::()); + } + drop(stream); + tokio::time::timeout(Duration::from_secs(2), async { + while budget.diagnostics().current_inflight != 0 + || tracker.resident_bytes.load(AtomicOrdering::SeqCst) != 0 + { + tokio::task::yield_now().await; + } + }) + .await + .expect("all buffer owners and row-group permits must be released"); + let peak = tracker.peak_resident_bytes.load(AtomicOrdering::SeqCst); + assert!( + peak > MIB && peak <= BUDGET, + "resident source buffers: {peak}" + ); + assert_eq!(budget.diagnostics().peak_inflight, 2); + let _permit = + tokio::time::timeout(Duration::from_secs(1), budget.acquire(BUDGET as u64)) + .await + .unwrap() + .unwrap(); } } diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs index b10570c0a..e74c22714 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -195,7 +195,15 @@ fn test_ivf_training_ranges_fall_back_for_a_single_range_sample() { .unwrap() .remove(0); - assert!(plan_ivf_training_ranges(&shard, 128).unwrap().is_none()); + for rows in [1, 127, 128] { + assert!(plan_ivf_training_ranges(&shard, rows).unwrap().is_none()); + } + let ranges = plan_ivf_training_ranges(&shard, 129).unwrap().unwrap(); + assert_eq!(ranges.iter().map(RowRange::count).sum::(), 129); + assert!(ranges + .iter() + .all(|range| range.from() >= 100 && range.to() < 1_100)); + assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); } #[test] @@ -917,6 +925,72 @@ fn vindex_build_logs_read_phases() { } } +#[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(); + assert_eq!(result.iter().map(RowRange::count).sum::(), 10); + // Equal-distance IDs need not have a stable order; all hits must be in the tail cluster. + assert!( + result + .iter() + .all(|range| range.from() >= 900 && range.to() < 1000), + "{result:?}" + ); +} + #[tokio::test] async fn vindex_upload_failure_preserves_committed_index() { use crate::io::multipart_test::{Fault, MultipartProvider}; From 8325dc9fe057c3bf6f67d958ca5f710243c77729 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 12:55:32 +0800 Subject: [PATCH 10/24] refactor(vindex): keep sparse build changes focused --- crates/paimon/src/arrow/format/parquet.rs | 2 +- crates/paimon/src/io/file_io.rs | 3 - .../paimon/src/io/file_io/multipart_test.rs | 296 ----------------- crates/paimon/src/table/data_file_reader.rs | 31 +- .../vindex_index_build_builder/extraction.rs | 2 +- .../table/vindex_index_build_builder/tests.rs | 307 ++++++++---------- .../vindex_index_build_builder/timing.rs | 45 +-- .../vindex_index_build_builder/writer.rs | 232 +++++-------- 8 files changed, 222 insertions(+), 696 deletions(-) delete mode 100644 crates/paimon/src/io/file_io/multipart_test.rs diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index ef506af84..65d003082 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -3938,7 +3938,7 @@ mod tests { 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 % 2 == 0 { + if (row / PAGE_ROWS).is_multiple_of(2) { 0 } else { rng.next_u32() as i32 diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs index f6c50ea9f..720225ebc 100644 --- a/crates/paimon/src/io/file_io.rs +++ b/crates/paimon/src/io/file_io.rs @@ -34,9 +34,6 @@ use snafu::ResultExt; use tokio_util::compat::FuturesAsyncWriteCompatExt; use url::Url; -#[cfg(test)] -pub(crate) mod multipart_test; - use super::cache::{CachedFileReader, LocalCache}; use super::Storage; diff --git a/crates/paimon/src/io/file_io/multipart_test.rs b/crates/paimon/src/io/file_io/multipart_test.rs deleted file mode 100644 index e6fa73edd..000000000 --- a/crates/paimon/src/io/file_io/multipart_test.rs +++ /dev/null @@ -1,296 +0,0 @@ -// 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::*; -use crate::io::FileIOBuilder; -use opendal::raw::{ - oio, OpCopier, OpCopy, OpCreateDir, OpList, OpPresign, OpRead, OpRename, OpStat, OpWrite, - RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, -}; -use opendal::{Buffer, Capability, Metadata, OperationContext}; -use std::collections::BTreeMap; -use std::sync::Mutex; -use tokio::sync::Notify; - -#[derive(Clone, Copy, Debug, Default, PartialEq)] -pub(crate) enum Fault { - #[default] - None, - Part, - Close, -} - -#[derive(Debug, Default)] -pub(crate) struct UploadState { - pub(crate) fault: Fault, - pub(crate) fail_on_index: usize, - pub(crate) index_writes: usize, - pub(crate) concurrency: Vec, - pub(crate) uploads: BTreeMap>, - pub(crate) completed_parts: Vec, - pub(crate) aborts: usize, -} - -/// Exercise OpenDAL's real multipart scheduler, storing completed objects in memory. -#[derive(Clone, Debug)] -pub(crate) struct MultipartProvider { - memory: Operator, - part_size: usize, - pub(crate) state: Arc>, -} - -impl MultipartProvider { - pub(crate) fn new(part_size: usize) -> Self { - Self { - memory: Operator::via_iter(opendal::services::MEMORY_SCHEME, []).unwrap(), - part_size, - state: Arc::default(), - } - } - - pub(crate) fn file_io(&self) -> FileIO { - FileIOBuilder::new("memory") - .build() - .unwrap() - .with_provider(Arc::new(self.clone())) - } -} - -#[async_trait::async_trait] -impl FileIOProvider for MultipartProvider { - async fn create(&self, path: &str) -> crate::Result<(Operator, String)> { - Ok(( - Operator::from_parts(OperationContext::default(), Arc::new(self.clone())), - Url::parse(path) - .unwrap() - .path() - .trim_start_matches('/') - .to_string(), - )) - } -} - -impl Service for MultipartProvider { - type Reader = oio::Reader; - type Writer = oio::Writer; - type Lister = oio::Lister; - type Deleter = oio::Deleter; - type Copier = oio::Copier; - - fn info(&self) -> ServiceInfo { - self.memory.service().info() - } - - fn capability(&self) -> Capability { - Capability { - write_can_multi: true, - write_multi_max_size: Some(self.part_size), - ..self.memory.service().capability() - } - } - - fn write( - &self, - ctx: &OperationContext, - path: &str, - args: OpWrite, - ) -> opendal::Result { - let mut state = self.state.lock().unwrap(); - state.concurrency.push(args.concurrent()); - if !path.ends_with(".index") { - return self.memory.service().write(ctx, path, args); - } - state.index_writes += 1; - let fault = if state.index_writes == state.fail_on_index { - state.fault - } else { - Fault::None - }; - Ok(Box::new(oio::MultipartWriter::new( - ctx.executor().clone(), - Upload { - provider: self.clone(), - path: path.to_string(), - fault, - reorder: args.concurrent() > 1, - second_part: Notify::new(), - }, - args.concurrent(), - ))) - } - - async fn create_dir( - &self, - ctx: &OperationContext, - path: &str, - args: OpCreateDir, - ) -> opendal::Result { - self.memory.service().create_dir(ctx, path, args).await - } - async fn stat( - &self, - ctx: &OperationContext, - path: &str, - args: OpStat, - ) -> opendal::Result { - self.memory.service().stat(ctx, path, args).await - } - fn read( - &self, - ctx: &OperationContext, - path: &str, - args: OpRead, - ) -> opendal::Result { - self.memory.service().read(ctx, path, args) - } - fn delete(&self, ctx: &OperationContext) -> opendal::Result { - self.memory.service().delete(ctx) - } - fn list( - &self, - ctx: &OperationContext, - path: &str, - args: OpList, - ) -> opendal::Result { - self.memory.service().list(ctx, path, args) - } - fn copy( - &self, - ctx: &OperationContext, - from: &str, - to: &str, - args: OpCopy, - opts: OpCopier, - ) -> opendal::Result { - self.memory.service().copy(ctx, from, to, args, opts) - } - async fn rename( - &self, - ctx: &OperationContext, - from: &str, - to: &str, - args: OpRename, - ) -> opendal::Result { - self.memory.service().rename(ctx, from, to, args).await - } - async fn presign( - &self, - ctx: &OperationContext, - path: &str, - args: OpPresign, - ) -> opendal::Result { - self.memory.service().presign(ctx, path, args).await - } -} - -struct Upload { - provider: MultipartProvider, - path: String, - fault: Fault, - reorder: bool, - second_part: Notify, -} - -fn injected_error() -> opendal::Error { - opendal::Error::new(opendal::ErrorKind::Unexpected, "injected multipart failure") -} - -impl oio::MultipartWrite for Upload { - async fn write_once(&self, size: u64, body: Buffer) -> opendal::Result { - assert_eq!( - self.fault, - Fault::None, - "failure test must exercise multipart" - ); - self.provider.memory.write(&self.path, body).await?; - Ok(Metadata::default().with_content_length(size)) - } - - async fn initiate_part(&self) -> opendal::Result { - self.provider - .state - .lock() - .unwrap() - .uploads - .insert(self.path.clone(), BTreeMap::new()); - Ok(self.path.clone()) - } - - async fn write_part( - &self, - upload_id: &str, - part_number: usize, - size: u64, - body: Buffer, - ) -> opendal::Result { - assert_eq!(size as usize, body.len()); - if self.reorder && part_number == 0 { - self.second_part.notified().await; - } - { - let mut state = self.provider.state.lock().unwrap(); - state.completed_parts.push(part_number); - state - .uploads - .get_mut(upload_id) - .unwrap() - .insert(part_number, body.to_bytes()); - } - if part_number == 1 { - self.second_part.notify_one(); - } - if self.fault == Fault::Part && part_number == 0 { - return Err(injected_error()); - } - Ok(oio::MultipartPart { - part_number, - etag: part_number.to_string(), - checksum: None, - size: Some(size), - }) - } - - async fn complete_part( - &self, - upload_id: &str, - parts: &[oio::MultipartPart], - ) -> opendal::Result { - if self.fault == Fault::Close { - return Err(injected_error()); - } - let mut bytes = Vec::new(); - { - let mut state = self.provider.state.lock().unwrap(); - let uploaded = state.uploads.remove(upload_id).unwrap(); - assert_eq!(parts.len(), uploaded.len()); - for (expected, part) in parts.iter().enumerate() { - assert_eq!(part.part_number, expected); - bytes.extend_from_slice(&uploaded[&part.part_number]); - } - } - let size = bytes.len() as u64; - self.provider.memory.write(&self.path, bytes).await?; - Ok(Metadata::default().with_content_length(size)) - } - - async fn abort_part(&self, upload_id: &str) -> opendal::Result<()> { - let mut state = self.provider.state.lock().unwrap(); - state.aborts += 1; - state.uploads.remove(upload_id); - Ok(()) - } -} diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 0e16766d6..7019e9bc4 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -44,8 +44,6 @@ use std::time::{Duration, Instant}; #[derive(Debug, Default)] pub(crate) struct DataFileReadTiming { file_read_nanos: AtomicU64, - file_read_bytes: AtomicU64, - file_read_requests: AtomicU64, parquet_decode_nanos: AtomicU64, file_schema_open_nanos: AtomicU64, first_batch_wait_nanos: AtomicU64, @@ -53,17 +51,11 @@ pub(crate) struct DataFileReadTiming { } impl DataFileReadTiming { - pub(crate) fn add_file_read(&self, duration: Duration) { + fn add_file_read(&self, duration: Duration) { self.file_read_nanos .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); } - fn add_file_io(&self, bytes: usize) { - self.file_read_bytes - .fetch_add(bytes as u64, Ordering::Relaxed); - self.file_read_requests.fetch_add(1, Ordering::Relaxed); - } - fn add_parquet_decode(&self, duration: Duration) { self.parquet_decode_nanos .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); @@ -87,20 +79,6 @@ impl DataFileReadTiming { Duration::from_nanos(self.file_read_nanos.load(Ordering::Relaxed)) } - pub(crate) fn file_io(&self) -> (u64, u64) { - ( - self.file_read_bytes.load(Ordering::Relaxed), - self.file_read_requests.load(Ordering::Relaxed), - ) - } - - pub(crate) fn wrap_reader(self: &Arc, inner: Box) -> Box { - Box::new(TimedFileRead { - inner, - timing: Arc::clone(self), - }) - } - pub(crate) fn parquet_decode(&self) -> Duration { Duration::from_nanos(self.parquet_decode_nanos.load(Ordering::Relaxed)) } @@ -124,8 +102,6 @@ impl FileRead for TimedFileRead { let start = Instant::now(); let result = self.inner.read(range).await; self.timing.add_file_read(start.elapsed()); - self.timing - .add_file_io(result.as_ref().map_or(0, bytes::Bytes::len)); result } } @@ -561,7 +537,10 @@ impl DataFileReader { timing.add_file_read(start.elapsed()); } let file_reader: Box = match read_timing.as_ref() { - Some(timing) => timing.wrap_reader(Box::new(file_reader)), + Some(timing) => Box::new(TimedFileRead { + inner: Box::new(file_reader), + timing: Arc::clone(timing), + }), None => Box::new(file_reader), }; let is_parquet = path_to_read.to_ascii_lowercase().ends_with(".parquet"); 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 ca48a3f4c..dcda47217 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs @@ -89,7 +89,7 @@ pub(super) fn validate_vector_batch_ranges<'a>( message: format!("vindex vector extraction got unexpected _ROW_ID {row_id}"), source: None, })?; - if row_id != *expected_row_id || row_id > range.to() { + if row_id != *expected_row_id { return Err(Error::DataInvalid { message: format!( "vindex vector extraction expected _ROW_ID {}, got {}", 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 e74c22714..20c7d6d13 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -770,16 +770,56 @@ async fn vindex_incremental_build_indexes_only_new_rows() { .collect(), ) .await; - let first_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .with_options(HashMap::from([( - "ivf-flat.train.sample-ratio".to_string(), - "0.2".to_string(), - )])) - .execute() + let build_options = + HashMap::from([("ivf-flat.train.sample-ratio".to_string(), "0.2".to_string())]); + let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); + let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) .await .unwrap(); + let core_options = CoreOptions::new(table.schema().options()); + let shards = plan_vindex_shards( + table_path, + table.schema().partition_keys(), + table.schema().fields(), + &core_options, + snapshot.id(), + entries, + core_options.global_index_row_count_per_shard().unwrap(), + &[], + ) + .unwrap(); + assert_eq!(shards.len(), 1); + let options = crate::vindex::VindexVectorIndexOptions::new( + table.schema().options(), + &build_options, + IVF_FLAT_IDENTIFIER, + find_index_field(&table, "embedding").unwrap(), + ) + .unwrap(); + let training_rows = paimon_vindex_core::autotune::default_training_vector_count( + checked_training_vector_count( + (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, + options.train_sample_ratio, + ) + .unwrap(), + options.config.nlist(), + ) + .unwrap(); + let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + builder + .with_index_column("embedding") + .with_options(build_options); + assert!(builder + .sparse_training_ranges(&shards[0], "embedding", training_rows) + .await + .unwrap() + .is_some()); + let first_built = builder.execute().await.unwrap(); assert!(first_built > 0, "first build must index the initial rows"); // First appended row-id, derived from the data manifest (never hard-coded). @@ -808,13 +848,51 @@ async fn vindex_incremental_build_indexes_only_new_rows() { ) .await; - // End-to-end: build #2 must SUCCEED and index the appended rows. - let second_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() + let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) .await .unwrap(); + let shards = plan_vindex_shards( + table_path, + table.schema().partition_keys(), + table.schema().fields(), + &core_options, + snapshot.id(), + entries, + core_options.global_index_row_count_per_shard().unwrap(), + &indexed_coverage, + ) + .unwrap(); + assert_eq!(shards.len(), 1); + let options = crate::vindex::VindexVectorIndexOptions::new( + table.schema().options(), + &HashMap::new(), + IVF_FLAT_IDENTIFIER, + find_index_field(&table, "embedding").unwrap(), + ) + .unwrap(); + let training_rows = paimon_vindex_core::autotune::default_training_vector_count( + checked_training_vector_count( + (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, + options.train_sample_ratio, + ) + .unwrap(), + options.config.nlist(), + ) + .unwrap(); + let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + builder.with_index_column("embedding"); + assert!(builder + .sparse_training_ranges(&shards[0], "embedding", training_rows) + .await + .unwrap() + .is_none()); + // End-to-end: build #2 must SUCCEED and index the appended rows. + let second_built = builder.execute().await.unwrap(); assert!(second_built > 0, "appended rows must be indexed"); let all_files = latest_vindex_index_files(&table).await; @@ -850,79 +928,46 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } -#[test] -fn vindex_build_logs_read_phases() { - // Run the real sparse and fallback builds in a separate process so the - // timing environment variable and stderr capture cannot race other tests. - let output = std::process::Command::new(std::env::current_exe().unwrap()) - .args([ - "--exact", - "table::vindex_index_build_builder::tests::vindex_incremental_build_indexes_only_new_rows", - "--nocapture", - ]) - .env("PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING", "1") - .output() +#[tokio::test] +async fn vindex_sparse_probe_errors_fall_back() { + let table_path = "memory:/test_vindex_probe_fallback"; + let table = vindex_e2e_table(table_path, "1024"); + let mut shard = plan( + vec![manifest_entry(data_file("broken.parquet", Some(0), 1024))], + 1024, + ) + .unwrap() + .remove(0); + shard.bucket_path = table_path.to_string(); + let builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); + assert!(plan_ivf_training_ranges(&shard, 192).unwrap().is_some()); + + // A failed open and an unreadable footer both disable only the optimization. + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .unwrap() + .is_none()); + let path = shard.files[0].data_file_path(&shard.bucket_path); + table + .file_io() + .new_output(&path) + .unwrap() + .write(vec![0; 128].into()) + .await .unwrap(); - let stderr = String::from_utf8(output.stderr).unwrap(); - assert!(output.status.success(), "{stderr}"); - let events = stderr - .lines() - .filter(|line| line.starts_with("event=paimon_vector_index_build")) - .map(|line| { - line.split_whitespace() - .filter_map(|field| field.split_once('=')) - .collect::>() - }) - .collect::>(); - let plans = events - .iter() - .filter(|event| event["event"] == "paimon_vector_index_build_plan") - .collect::>(); - assert_eq!(plans.len(), 2, "missing build provenance: {stderr}"); - assert_eq!(plans[0]["sparse"], "true"); - assert_eq!(plans[1]["sparse"], "false"); - for plan in &plans { - assert!(plan["snapshot_id"].parse::().unwrap() > 0); - assert!(plan["training_seed"].parse::().is_ok()); - assert_eq!(plan["parquet_row_group_parallelism"], "8"); - assert_eq!(plan["parquet_max_inflight_bytes"], "268435456"); - } - let phases = events - .iter() - .filter(|event| event["event"] == "paimon_vector_index_build_read") - .collect::>(); - assert_eq!( - phases - .iter() - .map(|event| event["phase"]) - .collect::>(), - ["probe", "sample", "full_scan", "full_scan"] - ); - for phase in &phases { - assert_eq!(phase["io_scope"], "file_read_wrapper"); - assert!(phase["read_bytes"].parse::().unwrap() > 0); - assert!(phase["read_calls"].parse::().unwrap() > 0); - assert!(phase["read_ms"].parse::().unwrap() >= 0.0); - } - let totals = events - .iter() - .filter(|event| event["event"] == "paimon_vector_index_build") - .collect::>(); - for (total, reads) in [(totals[0], &phases[1..3]), (totals[1], &phases[3..4])] { - for (total_key, phase_key) in [ - ("oss_read_bytes", "read_bytes"), - ("oss_range_requests", "read_calls"), - ] { - assert_eq!( - total[total_key].parse::().unwrap(), - reads - .iter() - .map(|event| event[phase_key].parse::().unwrap()) - .sum::(), - "phase counters must exclude probe and partition the existing total" - ); - } - } + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .unwrap() + .is_none()); + + // Invalid source metadata is not an optional probe failure. + shard.files[0].file_size = -1; + assert!(builder + .sparse_training_ranges(&shard, "embedding", 192) + .await + .is_err()); } #[tokio::test] @@ -991,98 +1036,6 @@ async fn vindex_small_training_sample_preserves_tail_cluster_recall() { ); } -#[tokio::test] -async fn vindex_upload_failure_preserves_committed_index() { - use crate::io::multipart_test::{Fault, MultipartProvider}; - - for fault in [Fault::Part, Fault::Close] { - let provider = MultipartProvider::new(128); - let table_path = "memory:/test_vindex_upload_failure"; - let table = test_table_with_io( - provider.file_io(), - table_path, - vindex_schema_builder(vindex_e2e_options("3")) - .build() - .unwrap(), - ); - setup_dirs(table.file_io(), table_path).await; - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], - ) - .await; - table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .unwrap(); - let existing = latest_vindex_index_files(&table).await; - let old_path = format!("{table_path}/{INDEX_DIR}/{}", existing[0].file_name); - let old_bytes = table - .file_io() - .new_input(&old_path) - .unwrap() - .read() - .await - .unwrap(); - let mut search = table.new_vector_search_builder(); - search - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(1); - let old_result = search.execute().await.unwrap(); - assert!(!old_result.is_empty()); - - write_vectors(&table, vec![4, 5, 6, 7, 8, 9], vec![vec![-1.0, 0.0]; 6]).await; - let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); - let before = snapshots.get_latest_snapshot().await.unwrap().unwrap(); - { - let mut state = provider.state.lock().unwrap(); - state.fault = fault; - state.fail_on_index = state.index_writes + 2; - state.concurrency.clear(); - } - let error = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .expect_err("injected upload must fail"); - assert!( - error.to_string().contains("injected multipart failure"), - "{error}" - ); - let after = snapshots.get_latest_snapshot().await.unwrap().unwrap(); - assert_eq!(before.id(), after.id()); - assert_eq!(before.index_manifest(), after.index_manifest()); - assert_eq!(latest_vindex_index_files(&table).await, existing); - assert_eq!( - table - .file_io() - .new_input(&old_path) - .unwrap() - .read() - .await - .unwrap(), - old_bytes - ); - assert_eq!(search.execute().await.unwrap(), old_result); - let files = table - .file_io() - .list_status(&format!("{table_path}/{INDEX_DIR}/")) - .await - .unwrap(); - assert_eq!(files.len(), 1); - assert!(table.file_io().exists(&old_path).await.unwrap()); - let state = provider.state.lock().unwrap(); - assert_eq!(state.concurrency, [1, 1]); - assert_eq!(state.uploads.len(), 1); - assert_eq!(state.aborts, 0); - } -} - #[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 2181ed5d7..26f7039dd 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -use super::planning::VindexIndexShard; -use crate::table::data_file_reader::DataFileReadTiming; use std::sync::OnceLock; use std::time::Duration; @@ -29,35 +27,10 @@ pub(super) fn vector_index_build_timing_enabled() -> bool { }) } -// Snapshot only at drained phase boundaries. These are completed FileRead -// calls (plus reader-open time), not OSS wire requests; concurrent read_ms -// is cumulative and must not be added to wall-clock stage durations. -pub(super) fn log_read_phase( - phase: &str, - shard: &VindexIndexShard, - timing: &DataFileReadTiming, - previous: (Duration, u64, u64), -) -> (Duration, u64, u64) { - let read = timing.file_read(); - let (bytes, calls) = timing.file_io(); - eprintln!( - "event=paimon_vector_index_build_read phase={phase} snapshot_id={} row_range_start={} row_range_end={} io_scope=file_read_wrapper read_ms={:.3} read_bytes={} read_calls={}", - shard.snapshot_id, - shard.row_range_start, - shard.row_range_end, - read.saturating_sub(previous.0).as_secs_f64() * 1000.0, - bytes.saturating_sub(previous.1), - calls.saturating_sub(previous.2), - ); - (read, bytes, calls) -} - pub(super) struct VectorIndexBuildTiming { pub(super) total_without_commit: Duration, pub(super) source_batch_wait: Duration, pub(super) oss_read: Duration, - pub(super) oss_read_bytes: u64, - pub(super) oss_range_requests: u64, pub(super) parquet_decode: Duration, pub(super) file_schema_open: Duration, pub(super) first_batch_wait: Duration, @@ -74,18 +47,11 @@ pub(super) struct VectorIndexBuildTiming { pub(super) raw_temp_reread: Duration, pub(super) index_add: Duration, pub(super) full_scan_add: Duration, - pub(super) pipeline_blocked: Duration, - pub(super) producer_blocked: Duration, - pub(super) consumer_validate_add: Duration, pub(super) serialize_upload: Duration, pub(super) rows: usize, pub(super) training_rows_seen: usize, pub(super) training_rows_retained: usize, pub(super) batch_count: usize, - pub(super) batch_bytes_min: usize, - pub(super) batch_bytes_max: usize, - pub(super) batch_bytes_total: usize, - pub(super) peak_ready_batches: usize, pub(super) raw_temp_bytes: usize, pub(super) index_bytes: u64, pub(super) data_file_count: usize, @@ -113,22 +79,17 @@ 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={} batch_bytes_min={} batch_bytes_max={} batch_bytes_total={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} oss_read_bytes={} oss_range_requests={} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} capability_check_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} pipeline_blocked_ms={:.3} producer_blocked_ms={:.3} consumer_validate_add_ms={:.3} data_file_count={} data_file_read_concurrency=1 peak_ready_batches={} total_ms={:.3} unattributed_ms={:.3}", + "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} capability_check_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} data_file_count={} data_file_read_concurrency=1 total_ms={:.3} unattributed_ms={:.3}", index_type, self.file_name, self.rows, self.training_rows_seen, self.training_rows_retained, self.batch_count, - self.batch_bytes_min, - self.batch_bytes_max, - self.batch_bytes_total, self.raw_temp_bytes, self.index_bytes, self.source_batch_wait.as_secs_f64() * 1000.0, self.oss_read.as_secs_f64() * 1000.0, - self.oss_read_bytes, - self.oss_range_requests, self.parquet_decode.as_secs_f64() * 1000.0, self.file_schema_open.as_secs_f64() * 1000.0, self.first_batch_wait.as_secs_f64() * 1000.0, @@ -147,11 +108,7 @@ impl VectorIndexBuildTiming { commit.as_secs_f64() * 1000.0, self.sample_read.as_secs_f64() * 1000.0, self.full_scan_add.as_secs_f64() * 1000.0, - self.pipeline_blocked.as_secs_f64() * 1000.0, - self.producer_blocked.as_secs_f64() * 1000.0, - self.consumer_validate_add.as_secs_f64() * 1000.0, self.data_file_count, - self.peak_ready_batches, total.as_secs_f64() * 1000.0, unattributed.as_secs_f64() * 1000.0, ); diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs index 3fa53a780..8ece02fe1 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -20,7 +20,7 @@ use super::extraction::{ validate_vector_batch_ranges, }; use super::planning::{ivf_training_seed, plan_ivf_training_ranges, VindexIndexShard}; -use super::timing::{log_read_phase, vector_index_build_timing_enabled, VectorIndexBuildTiming}; +use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, @@ -53,62 +53,24 @@ pub(super) struct BuiltIndexFile { } impl<'a> VindexIndexBuildBuilder<'a> { - pub(super) async fn build_index_file( + /// Plans sparse training reads, falling back when the optional file probe fails. + pub(super) async fn sparse_training_ranges( &self, shard: &VindexIndexShard, index_column: &str, - 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 raw_temp_write = Duration::ZERO; - 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 mut batch_count = 0usize; + training_rows_retained: usize, + ) -> Result>> { let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; let row_count_usize = usize::try_from(row_count).map_err(|e| Error::DataInvalid { message: format!("Invalid vindex row count: {row_count}"), source: Some(Box::new(e)), })?; - let dimension_usize = usize::try_from(dimension).map_err(|e| Error::DataInvalid { - message: format!("Invalid vindex dimension: {dimension}"), - source: Some(Box::new(e)), - })?; - if dimension_usize == 0 { - return Err(Error::DataInvalid { - message: "vindex vector dimension must be positive".to_string(), - source: None, - }); - } - let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; - let training_vector_count = - checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { - 0 - } else { - // This only gates sparse reads; errors must preserve the full-scan fallback. - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) - }; let mut sparse_ranges = if training_rows_retained > 0 && training_rows_retained < row_count_usize { plan_ivf_training_ranges(shard, training_rows_retained)? } else { None }; - let capability_start = (timing_enabled && sparse_ranges.is_some()).then(Instant::now); - let capability_read_timing = - capability_start.map(|_| Arc::new(DataFileReadTiming::default())); if let Some(ranges) = sparse_ranges.as_ref() { let mut checks = Vec::new(); let mut usable = true; @@ -157,27 +119,31 @@ impl<'a> VindexIndexBuildBuilder<'a> { .core_options() .parquet_row_group_parallelism()?; let file_io = self.table.file_io(); - let timing = capability_read_timing.as_ref(); let mut checks = futures::stream::iter(checks) .map(|(path, file_size, local_ranges)| async move { let input = file_io.new_input(&path)?; - let open_start = timing.map(|_| Instant::now()); let reader = Box::new(input.reader().await?); - if let (Some(timing), Some(start)) = (timing, open_start) { - timing.add_file_read(start.elapsed()); - } - let reader = match timing { - Some(timing) => timing.wrap_reader(reader), - None => reader, - }; has_beneficial_offset_index(reader, file_size, index_column, &local_ranges) .await }) .buffer_unordered(concurrency); - while let Some(file_usable) = checks.try_next().await? { - if !file_usable { - usable = false; - break; + while let Some(result) = checks.next().await { + match result { + Ok(true) => {} + Ok(false) => { + usable = false; + break; + } + Err(error) => { + log::warn!( + "vindex sparse training probe failed for column '{}' in shard [{}, {}]; falling back to a full scan: {}", + index_column, + shard.row_range_start, + shard.row_range_end, + error, + ); + return Ok(None); + } } } } @@ -191,10 +157,61 @@ impl<'a> VindexIndexBuildBuilder<'a> { sparse_ranges = None; } } - let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); - if let Some(timing) = capability_read_timing.as_ref() { - log_read_phase("probe", shard, timing, Default::default()); + Ok(sparse_ranges) + } + + pub(super) async fn build_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); + let mut source_batch_wait = Duration::ZERO; + let mut raw_temp_write = Duration::ZERO; + 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 mut batch_count = 0usize; + 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 = usize::try_from(dimension).map_err(|e| Error::DataInvalid { + message: format!("Invalid vindex dimension: {dimension}"), + source: Some(Box::new(e)), + })?; + if dimension_usize == 0 { + return Err(Error::DataInvalid { + message: "vindex vector dimension must be positive".to_string(), + source: None, + }); } + let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; + let training_vector_count = + checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; + let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { + 0 + } else { + // This only gates sparse reads; errors must preserve the full-scan fallback. + default_training_vector_count(training_vector_count, options.config.nlist()) + .unwrap_or(0) + }; + let capability_start = timing_enabled.then(Instant::now); + let sparse_ranges = self + .sparse_training_ranges(shard, index_column, training_rows_retained) + .await?; + let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); if let Some(budget) = parquet_read_budget.as_ref() { eprintln!( "event=paimon_vector_index_build_plan snapshot_id={} row_range_start={} row_range_end={} source_bucket={} sparse={} training_seed={} training_range_count={} parquet_row_group_parallelism={} parquet_max_inflight_bytes={}", @@ -217,19 +234,11 @@ impl<'a> VindexIndexBuildBuilder<'a> { })?; let mut sample_read = Duration::ZERO; let mut full_scan_add = Duration::ZERO; - let mut pipeline_blocked = Duration::ZERO; - let mut producer_blocked = Duration::ZERO; - let mut consumer_validate_add = Duration::ZERO; - let mut batch_bytes_min = 0usize; - let mut batch_bytes_max = 0usize; - let mut batch_bytes_total = 0usize; - let mut peak_ready_batches = 0usize; let raw_temp_reread; let index_add; let train_finish; let mut bytes_written = 0usize; let training_rows_seen; - let mut sample_io = Default::default(); let writer = if let Some(ranges) = sparse_ranges { let sample_start = timing_enabled.then(Instant::now); @@ -283,9 +292,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { } training_rows_seen = rows_seen; sample_read = sample_start.map_or(Duration::ZERO, |start| start.elapsed()); - if let Some(timing) = read_timing.as_ref() { - sample_io = log_read_phase("sample", shard, timing, Default::default()); - } let train_start = timing_enabled.then(Instant::now); let training = tokio::task::spawn_blocking(move || trainer.finish()) @@ -332,24 +338,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut expected_row_id = row_range_start; let mut rows_added = 0usize; let mut batches_added = 0usize; - let mut blocked = Duration::ZERO; - let mut validate_add = Duration::ZERO; - let mut batch_bytes_min = usize::MAX; - let mut batch_bytes_max = 0usize; - let mut batch_bytes_total = 0usize; let mut ids = Vec::new(); - loop { - let wait_start = timing_enabled.then(Instant::now); - let batch = receiver.blocking_recv(); - if let Some(start) = wait_start { - blocked = blocked.saturating_add(start.elapsed()); - } - let Some(batch) = batch else { break }; - let validate_add_start = timing_enabled.then(Instant::now); - let batch_bytes = batch.get_array_memory_size(); - batch_bytes_min = batch_bytes_min.min(batch_bytes); - batch_bytes_max = batch_bytes_max.max(batch_bytes); - batch_bytes_total = batch_bytes_total.saturating_add(batch_bytes); + while let Some(batch) = receiver.blocking_recv() { let vectors = validate_vector_batch( &batch, &index_column, @@ -375,27 +365,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { message: format!("Failed to add vectors to vindex index: {e}"), source: Some(Box::new(e)), })?; - if let Some(start) = validate_add_start { - validate_add = validate_add.saturating_add(start.elapsed()); - } rows_added = batch_end; batches_added += 1; } - Ok(( - writer, - rows_added, - expected_row_id, - batches_added, - blocked, - validate_add, - if batches_added == 0 { - 0 - } else { - batch_bytes_min - }, - batch_bytes_max, - batch_bytes_total, - )) + Ok((writer, rows_added, expected_row_id, batches_added)) }); let mut producer_error = None; @@ -413,33 +386,17 @@ impl<'a> VindexIndexBuildBuilder<'a> { break; } }; - let send_start = timing_enabled.then(Instant::now); - let send_result = sender.send(batch).await; - if let Some(start) = send_start { - producer_blocked = producer_blocked.saturating_add(start.elapsed()); - } - if send_result.is_err() { + if sender.send(batch).await.is_err() { break; } - peak_ready_batches = peak_ready_batches - .max(READ_ADD_QUEUE_CAPACITY.saturating_sub(sender.capacity())); } drop(sender); let consumer_result = consumer.await; - let ( - writer, - rows_added, - next_row_id, - batches_added, - blocked, - validate_add, - min_bytes, - max_bytes, - total_bytes, - ) = consumer_result.map_err(|e| Error::UnexpectedError { - message: format!("vindex add task failed: {e}"), - source: None, - })??; + let (writer, rows_added, next_row_id, batches_added) = + consumer_result.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })??; if let Some(error) = producer_error { return Err(error); } @@ -452,11 +409,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { }); } batch_count = batches_added; - pipeline_blocked = blocked; - consumer_validate_add = validate_add; - batch_bytes_min = min_bytes; - batch_bytes_max = max_bytes; - batch_bytes_total = total_bytes; full_scan_add = full_scan_start.map_or(Duration::ZERO, |start| start.elapsed()); raw_temp_reread = Duration::ZERO; index_add = Duration::ZERO; @@ -703,10 +655,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { result.0 }; - if let Some(timing) = read_timing.as_ref() { - log_read_phase("full_scan", shard, timing, sample_io); - } - let serialize_upload_start = timing_enabled.then(Instant::now); self.table .file_io() @@ -783,9 +731,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { .map_or((Duration::ZERO, Duration::ZERO), |timing| { (timing.file_read(), timing.parquet_decode()) }); - let (oss_read_bytes, oss_range_requests) = read_timing - .as_ref() - .map_or((0, 0), |timing| timing.file_io()); let (file_schema_open, first_batch_wait, remaining_batch_wait) = read_timing .as_ref() .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing| { @@ -798,8 +743,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { total_without_commit: start.elapsed(), source_batch_wait, oss_read, - oss_read_bytes, - oss_range_requests, parquet_decode, file_schema_open, first_batch_wait, @@ -816,18 +759,11 @@ impl<'a> VindexIndexBuildBuilder<'a> { raw_temp_reread, index_add, full_scan_add, - pipeline_blocked, - producer_blocked, - consumer_validate_add, serialize_upload, rows: row_count_usize, training_rows_seen, training_rows_retained, batch_count, - batch_bytes_min, - batch_bytes_max, - batch_bytes_total, - peak_ready_batches, raw_temp_bytes: bytes_written, index_bytes: status.size, data_file_count: shard.files.len(), From c866b823bf8fd099ce5339011e33a00473069f71 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 21:16:39 +0800 Subject: [PATCH 11/24] fix(vindex): scope sparse savings to current shard --- crates/paimon/src/arrow/format/parquet.rs | 110 +++++++++++------- .../vindex_index_build_builder/writer.rs | 18 ++- 2 files changed, 84 insertions(+), 44 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 65d003082..eb32e215a 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -71,7 +71,8 @@ pub(crate) async fn has_beneficial_offset_index( reader: Box, file_size: u64, column_name: &str, - row_ranges: &[RowRange], + sample_ranges: &[RowRange], + scan_ranges: &[RowRange], ) -> crate::Result { let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); let mut reader = ArrowFileReader::new(file_size, reader.into()); @@ -79,14 +80,16 @@ pub(crate) async fn has_beneficial_offset_index( Ok(metadata_has_beneficial_offset_index( &metadata, column_name, - row_ranges, + sample_ranges, + scan_ranges, )) } fn metadata_has_beneficial_offset_index( metadata: &ParquetMetaData, column_name: &str, - row_ranges: &[RowRange], + sample_ranges: &[RowRange], + scan_ranges: &[RowRange], ) -> bool { let columns = metadata .file_metadata() @@ -109,30 +112,27 @@ fn metadata_has_beneficial_offset_index( if columns.is_empty() || offset_index.len() != metadata.row_groups().len() { return false; } - let mut selection = build_row_ranges_selection(metadata.row_groups(), row_ranges); + let mut sample_selection = build_row_ranges_selection(metadata.row_groups(), sample_ranges); + let mut scan_selection = build_row_ranges_selection(metadata.row_groups(), scan_ranges); let mut checked = false; - let mut full_bytes = 0u64; - let mut selected_bytes = 0u64; + let mut sample_bytes = 0u64; + let mut scan_bytes = 0u64; for (row_group, indexes) in metadata.row_groups().iter().zip(offset_index) { let Ok(row_count) = usize::try_from(row_group.num_rows()) else { return false; }; - let row_group_selection = selection.split_off(row_count); - let mut selected_ranges = Vec::new(); + let sample_row_group_selection = sample_selection.split_off(row_count); + let scan_row_group_selection = scan_selection.split_off(row_count); + let sample_selected = sample_row_group_selection.selects_any(); + let scan_selected = scan_row_group_selection.selects_any(); + checked |= sample_selected; + if !sample_selected && !scan_selected { + continue; + } + let mut sample_byte_ranges = Vec::new(); + let mut scan_byte_ranges = Vec::new(); for index in &columns { let column = row_group.column(*index); - let Ok(column_bytes) = u64::try_from(column.compressed_size()) else { - return false; - }; - let Some(total) = full_bytes.checked_add(column_bytes) else { - return false; - }; - full_bytes = total; - - if !row_group_selection.selects_any() { - continue; - } - checked = true; let Some(page_locations) = indexes.get(*index).map(|index| index.page_locations()) else { return false; @@ -150,30 +150,40 @@ fn metadata_has_beneficial_offset_index( let Ok(first_page_offset) = u64::try_from(first_page.offset) else { return false; }; - if column_start < first_page_offset { - selected_ranges.push(column_start..first_page_offset); + for (selection, selected_ranges) in [ + (&sample_row_group_selection, &mut sample_byte_ranges), + (&scan_row_group_selection, &mut scan_byte_ranges), + ] { + if !selection.selects_any() { + continue; + } + if column_start < first_page_offset { + selected_ranges.push(column_start..first_page_offset); + } + selected_ranges.extend(selection.scan_ranges(page_locations)); } - selected_ranges.extend(row_group_selection.scan_ranges(page_locations)); } - if row_group_selection.selects_any() { - let Some(group_selected_bytes) = - merge_byte_ranges(&selected_ranges, RANGE_COALESCE_BYTES) - .into_iter() - .try_fold(0u64, |total, range| { - total.checked_add(range.end.checked_sub(range.start)?) - }) + for (selected_ranges, total_bytes) in [ + (sample_byte_ranges, &mut sample_bytes), + (scan_byte_ranges, &mut scan_bytes), + ] { + let Some(group_bytes) = merge_byte_ranges(&selected_ranges, RANGE_COALESCE_BYTES) + .into_iter() + .try_fold(0u64, |total, range| { + total.checked_add(range.end.checked_sub(range.start)?) + }) else { return false; }; - let Some(total) = selected_bytes.checked_add(group_selected_bytes) else { + let Some(total) = total_bytes.checked_add(group_bytes) else { return false; }; - selected_bytes = total; + *total_bytes = total; } } // Sparse training is followed by a full scan, so require it to skip at // least half of the projected bytes instead of accepting marginal savings. - checked && selected_bytes <= full_bytes / 2 + checked && sample_bytes <= scan_bytes / 2 } enum ParquetRowGroupMessage { @@ -3715,21 +3725,37 @@ mod tests { let bytes = write_multi_row_group_parquet(10, 30, EnabledStatistics::Chunk, false).await; let metadata = load_metadata_with_page_index(&bytes, true); assert!( - !metadata_has_beneficial_offset_index(&metadata, "value", &[RowRange::new(0, 19)]), + !metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(0, 19)], + &[RowRange::new(0, 29)], + ), "reading two of three row groups is not sufficiently sparse" ); let bytes = write_multi_page_parquet(10, 80).await; let metadata = load_metadata_with_page_index(&bytes, true); + assert!( + !metadata_has_beneficial_offset_index( + &metadata, + "value", + &[RowRange::new(20, 38)], + &[RowRange::new(20, 39)], + ), + "near-full reads within one shard must not use the whole file as the baseline" + ); assert!(!metadata_has_beneficial_offset_index( &metadata, "value", - &[RowRange::new(0, 69)] + &[RowRange::new(0, 69)], + &[RowRange::new(0, 79)], )); assert!(!metadata_has_beneficial_offset_index( &metadata, "value", - &[RowRange::new(0, 0), RowRange::new(79, 79)] + &[RowRange::new(0, 0), RowRange::new(79, 79)], + &[RowRange::new(0, 79)], )); let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; @@ -3738,18 +3764,21 @@ mod tests { assert!(!metadata_has_beneficial_offset_index( &metadata, "value", - &[RowRange::new(0, 19)] + &[RowRange::new(0, 19)], + &[RowRange::new(0, 19)], )); let sparse_ranges = [RowRange::new(0, 0)]; assert!(metadata_has_beneficial_offset_index( &metadata, "value", - &sparse_ranges + &sparse_ranges, + &[RowRange::new(0, 19)], )); assert!(!metadata_has_beneficial_offset_index( &metadata, "missing", - &sparse_ranges + &sparse_ranges, + &[RowRange::new(0, 19)], )); let bytes_without_index = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; @@ -3757,7 +3786,8 @@ mod tests { assert!(!metadata_has_beneficial_offset_index( &metadata_without_index, "value", - &sparse_ranges + &sparse_ranges, + &[RowRange::new(0, 19)], )); } 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 8ece02fe1..1dcfb705b 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -97,6 +97,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { if local_ranges.is_empty() { continue; } + let scan_ranges = vec![RowRange::new( + shard.row_range_start.max(file_start) - file_start, + shard.row_range_end.min(file_end) - file_start, + )]; let path = file.data_file_path(&shard.bucket_path); if !path.to_ascii_lowercase().ends_with(".parquet") { usable = false; @@ -109,7 +113,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { ), source: Some(Box::new(e)), })?; - checks.push((path, file_size, local_ranges)); + checks.push((path, file_size, local_ranges, scan_ranges)); } let found_vector_file = !checks.is_empty(); if usable && found_vector_file { @@ -120,11 +124,17 @@ impl<'a> VindexIndexBuildBuilder<'a> { .parquet_row_group_parallelism()?; let file_io = self.table.file_io(); let mut checks = futures::stream::iter(checks) - .map(|(path, file_size, local_ranges)| async move { + .map(|(path, file_size, local_ranges, scan_ranges)| async move { let input = file_io.new_input(&path)?; let reader = Box::new(input.reader().await?); - has_beneficial_offset_index(reader, file_size, index_column, &local_ranges) - .await + has_beneficial_offset_index( + reader, + file_size, + index_column, + &local_ranges, + &scan_ranges, + ) + .await }) .buffer_unordered(concurrency); while let Some(result) = checks.next().await { From 8ec4e02e3a0c67ec600fcda90cf8e12c3ab47b35 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 17 Sep 2026 13:08:16 +0800 Subject: [PATCH 12/24] feat(vindex): build IVF indexes with granule pipeline --- crates/paimon/src/arrow/format/parquet.rs | 313 +++--- crates/paimon/src/arrow/read_budget.rs | 4 - .../src/table/vindex_index_build_builder.rs | 1 + .../vindex_index_build_builder/extraction.rs | 10 + .../vindex_index_build_builder/pipeline.rs | 911 ++++++++++++++++++ .../vindex_index_build_builder/planning.rs | 85 +- .../table/vindex_index_build_builder/tests.rs | 214 +--- .../vindex_index_build_builder/writer.rs | 823 ++++------------ 8 files changed, 1284 insertions(+), 1077 deletions(-) create mode 100644 crates/paimon/src/table/vindex_index_build_builder/pipeline.rs diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index eb32e215a..56e91d1a0 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -67,30 +67,21 @@ impl ParquetFormatReader { } } -pub(crate) async fn has_beneficial_offset_index( +#[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, - sample_ranges: &[RowRange], - scan_ranges: &[RowRange], -) -> crate::Result { +) -> crate::Result<(Vec, bool)> { let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); let mut reader = ArrowFileReader::new(file_size, reader.into()); let metadata = reader.get_metadata(Some(&options)).await?; - Ok(metadata_has_beneficial_offset_index( - &metadata, - column_name, - sample_ranges, - scan_ranges, - )) -} - -fn metadata_has_beneficial_offset_index( - metadata: &ParquetMetaData, - column_name: &str, - sample_ranges: &[RowRange], - scan_ranges: &[RowRange], -) -> bool { let columns = metadata .file_metadata() .schema_descr() @@ -106,84 +97,105 @@ fn metadata_has_beneficial_offset_index( .then_some(index) }) .collect::>(); - let Some(offset_index) = metadata.offset_index() else { - return false; + 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)); }; - if columns.is_empty() || offset_index.len() != metadata.row_groups().len() { - return false; - } - let mut sample_selection = build_row_ranges_selection(metadata.row_groups(), sample_ranges); - let mut scan_selection = build_row_ranges_selection(metadata.row_groups(), scan_ranges); - let mut checked = false; - let mut sample_bytes = 0u64; - let mut scan_bytes = 0u64; + + 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 false; + return Ok((row_group_granules(&metadata, &columns), false)); }; - let sample_row_group_selection = sample_selection.split_off(row_count); - let scan_row_group_selection = scan_selection.split_off(row_count); - let sample_selected = sample_row_group_selection.selects_any(); - let scan_selected = scan_row_group_selection.selects_any(); - checked |= sample_selected; - if !sample_selected && !scan_selected { - continue; + let 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)); } - let mut sample_byte_ranges = Vec::new(); - let mut scan_byte_ranges = Vec::new(); - for index in &columns { - let column = row_group.column(*index); - let Some(page_locations) = indexes.get(*index).map(|index| index.page_locations()) - else { - return false; - }; - let Some(first_page) = page_locations.first() else { - return false; - }; - let Ok(column_start) = u64::try_from( - column - .dictionary_page_offset() - .unwrap_or_else(|| column.data_page_offset()), - ) else { - return false; - }; - let Ok(first_page_offset) = u64::try_from(first_page.offset) else { - return false; - }; - for (selection, selected_ranges) in [ - (&sample_row_group_selection, &mut sample_byte_ranges), - (&scan_row_group_selection, &mut scan_byte_ranges), - ] { - if !selection.selects_any() { - continue; - } - if column_start < first_page_offset { - selected_ranges.push(column_start..first_page_offset); + 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); + } } - selected_ranges.extend(selection.scan_ranges(page_locations)); + 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, + }); } - for (selected_ranges, total_bytes) in [ - (sample_byte_ranges, &mut sample_bytes), - (scan_byte_ranges, &mut scan_bytes), - ] { - let Some(group_bytes) = merge_byte_ranges(&selected_ranges, RANGE_COALESCE_BYTES) - .into_iter() - .try_fold(0u64, |total, range| { - total.checked_add(range.end.checked_sub(range.start)?) - }) - else { - return false; - }; - let Some(total) = total_bytes.checked_add(group_bytes) else { - return false; - }; - *total_bytes = total; - } + cursor += row_group.num_rows(); } - // Sparse training is followed by a full scan, so require it to skip at - // least half of the projected bytes instead of accepting marginal savings. - checked && sample_bytes <= scan_bytes / 2 + 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 { @@ -2366,6 +2378,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. @@ -2425,9 +2444,8 @@ fn split_ranges_for_concurrency(merged: Vec>, concurrency: usize) -> mod tests { use super::build_parquet_row_filter; use super::{ - forward_row_group_batches, metadata_has_beneficial_offset_index, 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, @@ -3720,77 +3738,6 @@ mod tests { assert_eq!(selection.row_count(), 10); } - #[tokio::test] - async fn test_sparse_row_selection_requires_offset_index_and_page_savings() { - let bytes = write_multi_row_group_parquet(10, 30, EnabledStatistics::Chunk, false).await; - let metadata = load_metadata_with_page_index(&bytes, true); - assert!( - !metadata_has_beneficial_offset_index( - &metadata, - "value", - &[RowRange::new(0, 19)], - &[RowRange::new(0, 29)], - ), - "reading two of three row groups is not sufficiently sparse" - ); - - let bytes = write_multi_page_parquet(10, 80).await; - let metadata = load_metadata_with_page_index(&bytes, true); - assert!( - !metadata_has_beneficial_offset_index( - &metadata, - "value", - &[RowRange::new(20, 38)], - &[RowRange::new(20, 39)], - ), - "near-full reads within one shard must not use the whole file as the baseline" - ); - assert!(!metadata_has_beneficial_offset_index( - &metadata, - "value", - &[RowRange::new(0, 69)], - &[RowRange::new(0, 79)], - )); - assert!(!metadata_has_beneficial_offset_index( - &metadata, - "value", - &[RowRange::new(0, 0), RowRange::new(79, 79)], - &[RowRange::new(0, 79)], - )); - - let bytes = write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, false).await; - let metadata = load_metadata_with_page_index(&bytes, true); - - assert!(!metadata_has_beneficial_offset_index( - &metadata, - "value", - &[RowRange::new(0, 19)], - &[RowRange::new(0, 19)], - )); - let sparse_ranges = [RowRange::new(0, 0)]; - assert!(metadata_has_beneficial_offset_index( - &metadata, - "value", - &sparse_ranges, - &[RowRange::new(0, 19)], - )); - assert!(!metadata_has_beneficial_offset_index( - &metadata, - "missing", - &sparse_ranges, - &[RowRange::new(0, 19)], - )); - let bytes_without_index = - write_multi_row_group_parquet(10, 20, EnabledStatistics::Chunk, true).await; - let metadata_without_index = load_metadata_with_page_index(&bytes_without_index, true); - assert!(!metadata_has_beneficial_offset_index( - &metadata_without_index, - "value", - &sparse_ranges, - &[RowRange::new(0, 19)], - )); - } - #[tokio::test] async fn test_row_group_selection_in_fails_open_on_unusable_stats() { let fields = vec![int_field("id"), int_field("value")]; @@ -3868,6 +3815,50 @@ 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", + ) + .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", + ) + .await + .unwrap(); + assert!(!page_level); + assert_eq!(granules.len(), 3); + assert_eq!( + granules + .iter() + .map(|granule| granule.row_count) + .sum::(), + 30 + ); + } + #[derive(Clone)] struct TrackingFileRead { data: Bytes, diff --git a/crates/paimon/src/arrow/read_budget.rs b/crates/paimon/src/arrow/read_budget.rs index 42aa4411f..334ac8e34 100644 --- a/crates/paimon/src/arrow/read_budget.rs +++ b/crates/paimon/src/arrow/read_budget.rs @@ -140,10 +140,6 @@ impl ReadBudget { self.parallelism } - pub(crate) fn max_inflight_bytes(&self) -> u64 { - self.max_inflight_bytes - } - pub(crate) fn enable_diagnostics(&self) { self.diagnostics.enabled.store(true, Ordering::Relaxed); } diff --git a/crates/paimon/src/table/vindex_index_build_builder.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 dcda47217..37a8bb515 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs @@ -47,9 +47,18 @@ pub(super) fn data_split_for_shard_ranges( pub(super) struct ValidatedVectorBatch<'a> { pub(super) values: &'a [f32], pub(super) bytes: &'a [u8], + pub(super) row_ids: &'a [i64], pub(super) row_count: usize, } +pub(super) fn extract_vector_batch<'a>( + batch: &'a RecordBatch, + index_column: &str, + dimension: usize, +) -> Result> { + validate_vector_batch_with(batch, index_column, dimension, |_| Ok(())) +} + pub(super) fn validate_vector_batch<'a>( batch: &'a RecordBatch, index_column: &str, @@ -245,6 +254,7 @@ fn validate_vector_batch_with<'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..b4fc98a14 --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -0,0 +1,911 @@ +// 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; +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}; +use crate::spec::ROW_ID_FIELD_NAME; +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::time::Instant; +use tokio::sync::mpsc; +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, + first_bytes: u64, + total_bytes: u64, + granule_count: usize, + first_granules: usize, + source: &'static str, +} + +struct SpillRecord { + ids: Vec, + bytes: Vec, +} + +enum AddItem { + Batch(RecordBatch, Vec), + Spilled(Vec, MutableBuffer), +} + +type SpillTask = JoinHandle>; +type ConsumerTask = JoinHandle>; +type ReplayTask = JoinHandle>; + +struct SpillWriter { + sender: mpsc::Sender, + task: SpillTask, +} + +impl SpillWriter { + async fn finish(self) -> Result<(std::fs::File, u64)> { + drop(self.sender); + join_spill(self.task).await + } +} + +struct LivePipeline { + sender: mpsc::Sender, + consumer: ConsumerTask, + replay: ReplayTask, +} + +fn spawn_spill_writer() -> 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 = tokio::task::spawn_blocking(move || -> std::io::Result<_> { + let mut writer = BufWriter::with_capacity(BUFFER_BYTES, file); + let mut bytes_written = 0u64; + while let Some(record) = receiver.blocking_recv() { + 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)?; + bytes_written += 8 + count * 8 + record.bytes.len() as u64; + } + writer.flush()?; + let file = writer.into_inner().map_err(|e| e.into_error())?; + Ok((file, bytes_written)) + }); + Ok(SpillWriter { sender, task }) +} + +async fn join_spill(task: SpillTask) -> Result<(std::fs::File, u64)> { + 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, +) -> ConsumerTask { + tokio::task::spawn_blocking(move || -> Result<_> { + let mut writer = writer; + let mut rows_added = 0usize; + let mut replay_rows = 0usize; + while let Some(item) = receiver.blocking_recv() { + 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(); + } + } + } + Ok((writer, rows_added, replay_rows)) + }) +} + +async fn join_consumer(task: ConsumerTask) -> Result<(VectorIndexWriter, usize, usize)> { + task.await.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })? +} + +async fn join_training( + task: JoinHandle>, +) -> Result { + 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: JoinHandle>, + spill: SpillWriter, + index_column: String, + dimension: usize, +) -> Result<(LivePipeline, u64)> { + let trained = join_training(training).await; + let spilled = spill.finish().await; + let trained = trained?; + let (file, spill_bytes) = spilled?; + let (sender, receiver) = mpsc::channel(QUEUE_CAPACITY); + let consumer = spawn_add_consumer( + VectorIndexWriter::new(trained), + receiver, + index_column, + dimension, + ); + let replay = spawn_replay(file, sender.clone(), dimension); + Ok(( + LivePipeline { + sender, + consumer, + replay, + }, + spill_bytes, + )) +} + +async fn finish_live_pipeline( + pipeline: LivePipeline, +) -> (Result<(VectorIndexWriter, usize, usize)>, Result) { + 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) +} + +fn spawn_replay( + mut file: std::fs::File, + sender: mpsc::Sender, + dimension: usize, +) -> ReplayTask { + tokio::task::spawn_blocking(move || -> std::io::Result { + file.seek(SeekFrom::Start(0))?; + 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 mut header = [0u8; 8]; + let read = reader.read(&mut header)?; + if read == 0 { + 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)?; + ids.extend( + id_bytes + .chunks_exact(8) + .map(|bytes| i64::from_le_bytes(bytes.try_into().unwrap())), + ); + 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..])?; + 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); + } + } + } + if !ids.is_empty() { + let _ = sender.blocking_send(AddItem::Spilled(ids, vectors)); + } + Ok(rows) + }) +} + +fn pick_indices(total: usize, count: usize) -> impl Iterator { + (0..count).map(move |index| ((2 * index + 1) * total) / (2 * count)) +} + +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], training_rows: usize) -> HashSet { + let target = training_rows + .div_ceil(ROWS_PER_STRATUM) + .max(MIN_STRATA) + .min(granules.len()); + let mut selected = pick_indices(granules.len(), target).collect::>(); + let mut rows = selected + .iter() + .map(|index| granules[*index].range.count() as usize) + .sum::(); + if rows < training_rows { + for index in 0..granules.len() { + if selected.insert(index) { + rows += granules[index].range.count() as usize; + if rows >= training_rows { + break; + } + } + } + } + let first_bytes = granule_bytes(granules, Some(&selected)); + let total_bytes = granule_bytes(granules, None); + if selected.len() < MIN_STRATA.min(granules.len()) + || rows < training_rows + || total_bytes == 0 + || first_bytes.saturating_mul(FIRST_BYTES_DENOMINATOR) + > total_bytes.saturating_mul(FIRST_BYTES_NUMERATOR) + { + (0..granules.len()).collect() + } else { + selected + } +} + +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, + ) -> 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])?; + read_builder.new_read()?.to_arrow(&[split]) + } + + pub(super) async fn plan_granules( + &self, + shard: &VindexIndexShard, + index_column: &str, + training_rows: usize, + ) -> Result { + let shard_range = RowRange::new(shard.row_range_start, shard.row_range_end); + let mut granules = Vec::new(); + let mut page_level = true; + 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 { + page_level = false; + granules.push(Granule { + range, + file_index, + byte_ranges: vec![0..file_size], + }); + } + } + + if !use_whole_shard && !parquet_files.is_empty() { + 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, pages) = + parquet_granules(reader, file_size, index_column).await?; + Ok::<_, Error>((file_index, file_start, file_end, granules, pages)) + }, + ) + .buffer_unordered(concurrency); + while let Some(result) = results.next().await { + match result { + Ok((file_index, file_start, file_end, file_granules, pages)) => { + page_level &= pages; + let covered = file_granules + .iter() + .map(|granule| granule.row_count) + .sum::(); + if covered != file_end - file_start + 1 { + use_whole_shard = true; + break; + } + 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, + }); + } + } + } + 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 { + let coverage = merge_row_ranges( + granules + .iter() + .map(|granule| granule.range.clone()) + .collect(), + ); + use_whole_shard = coverage.len() != 1 || coverage[0] != shard_range; + } + let source = if use_whole_shard || granules.is_empty() { + granules = vec![Granule { + range: shard_range, + file_index: 0, + byte_ranges: vec![ + 0..shard + .files + .iter() + .map(|file| file.file_size.max(0) as u64) + .sum(), + ], + }]; + "shard" + } else if page_level { + "page" + } else { + "row_group_or_file" + }; + + let selected = select_first(&granules, training_rows); + let mut first = Vec::with_capacity(selected.len()); + let mut rest = Vec::with_capacity(granules.len() - selected.len()); + let mut first_rows = 0usize; + let total_bytes = granule_bytes(&granules, None); + 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()); + } + } + let first_bytes = granule_bytes(&granules, Some(&selected)); + Ok(GranulePlan { + first: merge_row_ranges(first), + rest: merge_row_ranges(rest), + first_rows, + first_bytes, + total_bytes, + granule_count: granules.len(), + first_granules: selected.len(), + source, + }) + } + + #[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 total_start = Instant::now(); + 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, 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()?); + let training_rows = eligible.min(plan.first_rows); + 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)?; + let mut expected_row_id = plan.first[0].from(); + while let Some(batch) = stream.try_next().await? { + 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 sample = checked_training_sample_index( + next_training_sample, + plan.first_rows, + training_rows, + )?; + if sample >= batch_end { + break; + } + let offset = (sample - first_rows) * 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 train_start = Instant::now(); + let mut training: Option>> = + Some(tokio::task::spawn_blocking(move || trainer.finish())); + let mut live: Option = None; + let mut train_ms = 0.0; + let mut spill_bytes = 0u64; + let index_column = index_column.to_string(); + + macro_rules! go_live { + () => {{ + let (pipeline, bytes) = start_live_pipeline( + training.take().expect("training task"), + spill.take().expect("spill writer"), + index_column.clone(), + dimension, + ) + .await?; + train_ms = train_start.elapsed().as_secs_f64() * 1000.0; + spill_bytes = bytes; + 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)?; + let mut range_index = 0usize; + let mut expected_row_id = plan.rest[0].from(); + while let Some(batch) = stream.try_next().await? { + 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; + if let Err(consumer_error) = consumer { + return Err(consumer_error); + } + if let Err(replay_error) = replay { + return Err(replay_error); + } + } 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) = finish_live_pipeline(live.unwrap()).await; + let (writer, rows_added, consumer_replay_rows) = consumer?; + let replay_rows = 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 meta = self + .finish_index_file(writer, shard, index_field_id, index_meta, row_count) + .await?; + if vector_index_build_timing_enabled() { + eprintln!( + "event=paimon_vector_index_build_pipeline index_type={} rows={} dimension={} granule_source={} granules={} first_granules={} first_rows={} first_bytes={} total_bytes={} training_rows={} train_ms={:.1} spill_bytes={} total_ms={:.1}", + self.index_type, + row_count_usize, + dimension, + plan.source, + plan.granule_count, + plan.first_granules, + plan.first_rows, + plan.first_bytes, + plan.total_bytes, + training_rows, + train_ms, + spill_bytes, + total_start.elapsed().as_secs_f64() * 1000.0, + ); + } + Ok(BuiltIndexFile { meta, timing: None }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stratified_first_and_rest_cover_every_granule_once() { + let granules = (0..1_000) + .map(|row| Granule { + range: RowRange::new(row, row), + file_index: row as usize, + byte_ranges: vec![0..1], + }) + .collect::>(); + let selected = select_first(&granules, 256); + assert_eq!(selected.len(), 256); + assert!(selected.contains(&1)); + assert!(selected.contains(&998)); + } + + #[test] + fn coalesced_page_holes_force_all_first() { + let granules = (0..1_000) + .map(|row| Granule { + range: RowRange::new(row, row), + file_index: 0, + byte_ranges: vec![row as u64 * 2..row as u64 * 2 + 1], + }) + .collect::>(); + + assert_eq!(select_first(&granules, 256).len(), granules.len()); + } + + #[test] + fn coarse_granules_read_all_first() { + let granules = (0..16) + .map(|row| Granule { + range: RowRange::new(row, row), + file_index: 0, + byte_ranges: vec![row as u64..row as u64 + 1], + }) + .collect::>(); + assert_eq!(select_first(&granules, 8).len(), granules.len()); + } +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs index f4d35e126..3b4a1fb84 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/planning.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -17,14 +17,8 @@ use crate::spec::{CoreOptions, DataField, ManifestEntry}; use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; -use crate::table::{merge_row_ranges, RowRange}; -use crate::{Error, Result}; - -use super::validation::checked_row_count; - -// Keep samples short enough to avoid storage-order bias; fall back before range I/O explodes. -const MAX_IVF_TRAINING_RANGE_ROWS: usize = 128; -const MAX_IVF_TRAINING_RANGES: usize = 4_096; +use crate::table::RowRange; +use crate::Result; pub(crate) type VindexIndexShard = VectorIndexShard; @@ -51,78 +45,3 @@ pub(super) fn plan_vindex_shards( "vindex", ) } - -pub(super) fn plan_ivf_training_ranges( - shard: &VindexIndexShard, - training_rows: usize, -) -> Result>> { - let shard_rows = usize::try_from(checked_row_count( - shard.row_range_start, - shard.row_range_end, - )?) - .map_err(|error| Error::DataInvalid { - message: "vindex shard row count does not fit usize".to_string(), - source: Some(Box::new(error)), - })?; - if training_rows == 0 || training_rows > shard_rows { - return Err(Error::DataInvalid { - message: format!( - "Invalid IVF training row count: {training_rows}; shard contains {shard_rows} rows" - ), - source: None, - }); - } - let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS); - if range_count == 1 || range_count > MAX_IVF_TRAINING_RANGES { - return Ok(None); - } - let seed = ivf_training_seed(shard); - let mut cursor = shard.row_range_start; - let mut ranges = Vec::with_capacity(range_count); - - for range_index in 0..range_count { - let stratum_length = - shard_rows / range_count + usize::from(range_index < shard_rows % range_count); - let length = - training_rows / range_count + usize::from(range_index < training_rows % range_count); - debug_assert!(length <= stratum_length); - let available_offsets = stratum_length - length + 1; - let offset = mix_seed(seed ^ range_index as u64) as usize % available_offsets; - let start = checked_add_offset(cursor, offset, "training range")?; - let end = checked_add_offset(start, length - 1, "training range")?; - ranges.push(RowRange::new(start, end)); - cursor = checked_add_offset(cursor, stratum_length, "training stratum")?; - } - - Ok(Some(merge_row_ranges(ranges))) -} - -pub(super) fn ivf_training_seed(shard: &VindexIndexShard) -> u64 { - let mut seed = mix_seed( - (shard.snapshot_id as u64) - ^ (shard.row_range_start as u64).rotate_left(21) - ^ (shard.row_range_end as u64).rotate_left(42) - ^ (shard.source_bucket as u64).rotate_left(11), - ); - for byte in &shard.partition_bytes { - seed = mix_seed(seed ^ u64::from(*byte)); - } - seed -} - -fn checked_add_offset(value: i64, offset: usize, name: &str) -> Result { - let offset = i64::try_from(offset).map_err(|error| Error::DataInvalid { - message: format!("vindex {name} offset does not fit i64"), - source: Some(Box::new(error)), - })?; - value.checked_add(offset).ok_or_else(|| Error::DataInvalid { - message: format!("vindex {name} offset overflows i64"), - source: None, - }) -} - -fn mix_seed(mut value: u64) -> u64 { - value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); - value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); - value ^ (value >> 31) -} diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs index 20c7d6d13..89447ff79 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -16,7 +16,7 @@ // under the License. use super::extraction::{validate_vector_batch, validate_vector_batch_ranges}; -use super::planning::{plan_ivf_training_ranges, plan_vindex_shards, VindexIndexShard}; +use super::planning::{plan_vindex_shards, VindexIndexShard}; use super::validation::{ checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, find_index_field, validate_vector_field, @@ -143,130 +143,6 @@ fn test_planner_splits_single_file_across_shards() { ); } -#[test] -fn test_ivf_training_ranges_are_bounded_and_exact() { - let shard = plan( - vec![manifest_entry(data_file("a", Some(100), 1_000))], - 1_000, - ) - .unwrap() - .remove(0); - - for training_rows in [129, 200, 899] { - let ranges = plan_ivf_training_ranges(&shard, training_rows) - .unwrap() - .expect("sparse ranges expected"); - - assert!(ranges.len() <= 64); - assert_eq!( - ranges.iter().map(RowRange::count).sum::(), - training_rows as i64 - ); - assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); - assert!(ranges.first().unwrap().from() >= shard.row_range_start); - assert!(ranges.last().unwrap().to() <= shard.row_range_end); - assert_eq!( - ranges, - plan_ivf_training_ranges(&shard, training_rows) - .unwrap() - .expect("sparse ranges expected") - ); - } - - let ranges = plan_ivf_training_ranges(&shard, 200) - .unwrap() - .expect("sparse ranges expected"); - let mut other_snapshot = shard.clone(); - other_snapshot.snapshot_id += 1; - assert_ne!( - ranges, - plan_ivf_training_ranges(&other_snapshot, 200) - .unwrap() - .expect("sparse ranges expected") - ); -} - -#[test] -fn test_ivf_training_ranges_fall_back_for_a_single_range_sample() { - let shard = plan( - vec![manifest_entry(data_file("a", Some(100), 1_000))], - 1_000, - ) - .unwrap() - .remove(0); - - for rows in [1, 127, 128] { - assert!(plan_ivf_training_ranges(&shard, rows).unwrap().is_none()); - } - let ranges = plan_ivf_training_ranges(&shard, 129).unwrap().unwrap(); - assert_eq!(ranges.iter().map(RowRange::count).sum::(), 129); - assert!(ranges - .iter() - .all(|range| range.from() >= 100 && range.to() < 1_100)); - assert!(ranges.windows(2).all(|pair| pair[0].to() < pair[1].from())); -} - -#[test] -fn test_ivf_training_ranges_keep_segments_short() { - let shard = plan( - vec![manifest_entry(data_file("a", Some(0), 1_000_000))], - 1_000_000, - ) - .unwrap() - .remove(0); - - let ranges = plan_ivf_training_ranges(&shard, 65_536) - .unwrap() - .expect("sparse ranges expected"); - - assert_eq!(ranges.len(), 512); - assert!(ranges.iter().all(|range| range.count() <= 128)); - assert_eq!(ranges.iter().map(RowRange::count).sum::(), 65_536); -} - -#[test] -fn test_ivf_training_ranges_scale_to_keep_segments_short() { - let shard = plan( - vec![manifest_entry(data_file("a", Some(0), 2_000_000))], - 2_000_000, - ) - .unwrap() - .remove(0); - - let ranges = plan_ivf_training_ranges(&shard, 262_144) - .unwrap() - .expect("sparse ranges expected"); - - assert_eq!(ranges.len(), 2_048); - assert!(ranges.iter().all(|range| range.count() <= 128)); - assert_eq!(ranges.iter().map(RowRange::count).sum::(), 262_144); -} - -#[test] -fn test_ivf_training_ranges_fall_back_above_range_limit() { - let shard = plan( - vec![manifest_entry(data_file("a", Some(0), 2_000_000))], - 2_000_000, - ) - .unwrap() - .remove(0); - - assert!(plan_ivf_training_ranges(&shard, 524_289).unwrap().is_none()); -} - -#[test] -fn test_ivf_training_ranges_cover_full_shard_without_empty_sentinel() { - let shard = plan(vec![manifest_entry(data_file("a", Some(0), 1_000))], 1_000) - .unwrap() - .remove(0); - - let ranges = plan_ivf_training_ranges(&shard, 1_000) - .unwrap() - .expect("full range expected"); - - assert_eq!(ranges, vec![RowRange::new(0, 999)]); -} - #[test] fn test_planner_rejects_missing_first_row_id() { let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) @@ -377,7 +253,7 @@ fn test_extract_vectors_accepts_list_float32_and_row_ids() { } #[test] -fn test_sparse_vector_validation_accepts_gaps_across_batches() { +fn test_ranged_vector_validation_accepts_gaps_across_batches() { let ranges = vec![RowRange::new(10, 11), RowRange::new(15, 16)]; let batches = [ vector_batch( @@ -415,7 +291,7 @@ fn test_sparse_vector_validation_accepts_gaps_across_batches() { } #[test] -fn test_sparse_vector_validation_rejects_bad_row_ids() { +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)], @@ -794,31 +670,10 @@ async fn vindex_incremental_build_indexes_only_new_rows() { ) .unwrap(); assert_eq!(shards.len(), 1); - let options = crate::vindex::VindexVectorIndexOptions::new( - table.schema().options(), - &build_options, - IVF_FLAT_IDENTIFIER, - find_index_field(&table, "embedding").unwrap(), - ) - .unwrap(); - let training_rows = paimon_vindex_core::autotune::default_training_vector_count( - checked_training_vector_count( - (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, - options.train_sample_ratio, - ) - .unwrap(), - options.config.nlist(), - ) - .unwrap(); let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); builder .with_index_column("embedding") .with_options(build_options); - assert!(builder - .sparse_training_ranges(&shards[0], "embedding", training_rows) - .await - .unwrap() - .is_some()); let first_built = builder.execute().await.unwrap(); assert!(first_built > 0, "first build must index the initial rows"); @@ -868,29 +723,8 @@ async fn vindex_incremental_build_indexes_only_new_rows() { ) .unwrap(); assert_eq!(shards.len(), 1); - let options = crate::vindex::VindexVectorIndexOptions::new( - table.schema().options(), - &HashMap::new(), - IVF_FLAT_IDENTIFIER, - find_index_field(&table, "embedding").unwrap(), - ) - .unwrap(); - let training_rows = paimon_vindex_core::autotune::default_training_vector_count( - checked_training_vector_count( - (shards[0].row_range_end - shards[0].row_range_start + 1) as usize, - options.train_sample_ratio, - ) - .unwrap(), - options.config.nlist(), - ) - .unwrap(); let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); builder.with_index_column("embedding"); - assert!(builder - .sparse_training_ranges(&shards[0], "embedding", training_rows) - .await - .unwrap() - .is_none()); // End-to-end: build #2 must SUCCEED and index the appended rows. let second_built = builder.execute().await.unwrap(); assert!(second_built > 0, "appended rows must be indexed"); @@ -928,48 +762,6 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } -#[tokio::test] -async fn vindex_sparse_probe_errors_fall_back() { - let table_path = "memory:/test_vindex_probe_fallback"; - let table = vindex_e2e_table(table_path, "1024"); - let mut shard = plan( - vec![manifest_entry(data_file("broken.parquet", Some(0), 1024))], - 1024, - ) - .unwrap() - .remove(0); - shard.bucket_path = table_path.to_string(); - let builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); - assert!(plan_ivf_training_ranges(&shard, 192).unwrap().is_some()); - - // A failed open and an unreadable footer both disable only the optimization. - assert!(builder - .sparse_training_ranges(&shard, "embedding", 192) - .await - .unwrap() - .is_none()); - let path = shard.files[0].data_file_path(&shard.bucket_path); - table - .file_io() - .new_output(&path) - .unwrap() - .write(vec![0; 128].into()) - .await - .unwrap(); - assert!(builder - .sparse_training_ranges(&shard, "embedding", 192) - .await - .unwrap() - .is_none()); - - // Invalid source metadata is not an optional probe failure. - shard.files[0].file_size = -1; - assert!(builder - .sparse_training_ranges(&shard, "embedding", 192) - .await - .is_err()); -} - #[tokio::test] async fn vindex_small_training_sample_preserves_tail_cluster_recall() { let table_path = "memory:/test_vindex_small_sample_recall"; 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 1dcfb705b..b83c9f6dd 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -15,162 +15,67 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::{ - data_split_for_shard, data_split_for_shard_ranges, validate_vector_batch, - validate_vector_batch_ranges, -}; -use super::planning::{ivf_training_seed, plan_ivf_training_ranges, VindexIndexShard}; -use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; +use super::extraction::{data_split_for_shard, validate_vector_batch}; +use super::planning::VindexIndexShard; +use super::timing::VectorIndexBuildTiming; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, }; use super::VindexIndexBuildBuilder; -use crate::arrow::format::parquet::has_beneficial_offset_index; use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME}; -use crate::table::data_file_reader::DataFileReadTiming; -use crate::table::table_read::configured_parquet_read_budget; -use crate::table::RowRange; use crate::vindex::{VindexVectorIndexOptions, DISKANN_IDENTIFIER}; use crate::{Error, Result}; use arrow_buffer::MutableBuffer; -use futures::{StreamExt, TryStreamExt}; -use paimon_vindex_core::autotune::default_training_vector_count; +use futures::TryStreamExt; use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::PosWriter; use std::io::{Read, Seek, SeekFrom}; -use std::sync::Arc; -use std::time::{Duration, Instant}; use tokio::io::AsyncWriteExt; use tokio_util::io::SyncIoBridge; const INDEX_DIR: &str = "index"; const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024; -const READ_ADD_QUEUE_CAPACITY: usize = 2; + pub(super) struct BuiltIndexFile { pub(super) meta: IndexFileMeta, pub(super) timing: Option, } impl<'a> VindexIndexBuildBuilder<'a> { - /// Plans sparse training reads, falling back when the optional file probe fails. - pub(super) async fn sparse_training_ranges( + pub(super) async fn build_index_file( &self, shard: &VindexIndexShard, index_column: &str, - training_rows_retained: usize, - ) -> Result>> { - let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - let row_count_usize = usize::try_from(row_count).map_err(|e| Error::DataInvalid { - message: format!("Invalid vindex row count: {row_count}"), - source: Some(Box::new(e)), - })?; - let mut sparse_ranges = - if training_rows_retained > 0 && training_rows_retained < row_count_usize { - plan_ivf_training_ranges(shard, training_rows_retained)? - } else { - None - }; - if let Some(ranges) = sparse_ranges.as_ref() { - let mut checks = Vec::new(); - let mut usable = true; - for file in &shard.files { - if file - .write_cols - .as_ref() - .is_some_and(|columns| !columns.iter().any(|column| column == index_column)) - { - continue; - } - let Some((file_start, file_end)) = file.row_id_range() else { - usable = false; - break; - }; - let local_ranges = ranges - .iter() - .filter_map(|range| { - let from = range.from().max(file_start); - let to = range.to().min(file_end); - (from <= to).then(|| RowRange::new(from - file_start, to - file_start)) - }) - .collect::>(); - if local_ranges.is_empty() { - continue; - } - let scan_ranges = vec![RowRange::new( - shard.row_range_start.max(file_start) - file_start, - shard.row_range_end.min(file_end) - file_start, - )]; - let path = file.data_file_path(&shard.bucket_path); - if !path.to_ascii_lowercase().ends_with(".parquet") { - usable = false; - break; - } - let file_size = u64::try_from(file.file_size).map_err(|e| Error::DataInvalid { - message: format!( - "Invalid data file size for '{}': {}", - file.file_name, file.file_size - ), - source: Some(Box::new(e)), - })?; - checks.push((path, file_size, local_ranges, scan_ranges)); - } - let found_vector_file = !checks.is_empty(); - if usable && found_vector_file { - let concurrency = self - .table - .schema() - .core_options() - .parquet_row_group_parallelism()?; - let file_io = self.table.file_io(); - let mut checks = futures::stream::iter(checks) - .map(|(path, file_size, local_ranges, scan_ranges)| async move { - let input = file_io.new_input(&path)?; - let reader = Box::new(input.reader().await?); - has_beneficial_offset_index( - reader, - file_size, - index_column, - &local_ranges, - &scan_ranges, - ) - .await - }) - .buffer_unordered(concurrency); - while let Some(result) = checks.next().await { - match result { - Ok(true) => {} - Ok(false) => { - usable = false; - break; - } - Err(error) => { - log::warn!( - "vindex sparse training probe failed for column '{}' in shard [{}, {}]; falling back to a full scan: {}", - index_column, - shard.row_range_start, - shard.row_range_end, - error, - ); - return Ok(None); - } - } - } - } - if !usable || !found_vector_file { - log::warn!( - "vindex sparse training read is unavailable for column '{}' in shard [{}, {}]; falling back to a full scan (beneficial_offset_indexes={usable}, found_vector_file={found_vector_file})", + dimension: i32, + index_field_id: i32, + options: &VindexVectorIndexOptions, + index_meta: Vec, + ) -> Result { + if self.index_type != DISKANN_IDENTIFIER { + return self + .build_index_file_granule( + shard, index_column, - shard.row_range_start, - shard.row_range_end, - ); - sparse_ranges = None; - } + dimension, + index_field_id, + options, + index_meta, + ) + .await; } - Ok(sparse_ranges) + self.build_diskann_index_file( + shard, + index_column, + dimension, + index_field_id, + options, + index_meta, + ) + .await } - pub(super) async fn build_index_file( + async fn build_diskann_index_file( &self, shard: &VindexIndexShard, index_column: &str, @@ -179,493 +84,221 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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 raw_temp_write = Duration::ZERO; - 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 mut batch_count = 0usize; 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 = usize::try_from(dimension).map_err(|e| Error::DataInvalid { + let dimension = usize::try_from(dimension).map_err(|e| Error::DataInvalid { message: format!("Invalid vindex dimension: {dimension}"), source: Some(Box::new(e)), })?; - if dimension_usize == 0 { + if dimension == 0 { return Err(Error::DataInvalid { message: "vindex vector dimension must be positive".to_string(), source: None, }); } - let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; + let expected_bytes = checked_vector_bytes(row_count_usize, dimension)?; let training_vector_count = checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER { - 0 - } else { - // This only gates sparse reads; errors must preserve the full-scan fallback. - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) - }; - let capability_start = timing_enabled.then(Instant::now); - let sparse_ranges = self - .sparse_training_ranges(shard, index_column, training_rows_retained) - .await?; - let capability_check = capability_start.map_or(Duration::ZERO, |start| start.elapsed()); - if let Some(budget) = parquet_read_budget.as_ref() { - eprintln!( - "event=paimon_vector_index_build_plan snapshot_id={} row_range_start={} row_range_end={} source_bucket={} sparse={} training_seed={} training_range_count={} parquet_row_group_parallelism={} parquet_max_inflight_bytes={}", - shard.snapshot_id, - shard.row_range_start, - shard.row_range_end, - shard.source_bucket, - sparse_ranges.is_some(), - ivf_training_seed(shard), - sparse_ranges.as_ref().map_or(0, Vec::len), - budget.parallelism(), - budget.max_inflight_bytes(), - ); - } + let buffer_rows = (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension)?).max(1); + let buffer_floats = + buffer_rows + .checked_mul(dimension) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training buffer length overflows usize".to_string(), + source: None, + })?; 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 sample_read = Duration::ZERO; - let mut full_scan_add = Duration::ZERO; - let raw_temp_reread; - let index_add; - let train_finish; + let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { + message: format!("Failed to create temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let mut raw_file = tokio::fs::File::from_std(raw_file); + let split = data_split_for_shard(shard)?; + let mut read_builder = self.table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let mut batches = read_builder.new_read()?.to_arrow(&[split])?; + let mut expected_row_id = shard.row_range_start; + let mut rows_seen = 0usize; let mut bytes_written = 0usize; - let training_rows_seen; + let mut next_training_sample = 0usize; + let mut training_buffer = Vec::with_capacity(buffer_floats); - let writer = if let Some(ranges) = sparse_ranges { - let sample_start = timing_enabled.then(Instant::now); - let split = data_split_for_shard_ranges(shard, ranges.clone())?; - let mut read_builder = self.table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let read = match read_timing.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let mut range_index = 0usize; - let mut expected_row_id = ranges[0].from(); - let mut rows_seen = 0usize; - while let Some(batch) = batches.try_next().await? { - let vectors = validate_vector_batch_ranges( - &batch, - index_column, - dimension_usize, - &ranges, - &mut range_index, - &mut expected_row_id, - )?; - rows_seen = - rows_seen - .checked_add(vectors.row_count) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training row count overflows usize".to_string(), - source: None, - })?; + while let Some(batch) = batches.try_next().await? { + let vectors = + validate_vector_batch(&batch, index_column, dimension, &mut expected_row_id)?; + let batch_end = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + })?; + if training_vector_count == row_count_usize { trainer .add_training_vectors_mut(vectors.values, vectors.row_count) .map_err(|e| Error::DataInvalid { message: format!("Failed to add vindex training vectors: {e}"), source: Some(Box::new(e)), })?; - } - if rows_seen != training_rows_retained || range_index != ranges.len() { - return Err(Error::DataInvalid { - message: format!( - "vindex sparse training data mismatch: rows={rows_seen}/{training_rows_retained}, ranges={range_index}/{}", - ranges.len() - ), - source: None, - }); - } - training_rows_seen = rows_seen; - sample_read = sample_start.map_or(Duration::ZERO, |start| start.elapsed()); - - let train_start = timing_enabled.then(Instant::now); - let training = tokio::task::spawn_blocking(move || trainer.finish()) - .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)), - })?; - train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let writer = VectorIndexWriter::new(training); - - let split = data_split_for_shard(shard)?; - let mut read_builder = self.table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let read = match read_timing.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let expected_end = - shard - .row_range_end - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range end overflows i64".to_string(), - source: None, - })?; - let index_column = index_column.to_string(); - let row_range_start = shard.row_range_start; - let (sender, mut receiver) = - tokio::sync::mpsc::channel::(READ_ADD_QUEUE_CAPACITY); - let full_scan_start = timing_enabled.then(Instant::now); - let consumer = tokio::task::spawn_blocking(move || -> Result<_> { - let mut writer = writer; - let mut expected_row_id = row_range_start; - let mut rows_added = 0usize; - let mut batches_added = 0usize; - let mut ids = Vec::new(); - while let Some(batch) = receiver.blocking_recv() { - let vectors = validate_vector_batch( - &batch, - &index_column, - dimension_usize, - &mut expected_row_id, + } else { + while next_training_sample < training_vector_count { + let sample = checked_training_sample_index( + next_training_sample, + row_count_usize, + training_vector_count, )?; - let batch_end = rows_added.checked_add(vectors.row_count).ok_or_else(|| { - Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), - source: None, - } - })?; - ids.clear(); - for row in rows_added..batch_end { - ids.push(i64::try_from(row).map_err(|e| Error::DataInvalid { - message: "vindex row id does not fit i64".to_string(), - source: Some(Box::new(e)), - })?); - } - writer - .add_vectors(&ids, vectors.values, vectors.row_count) - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to add vectors to vindex index: {e}"), - source: Some(Box::new(e)), - })?; - rows_added = batch_end; - batches_added += 1; - } - Ok((writer, rows_added, expected_row_id, batches_added)) - }); - - let mut producer_error = None; - loop { - let source_start = timing_enabled.then(Instant::now); - let batch = batches.try_next().await; - if let Some(start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(start.elapsed()); - } - let batch = match batch { - Ok(Some(batch)) => batch, - Ok(None) => break, - Err(error) => { - producer_error = Some(error); + if sample >= batch_end { break; } - }; - if sender.send(batch).await.is_err() { - break; - } - } - drop(sender); - let consumer_result = consumer.await; - let (writer, rows_added, next_row_id, batches_added) = - consumer_result.map_err(|e| Error::UnexpectedError { - message: format!("vindex add task failed: {e}"), - source: None, - })??; - if let Some(error) = producer_error { - return Err(error); - } - if rows_added != row_count_usize || next_row_id != expected_end { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_added}/{row_count_usize}, next_row_id={next_row_id}/{expected_end}" - ), - source: None, - }); - } - batch_count = batches_added; - full_scan_add = full_scan_start.map_or(Duration::ZERO, |start| start.elapsed()); - raw_temp_reread = Duration::ZERO; - index_add = Duration::ZERO; - writer - } else { - let training_buffer_rows = - (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); - let training_buffer_floats = training_buffer_rows - .checked_mul(dimension_usize) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; - let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { - message: format!("Failed to create temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - let mut raw_file = tokio::fs::File::from_std(raw_file); - let split = data_split_for_shard(shard)?; - let mut read_builder = self.table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let read = match read_timing.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let mut expected_row_id = shard.row_range_start; - let mut rows_seen = 0usize; - let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(training_buffer_floats); - - loop { - let source_start = timing_enabled.then(Instant::now); - let batch = batches.try_next().await?; - if let Some(source_start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); - } - let Some(batch) = batch else { break }; - batch_count += 1; - let vectors = validate_vector_batch( - &batch, - index_column, - dimension_usize, - &mut expected_row_id, - )?; - let batch_end = - rows_seen - .checked_add(vectors.row_count) - .ok_or_else(|| Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), - source: None, - })?; - - if training_vector_count == row_count_usize { - trainer - .add_training_vectors_mut(vectors.values, vectors.row_count) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } else { - while next_training_sample < training_vector_count { - let sample_row = checked_training_sample_index( - next_training_sample, - row_count_usize, - training_vector_count, - )?; - if sample_row >= batch_end { - break; - } - let start = (sample_row - rows_seen) * dimension_usize; - training_buffer - .extend_from_slice(&vectors.values[start..start + dimension_usize]); - next_training_sample += 1; - if training_buffer.len() == training_buffer_floats { - trainer - .add_training_vectors_mut( - &training_buffer, - training_buffer.len() / dimension_usize, - ) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - training_buffer.clear(); - } + let offset = (sample - rows_seen) * dimension; + training_buffer.extend_from_slice(&vectors.values[offset..offset + dimension]); + next_training_sample += 1; + if training_buffer.len() == 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 raw_write_start = timing_enabled.then(Instant::now); - raw_file - .write_all(vectors.bytes) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to spill vindex vectors: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - bytes_written = - bytes_written - .checked_add(vectors.bytes.len()) - .ok_or_else(|| Error::DataInvalid { - message: "vindex spilled byte count overflows usize".to_string(), - source: None, - })?; - rows_seen = batch_end; } - - if !training_buffer.is_empty() { - trainer - .add_training_vectors_mut( - &training_buffer, - training_buffer.len() / dimension_usize, - ) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } - let expected_end = - shard - .row_range_end - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range end overflows i64".to_string(), - source: None, - })?; - if rows_seen != row_count_usize - || expected_row_id != expected_end - || (training_vector_count != row_count_usize - && next_training_sample != training_vector_count) - || bytes_written != expected_bytes - { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" - ), - source: None, - }); - } - training_rows_seen = training_vector_count; - let raw_write_start = timing_enabled.then(Instant::now); - raw_file.flush().await.map_err(|e| Error::UnexpectedError { - message: format!("Failed to flush temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - let raw_file_len = raw_file - .metadata() + raw_file + .write_all(vectors.bytes) .await .map_err(|e| Error::UnexpectedError { - message: format!("Failed to inspect temporary vindex vector file: {e}"), + message: format!("Failed to spill vindex vectors: {e}"), source: Some(Box::new(e)), - })? - .len(); - if raw_file_len != expected_bytes as u64 { - return Err(Error::DataInvalid { - message: format!( - "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" - ), + })?; + bytes_written = bytes_written + .checked_add(vectors.bytes.len()) + .ok_or_else(|| Error::DataInvalid { + message: "vindex spilled byte count overflows usize".to_string(), source: None, - }); - } - let raw_file = raw_file.into_std().await; + })?; + rows_seen = batch_end; + } - let result = tokio::task::spawn_blocking( - move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { - let train_start = timing_enabled.then(Instant::now); - let training = trainer.finish()?; - let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let mut writer = VectorIndexWriter::new(training); - let mut raw_temp_reread = Duration::ZERO; - let mut index_add = Duration::ZERO; - let mut raw_file = raw_file; - let reread_start = timing_enabled.then(Instant::now); - raw_file.seek(SeekFrom::Start(0))?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - let batch_rows = training_buffer_rows.min(row_count_usize); - let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; - let mut buffer = MutableBuffer::new(batch_bytes); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); - let reread_start = timing_enabled.then(Instant::now); - raw_file.read_exact(buffer.as_slice_mut())?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); - } - let add_start = timing_enabled.then(Instant::now); - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - if let Some(start) = add_start { - index_add = index_add.saturating_add(start.elapsed()); - } - rows_added += rows; - } - let mut trailing = [0u8; 1]; - let reread_start = timing_enabled.then(Instant::now); - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - Ok((writer, train_finish, raw_temp_reread, index_add)) - }, - ) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("vindex training task failed: {e}"), + 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)), + })?; + } + let expected_end = + shard + .row_range_end + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range end overflows i64".to_string(), + source: None, + })?; + if rows_seen != row_count_usize + || expected_row_id != expected_end + || (training_vector_count != row_count_usize + && next_training_sample != training_vector_count) + || bytes_written != expected_bytes + { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" + ), source: None, - })? + }); + } + raw_file.flush().await.map_err(|e| Error::UnexpectedError { + message: format!("Failed to flush temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let raw_file_len = raw_file + .metadata() + .await .map_err(|e| Error::UnexpectedError { - message: format!("Failed to train or add vectors to vindex index: {e}"), + message: format!("Failed to inspect temporary vindex vector file: {e}"), source: Some(Box::new(e)), - })?; - train_finish = result.1; - raw_temp_reread = result.2; - index_add = result.3; - result.0 - }; + })? + .len(); + if raw_file_len != expected_bytes as u64 { + return Err(Error::DataInvalid { + message: format!( + "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" + ), + source: None, + }); + } + let raw_file = raw_file.into_std().await; + let writer = tokio::task::spawn_blocking(move || -> std::io::Result<_> { + let training = trainer.finish()?; + let mut writer = VectorIndexWriter::new(training); + let mut raw_file = raw_file; + raw_file.seek(SeekFrom::Start(0))?; + let batch_rows = buffer_rows.min(row_count_usize); + let mut buffer = MutableBuffer::new(checked_std_vector_bytes(batch_rows, dimension)?); + let mut ids = Vec::with_capacity(batch_rows); + let mut rows_added = 0usize; + while rows_added < row_count_usize { + let rows = batch_rows.min(row_count_usize - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension)?, 0); + raw_file.read_exact(buffer.as_slice_mut())?; + ids.clear(); + for row in rows_added..rows_added + rows { + ids.push(i64::try_from(row).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "vindex row id does not fit i64", + ) + })?); + } + writer.add_vectors(&ids, buffer.typed_data::(), rows)?; + rows_added += rows; + } + let mut trailing = [0u8; 1]; + if raw_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "temporary vindex vector file contains trailing bytes", + )); + } + Ok(writer) + }) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex training task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to train or add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; + + let meta = self + .finish_index_file(writer, shard, index_field_id, index_meta, row_count) + .await?; + Ok(BuiltIndexFile { meta, timing: None }) + } - let serialize_upload_start = timing_enabled.then(Instant::now); + 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!( @@ -715,11 +348,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", @@ -735,50 +366,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, - capability_check, - sample_read, - train_finish, - raw_temp_reread, - index_add, - full_scan_add, - serialize_upload, - rows: row_count_usize, - training_rows_seen, - 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 }) + }) } } From 43022ec9a4ece217f4f964d9c0d0b3f7f01bd167 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 17 Sep 2026 14:16:20 +0800 Subject: [PATCH 13/24] fix(vindex): simplify granule build pipeline --- .../vindex_index_build_builder/pipeline.rs | 347 ++++++++++++------ .../table/vindex_index_build_builder/tests.rs | 91 +---- .../vindex_index_build_builder/timing.rs | 28 +- .../vindex_index_build_builder/writer.rs | 250 +++++++++---- 4 files changed, 434 insertions(+), 282 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index b4fc98a14..ce596acbc 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -19,15 +19,19 @@ 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; +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}; +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}; @@ -39,7 +43,8 @@ use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexTraining, VectorI use std::collections::{HashMap, HashSet}; use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}; use std::ops::Range; -use std::time::Instant; +use std::sync::Arc; +use std::time::{Duration, Instant}; use tokio::sync::mpsc; use tokio::task::JoinHandle; @@ -63,11 +68,6 @@ pub(super) struct GranulePlan { pub(super) first: Vec, pub(super) rest: Vec, first_rows: usize, - first_bytes: u64, - total_bytes: u64, - granule_count: usize, - first_granules: usize, - source: &'static str, } struct SpillRecord { @@ -80,8 +80,9 @@ enum AddItem { Spilled(Vec, MutableBuffer), } -type SpillTask = JoinHandle>; -type ConsumerTask = JoinHandle>; +type SpillTask = JoinHandle>; +type ConsumerTask = JoinHandle>; +type TrainingTask = JoinHandle>; type ReplayTask = JoinHandle>; struct SpillWriter { @@ -90,7 +91,7 @@ struct SpillWriter { } impl SpillWriter { - async fn finish(self) -> Result<(std::fs::File, u64)> { + async fn finish(self) -> Result { drop(self.sender); join_spill(self.task).await } @@ -110,7 +111,6 @@ fn spawn_spill_writer() -> Result { let (sender, mut receiver) = mpsc::channel::(QUEUE_CAPACITY); let task = tokio::task::spawn_blocking(move || -> std::io::Result<_> { let mut writer = BufWriter::with_capacity(BUFFER_BYTES, file); - let mut bytes_written = 0u64; while let Some(record) = receiver.blocking_recv() { let count = record.ids.len() as u64; writer.write_all(&count.to_le_bytes())?; @@ -118,16 +118,14 @@ fn spawn_spill_writer() -> Result { writer.write_all(&id.to_le_bytes())?; } writer.write_all(&record.bytes)?; - bytes_written += 8 + count * 8 + record.bytes.len() as u64; } writer.flush()?; - let file = writer.into_inner().map_err(|e| e.into_error())?; - Ok((file, bytes_written)) + writer.into_inner().map_err(|e| e.into_error()) }); Ok(SpillWriter { sender, task }) } -async fn join_spill(task: SpillTask) -> Result<(std::fs::File, u64)> { +async fn join_spill(task: SpillTask) -> Result { task.await .map_err(|e| Error::UnexpectedError { message: format!("vindex spill task failed: {e}"), @@ -144,12 +142,15 @@ fn spawn_add_consumer( mut receiver: mpsc::Receiver, index_column: String, dimension: usize, + timing_enabled: bool, ) -> ConsumerTask { tokio::task::spawn_blocking(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)?; @@ -185,21 +186,22 @@ fn spawn_add_consumer( replay_rows += ids.len(); } } + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } } - Ok((writer, rows_added, replay_rows)) + Ok((writer, rows_added, replay_rows, index_add)) }) } -async fn join_consumer(task: ConsumerTask) -> Result<(VectorIndexWriter, usize, usize)> { +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: JoinHandle>, -) -> Result { +async fn join_training(task: TrainingTask) -> Result<(VectorIndexTraining, Duration)> { task.await .map_err(|e| Error::UnexpectedError { message: format!("vindex training task failed: {e}"), @@ -212,21 +214,23 @@ async fn join_training( } async fn start_live_pipeline( - training: JoinHandle>, + training: TrainingTask, spill: SpillWriter, index_column: String, dimension: usize, -) -> Result<(LivePipeline, u64)> { + timing_enabled: bool, +) -> Result<(LivePipeline, Duration)> { let trained = join_training(training).await; let spilled = spill.finish().await; - let trained = trained?; - let (file, spill_bytes) = spilled?; + let (trained, train_finish) = trained?; + let file = 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); Ok(( @@ -235,13 +239,16 @@ async fn start_live_pipeline( consumer, replay, }, - spill_bytes, + train_finish, )) } async fn finish_live_pipeline( pipeline: LivePipeline, -) -> (Result<(VectorIndexWriter, usize, usize)>, Result) { +) -> ( + Result<(VectorIndexWriter, usize, usize, Duration)>, + Result, +) { let replay = pipeline .replay .await @@ -286,11 +293,9 @@ fn spawn_replay( })?; let mut id_bytes = vec![0u8; id_bytes_len]; reader.read_exact(&mut id_bytes)?; - ids.extend( - id_bytes - .chunks_exact(8) - .map(|bytes| i64::from_le_bytes(bytes.try_into().unwrap())), - ); + 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)) @@ -352,9 +357,9 @@ fn select_first(granules: &[Granule], training_rows: usize) -> HashSet { .map(|index| granules[*index].range.count() as usize) .sum::(); if rows < training_rows { - for index in 0..granules.len() { + for (index, granule) in granules.iter().enumerate() { if selected.insert(index) { - rows += granules[index].range.count() as usize; + rows += granule.range.count() as usize; if rows >= training_rows { break; } @@ -375,6 +380,37 @@ fn select_first(granules: &[Granule], training_rows: usize) -> HashSet { } } +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 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 { @@ -406,11 +442,22 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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])?; - read_builder.new_read()?.to_arrow(&[split]) + 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( @@ -421,7 +468,6 @@ impl<'a> VindexIndexBuildBuilder<'a> { ) -> Result { let shard_range = RowRange::new(shard.row_range_start, shard.row_range_end); let mut granules = Vec::new(); - let mut page_level = true; let mut use_whole_shard = false; let mut parquet_files = Vec::new(); @@ -451,11 +497,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { if path.to_ascii_lowercase().ends_with(".parquet") { parquet_files.push((file_index, path, file_size, file_start, file_end)); } else { - page_level = false; granules.push(Granule { range, file_index, - byte_ranges: vec![0..file_size], + byte_ranges: std::iter::once(0..file_size).collect(), }); } } @@ -473,16 +518,15 @@ impl<'a> VindexIndexBuildBuilder<'a> { |(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, pages) = + let (granules, _) = parquet_granules(reader, file_size, index_column).await?; - Ok::<_, Error>((file_index, file_start, file_end, granules, pages)) + 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, pages)) => { - page_level &= pages; + Ok((file_index, file_start, file_end, file_granules)) => { let covered = file_granules .iter() .map(|granule| granule.row_count) @@ -491,28 +535,13 @@ impl<'a> VindexIndexBuildBuilder<'a> { use_whole_shard = true; break; } - 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, - }); - } - } + append_shard_granules( + &mut granules, + file_index, + file_start, + &shard_range, + file_granules, + )?; } Err(error) => { log::warn!( @@ -535,30 +564,25 @@ impl<'a> VindexIndexBuildBuilder<'a> { ); use_whole_shard = coverage.len() != 1 || coverage[0] != shard_range; } - let source = if use_whole_shard || granules.is_empty() { + if use_whole_shard || granules.is_empty() { granules = vec![Granule { range: shard_range, file_index: 0, - byte_ranges: vec![ + byte_ranges: std::iter::once( 0..shard .files .iter() .map(|file| file.file_size.max(0) as u64) .sum(), - ], + ) + .collect(), }]; - "shard" - } else if page_level { - "page" - } else { - "row_group_or_file" - }; + } let selected = select_first(&granules, training_rows); let mut first = Vec::with_capacity(selected.len()); let mut rest = Vec::with_capacity(granules.len() - selected.len()); let mut first_rows = 0usize; - let total_bytes = granule_bytes(&granules, None); for (index, granule) in granules.iter().enumerate() { if selected.contains(&index) { first_rows += granule.range.count() as usize; @@ -567,16 +591,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { rest.push(granule.range.clone()); } } - let first_bytes = granule_bytes(&granules, Some(&selected)); Ok(GranulePlan { first: merge_row_ranges(first), rest: merge_row_ranges(rest), first_rows, - first_bytes, - total_bytes, - granule_count: granules.len(), - first_granules: selected.len(), - source, }) } @@ -590,7 +608,18 @@ impl<'a> VindexIndexBuildBuilder<'a> { options: &VindexVectorIndexOptions, index_meta: Vec, ) -> Result { - let total_start = Instant::now(); + 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}"), @@ -626,9 +655,22 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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)?; + 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(); - while let Some(batch) = stream.try_next().await? { + 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, @@ -711,25 +753,31 @@ impl<'a> VindexIndexBuildBuilder<'a> { }; } - let train_start = Instant::now(); - let mut training: Option>> = - Some(tokio::task::spawn_blocking(move || trainer.finish())); + 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_ms = 0.0; - let mut spill_bytes = 0u64; + let mut train_finish = Duration::ZERO; let index_column = index_column.to_string(); macro_rules! go_live { () => {{ - let (pipeline, bytes) = start_live_pipeline( + let (pipeline, duration) = start_live_pipeline( training.take().expect("training task"), spill.take().expect("spill writer"), index_column.clone(), dimension, + timing_enabled, ) .await?; - train_ms = train_start.elapsed().as_secs_f64() * 1000.0; - spill_bytes = bytes; + train_finish = duration; live = Some(pipeline); }}; } @@ -737,10 +785,23 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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)?; + 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(); - while let Some(batch) = stream.try_next().await? { + 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!(); } @@ -804,12 +865,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { if let Err(producer_error) = producer_result { if let Some(pipeline) = live.take() { let (consumer, replay) = finish_live_pipeline(pipeline).await; - if let Err(consumer_error) = consumer { - return Err(consumer_error); - } - if let Err(replay_error) = replay { - return Err(replay_error); - } + consumer?; + replay?; } else { let spill_result = match spill.take() { Some(spill) => spill.finish().await.map(|_| ()), @@ -829,7 +886,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { } let (consumer, replay) = finish_live_pipeline(live.unwrap()).await; - let (writer, rows_added, consumer_replay_rows) = consumer?; + let (writer, rows_added, consumer_replay_rows, index_add) = consumer?; let replay_rows = replay?; if rows_added != row_count_usize || replay_rows != consumer_replay_rows { return Err(Error::DataInvalid { @@ -840,28 +897,53 @@ 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?; - if vector_index_build_timing_enabled() { - eprintln!( - "event=paimon_vector_index_build_pipeline index_type={} rows={} dimension={} granule_source={} granules={} first_granules={} first_rows={} first_bytes={} total_bytes={} training_rows={} train_ms={:.1} spill_bytes={} total_ms={:.1}", - self.index_type, - row_count_usize, - dimension, - plan.source, - plan.granule_count, - plan.first_granules, - plan.first_rows, - plan.first_bytes, - plan.total_bytes, - training_rows, - train_ms, - spill_bytes, - total_start.elapsed().as_secs_f64() * 1000.0, - ); - } - Ok(BuiltIndexFile { meta, timing: None }) + 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, + train_finish, + raw_temp_reread: Duration::ZERO, + index_add, + serialize_upload, + rows: row_count_usize, + training_rows_seen: plan.first_rows, + training_rows_retained: training_rows, + batch_count, + raw_temp_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 }) } } @@ -875,7 +957,7 @@ mod tests { .map(|row| Granule { range: RowRange::new(row, row), file_index: row as usize, - byte_ranges: vec![0..1], + byte_ranges: std::iter::once(0..1).collect(), }) .collect::>(); let selected = select_first(&granules, 256); @@ -890,7 +972,7 @@ mod tests { .map(|row| Granule { range: RowRange::new(row, row), file_index: 0, - byte_ranges: vec![row as u64 * 2..row as u64 * 2 + 1], + byte_ranges: std::iter::once(row as u64 * 2..row as u64 * 2 + 1).collect(), }) .collect::>(); @@ -903,9 +985,30 @@ mod tests { .map(|row| Granule { range: RowRange::new(row, row), file_index: 0, - byte_ranges: vec![row as u64..row as u64 + 1], + byte_ranges: std::iter::once(row as u64..row as u64 + 1).collect(), }) .collect::>(); assert_eq!(select_first(&granules, 8).len(), granules.len()); } + + #[test] + fn near_full_shard_in_shared_file_reads_all_first() { + let file_granules = (0..1_200) + .map(|row| ParquetGranule { + first_row: row, + row_count: 1, + 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, 699); + let mut granules = Vec::new(); + + append_shard_granules(&mut granules, 0, 0, &shard_range, file_granules).unwrap(); + + assert_eq!(granules.len(), 300); + assert_eq!(select_first(&granules, 250).len(), granules.len()); + } } 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 89447ff79..c06731b35 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -623,58 +623,26 @@ async fn vindex_second_build_without_new_data_is_noop() { #[tokio::test] async fn vindex_incremental_build_indexes_only_new_rows() { let table_path = "memory:/test_vindex_incremental"; - let mut options = vindex_e2e_options("2048"); - options.insert("ivf-flat.dimension".to_string(), "4096".to_string()); - options.insert("ivf-flat.nlist".to_string(), "3".to_string()); - let table = test_table_with_io( - FileIOBuilder::new("memory").build().unwrap(), - table_path, - vindex_schema_builder(options).build().unwrap(), - ); + let table = vindex_e2e_table(table_path, "10"); setup_dirs(table.file_io(), table_path).await; // Build #1 over the initial batch via a real end-to-end build. write_vectors( &table, - (0..1024).collect(), - (0..1024) - .map(|id| { - (0..4096) - .map(|component| (id * 4096 + component) as f32) - .collect() - }) - .collect(), + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], ) .await; - let build_options = - HashMap::from([("ivf-flat.train.sample-ratio".to_string(), "0.2".to_string())]); - let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); - let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); - let entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) + let first_built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .with_options(HashMap::from([( + "ivf-flat.train.sample-ratio".to_string(), + "0.9".to_string(), + )])) + .execute() .await .unwrap(); - let core_options = CoreOptions::new(table.schema().options()); - let shards = plan_vindex_shards( - table_path, - table.schema().partition_keys(), - table.schema().fields(), - &core_options, - snapshot.id(), - entries, - core_options.global_index_row_count_per_shard().unwrap(), - &[], - ) - .unwrap(); - assert_eq!(shards.len(), 1); - let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); - builder - .with_index_column("embedding") - .with_options(build_options); - let first_built = builder.execute().await.unwrap(); assert!(first_built > 0, "first build must index the initial rows"); // First appended row-id, derived from the data manifest (never hard-coded). @@ -692,41 +660,18 @@ async fn vindex_incremental_build_indexes_only_new_rows() { // Append a second batch (new row-ids [n..]). write_vectors( &table, - vec![1024, 1025, 1026], - (1024..1027) - .map(|id| { - (0..4096) - .map(|component| (id * 4096 + component) as f32) - .collect() - }) - .collect(), + vec![4, 5, 6], + vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], ) .await; - let snapshot = snapshots.get_latest_snapshot().await.unwrap().unwrap(); - let entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) + // End-to-end: build #2 must SUCCEED and index the appended rows. + let second_built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() .await .unwrap(); - let shards = plan_vindex_shards( - table_path, - table.schema().partition_keys(), - table.schema().fields(), - &core_options, - snapshot.id(), - entries, - core_options.global_index_row_count_per_shard().unwrap(), - &indexed_coverage, - ) - .unwrap(); - assert_eq!(shards.len(), 1); - let mut builder = table.new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER); - builder.with_index_column("embedding"); - // End-to-end: build #2 must SUCCEED and index the appended rows. - let second_built = builder.execute().await.unwrap(); assert!(second_built > 0, "appended rows must be indexed"); let all_files = latest_vindex_index_files(&table).await; 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 26f7039dd..3221dc4c1 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/timing.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -41,12 +41,9 @@ pub(super) struct VectorIndexBuildTiming { pub(super) parquet_projected_bytes_total: u64, pub(super) parquet_peak_inflight_row_groups: usize, pub(super) raw_temp_write: Duration, - pub(super) capability_check: Duration, - pub(super) sample_read: Duration, pub(super) train_finish: Duration, pub(super) raw_temp_reread: Duration, pub(super) index_add: Duration, - pub(super) full_scan_add: Duration, pub(super) serialize_upload: Duration, pub(super) rows: usize, pub(super) training_rows_seen: usize, @@ -61,25 +58,17 @@ pub(super) struct VectorIndexBuildTiming { impl VectorIndexBuildTiming { pub(super) fn log(self, index_type: &str, commit: Duration) { let total = self.total_without_commit.saturating_add(commit); - let build = self - .capability_check - .saturating_add(if self.raw_temp_bytes != 0 { - self.source_batch_wait - .saturating_add(self.raw_temp_write) - .saturating_add(self.train_finish) - .saturating_add(self.raw_temp_reread) - .saturating_add(self.index_add) - } else { - self.sample_read - .saturating_add(self.train_finish) - .saturating_add(self.full_scan_add) - }); - let accounted = build + let accounted = self + .source_batch_wait + .saturating_add(self.raw_temp_write) + .saturating_add(self.train_finish) + .saturating_add(self.raw_temp_reread) + .saturating_add(self.index_add) .saturating_add(self.serialize_upload) .saturating_add(commit); let unattributed = total.saturating_sub(accounted); eprintln!( - "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} capability_check_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms={:.3} full_scan_add_ms={:.3} data_file_count={} data_file_read_concurrency=1 total_ms={:.3} unattributed_ms={:.3}", + "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}", index_type, self.file_name, self.rows, @@ -100,14 +89,11 @@ impl VectorIndexBuildTiming { self.parquet_projected_bytes_total, self.parquet_peak_inflight_row_groups, self.raw_temp_write.as_secs_f64() * 1000.0, - self.capability_check.as_secs_f64() * 1000.0, self.train_finish.as_secs_f64() * 1000.0, self.raw_temp_reread.as_secs_f64() * 1000.0, self.index_add.as_secs_f64() * 1000.0, self.serialize_upload.as_secs_f64() * 1000.0, commit.as_secs_f64() * 1000.0, - self.sample_read.as_secs_f64() * 1000.0, - self.full_scan_add.as_secs_f64() * 1000.0, self.data_file_count, total.as_secs_f64() * 1000.0, unattributed.as_secs_f64() * 1000.0, diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs index b83c9f6dd..f85425e5a 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -17,20 +17,25 @@ use super::extraction::{data_split_for_shard, validate_vector_batch}; use super::planning::VindexIndexShard; -use super::timing::VectorIndexBuildTiming; +use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, }; use super::VindexIndexBuildBuilder; use crate::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, DISKANN_IDENTIFIER}; use crate::{Error, Result}; use arrow_buffer::MutableBuffer; use futures::TryStreamExt; +use paimon_vindex_core::autotune::default_training_vector_count; use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::PosWriter; use std::io::{Read, Seek, SeekFrom}; +use std::sync::Arc; +use std::time::{Duration, Instant}; use tokio::io::AsyncWriteExt; use tokio_util::io::SyncIoBridge; @@ -84,32 +89,45 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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 raw_temp_write = Duration::ZERO; + 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 mut batch_count = 0usize; 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 { + let dimension_usize = usize::try_from(dimension).map_err(|e| Error::DataInvalid { message: format!("Invalid vindex dimension: {dimension}"), source: Some(Box::new(e)), })?; - if dimension == 0 { + if dimension_usize == 0 { return Err(Error::DataInvalid { message: "vindex vector dimension must be positive".to_string(), source: None, }); } - let expected_bytes = checked_vector_bytes(row_count_usize, dimension)?; + let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; let training_vector_count = checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let buffer_rows = (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension)?).max(1); - let buffer_floats = - buffer_rows - .checked_mul(dimension) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; + let training_buffer_rows = + (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); + let training_buffer_floats = training_buffer_rows + .checked_mul(dimension_usize) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training buffer length overflows usize".to_string(), + source: None, + })?; let mut trainer = VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { @@ -124,16 +142,32 @@ impl<'a> VindexIndexBuildBuilder<'a> { let split = data_split_for_shard(shard)?; let mut read_builder = self.table.new_read_builder(); read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let mut batches = read_builder.new_read()?.to_arrow(&[split])?; + let read = read_builder.new_read()?; + let read = match read_timing.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; let mut expected_row_id = shard.row_range_start; let mut rows_seen = 0usize; let mut bytes_written = 0usize; let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(buffer_floats); + let mut training_buffer = Vec::with_capacity(training_buffer_floats); - while let Some(batch) = batches.try_next().await? { + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await?; + if let Some(source_start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); + } + let Some(batch) = batch else { break }; + batch_count += 1; let vectors = - validate_vector_batch(&batch, index_column, dimension, &mut expected_row_id)?; + validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; let batch_end = rows_seen .checked_add(vectors.row_count) @@ -141,6 +175,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { message: "vindex streamed row count overflows usize".to_string(), source: None, })?; + if training_vector_count == row_count_usize { trainer .add_training_vectors_mut(vectors.values, vectors.row_count) @@ -150,22 +185,23 @@ impl<'a> VindexIndexBuildBuilder<'a> { })?; } else { while next_training_sample < training_vector_count { - let sample = checked_training_sample_index( + let sample_row = checked_training_sample_index( next_training_sample, row_count_usize, training_vector_count, )?; - if sample >= batch_end { + if sample_row >= batch_end { break; } - let offset = (sample - rows_seen) * dimension; - training_buffer.extend_from_slice(&vectors.values[offset..offset + dimension]); + let start = (sample_row - rows_seen) * dimension_usize; + training_buffer + .extend_from_slice(&vectors.values[start..start + dimension_usize]); next_training_sample += 1; - if training_buffer.len() == buffer_floats { + if training_buffer.len() == training_buffer_floats { trainer .add_training_vectors_mut( &training_buffer, - training_buffer.len() / dimension, + training_buffer.len() / dimension_usize, ) .map_err(|e| Error::DataInvalid { message: format!("Failed to add vindex training vectors: {e}"), @@ -175,6 +211,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { } } } + + let raw_write_start = timing_enabled.then(Instant::now); raw_file .write_all(vectors.bytes) .await @@ -182,6 +220,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { message: format!("Failed to spill vindex vectors: {e}"), source: Some(Box::new(e)), })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } bytes_written = bytes_written .checked_add(vectors.bytes.len()) .ok_or_else(|| Error::DataInvalid { @@ -193,22 +234,21 @@ impl<'a> VindexIndexBuildBuilder<'a> { if !training_buffer.is_empty() { trainer - .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension) + .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension_usize) .map_err(|e| Error::DataInvalid { message: format!("Failed to add vindex training vectors: {e}"), source: Some(Box::new(e)), })?; } - let expected_end = - shard - .row_range_end - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range end overflows i64".to_string(), - source: None, - })?; if rows_seen != row_count_usize - || expected_row_id != expected_end + || expected_row_id + != shard + .row_range_end + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range end overflows i64".to_string(), + source: None, + })? || (training_vector_count != row_count_usize && next_training_sample != training_vector_count) || bytes_written != expected_bytes @@ -220,10 +260,14 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: None, }); } + let raw_write_start = timing_enabled.then(Instant::now); raw_file.flush().await.map_err(|e| Error::UnexpectedError { message: format!("Failed to flush temporary vindex vector file: {e}"), source: Some(Box::new(e)), })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } let raw_file_len = raw_file .metadata() .await @@ -241,40 +285,71 @@ impl<'a> VindexIndexBuildBuilder<'a> { }); } let raw_file = raw_file.into_std().await; - let writer = tokio::task::spawn_blocking(move || -> std::io::Result<_> { - let training = trainer.finish()?; - let mut writer = VectorIndexWriter::new(training); - let mut raw_file = raw_file; - raw_file.seek(SeekFrom::Start(0))?; - let batch_rows = buffer_rows.min(row_count_usize); - let mut buffer = MutableBuffer::new(checked_std_vector_bytes(batch_rows, dimension)?); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension)?, 0); - raw_file.read_exact(buffer.as_slice_mut())?; - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); + // Diagnostics only: never fail the build for a timing log field. + let training_rows_retained = if timing_enabled { + default_training_vector_count(training_vector_count, options.config.nlist()) + .unwrap_or(0) + } else { + 0 + }; + + let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( + move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { + let train_start = timing_enabled.then(Instant::now); + let training = trainer.finish()?; + let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let mut writer = VectorIndexWriter::new(training); + let mut raw_temp_reread = Duration::ZERO; + let mut index_add = Duration::ZERO; + let mut raw_file = raw_file; + let reread_start = timing_enabled.then(Instant::now); + raw_file.seek(SeekFrom::Start(0))?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); } - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - rows_added += rows; - } - let mut trailing = [0u8; 1]; - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - Ok(writer) - }) + let batch_rows = training_buffer_rows.min(row_count_usize); + let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; + let mut buffer = MutableBuffer::new(batch_bytes); + let mut ids = Vec::with_capacity(batch_rows); + let mut rows_added = 0usize; + while rows_added < row_count_usize { + let rows = batch_rows.min(row_count_usize - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let reread_start = timing_enabled.then(Instant::now); + raw_file.read_exact(buffer.as_slice_mut())?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + ids.clear(); + for row in rows_added..rows_added + rows { + ids.push(i64::try_from(row).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "vindex row id does not fit i64", + ) + })?); + } + let add_start = timing_enabled.then(Instant::now); + writer.add_vectors(&ids, buffer.typed_data::(), rows)?; + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + rows_added += rows; + } + let mut trailing = [0u8; 1]; + let reread_start = timing_enabled.then(Instant::now); + if raw_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "temporary vindex vector file contains trailing bytes", + )); + } + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + Ok((writer, train_finish, raw_temp_reread, index_add)) + }, + ) .await .map_err(|e| Error::UnexpectedError { message: format!("vindex training task failed: {e}"), @@ -285,10 +360,53 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: Some(Box::new(e)), })?; + 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?; - Ok(BuiltIndexFile { meta, timing: None }) + 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, + 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: 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( From 07758fb500edb7391e19c543eff68cb7aabde42d Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 17 Sep 2026 16:35:31 +0800 Subject: [PATCH 14/24] test(vindex): use search result row ids --- .../paimon/src/table/vindex_index_build_builder/tests.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 c06731b35..7c24fc68d 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -763,12 +763,11 @@ async fn vindex_small_training_sample_preserves_tail_cluster_recall() { .execute() .await .unwrap(); - assert_eq!(result.iter().map(RowRange::count).sum::(), 10); + 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!( - result - .iter() - .all(|range| range.from() >= 900 && range.to() < 1000), + row_ids.iter().all(|row_id| (900..1000).contains(row_id)), "{result:?}" ); } From 799df506796687d2e44ecf12be05ee559fa8a1bb Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 17 Sep 2026 17:11:57 +0800 Subject: [PATCH 15/24] fix: handle overlapping vindex granule providers --- .../vindex_index_build_builder/pipeline.rs | 51 ++++++++++++++++--- 1 file changed, 44 insertions(+), 7 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index ce596acbc..3dfe5bd07 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -411,6 +411,22 @@ fn append_shard_granules( Ok(()) } +fn granules_partition_shard(granules: &[Granule], shard_range: &RowRange) -> bool { + if granules + .windows(2) + .any(|pair| pair[1].range.from() <= pair[0].range.to()) + { + return false; + } + let coverage = merge_row_ranges( + granules + .iter() + .map(|granule| granule.range.clone()) + .collect(), + ); + coverage.len() == 1 && coverage.first() == Some(shard_range) +} + 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 { @@ -556,13 +572,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { granules.sort_by_key(|granule| granule.range.from()); if !use_whole_shard { - let coverage = merge_row_ranges( - granules - .iter() - .map(|granule| granule.range.clone()) - .collect(), - ); - use_whole_shard = coverage.len() != 1 || coverage[0] != shard_range; + // 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 { @@ -1011,4 +1023,29 @@ mod tests { assert_eq!(granules.len(), 300); assert_eq!(select_first(&granules, 250).len(), granules.len()); } + + #[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)); + } } From 5fd19b7a0122b5d793f2ae39412551c771da6d42 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 17 Sep 2026 17:39:54 +0800 Subject: [PATCH 16/24] fix: report granule spill metrics --- .../vindex_index_build_builder/pipeline.rs | 96 +++++++++++++++---- .../vindex_index_build_builder/timing.rs | 8 +- .../vindex_index_build_builder/writer.rs | 3 + 3 files changed, 89 insertions(+), 18 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 3dfe5bd07..a99ba09f7 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -80,10 +80,10 @@ enum AddItem { Spilled(Vec, MutableBuffer), } -type SpillTask = JoinHandle>; +type SpillTask = JoinHandle>; type ConsumerTask = JoinHandle>; type TrainingTask = JoinHandle>; -type ReplayTask = JoinHandle>; +type ReplayTask = JoinHandle>; struct SpillWriter { sender: mpsc::Sender, @@ -91,7 +91,7 @@ struct SpillWriter { } impl SpillWriter { - async fn finish(self) -> Result { + async fn finish(self) -> Result<(std::fs::File, u64, Duration)> { drop(self.sender); join_spill(self.task).await } @@ -101,9 +101,11 @@ struct LivePipeline { sender: mpsc::Sender, consumer: ConsumerTask, replay: ReplayTask, + spill_bytes: u64, + spill_write: Duration, } -fn spawn_spill_writer() -> Result { +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)), @@ -111,21 +113,36 @@ fn spawn_spill_writer() -> Result { let (sender, mut receiver) = mpsc::channel::(QUEUE_CAPACITY); let task = tokio::task::spawn_blocking(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()?; - writer.into_inner().map_err(|e| e.into_error()) + 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 { +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}"), @@ -223,7 +240,7 @@ async fn start_live_pipeline( let trained = join_training(training).await; let spilled = spill.finish().await; let (trained, train_finish) = trained?; - let file = spilled?; + let (file, spill_bytes, spill_write) = spilled?; let (sender, receiver) = mpsc::channel(QUEUE_CAPACITY); let consumer = spawn_add_consumer( VectorIndexWriter::new(trained), @@ -232,12 +249,14 @@ async fn start_live_pipeline( dimension, timing_enabled, ); - let replay = spawn_replay(file, sender.clone(), dimension); + let replay = spawn_replay(file, sender.clone(), dimension, timing_enabled); Ok(( LivePipeline { sender, consumer, replay, + spill_bytes, + spill_write, }, train_finish, )) @@ -247,7 +266,9 @@ async fn finish_live_pipeline( pipeline: LivePipeline, ) -> ( Result<(VectorIndexWriter, usize, usize, Duration)>, - Result, + Result<(usize, Duration)>, + u64, + Duration, ) { let replay = pipeline .replay @@ -264,24 +285,34 @@ async fn finish_live_pipeline( }); drop(pipeline.sender); let consumer = join_consumer(pipeline.consumer).await; - (consumer, replay) + (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 { + 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..])?; @@ -308,6 +339,9 @@ fn spawn_replay( 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( @@ -315,14 +349,14 @@ fn spawn_replay( std::mem::replace(&mut vectors, MutableBuffer::new(REPLAY_TARGET_BYTES)), ); if sender.blocking_send(item).is_err() { - return Ok(rows); + return Ok((rows, spill_read)); } } } if !ids.is_empty() { let _ = sender.blocking_send(AddItem::Spilled(ids, vectors)); } - Ok(rows) + Ok((rows, spill_read)) }) } @@ -658,7 +692,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { message: format!("Failed to initialize vindex trainer: {e}"), source: Some(Box::new(e)), })?; - let mut spill = Some(spawn_spill_writer()?); + let mut spill = Some(spawn_spill_writer(timing_enabled)?); let training_rows = eligible.min(plan.first_rows); let training_buffer_rows = (BUFFER_BYTES / checked_vector_bytes(1, dimension)?).max(1); let training_buffer_floats = training_buffer_rows * dimension; @@ -876,7 +910,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { if let Err(producer_error) = producer_result { if let Some(pipeline) = live.take() { - let (consumer, replay) = finish_live_pipeline(pipeline).await; + let (consumer, replay, _, _) = finish_live_pipeline(pipeline).await; consumer?; replay?; } else { @@ -897,9 +931,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { go_live!(); } - let (consumer, replay) = finish_live_pipeline(live.unwrap()).await; + 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 = replay?; + let (replay_rows, granule_spill_read) = replay?; if rows_added != row_count_usize || replay_rows != consumer_replay_rows { return Err(Error::DataInvalid { message: format!( @@ -942,8 +977,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, raw_temp_write: Duration::ZERO, + granule_spill_write, train_finish, raw_temp_reread: Duration::ZERO, + granule_spill_read, index_add, serialize_upload, rows: row_count_usize, @@ -951,6 +988,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { training_rows_retained: training_rows, 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(), @@ -1048,4 +1086,28 @@ mod tests { granules.truncate(1); assert!(granules_partition_shard(&granules, &shard_range)); } + + #[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/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 f85425e5a..129ec307c 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -393,8 +393,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, raw_temp_write, + granule_spill_write: Duration::ZERO, train_finish, raw_temp_reread, + granule_spill_read: Duration::ZERO, index_add, serialize_upload, rows: row_count_usize, @@ -402,6 +404,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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(), From fc46acc0c44399888f1423ace18af16752969ef8 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 18 Sep 2026 11:45:04 +0800 Subject: [PATCH 17/24] feat(vindex): add granule build toggle --- .../table/vindex_index_build_builder/tests.rs | 4 + .../vindex_index_build_builder/writer.rs | 12 ++- crates/paimon/src/vindex/mod.rs | 74 +++++++++++++++++++ docs/src/sql.md | 1 + 4 files changed, 88 insertions(+), 3 deletions(-) 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 7c24fc68d..459b6562e 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -669,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(); 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 129ec307c..8c7fc1260 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -57,7 +57,13 @@ impl<'a> VindexIndexBuildBuilder<'a> { options: &VindexVectorIndexOptions, index_meta: Vec, ) -> Result { - if self.index_type != DISKANN_IDENTIFIER { + 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, @@ -69,7 +75,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { ) .await; } - self.build_diskann_index_file( + self.build_full_spill_index_file( shard, index_column, dimension, @@ -80,7 +86,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { .await } - async fn build_diskann_index_file( + async fn build_full_spill_index_file( &self, shard: &VindexIndexShard, index_column: &str, diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs index 8ec93a133..03d85e99f 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 { @@ -228,10 +230,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, }) } @@ -290,6 +294,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; } @@ -307,6 +314,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, @@ -831,6 +857,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"] { @@ -949,6 +1022,7 @@ mod tests { (IVF_FLAT_IDENTIFIER, "ivf-flat.pq.m"), (IVF_FLAT_IDENTIFIER, "diskann.max-degree"), (DISKANN_IDENTIFIER, "diskann.nlist"), + (DISKANN_IDENTIFIER, VINDEX_BUILD_GRANULE_ENABLED), ] { let user_options = HashMap::from([(key.to_string(), "2".to_string())]); let err = VindexVectorIndexOptions::new( diff --git a/docs/src/sql.md b/docs/src/sql.md index f35345dd3..56fdde442 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1245,6 +1245,7 @@ Supported vindex options: | `.distance.metric` | `inner_product` | all vindex types | Distance metric: `inner_product`, `cosine`, or `l2`. | | `.nlist` | `256` | all IVF types | Number of IVF lists. DiskANN rejects this option. | | `.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-rq.rq.bits` | `4`, or inferred | `ivf-rq` | Residual-quantization width in the range `1` to `8`. When omitted, `ivf-rq.max-bytes-per-vector` can select it. | From 407e39630abfe30428327d30bcbd913835558985 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 09:28:21 +0800 Subject: [PATCH 18/24] fix(vindex): jitter granule training samples --- .../vindex_index_build_builder/pipeline.rs | 195 +++++++++++++++--- .../table/vindex_index_build_builder/tests.rs | 78 +++++++ 2 files changed, 248 insertions(+), 25 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index a99ba09f7..a96ba4946 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -360,8 +360,40 @@ fn spawn_replay( }) } -fn pick_indices(total: usize, count: usize) -> impl Iterator { - (0..count).map(move |index| ((2 * index + 1) * total) / (2 * count)) +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 pick_indices(total: usize, count: usize, seed: u64) -> impl Iterator { + let mut state = seed; + (0..count).map(move |index| { + let start = index * total / count; + let end = (index + 1) * total / count; + start + (splitmix64(&mut state) % (end - start) as u64) as usize + }) } fn granule_bytes(granules: &[Granule], selected: Option<&HashSet>) -> u64 { @@ -380,26 +412,25 @@ fn granule_bytes(granules: &[Granule], selected: Option<&HashSet>) -> u64 .sum() } -fn select_first(granules: &[Granule], training_rows: usize) -> HashSet { - let target = training_rows +fn select_first(granules: &[Granule], training_rows: usize, seed: u64) -> HashSet { + let mut target = training_rows .div_ceil(ROWS_PER_STRATUM) .max(MIN_STRATA) .min(granules.len()); - let mut selected = pick_indices(granules.len(), target).collect::>(); - let mut rows = selected - .iter() - .map(|index| granules[*index].range.count() as usize) - .sum::(); - if rows < training_rows { - for (index, granule) in granules.iter().enumerate() { - if selected.insert(index) { - rows += granule.range.count() as usize; - if rows >= training_rows { - break; - } - } + let (selected, rows) = loop { + let selected = pick_indices(granules.len(), target, seed).collect::>(); + let rows = selected + .iter() + .map(|index| granules[*index].range.count() as usize) + .sum::(); + if rows >= training_rows || target == granules.len() { + break (selected, rows); } - } + target = target + .saturating_mul(training_rows) + .div_ceil(rows.max(1)) + .clamp(target + 1, granules.len()); + }; let first_bytes = granule_bytes(granules, Some(&selected)); let total_bytes = granule_bytes(granules, None); if selected.len() < MIN_STRATA.min(granules.len()) @@ -625,7 +656,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { }]; } - let selected = select_first(&granules, training_rows); + let selected = select_first(&granules, training_rows, shard_seed(shard)); let mut first = Vec::with_capacity(selected.len()); let mut rest = Vec::with_capacity(granules.len() - selected.len()); let mut first_rows = 0usize; @@ -1010,10 +1041,124 @@ mod tests { byte_ranges: std::iter::once(0..1).collect(), }) .collect::>(); - let selected = select_first(&granules, 256); + let selected = select_first(&granules, 256, 7); assert_eq!(selected.len(), 256); - assert!(selected.contains(&1)); - assert!(selected.contains(&998)); + for stratum in 0..256 { + let (start, end) = (stratum * 1_000 / 256, (stratum + 1) * 1_000 / 256); + assert_eq!( + (start..end) + .filter(|index| selected.contains(index)) + .count(), + 1, + "stratum {stratum}" + ); + } + assert_eq!(selected, select_first(&granules, 256, 7)); + assert_ne!(selected, select_first(&granules, 256, 8)); + } + + #[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, 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() { + for seed in 0..32u64 { + let picks = pick_indices(4_096, 512, seed).collect::>(); + assert!(picks.windows(2).all(|pair| pair[0] < pair[1])); + for period in [2usize, 4, 8, 16] { + let mut hits = vec![0usize; period]; + for pick in &picks { + hits[pick % period] += 1; + } + let expected = 512 / 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 mut next_row = 0i64; + let granules = (0..4_096usize) + .map(|index| { + let rows = if index % 8 == 4 { 16 } else { 256 }; + let range = RowRange::new(next_row, next_row + rows - 1); + next_row += rows; + Granule { + range, + file_index: index, + byte_ranges: std::iter::once(0..rows as u64).collect(), + } + }) + .collect::>(); + let short_share = 16.0 / (16.0 + 7.0 * 256.0); + for seed in 0..32u64 { + let selected = select_first(&granules, 65_536, seed); + assert!(selected.len() < granules.len() / 4, "seed {seed}"); + let (mut rows, mut short_rows) = (0i64, 0i64); + for index in &selected { + let count = granules[*index].range.count(); + rows += count; + if index % 8 == 4 { + short_rows += count; + } + } + assert!(rows >= 65_536, "seed {seed}: {rows}"); + let share = short_rows as f64 / rows as f64; + assert!( + share > short_share / 2.0 && share < short_share * 2.0, + "seed {seed}: {share}" + ); + } + } + + #[test] + fn short_granules_widen_the_pick_instead_of_filling_from_the_head() { + let granules = (0..4_096i64) + .map(|index| Granule { + range: RowRange::new(index * 64, index * 64 + 63), + file_index: index as usize, + byte_ranges: std::iter::once(0..64).collect(), + }) + .collect::>(); + let selected = select_first(&granules, 32_768, 7); + assert_eq!(selected.len(), 512); + for quarter in 0..4 { + let picked = (quarter * 1_024..(quarter + 1) * 1_024) + .filter(|index| selected.contains(index)) + .count(); + assert_eq!(picked, 128, "quarter {quarter}"); + } } #[test] @@ -1026,7 +1171,7 @@ mod tests { }) .collect::>(); - assert_eq!(select_first(&granules, 256).len(), granules.len()); + assert_eq!(select_first(&granules, 256, 7).len(), granules.len()); } #[test] @@ -1038,7 +1183,7 @@ mod tests { byte_ranges: std::iter::once(row as u64..row as u64 + 1).collect(), }) .collect::>(); - assert_eq!(select_first(&granules, 8).len(), granules.len()); + assert_eq!(select_first(&granules, 8, 7).len(), granules.len()); } #[test] @@ -1059,7 +1204,7 @@ mod tests { append_shard_granules(&mut granules, 0, 0, &shard_range, file_granules).unwrap(); assert_eq!(granules.len(), 300); - assert_eq!(select_first(&granules, 250).len(), granules.len()); + assert_eq!(select_first(&granules, 250, 7).len(), granules.len()); } #[test] 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 459b6562e..4420498d0 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -776,6 +776,84 @@ async fn vindex_small_training_sample_preserves_tail_cluster_recall() { ); } +/// 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_build_cleans_written_shards_when_later_shard_fails() { let table_path = "memory:/test_vindex_abort_written_shard"; From 1a12958101e641a579c18c4e514d442185db2510 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 10:00:39 +0800 Subject: [PATCH 19/24] fix(vindex): adapt granule reads to read budget API --- crates/paimon/src/arrow/format/parquet.rs | 8 ++++---- .../src/table/vindex_index_build_builder/pipeline.rs | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 56e91d1a0..7a4e17f8c 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -3046,7 +3046,7 @@ mod tests { RowRange::new(258, 263), RowRange::new(380, 383), ]; - let budget = Arc::new(ParquetReadBudget::new(2, 256 * 1024 * 1024).unwrap()); + let budget = Arc::new(ReadBudget::new(2, 256 * 1024 * 1024).unwrap()); budget.enable_diagnostics(); let batches = ParquetFormatReader::with_read_budget(Arc::clone(&budget)) @@ -3118,7 +3118,7 @@ mod tests { let projected_bytes = super::projected_row_group_bytes(&row_group, &projection); assert_eq!(projected_bytes, 308 * MIB as u64); - let budget = ParquetReadBudget::new(8, 256 * MIB as u64).unwrap(); + let budget = ReadBudget::new(8, 256 * MIB as u64).unwrap(); budget.enable_diagnostics(); let permit = budget.acquire(projected_bytes).await.unwrap(); assert!( @@ -3225,7 +3225,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; @@ -4001,7 +4001,7 @@ mod tests { .collect::>(); for cancel in [false, true] { let tracker = TrackingFileRead::new(data.clone()); - let budget = Arc::new(ParquetReadBudget::new(8, BUDGET as u64).unwrap()); + 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( diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index a96ba4946..5117a14fd 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -524,7 +524,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { ranges: Vec, index_column: &str, read_timing: Option<&Arc>, - parquet_read_budget: 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(); From a2ab2f8cf2f62a37f134321751b8b4bf1ef2aea4 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 10:53:03 +0800 Subject: [PATCH 20/24] fix(vindex): report actual trainer sample counts --- .../paimon/src/table/vindex_index_build_builder/pipeline.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 5117a14fd..156329b37 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -1015,8 +1015,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { index_add, serialize_upload, rows: row_count_usize, - training_rows_seen: plan.first_rows, - training_rows_retained: training_rows, + training_rows_seen: training_rows, + training_rows_retained: retained, batch_count, raw_temp_bytes: 0, granule_spill_bytes, From 2125a30800f5c79f80d268464e13a0d1097c4a8a Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 12:57:02 +0800 Subject: [PATCH 21/24] fix(vindex): sample granule training in logical row space --- .../vindex_index_build_builder/pipeline.rs | 544 +++++++++++++----- .../table/vindex_index_build_builder/tests.rs | 165 ++++++ 2 files changed, 581 insertions(+), 128 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 156329b37..08b868a5f 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -68,6 +68,8 @@ 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 { @@ -387,12 +389,50 @@ fn shard_seed(shard: &VindexIndexShard) -> u64 { seed } -fn pick_indices(total: usize, count: usize, seed: u64) -> impl Iterator { - let mut state = 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 * total / count; - let end = (index + 1) * total / count; - start + (splitmix64(&mut state) % (end - start) as u64) as usize + 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 }) } @@ -412,37 +452,104 @@ fn granule_bytes(granules: &[Granule], selected: Option<&HashSet>) -> u64 .sum() } -fn select_first(granules: &[Granule], training_rows: usize, seed: u64) -> HashSet { - let mut target = training_rows +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 ends = units + .iter() + .map(|unit| (granules[unit.end - 1].range.to() - range.from() + 1) as usize) + .collect::>(); + let strata = retained .div_ceil(ROWS_PER_STRATUM) .max(MIN_STRATA) - .min(granules.len()); - let (selected, rows) = loop { - let selected = pick_indices(granules.len(), target, seed).collect::>(); - let rows = selected - .iter() - .map(|index| granules[*index].range.count() as usize) - .sum::(); - if rows >= training_rows || target == granules.len() { - break (selected, rows); + .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) { + quotas[ends.partition_point(|end| *end <= row)] += 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()); } - target = target - .saturating_mul(training_rows) - .div_ceil(rows.max(1)) - .clamp(target + 1, granules.len()); - }; + 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() < MIN_STRATA.min(granules.len()) - || rows < training_rows + 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) { - (0..granules.len()).collect() - } else { - selected + 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( @@ -545,7 +652,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { &self, shard: &VindexIndexShard, index_column: &str, - training_rows: usize, + eligible: usize, + retained: usize, ) -> Result { let shard_range = RowRange::new(shard.row_range_start, shard.row_range_end); let mut granules = Vec::new(); @@ -656,23 +764,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { }]; } - let selected = select_first(&granules, training_rows, shard_seed(shard)); - let mut first = Vec::with_capacity(selected.len()); - let mut rest = Vec::with_capacity(granules.len() - selected.len()); - let mut first_rows = 0usize; - 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, - }) + select_first(&granules, eligible, retained, shard_seed(shard)) } #[allow(clippy::too_many_arguments)] @@ -716,7 +808,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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, retained).await?; + let plan = self + .plan_granules(shard, index_column, eligible, retained) + .await?; let mut trainer = VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { @@ -724,7 +818,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: Some(Box::new(e)), })?; let mut spill = Some(spawn_spill_writer(timing_enabled)?); - let training_rows = eligible.min(plan.first_rows); + 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); @@ -763,15 +857,28 @@ impl<'a> VindexIndexBuildBuilder<'a> { } })?; while next_training_sample < training_rows { - let sample = checked_training_sample_index( - next_training_sample, - plan.first_rows, - training_rows, - )?; - if sample >= batch_end { - break; - } - let offset = (sample - first_rows) * dimension; + 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 { @@ -1032,29 +1139,215 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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 = (0..1_000) - .map(|row| Granule { - range: RowRange::new(row, row), - file_index: row as usize, - byte_ranges: std::iter::once(0..1).collect(), - }) - .collect::>(); - let selected = select_first(&granules, 256, 7); - assert_eq!(selected.len(), 256); + 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 * 1_000 / 256, (stratum + 1) * 1_000 / 256); + let (start, end) = (stratum * 4_096, (stratum + 1) * 4_096); assert_eq!( - (start..end) - .filter(|index| selected.contains(index)) + training + .iter() + .filter(|row| start <= **row && **row < end) .count(), - 1, + 128, "stratum {stratum}" ); } - assert_eq!(selected, select_first(&granules, 256, 7)); - assert_ne!(selected, select_first(&granules, 256, 8)); + 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] @@ -1071,7 +1364,7 @@ mod tests { bucket_path: "memory:/t/bucket-0".to_string(), }; let picks = |shard: &VindexIndexShard| { - pick_indices(4_096, 512, shard_seed(shard)).collect::>() + pick_indices(4_096, 512, &mut shard_seed(shard)).collect::>() }; assert_eq!(picks(&shard), picks(&shard.clone())); @@ -1088,15 +1381,18 @@ mod tests { #[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 picks = pick_indices(4_096, 512, seed).collect::>(); - assert!(picks.windows(2).all(|pair| pair[0] < pair[1])); + 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 pick in &picks { - hits[pick % period] += 1; + for row in &training { + hits[*row as usize / 256 % period] += 1; } - let expected = 512 / period; + let expected = training.len() / period; assert!( hits.iter() .all(|count| *count * 4 > expected && *count < expected * 2), @@ -1108,33 +1404,20 @@ mod tests { #[test] fn unequal_granules_are_sampled_in_proportion() { - let mut next_row = 0i64; - let granules = (0..4_096usize) - .map(|index| { - let rows = if index % 8 == 4 { 16 } else { 256 }; - let range = RowRange::new(next_row, next_row + rows - 1); - next_row += rows; - Granule { - range, - file_index: index, - byte_ranges: std::iter::once(0..rows as u64).collect(), - } - }) - .collect::>(); + 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 selected = select_first(&granules, 65_536, seed); - assert!(selected.len() < granules.len() / 4, "seed {seed}"); - let (mut rows, mut short_rows) = (0i64, 0i64); - for index in &selected { - let count = granules[*index].range.count(); - rows += count; - if index % 8 == 4 { - short_rows += count; - } - } - assert!(rows >= 65_536, "seed {seed}: {rows}"); - let share = short_rows as f64 / rows as f64; + 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}" @@ -1143,35 +1426,37 @@ mod tests { } #[test] - fn short_granules_widen_the_pick_instead_of_filling_from_the_head() { - let granules = (0..4_096i64) - .map(|index| Granule { - range: RowRange::new(index * 64, index * 64 + 63), - file_index: index as usize, - byte_ranges: std::iter::once(0..64).collect(), - }) - .collect::>(); - let selected = select_first(&granules, 32_768, 7); - assert_eq!(selected.len(), 512); + 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 = (quarter * 1_024..(quarter + 1) * 1_024) - .filter(|index| selected.contains(index)) + let picked = training + .iter() + .filter(|row| **row / 65_536 == quarter) .count(); - assert_eq!(picked, 128, "quarter {quarter}"); + assert_eq!(picked, 8_192, "quarter {quarter}"); } } #[test] fn coalesced_page_holes_force_all_first() { - let granules = (0..1_000) - .map(|row| Granule { - range: RowRange::new(row, row), - file_index: 0, - byte_ranges: std::iter::once(row as u64 * 2..row as u64 * 2 + 1).collect(), - }) - .collect::>(); - - assert_eq!(select_first(&granules, 256, 7).len(), granules.len()); + 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] @@ -1183,28 +1468,31 @@ mod tests { byte_ranges: std::iter::once(row as u64..row as u64 + 1).collect(), }) .collect::>(); - assert_eq!(select_first(&granules, 8, 7).len(), granules.len()); + 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, - row_count: 1, + 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, 699); + 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_eq!(select_first(&granules, 250, 7).len(), granules.len()); + assert!(select_first(&granules, 32_768, 32_768, 7) + .unwrap() + .training + .is_none()); } #[test] 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 4420498d0..4ab0f759f 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -854,6 +854,171 @@ async fn vindex_granule_training_sees_file_periodic_cluster() { ); } +#[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", + ) + .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"; From de03ff2acb41f3feb5ebfd7325c3cd23d934b6a2 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 13:34:30 +0800 Subject: [PATCH 22/24] refactor(vindex): remove redundant granule planning work --- .../vindex_index_build_builder/pipeline.rs | 70 ++++++++++++++----- 1 file changed, 52 insertions(+), 18 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 08b868a5f..54760c6c9 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -474,10 +474,6 @@ fn select_first( return Ok(whole_shard()); } let units = sampling_units(granules); - let ends = units - .iter() - .map(|unit| (granules[unit.end - 1].range.to() - range.from() + 1) as usize) - .collect::>(); let strata = retained .div_ceil(ROWS_PER_STRATUM) .max(MIN_STRATA) @@ -490,7 +486,9 @@ fn select_first( })?; let mut quotas = vec![0usize; units.len()]; for row in pick_indices(rows, strata, &mut seed) { - quotas[ends.partition_point(|end| *end <= row)] += ROWS_PER_STRATUM; + 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) { @@ -584,19 +582,15 @@ fn append_shard_granules( } fn granules_partition_shard(granules: &[Granule], shard_range: &RowRange) -> bool { - if granules - .windows(2) - .any(|pair| pair[1].range.from() <= pair[0].range.to()) - { - return false; - } - let coverage = merge_row_ranges( - granules - .iter() - .map(|granule| granule.range.clone()) - .collect(), - ); - coverage.len() == 1 && coverage.first() == Some(shard_range) + 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> { @@ -1495,6 +1489,46 @@ mod tests { .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); From d1f86cd9ea82fe9db9fb9d007b712e8095fa39f6 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 14:15:54 +0800 Subject: [PATCH 23/24] fix(vindex): honor Parquet page index option in granule planning --- crates/paimon/src/arrow/format/parquet.rs | 34 ++++++++++++++++++- .../vindex_index_build_builder/pipeline.rs | 8 ++++- .../table/vindex_index_build_builder/tests.rs | 1 + 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 949a6a88b..7b3b125e8 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -106,8 +106,12 @@ pub(crate) async fn parquet_granules( reader: Box, file_size: u64, column_name: &str, + page_index_enabled: bool, ) -> crate::Result<(Vec, bool)> { - let options = ArrowReaderOptions::new().with_offset_index_policy(PageIndexPolicy::Optional); + 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 @@ -3855,6 +3859,7 @@ mod tests { Box::new(TrackingFileRead::new(bytes.clone())), bytes.len() as u64, "value", + true, ) .await .unwrap(); @@ -3878,6 +3883,7 @@ mod tests { Box::new(TrackingFileRead::new(bytes.clone())), bytes.len() as u64, "value", + true, ) .await .unwrap(); @@ -3892,6 +3898,32 @@ mod tests { ); } + #[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![ diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 54760c6c9..f24dd7b12 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -689,6 +689,11 @@ impl<'a> VindexIndexBuildBuilder<'a> { } 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() @@ -702,7 +707,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { let input = file_io.new_input(&path)?; let reader = Box::new(input.reader().await?); let (granules, _) = - parquet_granules(reader, file_size, index_column).await?; + parquet_granules(reader, file_size, index_column, page_index_enabled) + .await?; Ok::<_, Error>((file_index, file_start, file_end, granules)) }, ) 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 4ab0f759f..fb5c14990 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -959,6 +959,7 @@ async fn vindex_granule_training_sees_one_oversized_row_group() { Box::new(input.reader().await.unwrap()), large_file.file_size as u64, "embedding", + true, ) .await .unwrap(); From 8af080f6933e4783a4f64233ff07e8d1830660e9 Mon Sep 17 00:00:00 2001 From: yantian Date: Sun, 20 Sep 2026 19:03:31 +0800 Subject: [PATCH 24/24] fix(vindex): isolate channel workers from blocking pool --- .../vindex_index_build_builder/pipeline.rs | 100 ++++++++++++++++-- 1 file changed, 92 insertions(+), 8 deletions(-) diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index f24dd7b12..ebd2b8b3d 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -45,7 +45,7 @@ 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; +use tokio::sync::{mpsc, oneshot}; use tokio::task::JoinHandle; const MIN_STRATA: usize = 256; @@ -82,11 +82,29 @@ enum AddItem { Spilled(Vec, MutableBuffer), } -type SpillTask = JoinHandle>; -type ConsumerTask = JoinHandle>; +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, @@ -113,7 +131,7 @@ fn spawn_spill_writer(timing_enabled: bool) -> Result { source: Some(Box::new(e)), })?; let (sender, mut receiver) = mpsc::channel::(QUEUE_CAPACITY); - let task = tokio::task::spawn_blocking(move || -> std::io::Result<_> { + 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; @@ -140,7 +158,7 @@ fn spawn_spill_writer(timing_enabled: bool) -> Result { spill_write = spill_write.saturating_add(start.elapsed()); } Ok((file, spill_bytes, spill_write)) - }); + })?; Ok(SpillWriter { sender, task }) } @@ -162,8 +180,8 @@ fn spawn_add_consumer( index_column: String, dimension: usize, timing_enabled: bool, -) -> ConsumerTask { - tokio::task::spawn_blocking(move || -> Result<_> { +) -> Result { + spawn_channel_worker("paimon-vindex-add", move || -> Result<_> { let mut writer = writer; let mut rows_added = 0usize; let mut replay_rows = 0usize; @@ -250,7 +268,7 @@ async fn start_live_pipeline( index_column, dimension, timing_enabled, - ); + )?; let replay = spawn_replay(file, sender.clone(), dimension, timing_enabled); Ok(( LivePipeline { @@ -1560,6 +1578,72 @@ mod tests { 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();