diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 144c6dbfb..a81e2e627 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -416,7 +416,7 @@ mod tests { use arrow_array::record_batch; use arrow_array::{ Array, ArrayRef, BinaryArray, BooleanArray, FixedSizeListArray, Float32Array, Int32Array, - LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StructArray, + LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StructArray, UInt64Array, }; use arrow_data::ArrayDataBuilder; use arrow_schema::{DataType, Field, Schema}; @@ -436,6 +436,55 @@ mod tests { use crate::table::optimize::{CompactionOptions, OptimizeAction}; use lance_index::scalar::FullTextSearchQuery; + struct OrderPreservingShuffler; + + struct SinglePartitionReader { + batches: Vec, + } + + #[async_trait::async_trait] + impl lance_index::vector::v3::shuffler::Shuffler for OrderPreservingShuffler { + async fn shuffle( + &self, + data: Box, + ) -> lance_core::Result> { + Ok(Box::new(SinglePartitionReader { + batches: data.try_collect().await?, + })) + } + } + + #[async_trait::async_trait] + impl lance_index::vector::v3::shuffler::ShuffleReader for SinglePartitionReader { + async fn read_partition( + &self, + partition_id: usize, + ) -> lance_core::Result< + Option>, + > { + if partition_id != 0 || self.batches.is_empty() { + return Ok(None); + } + let schema = self.batches[0].schema(); + let stream = futures::stream::iter(self.batches.clone().into_iter().map(Ok)); + Ok(Some(Box::new( + lance_io::stream::RecordBatchStreamAdapter::new(schema, stream), + ))) + } + + fn partition_size(&self, partition_id: usize) -> lance_core::Result { + Ok(if partition_id == 0 { + self.batches.iter().map(RecordBatch::num_rows).sum() + } else { + 0 + }) + } + + fn total_loss(&self) -> Option { + Some(0.0) + } + } + fn create_fixed_size_list( values: T, list_size: i32, @@ -840,6 +889,135 @@ mod tests { assert_eq!(stats.distance_type, Some(crate::DistanceType::L2)); } + #[tokio::test] + async fn test_ivf_hnsw_sq_merge_preserves_shuffled_row_order() { + use lance::index::vector::builder::IvfIndexBuilder; + use lance_core::{ROW_ID, ROW_ID_FIELD, cache::LanceCache}; + use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; + use lance_file::reader::{FileReader, FileReaderOptions}; + use lance_index::INDEX_AUXILIARY_FILE_NAME; + use lance_index::vector::hnsw::{HNSW, builder::HnswBuildParams}; + use lance_index::vector::ivf::IvfBuildParams; + use lance_index::vector::sq::{ScalarQuantizer, builder::SQBuildParams}; + use lance_io::ReadBatchParams; + use lance_io::object_store::ObjectStore; + use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; + use lance_io::stream::RecordBatchStreamAdapter; + use lance_io::utils::CachedFileSize; + + const NUM_ROWS: usize = 64; + const DIMENSION: i32 = 16; + + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let vector_field = Field::new( + "embeddings", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIMENSION, + ), + false, + ); + let values = Float32Array::from_iter_values( + (0..NUM_ROWS * DIMENSION as usize).map(|value| (value % 97) as f32), + ); + let vectors = Arc::new(create_fixed_size_list(values, DIMENSION).unwrap()); + let table_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![vector_field.clone()])), + vec![vectors.clone()], + ) + .unwrap(); + let table = conn + .create_table("test", table_batch) + .execute() + .await + .unwrap(); + + let dataset_guard = table.as_native().unwrap().dataset.get().await.unwrap(); + let dataset = (*dataset_guard).clone(); + drop(dataset_guard); + + // Descending row IDs are a small sentinel for the removed merge-time + // sort. The old sort reordered this batch and used FixedSizeList::take, + // which overflowed for the large partitions reported in #3126. + let expected_row_ids = (0..NUM_ROWS as u64).rev().collect::>(); + let input_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ROW_ID_FIELD.clone(), vector_field])), + vec![ + Arc::new(UInt64Array::from(expected_row_ids.clone())), + vectors, + ], + ) + .unwrap(); + let input = RecordBatchStreamAdapter::new( + input_batch.schema(), + futures::stream::iter(vec![Ok(input_batch)]), + ); + + let index_dir = dataset.indices_dir().join(uuid::Uuid::new_v4().to_string()); + let mut builder = IvfIndexBuilder::::new( + dataset, + "embeddings".to_string(), + index_dir.clone(), + lance_linalg::distance::DistanceType::L2, + Box::new(OrderPreservingShuffler), + Some(IvfBuildParams::new(1)), + Some(SQBuildParams::default()), + HnswBuildParams::default(), + None, + ) + .unwrap(); + builder.shuffle_data_input(Some(input)); + builder.build().await.unwrap(); + + let object_store = Arc::new(ObjectStore::local()); + let scheduler = ScanScheduler::new(object_store, SchedulerConfig::default_for_testing()); + let auxiliary_path = index_dir.join(INDEX_AUXILIARY_FILE_NAME); + let reader = FileReader::try_open( + scheduler + .open_file(&auxiliary_path, &CachedFileSize::unknown()) + .await + .unwrap(), + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + let batches = reader + .read_stream( + ReadBatchParams::RangeFull, + u32::MAX, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let stored_row_ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + + assert_eq!(stored_row_ids, expected_row_ids); + } + #[tokio::test] async fn test_create_index_ivf_hnsw_pq() { use std::iter::repeat_with;