diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 8b92b9eddca..00eab6a1fa3 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -62,12 +62,26 @@ const BITMAP_PART_LOOKUP_SUFFIX: &str = "_bitmap_page_lookup.lance"; const EXPLICIT_SHARD_ID_TAG: u64 = 0; const IMPLICIT_FRAGMENT_ID_TAG: u64 = 1; -const MAX_BITMAP_ARRAY_LENGTH: usize = i32::MAX as usize - 1024 * 1024; // leave headroom +/// Maximum bytes a [`BitmapBatchWriter`] buffers before flushing a record +/// batch. +/// +/// Charged for the keys as well as the serialized bitmaps, so it limits this +/// writer's buffered state independently of how many keys the index has. A +/// flush temporarily makes another copy to build the Arrow arrays, and a single +/// entry can exceed the threshold because it is checked after serialization. +/// Memory held by the caller, input pipeline, caches, or merge state is outside +/// this writer limit. +/// +/// It also keeps both output columns far below the `i32` offset ceiling of +/// Arrow's `Binary`/`Utf8` layouts. The previous threshold was that ceiling +/// itself, which charged the bitmap column only: a high-cardinality column with +/// tiny bitmaps could overflow the keys column's offsets before it ever tripped. +const MAX_BUFFERED_BYTES: usize = 32 * 1024 * 1024; const MAX_ROWS_PER_CHUNK: usize = 2 * 1024; -// Smaller than MAX_ROWS_PER_CHUNK to bound the per-cursor in-memory batch -// footprint during a k-way merge (N cursors × chunk), while still amortising -// I/O over a reasonable number of rows per read. +// Smaller than MAX_ROWS_PER_CHUNK to cap the rows retained per cursor during a +// k-way merge (N cursors x chunk), while still amortising I/O over a reasonable +// number of rows per read. This is not a byte limit because bitmap sizes vary. const MERGE_ROWS_PER_CHUNK: usize = 512; const BITMAP_INDEX_VERSION: u32 = 0; @@ -524,26 +538,6 @@ impl BitmapIndex { pub(crate) fn value_type(&self) -> &DataType { &self.value_type } - - /// Loads the current bitmap index into an in-memory value-to-row-id map. - pub(crate) async fn load_bitmap_index_state( - &self, - ) -> Result> { - let mut state = HashMap::new(); - - for key in self.index_map.keys() { - let bitmap = self.load_bitmap(key, None).await?; - state.insert(key.0.clone(), (*bitmap).clone()); - } - - if !self.null_map.is_empty() { - let existing_null = new_null_array(&self.value_type, 1); - let existing_null = ScalarValue::try_from_array(existing_null.as_ref(), 0)?; - state.insert(existing_null, (*self.null_map).clone()); - } - - Ok(state) - } } impl DeepSizeOf for BitmapIndex { @@ -843,11 +837,10 @@ impl ScalarIndex for BitmapIndex { mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { - let state = self.load_bitmap_index_state().await?; - let remapped_state = BitmapIndexPlugin::remap_bitmap_state(state, mapping); - let file = - BitmapIndexPlugin::write_bitmap_index(remapped_state, dest_store, &self.value_type) - .await?; + let mut writer = + new_bitmap_batch_writer(dest_store, BITMAP_LOOKUP_NAME, &self.value_type).await?; + remap_index_map(self, mapping, &mut writer).await?; + let file = writer.finish().await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::BitmapIndexDetails::default()) @@ -891,13 +884,23 @@ impl ScalarIndex for BitmapIndex { } /// Buffers serialized (key, bitmap) pairs and flushes them as record batches -/// to the index file, respecting the MAX_BITMAP_ARRAY_LENGTH limit. -struct BitmapBatchWriter { +/// to the index file once they reach [`MAX_BUFFERED_BYTES`]. +pub(crate) struct BitmapBatchWriter { file: Box, keys: Vec, serialized: Vec>, bytes: usize, num_bitmaps: usize, + /// Flush threshold. A field rather than [`MAX_BUFFERED_BYTES`] directly only + /// so that tests can drive the multi-batch path without writing 32 MiB. + max_buffered_bytes: usize, + /// Record batches handed to `file` so far, so tests can assert the writer + /// actually flushed rather than buffering everything. + #[cfg(test)] + batches_written: usize, + /// Global-buffer keys and the buffer index each was written to. Recorded + /// as file metadata at finish so readers can find them. + buffer_indices: HashMap, } impl BitmapBatchWriter { @@ -908,17 +911,45 @@ impl BitmapBatchWriter { serialized: Vec::new(), bytes: 0, num_bitmaps: 0, + max_buffered_bytes: MAX_BUFFERED_BYTES, + #[cfg(test)] + batches_written: 0, + buffer_indices: HashMap::new(), } } + #[cfg(test)] + fn with_max_buffered_bytes(mut self, bytes: usize) -> Self { + self.max_buffered_bytes = bytes; + self + } + + #[cfg(test)] + pub(crate) fn batches_written(&self) -> usize { + self.batches_written + } + + /// Attach a global buffer to the file, recording its index under `key` so + /// that a reader can find it from the file metadata. + /// + /// Callable at any point: the underlying writer records the current offset + /// and writes the buffer immediately, so its position relative to the data + /// pages does not matter. Callers here do it first only to keep the metadata + /// setup in one place. + pub(crate) async fn add_global_buffer(&mut self, key: String, data: Bytes) -> Result<()> { + let buffer_idx = self.file.add_global_buffer(data).await?; + self.buffer_indices.insert(key, buffer_idx.to_string()); + Ok(()) + } + /// Serialize and buffer a single (key, bitmap) pair, flushing the current - /// batch to disk if adding it would exceed MAX_BITMAP_ARRAY_LENGTH. - async fn emit(&mut self, key: ScalarValue, bitmap: &RowAddrTreeMap) -> Result<()> { + /// batch to disk if adding it would exceed [`MAX_BUFFERED_BYTES`]. + pub(crate) async fn emit(&mut self, key: ScalarValue, bitmap: &RowAddrTreeMap) -> Result<()> { let mut buf = Vec::new(); bitmap.serialize_into(&mut buf).unwrap(); - let size = buf.len(); + let size = buf.len() + key.size(); - if self.bytes + size > MAX_BITMAP_ARRAY_LENGTH { + if self.bytes + size > self.max_buffered_bytes { self.flush().await?; } @@ -945,17 +976,22 @@ impl BitmapBatchWriter { let batch = BitmapIndexPlugin::get_batch_from_arrays(keys_array, bitmaps_array)?; self.file.write_record_batch(batch).await?; self.bytes = 0; + #[cfg(test)] + { + self.batches_written += 1; + } Ok(()) } - /// Flush any remaining data, write index statistics, and finalize the file. - async fn finish(mut self) -> Result { + /// Flush any remaining data, write index statistics and any global-buffer + /// indices, and finalize the file. + pub(crate) async fn finish(mut self) -> Result { self.flush().await?; let stats_json = serde_json::to_string(&BitmapStatistics { num_bitmaps: self.num_bitmaps, }) .map_err(|e| Error::internal(format!("failed to serialize bitmap statistics: {e}")))?; - let mut metadata = HashMap::new(); + let mut metadata = std::mem::take(&mut self.buffer_indices); metadata.insert(INDEX_STATS_METADATA_KEY.to_string(), stats_json); self.file.finish_with_metadata(metadata).await } @@ -1019,7 +1055,7 @@ fn deserialize_bitmap(bitmap_bytes: &[u8], file_name: &str) -> Result, total_rows: usize, @@ -1181,6 +1217,72 @@ async fn drain_same_key_bitmaps( Ok((merged_key.0, merged_bitmap)) } +/// Open a set of key-sorted bitmap files as merge cursors, seeded into a +/// min-heap on their first key, and confirm every file shares a value type. +/// +/// Returns `None` for the value type only when every file was empty. +pub(crate) async fn open_sorted_bitmap_cursors( + store: &dyn IndexStore, + files: &[String], +) -> Result<( + Vec, + BinaryHeap>, + Option, +)> { + let mut cursors = Vec::with_capacity(files.len()); + let mut heap = BinaryHeap::with_capacity(files.len()); + let mut value_type: Option = None; + + for file_name in files { + let reader = store.open_index_file(file_name).await?; + let file_value_type = reader.schema().fields[0].data_type().clone(); + if let Some(existing_type) = &value_type { + if existing_type != &file_value_type { + return Err(Error::invalid_input(format!( + "Bitmap shard {} has value type {:?}, expected {:?}", + file_name, file_value_type, existing_type + ))); + } + } else { + value_type = Some(file_value_type); + } + if let Some(cursor) = BitmapShardCursor::try_new(file_name.clone(), reader).await? { + let key = cursor.peek_key()?; + let shard_idx = cursors.len(); + cursors.push(cursor); + heap.push(Reverse(BitmapHeapItem { key, shard_idx })); + } + } + + Ok((cursors, heap, value_type)) +} + +/// Drain cursors opened by [`open_sorted_bitmap_cursors`] into `writer`, +/// emitting each key once in ascending order with the row sets of duplicate +/// keys unioned. +/// +/// The merge's working state is one row-bounded record batch per cursor plus the +/// bitmap currently being merged, independent of the total number of keys. This +/// does not include the output writer or other state retained by the caller, and +/// the cursor batches are not byte-bounded. +pub(crate) async fn drain_sorted_bitmap_cursors( + cursors: &mut [BitmapShardCursor], + heap: &mut BinaryHeap>, + writer: &mut BitmapBatchWriter, + progress: Option<(&dyn IndexBuildProgress, &str)>, +) -> Result<()> { + let mut merged_keys = 0u64; + while let Some(Reverse(item)) = heap.pop() { + let (key, merged_bitmap) = drain_same_key_bitmaps(cursors, heap, item).await?; + writer.emit(key, &merged_bitmap).await?; + merged_keys += 1; + if let Some((progress, stage)) = progress { + progress.stage_progress(stage, merged_keys).await?; + } + } + Ok(()) +} + async fn list_bitmap_shard_files( object_store: &ObjectStore, index_dir: &Path, @@ -1272,118 +1374,6 @@ impl BitmapIndexPlugin { Ok(RecordBatch::try_new(schema, columns)?) } - async fn write_bitmap_index( - state: HashMap, - index_store: &dyn IndexStore, - value_type: &DataType, - ) -> Result { - Self::write_bitmap_index_with_extras( - state, - index_store, - value_type, - HashMap::new(), - Vec::new(), - ) - .await - } - - /// Writes a bitmap index and attaches extra metadata and global buffers. - pub(crate) async fn write_bitmap_index_with_extras( - state: HashMap, - index_store: &dyn IndexStore, - value_type: &DataType, - mut metadata: HashMap, - global_buffers: Vec<(String, Bytes)>, - ) -> Result { - let num_bitmaps = state.len(); - let schema = Arc::new(Schema::new(vec![ - Field::new("keys", value_type.clone(), true), - Field::new("bitmaps", DataType::Binary, true), - ])); - - let mut bitmap_index_file = index_store - .new_index_file(BITMAP_LOOKUP_NAME, schema) - .await?; - - for (metadata_key, data) in global_buffers { - let buffer_idx = bitmap_index_file.add_global_buffer(data).await?; - metadata.insert(metadata_key, buffer_idx.to_string()); - } - - let mut cur_keys = Vec::new(); - let mut cur_bitmaps = Vec::new(); - let mut cur_bytes = 0; - - for (key, bitmap) in state.into_iter() { - let mut bytes = Vec::new(); - bitmap.serialize_into(&mut bytes).unwrap(); - let bitmap_size = bytes.len(); - - if cur_bytes + bitmap_size > MAX_BITMAP_ARRAY_LENGTH { - let keys_array = ScalarValue::iter_to_array(cur_keys.clone()).unwrap(); - let mut binary_builder = BinaryBuilder::new(); - for b in &cur_bitmaps { - binary_builder.append_value(b); - } - let bitmaps_array = Arc::new(binary_builder.finish()) as Arc; - - let record_batch = Self::get_batch_from_arrays(keys_array, bitmaps_array)?; - bitmap_index_file.write_record_batch(record_batch).await?; - - cur_keys.clear(); - cur_bitmaps.clear(); - cur_bytes = 0; - } - - cur_keys.push(key); - cur_bitmaps.push(bytes); - cur_bytes += bitmap_size; - } - - // Flush any remaining - if !cur_keys.is_empty() { - let keys_array = ScalarValue::iter_to_array(cur_keys).unwrap(); - let mut binary_builder = BinaryBuilder::new(); - for b in &cur_bitmaps { - binary_builder.append_value(b); - } - let bitmaps_array = Arc::new(binary_builder.finish()) as Arc; - - let record_batch = Self::get_batch_from_arrays(keys_array, bitmaps_array)?; - bitmap_index_file.write_record_batch(record_batch).await?; - } - - // Finish file with metadata that allows lightweight statistics reads - let stats_json = serde_json::to_string(&BitmapStatistics { num_bitmaps }) - .map_err(|e| Error::internal(format!("failed to serialize bitmap statistics: {e}")))?; - metadata.insert(INDEX_STATS_METADATA_KEY.to_string(), stats_json); - - bitmap_index_file.finish_with_metadata(metadata).await - } - - /// Builds bitmap index state from a `(value, row_id)` stream without writing it. - pub(crate) async fn build_bitmap_index_state( - mut data_source: SendableRecordBatchStream, - mut state: HashMap, - ) -> Result<(HashMap, DataType)> { - let value_type = data_source.schema().field(0).data_type().clone(); - while let Some(batch) = data_source.try_next().await? { - let values = batch.column_by_name(VALUE_COLUMN_NAME).expect_ok()?; - let row_ids = batch.column_by_name(ROW_ID).expect_ok()?; - debug_assert_eq!(row_ids.data_type(), &DataType::UInt64); - - let row_id_column = row_ids.as_any().downcast_ref::().unwrap(); - - for i in 0..values.len() { - let row_id = row_id_column.value(i); - let key = ScalarValue::try_from_array(values.as_ref(), i)?; - state.entry(key.clone()).or_default().insert(row_id); - } - } - - Ok((state, value_type)) - } - pub async fn train_bitmap_index( data: SendableRecordBatchStream, index_store: &dyn IndexStore, @@ -1410,8 +1400,9 @@ impl BitmapIndexPlugin { } /// Builds and writes a bitmap index in a streaming fashion from value-sorted - /// input. Only one value's bitmap is in memory at a time, reducing peak memory - /// from O(unique_values * avg_bitmap) to O(largest_single_bitmap). + /// input. Only one new value's aggregate bitmap is held at a time instead of + /// an aggregate map containing every value. The input pipeline, an existing + /// index and its cache, and the output writer retain separate memory. /// /// If `old_index` is provided, its existing bitmaps are merged with the new /// data via a sorted merge-join (the old index_map is a BTreeMap, already @@ -1587,24 +1578,6 @@ impl BitmapIndexPlugin { Ok(()) } - /// Remaps every bitmap in a materialized bitmap-index state using row-id mappings. - pub(crate) fn remap_bitmap_state( - state: HashMap, - mapping: &RowAddrRemap, - ) -> HashMap { - state - .into_iter() - .map(|(key, bitmap)| { - let remapped_bitmap = - RowAddrTreeMap::from_iter(bitmap.row_addrs().unwrap().filter_map(|addr| { - let addr_as_u64 = u64::from(addr); - mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64)) - })); - (key, remapped_bitmap) - }) - .collect() - } - /// Merge per-shard bitmap lookup files into a single bitmap index file. /// /// Each shard file is already sorted by key and can contain many distinct keys. @@ -1619,8 +1592,9 @@ impl BitmapIndexPlugin { /// - advance only those shards that participated in the union and push their next /// keys back into the heap /// - /// This keeps memory usage proportional to the number of shards plus the bitmaps - /// currently being merged, instead of the total number of keys across all shards. + /// The merge-specific working state is proportional to the number of shards + /// plus the bitmaps currently being merged, instead of the total number of + /// keys across all shards. This is not a total-memory or byte-bound claim. async fn merge_shards( store: &dyn IndexStore, shard_files: &[String], @@ -1630,46 +1604,21 @@ impl BitmapIndexPlugin { .stage_start("merge_bitmap_shards", None, "bitmaps") .await?; - let mut cursors = Vec::with_capacity(shard_files.len()); - let mut heap = BinaryHeap::with_capacity(shard_files.len()); - let mut value_type: Option = None; - - for file_name in shard_files { - let reader = store.open_index_file(file_name).await?; - let shard_value_type = reader.schema().fields[0].data_type().clone(); - if let Some(existing_type) = &value_type { - if existing_type != &shard_value_type { - return Err(Error::invalid_input(format!( - "Bitmap shard {} has value type {:?}, expected {:?}", - file_name, shard_value_type, existing_type - ))); - } - } else { - value_type = Some(shard_value_type); - } - if let Some(cursor) = BitmapShardCursor::try_new(file_name.clone(), reader).await? { - let key = cursor.peek_key()?; - let shard_idx = cursors.len(); - cursors.push(cursor); - heap.push(Reverse(BitmapHeapItem { key, shard_idx })); - } - } + let (mut cursors, mut heap, value_type) = + open_sorted_bitmap_cursors(store, shard_files).await?; let value_type = value_type.ok_or_else(|| { Error::invalid_input("Bitmap shard merge requires at least one shard file".to_string()) })?; let mut writer = new_bitmap_batch_writer(store, BITMAP_LOOKUP_NAME, &value_type).await?; - let mut merged_keys = 0u64; - - while let Some(Reverse(item)) = heap.pop() { - let (key, merged_bitmap) = - drain_same_key_bitmaps(&mut cursors, &mut heap, item).await?; - writer.emit(key, &merged_bitmap).await?; - merged_keys += 1; - progress - .stage_progress("merge_bitmap_shards", merged_keys) - .await?; - } + + drain_sorted_bitmap_cursors( + &mut cursors, + &mut heap, + &mut writer, + Some((progress.as_ref(), "merge_bitmap_shards")), + ) + .await?; progress.stage_complete("merge_bitmap_shards").await?; progress @@ -1699,6 +1648,162 @@ pub async fn merge_index_files( Ok(()) } +/// Apply `mapping` to every row address in `index` without materializing every +/// bitmap payload at once, writing the result through `writer`. +/// +/// This helper's transient aggregation state is one bitmap at a time. The +/// loaded `index_map`, index cache, mapping, and output writer remain resident +/// separately. Nulls live outside `index_map`, in `null_map`, so they are +/// remapped separately and emitted first -- a null sorts below every value. +pub(crate) async fn remap_index_map( + index: &BitmapIndex, + mapping: &RowAddrRemap, + writer: &mut BitmapBatchWriter, +) -> Result<()> { + if !index.null_map.is_empty() { + let null_key = new_null_array(index.value_type(), 1); + let null_key = ScalarValue::try_from_array(null_key.as_ref(), 0)?; + writer + .emit(null_key, &remap_row_addrs(&index.null_map, mapping)) + .await?; + } + + for key in index.index_map.keys() { + let bitmap = index.load_bitmap(key, None).await?; + writer + .emit(key.0.clone(), &remap_row_addrs(&bitmap, mapping)) + .await?; + } + + Ok(()) +} + +pub(crate) fn remap_row_addrs(bitmap: &RowAddrTreeMap, mapping: &RowAddrRemap) -> RowAddrTreeMap { + RowAddrTreeMap::from_iter(bitmap.row_addrs().unwrap().filter_map(|addr| { + let addr_as_u64 = u64::from(addr); + mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64)) + })) +} + +/// Total source entries a merge of `sources` will consume. +/// +/// The exact denominator for the merge's progress: known before the merge +/// starts, at no I/O cost since every `index_map` is already loaded, and reached +/// exactly, because the merge drains each source's keys and advances every one +/// of them once. Counting entries rather than emitted keys is what makes it +/// exact -- a key held by three sources costs three loads and yields one output +/// row, and a key whose rows are all retired by `old_data_filter` yields none. +/// +/// Nulls are excluded, matching how the merge treats them everywhere else: they +/// live in `null_map`, outside `index_map`, and are unioned in one step. +pub(crate) fn merge_source_entry_count(sources: &[Arc]) -> u64 { + sources.iter().map(|s| s.index_map.len() as u64).sum() +} + +/// Merge loaded bitmap indexes into `writer` without materializing all source +/// bitmap payloads at once. +/// +/// Drives each source through its `index_map` -- a sorted `BTreeMap` rebuilt at +/// load time -- rather than through the rows of its file. LabelList index files +/// written before spill-based builds landed are unsorted on disk, so file order +/// cannot be trusted for them; `index_map` order can, for old and new files +/// alike, which is what makes this work without an index version bump. +/// +/// Null keys live outside `index_map`, in each source's `null_map`, so they are +/// unioned separately and emitted first -- a null sorts below every value. +/// +/// The merge's transient aggregation state is the merged bitmap for the current +/// key plus one loaded bitmap per participating source. Each source `index_map`, +/// any bitmaps retained by the index cache, and the output writer remain outside +/// that state. +/// +/// `progress` reports source entries consumed, against the total from +/// [`merge_source_entry_count`]. Not segments: the merge is key-driven and +/// touches every source on every key, so no source is ever "done" to report. +/// Not emitted keys either: those have no denominator that can be known up front +/// or reached exactly. Entries have both, and are the unit the merge's cost is +/// actually in, since it loads one bitmap per entry. +pub(crate) async fn merge_index_maps( + sources: &[Arc], + old_data_filter: Option<&super::OldIndexDataFilter>, + writer: &mut BitmapBatchWriter, + progress: Option<(&dyn IndexBuildProgress, &str)>, +) -> Result<()> { + let Some(first) = sources.first() else { + return Ok(()); + }; + let value_type = first.value_type().clone(); + + let mut merged_nulls = RowAddrTreeMap::default(); + for source in sources { + merged_nulls |= source.null_map.as_ref(); + } + let merged_nulls = retain_valid(merged_nulls, old_data_filter); + if !merged_nulls.is_empty() { + let null_key = new_null_array(&value_type, 1); + let null_key = ScalarValue::try_from_array(null_key.as_ref(), 0)?; + writer.emit(null_key, &merged_nulls).await?; + } + + let mut consumed = 0u64; + + let mut key_iters: Vec<_> = sources + .iter() + .map(|source| source.index_map.keys()) + .collect(); + + // Every source is sorted, so the smallest key any of them is currently + // positioned on is the next key overall. A min-heap holding one entry per + // live source finds it in `log(sources)`, where scanning every source's + // current key -- twice, once to select and once to consume -- cost + // `O(keys x sources)`: with one segment per fragment that turns a linear + // merge into billions of comparisons. It is the same merge + // `drain_sorted_bitmap_cursors` runs over file-backed cursors. + // + // Entries borrow their key from the source's `index_map` rather than cloning + // it. Seeding the heap touches every source key, so cloning here would cost + // more than the scan it replaces whenever there are only a few sources. + let mut heap: BinaryHeap> = + BinaryHeap::with_capacity(key_iters.len()); + for (source_idx, keys) in key_iters.iter_mut().enumerate() { + if let Some(key) = keys.next() { + heap.push(Reverse((key, source_idx))); + } + } + + while let Some(Reverse((next_key, _))) = heap.peek().copied() { + let mut merged = RowAddrTreeMap::default(); + + // Drain the sources positioned on this key -- only those, where the + // previous scan visited every source on every key. A source's next key is + // strictly greater, so re-pushing it cannot re-enter this loop. + while let Some(Reverse((key, source_idx))) = heap.peek().copied() { + if key != next_key { + break; + } + heap.pop(); + consumed += 1; + merged |= sources[source_idx].load_bitmap(key, None).await?.as_ref(); + if let Some(next) = key_iters[source_idx].next() { + heap.push(Reverse((next, source_idx))); + } + } + + let merged = retain_valid(merged, old_data_filter); + if !merged.is_empty() { + writer.emit(next_key.0.clone(), &merged).await?; + } + + // Reported outside the guard above: a key the filter emptied still + // consumed its source entries, and skipping it would stall the count. + if let Some((progress, stage)) = progress { + progress.stage_progress(stage, consumed).await?; + } + } + + Ok(()) +} + pub async fn merge_bitmap_indices( source_indices: &[Arc], dest_store: &dyn IndexStore, @@ -1711,16 +1816,15 @@ pub async fn merge_bitmap_indices( } let value_type = source_indices[0].value_type().clone(); - let mut merged_state = HashMap::::new(); progress .stage_start( "merge_bitmap_segments", - Some(source_indices.len() as u64), - "segments", + Some(merge_source_entry_count(source_indices)), + "index entries", ) .await?; - for (idx, source_index) in source_indices.iter().enumerate() { + for source_index in source_indices.iter() { if source_index.value_type() != &value_type { return Err(Error::invalid_input(format!( "Bitmap segment has value type {:?}, expected {:?}", @@ -1728,24 +1832,22 @@ pub async fn merge_bitmap_indices( value_type ))); } - - let state = source_index.load_bitmap_index_state().await?; - for (key, bitmap) in state { - merged_state - .entry(key) - .and_modify(|existing| *existing |= &bitmap) - .or_insert(bitmap); - } - progress - .stage_progress("merge_bitmap_segments", (idx + 1) as u64) - .await?; } + + let mut writer = new_bitmap_batch_writer(dest_store, BITMAP_LOOKUP_NAME, &value_type).await?; + merge_index_maps( + source_indices, + None, + &mut writer, + Some((progress.as_ref(), "merge_bitmap_segments")), + ) + .await?; progress.stage_complete("merge_bitmap_segments").await?; progress .stage_start("write_bitmap_index", Some(1), "files") .await?; - let file = BitmapIndexPlugin::write_bitmap_index(merged_state, dest_store, &value_type).await?; + let file = writer.finish().await?; progress.stage_progress("write_bitmap_index", 1).await?; progress.stage_complete("write_bitmap_index").await?; @@ -1919,6 +2021,75 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { } } +/// Fixtures shared by the tests of every module that writes this file format. +#[cfg(test)] +pub(crate) mod test_util { + use std::sync::Arc; + + use arrow_array::{Array, BinaryArray}; + use datafusion_common::ScalarValue; + use lance_core::cache::LanceCache; + use lance_core::utils::tempfile::TempObjDir; + use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; + + use crate::scalar::IndexStore; + use crate::scalar::lance_format::LanceIndexStore; + + /// A local index store in a fresh temporary directory. The directory is + /// returned because dropping it deletes the store. + pub fn index_store() -> (TempObjDir, Arc) { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + (tmpdir, store) + } + + /// Every `(key, bitmap)` row of a bitmap-shaped index file, in file order. + /// + /// File order matters to callers: the build path emits keys ascending, while + /// indexes written before spill-based builds are in arbitrary order, and some + /// tests assert on which of the two they are looking at. + pub async fn read_key_bitmaps( + store: &dyn IndexStore, + file_name: &str, + ) -> Vec<(Option, RowAddrTreeMap)> { + let reader = store.open_index_file(file_name).await.unwrap(); + let total = reader.num_rows(); + if total == 0 { + return Vec::new(); + } + let batch = reader.read_range(0..total, None).await.unwrap(); + let bitmaps = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|idx| { + let key = match ScalarValue::try_from_array(batch.column(0), idx).unwrap() { + ScalarValue::Utf8(value) => value, + other => panic!("unexpected key type {other:?}"), + }; + let bitmap = RowAddrTreeMap::deserialize_from(bitmaps.value(idx)).unwrap(); + (key, bitmap) + }) + .collect() + } + + /// A bitmap's row addresses, ascending. Empty for a bitmap with no + /// enumerable addresses. + pub fn row_addrs(bitmap: &RowAddrTreeMap) -> Vec { + bitmap + .row_addrs() + .map(|iter| iter.map(u64::from).collect()) + .unwrap_or_default() + } +} + #[cfg(test)] mod tests { use super::*; @@ -2366,7 +2537,6 @@ mod tests { use lance_core::cache::LanceCache; use lance_io::object_store::ObjectStore; use lance_select::RowAddrTreeMap; - use std::collections::HashMap; use std::sync::Arc; // Adjust these numbers so that: @@ -2377,15 +2547,6 @@ mod tests { let m: u32 = 2_500_000; let per_bitmap_size = 1000; // assumed bytes per bitmap - let mut state = HashMap::new(); - for i in 0..m { - // Create a bitmap that contains, say, 1000 row IDs. - let bitmap = RowAddrTreeMap::from_iter(0..per_bitmap_size); - - let key = ScalarValue::UInt32(Some(i)); - state.insert(key, bitmap); - } - // Create a temporary store. let tmpdir = TempObjDir::default(); let test_store = LanceIndexStore::new( @@ -2394,10 +2555,22 @@ mod tests { Arc::new(LanceCache::no_cache()), ); - // This call should never trigger a "byte array offset overflow" error since now the code supports - // read by chunks - let result = - BitmapIndexPlugin::write_bitmap_index(state, &test_store, &DataType::UInt32).await; + // This should never trigger a "byte array offset overflow" error, since + // the writer flushes a batch once it reaches MAX_BUFFERED_BYTES, which is + // far below the i32 offset ceiling of either output column. + let mut writer = + new_bitmap_batch_writer(&test_store, BITMAP_LOOKUP_NAME, &DataType::UInt32) + .await + .unwrap(); + for i in 0..m { + // Create a bitmap that contains, say, 1000 row IDs. + let bitmap = RowAddrTreeMap::from_iter(0..per_bitmap_size); + writer + .emit(ScalarValue::UInt32(Some(i)), &bitmap) + .await + .unwrap(); + } + let result = writer.finish().await; assert!( result.is_ok(), @@ -2787,6 +2960,124 @@ mod tests { } } + /// Remap must emit exactly what the pre-streaming path did: one row per + /// source key, nulls included, every address put through the same mapping. + /// + /// The old path materialized the index into a + /// `HashMap`, remapped each entry and wrote the + /// whole map, so a key whose rows were all deleted still produced a row with + /// an empty bitmap. `remap_index_map` streams key-by-key instead and emits + /// unconditionally to preserve that -- deliberately unlike `merge_index_maps`, + /// which drops keys its filter empties. + #[tokio::test] + async fn test_bitmap_remap_matches_materialized_path() { + // frag 1 - { 0: null, 1: "a", 2: "b" } + // frag 2 - { 0: "a", 1: "c", 2: null } + let addrs: Vec = [(1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)] + .into_iter() + .map(|(frag, offset)| RowAddress::new_from_parts(frag, offset).into()) + .collect(); + let values = [None, Some("a"), Some("b"), Some("a"), Some("c"), None]; + + let (_src_dir, src_store) = test_util::index_store(); + let schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from_iter(values)), + Arc::new(UInt64Array::from(addrs.clone())), + ], + ) + .unwrap(); + let batch = sort_batch_by_value(&batch); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::once(async move { Ok(batch) }), + )); + BitmapIndexPlugin::train_bitmap_index(stream, src_store.as_ref()) + .await + .unwrap(); + let index = BitmapIndex::load(src_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + // One key per arm of the mapping's tri-state: "a" keeps a remapped row + // and loses a deleted one, "b" loses its only row, "c" is absent from the + // mapping and passes through unchanged, and the null bitmap -- which lives + // outside `index_map` -- sees both a remap and a delete. + let mapping = RowAddrRemap::direct(HashMap::from([ + (addrs[0], Some(RowAddress::new_from_parts(3, 0).into())), + (addrs[1], None), + (addrs[2], None), + (addrs[3], Some(RowAddress::new_from_parts(3, 1).into())), + (addrs[5], None), + ])); + + // The old path, restated: materialize every key including the null one, + // remap each bitmap, write them all out. + let mut old_path = Vec::new(); + if !index.null_map.is_empty() { + old_path.push((None, remap_row_addrs(&index.null_map, &mapping))); + } + for key in index.index_map.keys() { + let bitmap = index.load_bitmap(key, None).await.unwrap(); + let ScalarValue::Utf8(key) = key.0.clone() else { + panic!("keys are utf8") + }; + old_path.push((key, remap_row_addrs(&bitmap, &mapping))); + } + let mut old_path: Vec<(Option, Vec)> = old_path + .into_iter() + .map(|(key, bitmap)| (key, test_util::row_addrs(&bitmap))) + .collect(); + old_path.sort(); + + // Guard the oracle: a remap that emitted nothing, or that dropped the + // emptied key, would otherwise agree with an equally broken expectation. + let frag_3 = + |offset: u32| -> Vec { vec![RowAddress::new_from_parts(3, offset).into()] }; + assert_eq!( + old_path, + vec![ + (None, frag_3(0)), + (Some("a".to_string()), frag_3(1)), + (Some("b".to_string()), Vec::new()), + (Some("c".to_string()), vec![addrs[4]]), + ] + ); + + let (_dest_dir, dest_store) = test_util::index_store(); + index.remap(&mapping, dest_store.as_ref()).await.unwrap(); + assert_eq!(old_path, read_bitmap_contents(dest_store.as_ref()).await); + + // The old path wrote a `HashMap`, in no particular order. The streaming + // one emits the null key first and then ascending keys. + let written_keys: Vec> = + test_util::read_key_bitmaps(dest_store.as_ref(), BITMAP_LOOKUP_NAME) + .await + .into_iter() + .map(|(key, _)| key) + .collect(); + assert_eq!( + written_keys, + vec![ + None, + Some("a".to_string()), + Some("b".to_string()), + Some("c".to_string()) + ] + ); + + // The emptied key survives the round trip as a key rather than vanishing. + let reloaded = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + assert_eq!(reloaded.index_map.len(), 3); + } + #[tokio::test] async fn test_bitmap_null_handling_in_queries() { // Test that bitmap index correctly returns null_list for queries @@ -2926,4 +3217,137 @@ mod tests { _ => panic!("Expected Exact search result"), } } + + /// Merging bitmap segments must equal a single build over the same rows, + /// including the null bitmap, which lives outside `index_map`. + #[tokio::test] + async fn test_bitmap_segment_merge_matches_single_build() { + let values: Vec> = (0..600) + .map(|i| { + if i % 13 == 0 { + None + } else { + Some(format!("v-{:03}", i % 50)) + } + }) + .collect(); + + async fn build( + values: &[Option], + offset: u64, + ) -> (TempObjDir, Arc) { + let (tmpdir, store) = test_util::index_store(); + let schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from_iter(values.iter().cloned())), + Arc::new(UInt64Array::from_iter_values( + (0..values.len() as u64).map(|i| i + offset), + )), + ], + ) + .unwrap(); + let batch = sort_batch_by_value(&batch); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::once(async move { Ok(batch) }), + )); + BitmapIndexPlugin::train_bitmap_index(stream, store.as_ref()) + .await + .unwrap(); + (tmpdir, store) + } + + let (_all_dir, all_store) = build(&values, 0).await; + let expected = read_bitmap_contents(all_store.as_ref()).await; + assert!( + expected.iter().any(|(key, _)| key.is_none()), + "fixture must exercise the null bitmap" + ); + + let (_left_dir, left_store) = build(&values[..300], 0).await; + let (_right_dir, right_store) = build(&values[300..], 300).await; + let left = BitmapIndex::load(left_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + let right = BitmapIndex::load(right_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + let (_dest_dir, dest_store) = test_util::index_store(); + merge_bitmap_indices( + &[left, right], + dest_store.as_ref(), + crate::progress::noop_progress(), + ) + .await + .unwrap(); + + assert_eq!(expected, read_bitmap_contents(dest_store.as_ref()).await); + } + + /// The keys column counts toward the flush threshold, not just the bitmaps. + /// A column with a very large number of tiny bitmaps used to buffer without + /// limit: the old threshold charged the bitmap column only, so the keys could + /// grow until the writer held gigabytes and their Arrow i32 offsets overflowed. + #[tokio::test] + async fn test_batch_writer_charges_keys_against_flush_threshold() { + // Small enough to cross with a handful of keys; the production threshold + // is MAX_BUFFERED_BYTES, which no test wants to write out. + const THRESHOLD: usize = 4 * 1024; + const NUM_KEYS: u64 = 64; + + let (_tmpdir, store) = test_util::index_store(); + let mut writer = + new_bitmap_batch_writer(store.as_ref(), BITMAP_LOOKUP_NAME, &DataType::Utf8) + .await + .unwrap() + .with_max_buffered_bytes(THRESHOLD); + + // Every bitmap holds one row, so all the bitmaps together stay under the + // threshold; only the keys can push the writer past it. + for i in 0..NUM_KEYS { + let key = format!("{i:08}{}", "k".repeat(512)); + writer + .emit( + ScalarValue::Utf8(Some(key)), + &RowAddrTreeMap::from_iter([i]), + ) + .await + .unwrap(); + } + assert!( + writer.batches_written() > 1, + "keys must be charged against the flush threshold, but the writer \ + buffered all {NUM_KEYS} keys in {} batch(es)", + writer.batches_written() + ); + writer.finish().await.unwrap(); + + // Flushing part-way through must not disturb the file's contents. + let contents = read_bitmap_contents(store.as_ref()).await; + assert_eq!(contents.len() as u64, NUM_KEYS); + for (idx, (key, addrs)) in contents.iter().enumerate() { + let key = key.as_deref().expect("keys are non-null"); + assert_eq!(&key[..8], format!("{idx:08}"), "keys must stay ascending"); + assert_eq!(addrs, &vec![idx as u64]); + } + } + + /// Key to sorted row addresses, read from the file and sorted so the + /// comparison does not depend on the order keys happen to be written in. + async fn read_bitmap_contents(store: &dyn IndexStore) -> Vec<(Option, Vec)> { + let mut out: Vec<(Option, Vec)> = + test_util::read_key_bitmaps(store, BITMAP_LOOKUP_NAME) + .await + .into_iter() + .map(|(key, bitmap)| (key, test_util::row_addrs(&bitmap))) + .collect(); + out.sort(); + out + } } diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index d4307d62670..b95649d2087 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -4,7 +4,6 @@ use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, - collections::HashMap, fmt::Debug, pin::Pin, sync::{Arc, Mutex}, @@ -37,7 +36,10 @@ use super::{ use super::{BuiltinIndexType, SargableQuery, ScalarIndexParams}; use super::{MetricsCollector, SearchResult}; use crate::pbold; -use crate::scalar::bitmap::{BitmapIndexPlugin, BitmapIndexState}; +use crate::scalar::bitmap::{ + BitmapIndexState, merge_index_maps, merge_source_entry_count, new_bitmap_batch_writer, + remap_index_map, remap_row_addrs, +}; use crate::scalar::expression::{LabelListQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ BasicTrainer, DefaultTrainingRequest, ScalarIndexLoad, ScalarIndexPlugin, TrainingCriteria, @@ -46,6 +48,8 @@ use crate::scalar::registry::{ use crate::scalar::{CreatedIndex, RowIdRemapper, UpdateCriteria}; use crate::{Index, IndexType}; +mod spill; + pub const BITMAP_LOOKUP_NAME: &str = "bitmap_page_lookup.lance"; pub const LABEL_LIST_NULLS_METADATA_KEY: &str = "lance:label_list_nulls"; pub const LABEL_LIST_NULLS_MIN_VERSION: i32 = 1; @@ -217,20 +221,21 @@ impl ScalarIndex for LabelListIndex { mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { - let state = self.values_index.load_bitmap_index_state().await?; - let remapped_state = BitmapIndexPlugin::remap_bitmap_state(state, mapping); - let remapped_nulls = - RowAddrTreeMap::from_iter(self.list_nulls.row_addrs().unwrap().filter_map(|addr| { - let addr_as_u64 = u64::from(addr); - mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64)) - })); - let file = write_label_list_bitmap_index( - remapped_state, + let remapped_nulls = remap_row_addrs(&self.list_nulls, mapping); + let mut writer = new_bitmap_batch_writer( dest_store, + BITMAP_LOOKUP_NAME, self.values_index.value_type(), - &remapped_nulls, ) .await?; + writer + .add_global_buffer( + LABEL_LIST_NULLS_METADATA_KEY.to_string(), + serialize_list_nulls(&remapped_nulls)?, + ) + .await?; + remap_index_map(&self.values_index, mapping, &mut writer).await?; + let file = writer.finish().await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::LabelListIndexDetails::default()) @@ -247,20 +252,19 @@ impl ScalarIndex for LabelListIndex { dest_store: &dyn IndexStore, old_data_filter: Option, ) -> Result { - let state = self.values_index.load_bitmap_index_state().await?; - let list_nulls = Arc::new(Mutex::new(RowAddrTreeMap::new())); - let new_data = track_list_nulls(new_data, list_nulls.clone()); - let (merged_state, value_type) = - BitmapIndexPlugin::build_bitmap_index_state(unnest_chunks(new_data)?, state).await?; + // Not applied, matching every other derived-key scalar index (ngram, + // fmindex, bloomfilter, zonemap, rtree). Only btree and bitmap -- one + // key per row -- prune retired rows here. Whether the derived-key class + // should honour the filter is unresolved and tracked separately; see + // OSS-2032. let _ = old_data_filter; - let mut merged_nulls = (*self.list_nulls).clone(); - let new_nulls = list_nulls.lock().unwrap().clone(); - if !new_nulls.is_empty() { - merged_nulls |= &new_nulls; - } - let file = - write_label_list_bitmap_index(merged_state, dest_store, &value_type, &merged_nulls) - .await?; + let file = update_label_list_index( + self, + new_data, + dest_store, + spill::default_spill_budget_bytes()?, + ) + .await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::LabelListIndexDetails::default()) @@ -468,34 +472,137 @@ fn serialize_list_nulls(null_map: &RowAddrTreeMap) -> Result { Ok(Bytes::from(bytes)) } -async fn write_label_list_bitmap_index( - state: HashMap, +/// Drain `data` into `builder` as one `(label, row address)` pair per unnested +/// list element. +async fn accumulate_labels( + mut data: SendableRecordBatchStream, + builder: &mut spill::LabelListSpillBuilder, +) -> Result<()> { + while let Some(batch) = data.try_next().await? { + let values = batch.column_by_name(VALUE_COLUMN_NAME).expect_ok()?; + let row_addrs = batch.column_by_name(ROW_ID).expect_ok()?; + debug_assert_eq!(row_addrs.data_type(), &DataType::UInt64); + let row_addrs = row_addrs.as_any().downcast_ref::().unwrap(); + for i in 0..values.len() { + let key = ScalarValue::try_from_array(values.as_ref(), i)?; + builder.insert(key, row_addrs.value(i)).await?; + } + } + Ok(()) +} + +/// Merge `spills` into a LabelList index file in `store`, carrying `list_nulls` +/// in the file's global buffer. +async fn write_label_list_index( store: &dyn IndexStore, value_type: &DataType, list_nulls: &RowAddrTreeMap, + spills: &mut spill::LabelListSpills, ) -> Result { - BitmapIndexPlugin::write_bitmap_index_with_extras( - state, - store, - value_type, - HashMap::new(), - vec![( + let mut writer = new_bitmap_batch_writer(store, BITMAP_LOOKUP_NAME, value_type).await?; + writer + .add_global_buffer( LABEL_LIST_NULLS_METADATA_KEY.to_string(), serialize_list_nulls(list_nulls)?, - )], - ) - .await + ) + .await?; + spills.merge_into(&mut writer).await?; + writer.finish().await +} + +/// Build a LabelList index from an unnested-and-tracked `(value, row addr)` +/// stream while limiting the aggregation state to `spill_budget_bytes`. +/// +/// The label to row-set map is accumulated in a byte-budgeted sorted map that +/// spills to local scratch when it grows too large; the spills are then k-way +/// merged straight into the index file. The written file is identical in format +/// to the one the previous in-memory build produced, except that its keys are +/// now in ascending order. +/// +/// The budget applies only to the mutable label-to-row-set map. Scan batches, +/// `list_nulls`, merge cursors and bitmaps, the output writer, caches, and +/// concurrent operations are outside it. +async fn train_label_list_index( + data: SendableRecordBatchStream, + index_store: &dyn IndexStore, + spill_budget_bytes: usize, +) -> Result { + let list_nulls = Arc::new(Mutex::new(RowAddrTreeMap::new())); + let data = track_list_nulls(data, list_nulls.clone()); + let data = unnest_chunks(data)?; + + let value_type = data.schema().field(0).data_type().clone(); + let mut builder = + spill::LabelListSpillBuilder::new_local(value_type.clone(), spill_budget_bytes)?; + accumulate_labels(data, &mut builder).await?; + + let mut spills = builder.finish().await?; + let list_nulls = list_nulls.lock().unwrap().clone(); + write_label_list_index(index_store, &value_type, &list_nulls, &mut spills).await +} + +/// Add `new_data` to an existing LabelList index while limiting the new-data +/// aggregation state to `spill_budget_bytes`, writing the result to `dest_store`. +/// +/// The existing index's bitmap payload is read one key at a time rather than +/// materialized in full. It is rewritten into one more sorted scratch file, +/// which then joins the new data's spill files in the same k-way merge. The +/// source `index_map`, index cache, `list_nulls`, merge working state, and output +/// writer remain outside the aggregation budget. The rewrite costs a full pass +/// of the index through scratch; see +/// [`spill::LabelListSpills::add_existing_index`] for why it is not merged from +/// its own file directly. +async fn update_label_list_index( + existing: &LabelListIndex, + new_data: SendableRecordBatchStream, + dest_store: &dyn IndexStore, + spill_budget_bytes: usize, +) -> Result { + let list_nulls = Arc::new(Mutex::new(RowAddrTreeMap::new())); + let new_data = track_list_nulls(new_data, list_nulls.clone()); + let new_data = unnest_chunks(new_data)?; + + // The spill and destination file schemas are declared from the existing + // index, but every key written comes from the new stream. If the two + // disagree the written batches would not match the schema they are declared + // under, so reject it here rather than emit a file whose schema lies about + // its key type. `open_sorted_bitmap_cursors` makes the same check across its + // inputs. + let value_type = existing.values_index.value_type().clone(); + let new_value_type = new_data.schema().field(0).data_type().clone(); + if new_value_type != value_type { + return Err(Error::invalid_input(format!( + "Cannot update a LabelList index with value type {value_type} \ + from new data whose list items are {new_value_type}" + ))); + } + + let mut builder = + spill::LabelListSpillBuilder::new_local(value_type.clone(), spill_budget_bytes)?; + accumulate_labels(new_data, &mut builder).await?; + + let mut spills = builder.finish().await?; + spills.add_existing_index(&existing.values_index).await?; + + let mut merged_nulls = (*existing.list_nulls).clone(); + let new_nulls = list_nulls.lock().unwrap().clone(); + if !new_nulls.is_empty() { + merged_nulls |= &new_nulls; + } + + write_label_list_index(dest_store, &value_type, &merged_nulls, &mut spills).await } /// Merge multiple LabelList index segments into a single index. /// /// A [`LabelListIndex`] is a [`BitmapIndex`] over the unnested list values plus a /// separate `list_nulls` row set. Because distributed segments cover disjoint rows -/// (distinct fragments), merging is a cheap union of the underlying bitmap states -/// and of the `list_nulls` sets — no re-scan of source data is required. This -/// mirrors [`crate::scalar::bitmap::merge_bitmap_indices`] but also carries the -/// per-segment `list_nulls`. When `old_data_filter` is provided, rows from -/// retired fragments are removed from both the value bitmaps and `list_nulls`. +/// (distinct fragments), merging streams and unions the bitmap payloads by key +/// and separately unions the `list_nulls` sets; no source-data re-scan is +/// required. This mirrors [`crate::scalar::bitmap::merge_bitmap_indices`] but +/// also carries the per-segment `list_nulls`. When `old_data_filter` is provided, +/// rows from retired fragments are removed from both the value bitmaps and +/// `list_nulls`. pub async fn merge_label_list_indices( source_indices: &[Arc], dest_store: &dyn IndexStore, @@ -509,17 +616,10 @@ pub async fn merge_label_list_indices( } let value_type = source_indices[0].values_index.value_type().clone(); - let mut merged_state = HashMap::::new(); let mut merged_nulls = RowAddrTreeMap::new(); - progress - .stage_start( - "merge_label_list_segments", - Some(source_indices.len() as u64), - "segments", - ) - .await?; - for (idx, source_index) in source_indices.iter().enumerate() { + let mut values_indices = Vec::with_capacity(source_indices.len()); + for source_index in source_indices.iter() { if source_index.values_index.value_type() != &value_type { return Err(Error::invalid_input(format!( "LabelList segment has value type {:?}, expected {:?}", @@ -527,36 +627,46 @@ pub async fn merge_label_list_indices( value_type ))); } + values_indices.push(source_index.values_index.clone()); - let state = source_index.values_index.load_bitmap_index_state().await?; - for (key, mut bitmap) in state { - if let Some(filter) = old_data_filter.as_ref() { - filter.retain_old_rows(&mut bitmap); - } - if bitmap.is_empty() { - continue; - } - merged_state - .entry(key) - .and_modify(|existing| *existing |= &bitmap) - .or_insert(bitmap); - } + // `list_nulls` records whole rows whose list was null, so it is a row + // set per segment rather than a per-label bitmap. It stays outside both + // the key merge and the aggregation-state budget. let mut list_nulls = source_index.list_nulls.as_ref().clone(); if let Some(filter) = old_data_filter.as_ref() { filter.retain_old_rows(&mut list_nulls); } merged_nulls |= &list_nulls; - progress - .stage_progress("merge_label_list_segments", (idx + 1) as u64) - .await?; } + + progress + .stage_start( + "merge_label_list_segments", + Some(merge_source_entry_count(&values_indices)), + "index entries", + ) + .await?; + + let mut writer = new_bitmap_batch_writer(dest_store, BITMAP_LOOKUP_NAME, &value_type).await?; + writer + .add_global_buffer( + LABEL_LIST_NULLS_METADATA_KEY.to_string(), + serialize_list_nulls(&merged_nulls)?, + ) + .await?; + merge_index_maps( + &values_indices, + old_data_filter.as_ref(), + &mut writer, + Some((progress.as_ref(), "merge_label_list_segments")), + ) + .await?; progress.stage_complete("merge_label_list_segments").await?; progress .stage_start("write_label_list_index", Some(1), "files") .await?; - let file = - write_label_list_bitmap_index(merged_state, dest_store, &value_type, &merged_nulls).await?; + let file = writer.finish().await?; progress.stage_progress("write_label_list_index", 1).await?; progress.stage_complete("write_label_list_index").await?; @@ -752,14 +862,8 @@ impl BasicTrainer for LabelListIndexPlugin { validate_label_list_data_type(field.data_type())?; - let list_nulls = Arc::new(Mutex::new(RowAddrTreeMap::new())); - let data = track_list_nulls(data, list_nulls.clone()); - let data = unnest_chunks(data)?; - let (state, value_type) = - BitmapIndexPlugin::build_bitmap_index_state(data, HashMap::new()).await?; - let list_nulls = list_nulls.lock().unwrap().clone(); let file = - write_label_list_bitmap_index(state, index_store, &value_type, &list_nulls).await?; + train_label_list_index(data, index_store, spill::default_spill_budget_bytes()?).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::LabelListIndexDetails::default()) .unwrap(), @@ -869,6 +973,8 @@ mod tests { use super::super::bitmap::BitmapIndexState; use super::super::btree::OrderableScalarValue; use super::*; + use crate::scalar::bitmap::test_util::{self, row_addrs}; + use lance_core::utils::tempfile::TempObjDir; #[rstest] #[case::list(DataType::List(Arc::new(Field::new( @@ -974,4 +1080,733 @@ mod tests { } } } + + // ---- shared fixtures for the spill-build tests ------------------------ + + /// One test row: its address, and either a null list or a list whose + /// elements may individually be null. Both null shapes matter — a null list + /// is recorded in `list_nulls`, a null element survives unnesting as a null + /// index key. + type SampleRow = (u64, Option>>); + + /// Deterministic rows with heavy label reuse, a null list every 17th row, + /// and a null element every 11th, so every code path sees both. + fn sample_label_list_rows(count: usize, addr_offset: u64) -> Vec { + (0..count) + .map(|i| { + let addr = addr_offset + i as u64; + if i % 17 == 0 { + return (addr, None); + } + let mut labels: Vec> = (0..(i % 4) + 1) + .map(|k| Some(format!("label-{:04}", (i * 7 + k * 13) % 200))) + .collect(); + if i % 11 == 0 { + labels.push(None); + } + (addr, Some(labels)) + }) + .collect() + } + + fn label_list_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + VALUE_COLUMN_NAME, + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + Field::new(ROW_ID, DataType::UInt64, false), + ])) + } + + /// Chunk the rows into several batches so the builder crosses batch + /// boundaries the way a real scan does. + fn sample_rows_to_stream(rows: &[SampleRow]) -> SendableRecordBatchStream { + let schema = label_list_schema(); + let batches: Vec<_> = rows + .chunks(64) + .map(|chunk| { + let mut builder = arrow_array::builder::ListBuilder::new( + arrow_array::builder::StringBuilder::new(), + ); + for (_, labels) in chunk { + match labels { + Some(labels) => { + for label in labels { + builder.values().append_option(label.as_deref()); + } + builder.append(true); + } + None => builder.append(false), + } + } + let addrs = UInt64Array::from_iter_values(chunk.iter().map(|(addr, _)| *addr)); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(builder.finish()), Arc::new(addrs)], + ) + .unwrap() + }) + .collect(); + + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches.into_iter().map(Ok)), + )) + } + + /// The index as a comparable value: labels (nulls included) each mapped to + /// their sorted row addresses, plus the separately-stored `list_nulls`. + /// Read from the file rather than the loaded index so that the on-disk + /// format itself is what the assertions compare. + #[derive(Debug, PartialEq, Eq)] + struct IndexContents { + labels: Vec<(Option, Vec)>, + list_nulls: Vec, + } + + async fn read_index_contents(store: &dyn IndexStore) -> IndexContents { + let mut labels: Vec<(Option, Vec)> = + test_util::read_key_bitmaps(store, BITMAP_LOOKUP_NAME) + .await + .into_iter() + .map(|(key, bitmap)| (key, row_addrs(&bitmap))) + .collect(); + // The build path emits keys in sorted order; older indexes on disk are + // in arbitrary order. Sort so the comparison is order-independent. + labels.sort(); + + let list_nulls = read_list_nulls(store.clone_arc(), None).await.unwrap(); + IndexContents { + labels, + list_nulls: row_addrs(&list_nulls), + } + } + + async fn build_label_list_index( + rows: &[SampleRow], + spill_budget_bytes: usize, + ) -> IndexContents { + let (_tmpdir, store) = test_util::index_store(); + train_label_list_index( + sample_rows_to_stream(rows), + store.as_ref(), + spill_budget_bytes, + ) + .await + .unwrap(); + read_index_contents(store.as_ref()).await + } + + /// Same input, different budgets: the emitted index must be identical no + /// matter how many times the builder spilled. + #[rstest] + #[case::spills_often(1024)] + #[case::spills_every_key(1)] + #[tokio::test] + async fn test_train_index_identical_across_spill_budgets(#[case] spill_budget_bytes: usize) { + let rows = sample_label_list_rows(400, 0); + let no_spill = build_label_list_index(&rows, usize::MAX).await; + + assert!( + no_spill.labels.iter().any(|(key, _)| key.is_none()), + "fixture must exercise null labels" + ); + assert!( + !no_spill.list_nulls.is_empty(), + "fixture must exercise null lists" + ); + + assert_eq!( + no_spill, + build_label_list_index(&rows, spill_budget_bytes).await + ); + } + + /// The build must leave nothing behind in the index directory but the index. + #[tokio::test] + async fn test_train_index_leaves_no_spill_files() { + let (_tmpdir, store) = test_util::index_store(); + train_label_list_index( + sample_rows_to_stream(&sample_label_list_rows(500, 0)), + store.as_ref(), + 1, + ) + .await + .unwrap(); + + let files: Vec = store + .list_files_with_sizes() + .await + .unwrap() + .into_iter() + .map(|file| file.path) + .collect(); + assert_eq!(files, vec![BITMAP_LOOKUP_NAME.to_string()]); + } + + async fn build_label_list_segment(rows: &[SampleRow]) -> (TempObjDir, Arc) { + let (tmpdir, store) = test_util::index_store(); + train_label_list_index(sample_rows_to_stream(rows), store.as_ref(), usize::MAX) + .await + .unwrap(); + let index = LabelListIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + (tmpdir, index) + } + + /// `remap` carries `list_nulls` in a global buffer written through the + /// streaming writer, and nothing else covers that path. If the buffer or its + /// metadata key were dropped, `read_list_nulls` would return an empty set and + /// the index would be classified as pre-nulls, so `NOT` filters over nullable + /// list columns would return wrong rows with only a warning -- no error. + #[tokio::test] + async fn test_remap_rewrites_addresses_and_preserves_list_nulls() { + const OLD_FRAGMENT: u64 = 1 << 32; + const NEW_FRAGMENT: u64 = 3 << 32; + const ROWS: usize = 120; + + let rows = sample_label_list_rows(ROWS, OLD_FRAGMENT); + // Built here rather than via build_label_list_segment so the source + // store stays reachable, to compare against after the remap. + let (_src_dir, src_store) = test_util::index_store(); + train_label_list_index(sample_rows_to_stream(&rows), src_store.as_ref(), usize::MAX) + .await + .unwrap(); + let index = LabelListIndex::load(src_store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + let before = read_index_contents(src_store.as_ref()).await; + assert!( + !before.list_nulls.is_empty(), + "fixture must exercise null lists, or the buffer is not under test" + ); + assert!( + before.labels.iter().any(|(key, _)| key.is_none()), + "fixture must exercise null labels" + ); + + // Move every row to a new fragment, and delete one so the dropped-row + // path is covered too. + let deleted = OLD_FRAGMENT + 1; + let mapping = RowAddrRemap::direct( + (0..ROWS as u64) + .map(|i| { + let old = OLD_FRAGMENT + i; + (old, (old != deleted).then_some(NEW_FRAGMENT + i)) + }) + .collect(), + ); + + let (_dest_dir, dest_store) = test_util::index_store(); + index.remap(&mapping, dest_store.as_ref()).await.unwrap(); + let after = read_index_contents(dest_store.as_ref()).await; + + // The nulls buffer must survive, remapped, not silently become empty. + assert_eq!( + after.list_nulls, + before + .list_nulls + .iter() + .filter(|addr| **addr != deleted) + .map(|addr| addr - OLD_FRAGMENT + NEW_FRAGMENT) + .collect::>(), + "list_nulls must be carried through remap with its addresses rewritten" + ); + + assert_eq!( + after.labels.len(), + before.labels.len(), + "remap must not add or drop labels" + ); + for ((old_key, old_addrs), (new_key, new_addrs)) in before.labels.iter().zip(&after.labels) + { + assert_eq!(old_key, new_key, "labels themselves must be unchanged"); + assert_eq!( + new_addrs, + &old_addrs + .iter() + .filter(|addr| **addr != deleted) + .map(|addr| addr - OLD_FRAGMENT + NEW_FRAGMENT) + .collect::>(), + "every row address for {old_key:?} must be rewritten to the new fragment" + ); + } + } + + /// A key whose rows are all retired emits nothing but still consumed its + /// source entries. Progress has to count it, or a merge that drops a + /// segment's exclusive labels ends below the total it declared -- the + /// stalled bar this reporting exists to remove. + #[tokio::test] + async fn test_segment_merge_progress_reaches_total_when_keys_are_filtered_out() { + const FRAGMENT_1: u64 = 1 << 32; + + // The retired segment owns labels the kept one does not, so the filter + // empties those keys outright rather than merely trimming them. + let kept: Vec = (0..40u64) + .map(|i| (i, Some(vec![Some(format!("kept-{i:04}"))]))) + .collect(); + let retired: Vec = (0..40u64) + .map(|i| (FRAGMENT_1 + i, Some(vec![Some(format!("retired-{i:04}"))]))) + .collect(); + + let (_kept_dir, kept_index) = build_label_list_segment(&kept).await; + let (_retired_dir, retired_index) = build_label_list_segment(&retired).await; + let (_dest_dir, dest_store) = test_util::index_store(); + + let progress = Arc::new(RecordingProgress::default()); + merge_label_list_indices( + &[kept_index, retired_index], + dest_store.as_ref(), + Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::from_iter([0u32]), + to_remove: RoaringBitmap::from_iter([1u32]), + }), + progress.clone(), + ) + .await + .unwrap(); + + let merged = read_index_contents(dest_store.as_ref()).await; + assert_eq!( + merged.labels.len(), + kept.len(), + "only the kept segment's labels survive, so half the entries emit nothing" + ); + + let declared_total = progress.declared_total("merge_label_list_segments"); + assert_eq!( + declared_total as usize, + kept.len() + retired.len(), + "both segments' entries count toward the total" + ); + assert_eq!( + progress.counts("merge_label_list_segments").last().copied(), + Some(declared_total), + "entries consumed by keys the filter emptied must still be reported" + ); + } + + /// Many segments whose keys interleave without overlapping is the case key + /// selection must handle in `log(segments)` rather than by scanning every + /// segment for every key. Covers both shapes at once: a vocabulary unique to + /// each segment, plus one label every segment holds. + #[tokio::test] + async fn test_segment_merge_over_many_interleaved_segments() { + const NUM_SEGMENTS: usize = 5; + const ROWS_PER_SEGMENT: usize = 40; + const SHARED_LABEL: &str = "shared"; + + fn segment_rows(segment: usize) -> Vec { + (0..ROWS_PER_SEGMENT) + .map(|i| { + let addr = ((segment as u64) << 32) | i as u64; + // Striding by the segment count interleaves the vocabularies, + // so each segment's keys fall between its neighbours' rather + // than alongside them. + let labels = vec![ + Some(format!("label-{:04}", i * NUM_SEGMENTS + segment)), + Some(SHARED_LABEL.to_string()), + ]; + (addr, Some(labels)) + }) + .collect() + } + + let mut all_rows = Vec::new(); + let mut segments = Vec::new(); + let mut _dirs = Vec::new(); + for segment in 0..NUM_SEGMENTS { + let rows = segment_rows(segment); + all_rows.extend(rows.clone()); + let (dir, index) = build_label_list_segment(&rows).await; + _dirs.push(dir); + segments.push(index); + } + + let (_dest_dir, dest_store) = test_util::index_store(); + merge_label_list_indices( + &segments, + dest_store.as_ref(), + None, + crate::progress::noop_progress(), + ) + .await + .unwrap(); + + let merged = read_index_contents(dest_store.as_ref()).await; + assert_eq!( + merged, + build_label_list_index(&all_rows, usize::MAX).await, + "merging {NUM_SEGMENTS} interleaved segments must equal a single build" + ); + + let shared = merged + .labels + .iter() + .find(|(key, _)| key.as_deref() == Some(SHARED_LABEL)) + .expect("the shared label must survive the merge"); + assert_eq!( + shared.1.len(), + NUM_SEGMENTS * ROWS_PER_SEGMENT, + "a key held by every segment must union all their rows" + ); + } + + /// Records what a build reported, so progress can be asserted on. + #[derive(Debug, Default)] + struct RecordingProgress { + /// `(stage, completed)` for every progress call, in order. + calls: std::sync::Mutex>, + /// `(stage, total)` for every stage start. + totals: std::sync::Mutex)>>, + } + + impl RecordingProgress { + fn counts(&self, stage: &str) -> Vec { + self.calls + .lock() + .unwrap() + .iter() + .filter(|(reported, _)| reported == stage) + .map(|(_, completed)| *completed) + .collect() + } + + fn declared_total(&self, stage: &str) -> u64 { + self.totals + .lock() + .unwrap() + .iter() + .find(|(reported, _)| reported == stage) + .unwrap_or_else(|| panic!("stage {stage} was never started")) + .1 + .unwrap_or_else(|| panic!("stage {stage} declared no total")) + } + } + + #[async_trait] + impl crate::progress::IndexBuildProgress for RecordingProgress { + async fn stage_start(&self, stage: &str, total: Option, _: &str) -> Result<()> { + self.totals.lock().unwrap().push((stage.to_string(), total)); + Ok(()) + } + async fn stage_progress(&self, stage: &str, completed: u64) -> Result<()> { + self.calls + .lock() + .unwrap() + .push((stage.to_string(), completed)); + Ok(()) + } + async fn stage_complete(&self, _: &str) -> Result<()> { + Ok(()) + } + } + + /// The segment merge is key-driven, so it must report progress as it goes -- + /// reporting once on return leaves a long merge indistinguishable from a hung + /// one -- and against a denominator it actually reaches. Source entries give + /// one; emitted keys do not, since sources sharing a label emit one row for + /// several entries. + #[tokio::test] + async fn test_segment_merge_reports_progress_while_merging() { + let rows = sample_label_list_rows(200, 0); + let (left, right) = rows.split_at(100); + let (_left_dir, left_index) = build_label_list_segment(left).await; + let (_right_dir, right_index) = build_label_list_segment(right).await; + let (_dest_dir, dest_store) = test_util::index_store(); + + let progress = Arc::new(RecordingProgress::default()); + merge_label_list_indices( + &[left_index, right_index], + dest_store.as_ref(), + None, + progress.clone(), + ) + .await + .unwrap(); + + let merge_counts = progress.counts("merge_label_list_segments"); + let declared_total = progress.declared_total("merge_label_list_segments"); + + let merged_keys = read_index_contents(dest_store.as_ref()).await.labels.len() as u64; + assert!( + merged_keys > 1, + "fixture must merge more than one key, got {merged_keys}" + ); + // The two segments share a vocabulary, so entries outnumber output rows. + // This is what an emitted-key denominator cannot express, and reverting + // to one fails here rather than silently under-reporting. + assert!( + declared_total > merged_keys, + "fixture must have segments sharing labels so entries ({declared_total}) \ + exceed emitted keys ({merged_keys})" + ); + + assert!( + merge_counts.windows(2).all(|pair| pair[0] <= pair[1]), + "reported progress must never go backwards: {merge_counts:?}" + ); + assert_eq!( + merge_counts.last().copied(), + Some(declared_total), + "the merge must finish on exactly the total it declared" + ); + } + + /// Merging N segments must equal a single build over the same rows. + #[tokio::test] + async fn test_segment_merge_matches_single_build() { + let rows = sample_label_list_rows(1500, 0); + let (left, right) = rows.split_at(750); + + let single = build_label_list_index(&rows, usize::MAX).await; + + let (_left_dir, left_index) = build_label_list_segment(left).await; + let (_right_dir, right_index) = build_label_list_segment(right).await; + let (_dest_dir, dest_store) = test_util::index_store(); + merge_label_list_indices( + &[left_index, right_index], + dest_store.as_ref(), + None, + crate::progress::noop_progress(), + ) + .await + .unwrap(); + + assert_eq!(single, read_index_contents(dest_store.as_ref()).await); + } + + /// Segment merge prunes rows from retired fragments. This is existing + /// behaviour and the streaming merge must not quietly drop it. + #[tokio::test] + async fn test_segment_merge_applies_old_data_filter() { + const FRAGMENT_1: u64 = 1 << 32; + let kept = sample_label_list_rows(400, 0); + let retired = sample_label_list_rows(400, FRAGMENT_1); + + let expected = build_label_list_index(&kept, usize::MAX).await; + + let (_kept_dir, kept_index) = build_label_list_segment(&kept).await; + let (_retired_dir, retired_index) = build_label_list_segment(&retired).await; + let (_dest_dir, dest_store) = test_util::index_store(); + merge_label_list_indices( + &[kept_index, retired_index], + dest_store.as_ref(), + Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::from_iter([0u32]), + to_remove: RoaringBitmap::from_iter([1u32]), + }), + crate::progress::noop_progress(), + ) + .await + .unwrap(); + + let merged = read_index_contents(dest_store.as_ref()).await; + assert_eq!( + expected, merged, + "rows from the retired fragment must not survive the merge" + ); + } + + /// Updating an existing index with new rows must equal a full rebuild over + /// the union of old and new rows, at any spill budget. + #[rstest] + #[case::no_spill(usize::MAX)] + #[case::spills_often(1024)] + #[case::spills_every_key(1)] + #[tokio::test] + async fn test_update_matches_full_rebuild(#[case] spill_budget_bytes: usize) { + const FRAGMENT_1: u64 = 1 << 32; + let initial = sample_label_list_rows(300, 0); + let additional = sample_label_list_rows(150, FRAGMENT_1); + + let mut all = initial.clone(); + all.extend(additional.clone()); + let rebuilt = build_label_list_index(&all, usize::MAX).await; + + let (_src_dir, index) = build_label_list_segment(&initial).await; + let (_dest_dir, dest_store) = test_util::index_store(); + update_label_list_index( + index.as_ref(), + sample_rows_to_stream(&additional), + dest_store.as_ref(), + spill_budget_bytes, + ) + .await + .unwrap(); + + assert_eq!(rebuilt, read_index_contents(dest_store.as_ref()).await); + } + + /// The spill and destination file schemas are declared from the existing + /// index while the keys come from the new stream, so a disagreement must be + /// rejected rather than written as a file whose schema lies about its keys. + #[tokio::test] + async fn test_update_rejects_a_mismatched_value_type() { + let (_src_dir, index) = build_label_list_segment(&sample_label_list_rows(64, 0)).await; + let (_dest_dir, dest_store) = test_util::index_store(); + + // The index above is List; feed the update List. + let mut list_builder = + arrow_array::builder::ListBuilder::new(arrow_array::builder::LargeStringBuilder::new()); + list_builder.values().append_value("a"); + list_builder.append(true); + let labels = list_builder.finish(); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, labels.data_type().clone(), true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(labels), + Arc::new(UInt64Array::from(vec![1u64 << 32])), + ], + ) + .unwrap(); + let new_data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(vec![Ok(batch)]), + )); + + let error = + update_label_list_index(index.as_ref(), new_data, dest_store.as_ref(), usize::MAX) + .await + .expect_err("a value-type mismatch must be rejected"); + let message = error.to_string(); + assert!( + message.contains("value type Utf8") && message.contains("items are LargeUtf8"), + "the error must name both value types, got: {message}" + ); + } + + /// An update whose new rows reuse labels already in the index must union + /// the row sets rather than replace them. + #[tokio::test] + async fn test_update_unions_rows_for_existing_labels() { + const FRAGMENT_1: u64 = 1 << 32; + let initial = sample_label_list_rows(200, 0); + // Same label vocabulary, different addresses. + let additional = sample_label_list_rows(200, FRAGMENT_1); + + let (_src_dir, index) = build_label_list_segment(&initial).await; + let (_dest_dir, dest_store) = test_util::index_store(); + update_label_list_index( + index.as_ref(), + sample_rows_to_stream(&additional), + dest_store.as_ref(), + 1, + ) + .await + .unwrap(); + + let updated = read_index_contents(dest_store.as_ref()).await; + let shared = updated + .labels + .iter() + .find(|(key, _)| key.as_deref() == Some("label-0000")) + .expect("shared label must be present"); + assert!( + shared.1.iter().any(|addr| *addr < FRAGMENT_1) + && shared.1.iter().any(|addr| *addr >= FRAGMENT_1), + "a label present in both old and new data must keep both row sets" + ); + } + + /// Every LabelList index written before spill-based builds landed has its + /// keys in arbitrary order on disk, because the old writer iterated a HashMap. + /// Nothing added here may assume sorted files: reading, updating and + /// merging such an index must all still be correct. + #[tokio::test] + async fn test_reads_updates_and_merges_an_unsorted_legacy_index() { + const FRAGMENT_1: u64 = 1 << 32; + + /// Write the same index the old HashMap-consuming writer would have: + /// identical contents, keys in an order that is not sorted. + async fn write_unsorted_index(rows: &[SampleRow], store: &dyn IndexStore) { + let sorted = build_label_list_index(rows, usize::MAX).await; + let mut writer = new_bitmap_batch_writer(store, BITMAP_LOOKUP_NAME, &DataType::Utf8) + .await + .unwrap(); + writer + .add_global_buffer( + LABEL_LIST_NULLS_METADATA_KEY.to_string(), + serialize_list_nulls(&RowAddrTreeMap::from_iter( + sorted.list_nulls.iter().copied(), + )) + .unwrap(), + ) + .await + .unwrap(); + // Reverse order is sorted-ness's clearest counterexample. + for (key, addrs) in sorted.labels.iter().rev() { + writer + .emit( + ScalarValue::Utf8(key.clone()), + &RowAddrTreeMap::from_iter(addrs.iter().copied()), + ) + .await + .unwrap(); + } + writer.finish().await.unwrap(); + } + + let initial = sample_label_list_rows(600, 0); + let expected_initial = build_label_list_index(&initial, usize::MAX).await; + + let (_legacy_dir, legacy_store) = test_util::index_store(); + write_unsorted_index(&initial, legacy_store.as_ref()).await; + + // Reading: the loaded index must match one built by the current path. + assert_eq!( + expected_initial, + read_index_contents(legacy_store.as_ref()).await, + "an unsorted file must carry the same contents" + ); + let legacy = LabelListIndex::load(legacy_store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + // Updating: streaming the unsorted index in as a merge input must still + // produce the same index as a full rebuild. + let additional = sample_label_list_rows(300, FRAGMENT_1); + let mut all = initial.clone(); + all.extend(additional.clone()); + let (_updated_dir, updated_store) = test_util::index_store(); + update_label_list_index( + legacy.as_ref(), + sample_rows_to_stream(&additional), + updated_store.as_ref(), + 1, + ) + .await + .unwrap(); + assert_eq!( + build_label_list_index(&all, usize::MAX).await, + read_index_contents(updated_store.as_ref()).await, + "update over an unsorted index diverged from a full rebuild" + ); + + // Merging: an unsorted segment must merge correctly with a sorted one. + let (_other_dir, other_index) = build_label_list_segment(&additional).await; + let (_merged_dir, merged_store) = test_util::index_store(); + merge_label_list_indices( + &[legacy, other_index], + merged_store.as_ref(), + None, + crate::progress::noop_progress(), + ) + .await + .unwrap(); + assert_eq!( + build_label_list_index(&all, usize::MAX).await, + read_index_contents(merged_store.as_ref()).await, + "merging an unsorted segment diverged from a full build" + ); + } } diff --git a/rust/lance-index/src/scalar/label_list/spill.rs b/rust/lance-index/src/scalar/label_list/spill.rs new file mode 100644 index 00000000000..e57f19b6b37 --- /dev/null +++ b/rust/lance-index/src/scalar/label_list/spill.rs @@ -0,0 +1,787 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Spillable aggregation for LabelList index builds. +//! +//! A LabelList index is a bitmap index over unnested list elements, so its keys +//! are derived after the scan and the scanner cannot pre-sort by them (hence +//! `TrainingOrdering::None`). That rules out the run-length streaming writer +//! plain bitmap indexes use, which requires value-sorted input. This builder +//! establishes the ordering itself: it accumulates into a sorted map, writes a +//! sorted spill file whenever the estimated size exceeds a byte budget, and +//! leaves the caller with a set of sorted files to k-way merge, reduced to a +//! bounded fan-in first so that the number of open cursors does not grow with +//! the number of spills. Modelled on +//! `NGramIndexBuilder`, which solves the same problem for the same reason; +//! see [`MAX_MERGE_FAN_IN`] for where the two diverge. +//! +//! # Aggregation-state contract +//! +//! `LANCE_LABEL_LIST_SPILL_BYTES` limits the estimated memory held by the mutable +//! label-to-row-set aggregation map. After an insertion brings the estimate to or +//! above the configured budget, the map is flushed and cleared. The map may +//! therefore exceed the budget by the most recent insertion, and the estimate is +//! deliberately conservative rather than an allocator-exact measurement. +//! +//! This is not an end-to-end build-memory limit. Memory outside the aggregation +//! budget includes scan batches, the `list_nulls` row set, source indexes and +//! their caches during update or segment merge, merge cursor batches, the bitmap +//! currently being merged, and the output writer. Concurrent builds add their +//! own working state as well. +//! +//! [`MAX_MERGE_FAN_IN`] separately bounds the number of cursor batches held by a +//! merge, not their total bytes: a batch can contain large serialized bitmaps. +//! `MAX_BUFFERED_BYTES` separately limits the output writer's buffered keys and +//! bitmaps. Neither is charged to the aggregation budget. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use arrow_schema::DataType; +use datafusion_common::ScalarValue; +use lance_core::cache::LanceCache; +use lance_core::utils::address::RowAddress; +use lance_core::utils::tempfile::TempDir; +use lance_core::{Error, Result}; +use lance_io::object_store::ObjectStore; +use lance_select::RowAddrTreeMap; +use tracing::warn; + +use crate::scalar::IndexStore; +use crate::scalar::bitmap::{ + BitmapBatchWriter, BitmapIndex, drain_sorted_bitmap_cursors, merge_index_maps, + new_bitmap_batch_writer, open_sorted_bitmap_cursors, +}; +use crate::scalar::btree::OrderableScalarValue; +use crate::scalar::lance_format::LanceIndexStore; + +/// Default estimated aggregation-state budget before spilling. +/// +/// This prevents the mutable label-to-row-set map from growing without limit as +/// label cardinality increases. It is not a limit on total build memory; see the +/// module-level aggregation-state contract. +const DEFAULT_SPILL_BUDGET_BYTES: usize = 512 * 1024 * 1024; + +/// Fixed cost charged for each distinct label, on top of the key's own size. +/// +/// Dominated by the inner `BTreeMap` inside the label's +/// `RowAddrTreeMap`: Rust's B=6 gives a leaf capacity of 11, and the whole node +/// is allocated on the first insert, so a label with a single row address pays +/// `16 + 11 * (4 + size_of::())` ~= 324 bytes for one entry. +/// On top of that come the outer map's value slot (~25 bytes amortised over a +/// full node), the roaring `containers` Vec, and its first container's store. +/// Rounded up from ~410 to leave room for allocator rounding, which none of the +/// above measures. +/// +/// Covers the label's *first* fragment, since the leaf node and roaring bitmap +/// above are what its first row address allocates. Every fragment after that is +/// charged [`FRAGMENT_OVERHEAD_BYTES`]. +/// +/// Deliberately generous, and measured rather than guessed: LabelList's problem +/// case is a very large number of labels with few rows each, so under-charging +/// the per-label overhead is exactly how a byte budget silently fails to +/// constrain the accumulator. `deep_size_of` is not usable here -- it walks +/// logical contents, not allocation capacity, and reports about half the real +/// cost for this shape. +/// `test_label_overhead_covers_a_single_row_labels_allocations` pins the +/// arithmetic so a change to `RowAddrSelection` cannot silently invalidate it. +const LABEL_OVERHEAD_BYTES: usize = 512; + +/// Cost charged for a row address landing in a fragment the label already has. +/// +/// Roaring's sparse store is a `Vec`, so two bytes per value, and a `Vec` +/// grows by doubling, so its capacity can be twice its length -- four bytes is +/// therefore the ceiling, not an average. A dense set costs far less: it becomes +/// a bitmap store, a fixed 8 KiB per 65,536-value block, or about an eighth of a +/// byte per row. The estimate errs toward spilling early rather than late. +const ROW_ADDR_COST_BYTES: usize = 4; + +/// Cost charged for a row address landing in a fragment the label has not seen. +/// +/// A `RowAddrTreeMap` is a `BTreeMap` keyed by fragment, +/// so a new fragment is not a cheap increment: it takes a slot in that map, +/// roughly 29 bytes amortised over a full leaf node, and allocates a whole new +/// `RoaringBitmap` -- its `containers` `Vec`, plus that container's own store. +/// Measured at 70 bytes per fragment by `deep_size_of`, which counts neither +/// allocation capacity nor malloc overhead; scaled by the same factor +/// [`LABEL_OVERHEAD_BYTES`] uses over its own measurement, and rounded. +/// +/// This is the difference between a label whose rows sit in one fragment and one +/// spread across thousands. Charging [`ROW_ADDR_COST_BYTES`] for both under-counts +/// the spread case by more than an order of magnitude, which would make the +/// aggregation-state budget ineffective. +const FRAGMENT_OVERHEAD_BYTES: usize = 192; + +/// Name of the spill file the index being updated is streamed into. Distinct +/// from the numbered build spills and from the intermediate merge files so the +/// three never collide. +const EXISTING_INDEX_FILE_NAME: &str = "label-list-existing.lance"; + +/// Prefix for the intermediate files written when the spill count exceeds +/// [`MAX_MERGE_FAN_IN`]. +const MERGE_FILE_PREFIX: &str = "label-list-merge-"; + +/// Maximum number of spill files merged in a single k-way pass. +/// +/// Without a cap, the number of spill files is `total_bytes / budget_bytes`, and +/// merging them all at once holds one `MERGE_ROWS_PER_CHUNK` batch open per file. +/// Capping fan-in makes the cursor count independent of the spill count, at the +/// cost of extra passes over the spill. It does not impose a byte limit on those +/// cursor batches. +/// +/// The bound is in rows, not bytes: 64 files x 512 serialized bitmaps. That is +/// small whenever the bitmaps are (the high-cardinality case, which is also the +/// case that produces many files), but a column with a few very popular labels +/// can serialize those to megabytes each, and a run of them landing in one chunk +/// is not bounded by `budget_bytes`. Making the chunk byte-aware is the fix for +/// that; it belongs in `BitmapShardCursor`, which every merge shares. +/// +/// A pass rewrites the spilled bytes exactly once and divides the file count by +/// the fan-in, so reducing any realistic spill takes one or two passes. +/// `NGramIndexBuilder` instead merges each flush back into a single per-worker +/// file, which holds fan-in at one but rewrites the whole spill on every flush +/// -- quadratic in the number of flushes, where this is logarithmic. +const MAX_MERGE_FAN_IN: usize = 64; + +/// Read the configured budget, rejecting a value that cannot be honoured. +/// +/// Rejected rather than defaulted: someone who sets this knob is limiting +/// aggregation state deliberately, and silently substituting +/// [`DEFAULT_SPILL_BUDGET_BYTES`] would use a different limit without warning. +/// Zero is rejected by [`LabelListSpillBuilder::new_local`], which every budget +/// reaches. +pub(super) fn default_spill_budget_bytes() -> Result { + match std::env::var("LANCE_LABEL_LIST_SPILL_BYTES") { + Ok(raw) => parse_spill_budget_bytes(&raw), + Err(std::env::VarError::NotPresent) => Ok(DEFAULT_SPILL_BUDGET_BYTES), + Err(std::env::VarError::NotUnicode(_)) => Err(Error::invalid_input( + "LANCE_LABEL_LIST_SPILL_BYTES is not valid unicode; expected a whole number of bytes" + .to_string(), + )), + } +} + +/// Split out of [`default_spill_budget_bytes`] so that the rejection is testable +/// without mutating the process environment, which is unsound to do while other +/// tests are running. +fn parse_spill_budget_bytes(raw: &str) -> Result { + raw.trim().parse().map_err(|_| { + Error::invalid_input(format!( + "LANCE_LABEL_LIST_SPILL_BYTES must be a whole number of bytes, got '{raw}'" + )) + }) +} + +/// The sorted spill files produced by a [`LabelListSpillBuilder`], together +/// with the store they live in. +pub(super) struct LabelListSpills { + store: Arc, + value_type: DataType, + files: Vec, + /// Serial number for the next intermediate merge file, so that repeated + /// fan-in reduction passes never reuse a name. + next_merge_id: usize, + /// Kept alive so the spill files outlive the builder that wrote them. + _tmpdir: Option, +} + +impl LabelListSpills { + #[cfg(test)] + pub(super) fn files(&self) -> &[String] { + &self.files + } + + #[cfg(test)] + pub(super) fn is_empty(&self) -> bool { + self.files.is_empty() + } + + /// Add the index being updated as one more sorted merge input. + /// + /// It is streamed out through its `index_map`, one bitmap at a time, so the + /// bitmap payload for every key is not materialized at once. The source + /// `index_map` and any bitmaps retained by its cache remain outside the + /// aggregation-state budget. + /// `NGramIndexBuilder::merge_old_index` feeds the previous index into its + /// merge the same way. + /// + /// This rewrites the index into scratch rather than merging from the index + /// file in place, which would save a full write and read. Two things block + /// that: [`open_sorted_bitmap_cursors`] opens every input from one store, + /// and the old index lives in the index store rather than local scratch; and + /// cursors require key-sorted files, which LabelList indexes written before + /// spill-based builds are not. The rewrite is also not where this path's cost + /// sits -- `merge_index_maps` issues one single-row read per key of the old + /// index, which dominates either way. + pub(super) async fn add_existing_index(&mut self, index: &Arc) -> Result<()> { + let file_name = EXISTING_INDEX_FILE_NAME.to_string(); + let mut writer = + new_bitmap_batch_writer(self.store.as_ref(), &file_name, index.value_type()).await?; + merge_index_maps(std::slice::from_ref(index), None, &mut writer, None).await?; + writer.finish().await?; + self.files.push(file_name); + Ok(()) + } + + /// Merge every spill file into `writer` as one ascending `(key, bitmap)` + /// stream, unioning the row sets of labels that spilled more than once. + /// + /// Reduces the file set to at most [`MAX_MERGE_FAN_IN`] entries first, so the + /// number of cursors open at once is independent of how many times the + /// builder spilled. This bounds cursor count, not cursor bytes or total + /// operation memory. + pub(super) async fn merge_into(&mut self, writer: &mut BitmapBatchWriter) -> Result<()> { + while self.files.len() > MAX_MERGE_FAN_IN { + let inputs = std::mem::take(&mut self.files); + for group in inputs.chunks(MAX_MERGE_FAN_IN) { + // A trailing group of one is already a sorted file; merging it + // into a new one would only copy it. + if let [only] = group { + self.files.push(only.clone()); + continue; + } + let merged = self.merge_group(group).await?; + self.files.push(merged); + } + } + + let (mut cursors, mut heap, _) = + open_sorted_bitmap_cursors(self.store.as_ref(), &self.files).await?; + drain_sorted_bitmap_cursors(&mut cursors, &mut heap, writer, None).await + } + + /// Merge `group` into one new sorted spill file and delete the inputs. + async fn merge_group(&mut self, group: &[String]) -> Result { + let file_name = format!("{MERGE_FILE_PREFIX}{}.lance", self.next_merge_id); + self.next_merge_id += 1; + + let mut writer = + new_bitmap_batch_writer(self.store.as_ref(), &file_name, &self.value_type).await?; + let (mut cursors, mut heap, _) = + open_sorted_bitmap_cursors(self.store.as_ref(), group).await?; + drain_sorted_bitmap_cursors(&mut cursors, &mut heap, &mut writer, None).await?; + writer.finish().await?; + // Release the readers before unlinking what they were reading. + drop(cursors); + + // The merged file supersedes its inputs, so dropping them now holds + // scratch usage at roughly one copy of the spill rather than one per pass. + for name in group { + if let Err(error) = self.store.delete_index_file(name).await { + warn!( + "Failed to delete intermediate label list spill file '{}': {}. \ + This does not affect the built index, but the spill file \ + may need manual cleanup.", + name, error + ); + } + } + + Ok(file_name) + } +} + +/// Accumulates `(label, row address)` pairs into sorted spill files while +/// limiting the estimated memory held by the in-memory aggregation map. +pub(super) struct LabelListSpillBuilder { + spill_store: Arc, + tmpdir: Option, + value_type: DataType, + budget_bytes: usize, + state: BTreeMap, + /// Running estimate of the memory held by `state`. + estimated_bytes: usize, + spill_files: Vec, +} + +impl LabelListSpillBuilder { + /// Build against a caller-supplied store. Used by tests; production callers + /// want [`Self::new_local`] so that spilling stays off object storage. + #[cfg(test)] + pub(super) fn new( + spill_store: Arc, + value_type: DataType, + budget_bytes: usize, + ) -> Self { + Self { + spill_store, + tmpdir: None, + value_type, + budget_bytes, + state: BTreeMap::new(), + estimated_bytes: 0, + spill_files: Vec::new(), + } + } + + /// Build against a local temporary directory. + /// + /// Spill files are scratch: written once, read back once, then removed with + /// the directory when the resulting [`LabelListSpills`] drops. Putting them + /// in the index store would put them on object storage, where the round trip + /// costs far more than the local scratch disk the indexer is already + /// provisioned with. `NGramIndexBuilder` spills the same way for the same + /// reason. + /// + /// The directory comes from `std::env::temp_dir()`, so `TMPDIR` chooses it, + /// and it needs real capacity: a build spills roughly one copy of the + /// label-to-row-set map, and an update spills that plus a full copy of the + /// index being updated (see [`LabelListSpills::add_existing_index`]). Where + /// `/tmp` is a memory-backed tmpfs -- common on container images -- scratch + /// still consumes system memory and defeats the purpose of moving aggregation + /// state out of RAM, so point `TMPDIR` at disk. + pub(super) fn new_local(value_type: DataType, budget_bytes: usize) -> Result { + if budget_bytes == 0 { + return Err(Error::invalid_input( + "LabelList spill budget must be at least one byte, got 0".to_string(), + )); + } + let tmpdir = TempDir::try_new()?; + let spill_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + Ok(Self { + spill_store, + tmpdir: Some(tmpdir), + value_type, + budget_bytes, + state: BTreeMap::new(), + estimated_bytes: 0, + spill_files: Vec::new(), + }) + } + + fn spill_filename(id: usize) -> String { + format!("label-list-spill-{id}.lance") + } + + pub(super) async fn insert(&mut self, key: ScalarValue, row_addr: u64) -> Result<()> { + let fragment = RowAddress::from(row_addr).fragment_id(); + + // `entry` takes the key by value and descends once. A `get_mut` miss + // followed by `insert` walks the tree twice, on the path that dominates + // the high-cardinality case this builder exists for. + match self.state.entry(OrderableScalarValue(key)) { + std::collections::btree_map::Entry::Occupied(mut entry) => { + let bitmap = entry.get_mut(); + let new_fragment = bitmap.get(&fragment).is_none(); + bitmap.insert(row_addr); + self.estimated_bytes += if new_fragment { + FRAGMENT_OVERHEAD_BYTES + ROW_ADDR_COST_BYTES + } else { + ROW_ADDR_COST_BYTES + }; + } + std::collections::btree_map::Entry::Vacant(entry) => { + self.estimated_bytes += + entry.key().0.size() + LABEL_OVERHEAD_BYTES + ROW_ADDR_COST_BYTES; + let mut bitmap = RowAddrTreeMap::default(); + bitmap.insert(row_addr); + entry.insert(bitmap); + } + } + + if self.estimated_bytes >= self.budget_bytes { + self.flush().await?; + } + Ok(()) + } + + async fn flush(&mut self) -> Result<()> { + if self.state.is_empty() { + return Ok(()); + } + let file_name = Self::spill_filename(self.spill_files.len()); + let mut writer = + new_bitmap_batch_writer(self.spill_store.as_ref(), &file_name, &self.value_type) + .await?; + + // `BTreeMap` iteration is already in key order, which is exactly what + // the downstream k-way merge requires of each spill file. + for (key, bitmap) in std::mem::take(&mut self.state) { + writer.emit(key.0, &bitmap).await?; + } + writer.finish().await?; + + self.estimated_bytes = 0; + self.spill_files.push(file_name); + Ok(()) + } + + pub(super) async fn finish(mut self) -> Result { + self.flush().await?; + Ok(LabelListSpills { + store: self.spill_store, + value_type: self.value_type, + files: self.spill_files, + next_merge_id: 0, + _tmpdir: self.tmpdir, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_schema::DataType; + use datafusion_common::ScalarValue; + use lance_select::RowAddrTreeMap; + use rstest::rstest; + + use crate::scalar::bitmap::test_util::{self, row_addrs}; + + /// Read every spill file back in file order, which is the order the + /// downstream k-way merge relies on being sorted. + async fn read_all_sorted( + store: &dyn IndexStore, + files: &[String], + ) -> Vec<(String, RowAddrTreeMap)> { + let mut out = Vec::new(); + for name in files { + out.extend( + test_util::read_key_bitmaps(store, name) + .await + .into_iter() + .map(|(key, bitmap)| { + ( + key.expect("spill files under test have no null keys"), + bitmap, + ) + }), + ); + } + out + } + + /// A budget of one byte forces a flush on essentially every key, which is + /// the cheapest way to exercise the multi-spill path deterministically. + #[tokio::test] + async fn test_spills_when_over_budget() { + let (_tmpdir, store) = test_util::index_store(); + let mut builder = LabelListSpillBuilder::new(store.clone(), DataType::Utf8, 1); + + for (key, row) in [("b", 1u64), ("a", 2), ("b", 3), ("c", 4)] { + builder + .insert(ScalarValue::Utf8(Some(key.to_string())), row) + .await + .unwrap(); + } + + let spills = builder.finish().await.unwrap(); + let files = spills.files(); + assert!( + files.len() > 1, + "expected multiple spill files, got {files:?}" + ); + + // Every individual spill file must be sorted by key on its own. + for name in files { + let contents = read_all_sorted(store.as_ref(), std::slice::from_ref(name)).await; + let keys: Vec<&String> = contents.iter().map(|(k, _)| k).collect(); + let mut sorted = keys.clone(); + sorted.sort(); + assert_eq!(keys, sorted, "spill file {name} is not sorted by key"); + } + } + + #[tokio::test] + async fn test_single_spill_when_under_budget() { + let (_tmpdir, store) = test_util::index_store(); + let mut builder = LabelListSpillBuilder::new(store.clone(), DataType::Utf8, 1 << 30); + + for (key, row) in [("b", 1u64), ("a", 2), ("b", 3)] { + builder + .insert(ScalarValue::Utf8(Some(key.to_string())), row) + .await + .unwrap(); + } + + let spills = builder.finish().await.unwrap(); + assert_eq!(spills.files().len(), 1); + + let contents = read_all_sorted(store.as_ref(), spills.files()).await; + let keys: Vec = contents.iter().map(|(k, _)| k.clone()).collect(); + assert_eq!(keys, vec!["a", "b"], "keys are unioned within a spill"); + assert_eq!( + row_addrs(&contents[1].1), + vec![1, 3], + "duplicate keys must union their row sets" + ); + } + + #[tokio::test] + async fn test_no_spill_files_when_nothing_inserted() { + let (_tmpdir, store) = test_util::index_store(); + let builder = LabelListSpillBuilder::new(store, DataType::Utf8, 1 << 30); + assert!(builder.finish().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn test_merge_spills_unions_duplicate_keys() { + let (_tmpdir, store) = test_util::index_store(); + let mut builder = LabelListSpillBuilder::new(store.clone(), DataType::Utf8, 1); + for (key, row) in [("b", 1u64), ("a", 2), ("b", 3), ("c", 4)] { + builder + .insert(ScalarValue::Utf8(Some(key.to_string())), row) + .await + .unwrap(); + } + let mut spills = builder.finish().await.unwrap(); + assert!(spills.files().len() > 1); + + let (_out_tmpdir, out_store) = test_util::index_store(); + let mut writer = + new_bitmap_batch_writer(out_store.as_ref(), "merged.lance", &DataType::Utf8) + .await + .unwrap(); + spills.merge_into(&mut writer).await.unwrap(); + writer.finish().await.unwrap(); + + let merged = read_all_sorted(out_store.as_ref(), &["merged.lance".to_string()]).await; + let keys: Vec = merged.iter().map(|(k, _)| k.clone()).collect(); + assert_eq!( + keys, + vec!["a", "b", "c"], + "each key must appear exactly once, in ascending order" + ); + assert_eq!(row_addrs(&merged[1].1), vec![1, 3]); + } + + /// The per-fragment charge must cover what touching a new fragment allocates: + /// a slot in the label's inner map plus a whole new roaring bitmap. Pinned + /// alongside the per-label arithmetic for the same reason -- the aggregation + /// estimate is useful only while these constants track the structures they + /// stand for. + #[test] + fn test_fragment_overhead_covers_a_new_fragments_allocations() { + const BTREE_LEAF_CAPACITY: usize = 11; + // One slot in the label's fragment map, amortised over a full leaf node. + let inner_slot = (2 * std::mem::size_of::() + + BTREE_LEAF_CAPACITY + * (std::mem::size_of::() + + std::mem::size_of::())) + / BTREE_LEAF_CAPACITY; + assert!( + FRAGMENT_OVERHEAD_BYTES > inner_slot, + "FRAGMENT_OVERHEAD_BYTES={FRAGMENT_OVERHEAD_BYTES} must exceed the map \ + slot alone ({inner_slot}); a new fragment also allocates a RoaringBitmap" + ); + } + + /// A label whose rows are spread across many fragments must be charged for + /// them. The estimate used to charge ROW_ADDR_COST_BYTES per row address + /// however it landed, so one hot label over N fragments was under-counted by + /// more than an order of magnitude and ran the aggregation map past its + /// configured target. + #[tokio::test] + async fn test_budget_accounts_for_a_label_spread_across_fragments() { + use lance_core::deepsize::DeepSizeOf; + + const FRAGMENTS: u64 = 5_000; + + let (_tmpdir, store) = test_util::index_store(); + // A budget nothing can reach, so `estimated_bytes` is the running + // accounting for everything the map still holds. + let mut builder = LabelListSpillBuilder::new(store, DataType::Utf8, usize::MAX); + for fragment in 0..FRAGMENTS { + builder + .insert( + ScalarValue::Utf8(Some("hot".to_string())), + (fragment << 32) | 1, + ) + .await + .unwrap(); + } + + let real: usize = builder + .state + .values() + .map(|bitmap| bitmap.deep_size_of()) + .sum(); + assert!( + real > 0, + "the fixture must actually hold something to measure" + ); + // `deep_size_of` counts neither allocation capacity nor malloc overhead, + // so it is a floor on the true cost. The estimate must clear it. + assert!( + builder.estimated_bytes >= real, + "estimated {} bytes for a label across {FRAGMENTS} fragments, but it \ + holds at least {real}", + builder.estimated_bytes + ); + } + + /// The per-label charge must cover what a single-row label really allocates, + /// the largest part of which is the inner `BTreeMap`'s leaf node -- allocated + /// at full capacity to hold one entry. Pinned so that a change to + /// `RowAddrSelection`'s size cannot silently invalidate the budget. + #[test] + fn test_label_overhead_covers_a_single_row_labels_allocations() { + // Rust's BTreeMap uses B=6, so a leaf holds 2B-1 = 11 entries and the + // whole node is allocated on the first insert. + const BTREE_LEAF_CAPACITY: usize = 11; + let inner_leaf_node = 2 * std::mem::size_of::() + + BTREE_LEAF_CAPACITY + * (std::mem::size_of::() + + std::mem::size_of::()); + let outer_value_slot = std::mem::size_of::(); + assert!( + LABEL_OVERHEAD_BYTES >= inner_leaf_node + outer_value_slot, + "LABEL_OVERHEAD_BYTES={LABEL_OVERHEAD_BYTES} must cover the inner leaf \ + node ({inner_leaf_node}) plus the outer value slot \ + ({outer_value_slot}); a budget that under-charges these does not \ + constrain the accumulator" + ); + } + + /// A malformed budget must be rejected rather than silently replaced by the + /// default: someone who set the knob specifically to limit aggregation state + /// should not unknowingly get `DEFAULT_SPILL_BUDGET_BYTES` instead. + #[rstest] + #[case::size_suffix("64M")] + #[case::not_a_number("abc")] + #[case::negative("-1")] + #[case::empty("")] + fn test_spill_budget_rejects_a_malformed_value(#[case] raw: &str) { + let error = parse_spill_budget_bytes(raw).expect_err("must be rejected"); + assert!( + error.to_string().contains("LANCE_LABEL_LIST_SPILL_BYTES"), + "the error must name the variable so it can be fixed, got: {error}" + ); + } + + #[test] + fn test_spill_budget_accepts_a_plain_byte_count() { + assert_eq!(parse_spill_budget_bytes(" 4096 ").unwrap(), 4096); + } + + /// Zero constrains no aggregation state, so it is an error rather than + /// something to clamp -- and it does not merely spill often: without the + /// guard it reaches `lance-io` and panics there. + #[test] + fn test_spill_builder_rejects_a_zero_budget() { + let Err(error) = LabelListSpillBuilder::new_local(DataType::Utf8, 0) else { + panic!("a zero budget must be rejected") + }; + assert!( + error.to_string().contains("at least one byte"), + "got: {error}" + ); + } + + /// More spill files than the merge may open at once must be reduced to a + /// bounded fan-in first, otherwise the number of resident cursor batches + /// grows with the number of spills. + #[tokio::test] + async fn test_merge_reduces_fan_in_to_a_bound() { + let (_tmpdir, store) = test_util::index_store(); + let mut builder = LabelListSpillBuilder::new(store.clone(), DataType::Utf8, 1); + + // A one-byte budget flushes on every insert, so this is one spill file + // per key -- comfortably past MAX_MERGE_FAN_IN. + let keys: Vec = (0..MAX_MERGE_FAN_IN * 3) + .map(|i| format!("label-{i:04}")) + .collect(); + for (row, key) in keys.iter().enumerate() { + builder + .insert(ScalarValue::Utf8(Some(key.clone())), row as u64) + .await + .unwrap(); + } + + let mut spills = builder.finish().await.unwrap(); + assert!( + spills.files().len() > MAX_MERGE_FAN_IN, + "fixture must spill more files than the merge may open at once, got {}", + spills.files().len() + ); + + let (_out_tmpdir, out_store) = test_util::index_store(); + let mut writer = + new_bitmap_batch_writer(out_store.as_ref(), "merged.lance", &DataType::Utf8) + .await + .unwrap(); + spills.merge_into(&mut writer).await.unwrap(); + writer.finish().await.unwrap(); + + assert!( + spills.files().len() <= MAX_MERGE_FAN_IN, + "the final merge opened {} files at once, above the {MAX_MERGE_FAN_IN} bound", + spills.files().len() + ); + + // Each pass deletes the files it consumed, so scratch holds one copy of + // the spill rather than one per pass. + let live: Vec = store + .list_files_with_sizes() + .await + .unwrap() + .into_iter() + .map(|file| file.path) + .collect(); + assert_eq!( + live.len(), + spills.files().len(), + "consumed spill files must be deleted, found {live:?}" + ); + + let merged = read_all_sorted(out_store.as_ref(), &["merged.lance".to_string()]).await; + let merged_keys: Vec = merged.iter().map(|(key, _)| key.clone()).collect(); + assert_eq!( + merged_keys, keys, + "every key must survive the multi-pass merge, in ascending order" + ); + for (row, (_, bitmap)) in merged.iter().enumerate() { + assert_eq!(row_addrs(bitmap), vec![row as u64]); + } + } + + /// Nulls are legitimate label values: a null element inside a list survives + /// unnesting as a null key, so the merge must carry it like any other. + #[tokio::test] + async fn test_merge_spills_preserves_null_keys() { + let (_tmpdir, store) = test_util::index_store(); + let mut builder = LabelListSpillBuilder::new(store.clone(), DataType::Utf8, 1); + builder + .insert(ScalarValue::Utf8(Some("a".to_string())), 1) + .await + .unwrap(); + builder.insert(ScalarValue::Utf8(None), 2).await.unwrap(); + builder.insert(ScalarValue::Utf8(None), 3).await.unwrap(); + let mut spills = builder.finish().await.unwrap(); + + let (_out_tmpdir, out_store) = test_util::index_store(); + let mut writer = + new_bitmap_batch_writer(out_store.as_ref(), "merged.lance", &DataType::Utf8) + .await + .unwrap(); + spills.merge_into(&mut writer).await.unwrap(); + writer.finish().await.unwrap(); + + let reader = out_store.open_index_file("merged.lance").await.unwrap(); + let batch = reader.read_range(0..reader.num_rows(), None).await.unwrap(); + let keys: Vec = (0..batch.num_rows()) + .map(|idx| ScalarValue::try_from_array(batch.column(0), idx).unwrap()) + .collect(); + assert_eq!( + keys, + vec![ + ScalarValue::Utf8(None), + ScalarValue::Utf8(Some("a".to_string())) + ], + "null sorts first and must survive the merge" + ); + + let bitmaps = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let null_rows = RowAddrTreeMap::deserialize_from(bitmaps.value(0)).unwrap(); + assert_eq!(row_addrs(&null_rows), vec![2, 3]); + } +}