From 2fac200c578aa20948290a8efae6dd66638a90e1 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 9 Sep 2026 17:28:43 +0800 Subject: [PATCH 01/11] feat: optimize IVF index build reads --- crates/paimon/src/arrow/format/parquet.rs | 71 +- .../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, 783 insertions(+), 240 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 7aab8a074..81317aceb 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), @@ -2231,8 +2272,8 @@ fn split_ranges_for_concurrency(merged: Vec>, concurrency: usize) -> mod tests { use super::build_parquet_row_filter; use super::{ - forward_row_group_batches, FilePredicates, ParquetFormatReader, ParquetFormatWriter, - ParquetRowGroupMessage, + forward_row_group_batches, metadata_has_usable_offset_index, FilePredicates, + ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage, }; use super::{ AsyncArrowWriter, Bytes, PageIndexPolicy, ParquetMetaDataReader, Predicate, @@ -2904,7 +2945,7 @@ mod tests { #[tokio::test] async fn test_parquet_diagnostics_include_reads_with_row_selection() { - let data = write_multi_row_group_parquet(32, 64, EnabledStatistics::Chunk).await; + let data = write_multi_row_group_parquet(32, 64, EnabledStatistics::Chunk, false).await; let budget = Arc::new(ParquetReadBudget::new(8, 256 * 1024 * 1024).unwrap()); budget.enable_diagnostics(); let file_size = data.len() as u64; @@ -3367,11 +3408,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(); @@ -3387,7 +3430,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()); @@ -3415,19 +3458,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 52af8c96fdede27b41c98625a8643095e69b7987 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 12:52:30 +0800 Subject: [PATCH 02/11] 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 81317aceb..e13d2f3a7 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 { @@ -3463,14 +3488,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 1257f6328..c5630ef5e 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 a6314a5ec..b44791d91 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 3c071ed55a8fca4370f4713dd8e3dcd557db4684 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 16:01:49 +0800 Subject: [PATCH 03/11] 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 f51aaf1d4c614815172ff0a8853bb5c4965691ce Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 16:43:00 +0800 Subject: [PATCH 04/11] 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 dfe29d90b6e8be2a3cdd43911496f6dbbfddd55c Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 17:36:40 +0800 Subject: [PATCH 05/11] 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 e8fabff64300ba45c10336ea1c9f46827a35dba3 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 10:19:41 +0800 Subject: [PATCH 06/11] perf(vindex): parallelize sparse parquet reads --- crates/paimon/src/arrow/format/parquet.rs | 121 +++++++++++++++--- .../paimon/src/arrow/parquet_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 e13d2f3a7..62607ae54 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -553,12 +553,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 @@ -566,39 +566,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; }; @@ -623,6 +637,7 @@ impl FormatFileReader for ParquetFormatReader { row_group_mask, row_group_index, batch_size, + selection, permit, batch_tx, )); @@ -726,6 +741,7 @@ async fn read_row_group( projection: ProjectionMask, row_group_index: usize, batch_size: Option, + selection: Option, _permit: ParquetReadPermit, sender: mpsc::Sender, ) { @@ -735,6 +751,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); } @@ -2881,6 +2900,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; @@ -2994,7 +3073,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/parquet_read_budget.rs b/crates/paimon/src/arrow/parquet_read_budget.rs index bef9c6ebd..a11111759 100644 --- a/crates/paimon/src/arrow/parquet_read_budget.rs +++ b/crates/paimon/src/arrow/parquet_read_budget.rs @@ -108,6 +108,10 @@ impl ParquetReadBudget { self.parallelism } + pub(crate) fn max_inflight_bytes(&self) -> u64 { + self.max_inflight_bytes + } + pub(crate) fn enable_diagnostics(&self) { self.diagnostics.enabled.store(true, Ordering::Relaxed); } diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index b44791d91..8fdbb59a0 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)) } @@ -529,10 +536,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 28d6670abe021356d8dadd09c33f6a1eefba848b Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 13:07:00 +0800 Subject: [PATCH 07/11] 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 62607ae54..64abf90ef 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; @@ -572,6 +573,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() @@ -587,11 +589,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 @@ -723,13 +729,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) } @@ -2960,6 +3003,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 5c4affaa5ee4253aef5d2e12f2f958f9748b81cf Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 14:56:49 +0800 Subject: [PATCH 08/11] 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 64abf90ef..26fcbf256 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 { @@ -573,7 +613,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() @@ -589,14 +628,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::>(); @@ -729,50 +761,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) } @@ -2359,7 +2354,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, FilePredicates, + forward_row_group_batches, metadata_has_beneficial_offset_index, FilePredicates, ParquetFormatReader, ParquetFormatWriter, ParquetRowGroupMessage, }; use super::{ @@ -3004,7 +2999,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; @@ -3029,30 +3024,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] @@ -3658,24 +3632,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 b77277c82506c4f632d787d394b58219d1f997bb Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 11:49:57 +0800 Subject: [PATCH 09/11] 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 26fcbf256..0cc278206 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 { @@ -3027,6 +3029,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] @@ -3633,8 +3651,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", @@ -3751,6 +3781,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 { @@ -3758,6 +3808,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)), } } @@ -3779,7 +3831,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 c4dd69774ecbce1dc652cf2b8072a8548e742d10 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 12:55:32 +0800 Subject: [PATCH 10/11] 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 0cc278206..3fb1ff0b1 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -3877,7 +3877,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 c5630ef5e..1257f6328 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 8fdbb59a0..a6314a5ec 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 } } @@ -536,7 +512,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 6d42259610aba6bd6ab6657ff5fedad92c2cd563 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 21:16:39 +0800 Subject: [PATCH 11/11] 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 a77809b3d..45e920c91 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 { @@ -3671,21 +3681,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; @@ -3694,18 +3720,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; @@ -3713,7 +3742,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 {