From f4a4ba81aeb1a9df1e724de02c5b7d5f412b806e Mon Sep 17 00:00:00 2001 From: Vivek Date: Thu, 27 Aug 2026 09:52:16 -0700 Subject: [PATCH] fix(label_list): bound memory in index build, update, and merge LabelList materialized its whole label-to-bitmap map in memory during build, update, merge and remap, so peak memory scaled with distinct-label cardinality and large list columns led to out of memory errors. Build and update now accumulate into a byte-budgeted sorted map (LANCE_LABEL_LIST_SPILL_BYTES, default 512 MiB), spill sorted files to local scratch and k-way merge them; update rewrites the existing index into one more sorted scratch file so it joins the same merge. This also changes the BITMAP index. The shared bitmap writer's flush threshold drops from ~2 GiB to 32 MiB and now charges keys as well as serialized bitmaps, fixing a latent overflow where a high-cardinality column with tiny bitmaps could exceed the keys column's i32 offsets before ever tripping a flush. --- rust/lance-index/src/scalar/bitmap.rs | 904 +++++++++++----- rust/lance-index/src/scalar/label_list.rs | 985 ++++++++++++++++-- .../src/scalar/label_list/spill.rs | 787 ++++++++++++++ 3 files changed, 2361 insertions(+), 315 deletions(-) create mode 100644 rust/lance-index/src/scalar/label_list/spill.rs 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]); + } +}