From 94a8eb0dd4681b6fe471f6e0c33e1e016499d187 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 10 Oct 2026 17:17:14 +0800 Subject: [PATCH] Support sparse vindex builds with Java-compatible training semantics --- .../src/table/pk_vector_search_params.rs | 18 +- .../table/pk_vector_search_params/tests.rs | 34 ++ .../src/table/vindex_index_build_builder.rs | 21 +- .../vindex_index_build_builder/extraction.rs | 200 ++++-- .../vindex_index_build_builder/pipeline.rs | 142 +++-- .../sparse_tests.rs | 534 ++++++++++++++++ .../table/vindex_index_build_builder/tests.rs | 5 +- .../vindex_index_build_builder/writer.rs | 404 ++++++------ crates/paimon/src/vindex/mod.rs | 578 +++++++++--------- docs/src/sql.md | 19 +- 10 files changed, 1338 insertions(+), 617 deletions(-) create mode 100644 crates/paimon/src/table/vindex_index_build_builder/sparse_tests.rs diff --git a/crates/paimon/src/table/pk_vector_search_params.rs b/crates/paimon/src/table/pk_vector_search_params.rs index 954c02851..fbd751db9 100644 --- a/crates/paimon/src/table/pk_vector_search_params.rs +++ b/crates/paimon/src/table/pk_vector_search_params.rs @@ -193,15 +193,31 @@ fn pk_vector_query_dimension( for key in [ "dimension".to_string(), format!("{index_type}.dimension"), + format!("{index_type}.index.dimension"), format!("fields.{}.dimension", vector_field.name()), + format!("fields.{}.index.dimension", vector_field.name()), ] { if let Some(value) = query_options.get(&key) { + let key = if key == "dimension" { + format!("{index_type}.dimension") + } else { + key + }; dimension_options.insert(key, value.clone()); } } + let mut table_dimension_options = table_options.clone(); + if let Some(value) = table_options.get("dimension") { + if !table_options.contains_key(&format!("{index_type}.dimension")) + && !table_options.contains_key(&format!("{index_type}.index.dimension")) + { + table_dimension_options + .insert(format!("{index_type}.dimension"), value.clone()); + } + } Ok(Some( VindexVectorIndexOptions::new( - table_options, + &table_dimension_options, &dimension_options, index_type, vector_field, diff --git a/crates/paimon/src/table/pk_vector_search_params/tests.rs b/crates/paimon/src/table/pk_vector_search_params/tests.rs index 061c09d4b..f01a42740 100644 --- a/crates/paimon/src/table/pk_vector_search_params/tests.rs +++ b/crates/paimon/src/table/pk_vector_search_params/tests.rs @@ -39,3 +39,37 @@ fn vindex_array_dimension_accepts_diskann_search_options() { Some(8) ); } + +#[test] +fn raw_query_dimensions_remain_supported_with_java_build_option_mapping() { + let field = DataField::new( + 1, + "embedding".into(), + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + for index_type in ["ivf-flat", "ivf-pq", "diskann"] { + let table = HashMap::from([("dimension".into(), "8".into())]); + assert_eq!( + pk_vector_query_dimension(&table, &HashMap::new(), index_type, &field).unwrap(), + Some(8) + ); + let query = HashMap::from([("dimension".into(), "2".into())]); + assert_eq!( + pk_vector_query_dimension(&table, &query, index_type, &field).unwrap(), + Some(2) + ); + let table = HashMap::from([ + ("dimension".into(), "8".into()), + (format!("{index_type}.index.dimension"), "16".into()), + ("fields.embedding.index.dimension".into(), "32".into()), + ]); + assert_eq!( + pk_vector_query_dimension(&table, &HashMap::new(), index_type, &field).unwrap(), + Some(32) + ); + assert_eq!( + pk_vector_query_dimension(&table, &query, index_type, &field).unwrap(), + Some(2) + ); + } +} diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs b/crates/paimon/src/table/vindex_index_build_builder.rs index ca049b193..10a4637ea 100644 --- a/crates/paimon/src/table/vindex_index_build_builder.rs +++ b/crates/paimon/src/table/vindex_index_build_builder.rs @@ -92,15 +92,27 @@ impl<'a> VindexIndexBuildBuilder<'a> { let core_options = CoreOptions::new(&merged_options); validate_table_options(self.table, &core_options)?; let rows_per_shard = core_options.global_index_row_count_per_shard()?; + if let Some(value) = merged_options.get("global-index.build.parallelism") { + let parallelism = value + .trim() + .parse::() + .map_err(|_| Error::ConfigInvalid { + message: "Option 'global-index.build.parallelism' must be greater than 0." + .into(), + })?; + if parallelism == 0 { + return Err(Error::ConfigInvalid { + message: "Option 'global-index.build.parallelism' must be greater than 0." + .into(), + }); + } + } let index_field = find_index_field(self.table, index_column)?; validate_vector_field(index_field)?; - // The builder consumes shard sizing; it is not a native writer option. - let mut writer_options = self.options.clone(); - writer_options.remove("global-index.row-count-per-shard"); let vindex_options = VindexVectorIndexOptions::new( self.table.schema().options(), - &writer_options, + &self.options, &self.index_type, index_field, )?; @@ -194,6 +206,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { return Err(error); } }; + let Some(built) = built else { continue }; let mut message = CommitMessage::new(shard.partition_bytes.clone(), 0, vec![]); message.new_index_files = vec![built.meta]; messages.push(message); 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 37a8bb515..b8aabaf7e 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs @@ -21,6 +21,7 @@ use crate::spec::ROW_ID_FIELD_NAME; use crate::table::{DataSplit, DataSplitBuilder, RowRange}; use crate::{Error, Result}; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; +use arrow_buffer::ScalarBuffer; pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result { data_split_for_shard_ranges( @@ -44,27 +45,40 @@ pub(super) fn data_split_for_shard_ranges( .build() } -pub(super) struct ValidatedVectorBatch<'a> { - pub(super) values: &'a [f32], - pub(super) bytes: &'a [u8], - pub(super) row_ids: &'a [i64], - pub(super) row_count: usize, +pub(super) struct ValidatedVectorBatch { + pub(super) values: ScalarBuffer, + pub(super) row_ids: ScalarBuffer, + /// Number of non-null vectors, not the number of source rows. + pub(super) vector_count: usize, + pub(super) source_rows: usize, } -pub(super) fn extract_vector_batch<'a>( - batch: &'a RecordBatch, +impl ValidatedVectorBatch { + pub(super) fn bytes(&self) -> &[u8] { + self.values.inner().as_slice() + } +} + +pub(super) fn contains_null_vectors(batch: &RecordBatch, index_column: &str) -> bool { + batch + .column_by_name(index_column) + .is_some_and(|column| column.null_count() != 0) +} + +pub(super) fn extract_vector_batch( + batch: &RecordBatch, index_column: &str, dimension: usize, -) -> Result> { +) -> Result { validate_vector_batch_with(batch, index_column, dimension, |_| Ok(())) } -pub(super) fn validate_vector_batch<'a>( - batch: &'a RecordBatch, +pub(super) fn validate_vector_batch( + batch: &RecordBatch, index_column: &str, dimension: usize, expected_row_id: &mut i64, -) -> Result> { +) -> Result { validate_vector_batch_with(batch, index_column, dimension, |row_id| { if row_id != *expected_row_id { return Err(Error::DataInvalid { @@ -85,14 +99,14 @@ pub(super) fn validate_vector_batch<'a>( }) } -pub(super) fn validate_vector_batch_ranges<'a>( - batch: &'a RecordBatch, +pub(super) fn validate_vector_batch_ranges( + batch: &RecordBatch, index_column: &str, dimension: usize, ranges: &[RowRange], range_index: &mut usize, expected_row_id: &mut i64, -) -> Result> { +) -> 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}"), @@ -126,12 +140,12 @@ pub(super) fn validate_vector_batch_ranges<'a>( }) } -fn validate_vector_batch_with<'a>( - batch: &'a RecordBatch, +fn validate_vector_batch_with( + batch: &RecordBatch, index_column: &str, dimension: usize, mut validate_row_id: impl FnMut(i64) -> Result<()>, -) -> Result> { +) -> Result { let vector_index = batch .schema() .index_of(index_column) @@ -147,14 +161,30 @@ fn validate_vector_batch_with<'a>( message: format!("_ROW_ID column not found in read batch: {e}"), source: None, })?; + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "vindex vector extraction requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + if row_ids.null_count() != 0 { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null _ROW_ID".to_string(), + source: None, + }); + } + for row_id in row_ids.values() { + validate_row_id(*row_id)?; + } + let column = batch.column(vector_index); + if column.null_count() != 0 { + return compact_nullable_vectors(column.as_ref(), row_ids, dimension, batch.num_rows()); + } + let (values, start, end) = if let Some(array) = column.as_any().downcast_ref::() { - if array.null_count() != 0 { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null vector row".to_string(), - source: None, - }); - } let offsets = array.value_offsets(); for offsets in offsets.windows(2) { let actual = offsets[1] - offsets[0]; @@ -192,12 +222,6 @@ fn validate_vector_batch_with<'a>( source: None, }); } - if array.null_count() != 0 { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null vector row".to_string(), - source: None, - }); - } let end = batch .num_rows() .checked_mul(dimension) @@ -231,30 +255,112 @@ fn validate_vector_batch_with<'a>( source: None, }); } - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "vindex vector extraction requires non-null Int64 _ROW_ID".to_string(), + checked_vector_bytes(end - start, 1)?; + Ok(ValidatedVectorBatch { + values: values.values().slice(start, end - start), + row_ids: row_ids.values().clone(), + vector_count: batch.num_rows(), + source_rows: batch.num_rows(), + }) +} + +fn compact_nullable_vectors( + column: &dyn Array, + row_ids: &Int64Array, + dimension: usize, + source_rows: usize, +) -> Result { + let list = column.as_any().downcast_ref::(); + let fixed = column.as_any().downcast_ref::(); + if list.is_none() && fixed.is_none() { + return Err(Error::DataInvalid { + message: + "vindex vector extraction requires Arrow List or FixedSizeList" + .into(), source: None, - })?; - if row_ids.null_count() != 0 { + }); + } + if let Some(fixed) = fixed { + if fixed.value_length() as usize != dimension { + return Err(Error::DataInvalid { + message: format!( + "vindex vector dimension mismatch: expected {dimension}, got {}", + fixed.value_length() + ), + source: None, + }); + } + } + let child = list + .map(|list| list.values()) + .or_else(|| fixed.map(|list| list.values())) + .unwrap(); + if !child.as_any().is::() { return Err(Error::DataInvalid { - message: "vindex vector extraction found null _ROW_ID".to_string(), + message: "vindex vector extraction requires Float32 vector elements".into(), source: None, }); } - for row_id in row_ids.values() { - validate_row_id(*row_id)?; + let count = source_rows - column.null_count(); + checked_vector_bytes(count, dimension)?; + let mut values = Vec::with_capacity(count * dimension); + let mut ids = Vec::with_capacity(count); + for row in 0..source_rows { + if column.is_null(row) { + continue; + } + let vector = match list { + Some(list) => list.value(row), + None => fixed.unwrap().value(row), + }; + if vector.len() != dimension { + return Err(Error::DataInvalid { + message: format!( + "vindex vector dimension mismatch: expected {dimension}, got {}", + vector.len() + ), + source: None, + }); + } + let vector = vector.as_any().downcast_ref::().unwrap(); + if vector.null_count() != 0 { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null vector element".into(), + source: None, + }); + } + values.extend_from_slice(vector.values()); + ids.push(row_ids.value(row)); } - - let byte_start = checked_vector_bytes(start, 1)?; - let byte_end = checked_vector_bytes(end, 1)?; Ok(ValidatedVectorBatch { - values: &values.values()[start..end], - bytes: &values.values().inner().as_slice()[byte_start..byte_end], - row_ids: row_ids.values(), - row_count: batch.num_rows(), + values: values.into(), + row_ids: ids.into(), + vector_count: count, + source_rows, }) } + +pub(super) fn local_ids(row_ids: &[i64], start: i64, row_count: usize) -> Result> { + let end = start + .checked_add(i64::try_from(row_count).map_err(|e| Error::DataInvalid { + message: "vindex row count does not fit i64".to_string(), + source: Some(Box::new(e)), + })?) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range overflows i64".to_string(), + source: None, + })?; + row_ids + .iter() + .map(|row_id| { + if *row_id < start || *row_id >= end { + Err(Error::DataInvalid { + message: format!("vindex row id {row_id} is outside shard [{start}, {end})"), + source: None, + }) + } else { + Ok(*row_id - start) + } + }) + .collect() +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs index 26088ca35..5848a9c29 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/pipeline.rs @@ -16,7 +16,8 @@ // under the License. use super::extraction::{ - data_split_for_shard_ranges, extract_vector_batch, validate_vector_batch_ranges, + contains_null_vectors, data_split_for_shard_ranges, extract_vector_batch, local_ids, + validate_vector_batch_ranges, }; use super::planning::VindexIndexShard; use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; @@ -56,6 +57,18 @@ const QUEUE_CAPACITY: usize = 2; const BUFFER_BYTES: usize = 8 * 1024 * 1024; const REPLAY_TARGET_BYTES: usize = 32 * 1024 * 1024; +pub(super) enum GranuleBuildOutcome { + Built(Box), + Sparse, +} + +#[derive(PartialEq)] +enum SourceLayout { + Dense, + Sparse, +} + +#[derive(Clone, Debug)] struct Granule { range: RowRange, file_index: usize, @@ -189,19 +202,19 @@ fn spawn_add_consumer( match item { AddItem::Batch(batch, ids) => { let vectors = extract_vector_batch(&batch, &index_column, dimension)?; - if ids.len() != vectors.row_count { + if ids.len() != vectors.vector_count { return Err(Error::DataInvalid { message: "vindex add batch id count mismatch".to_string(), source: None, }); } writer - .add_vectors(&ids, vectors.values, vectors.row_count) + .add_vectors(&ids, &vectors.values, vectors.vector_count) .map_err(|e| Error::UnexpectedError { message: format!("Failed to add vectors to vindex index: {e}"), source: Some(Box::new(e)), })?; - rows_added += vectors.row_count; + rows_added += vectors.vector_count; } AddItem::Spilled(ids, buffer) => { let values = buffer.typed_data::(); @@ -610,31 +623,6 @@ fn granules_partition_shard(granules: &[Granule], shard_range: &RowRange) -> boo .all(|pair| pair[0].range.to().checked_add(1) == Some(pair[1].range.from())) } -fn local_ids(row_ids: &[i64], start: i64, row_count: usize) -> Result> { - let end = start - .checked_add(i64::try_from(row_count).map_err(|e| Error::DataInvalid { - message: "vindex row count does not fit i64".to_string(), - source: Some(Box::new(e)), - })?) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range overflows i64".to_string(), - source: None, - })?; - row_ids - .iter() - .map(|row_id| { - if *row_id < start || *row_id >= end { - Err(Error::DataInvalid { - message: format!("vindex row id {row_id} is outside shard [{start}, {end})"), - source: None, - }) - } else { - Ok(*row_id - start) - } - }) - .collect() -} - impl<'a> VindexIndexBuildBuilder<'a> { fn open_vector_stream( &self, @@ -786,7 +774,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { index_field_id: i32, options: &VindexVectorIndexOptions, index_meta: Vec, - ) -> Result { + ) -> Result { let timing_enabled = vector_index_build_timing_enabled(); let total_start = timing_enabled.then(Instant::now); let mut source_batch_wait = Duration::ZERO; @@ -816,14 +804,15 @@ impl<'a> VindexIndexBuildBuilder<'a> { } checked_vector_bytes(row_count_usize, dimension)?; let eligible = checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; + let training_config = options.training_config(row_count_usize)?; let retained = - default_training_vector_count(eligible, options.config.nlist()).unwrap_or(eligible); + default_training_vector_count(eligible, training_config.nlist()).unwrap_or(eligible); let plan = self .plan_granules(shard, index_column, eligible, retained) .await?; let mut trainer = - VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { + VectorIndexTrainer::new(training_config).map_err(|e| Error::DataInvalid { message: format!("Failed to initialize vindex trainer: {e}"), source: Some(Box::new(e)), })?; @@ -835,7 +824,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut next_training_sample = 0usize; let mut first_rows = 0usize; let mut range_index = 0usize; - let first_result: Result<()> = async { + let first_result: Result = async { let mut stream = self.open_vector_stream( shard, plan.first.clone(), @@ -852,6 +841,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { } let Some(batch) = batch else { break }; batch_count += 1; + if contains_null_vectors(&batch, index_column) { return Ok(SourceLayout::Sparse); } let vectors = validate_vector_batch_ranges( &batch, index_column, @@ -860,7 +850,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { &mut range_index, &mut expected_row_id, )?; - let batch_end = first_rows.checked_add(vectors.row_count).ok_or_else(|| { + let batch_end = first_rows.checked_add(vectors.vector_count).ok_or_else(|| { Error::DataInvalid { message: "vindex first-batch row count overflows usize".to_string(), source: None, @@ -871,7 +861,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { let sample = training[next_training_sample]; match vectors.row_ids.binary_search(&sample) { Ok(row) => row, - Err(row) if row == vectors.row_count => break, + Err(row) if row == vectors.vector_count => break, Err(_) => return Err(Error::DataInvalid { message: format!("Missing vindex training row {sample}"), source: None, @@ -905,8 +895,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { } } let record = SpillRecord { - ids: local_ids(vectors.row_ids, shard.row_range_start, row_count_usize)?, - bytes: vectors.bytes.to_vec(), + ids: local_ids(&vectors.row_ids, shard.row_range_start, row_count_usize)?, + bytes: vectors.bytes().to_vec(), }; first_rows = batch_end; if spill.as_ref().unwrap().sender.send(record).await.is_err() { @@ -937,14 +927,22 @@ impl<'a> VindexIndexBuildBuilder<'a> { source: None, }); } - Ok(()) + Ok(SourceLayout::Dense) } .await; - if let Err(error) = first_result { - return match spill.take().unwrap().finish().await { - Ok(_) => Err(error), - Err(spill_error) => Err(spill_error), - }; + let first_layout = match first_result { + Ok(layout) => layout, + Err(error) => { + return match spill.take().unwrap().finish().await { + Ok(_) => Err(error), + Err(spill_error) => Err(spill_error), + }; + } + }; + + if first_layout == SourceLayout::Sparse { + spill.take().unwrap().finish().await?; + return Ok(GranuleBuildOutcome::Sparse); } let mut training: Option = Some(tokio::task::spawn_blocking( @@ -977,7 +975,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { } let mut rest_rows = 0usize; - let producer_result: Result<()> = async { + let producer_result: Result = async { if !plan.rest.is_empty() { let mut stream = self.open_vector_stream( shard, @@ -999,6 +997,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { if live.is_none() && training.as_ref().is_some_and(|task| task.is_finished()) { go_live!(); } + if contains_null_vectors(&batch, &index_column) { return Ok(SourceLayout::Sparse); } let vectors = validate_vector_batch_ranges( &batch, &index_column, @@ -1007,8 +1006,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { &mut range_index, &mut expected_row_id, )?; - let ids = local_ids(vectors.row_ids, shard.row_range_start, row_count_usize)?; - rest_rows = rest_rows.checked_add(vectors.row_count).ok_or_else(|| { + let ids = local_ids(&vectors.row_ids, shard.row_range_start, row_count_usize)?; + rest_rows = rest_rows.checked_add(vectors.vector_count).ok_or_else(|| { Error::DataInvalid { message: "vindex remaining row count overflows usize".to_string(), source: None, @@ -1029,7 +1028,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { } else { let record = SpillRecord { ids, - bytes: vectors.bytes.to_vec(), + bytes: vectors.bytes().to_vec(), }; if spill.as_ref().unwrap().sender.send(record).await.is_err() { return Err(Error::UnexpectedError { @@ -1052,28 +1051,46 @@ impl<'a> VindexIndexBuildBuilder<'a> { }); } } - Ok(()) + Ok(SourceLayout::Dense) } .await; - if let Err(producer_error) = producer_result { + let rest_layout = match producer_result { + Ok(layout) => layout, + Err(producer_error) => { + if let Some(pipeline) = live.take() { + let (consumer, replay, _, _) = finish_live_pipeline(pipeline).await; + consumer?; + replay?; + } else { + let spill_result = match spill.take() { + Some(spill) => spill.finish().await.map(|_| ()), + None => Ok(()), + }; + let training_result = match training.take() { + Some(training) => join_training(training).await.map(|_| ()), + None => Ok(()), + }; + spill_result?; + training_result?; + } + return Err(producer_error); + } + }; + if rest_layout == SourceLayout::Sparse { if let Some(pipeline) = live.take() { let (consumer, replay, _, _) = finish_live_pipeline(pipeline).await; consumer?; replay?; } else { - let spill_result = match spill.take() { - Some(spill) => spill.finish().await.map(|_| ()), - None => Ok(()), - }; - let training_result = match training.take() { - Some(training) => join_training(training).await.map(|_| ()), - None => Ok(()), - }; - spill_result?; - training_result?; + if let Some(spill) = spill.take() { + spill.finish().await?; + } + if let Some(training) = training.take() { + join_training(training).await?; + } } - return Err(producer_error); + return Ok(GranuleBuildOutcome::Sparse); } if live.is_none() { go_live!(); @@ -1141,7 +1158,10 @@ impl<'a> VindexIndexBuildBuilder<'a> { data_file_count: shard.files.len(), file_name: meta.file_name.clone(), }); - Ok(BuiltIndexFile { meta, timing }) + Ok(GranuleBuildOutcome::Built(Box::new(BuiltIndexFile { + meta, + timing, + }))) } } diff --git a/crates/paimon/src/table/vindex_index_build_builder/sparse_tests.rs b/crates/paimon/src/table/vindex_index_build_builder/sparse_tests.rs new file mode 100644 index 000000000..6a8def68a --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/sparse_tests.rs @@ -0,0 +1,534 @@ +// 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::table::vindex_index_build_builder::extraction::{extract_vector_batch, local_ids}; +use arrow_array::Array; +use arrow_buffer::NullBuffer; +use paimon_vindex_core::index::VectorIndexReader; +use std::io::Cursor; + +fn nullable_batch(rows: Vec>>>, start: i64) -> RecordBatch { + let mut vectors = ListBuilder::new(Float32Builder::new()); + for row in &rows { + if let Some(row) = row { + for value in row { + vectors.values().append_option(*value); + } + vectors.append(true); + } else { + vectors.append(false); + } + } + let vectors = vectors.finish(); + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("embedding", vectors.data_type().clone(), true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])), + vec![ + Arc::new(vectors), + Arc::new(Int64Array::from( + (start..start + rows.len() as i64).collect::>(), + )), + ], + ) + .unwrap() +} + +#[test] +fn null_rows_are_compacted_without_renumbering_or_accepting_null_ids() { + let batch = nullable_batch( + vec![ + None, + Some(vec![Some(1.), Some(2.)]), + None, + Some(vec![Some(3.), Some(4.)]), + None, + ], + 41, + ) + .slice(1, 3); + let mut expected = 42; + let result = validate_vector_batch(&batch, "embedding", 2, &mut expected).unwrap(); + assert_eq!(expected, 45); + assert_eq!(result.source_rows, 3); + assert_eq!(result.vector_count, 2); + assert_eq!(result.values.as_ref(), &[1., 2., 3., 4.]); + assert_eq!(result.row_ids.as_ref(), &[42, 44]); + assert_eq!(local_ids(&result.row_ids, 40, 10).unwrap(), vec![2, 4]); + assert_eq!(result.bytes().len(), 16); + let bad = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + batch.schema().field(0).clone(), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])), + vec![ + batch.column(0).clone(), + Arc::new(Int64Array::from(vec![Some(42), None, Some(44)])), + ], + ) + .unwrap(); + assert!(extract_vector_batch(&bad, "embedding", 2).is_err()); +} + +#[test] +fn fixed_size_null_rows_ignore_hidden_child_nulls_and_preserve_slices() { + let array = arrow_array::FixedSizeListArray::new( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + 2, + Arc::new(arrow_array::Float32Array::from(vec![ + None, + None, + Some(1.), + Some(2.), + None, + None, + Some(3.), + Some(4.), + None, + None, + ])), + Some(NullBuffer::from(vec![false, true, false, true, false])), + ); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("embedding", array.data_type().clone(), true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])), + vec![ + Arc::new(array), + Arc::new(Int64Array::from(vec![10, 11, 12, 13, 14])), + ], + ) + .unwrap() + .slice(1, 3); + let vectors = extract_vector_batch(&batch, "embedding", 2).unwrap(); + assert_eq!(vectors.values.as_ref(), &[1., 2., 3., 4.]); + assert_eq!(vectors.row_ids.as_ref(), &[11, 13]); + assert_eq!(vectors.source_rows, 3); + assert!(extract_vector_batch(&batch, "embedding", 3).is_err()); +} + +#[test] +fn null_rows_do_not_hide_invalid_dimensions_elements_or_row_ranges() { + for vector in [vec![Some(1.)], vec![Some(1.), None]] { + let batch = nullable_batch(vec![None, Some(vector), None], 100); + assert!(extract_vector_batch(&batch, "embedding", 2).is_err()); + } + let batch = nullable_batch(vec![None, Some(vec![Some(1.), Some(2.)])], 100); + let mut expected = 99; + assert!(validate_vector_batch(&batch, "embedding", 2, &mut expected).is_err()); + let mut range_index = 0; + let mut expected = 100; + assert!(validate_vector_batch_ranges( + &batch, + "embedding", + 2, + &[RowRange::new(100, 100)], + &mut range_index, + &mut expected + ) + .is_err()); +} + +#[test] +fn dense_batches_keep_borrowed_arrow_storage() { + let batch = nullable_batch( + vec![ + Some(vec![Some(1.), Some(2.)]), + Some(vec![Some(3.), Some(4.)]), + ], + 0, + ); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let vectors = extract_vector_batch(&batch, "embedding", 2).unwrap(); + assert_eq!(vectors.values.as_ptr(), values.values().as_ptr()); +} + +async fn append_nullable(table: &Table, rows: Vec>>>) { + let batch = nullable_batch(rows, 0); + let ids: ArrayRef = Arc::new(Int32Array::from( + (0..batch.num_rows() as i32).collect::>(), + )); + let input = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", batch.column(0).data_type().clone(), true), + ])), + vec![ids, batch.column(0).clone()], + ) + .unwrap(); + let mut writer = TableWrite::new(table, "sparse".into()).unwrap(); + writer.write_arrow_batch(&input).await.unwrap(); + TableCommit::new(table.clone(), "sparse".into()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); +} + +async fn metadata( + table: &Table, + file: &IndexFileMeta, +) -> paimon_vindex_core::index::VectorIndexMetadata { + let path = file + .external_path + .clone() + .unwrap_or_else(|| format!("{}/index/{}", table.location(), file.file_name)); + let input = table.file_io().new_input(&path).unwrap(); + VectorIndexReader::open(Cursor::new(input.read().await.unwrap())) + .unwrap() + .metadata() +} + +#[tokio::test] +async fn sparse_and_all_null_shards_obey_prepare_and_cardinality_contract() { + for kind in ["ivf-flat", "ivf-pq", "ivf-sq", "ivf-rq", "diskann"] { + for granule in [false, true] { + let table = vindex_e2e_table(&format!("memory:/sparse_{kind}_{granule}"), "10"); + setup_dirs(table.file_io(), table.location()).await; + append_nullable( + &table, + vec![ + None, + Some(vec![Some(1.), Some(0.)]), + None, + Some(vec![Some(0.), Some(1.)]), + None, + ], + ) + .await; + let mut builder = table.new_global_index_build_builder(); + let mut options = HashMap::from([(format!("{kind}.dimension"), "2".into())]); + if kind != "diskann" { + options.insert(format!("{kind}.nlist"), "1".into()); + options.insert("vindex.build.granule.enabled".into(), granule.to_string()); + } + builder + .with_index_column("embedding") + .with_index_type(kind) + .with_options(options); + let messages = builder.build().await.unwrap(); + assert_eq!(messages.len(), 1, "{kind}, granule={granule}"); + let file = &messages[0].new_index_files[0]; + assert_eq!(file.row_count, 5); + assert_eq!(file.global_index_meta.as_ref().unwrap().row_range_start, 0); + assert_eq!(metadata(&table, file).await.total_vectors, 2); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(), + 1 + ); + TableCommit::new(table.clone(), "sparse-index".into()) + .commit(messages) + .await + .unwrap(); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1., 1.]) + .with_limit(10) + .execute() + .await + .unwrap(); + let mut ids = result.row_ids().unwrap().row_ids.clone(); + ids.sort_unstable(); + assert_eq!(ids, vec![1, 3], "{kind}, granule={granule}"); + assert!(builder.build().await.unwrap().is_empty()); + append_nullable(&table, vec![None, None, None]).await; + let before = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(); + assert!(builder.build().await.unwrap().is_empty()); + assert_eq!(builder.execute().await.unwrap(), 0); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(), + before + ); + } + } +} + +#[tokio::test] +async fn automatic_ivf_sizing_uses_valid_vectors_and_declared_row_count_includes_nulls() { + let table = vindex_e2e_table("memory:/sparse_auto", "10000"); + setup_dirs(table.file_io(), table.location()).await; + let mut rows = vec![None; 1000]; + rows[31] = Some(vec![Some(1.), Some(0.)]); + rows[901] = Some(vec![Some(0.), Some(1.)]); + append_nullable(&table, rows).await; + let messages = table + .new_global_index_build_builder() + .with_index_column("embedding") + .with_index_type("ivf-flat") + .with_options(HashMap::from([ + ("ivf-flat.nlist".into(), "auto".into()), + ("ivf-flat.train.sample-ratio".into(), "0.5".into()), + ])) + .build() + .await + .unwrap(); + let file = &messages[0].new_index_files[0]; + assert_eq!(file.row_count, 1000); + let meta = metadata(&table, file).await; + assert_eq!(meta.total_vectors, 2); + assert_eq!( + meta.nlist, + paimon_vindex_core::autotune::infer_ivf_nlist(2).unwrap() + ); +} + +#[tokio::test] +async fn null_in_rest_granules_discards_private_training_then_builds_sparse_index() { + use super::super::pipeline::GranuleBuildOutcome; + use super::super::VindexIndexBuildBuilder; + use crate::vindex::VindexVectorIndexOptions; + use parquet::arrow::AsyncArrowWriter; + use parquet::basic::Compression; + use parquet::file::properties::WriterProperties; + + const FILES: usize = 1024; + const ROWS: usize = 256; + let table = vindex_e2e_table("memory:/sparse_rest", "1000000"); + setup_dirs(table.file_io(), table.location()).await; + let dense = nullable_batch(vec![Some(vec![Some(1.), Some(0.)]); ROWS], 0); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(ROWS)) + .set_offset_index_disabled(true) + .set_dictionary_enabled(false) + .set_compression(Compression::UNCOMPRESSED) + .build(); + let serialize = |batch: RecordBatch| { + let properties = properties.clone(); + async move { + let mut bytes = Vec::new(); + let mut writer = + AsyncArrowWriter::try_new(&mut bytes, batch.schema(), Some(properties)).unwrap(); + writer.write(&batch).await.unwrap(); + writer.close().await.unwrap(); + bytes::Bytes::from(bytes) + } + }; + let dense_bytes = serialize(dense).await; + let mut files = Vec::new(); + for i in 0..FILES { + let name = format!("data-{i:04}.parquet"); + table + .file_io() + .new_output(&format!("{}/bucket-0/{name}", table.location())) + .unwrap() + .write(dense_bytes.clone()) + .await + .unwrap(); + let mut file = data_file(&name, Some((i * ROWS) as i64), ROWS as i64); + file.file_size = dense_bytes.len() as i64; + file.file_source = Some(0); + files.push(file); + } + let plan = |files: Vec| { + plan_vindex_shards( + table.location(), + table.schema().partition_keys(), + table.schema().fields(), + &CoreOptions::new(table.schema().options()), + 1, + files + .into_iter() + .map(|file| { + ManifestEntry::new( + FileKind::Add, + BinaryRow::new(0).to_serialized_bytes(), + 0, + 1, + file, + 2, + ) + }) + .collect(), + 1_000_000, + &[], + ) + .unwrap() + }; + let builder = VindexIndexBuildBuilder::new(&table, IVF_FLAT_IDENTIFIER); + let dense_plan = builder + .plan_granules( + &plan(files.clone())[0], + "embedding", + FILES * ROWS / 8, + FILES * ROWS / 8, + ) + .await + .unwrap(); + assert!(!dense_plan.rest.is_empty()); + let null_row = dense_plan.rest.last().unwrap().to(); + let file = &mut files[null_row as usize / ROWS]; + let mut rows = vec![Some(vec![Some(1.), Some(0.)]); ROWS]; + rows[null_row as usize % ROWS] = None; + let sparse_bytes = serialize(nullable_batch(rows, 0)).await; + table + .file_io() + .new_output(&format!("{}/bucket-0/{}", table.location(), file.file_name)) + .unwrap() + .write(sparse_bytes.clone()) + .await + .unwrap(); + file.file_size = sparse_bytes.len() as i64; + for file in &mut files { + file.first_row_id = None; + } + TableCommit::new(table.clone(), "sparse-rest".into()) + .commit(vec![CommitMessage::new( + BinaryRow::new(0).to_serialized_bytes(), + 0, + files, + )]) + .await + .unwrap(); + let snapshot = table + .snapshot_manager() + .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.location(), + table.schema().partition_keys(), + table.schema().fields(), + &CoreOptions::new(table.schema().options()), + snapshot.id(), + entries, + 1_000_000, + &[], + ) + .unwrap(); + let shard = &shards[0]; + let actual_plan = builder + .plan_granules(shard, "embedding", FILES * ROWS / 8, FILES * ROWS / 8) + .await + .unwrap(); + assert_eq!(actual_plan.first, dense_plan.first); + assert_eq!(actual_plan.rest, dense_plan.rest); + assert!(actual_plan + .first + .iter() + .all(|range| !(range.from()..=range.to()).contains(&null_row))); + let user_options = HashMap::from([("ivf-flat.train.sample-ratio".into(), "0.125".into())]); + let field = find_index_field(&table, "embedding").unwrap(); + let options = VindexVectorIndexOptions::new( + table.schema().options(), + &user_options, + IVF_FLAT_IDENTIFIER, + field, + ) + .unwrap(); + let outcome = builder + .build_index_file_granule( + shard, + "embedding", + 2, + field.id(), + &options, + serde_json::to_vec(&options.native_options).unwrap(), + ) + .await + .unwrap(); + assert!(matches!(outcome, GranuleBuildOutcome::Sparse)); + let root = format!("{}/index", table.location()); + assert!(table.file_io().list_status(&root).await.unwrap().is_empty()); + let messages = table + .new_global_index_build_builder() + .with_index_column("embedding") + .with_index_type(IVF_FLAT_IDENTIFIER) + .with_options(user_options) + .build() + .await + .unwrap(); + let file = &messages[0].new_index_files[0]; + assert_eq!(file.row_count, (FILES * ROWS) as i64); + assert_eq!( + metadata(&table, file).await.total_vectors, + (FILES * ROWS - 1) as i64 + ); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(), + 1 + ); +} + +#[tokio::test] +async fn automatic_pq_budget_is_validated_against_non_null_cardinality() { + let table = vindex_e2e_table("memory:/sparse_budget", "10000"); + setup_dirs(table.file_io(), table.location()).await; + append_nullable( + &table, + (0..2000) + .map(|i| (i % 2 == 1).then_some(vec![Some(1.), Some(0.)])) + .collect(), + ) + .await; + let messages = table + .new_global_index_build_builder() + .with_index_column("embedding") + .with_index_type("ivf-pq") + .with_options(HashMap::from([ + ("ivf-pq.dimension".into(), "2".into()), + ("ivf-pq.max-bytes-per-vector".into(), "64".into()), + ])) + .build() + .await + .unwrap(); + let file = &messages[0].new_index_files[0]; + assert_eq!(file.row_count, 2000); + assert_eq!(metadata(&table, file).await.total_vectors, 1000); +} 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 97fad54d8..998791ef3 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -222,7 +222,7 @@ fn extract_vectors_from_batches( let mut vectors = Vec::new(); for batch in batches { vectors.extend_from_slice( - validate_vector_batch(batch, index_column, dimension, &mut expected_row_id)?.values, + &validate_vector_batch(batch, index_column, dimension, &mut expected_row_id)?.values, ); } if expected_row_id - row_range_start != expected_row_count { @@ -1261,3 +1261,6 @@ fn vindex_plan_incremental_prefix_leaves_suffix() { ); } } + +#[path = "sparse_tests.rs"] +mod sparse_tests; 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 bc4c77d25..7b77cff4b 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -use super::extraction::{data_split_for_shard, validate_vector_batch}; +use super::extraction::{data_split_for_shard, local_ids, validate_vector_batch}; +use super::pipeline::GranuleBuildOutcome; use super::planning::VindexIndexShard; use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; use super::validation::{ @@ -31,7 +32,7 @@ use crate::{Error, Result}; use arrow_buffer::MutableBuffer; use futures::TryStreamExt; use paimon_vindex_core::autotune::default_training_vector_count; -use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; +use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::PosWriter; use std::io::{Read, Seek, SeekFrom}; use std::sync::Arc; @@ -55,24 +56,47 @@ impl<'a> VindexIndexBuildBuilder<'a> { index_field_id: i32, options: &VindexVectorIndexOptions, index_meta: Vec, - ) -> Result { - let use_granule = self.index_type != DISKANN_IDENTIFIER && options.granule_build_enabled; + ) -> Result> { + let rows = usize::try_from(checked_row_count( + shard.row_range_start, + shard.row_range_end, + )?) + .map_err(|error| Error::DataInvalid { + message: "vindex row count does not fit usize".into(), + source: Some(Box::new(error)), + })?; + let use_granule = self.index_type != DISKANN_IDENTIFIER + && options.granule_build_enabled + && !options.needs_vector_count() + // Invalid training options still produce no file for an all-NULL + // shard, like Java. Full spill determines whether training is needed. + && options.training_config(rows).is_ok(); + // Auto IVF sizing needs the complete non-null cardinality, like Java. + // Explicit nlist/expected-count builds retain the granule fast path. log::info!( "vindex build strategy: index_type={}, strategy={}", self.index_type, if use_granule { "granule" } else { "full-spill" } ); if use_granule { - return self + match self .build_index_file_granule( shard, index_column, dimension, index_field_id, options, - index_meta, + index_meta.clone(), ) - .await; + .await? + { + GranuleBuildOutcome::Built(built) => return Ok(Some(*built)), + GranuleBuildOutcome::Sparse => { + // Only private in-memory/spilled preparation has happened. The + // pinned snapshot is reread with sparse row IDs and valid-only training. + log::info!("vindex granule source contains NULL vectors; use sparse full-spill preparation"); + } + } } self.build_full_spill_index_file( shard, @@ -93,7 +117,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { index_field_id: i32, options: &VindexVectorIndexOptions, index_meta: Vec, - ) -> Result { + ) -> Result> { let timing_enabled = vector_index_build_timing_enabled(); let total_start = timing_enabled.then(Instant::now); let mut source_batch_wait = Duration::ZERO; @@ -122,28 +146,9 @@ impl<'a> VindexIndexBuildBuilder<'a> { 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_buffer_rows = - (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); - let training_buffer_floats = training_buffer_rows - .checked_mul(dimension_usize) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; - - let mut trainer = - VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { - 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); + checked_vector_bytes(row_count_usize, dimension_usize)?; + let mut raw_file = temporary_vector_file("vectors")?; + let mut id_file = temporary_vector_file("row IDs")?; 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])?; @@ -159,211 +164,115 @@ impl<'a> VindexIndexBuildBuilder<'a> { let mut batches = read.to_arrow(&[split])?; let mut expected_row_id = shard.row_range_start; let mut rows_seen = 0usize; + let mut vector_count = 0usize; let mut bytes_written = 0usize; - let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(training_buffer_floats); - loop { let source_start = timing_enabled.then(Instant::now); let batch = batches.try_next().await?; - if let Some(source_start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); + if let Some(start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(start.elapsed()); } let Some(batch) = batch else { break }; batch_count += 1; let vectors = validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; - let batch_end = + rows_seen = rows_seen - .checked_add(vectors.row_count) + .checked_add(vectors.source_rows) .ok_or_else(|| Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), + message: "vindex streamed row count overflows usize".into(), 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(); - } - } - } - + vector_count = vector_count + .checked_add(vectors.vector_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex streamed vector count overflows usize".into(), + source: None, + })?; + let ids: arrow_buffer::ScalarBuffer = + local_ids(&vectors.row_ids, shard.row_range_start, row_count_usize)?.into(); let raw_write_start = timing_enabled.then(Instant::now); raw_file - .write_all(vectors.bytes) + .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()); + .map_err(spill_error)?; + id_file + .write_all(ids.inner().as_slice()) + .await + .map_err(spill_error)?; + if let Some(start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(start.elapsed()); } bytes_written = bytes_written - .checked_add(vectors.bytes.len()) + .checked_add(vectors.bytes().len()) .ok_or_else(|| Error::DataInvalid { - message: "vindex spilled byte count overflows usize".to_string(), + message: "vindex spilled byte count overflows usize".into(), 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".into(), + source: None, })?; - } if rows_seen != row_count_usize - || expected_row_id - != shard - .row_range_end - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range end overflows i64".to_string(), - source: None, - })? - || (training_vector_count != row_count_usize - && next_training_sample != training_vector_count) - || bytes_written != expected_bytes + || expected_row_id != expected_end + || bytes_written != checked_vector_bytes(vector_count, dimension_usize)? { 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, + message: format!("vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, vectors={vector_count}, bytes={bytes_written}"), 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()); + // Like Java NativeVectorGlobalIndexWriter.finish: no vectors means no file. + if vector_count == 0 { + return Ok(None); } - 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 { + raw_file.flush().await.map_err(spill_error)?; + id_file.flush().await.map_err(spill_error)?; + if raw_file.metadata().await.map_err(spill_error)?.len() != bytes_written as u64 + || id_file.metadata().await.map_err(spill_error)?.len() + != checked_vector_bytes(vector_count, 2)? as u64 + { return Err(Error::DataInvalid { - message: format!( - "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" - ), + message: "temporary vindex vector or row ID file size mismatch".into(), source: None, }); } let raw_file = raw_file.into_std().await; - // Diagnostics only: never fail the build for a timing log field. + let id_file = id_file.into_std().await; + let training_vector_count = + checked_training_vector_count(vector_count, options.train_sample_ratio)?; + let config = options.training_config(vector_count)?; let training_rows_retained = if timing_enabled { - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) + default_training_vector_count(training_vector_count, config.nlist()).unwrap_or(0) } else { 0 }; - - let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( - move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { - let train_start = timing_enabled.then(Instant::now); - let training = trainer.finish()?; - let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let mut writer = VectorIndexWriter::new(training); - let mut raw_temp_reread = Duration::ZERO; - let mut index_add = Duration::ZERO; - let mut raw_file = raw_file; - let reread_start = timing_enabled.then(Instant::now); - raw_file.seek(SeekFrom::Start(0))?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - let batch_rows = training_buffer_rows.min(row_count_usize); - let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; - let mut buffer = MutableBuffer::new(batch_bytes); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); - let reread_start = timing_enabled.then(Instant::now); - raw_file.read_exact(buffer.as_slice_mut())?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); - } - let add_start = timing_enabled.then(Instant::now); - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - if let Some(start) = add_start { - index_add = index_add.saturating_add(start.elapsed()); - } - rows_added += rows; - } - let mut trailing = [0u8; 1]; - let reread_start = timing_enabled.then(Instant::now); - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - Ok((writer, train_finish, raw_temp_reread, index_add)) - }, - ) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("vindex training task failed: {e}"), - source: None, - })? - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to train or add vectors to vindex index: {e}"), - source: Some(Box::new(e)), - })?; + let ratio = options.train_sample_ratio; + let (writer, train_finish, raw_temp_reread, index_add) = + tokio::task::spawn_blocking(move || { + train_and_add_spilled_vectors( + raw_file, + id_file, + config, + vector_count, + dimension_usize, + ratio, + timing_enabled, + ) + }) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex training task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to train or add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; let serialize_upload_start = timing_enabled.then(Instant::now); let meta = self @@ -414,7 +323,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { data_file_count: shard.files.len(), file_name: meta.file_name.clone(), }); - Ok(BuiltIndexFile { meta, timing }) + Ok(Some(BuiltIndexFile { meta, timing })) } pub(super) async fn finish_index_file( @@ -490,3 +399,102 @@ impl<'a> VindexIndexBuildBuilder<'a> { }) } } + +fn spill_error(error: std::io::Error) -> Error { + Error::UnexpectedError { + message: format!("Failed to spill vindex vectors: {error}"), + source: Some(Box::new(error)), + } +} + +fn temporary_vector_file(kind: &str) -> Result { + tempfile::tempfile() + .map(tokio::fs::File::from_std) + .map_err(|error| Error::UnexpectedError { + message: format!("Failed to create temporary vindex {kind} file: {error}"), + source: Some(Box::new(error)), + }) +} + +/// Java samples the compacted non-null vector stream, then adds the original relative IDs. +fn train_and_add_spilled_vectors( + mut raw_file: std::fs::File, + mut id_file: std::fs::File, + config: VectorIndexConfig, + vector_count: usize, + dimension: usize, + ratio: f64, + timing_enabled: bool, +) -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { + let train_start = timing_enabled.then(Instant::now); + let samples = + checked_training_vector_count(vector_count, ratio).map_err(std::io::Error::other)?; + let batch_rows = (VECTOR_BUFFER_BYTES / checked_std_vector_bytes(1, dimension)?) + .max(1) + .min(vector_count); + let mut buffer = MutableBuffer::new(checked_std_vector_bytes(batch_rows, dimension)?); + let mut sample_buffer = Vec::with_capacity(batch_rows * dimension); + let mut trainer = VectorIndexTrainer::new(config)?; + raw_file.seek(SeekFrom::Start(0))?; + let mut rows_seen = 0; + let mut selected = 0; + while rows_seen < vector_count { + let rows = batch_rows.min(vector_count - rows_seen); + buffer.resize(checked_std_vector_bytes(rows, dimension)?, 0); + raw_file.read_exact(buffer.as_slice_mut())?; + let values = buffer.typed_data::(); + while selected < samples { + let sample = checked_training_sample_index(selected, vector_count, samples) + .map_err(std::io::Error::other)?; + if sample >= rows_seen + rows { + break; + } + let offset = (sample - rows_seen) * dimension; + sample_buffer.extend_from_slice(&values[offset..offset + dimension]); + selected += 1; + } + if !sample_buffer.is_empty() { + trainer.add_training_vectors_mut(&sample_buffer, sample_buffer.len() / dimension)?; + sample_buffer.clear(); + } + rows_seen += rows; + } + if selected != samples { + return Err(std::io::Error::other( + "vindex training sample count mismatch", + )); + } + 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; + raw_file.seek(SeekFrom::Start(0))?; + id_file.seek(SeekFrom::Start(0))?; + let mut ids = MutableBuffer::new(batch_rows * std::mem::size_of::()); + let mut rows_added = 0; + while rows_added < vector_count { + let rows = batch_rows.min(vector_count - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension)?, 0); + ids.resize(rows * std::mem::size_of::(), 0); + let read_start = timing_enabled.then(Instant::now); + raw_file.read_exact(buffer.as_slice_mut())?; + id_file.read_exact(ids.as_slice_mut())?; + if let Some(start) = read_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + let add_start = timing_enabled.then(Instant::now); + writer.add_vectors(ids.typed_data::(), buffer.typed_data::(), rows)?; + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + rows_added += rows; + } + let mut trailing = [0u8; 1]; + if raw_file.read(&mut trailing)? != 0 || id_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::other( + "temporary vindex vector or row ID file contains trailing bytes", + )); + } + Ok((writer, train_finish, raw_temp_reread, index_add)) +} diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs index 90d6db1fc..425a0eec8 100644 --- a/crates/paimon/src/vindex/mod.rs +++ b/crates/paimon/src/vindex/mod.rs @@ -36,9 +36,6 @@ pub const DISKANN_IDENTIFIER: &str = "diskann"; const DEFAULT_DIMENSION: &str = "128"; const DEFAULT_METRIC: &str = "inner_product"; -const DEFAULT_NLIST: &str = "256"; -const DEFAULT_PQ_M: &str = "16"; -const DEFAULT_PQ_USE_OPQ: &str = "false"; const DEFAULT_TRAIN_SAMPLE_RATIO: f64 = 1.0; const VINDEX_BUILD_GRANULE_ENABLED: &str = "vindex.build.granule.enabled"; const VECTOR_SEARCH_TIMING_ENV: &str = "PAIMON_LOG_VECTOR_SEARCH_TIMING"; @@ -112,7 +109,7 @@ pub(crate) fn native_index_type(index_type: &str) -> Option<&'static str> { #[derive(Debug)] pub(crate) struct VindexVectorIndexOptions { - pub config: VectorIndexConfig, + dimension: usize, pub native_options: HashMap, pub train_sample_ratio: f64, pub granule_build_enabled: bool, @@ -132,205 +129,272 @@ impl VindexVectorIndexOptions { })?; validate_user_option_keys(user_options, index_type, field.name())?; - validate_index_type_option(table_options, user_options, native_index_type)?; - + validate_index_type_option(user_options, native_index_type)?; let mut native_options = HashMap::new(); - native_options.insert("index.type".to_string(), native_index_type.to_string()); - native_options.insert( - "dimension".to_string(), - resolve_dimension(table_options, user_options, index_type, field)?, - ); - if index_type != DISKANN_IDENTIFIER { - native_options.insert( - "nlist".to_string(), - option_value( - table_options, - user_options, - field.name(), - index_type, - "nlist", - "nlist", - DEFAULT_NLIST, - ), - ); - } - native_options.insert( - "metric".to_string(), - normalize_metric(&option_value( - table_options, - user_options, - field.name(), - index_type, - "metric", - "distance.metric", - DEFAULT_METRIC, - )), - ); - - if index_type == IVF_PQ_IDENTIFIER { - native_options.insert( - "pq.m".to_string(), - option_value( - table_options, - user_options, - field.name(), - index_type, - "pq.m", - "pq.m", - DEFAULT_PQ_M, - ), - ); - native_options.insert( - "use-opq".to_string(), - option_value( - table_options, - user_options, - field.name(), - index_type, - "use-opq", - "pq.use-opq", - DEFAULT_PQ_USE_OPQ, - ), - ); - } - if index_type == IVF_RQ_IDENTIFIER { - for key in ["rq.bits", "max-bytes-per-vector"] { - if let Some(value) = optional_value( - table_options, - user_options, - field.name(), - index_type, - key, - key, - ) { - native_options.insert(key.to_string(), value); - } - } - } - if index_type == DISKANN_IDENTIFIER { - for &(native_key, paimon_suffix) in DISKANN_OPTION_KEYS { - if let Some(value) = optional_value( - table_options, - user_options, - field.name(), - index_type, - native_key, - paimon_suffix, - ) { - native_options.insert(native_key.to_string(), value); - } - } - } - for key in [ - "ivf.coarse-assignment", - "ivf.pq-encoding", - "ivf.train.max-points-per-centroid", - "pq.train.max-points-per-centroid", - ] { - if is_allowed_native_key(key, index_type) { - if let Some(value) = optional_value( - table_options, - user_options, - field.name(), - index_type, - key, - key, - ) { - native_options.insert(key.to_string(), value); - } - } - } - - let config = VectorIndexConfig::from_options(&native_options).map_err(|e| { - crate::Error::DataInvalid { - message: format!("Invalid vindex options: {e}"), - source: Some(Box::new(e)), - } - })?; + // Source precedence is build > table; within each source, field > index + // > native keys. Table options follow Java; explicit builder options + // also accept the documented Rust/SQL native aliases. + collect_native_options( + &mut native_options, + table_options, + index_type, + field.name(), + false, + )?; + collect_native_options( + &mut native_options, + user_options, + index_type, + field.name(), + true, + )?; + native_options.insert("index.type".into(), native_index_type.into()); + let dimension = match field.data_type() { + DataType::Vector(vector) => vector.length().to_string(), + _ => native_options + .get("dimension") + .cloned() + .unwrap_or_else(|| DEFAULT_DIMENSION.into()), + }; + native_options.insert("dimension".into(), dimension); + let metric = native_options + .entry("metric".into()) + .or_insert_with(|| DEFAULT_METRIC.into()); + *metric = normalize_metric(metric); + + // Like Java, compile the native configuration only when training begins. + // A synthetic count can reject budgets satisfied by the real vector count. + let dimension = native_options["dimension"] + .trim() + .parse::() + .ok() + .filter(|dimension| *dimension > 0) + .ok_or_else(|| crate::Error::DataInvalid { + message: "vindex vector dimension must be a positive integer".into(), + source: None, + })?; let train_sample_ratio = resolve_train_sample_ratio(table_options, user_options, index_type, field.name())?; let granule_build_enabled = resolve_granule_build_enabled(table_options, user_options)?; Ok(Self { - config, + dimension, native_options, train_sample_ratio, granule_build_enabled, }) } + pub(crate) fn needs_vector_count(&self) -> bool { + self.native_options + .get("index.type") + .is_some_and(|kind| kind.starts_with("ivf_")) + && self + .native_options + .get("nlist") + .is_none_or(|value| value.trim() == "auto") + && !self.native_options.contains_key("expected-vector-count") + } + + pub(crate) fn training_config(&self, vector_count: usize) -> crate::Result { + compile_training_config(&self.native_options, vector_count) + } + pub fn dimension(&self) -> usize { - self.config.dimension() + self.dimension } } -fn validate_index_type_option( - table_options: &HashMap, - user_options: &HashMap, - expected_native: &str, -) -> crate::Result<()> { - for options in [table_options, user_options] { - if let Some(value) = options.get("index.type") { - let normalized = value.trim().to_ascii_lowercase().replace('-', "_"); - if normalized != expected_native { - return Err(crate::Error::ConfigInvalid { - message: format!( - "Option 'index.type' is '{}', but procedure index_type resolves to '{}'. \ - Remove 'index.type' from options or set it to '{}'.", - value, expected_native, expected_native - ), - }); - } - } +fn compile_training_config( + native_options: &HashMap, + vector_count: usize, +) -> crate::Result { + let mut options = native_options.clone(); + // Java NativeVectorGlobalIndexWriter.trainingOptions uses valid vectors. + if options + .get("index.type") + .is_some_and(|kind| kind.starts_with("ivf_")) + && options + .get("nlist") + .is_none_or(|value| value.trim() == "auto") + { + options + .entry("expected-vector-count".into()) + .or_insert_with(|| vector_count.to_string()); } - Ok(()) + VectorIndexConfig::from_options(&options).map_err(|error| crate::Error::DataInvalid { + message: format!("Invalid vindex options: {error}"), + source: Some(Box::new(error)), + }) } fn validate_user_option_keys( - user_options: &HashMap, + options: &HashMap, index_type: &str, - field_name: &str, + field: &str, ) -> crate::Result<()> { - let mut unknown = user_options + let mut unknown = options .keys() - .filter(|key| !is_supported_user_option_key(key, index_type, field_name)) + .filter(|key| { + if matches!( + key.as_str(), + "index.type" + | "global-index.row-count-per-shard" + | "global-index.build.parallelism" + ) { + return false; + } + if key.as_str() == VINDEX_BUILD_GRANULE_ENABLED { + return index_type == DISKANN_IDENTIFIER; + } + let suffix = key + .strip_prefix(&format!("{index_type}.")) + .or_else(|| key.strip_prefix(&format!("fields.{field}."))) + .unwrap_or(key); + if suffix == "train.sample-ratio" { + return key.as_str() == suffix; + } + native_option_key(suffix).is_none_or(|key| !native_option_applies(key, index_type)) + }) .cloned() .collect::>(); if unknown.is_empty() { return Ok(()); } - unknown.sort(); Err(crate::Error::ConfigInvalid { message: format!( - "Unknown vindex option(s) for index_type '{}': {}", - index_type, + "Unknown vindex option(s) for index_type '{index_type}': {}", unknown.join(", ") ), }) } -fn is_supported_user_option_key(key: &str, index_type: &str, field_name: &str) -> bool { - if key == "index.type" { - return true; - } - if key == VINDEX_BUILD_GRANULE_ENABLED { - return index_type != DISKANN_IDENTIFIER; +fn native_option_applies(key: &str, index_type: &str) -> bool { + match key { + "dimension" + | "metric" + | "expected-vector-count" + | "target-recall" + | "deployment-profile" => true, + "nlist" => index_type != DISKANN_IDENTIFIER, + "use-opq" => index_type == IVF_PQ_IDENTIFIER, + "rq.bits" => index_type == IVF_RQ_IDENTIFIER, + "pq.m" | "pq.code-ratio" => matches!(index_type, IVF_PQ_IDENTIFIER | DISKANN_IDENTIFIER), + "max-bytes-per-vector" => true, + key if is_native_build_option(key) => build_option_applies(key, index_type), + key => { + index_type == DISKANN_IDENTIFIER + && DISKANN_OPTION_KEYS.iter().any(|(native, _)| key == *native) + } } - if is_allowed_native_key(key, index_type) { - return true; +} + +fn validate_index_type_option( + user_options: &HashMap, + expected_native: &str, +) -> crate::Result<()> { + for options in [user_options] { + if let Some(value) = options.get("index.type") { + let normalized = value.trim().to_ascii_lowercase().replace('-', "_"); + if normalized != expected_native { + return Err(crate::Error::ConfigInvalid { + message: format!( + "Option 'index.type' is '{}', but procedure index_type resolves to '{}'. \ + Remove 'index.type' from options or set it to '{}'.", + value, expected_native, expected_native + ), + }); + } + } } + Ok(()) +} - let index_prefix = format!("{index_type}."); - if let Some(suffix) = key.strip_prefix(&index_prefix) { - return is_allowed_paimon_suffix(suffix, index_type); +fn native_option_key(key: &str) -> Option<&str> { + match key { + "index.dimension" | "dimension" => Some("dimension"), + "distance.metric" | "metric" => Some("metric"), + "pq.use-opq" | "use-opq" => Some("use-opq"), + "nlist" + | "expected-vector-count" + | "ivf.coarse-assignment" + | "ivf.pq-encoding" + | "ivf.train.max-points-per-centroid" + | "pq.train.max-points-per-centroid" + | "pq.m" + | "pq.code-ratio" + | "pq.bits" + | "rq.bits" + | "target-recall" + | "max-bytes-per-vector" + | "deployment-profile" => Some(key), + _ => DISKANN_OPTION_KEYS + .iter() + .find_map(|(native, suffix)| (key == *native || key == *suffix).then_some(*native)), } +} - let field_prefix = format!("fields.{field_name}."); - if let Some(suffix) = key.strip_prefix(&field_prefix) { - return is_allowed_paimon_suffix(suffix, index_type); +fn is_native_build_option(key: &str) -> bool { + matches!( + key, + "ivf.coarse-assignment" + | "ivf.pq-encoding" + | "ivf.train.max-points-per-centroid" + | "pq.train.max-points-per-centroid" + ) +} + +fn build_option_applies(key: &str, index_type: &str) -> bool { + match key { + "ivf.pq-encoding" => index_type == IVF_PQ_IDENTIFIER, + "pq.train.max-points-per-centroid" => { + matches!(index_type, IVF_PQ_IDENTIFIER | DISKANN_IDENTIFIER) + } + _ => index_type != DISKANN_IDENTIFIER, } +} - false +fn collect_native_options( + result: &mut HashMap, + options: &HashMap, + index_type: &str, + field_name: &str, + validate: bool, +) -> crate::Result<()> { + for prefix in [ + String::new(), + format!("{index_type}."), + format!("fields.{field_name}."), + ] { + for (key, value) in options { + let Some(suffix) = key.strip_prefix(&prefix) else { + continue; + }; + let Some(native_key) = native_option_key(suffix) else { + continue; + }; + // Schema options follow Java. Explicit Rust/SQL builder options also + // accept the documented native aliases. + if prefix.is_empty() + && !validate + && (key != native_key || !is_native_build_option(native_key)) + { + continue; + } + if is_native_build_option(native_key) && !build_option_applies(native_key, index_type) { + if validate { + return Err(crate::Error::ConfigInvalid { + message: format!( + "Option '{key}' is not supported for index type '{index_type}'." + ), + }); + } + continue; + } + result.insert(native_key.into(), value.clone()); + } + } + Ok(()) } fn resolve_granule_build_enabled( @@ -352,56 +416,6 @@ fn resolve_granule_build_enabled( } } -fn is_allowed_native_key(key: &str, index_type: &str) -> bool { - match key { - "dimension" | "metric" => true, - "nlist" => index_type != DISKANN_IDENTIFIER, - "use-opq" => index_type == IVF_PQ_IDENTIFIER, - "rq.bits" => index_type == IVF_RQ_IDENTIFIER, - "max-bytes-per-vector" => { - matches!(index_type, IVF_RQ_IDENTIFIER | DISKANN_IDENTIFIER) - } - "pq.m" if index_type == IVF_PQ_IDENTIFIER => true, - "ivf.coarse-assignment" | "ivf.train.max-points-per-centroid" => { - index_type != DISKANN_IDENTIFIER - } - "ivf.pq-encoding" => index_type == IVF_PQ_IDENTIFIER, - "pq.train.max-points-per-centroid" => { - matches!(index_type, IVF_PQ_IDENTIFIER | DISKANN_IDENTIFIER) - } - _ => { - index_type == DISKANN_IDENTIFIER - && DISKANN_OPTION_KEYS - .iter() - .any(|(native_key, _)| *native_key == key) - } - } -} - -fn is_allowed_paimon_suffix(suffix: &str, index_type: &str) -> bool { - match suffix { - "dimension" | "distance.metric" => true, - "nlist" => index_type != DISKANN_IDENTIFIER, - "train.sample-ratio" => true, - "pq.use-opq" => index_type == IVF_PQ_IDENTIFIER, - "rq.bits" => index_type == IVF_RQ_IDENTIFIER, - "max-bytes-per-vector" => { - matches!(index_type, IVF_RQ_IDENTIFIER | DISKANN_IDENTIFIER) - } - "pq.m" if index_type == IVF_PQ_IDENTIFIER => true, - "ivf.coarse-assignment" - | "ivf.pq-encoding" - | "ivf.train.max-points-per-centroid" - | "pq.train.max-points-per-centroid" => is_allowed_native_key(suffix, index_type), - _ => { - index_type == DISKANN_IDENTIFIER - && DISKANN_OPTION_KEYS - .iter() - .any(|(_, paimon_suffix)| *paimon_suffix == suffix) - } - } -} - fn resolve_train_sample_ratio( table_options: &HashMap, user_options: &HashMap, @@ -428,6 +442,7 @@ fn resolve_train_sample_ratio( return Ok(DEFAULT_TRAIN_SAMPLE_RATIO); }; let ratio = value + .trim() .parse::() .map_err(|_| crate::Error::ConfigInvalid { message: format!("Invalid vindex train.sample-ratio: '{value}'"), @@ -442,69 +457,6 @@ fn resolve_train_sample_ratio( Ok(ratio) } -fn resolve_dimension( - table_options: &HashMap, - user_options: &HashMap, - index_type: &str, - field: &DataField, -) -> crate::Result { - if let DataType::Vector(vector) = field.data_type() { - return Ok(vector.length().to_string()); - } - - Ok(option_value( - table_options, - user_options, - field.name(), - index_type, - "dimension", - "dimension", - DEFAULT_DIMENSION, - )) -} - -fn option_value( - table_options: &HashMap, - user_options: &HashMap, - field_name: &str, - index_type: &str, - native_key: &str, - paimon_suffix: &str, - default_value: &str, -) -> String { - optional_value( - table_options, - user_options, - field_name, - index_type, - native_key, - paimon_suffix, - ) - .unwrap_or_else(|| default_value.to_string()) -} - -fn optional_value( - table_options: &HashMap, - user_options: &HashMap, - field_name: &str, - index_type: &str, - native_key: &str, - paimon_suffix: &str, -) -> Option { - for options in [user_options, table_options] { - for key in [ - format!("fields.{field_name}.{paimon_suffix}"), - format!("{index_type}.{paimon_suffix}"), - native_key.to_string(), - ] { - if let Some(value) = options.get(&key) { - return Some(value.clone()); - } - } - } - None -} - fn normalize_metric(metric: &str) -> String { metric.trim().to_ascii_lowercase().replace('-', "_") } @@ -540,7 +492,7 @@ mod tests { ) .unwrap(); let n = 512; - let dimension = options.config.dimension(); + let dimension = options.training_config(512).unwrap().dimension(); let data = (0..n * dimension) .map(|offset| { let row = offset / dimension; @@ -548,7 +500,8 @@ mod tests { (row % 4) as f32 * 20.0 + column as f32 * 0.01 + row as f32 * 0.0001 }) .collect::>(); - let training = VectorIndexTrainer::train(options.config, &data, n).unwrap(); + let training = + VectorIndexTrainer::train(options.training_config(n).unwrap(), &data, n).unwrap(); let mut writer = VectorIndexWriter::new(training); writer .add_vectors(&(0..n as i64).collect::>(), &data, n) @@ -656,7 +609,7 @@ mod tests { .map(String::as_str), Some("canonical") ); - let resolved = options.config.resolved(); + let resolved = options.training_config(512).unwrap().resolved(); assert!(!resolved.use_approximate_coarse_assignment); assert!(resolved.canonical_pq_encoding); assert_eq!(resolved.ivf_train_max_points_per_centroid, Some(32)); @@ -673,7 +626,11 @@ mod tests { ) .unwrap(); assert_eq!( - diskann.config.resolved().pq_train_max_points_per_centroid, + diskann + .training_config(512) + .unwrap() + .resolved() + .pq_train_max_points_per_centroid, Some(16) ); } @@ -698,7 +655,10 @@ mod tests { options.native_options.get("rq.bits").map(String::as_str), Some("3") ); - assert_eq!(options.config.resolved().rq_bits, Some(3)); + assert_eq!( + options.training_config(512).unwrap().resolved().rq_bits, + Some(3) + ); let defaults = VindexVectorIndexOptions::new( &HashMap::new(), @@ -707,7 +667,10 @@ mod tests { &array_float_field(), ) .unwrap(); - assert_eq!(defaults.config.resolved().rq_bits, Some(4)); + assert_eq!( + defaults.training_config(512).unwrap().resolved().rq_bits, + Some(4) + ); assert!(!defaults.native_options.contains_key("rq.bits")); let capacity_goal = HashMap::from([ @@ -722,7 +685,10 @@ mod tests { &array_float_field(), ) .unwrap(); - assert_eq!(inferred.config.resolved().rq_bits, Some(3)); + assert_eq!( + inferred.training_config(512).unwrap().resolved().rq_bits, + Some(3) + ); assert_eq!( inferred .native_options @@ -739,6 +705,8 @@ mod tests { IVF_RQ_IDENTIFIER, &array_float_field(), ) + .unwrap() + .training_config(512) .expect_err("rq.bits outside 1..=8 must be rejected"); assert!(error.to_string().contains("rq.bits")); } @@ -790,7 +758,7 @@ mod tests { &array_float_field(), ) .unwrap(); - let resolved = options.config.resolved(); + let resolved = options.training_config(512).unwrap().resolved(); assert_eq!(resolved.pq_m, Some(16)); assert_eq!(resolved.pq_bits, Some(4)); @@ -1083,6 +1051,8 @@ mod tests { IVF_PQ_IDENTIFIER, &array_float_field(), ) + .unwrap() + .training_config(512) .expect_err("invalid native config should be rejected"); assert!( @@ -1091,23 +1061,29 @@ mod tests { } #[test] - fn test_vindex_options_reject_unknown_user_options() { - let table_options = HashMap::new(); - let user_options = HashMap::from([ - ("ivf-flat.dimension".to_string(), "8".to_string()), - ("ivf-flat.nlsit".to_string(), "4".to_string()), - ]); - - let err = VindexVectorIndexOptions::new( - &table_options, - &user_options, + fn test_vindex_table_options_ignore_unprefixed_keys_like_java() { + let options = VindexVectorIndexOptions::new( + &HashMap::from([ + ("dimension".into(), "7".into()), + ("metric".into(), "invalid".into()), + ("index.type".into(), "diskann".into()), + ("nlist".into(), "0".into()), + ("ivf-flat.nlsit".into(), "4".into()), + ("global-index.build.parallelism".into(), "2".into()), + ]), + &HashMap::new(), IVF_FLAT_IDENTIFIER, &array_float_field(), ) - .expect_err("unknown user option should be rejected"); - - assert!( - matches!(err, crate::Error::ConfigInvalid { message } if message.contains("ivf-flat.nlsit")) + .unwrap(); + assert_eq!(options.dimension(), 128); + assert_eq!(options.native_options["metric"], "inner_product"); + assert_eq!(options.native_options["index.type"], "ivf_flat"); + assert!(!options.native_options.contains_key("nlist")); + assert!(options.needs_vector_count()); + assert_eq!( + options.training_config(1000).unwrap().nlist(), + paimon_vindex_core::autotune::infer_ivf_nlist(1000).unwrap() ); } @@ -1115,14 +1091,16 @@ mod tests { fn test_vindex_options_reject_non_applicable_user_options() { for (index_type, key) in [ (IVF_FLAT_IDENTIFIER, "ivf-flat.pq.m"), + (IVF_FLAT_IDENTIFIER, "ivf-pq.dimension"), + (IVF_FLAT_IDENTIFIER, "ivf-flat.nlsit"), + (IVF_FLAT_IDENTIFIER, "diskann.max-degree"), + (DISKANN_IDENTIFIER, "diskann.nlist"), + (DISKANN_IDENTIFIER, VINDEX_BUILD_GRANULE_ENABLED), (IVF_FLAT_IDENTIFIER, "ivf-flat.ivf.pq-encoding"), ( IVF_FLAT_IDENTIFIER, "ivf-flat.pq.train.max-points-per-centroid", ), - (IVF_FLAT_IDENTIFIER, "diskann.max-degree"), - (DISKANN_IDENTIFIER, "diskann.nlist"), - (DISKANN_IDENTIFIER, VINDEX_BUILD_GRANULE_ENABLED), (DISKANN_IDENTIFIER, "diskann.ivf.coarse-assignment"), ( DISKANN_IDENTIFIER, @@ -1164,7 +1142,7 @@ mod tests { ); assert_eq!( options.native_options.get("nlist").map(String::as_str), - Some("256") + None ); } diff --git a/docs/src/sql.md b/docs/src/sql.md index 6de93f0ea..bd8d4a8fc 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1342,13 +1342,17 @@ Supported vindex options: |---|---:|---|---| | `.dimension` | `128` | all vindex types | Vector dimension for `ARRAY` columns. Existing `VECTOR` columns use `N` from the type. | | `.distance.metric` | `inner_product` | all vindex types | Distance metric: `inner_product`, `cosine`, or `l2`. | -| `.nlist` | `256` | all IVF types | Number of IVF lists. DiskANN rejects this option. | +| `.nlist` | automatic | all IVF types | Number of IVF lists. When omitted or `auto`, infer it from `expected-vector-count`. DiskANN rejects this option. | +| `.expected-vector-count` | unset, or inferred | all vindex types | Expected number of vectors used for automatic sizing. When omitted, an automatic IVF build counts the shard's non-NULL vectors before training. | +| `.target-recall` | unset | all vindex types | Optional tuning objective in `[0, 1]`; IVF-PQ enables automatic OPQ at values of at least `0.9`. | +| `.max-bytes-per-vector` | unset | all vindex types | Optional persisted-size tuning objective; PQ and RQ can use it to infer encoding width. | | `.ivf.coarse-assignment` | `auto` | all IVF types | Build-time list assignment mode: `auto` uses Vamana for large centroid matrices; `exact` always uses exact assignment. | | `.ivf.train.max-points-per-centroid` | `256` | all IVF types | Positive limit on coarse K-means training data: at most `nlist × value` vectors. | -| `.train.sample-ratio` or `fields..train.sample-ratio` | `1.0` | all vindex types | Fraction of shard rows selected evenly for training. Must be in `(0, 1]`; all rows are still added to the index. The field-specific option takes precedence. | -| `vindex.build.granule.enabled` | `true` | all IVF types | Use the granule build pipeline. Set to `false` to use the full-spill build path for future index builds. | -| `.pq.m` | `16` | `ivf-pq` | Number of product-quantization sub-vectors. The dimension must be divisible by this value. | -| `.pq.use-opq` | `false` | `ivf-pq` | Whether to enable OPQ before PQ encoding. | +| `.train.sample-ratio` or `fields..train.sample-ratio` | `1.0` | all vindex types | Fraction of non-NULL vectors selected for training. Must be in `(0, 1]`; all non-NULL vectors are still added to the index. The field-specific option takes precedence. | +| `vindex.build.granule.enabled` | `true` | all IVF types | Use the granule build pipeline when `nlist` or `expected-vector-count` is explicit. NULL vectors switch to full spill. Set to `false` to always use full spill. | +| `.pq.m` | automatic | `ivf-pq` | Number of product-quantization sub-vectors, inferred from dimension and the code budget when omitted. The dimension must be divisible by an explicit value. | +| `.pq.code-ratio` | `0.0625` | `ivf-pq` | Ratio of PQ-code bytes to raw `f32` vector bytes, in `(0, 0.25]`. An explicit `pq.m` overrides it. | +| `.pq.use-opq` | inferred | `ivf-pq` | Whether to enable OPQ before PQ encoding. When omitted or `auto`, enable it if `target-recall` is at least `0.9`; otherwise disable it. | | `ivf-pq.ivf.pq-encoding` | `auto` | `ivf-pq` | Build-time PQ encoding mode: `auto` selects an accelerated backend when supported; `canonical` uses the canonical encoder. | | `.pq.train.max-points-per-centroid` | `256` | `ivf-pq`, `diskann` | Positive limit on PQ training data: at most `2^pq.bits × value` vectors per subquantizer. | | `ivf-rq.rq.bits` | `4`, or inferred | `ivf-rq` | Residual-quantization width in the range `1` to `8`. When omitted, `ivf-rq.max-bytes-per-vector` can select it. | @@ -1384,6 +1388,11 @@ vindex aliases are also accepted in the `options` string: `dimension`, `metric`, `diskann.*` build keys listed in the table. Build options for another index family are rejected rather than ignored. +NULL vector rows are skipped, while index manifest `row_count` and the global +row range still include them. Non-NULL vectors retain their original relative +row IDs. An all-NULL shard produces no index file. Full-spill training samples +the compacted non-NULL vector stream at evenly spaced positions, matching Java. + Inspect committed index files with the `$table_indexes` system table: ```sql