From e3dce2cbad45c14422a077dfb34ee68050173827 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:27:27 +0800 Subject: [PATCH] fix(rust): prevent oversized SQ remap panics (#3861) ## Summary - preflight IVF-SQ and IVF-HNSW-SQ storage segments touched by compaction - preserve unknown legacy fragment coverage as conservatively affected - cover both fixed-size vector and multivector SQ storage dimensions - reject only affected partitions whose fixed-size-list child exceeds Arrow's `u32` take-index limit - return an actionable `InvalidInput` error before Arrow can panic ## Root cause Arrow 58 builds fixed-size-list child indices with `u32`. For a large SQ partition, `partition_rows * vector_dimension` can exceed `u32::MAX`; the child range then wraps inside `take_fixed_size_list`, producing the reported `ArrayData::slice` assertion during index remapping. The preflight now reads raw segment metadata, mirrors Lance's fragment-bitmap selection (including `None` as affected/unknown), opens only affected vector segments, and checks their actual storage quantizer and partition sizes. It derives the stored dimension for both `FixedSizeList` vectors and flattened `List` multivectors. Stable-row-ID datasets, deferred remapping, custom remappers, unaffected segments, non-SQ indexes, safe partitions, and no-op compactions are unchanged. Because the upstream Arrow kernel still has the limit, an affected index must be recreated with default IVF partitioning (or enough partitions to stay below the reported maximum) before compaction. This changes the failure from a process panic to a targeted recovery error; it does not silently alter index structure. ## Validation - `cargo fmt --all -- --check` - `cargo test --quiet --features aws,remote -p lancedb table::optimize::tests --lib` - `cargo test --quiet --features remote -p lancedb table::create_index::tests::test_create_index_ivf_hnsw_sq` - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples -- -A clippy::redundant-pub-crate` (the allowance is for pre-existing Rust 1.97 warnings in remote-table code) Fixes #2866 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- Cargo.lock | 1 + Cargo.toml | 1 + docs/web/maintenance/compaction.mdx | 18 ++ rust/lancedb/Cargo.toml | 1 + rust/lancedb/src/table/optimize.rs | 351 +++++++++++++++++++++++++++- 5 files changed, 363 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bbd41c0b1..e4a2eaada 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5789,6 +5789,7 @@ dependencies = [ "tokenizers", "tokio", "tonic", + "tracing", "url", "urlencoding", "uuid", diff --git a/Cargo.toml b/Cargo.toml index 3027ee065..22dec8c0c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,6 +82,7 @@ serde_json = { version = "1", features = ["preserve_order"] } tempfile = "3.5.0" tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] } tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] } +tracing = "0.1" uuid = { version = "1.7.0", features = ["v4", "v7"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } diff --git a/docs/web/maintenance/compaction.mdx b/docs/web/maintenance/compaction.mdx index 5d2a9c42f..8b0a8c86f 100644 --- a/docs/web/maintenance/compaction.mdx +++ b/docs/web/maintenance/compaction.mdx @@ -23,6 +23,24 @@ three; a deployment handles them in the background and does not support Run optimization as regular table maintenance, particularly after many small writes or a large ingest. +{/* +Draft for the OSS compaction SQ partition check release. Keep this section +unpublished until that behavior ships. + +### Recover from an oversized SQ partition {#oversized-sq-partition} + +Compacting a table with an IVF scalar-quantized (SQ) vector index can stop +before rewriting files if an affected partition is too large to remap. The +error identifies the index, segment, partition size, and vector dimension +instead of letting the remap fail at Arrow's fixed-size-list `u32` +child-array limit. + +Rebuild the affected vector index with enough IVF partitions to keep each +partition at or below `floor(u32::MAX / dimension)` vectors, then retry the +optimize. Rebuilding with the default IVF partitioning can also choose enough +partitions. +*/} + ## Optimize a table {#what-it-does} ```python diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index c18db11fe..4a2e2318e 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -55,6 +55,7 @@ moka = { workspace = true } pin-project = { workspace = true } tokio = { workspace = true, features = ["io-util", "net", "time"] } log.workspace = true +tracing.workspace = true async-trait = { workspace = true } bytes = { workspace = true } futures.workspace = true diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index ee86e7989..657385762 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -6,20 +6,183 @@ //! This module contains the implementation of optimization operations that help //! maintain good performance for LanceDB tables. -use std::sync::Arc; +use std::{ + collections::HashSet, + sync::{Arc, Mutex}, +}; +use arrow_schema::DataType; use chrono::{DateTime, Utc}; use lance::dataset::cleanup::{CleanupPolicyBuilder, RemovalStats}; -use lance::dataset::optimize::{CompactionMetrics, IndexRemapperOptions, compact_files}; -use lance::index::DatasetIndexExt; +use lance::dataset::optimize::{ + CompactionMetrics, CompactionPlan, CompactionPlanner, IndexRemapperOptions, + compact_files_with_planner, plan_compaction, +}; +use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; +use lance_core::utils::tracing::{DATASET_COMPACTING_EVENT, TRACE_DATASET_EVENTS}; +use lance_index::IndexType; +use lance_index::metrics::NoOpMetricsCollector; use lance_index::optimize::OptimizeOptions; -use log::info; +use log::{debug, info}; +use super::NativeTable; +use crate::error::{Error, Result}; pub use chrono::Duration; pub use lance::dataset::optimize::{CompactionMode, CompactionOptions, IndexRemapMode}; -use super::NativeTable; -use crate::error::Result; +const MAX_ARROW_FIXED_SIZE_LIST_CHILD_INDEX: u64 = u32::MAX as u64; + +struct PrecomputedCompactionPlanner { + plan: Mutex>, +} + +impl PrecomputedCompactionPlanner { + fn new(plan: CompactionPlan) -> Self { + Self { + plan: Mutex::new(Some(plan)), + } + } +} + +#[async_trait::async_trait] +impl CompactionPlanner for PrecomputedCompactionPlanner { + async fn plan(&self, dataset: &lance::Dataset) -> lance::Result { + let plan = self + .plan + .lock() + .map_err(|_| lance_core::Error::internal("precomputed compaction plan lock poisoned"))? + .take() + .ok_or_else(|| { + lance_core::Error::internal("precomputed compaction plan was already consumed") + })?; + debug_assert_eq!(dataset.manifest().version, plan.read_version()); + Ok(plan) + } +} + +fn sq_vector_dimension(data_type: &DataType) -> Option { + match data_type { + DataType::FixedSizeList(_, dimension) => Some(*dimension), + DataType::List(field) => match field.data_type() { + DataType::FixedSizeList(_, dimension) => Some(*dimension), + _ => None, + }, + _ => None, + } +} + +fn segment_touches_fragments( + fragment_ids: Option>, + affected_fragments: &HashSet, +) -> bool { + fragment_ids.is_none_or(|fragment_ids| { + fragment_ids + .into_iter() + .any(|fragment| affected_fragments.contains(&(fragment as u64))) + }) +} + +fn oversized_sq_partition( + partition_sizes: impl IntoIterator, + dimension: u64, +) -> Option { + partition_sizes.into_iter().find(|partition_size| { + partition_size + .checked_mul(dimension) + .is_none_or(|child_len| child_len > MAX_ARROW_FIXED_SIZE_LIST_CHILD_INDEX) + }) +} + +async fn validate_sq_index_remapping( + dataset: &lance::Dataset, + options: &CompactionOptions, + has_custom_remapper: bool, +) -> Result { + let plan = plan_compaction(dataset, options).await?; + if options.defer_index_remap || has_custom_remapper || dataset.manifest().uses_stable_row_ids() + { + return Ok(plan); + } + + if plan.tasks.is_empty() { + return Ok(plan); + } + let affected_fragments: HashSet = plan + .tasks + .iter() + .flat_map(|task| task.fragments.iter().map(|fragment| fragment.id)) + .collect(); + + for segment in dataset.load_indices().await?.iter() { + if !segment_touches_fragments( + segment + .fragment_bitmap + .as_ref() + .map(|fragment_bitmap| fragment_bitmap.iter()), + &affected_fragments, + ) { + continue; + } + let Some(field_id) = segment.fields.first() else { + continue; + }; + let Some(field) = dataset.schema().field_by_id(*field_id) else { + continue; + }; + let data_type = field.data_type(); + let Some(dimension) = sq_vector_dimension(&data_type) else { + continue; + }; + let dimension = u64::try_from(dimension).map_err(|_| Error::InvalidInput { + message: format!( + "SQ index '{}' has an invalid vector dimension of {}", + segment.name, dimension + ), + })?; + + let field_path = dataset.schema().field_path(*field_id)?; + let vector_index = match dataset + .open_vector_index(&field_path, &segment.uuid, &NoOpMetricsCollector) + .await + { + Ok(vector_index) => vector_index, + Err(error) => { + // The default remapper also treats an index it cannot open as + // unusable and drops it, so it cannot reach the SQ take kernel. + debug!( + "Skipping remap preflight for index segment {} because it could not be opened: {}", + segment.uuid, error + ); + continue; + } + }; + if !matches!( + vector_index.index_type(), + IndexType::IvfSq | IndexType::IvfHnswSq + ) { + continue; + } + + let partition_sizes = (0..vector_index.total_partitions()) + .map(|partition_id| vector_index.partition_size(partition_id) as u64); + if let Some(partition_size) = oversized_sq_partition(partition_sizes, dimension) { + let max_partition_size = MAX_ARROW_FIXED_SIZE_LIST_CHILD_INDEX / dimension; + return Err(Error::InvalidInput { + message: format!( + "Cannot compact table because SQ index '{}' segment {} has a partition with {} vectors of dimension {}. Arrow's fixed-size-list take kernel cannot remap partitions whose child array exceeds {} values. Recreate the index with default IVF partitioning, or enough partitions to keep each partition at or below {} vectors, before compacting.", + segment.name, + segment.uuid, + partition_size, + dimension, + MAX_ARROW_FIXED_SIZE_LIST_CHILD_INDEX, + max_partition_size, + ), + }); + } + } + + Ok(plan) +} /// Optimize the dataset. /// @@ -193,7 +356,7 @@ pub(crate) async fn cleanup_old_versions_before( /// This can be run after making several small appends to optimize the table /// for faster reads. /// -/// This calls into [lance::dataset::optimize::compact_files]. +/// This calls into [lance::dataset::optimize::compact_files_with_planner]. pub(crate) async fn compact_files_impl( table: &NativeTable, options: CompactionOptions, @@ -201,7 +364,10 @@ pub(crate) async fn compact_files_impl( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); - let metrics = compact_files(&mut dataset, options, remap_options).await?; + tracing::info!(target: TRACE_DATASET_EVENTS, event=DATASET_COMPACTING_EVENT, uri = dataset.uri()); + let plan = validate_sq_index_remapping(&dataset, &options, remap_options.is_some()).await?; + let planner = PrecomputedCompactionPlanner::new(plan); + let metrics = compact_files_with_planner(&mut dataset, remap_options, &planner).await?; table.dataset.update(dataset); Ok(metrics) } @@ -263,13 +429,14 @@ pub(crate) async fn execute_optimize( #[cfg(test)] mod tests { + use super::{oversized_sq_partition, segment_touches_fragments, sq_vector_dimension}; use arrow_array::{ Array, FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, }; use arrow_schema::{DataType, Field, Schema}; use lance_arrow::FixedSizeListArrayExt; use rstest::rstest; - use std::sync::Arc; + use std::{collections::HashSet, sync::Arc}; use crate::connect; use crate::database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS; @@ -279,6 +446,172 @@ mod tests { use crate::table::{CompactionOptions, OptimizeAction, OptimizeStats}; use futures::TryStreamExt; + /// Regression test for https://github.com/lancedb/lancedb/issues/2866. + #[test] + fn test_detect_oversized_sq_partition() { + const DIMENSION: u64 = 4095; + let safe_size = u32::MAX as u64 / DIMENSION; + + assert_eq!( + oversized_sq_partition([safe_size, safe_size + 1], DIMENSION), + Some(safe_size + 1) + ); + assert_eq!( + oversized_sq_partition([safe_size / 2, safe_size / 2 + 1], DIMENSION), + None + ); + } + + #[test] + fn test_segment_touches_fragments_includes_unknown_coverage() { + let affected_fragments = HashSet::from([7]); + + assert!(segment_touches_fragments( + Some([7_u32]), + &affected_fragments + )); + assert!(!segment_touches_fragments( + Some([8_u32]), + &affected_fragments + )); + assert!(segment_touches_fragments( + None::<[u32; 0]>, + &affected_fragments + )); + } + + #[test] + fn test_sq_vector_dimension_includes_multivectors() { + const DIMENSION: i32 = 4095; + let vector = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, false)), + DIMENSION, + ); + let multivector = DataType::List(Arc::new(Field::new("item", vector.clone(), false))); + + assert_eq!(sq_vector_dimension(&vector), Some(DIMENSION)); + assert_eq!(sq_vector_dimension(&multivector), Some(DIMENSION)); + assert_eq!(sq_vector_dimension(&DataType::Float32), None); + } + + #[tokio::test] + async fn test_compact_legacy_vector_index_with_unknown_fragment_coverage() { + use lance::Dataset; + use lance::index::vector::{IndexFileVersion, VectorIndexParams}; + use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_table::io::commit::write_manifest_file_to_path; + use object_store::ObjectStoreExt; + + const INDEX_NAME: &str = "legacy_ivf_pq"; + const ROWS: i32 = 128; + const DIMENSION: i32 = 8; + + let tmpdir = tempfile::tempdir().unwrap(); + let conn = connect(tmpdir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from_iter_values((0..ROWS).flat_map(|row| { + (0..DIMENSION).map(move |offset| (row * DIMENSION + offset) as f32) + })), + DIMENSION, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "vector", + vectors.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(vectors)]).unwrap(); + let table = conn + .create_table("test_legacy_vector_compact", batch.clone()) + .execute() + .await + .unwrap(); + table.add(batch.clone()).execute().await.unwrap(); + table.add(batch).execute().await.unwrap(); + let mut dataset = (*table.dataset().unwrap().get().await.unwrap()).clone(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 1, MetricType::L2, 10); + params.version(IndexFileVersion::Legacy); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + table.dataset().unwrap().update(dataset); + + let dataset = table.dataset().unwrap().get().await.unwrap(); + let mut segments = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(segments.len(), 1); + let original = segments.pop().unwrap(); + let field_path = dataset.schema().field_path(original.fields[0]).unwrap(); + let vector_index = dataset + .open_vector_index( + &field_path, + &original.uuid, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + assert_eq!(vector_index.index_type(), IndexType::IvfPq); + assert_eq!( + vector_index.statistics().unwrap()["index_file_version"], + "Legacy" + ); + let mut legacy = original.clone(); + legacy.fragment_bitmap = None; + let object_store = dataset.object_store(None).await.unwrap(); + let mut manifest = dataset.manifest().clone(); + manifest.index_section = None; + manifest.transaction_section = None; + let manifest_path = dataset.manifest_location().path.clone(); + object_store.inner.delete(&manifest_path).await.unwrap(); + write_manifest_file_to_path( + object_store.as_ref(), + &mut manifest, + Some(vec![legacy]), + &manifest_path, + None, + ) + .await + .unwrap(); + let legacy_dataset = Dataset::open(dataset.uri()).await.unwrap(); + table.dataset().unwrap().update(legacy_dataset); + + let dataset = table.dataset().unwrap().get().await.unwrap(); + let description_error = dataset + .describe_indices(None) + .await + .err() + .expect("legacy coverage should be unknown to index descriptions"); + assert!( + description_error + .to_string() + .contains("Fragment bitmap is required") + ); + + let stats = table + .optimize(OptimizeAction::Compact { + options: CompactionOptions { + target_rows_per_fragment: 1_000, + ..Default::default() + }, + remap_options: None, + }) + .await + .unwrap(); + assert!(stats.compaction.unwrap().fragments_removed > 0); + assert_eq!(table.count_rows(None).await.unwrap(), ROWS as usize * 3); + } + #[tokio::test] async fn test_optimize_compact_simple() { let conn = connect("memory://").execute().await.unwrap();