From 1ae56870c8246614e2cb7260889b331a7358c9a2 Mon Sep 17 00:00:00 2001 From: Zeljko Mihaljcic <7150613+zehiko@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:09:51 +0200 Subject: [PATCH 1/7] perf(index): merge Bitmap index segments instead of rescanning the dataset `merge_scalar_indices` lists only `IndexType::BTree` and `IndexType::NGram` in `has_segment_merge_primitive`, so a Bitmap index with more than one selected segment falls through to `rebuild_scalar_segment` -- a full rescan and retrain over every covered fragment. Bitmap declares `TrainingCriteria::new(TrainingOrdering::Values)`, so that rescan is wrapped in a blocking spilling `SortExec` against `FairSpillPool(LANCE_MEM_POOL_SIZE)`. On a large low-cardinality column that means sorting every row of the column to produce an index of a few megabytes, and failing outright once the pool is exhausted. Wire Bitmap into the same N:1 segment-merge path BTree uses. The old side of `build_index_map` becomes `OldSegments`: K source segments presented as one ascending-key stream of postings, a min-heap over each segment's in-memory `index_map` keys, with each segment's bitmaps read one key at a time through `load_bitmap`. `merge_index_maps` is re-expressed over the same type, so the single-segment update path, the LabelList merge and the new `BitmapIndex::merge_segments` share one merge-join. Two choices are load-bearing: - Keys come from `index_map`, a `BTreeMap`, so key order does not depend on the on-disk row order. Bitmap lookup files are not reliably key-sorted -- an update appends an old-only null row last, and LabelList files written before spill-based builds are unsorted -- so a merge that read the files sequentially and assumed sorted input would silently drop postings. - `load_bitmap` is the accessor that applies fragment-reuse remapping. Reading the lookup file's raw bytes would bypass it and emit row addresses pointing at fragments compaction has already retired, which under deferred index remapping is silent data loss. Peak memory on the old side is one posting per segment for the current key plus the union being emitted, on top of each segment's resident `index_map`. Per-segment `OldIndexDataFilter`s are applied to each source posting, so compaction-retired fragments and stable-row-id rewrites are dropped exactly as on the BTree path, and a filter that keeps nothing skips the segment's bitmap reads (its `index_map` is still loaded when the segment is opened). `Dataset::merge_existing_index_segments` now routes bitmap groups through the same primitive, which gives it the per-segment filtering it lacked. `merge_bitmap_indices`, which it replaces, is deprecated and forwards to `merge_segments` with no new data and no filters. --- rust/lance-index/src/scalar/bitmap.rs | 801 +++++++++++++++++----- rust/lance-index/src/scalar/btree.rs | 2 +- rust/lance-index/src/scalar/label_list.rs | 17 +- rust/lance/src/index/append.rs | 348 +++++++++- rust/lance/src/index/scalar/bitmap.rs | 69 +- rust/lance/src/index/scalar/btree.rs | 6 +- 6 files changed, 1029 insertions(+), 214 deletions(-) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 8f5b12feced..ae3ff7f7f2b 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -18,6 +18,7 @@ use arrow_schema::{DataType, Field, Schema}; use async_trait::async_trait; use bytes::Bytes; use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_common::ScalarValue; use futures::{StreamExt, TryStreamExt, stream}; use lance_core::deepsize::DeepSizeOf; @@ -39,7 +40,8 @@ use tracing::{instrument, warn}; use super::{AnyQuery, IndexFile, IndexStore, ScalarIndex, SearchOptions}; use super::{ - BuiltinIndexType, SargableQuery, ScalarIndexParams, SearchResult, btree::OrderableScalarValue, + BuiltinIndexType, SargableQuery, ScalarIndexParams, SearchResult, + btree::{OrderableScalarValue, filter_keeps_nothing}, }; use crate::pbold; use crate::{Index, IndexType, metrics::MetricsCollector}; @@ -537,6 +539,80 @@ impl BitmapIndex { pub(crate) fn value_type(&self) -> &DataType { &self.value_type } + + /// Merge N source bitmap segments plus an additional `new_data` stream into a + /// single bitmap index under `dest_store`, without re-reading the dataset. + /// + /// `old_data_filters` carries one filter per source segment, in the same + /// order. A segment whose filter keeps nothing contributes no postings, so + /// none of its bitmaps are read. + pub async fn merge_segments( + segments: &[Arc], + new_data: SendableRecordBatchStream, + dest_store: &dyn IndexStore, + old_data_filters: &[Option], + ) -> Result { + let Some(first) = segments.first() else { + return Err(Error::invalid_input( + "cannot merge bitmap index without at least one source segment".to_string(), + )); + }; + + if old_data_filters.len() != segments.len() { + return Err(Error::invalid_input(format!( + "Bitmap merge: expected one old-data filter per source segment \ + (segments={}, filters={})", + segments.len(), + old_data_filters.len() + ))); + } + + for segment in segments.iter().skip(1) { + if segment.value_type != first.value_type { + return Err(Error::invalid_input(format!( + "cannot merge bitmap segments with different value types ({:?} vs {:?})", + first.value_type, segment.value_type + ))); + } + } + + let new_schema = new_data.schema(); + let new_value_type = new_schema + .field(new_schema.index_of(VALUE_COLUMN_NAME)?) + .data_type(); + if new_value_type != &first.value_type { + return Err(Error::invalid_input(format!( + "Bitmap merge: new_data value column type {:?} does not match \ + segment value type {:?}", + new_value_type, first.value_type + ))); + } + + let old_segments = segments + .iter() + .zip(old_data_filters) + .filter(|(_, filter)| !filter_keeps_nothing(filter)) + .map(|(segment, filter)| OldSegment { + index: segment.as_ref(), + filter: filter.as_ref(), + }) + .collect(); + + let file = BitmapIndexPlugin::streaming_build_and_write( + new_data, + old_segments, + dest_store, + BITMAP_LOOKUP_NAME, + ) + .await?; + + Ok(CreatedIndex { + index_details: prost_types::Any::from_msg(&pbold::BitmapIndexDetails::default()) + .unwrap(), + index_version: BITMAP_INDEX_VERSION, + files: vec![file], + }) + } } impl DeepSizeOf for BitmapIndex { @@ -858,10 +934,12 @@ impl ScalarIndex for BitmapIndex { ) -> Result { let file = BitmapIndexPlugin::streaming_build_and_write( new_data, - Some(self), + vec![OldSegment { + index: self, + filter: old_data_filter.as_ref(), + }], dest_store, BITMAP_LOOKUP_NAME, - old_data_filter.as_ref(), ) .await?; @@ -1316,6 +1394,135 @@ fn retain_valid<'a>( } } +/// One source segment on the old side of a bitmap build, with the filter that +/// decides which of its rows survive. `None` keeps every row. +pub(crate) struct OldSegment<'a> { + pub(crate) index: &'a BitmapIndex, + pub(crate) filter: Option<&'a super::OldIndexDataFilter>, +} + +/// The already-indexed side of a bitmap build: K source segments presented as a +/// single ascending-key stream of filtered postings. +/// +/// Drives each source through its `index_map` -- a sorted `BTreeMap` rebuilt at +/// load time -- rather than through the rows of its file. Bitmap lookup files are +/// not reliably key-sorted on disk: LabelList files written before spill-based +/// builds landed are unsorted, and an update appends an old-only null row last. +/// `index_map` order can be trusted for old and new files alike, which is what +/// makes this work without an index version bump. +/// +/// Bitmaps are read one key at a time through [`BitmapIndex::load_bitmap`], the +/// accessor that applies fragment-reuse remapping. Reading the lookup file's raw +/// bytes would bypass it and emit row addresses pointing at fragments compaction +/// has already retired, which under deferred index remapping is silent data loss. +/// +/// 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)`. Entries borrow their key from the source's +/// `index_map` rather than cloning it: seeding the heap touches every source key, +/// so cloning would cost an allocation per key per source for string keys. +/// +/// The transient state is the merged bitmap for the current key plus one loaded +/// bitmap per participating source. Each source's `index_map`, any bitmaps the +/// index cache retains, and the output writer stay outside it. +struct OldSegments<'a> { + segments: Vec>, + key_iters: Vec>, + heap: BinaryHeap>, + /// Source entries drained so far: one per (source, key) pair, the unit + /// [`merge_source_entry_count`] declares up front. + consumed: u64, + null_taken: bool, +} + +impl<'a> OldSegments<'a> { + fn new(segments: Vec>) -> Self { + let mut key_iters: Vec<_> = segments + .iter() + .map(|segment| segment.index.index_map.keys()) + .collect(); + let mut heap = BinaryHeap::with_capacity(segments.len()); + for (source_idx, keys) in key_iters.iter_mut().enumerate() { + if let Some(key) = keys.next() { + heap.push(Reverse((key, source_idx))); + } + } + Self { + segments, + key_iters, + heap, + consumed: 0, + null_taken: false, + } + } + + fn consumed(&self) -> u64 { + self.consumed + } + + /// Union every source's filtered posting for the smallest pending key, if + /// that key satisfies `predicate`. `None` when no source has a pending key or + /// the smallest one fails the predicate; the postings stay pending then. + /// + /// Drains only the sources positioned on that key. A source's next key is + /// strictly greater, so re-pushing it cannot re-enter the same key. + async fn take_smallest_if( + &mut self, + predicate: impl FnOnce(&OrderableScalarValue) -> bool, + ) -> Result> { + let Some(Reverse((key, _))) = self.heap.peek().copied() else { + return Ok(None); + }; + if !predicate(key) { + return Ok(None); + } + + let mut merged = RowAddrTreeMap::default(); + while let Some(Reverse((next, source_idx))) = self.heap.peek().copied() { + if next != key { + break; + } + self.heap.pop(); + self.consumed += 1; + let segment = &self.segments[source_idx]; + let loaded = segment.index.load_bitmap(key, None).await?; + merged |= &*retain_valid(Cow::Borrowed(loaded.as_ref()), segment.filter); + if let Some(next) = self.key_iters[source_idx].next() { + self.heap.push(Reverse((next, source_idx))); + } + } + + // The one clone per emitted key, because the writer takes an owned key. + Ok(Some((key.0.clone(), merged))) + } + + /// Union of every source's filtered null posting, or `None` once taken or if + /// no source stores any. Nulls live in `null_map`, outside `index_map`, so + /// they are not part of the key merge and are unioned in one step. + fn take_null_bitmap(&mut self) -> Option { + if self.null_taken { + return None; + } + self.null_taken = true; + + let mut merged: Option = None; + for segment in &self.segments { + if segment.index.null_map.is_empty() { + continue; + } + let nulls = retain_valid( + Cow::Borrowed(segment.index.null_map.as_ref()), + segment.filter, + ); + match &mut merged { + Some(acc) => *acc |= &*nulls, + None => merged = Some(nulls.into_owned()), + } + } + merged + } +} + impl BitmapIndexPlugin { fn get_batch_from_arrays( keys: Arc, @@ -1335,7 +1542,7 @@ impl BitmapIndexPlugin { data: SendableRecordBatchStream, index_store: &dyn IndexStore, ) -> Result { - Self::streaming_build_and_write(data, None, index_store, BITMAP_LOOKUP_NAME, None).await + Self::streaming_build_and_write(data, Vec::new(), index_store, BITMAP_LOOKUP_NAME).await } async fn train_bitmap_shard( @@ -1351,77 +1558,68 @@ impl BitmapIndexPlugin { .stage_start("build_bitmap_shard", None, "rows") .await?; let file = - Self::streaming_build_and_write(data, None, index_store, &file_name, None).await?; + Self::streaming_build_and_write(data, Vec::new(), index_store, &file_name).await?; progress.stage_complete("build_bitmap_shard").await?; Ok(file) } /// Builds and writes a bitmap index in a streaming fashion from value-sorted /// 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. + /// an aggregate map containing every value. The input pipeline, the existing + /// segments and their 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 - /// sorted by value). + /// `old_segments` are merged with the new data via a k-way sorted merge-join + /// (each segment's `index_map` is a `BTreeMap`, already sorted by value), so + /// the old side holds one posting per segment for the key being unioned, plus + /// that union, rather than all of them at once. async fn streaming_build_and_write( data_source: SendableRecordBatchStream, - old_index: Option<&BitmapIndex>, + old_segments: Vec>, index_store: &dyn IndexStore, output_file_name: &str, - old_data_filter: Option<&super::OldIndexDataFilter>, ) -> Result { let value_type = data_source.schema().field(0).data_type().clone(); let mut writer = new_bitmap_batch_writer(index_store, output_file_name, &value_type).await?; - build_index_map(data_source, old_index, old_data_filter, &mut writer).await?; + build_index_map(data_source, old_segments, &mut writer).await?; writer.finish().await } /// Flush a completed value-run from the new data stream, emitting any - /// old-only entries that sort before it and merging the old bitmap if the - /// key exists in both old and new. + /// old-only entries that sort before it and merging the old postings if the + /// key exists on both sides. async fn finish_run( key: ScalarValue, bitmap: &mut RowAddrTreeMap, - old_index: Option<&BitmapIndex>, - old_keys: &mut std::iter::Peekable< - std::collections::btree_map::Keys<'_, OrderableScalarValue, usize>, - >, + old: &mut OldSegments<'_>, emitted_null: &mut bool, writer: &mut BitmapBatchWriter, - old_data_filter: Option<&super::OldIndexDataFilter>, ) -> Result<()> { if key.is_null() { - // Null values are stored separately in the old index's null_map. - if let Some(idx) = old_index - && !idx.null_map.is_empty() - { - *bitmap |= &*retain_valid(Cow::Borrowed(idx.null_map.as_ref()), old_data_filter); + // Null values are stored separately in each old segment's null_map. + if let Some(null_bitmap) = old.take_null_bitmap() { + *bitmap |= &null_bitmap; } *emitted_null = true; - writer.emit(key, bitmap).await?; - } else if let Some(idx) = old_index { + } else { let orderable = OrderableScalarValue(key.clone()); // Emit old-only entries that sort before this key. - while let Some(old_key) = old_keys.next_if(|old| **old < orderable) { - let loaded = idx.load_bitmap(old_key, None).await?; - let old_bitmap = retain_valid(Cow::Borrowed(loaded.as_ref()), old_data_filter); - writer.emit(old_key.0.clone(), &old_bitmap).await?; + while let Some((old_key, old_bitmap)) = + old.take_smallest_if(|old_key| *old_key < orderable).await? + { + writer.emit(old_key, &old_bitmap).await?; } - // If the old index also has this key, merge its bitmap. - if let Some(old_key) = old_keys.next_if(|old| **old == orderable) { - let loaded = idx.load_bitmap(old_key, None).await?; - *bitmap |= &*retain_valid(Cow::Borrowed(loaded.as_ref()), old_data_filter); + // If the old side also has this key, merge its postings. + if let Some((_, old_bitmap)) = old + .take_smallest_if(|old_key| *old_key == orderable) + .await? + { + *bitmap |= &old_bitmap; } - - writer.emit(key, bitmap).await?; - } else { - writer.emit(key, bitmap).await?; } - Ok(()) + writer.emit(key, bitmap).await } /// Merge per-shard bitmap lookup files into a single bitmap index file. @@ -1526,27 +1724,18 @@ pub async fn merge_index_files( /// exactly that for its `list_nulls` set. /// /// Input must be sorted by value with nulls first. This function's transient -/// aggregation state is one key's bitmap at a time. When `old_index` is given, -/// its entries are merge-joined in, loading each old bitmap on demand; its -/// already-loaded `index_map` remains resident separately, outside that state. +/// aggregation state is one key's bitmap at a time. `old_segments` are +/// merge-joined in through [`OldSegments`], loading each old bitmap on demand; +/// their already-loaded `index_map`s remain resident separately, outside that +/// state. pub(crate) async fn build_index_map( mut data_source: SendableRecordBatchStream, - old_index: Option<&BitmapIndex>, - old_data_filter: Option<&super::OldIndexDataFilter>, + old_segments: Vec>, writer: &mut BitmapBatchWriter, ) -> Result<()> { let value_type = data_source.schema().field(0).data_type().clone(); - // Borrowed from the source's `index_map` rather than collected into a Vec. - // That map already holds every key, so collecting made a second full copy: - // a 64-byte `ScalarValue` per key plus its label text, which at 10M labels - // is most of a gigabyte for nothing. - let empty = BTreeMap::new(); - let mut old_keys = old_index - .map(|idx| idx.index_map.as_ref()) - .unwrap_or(&empty) - .keys() - .peekable(); + let mut old = OldSegments::new(old_segments); // Current value being accumulated from the new data stream. let mut current_key: Option = None; @@ -1593,11 +1782,9 @@ pub(crate) async fn build_index_map( BitmapIndexPlugin::finish_run( prev_key, &mut prev_bitmap, - old_index, - &mut old_keys, + &mut old, &mut emitted_null, writer, - old_data_filter, ) .await?; } @@ -1615,32 +1802,22 @@ pub(crate) async fn build_index_map( BitmapIndexPlugin::finish_run( last_key, &mut last_bitmap, - old_index, - &mut old_keys, + &mut old, &mut emitted_null, writer, - old_data_filter, ) .await?; } // Emit any remaining old-only entries. - if let Some(idx) = old_index { - for old_key in old_keys { - let loaded = idx.load_bitmap(old_key, None).await?; - let old_bitmap = retain_valid(Cow::Borrowed(loaded.as_ref()), old_data_filter); - writer.emit(old_key.0.clone(), &old_bitmap).await?; - } + while let Some((key, bitmap)) = old.take_smallest_if(|_| true).await? { + writer.emit(key, &bitmap).await?; } // Emit old null bitmap if we didn't already merge it with new nulls. - if !emitted_null - && let Some(idx) = old_index - && !idx.null_map.is_empty() - { + if !emitted_null && let Some(null_bitmap) = old.take_null_bitmap() { let null_key = new_null_array(&value_type, 1); let null_key = ScalarValue::try_from_array(null_key.as_ref(), 0)?; - let null_bitmap = retain_valid(Cow::Borrowed(idx.null_map.as_ref()), old_data_filter); writer.emit(null_key, &null_bitmap).await?; } @@ -1723,22 +1900,12 @@ pub(crate) fn merge_source_entry_count(sources: &[Arc]) -> u64 { } /// 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. +/// bitmap payloads at once, applying the same `old_data_filter` to every source. +/// See [`OldSegments`] for the merge itself. /// /// 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. @@ -1763,18 +1930,24 @@ pub(crate) async fn merge_index_maps( }; 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(Cow::Owned(merged_nulls), old_data_filter); - if !merged_nulls.is_empty() { + let mut old = OldSegments::new( + sources + .iter() + .map(|source| OldSegment { + index: source.as_ref(), + filter: old_data_filter, + }) + .collect(), + ); + + if let Some(merged_nulls) = old.take_null_bitmap() + && !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; // Cap reporting at roughly a hundred times across the merge. `stage_progress` // is `#[async_trait]`, so every call boxes a future, and a real reporter does // more: the Python one allocates and sends, the Java one makes a JNI upcall. @@ -1784,56 +1957,14 @@ pub(crate) async fn merge_index_maps( let report_every = (merge_source_entry_count(sources) / 100).max(1); let mut last_reported = 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)`. What this replaced was not a - // cheaper comparison strategy but full materialization: the previous - // `merge_bitmap_indices` built a `HashMap` - // holding every key of every source at once. The win here is bounded - // memory, not fewer comparisons. It is the same merge `merge_shards` 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(Cow::Owned(merged), old_data_filter); + while let Some((key, merged)) = old.take_smallest_if(|_| true).await? { if !merged.is_empty() { - writer.emit(next_key.0.clone(), &merged).await?; + writer.emit(key, &merged).await?; } // Reported outside the guard above: a key the filter emptied still // consumed its source entries, and skipping it would stall the count. + let consumed = old.consumed(); if consumed - last_reported >= report_every && let Some((progress, stage)) = progress { @@ -1846,64 +1977,55 @@ pub(crate) async fn merge_index_maps( // throttle above skipped -- including when there were no keys at all (every // list null or empty), where the loop never runs. if let Some((progress, stage)) = progress { - progress.stage_progress(stage, consumed).await?; + progress.stage_progress(stage, old.consumed()).await?; } Ok(()) } +/// Merge bitmap segments with no new data and no old-data filters. +/// +/// Forwards to [`BitmapIndex::merge_segments`], which is the replacement: it +/// also takes a new-data stream and one filter per source segment. Kept with +/// its original signature and progress stages; progress for the merge stage is +/// reported once, at completion, rather than as entries are consumed. +#[deprecated(note = "use BitmapIndex::merge_segments")] pub async fn merge_bitmap_indices( source_indices: &[Arc], dest_store: &dyn IndexStore, progress: Arc, ) -> Result { - if source_indices.is_empty() { + let Some(first) = source_indices.first() else { return Err(Error::invalid_input( "Bitmap segment merge requires at least one source segment".to_string(), )); - } + }; + let entry_count = merge_source_entry_count(source_indices); + progress + .stage_start("merge_bitmap_segments", Some(entry_count), "entries") + .await?; - let value_type = source_indices[0].value_type().clone(); + let schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, first.value_type().clone(), true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let no_new_data: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new(schema, stream::empty())); + let no_filters = vec![None; source_indices.len()]; + let created_index = + BitmapIndex::merge_segments(source_indices, no_new_data, dest_store, &no_filters).await?; progress - .stage_start( - "merge_bitmap_segments", - Some(merge_source_entry_count(source_indices)), - "entries", - ) + .stage_progress("merge_bitmap_segments", entry_count) .await?; - 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 {:?}", - source_index.value_type(), - value_type - ))); - } - } - - 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 = writer.finish().await?; progress.stage_progress("write_bitmap_index", 1).await?; progress.stage_complete("write_bitmap_index").await?; - Ok(CreatedIndex { - index_details: prost_types::Any::from_msg(&pbold::BitmapIndexDetails::default()).unwrap(), - index_version: BITMAP_INDEX_VERSION, - files: vec![file], - }) + Ok(created_index) } #[async_trait] @@ -2155,8 +2277,9 @@ pub(crate) mod test_util { mod tests { use super::*; use crate::metrics::{LocalMetricsCollector, NoOpMetricsCollector}; + use crate::scalar::OldIndexDataFilter; use crate::scalar::lance_format::LanceIndexStore; - use arrow_array::{RecordBatch, StringArray, UInt64Array, record_batch}; + use arrow_array::{BooleanArray, RecordBatch, StringArray, UInt64Array, record_batch}; use arrow_schema::{DataType, Field, Schema}; /// Sort a (value, row_id) RecordBatch by the value column so that unit tests @@ -3239,8 +3362,10 @@ mod tests { } /// Merging bitmap segments must equal a single build over the same rows, - /// including the null bitmap, which lives outside `index_map`. + /// including the null bitmap, which lives outside `index_map`. The + /// deprecated `merge_bitmap_indices` must produce the same file. #[tokio::test] + #[allow(deprecated)] async fn test_bitmap_segment_merge_matches_single_build() { let values: Vec> = (0..600) .map(|i| { @@ -3284,15 +3409,335 @@ mod tests { .unwrap(); let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[left.clone(), right.clone()], + value_row_id_stream(&[]), + dest_store.as_ref(), + &[None, None], + ) + .await + .unwrap(); + assert_eq!(expected, read_bitmap_contents(dest_store.as_ref()).await); + + let (_shim_dir, shim_store) = test_util::index_store(); merge_bitmap_indices( &[left, right], - dest_store.as_ref(), + shim_store.as_ref(), crate::progress::noop_progress(), ) .await .unwrap(); + assert_eq!(expected, read_bitmap_contents(shim_store.as_ref()).await); + } - assert_eq!(expected, read_bitmap_contents(dest_store.as_ref()).await); + fn addr(fragment: u32, offset: u32) -> u64 { + RowAddress::new_from_parts(fragment, offset).into() + } + + /// `utf8_value_stream` over `(value, row_addr)` pairs. + fn value_row_id_stream(rows: &[(Option<&str>, u64)]) -> SendableRecordBatchStream { + utf8_value_stream( + rows.iter().map(|(value, _)| *value), + rows.iter().map(|(_, row_addr)| *row_addr), + ) + } + + /// The Boolean counterpart of [`utf8_value_stream`], likewise sorted. + fn bool_value_row_id_stream(rows: &[(Option, u64)]) -> SendableRecordBatchStream { + let schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, DataType::Boolean, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(BooleanArray::from_iter( + rows.iter().map(|(value, _)| *value), + )), + Arc::new(UInt64Array::from_iter_values( + rows.iter().map(|(_, row_addr)| *row_addr), + )), + ], + ) + .unwrap(); + let batch = sort_batch_by_value(&batch); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::once(async move { Ok(batch) }), + )) + } + + /// Train a standalone bitmap segment over `data`, keeping its temp dir alive. + async fn train_bitmap_segment_from( + data: SendableRecordBatchStream, + ) -> (TempObjDir, Arc) { + let (tmpdir, store) = test_util::index_store(); + BitmapIndexPlugin::train_bitmap_index(data, store.as_ref()) + .await + .unwrap(); + let index = BitmapIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + (tmpdir, index) + } + + async fn train_bitmap_segment(rows: &[(Option<&str>, u64)]) -> (TempObjDir, Arc) { + train_bitmap_segment_from(value_row_id_stream(rows)).await + } + + async fn search_addrs_for(index: &BitmapIndex, query: SargableQuery) -> Vec { + let SearchResult::Exact(selection) = + index.search(&query, &NoOpMetricsCollector).await.unwrap() + else { + panic!("expected an exact bitmap search result"); + }; + let mut addrs = selection + .true_rows() + .row_addrs() + .unwrap() + .map(u64::from) + .collect::>(); + addrs.sort_unstable(); + addrs + } + + async fn search_addrs(index: &BitmapIndex, value: Option<&str>) -> Vec { + let query = match value { + Some(value) => SargableQuery::Equals(ScalarValue::Utf8(Some(value.to_string()))), + None => SargableQuery::IsNull(), + }; + search_addrs_for(index, query).await + } + + async fn search_bool_addrs(index: &BitmapIndex, value: Option) -> Vec { + let query = match value { + Some(value) => SargableQuery::Equals(ScalarValue::Boolean(Some(value))), + None => SargableQuery::IsNull(), + }; + search_addrs_for(index, query).await + } + + /// A 4-way segment merge: three contributing segments plus one whose filter + /// keeps nothing, one deleted row masked by a row-id allow-list, nulls on + /// two segments, and a new-data stream on top. + #[tokio::test] + async fn test_bitmap_merge_k_segments() { + let (_dir0, seg0) = train_bitmap_segment(&[ + (Some("red"), addr(0, 0)), + (Some("red"), addr(0, 1)), + (Some("blue"), addr(0, 2)), + (None, addr(0, 3)), + ]) + .await; + let (_dir1, seg1) = train_bitmap_segment(&[ + (Some("blue"), addr(1, 0)), + (Some("green"), addr(1, 1)), + (Some("red"), addr(1, 2)), + ]) + .await; + let (_dir2, seg2) = train_bitmap_segment(&[ + (Some("yellow"), addr(2, 0)), + (Some("red"), addr(2, 1)), + (None, addr(2, 2)), + ]) + .await; + let (_dir3, seg3) = train_bitmap_segment(&[(Some("purple"), addr(4, 0))]).await; + + // seg1's row (1,2) is deleted: the allow-list omits it. + let mut still_valid = RowAddrTreeMap::new(); + still_valid.insert(addr(1, 0)); + still_valid.insert(addr(1, 1)); + // seg3 covers a compacted-away fragment, so it keeps nothing at all. + let filters = vec![ + None, + Some(OldIndexDataFilter::RowIds(still_valid)), + None, + Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::new(), + to_remove: RoaringBitmap::from_iter([4u32]), + }), + ]; + + let new_rows = [(Some("green"), addr(3, 0)), (Some("blue"), addr(3, 1))]; + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[seg0.clone(), seg1.clone(), seg2.clone(), seg3.clone()], + value_row_id_stream(&new_rows), + dest_store.as_ref(), + &filters, + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + // Every value the sources knew about, minus the two filtered-out sources' + // contributions, plus the new data. + for value in [ + Some("red"), + Some("blue"), + Some("green"), + Some("yellow"), + None, + ] { + let mut expected = Vec::new(); + for segment in [&seg0, &seg1, &seg2] { + expected.extend(search_addrs(segment, value).await); + } + expected.retain(|a| *a != addr(1, 2)); + for (new_value, new_addr) in new_rows { + if new_value == value { + expected.push(new_addr); + } + } + expected.sort_unstable(); + expected.dedup(); + assert_eq!( + search_addrs(&merged, value).await, + expected, + "merged postings differ for {value:?}" + ); + } + + // The deleted row and the kept-nothing segment leave no trace. + assert!( + !search_addrs(&merged, Some("red")) + .await + .contains(&addr(1, 2)) + ); + assert!(search_addrs(&merged, Some("purple")).await.is_empty()); + assert_eq!(merged.index_map.len(), 4); + } + + /// A K-way merge with no new data at all: pure consolidation. + #[tokio::test] + async fn test_bitmap_merge_k_segments_without_new_data() { + let (_dir0, seg0) = train_bitmap_segment(&[(Some("a"), addr(0, 0))]).await; + let (_dir1, seg1) = train_bitmap_segment(&[(Some("b"), addr(1, 0))]).await; + let (_dir2, seg2) = + train_bitmap_segment(&[(Some("a"), addr(2, 0)), (None, addr(2, 1))]).await; + + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[seg0, seg1, seg2], + value_row_id_stream(&[]), + dest_store.as_ref(), + &[None, None, None], + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + assert_eq!( + search_addrs(&merged, Some("a")).await, + vec![addr(0, 0), addr(2, 0)] + ); + assert_eq!(search_addrs(&merged, Some("b")).await, vec![addr(1, 0)]); + assert_eq!(search_addrs(&merged, None).await, vec![addr(2, 1)]); + } + + /// A K-way merge where the NEW data stream also contains nulls: the null run + /// (sorted first) must pull in every old segment's null_map exactly once, + /// respecting per-segment filters, and not emit a second null entry. + #[tokio::test] + async fn test_bitmap_merge_k_segments_with_new_nulls() { + let (_dir0, seg0) = train_bitmap_segment(&[ + (Some("red"), addr(0, 0)), + (None, addr(0, 1)), + (None, addr(0, 2)), + ]) + .await; + let (_dir1, seg1) = + train_bitmap_segment(&[(Some("blue"), addr(1, 0)), (None, addr(1, 1))]).await; + // seg2's null survives, so two old segments contribute nulls and the fold + // cannot pass by reading only the first. + let (_dir2, seg2) = + train_bitmap_segment(&[(Some("green"), addr(2, 0)), (None, addr(2, 1))]).await; + + // seg1's null row (1,1) is deleted: the allow-list omits it. + let mut still_valid = RowAddrTreeMap::new(); + still_valid.insert(addr(1, 0)); + let filters = vec![None, Some(OldIndexDataFilter::RowIds(still_valid)), None]; + + let new_rows = [ + (None, addr(3, 0)), + (Some("red"), addr(3, 1)), + (None, addr(3, 2)), + ]; + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[seg0, seg1, seg2], + value_row_id_stream(&new_rows), + dest_store.as_ref(), + &filters, + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + assert_eq!( + search_addrs(&merged, None).await, + vec![addr(0, 1), addr(0, 2), addr(2, 1), addr(3, 0), addr(3, 2)] + ); + assert_eq!( + search_addrs(&merged, Some("red")).await, + vec![addr(0, 0), addr(3, 1)] + ); + assert_eq!(search_addrs(&merged, Some("blue")).await, vec![addr(1, 0)]); + assert_eq!(search_addrs(&merged, Some("green")).await, vec![addr(2, 0)]); + assert_eq!(merged.index_map.len(), 3, "nulls must not enter index_map"); + } + + /// A K-way merge over a Boolean key column, the shape of a flag column whose + /// full rescan is the most expensive relative to the index it produces. + #[tokio::test] + async fn test_bitmap_merge_k_segments_boolean_keys() { + let (_dir0, seg0) = train_bitmap_segment_from(bool_value_row_id_stream(&[ + (Some(true), addr(0, 0)), + (Some(false), addr(0, 1)), + (None, addr(0, 2)), + ])) + .await; + let (_dir1, seg1) = train_bitmap_segment_from(bool_value_row_id_stream(&[ + (Some(false), addr(1, 0)), + (Some(true), addr(1, 1)), + ])) + .await; + let (_dir2, seg2) = + train_bitmap_segment_from(bool_value_row_id_stream(&[(Some(true), addr(2, 0))])).await; + + let new_rows = [(Some(false), addr(3, 0)), (None, addr(3, 1))]; + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[seg0, seg1, seg2], + bool_value_row_id_stream(&new_rows), + dest_store.as_ref(), + &[None, None, None], + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + assert_eq!( + search_bool_addrs(&merged, Some(true)).await, + vec![addr(0, 0), addr(1, 1), addr(2, 0)] + ); + assert_eq!( + search_bool_addrs(&merged, Some(false)).await, + vec![addr(0, 1), addr(1, 0), addr(3, 0)] + ); + assert_eq!( + search_bool_addrs(&merged, None).await, + vec![addr(0, 2), addr(3, 1)] + ); + assert_eq!(merged.value_type, DataType::Boolean); } /// The keys column counts toward the flush threshold, not just the bitmaps. @@ -3356,7 +3801,7 @@ mod tests { new_bitmap_batch_writer(store.as_ref(), BITMAP_LOOKUP_NAME, &DataType::Utf8) .await .unwrap(); - build_index_map(stream, None, None, &mut writer) + build_index_map(stream, Vec::new(), &mut writer) .await .unwrap(); writer @@ -3432,8 +3877,10 @@ mod tests { .unwrap(); build_index_map( utf8_value_stream([Some("b"), Some("d")], [20u64, 30]), - Some(old_index.as_ref()), - None, + vec![OldSegment { + index: old_index.as_ref(), + filter: None, + }], &mut writer, ) .await diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index cb2c842b906..9d0c402a357 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1980,7 +1980,7 @@ fn filter_row_ids( /// True if `filter` would keep no rows at all (its keep-set is empty), letting /// the merge skip reading the segment entirely. -fn filter_keeps_nothing(filter: &Option) -> bool { +pub(super) fn filter_keeps_nothing(filter: &Option) -> bool { match filter { Some(OldIndexDataFilter::Fragments { to_keep, .. }) => to_keep.is_empty(), Some(OldIndexDataFilter::RowIds(valid)) => valid.is_empty(), diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index 1f502035bcd..e7a0eb4f76c 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -44,7 +44,7 @@ use super::{BuiltinIndexType, SargableQuery, ScalarIndexParams}; use super::{MetricsCollector, SearchResult}; use crate::pbold; use crate::scalar::bitmap::{ - BitmapIndexState, build_index_map, merge_index_maps, merge_source_entry_count, + BitmapIndexState, OldSegment, build_index_map, merge_index_maps, merge_source_entry_count, new_bitmap_batch_writer, remap_index_map, remap_row_addrs, }; use crate::scalar::expression::{LabelListQueryParser, ScalarQueryParser}; @@ -570,8 +570,15 @@ async fn write_label_list_index( list_nulls: impl FnOnce() -> Result, ) -> Result { let mut writer = new_bitmap_batch_writer(store, BITMAP_LOOKUP_NAME, value_type).await?; - // `None`: LabelList does not apply the old-data filter. See `update`. - build_index_map(sorted_labels, old_index, None, &mut writer).await?; + // `filter: None`: LabelList does not apply the old-data filter. See `update`. + let old_segments = old_index + .map(|index| OldSegment { + index, + filter: None, + }) + .into_iter() + .collect(); + build_index_map(sorted_labels, old_segments, &mut writer).await?; writer .add_global_buffer( LABEL_LIST_NULLS_METADATA_KEY.to_string(), @@ -676,8 +683,8 @@ async fn update_label_list_index( /// separate `list_nulls` row set. Because distributed segments cover disjoint rows /// (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, +/// required. This mirrors [`crate::scalar::bitmap::BitmapIndex::merge_segments`] +/// 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( diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 245c861de72..8628073c5e4 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -484,7 +484,10 @@ async fn merge_scalar_indices<'a>( // Scalar Index that expos an N:1 segment-merge primitive reachable without // rescanning the dataset - let has_segment_merge_primitive = matches!(index_type, IndexType::BTree | IndexType::NGram); + let has_segment_merge_primitive = matches!( + index_type, + IndexType::BTree | IndexType::Bitmap | IndexType::NGram + ); let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?; let ngram_requires_rebuild = index_type == IndexType::NGram && frag_reuse_index.as_ref().is_some_and(|frag_reuse_index| { @@ -571,6 +574,19 @@ async fn merge_scalar_indices<'a>( ) .await? } + IndexType::Bitmap => { + let (_, old_data_filters) = + build_per_segment_filters(dataset.as_ref(), &selected_old_indices).await?; + crate::index::scalar::bitmap::open_and_merge_segments( + dataset.as_ref(), + field_path, + &selected_old_indices, + new_data_stream, + &new_store, + &old_data_filters, + ) + .await? + } IndexType::NGram => { let (_, old_data_filters) = build_per_segment_filters(dataset.as_ref(), &selected_old_indices).await?; @@ -4790,4 +4806,334 @@ mod tests { .num_rows(); assert_eq!(total, 150, "no rows may be lost across compaction + merge"); } + + /// Build one Bitmap segment per fragment of `dataset` under `index_name`. + async fn commit_bitmap_segment_per_fragment(dataset: &mut Dataset, index_name: &str) { + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap); + let frag_ids = dataset + .get_fragments() + .iter() + .map(|frag| frag.id() as u32) + .collect::>(); + let mut staged = Vec::with_capacity(frag_ids.len()); + for frag_id in frag_ids { + staged.push( + CreateIndexBuilder::new(dataset, &["cat"], IndexType::Bitmap, ¶ms) + .name(index_name.into()) + .fragments(vec![frag_id]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments(index_name, "cat", staged) + .await + .unwrap(); + } + + fn id_cat_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("cat", DataType::Utf8, false), + ])) + } + + /// `id` ascending, `cat` cycling through A/B/C. + fn id_cat_batch(schema: &Arc, range: std::ops::Range) -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(range.clone())), + Arc::new(StringArray::from_iter_values( + range.map(|i| ["A", "B", "C"][(i % 3) as usize]), + )), + ], + ) + .unwrap() + } + + async fn count_cat(dataset: &Dataset, cat: &str) -> usize { + dataset + .scan() + .filter(&format!("cat = '{cat}'")) + .unwrap() + .count_rows() + .await + .unwrap() as usize + } + + /// A 200-way merge over three Bitmap segments plus an unindexed fragment must + /// consolidate into one segment (via the N:1 segment-merge primitive) and keep + /// every posting. + #[tokio::test] + async fn test_optimize_bitmap_multi_segment_merge_consolidates() { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let schema = id_cat_schema(); + + let reader = RecordBatchIterator::new( + vec![ + Ok(id_cat_batch(&schema, 0..50)), + Ok(id_cat_batch(&schema, 50..100)), + Ok(id_cat_batch(&schema, 100..150)), + ], + schema.clone(), + ); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 50, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + + commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await; + assert_eq!( + dataset.load_indices_by_name("cat_idx").await.unwrap().len(), + 3 + ); + + dataset + .append( + RecordBatchIterator::new(vec![Ok(id_cat_batch(&schema, 150..200))], schema.clone()), + None, + ) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::merge(200)) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + let segments = dataset.load_indices_by_name("cat_idx").await.unwrap(); + assert_eq!( + segments.len(), + 1, + "a 200-way merge must consolidate every bitmap segment, got {segments:?}" + ); + let expected_coverage = dataset + .get_fragments() + .iter() + .map(|frag| frag.id() as u32) + .collect::(); + assert_eq!( + segments[0].fragment_bitmap.as_ref(), + Some(&expected_coverage), + "merged bitmap segment must cover every dataset fragment" + ); + + for (cat, expected) in [("A", 67), ("B", 67), ("C", 66)] { + assert_eq!( + count_cat(&dataset, cat).await, + expected, + "wrong row count for cat = {cat} after multi-segment bitmap merge" + ); + } + } + + /// A wide consolidation: 26 Bitmap delta segments merged into one by a single + /// `optimize_indices` call, with no unindexed data on top. + #[tokio::test] + async fn test_optimize_bitmap_wide_consolidation() { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let schema = id_cat_schema(); + + let num_segments = 26; + let rows_per_fragment = 12; + let total_rows = num_segments * rows_per_fragment; + let batches = (0..num_segments) + .map(|i| { + Ok(id_cat_batch( + &schema, + i * rows_per_fragment..(i + 1) * rows_per_fragment, + )) + }) + .collect::>(); + let reader = RecordBatchIterator::new(batches, schema.clone()); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: rows_per_fragment as usize, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), num_segments as usize); + + commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await; + assert_eq!( + dataset.load_indices_by_name("cat_idx").await.unwrap().len(), + num_segments as usize + ); + + dataset + .optimize_indices(&OptimizeOptions::merge(num_segments as usize)) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + let segments = dataset.load_indices_by_name("cat_idx").await.unwrap(); + assert_eq!( + segments.len(), + 1, + "a {num_segments}-way merge must consolidate every bitmap segment, got {segments:?}" + ); + let expected_coverage = dataset + .get_fragments() + .iter() + .map(|frag| frag.id() as u32) + .collect::(); + assert_eq!( + segments[0].fragment_bitmap.as_ref(), + Some(&expected_coverage) + ); + + let per_cat = (total_rows / 3) as usize; + for cat in ["A", "B", "C"] { + assert_eq!( + count_cat(&dataset, cat).await, + per_cat, + "wrong row count for cat = {cat} after wide bitmap consolidation" + ); + } + } + + /// Deferred-remap compaction leaves the bitmap segments pointing at retired + /// fragment ids; a K-way segment merge must remap them through the + /// FragReuseIndex instead of dropping or misattributing the rows. + #[tokio::test] + async fn test_optimize_bitmap_merge_remaps_deferred_compaction() { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let schema = id_cat_schema(); + + let reader = RecordBatchIterator::new( + vec![ + Ok(id_cat_batch(&schema, 0..50)), + Ok(id_cat_batch(&schema, 50..100)), + Ok(id_cat_batch(&schema, 100..150)), + ], + schema.clone(), + ); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 50, + ..Default::default() + }), + ) + .await + .unwrap(); + commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await; + + compact_files( + &mut dataset, + CompactionOptions { + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(id_cat_batch(&schema, 150..200))], schema.clone()), + None, + ) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::merge(200)) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + for (cat, expected) in [("A", 67), ("B", 67), ("C", 66)] { + assert_eq!( + count_cat(&dataset, cat).await, + expected, + "cat = {cat} lost or gained rows across deferred compaction + bitmap merge" + ); + } + assert_eq!( + dataset.scan().count_rows().await.unwrap(), + 200, + "no rows may be lost across compaction + bitmap merge" + ); + } + + /// Stable-row-id update against an older Bitmap segment's fragment: the K-way + /// merge must drop that segment's stale postings, not just the tail segment's. + #[tokio::test] + async fn test_optimize_bitmap_drops_stale_rows_across_segments_after_update() { + use crate::dataset::UpdateBuilder; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let schema = id_cat_schema(); + + let reader = RecordBatchIterator::new( + vec![ + Ok(id_cat_batch(&schema, 0..50)), + Ok(id_cat_batch(&schema, 50..100)), + ], + schema.clone(), + ); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 50, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await; + + // Rows 0..25 live in the *older* segment's fragment; rewrite their cat. + let res = UpdateBuilder::new(Arc::new(dataset.clone())) + .update_where("id < 25") + .unwrap() + .set("cat", "'Z'") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + let mut dataset = res.new_dataset.as_ref().clone(); + + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + // ids 0..25 are now 'Z'; the remaining A/B/C counts come from ids 25..100. + assert_eq!(count_cat(&dataset, "Z").await, 25, "updated rows missing"); + for (cat, expected) in [("A", 25), ("B", 25), ("C", 25)] { + assert_eq!( + count_cat(&dataset, cat).await, + expected, + "bitmap merge returned stale rows for cat = {cat}" + ); + } + } } diff --git a/rust/lance/src/index/scalar/bitmap.rs b/rust/lance/src/index/scalar/bitmap.rs index f9b8bf69e92..8375129058b 100644 --- a/rust/lance/src/index/scalar/bitmap.rs +++ b/rust/lance/src/index/scalar/bitmap.rs @@ -1,17 +1,47 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use datafusion::execution::SendableRecordBatchStream; use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::bitmap::BitmapIndex; use lance_index::scalar::index_files_to_table; use lance_index::scalar::lance_format::LanceIndexStore; +use lance_index::scalar::{CreatedIndex, OldIndexDataFilter}; use lance_table::format::IndexMetadata; -use roaring::RoaringBitmap; use std::sync::Arc; use uuid::Uuid; use crate::{Dataset, Error, Result, dataset::index::LanceIndexStoreExt}; +/// Open the given bitmap `segments` and k-way merge their postings, together +/// with `new_data`, into a single bitmap index written to `new_store`. +pub(in crate::index) async fn open_and_merge_segments( + dataset: &Dataset, + field_path: &str, + segments: &[&IndexMetadata], + new_data: SendableRecordBatchStream, + new_store: &LanceIndexStore, + old_data_filters: &[Option], +) -> Result { + let mut source_indices = Vec::with_capacity(segments.len()); + for &segment in segments { + let scalar_index = + super::open_scalar_index(dataset, field_path, segment, &NoOpMetricsCollector).await?; + let bitmap_index = scalar_index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::index(format!( + "Bitmap merge: expected bitmap segment {}, got {:?}", + segment.uuid, + scalar_index.index_type() + )) + })?; + source_indices.push(Arc::new(bitmap_index.clone())); + } + BitmapIndex::merge_segments(&source_indices, new_data, new_store, old_data_filters).await +} + /// Merge one caller-defined group of source bitmap segments into a single segment. pub(in crate::index) async fn merge_segments( dataset: &Dataset, @@ -29,36 +59,21 @@ pub(in crate::index) async fn merge_segments( })?; let field_path = dataset.schema().field_path(field_id)?; - let mut source_indices = Vec::with_capacity(segments.len()); - let mut fragment_bitmap = RoaringBitmap::new(); - for segment in &segments { - fragment_bitmap |= segment.fragment_bitmap.as_ref().cloned().ok_or_else(|| { - Error::invalid_input(format!( - "CreateIndex: segment {} is missing fragment coverage", - segment.uuid - )) - })?; - let scalar_index = - super::open_scalar_index(dataset, &field_path, segment, &NoOpMetricsCollector).await?; - let bitmap_index = scalar_index - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::index(format!( - "merge_existing_index_segments: expected bitmap segment {}, got {:?}", - segment.uuid, - scalar_index.index_type() - )) - })?; - source_indices.push(Arc::new(bitmap_index.clone())); - } + let segment_refs: Vec<&IndexMetadata> = segments.iter().collect(); + let (fragment_bitmap, old_data_filters) = + crate::index::append::build_per_segment_filters(dataset, &segment_refs).await?; let new_uuid = Uuid::new_v4(); let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?; - let created_index = lance_index::scalar::bitmap::merge_bitmap_indices( - &source_indices, + // Pure segment consolidation: no dataset scan, so `new_data` is an empty stream. + let empty_new_data = super::btree::empty_value_row_id_stream(dataset, field_id)?; + let created_index = open_and_merge_segments( + dataset, + &field_path, + &segment_refs, + empty_new_data, &new_store, - lance_index::progress::noop_progress(), + &old_data_filters, ) .await?; diff --git a/rust/lance/src/index/scalar/btree.rs b/rust/lance/src/index/scalar/btree.rs index 51f3cfc7cd1..684da278c02 100644 --- a/rust/lance/src/index/scalar/btree.rs +++ b/rust/lance/src/index/scalar/btree.rs @@ -21,8 +21,8 @@ use uuid::Uuid; use crate::{Dataset, Error, Result, dataset::index::LanceIndexStoreExt}; -/// Build a row-empty `new_data` stream for the BTree merge API. -fn empty_btree_update_stream( +/// Build a row-empty `new_data` stream for the segment-merge APIs. +pub(super) fn empty_value_row_id_stream( dataset: &Dataset, field_id: i32, ) -> Result { @@ -125,7 +125,7 @@ pub(crate) async fn merge_segments( let new_store = LanceIndexStore::from_dataset_for_new(dataset, &output_uuid)?; // Pure segment consolidation: no dataset scan, so `new_data` is an empty // stream and the merge is driven entirely by the source page data. - let empty_new_data = empty_btree_update_stream(dataset, field_id)?; + let empty_new_data = empty_value_row_id_stream(dataset, field_id)?; let created_index = open_and_merge_segments( dataset, &field_path, From f8fc1d7a8d510d83ed96d61388802fb8d6f337b7 Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 13:12:07 +0200 Subject: [PATCH 2/7] test(index): check the Bitmap segment merge against a brute-force union `BitmapIndex::merge_segments` promises that its output is the union of its filtered inputs, and nothing asserted that directly. The existing tests pin hand-built shapes; this covers the contract itself. 200 fixed seeds, each building 1-4 segments over a 5-value vocabulary and a 4-fragment address space so keys and row addresses collide across segments, every `OldIndexDataFilter` arm including ones that keep nothing, and nulls on either side. `oracle_keeps` reimplements `retain_old_rows` rather than calling it, so the oracle is independent of the code under test. Assertions read `index_map`/`null_map` rather than `search`, comparing what the merge wrote instead of how a query later resolves a row that is both null and non-null. A key whose rows are all filtered out may still be materialised with an empty bitmap, so only a key outside the vocabulary counts as fabricated. Runs in about a second. --- rust/lance-index/src/scalar/bitmap.rs | 124 ++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index ae3ff7f7f2b..94fb1b542ba 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -2302,6 +2302,7 @@ mod tests { use lance_core::utils::{address::RowAddress, tempfile::TempObjDir}; use lance_io::object_store::ObjectStore; use lance_select::RowSetOps; + use rand::{Rng, SeedableRng, rngs::SmallRng}; use rstest::rstest; fn assert_state_roundtrips(state: &BitmapIndexState) { @@ -3740,6 +3741,129 @@ mod tests { assert_eq!(merged.value_type, DataType::Boolean); } + /// Does `filter` keep `row_addr`? Mirrors + /// [`OldIndexDataFilter::retain_old_rows`], deliberately reimplemented so + /// the oracle below is independent of the merge. + fn oracle_keeps(filter: &Option, row_addr: u64) -> bool { + match filter { + Some(OldIndexDataFilter::Fragments { to_keep, .. }) => { + to_keep.contains(RowAddress::from(row_addr).fragment_id()) + } + Some(OldIndexDataFilter::RowIds(valid)) => valid.contains(row_addr), + None => true, + } + } + + /// The merge's contract is that its output is the union of its filtered inputs. + /// Check exactly that against a brute-force union over many pseudo-random + /// shapes: 1-4 segments over a small vocabulary and address space so keys and + /// row addresses collide across segments, every filter variant including ones + /// that keep nothing, and nulls on either side. + /// + /// Assertions read `index_map`/`null_map` rather than `search`, to compare what + /// the merge wrote and not how a query later resolves a row that is both null + /// and non-null. + #[tokio::test] + async fn test_bitmap_merge_matches_brute_force_union() { + const VOCAB: [&str; 5] = ["a", "b", "c", "d", "e"]; + const FRAGMENTS: u32 = 4; + + for seed in 0..200u64 { + let mut rng = SmallRng::seed_from_u64(seed); + let rand_row = |rng: &mut SmallRng| { + let value = rng + .random_bool(0.8) + .then(|| VOCAB[rng.random_range(0..VOCAB.len())]); + let row_addr = addr(rng.random_range(0..FRAGMENTS), rng.random_range(0..8u32)); + (value, row_addr) + }; + + let num_segments = rng.random_range(1..=4usize); + let mut segment_rows = Vec::with_capacity(num_segments); + let mut filters = Vec::with_capacity(num_segments); + for _ in 0..num_segments { + let rows = (0..rng.random_range(1..=8)) + .map(|_| rand_row(&mut rng)) + .collect::>(); + filters.push(match rng.random_range(0..4) { + 0 => None, + 1 => Some(OldIndexDataFilter::Fragments { + // The bitmap merge reads only `to_keep`. + to_keep: (0..FRAGMENTS).filter(|_| rng.random_bool(0.6)).collect(), + to_remove: RoaringBitmap::new(), + }), + 2 => Some(OldIndexDataFilter::RowIds( + rows.iter() + .filter(|_| rng.random_bool(0.6)) + .map(|(_, row_addr)| *row_addr) + .collect(), + )), + _ => Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::new(), + to_remove: RoaringBitmap::new(), + }), + }); + segment_rows.push(rows); + } + let new_rows = (0..rng.random_range(0..=6)) + .map(|_| rand_row(&mut rng)) + .collect::>(); + + // Brute-force union: surviving old rows plus every new row. + let mut expected: HashMap, RowAddrTreeMap> = HashMap::new(); + for (rows, filter) in segment_rows.iter().zip(&filters) { + for (value, row_addr) in rows { + if oracle_keeps(filter, *row_addr) { + expected.entry(*value).or_default().insert(*row_addr); + } + } + } + for (value, row_addr) in &new_rows { + expected.entry(*value).or_default().insert(*row_addr); + } + + let mut segments = Vec::with_capacity(num_segments); + let mut _dirs = Vec::with_capacity(num_segments); + for rows in &segment_rows { + let (dir, segment) = train_bitmap_segment(rows).await; + _dirs.push(dir); + segments.push(segment); + } + + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &segments, + value_row_id_stream(&new_rows), + dest_store.as_ref(), + &filters, + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + // A key whose rows were all filtered out may still be materialised + // with an empty bitmap, so only a key outside the vocabulary is a + // fabrication. + for value in VOCAB { + let key = OrderableScalarValue(ScalarValue::Utf8(Some(value.to_string()))); + let want = expected.get(&Some(value)).cloned().unwrap_or_default(); + let got = merged.load_bitmap(&key, None).await.unwrap(); + assert_eq!(*got, want, "seed {seed}: row sets differ for {value:?}"); + } + assert!( + merged.index_map.len() <= VOCAB.len(), + "seed {seed}: merge materialised keys outside the vocabulary" + ); + assert_eq!( + *merged.null_map, + expected.get(&None).cloned().unwrap_or_default(), + "seed {seed}: null row sets differ" + ); + } + } + /// 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 From 1531c58c1bc1b300be181d50115ca6e12c11ccab Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 13:14:24 +0200 Subject: [PATCH 3/7] test(index): cover the Bitmap merge_existing_index_segments path Every other Bitmap merge test goes through `optimize_indices` and builds one segment per fragment, so the distributed-build entry point had no coverage and neither did a segment covering more than one fragment -- which is the only way a per-segment filter is a strict subset of what the segment holds. `test_bitmap_merge_existing_index_segments_multi_fragment` commits two segments over four fragments, two fragments each, deletes one row from each segment's coverage, then merges and opens the result directly. Stable row ids make the filter an exact row-id allow-list, so the deletions reach it instead of being masked at scan time. `test_optimize_bitmap_drops_stale_rows_across_segments_after_update` now asserts one segment remains. It selects two today, but nothing pinned that, so a change to `select_segments_to_merge` could have quietly moved it onto the single-segment `update` path while it kept its name. --- rust/lance/src/index/append.rs | 7 +++ rust/lance/src/index/create.rs | 100 +++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 8628073c5e4..2b8b8376f3a 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -5126,6 +5126,13 @@ mod tests { .unwrap(); let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + // Two segments went in, so this exercised the K-way path rather than the + // single-segment `update` this change does not touch. + assert_eq!( + dataset.load_indices_by_name("cat_idx").await.unwrap().len(), + 1, + "the optimize must have consolidated both segments" + ); // ids 0..25 are now 'Z'; the remaining A/B/C counts come from ids 25..100. assert_eq!(count_cat(&dataset, "Z").await, 25, "updated rows missing"); for (cat, expected) in [("A", 25), ("B", 25), ("C", 25)] { diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 8cd006a59a1..0c8ccbdbe9a 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -3329,6 +3329,106 @@ mod tests { (schema, batch) } + /// `merge_existing_index_segments` for Bitmap, with segments that each cover + /// two fragments and an old-data filter that actually removes rows. + /// + /// The other Bitmap merge tests go through `optimize_indices` and build one + /// segment per fragment, so this is the only coverage of the distributed-build + /// entry point, and of a segment whose coverage is wider than one fragment. + /// Stable row ids make the filter an exact row-id allow-list, so the deleted + /// rows reach it rather than being masked at scan time. + #[tokio::test] + async fn test_bitmap_merge_existing_index_segments_multi_fragment() { + async fn count_value(dataset: &Dataset, segment: &IndexMetadata, value: &str) -> usize { + let field_path = dataset.schema().field_path(segment.fields[0]).unwrap(); + let index = crate::index::scalar::open_scalar_index( + dataset, + &field_path, + segment, + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let query = SargableQuery::Equals(ScalarValue::Utf8(Some(value.to_string()))); + match index.search(&query, &NoOpMetricsCollector).await.unwrap() { + SearchResult::Exact(row_ids) => row_ids.true_rows().row_addrs().unwrap().count(), + other => panic!("expected exact result, got {other:?}"), + } + } + + let tmpdir = TempStrDir::default(); + let dataset_uri = format!("file://{}", tmpdir.as_str()); + + // 16 rows over four 4-row fragments; `cat` cycles A/B/C/D, so four each. + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("cat", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..16)), + Arc::new(StringArray::from_iter_values( + (0..16).map(|i| ["A", "B", "C", "D"][(i % 4) as usize]), + )), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write( + reader, + &dataset_uri, + Some(WriteParams { + max_rows_per_file: 4, + mode: WriteMode::Overwrite, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 4); + + // Two segments, each covering two fragments. + let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::Bitmap); + let mut staged = Vec::with_capacity(2); + for fragments in [vec![0u32, 1], vec![2, 3]] { + staged.push( + CreateIndexBuilder::new(&mut dataset, &["cat"], IndexType::Bitmap, ¶ms) + .name("cat_idx".to_string()) + .fragments(fragments) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("cat_idx", "cat", staged) + .await + .unwrap(); + + // One B from each segment's coverage, so both filters have work to do. + dataset.delete("id = 1 OR id = 9").await.unwrap(); + + let merged = dataset + .merge_existing_index_segments(dataset.load_indices_by_name("cat_idx").await.unwrap()) + .await + .unwrap(); + assert_eq!( + merged.fragment_bitmap.as_ref(), + Some(&(0..4u32).collect::()), + "the merged segment must cover every fragment the sources did" + ); + + for (value, expected) in [("A", 4), ("B", 2), ("C", 4), ("D", 4)] { + assert_eq!( + count_value(&dataset, &merged, value).await, + expected, + "wrong row count for cat = {value} after merging multi-fragment segments" + ); + } + } + #[tokio::test] async fn test_label_list_merge_existing_index_segments_drops_retired_fragments() { use lance_index::scalar::{LabelListQuery, SearchResult}; From 904af1cb1b7794b58bfabdce0fe9c7f73c199fb8 Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 13:18:31 +0200 Subject: [PATCH 4/7] fix(index): reject bitmap input that is not sorted by value `build_index_map` groups rows into runs by comparing each value with the previous one only, so a value that reappears after another value opens a second run for a key already written. `BitmapIndex::load` keys `index_map` by value and keeps the last file offset, so the earlier run's row set was dropped with no error or log line -- a `debug_assert!` caught it in debug builds and nothing did in release builds. Nulls have the same exposure through `null_map`, which takes the last null entry in the file. Check the ordering at each value change and return `Error::invalid_input` naming both keys, in release builds too. The cost is one `OrderableScalarValue` comparison and two scalar clones per distinct value -- per run, not per row -- against a build that already serializes a bitmap per run; unsorted input otherwise silently produces a corrupt index. `BitmapIndex::merge_segments` and `BitmapIndexPlugin::train_bitmap_index` take a caller-supplied stream and validated everything about it except the ordering the build depends on. Only the non-null keys have to ascend. Nulls are collected into `null_map` rather than merge-joined by value, so one null run is correct wherever it falls, and both `asc_nulls_first` and `asc_nulls_last` input stay valid; a second null run is what gets rejected. Cover nulls end to end as well. Every `optimize_indices` test used a non-nullable column, so no null reached the scanner's ordering, the per-segment filters or the commit. `test_optimize_bitmap_multi_segment_merge_keeps_nulls` consolidates three segments plus an unindexed tail over a column that is null every fourth row, and checks `IS NULL` and `IS NOT NULL` against the value row sets. --- rust/lance-index/src/scalar/bitmap.rs | 115 +++++++++++++++++----- rust/lance-index/src/scalar/label_list.rs | 10 +- rust/lance/src/index/append.rs | 98 +++++++++++++++++- 3 files changed, 194 insertions(+), 29 deletions(-) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 94fb1b542ba..d0b2c132ed8 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -1716,6 +1716,47 @@ pub async fn merge_index_files( Ok(()) } +/// Enforce the ordering a bitmap build requires, once per value change. +/// +/// A run is a maximal group of equal keys, so reopening a run for a key already +/// written loses rows silently: the old-key cursor has advanced past it on an +/// earlier, larger run, so it leaves as an old-only row and returns later +/// carrying only new rows, and `BitmapIndex::load` keys `index_map` by value +/// and keeps only the last of the two file offsets. The same holds for nulls, +/// which `load` funnels into `null_map` from the last null entry it sees. +/// +/// Only the non-null keys have to ascend. Nulls are collected separately rather +/// than merge-joined by value, so a single null run is correct wherever it +/// falls, which lets a caller sort nulls first or last. `null_run_closed` says +/// whether one has already been flushed, making a second run detectable. +/// +/// Both keys come from the same column, so they share a `ScalarValue` variant +/// and `OrderableScalarValue`'s `Ord` cannot panic comparing them. +fn check_run_order( + previous: &ScalarValue, + next: &ScalarValue, + null_run_closed: bool, +) -> Result<()> { + if next.is_null() { + if null_run_closed { + return Err(Error::invalid_input( + "bitmap index: input is not sorted by value, it has more than one run of nulls" + .to_string(), + )); + } + return Ok(()); + } + if previous.is_null() { + return Ok(()); + } + if OrderableScalarValue(next.clone()) <= OrderableScalarValue(previous.clone()) { + return Err(Error::invalid_input(format!( + "bitmap index: input must be sorted by value, but {next} follows {previous}" + ))); + } + Ok(()) +} + /// Build a bitmap index map from value-sorted `(value, row_id)` input, emitting /// one key at a time into `writer`. /// @@ -1723,11 +1764,12 @@ pub async fn merge_index_files( /// own global buffers in the same file can reuse this. `LabelListIndex` does /// exactly that for its `list_nulls` set. /// -/// Input must be sorted by value with nulls first. This function's transient -/// aggregation state is one key's bitmap at a time. `old_segments` are -/// merge-joined in through [`OldSegments`], loading each old bitmap on demand; -/// their already-loaded `index_map`s remain resident separately, outside that -/// state. +/// Input must be sorted ascending by value, with nulls in a single run that may +/// lead or trail; see [`check_run_order`] for what is rejected. This function's +/// transient aggregation state is one key's bitmap at a time. `old_segments` +/// are merge-joined in through [`OldSegments`], loading each old bitmap on +/// demand; their already-loaded `index_map`s remain resident separately, +/// outside that state. pub(crate) async fn build_index_map( mut data_source: SendableRecordBatchStream, old_segments: Vec>, @@ -1761,23 +1803,7 @@ pub(crate) async fn build_index_map( _ => { // Value changed — flush the previous run. if let Some(prev_key) = current_key.take() { - // This function assumes value-sorted, nulls-first input - // and does not check it in release builds. Violated - // input emits one key twice -- the old-key cursor - // advances past it on an earlier, larger run, so it - // leaves as an old-only row and returns later carrying - // only new rows. Neither row is complete, and - // `BitmapIndex::load` inserts both into `index_map`, - // so the later one shadows the earlier and its rows - // are lost. - debug_assert!( - OrderableScalarValue(key.clone()) - > OrderableScalarValue(prev_key.clone()), - "build_index_map input must be sorted ascending by value \ - with nulls first; got key {:?} after {:?}", - key, - prev_key - ); + check_run_order(&prev_key, &key, emitted_null)?; let mut prev_bitmap = std::mem::take(&mut current_bitmap); BitmapIndexPlugin::finish_run( prev_key, @@ -3864,6 +3890,51 @@ mod tests { } } + /// Unsorted input must be rejected. It would reopen a run for a key already + /// written, and `load` keeps only the last file offset per key, so the earlier + /// row set would vanish silently. Nulls may lead or trail, but only once. + #[rstest] + #[case::value_reappears(vec![Some("a"), Some("b"), Some("a")], "a follows b")] + #[case::two_null_runs(vec![None, Some("a"), None], "more than one run of nulls")] + #[tokio::test] + async fn test_bitmap_build_rejects_unsorted_input( + #[case] values: Vec>, + #[case] expected: &str, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let row_addrs = (0..values.len() as u32) + .map(|i| addr(0, i)) + .collect::>(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(values)), + Arc::new(UInt64Array::from(row_addrs)), + ], + ) + .unwrap(); + // Deliberately not sorted, unlike `utf8_value_stream`. + let unsorted: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::once(async move { Ok(batch) }), + )); + + let (_tmpdir, store) = test_util::index_store(); + let Err(err) = BitmapIndexPlugin::train_bitmap_index(unsorted, store.as_ref()).await else { + panic!("expected unsorted input to be rejected"); + }; + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + let message = err.to_string(); + assert!(message.contains("sorted by value"), "{message}"); + assert!(message.contains(expected), "{message}"); + } + /// 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 diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index e7a0eb4f76c..10dba71a58a 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -485,12 +485,10 @@ fn serialize_list_nulls(null_map: &RowAddrTreeMap) -> Result { /// buffers the data volume at all. /// /// Nulls sort first because `OrderableScalarValue` orders them below every -/// value, so [`build_index_map`]'s ascending-input `debug_assert!` rejects a -/// stream that puts them last. Its runtime path would in fact tolerate a null -/// run anywhere -- `finish_run` never advances the old-keys cursor for a null -/// key -- but the assert is the contract, and null-first is also what -/// [`remap_index_map`] emits and what the plain bitmap index's training scan -/// produces. +/// value. [`build_index_map`] would accept a single null run anywhere -- nulls +/// are collected separately rather than merge-joined by value -- but null-first +/// is also what [`remap_index_map`] emits and what the plain bitmap index's +/// training scan produces. /// /// `mem_pool_size`, when set, overrides the session's memory pool for this /// sort instead of leaving it to `LANCE_MEM_POOL_SIZE`. Production callers diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 2b8b8376f3a..27b5b728334 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -4854,15 +4854,111 @@ mod tests { } async fn count_cat(dataset: &Dataset, cat: &str) -> usize { + count_rows_where(dataset, &format!("cat = '{cat}'")).await + } + + async fn count_rows_where(dataset: &Dataset, predicate: &str) -> usize { dataset .scan() - .filter(&format!("cat = '{cat}'")) + .filter(predicate) .unwrap() .count_rows() .await .unwrap() as usize } + /// `id` ascending, `cat` cycling through A/B/C/NULL. + fn id_cat_nullable_batch(schema: &Arc, range: std::ops::Range) -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(range.clone())), + Arc::new(StringArray::from_iter( + range.map(|i| ["A", "B", "C"].get((i % 4) as usize).copied()), + )), + ], + ) + .unwrap() + } + + /// The K-way merge over a nullable column: nulls must survive consolidation of + /// several segments plus an unindexed tail, and `IS NULL` is served from + /// `null_map`, which is filled separately from the value row sets. + #[tokio::test] + async fn test_optimize_bitmap_multi_segment_merge_keeps_nulls() { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("cat", DataType::Utf8, true), + ])); + + // 36 rows over three 12-row fragments; every fourth row is null. + let reader = RecordBatchIterator::new( + vec![ + Ok(id_cat_nullable_batch(&schema, 0..12)), + Ok(id_cat_nullable_batch(&schema, 12..24)), + Ok(id_cat_nullable_batch(&schema, 24..36)), + ], + schema.clone(), + ); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 12, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + + commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await; + dataset + .append( + RecordBatchIterator::new( + vec![Ok(id_cat_nullable_batch(&schema, 36..48))], + schema.clone(), + ), + None, + ) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::merge(200)) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + assert_eq!( + dataset.load_indices_by_name("cat_idx").await.unwrap().len(), + 1, + "the merge must consolidate every bitmap segment" + ); + + // 48 rows cycling A/B/C/NULL: 12 of each. + for cat in ["A", "B", "C"] { + assert_eq!( + count_cat(&dataset, cat).await, + 12, + "wrong row count for cat = {cat} after merging a nullable column" + ); + } + assert_eq!( + count_rows_where(&dataset, "cat IS NULL").await, + 12, + "nulls lost across the bitmap merge" + ); + assert_eq!( + count_rows_where(&dataset, "cat IS NOT NULL").await, + 36, + "IS NOT NULL disagrees with the value row sets" + ); + assert_eq!(dataset.scan().count_rows().await.unwrap(), 48); + } + /// A 200-way merge over three Bitmap segments plus an unindexed fragment must /// consolidate into one segment (via the N:1 segment-merge primitive) and keep /// every posting. From 586dfc726daa0fcde986af8766081e8e3e7f1a82 Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 13:22:12 +0200 Subject: [PATCH 5/7] perf(index): keep the bitmap segment merge out of the index cache `OldSegments` read each source row set through `load_bitmap`, which caches what it reads. A merge touches every key of every source segment exactly once, and the commit that follows retires those segments, so every entry it inserted was dead on arrival while evicting entries for indices that still exist. The cache holds up to 6 GiB by default, so a wide consolidation can displace a lot. Each row set was also copied three times: deserialized, cloned for the cache insert, then cloned again to get something the filter could run on in place. Split the read at the cache boundary. `read_bitmap_at` does the read, deserialize and fragment-reuse remap; `load_bitmap` keeps its null, absent-key and cache handling around it, unchanged for `search` and for `remap_index_map`. `read_bitmap_uncached` gives the merge an owned row set with no cache interaction, which `OldSegments` then filters in place, leaving one copy instead of three. The same path serves the single-segment `update` and the LabelList merge, whose sources are retired just the same. Going through `load_bitmap` was never gratuitous -- it is what applies the fragment-reuse remap, and reading the file directly would resurrect row addresses pointing at retired fragments. `read_bitmap_at` keeps that remap, so only the caching is bypassed. --- rust/lance-index/src/scalar/bitmap.rs | 101 +++++++++++++++++++++++--- 1 file changed, 89 insertions(+), 12 deletions(-) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index d0b2c132ed8..096c41686fe 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -512,6 +512,19 @@ impl BitmapIndex { metrics.record_part_load(); } + let bitmap = self.read_bitmap_at(row_offset).await?; + + self.index_cache + .insert_with_key(&cache_key, Arc::new(bitmap.clone())) + .await; + + Ok(Arc::new(bitmap)) + } + + /// Read one row set straight from the lookup file, remapped through the + /// fragment-reuse index but neither served from nor written to the index + /// cache. + async fn read_bitmap_at(&self, row_offset: usize) -> Result { let page_lookup_file = self.lazy_reader.get().await?; let batch = page_lookup_file .read_range(row_offset..row_offset + 1, Some(&["bitmaps"])) @@ -529,11 +542,25 @@ impl BitmapIndex { bitmap = fri.remap_row_addrs_tree_map(&bitmap); } - self.index_cache - .insert_with_key(&cache_key, Arc::new(bitmap.clone())) - .await; + Ok(bitmap) + } - Ok(Arc::new(bitmap)) + /// Owned row set for `key`, bypassing the index cache. + /// + /// A merge reads every key of every source segment exactly once, and the + /// commit that follows retires those segments, so routing the reads through + /// the cache would only evict entries for indices that still exist. Returning + /// an owned value also lets the caller filter in place, where + /// [`Self::load_bitmap`] hands back an `Arc` to clone. + /// + /// Nulls are not in `index_map` -- [`OldSegments::take_null_bitmap`] reads + /// `null_map` for those -- so unlike `load_bitmap` this does not special-case + /// them and would return an empty set for a null key. + async fn read_bitmap_uncached(&self, key: &OrderableScalarValue) -> Result { + match self.index_map.get(key) { + Some(row_offset) => self.read_bitmap_at(*row_offset).await, + None => Ok(RowAddrTreeMap::default()), + } } pub(crate) fn value_type(&self) -> &DataType { @@ -1411,10 +1438,14 @@ pub(crate) struct OldSegment<'a> { /// `index_map` order can be trusted for old and new files alike, which is what /// makes this work without an index version bump. /// -/// Bitmaps are read one key at a time through [`BitmapIndex::load_bitmap`], the -/// accessor that applies fragment-reuse remapping. Reading the lookup file's raw -/// bytes would bypass it and emit row addresses pointing at fragments compaction -/// has already retired, which under deferred index remapping is silent data loss. +/// Bitmaps are read one key at a time through +/// [`BitmapIndex::read_bitmap_uncached`], which applies fragment-reuse remapping +/// like the query path's `load_bitmap` but skips the index cache: every source +/// entry is read exactly once and the sources are retired right after, so caching +/// them would only evict live entries. Reading the lookup file's raw bytes instead +/// would skip the remap too and emit row addresses pointing at fragments +/// compaction has already retired, which under deferred index remapping is silent +/// data loss. /// /// 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 @@ -1423,8 +1454,8 @@ pub(crate) struct OldSegment<'a> { /// so cloning would cost an allocation per key per source for string keys. /// /// The transient state is the merged bitmap for the current key plus one loaded -/// bitmap per participating source. Each source's `index_map`, any bitmaps the -/// index cache retains, and the output writer stay outside it. +/// bitmap per participating source. Each source's `index_map` and the output +/// writer stay outside it. struct OldSegments<'a> { segments: Vec>, key_iters: Vec>, @@ -1485,8 +1516,11 @@ impl<'a> OldSegments<'a> { self.heap.pop(); self.consumed += 1; let segment = &self.segments[source_idx]; - let loaded = segment.index.load_bitmap(key, None).await?; - merged |= &*retain_valid(Cow::Borrowed(loaded.as_ref()), segment.filter); + let mut bitmap = segment.index.read_bitmap_uncached(key).await?; + if let Some(filter) = segment.filter { + filter.retain_old_rows(&mut bitmap); + } + merged |= &bitmap; if let Some(next) = self.key_iters[source_idx].next() { self.heap.push(Reverse((next, source_idx))); } @@ -3890,6 +3924,49 @@ mod tests { } } + /// A merge must not populate the index cache with the source segments' row + /// sets: it reads each one once, and the commit that follows retires those + /// segments, so the entries would only evict live ones. + #[tokio::test] + async fn test_bitmap_merge_does_not_cache_source_bitmaps() { + let cache = LanceCache::with_capacity(64 * 1024 * 1024); + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(cache.clone()), + )); + BitmapIndexPlugin::train_bitmap_index( + value_row_id_stream(&[(Some("red"), addr(0, 0)), (Some("blue"), addr(0, 1))]), + store.as_ref(), + ) + .await + .unwrap(); + let segment = BitmapIndex::load(store, None, &cache).await.unwrap(); + assert_eq!(cache.size().await, 0, "loading must not read row sets"); + + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + std::slice::from_ref(&segment), + value_row_id_stream(&[(Some("green"), addr(1, 0))]), + dest_store.as_ref(), + &[None], + ) + .await + .unwrap(); + assert_eq!( + cache.size().await, + 0, + "the merge cached the source segment's row sets" + ); + + // A query on the same segment still caches, so the bypass is scoped to + // the merge rather than disabling caching for the index. + let query = SargableQuery::Equals(ScalarValue::Utf8(Some("red".to_string()))); + segment.search(&query, &NoOpMetricsCollector).await.unwrap(); + assert!(cache.size().await > 0, "queries must still cache row sets"); + } + /// Unsorted input must be rejected. It would reopen a run for a key already /// written, and `load` keeps only the last file offset per key, so the earlier /// row set would vanish silently. Nulls may lead or trail, but only once. From ee8df2323ba65181ab3370a093b33115fbe291fe Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 13:25:55 +0200 Subject: [PATCH 6/7] fix(index): stop writing bitmap keys whose rows were all filtered out An old-data filter or a remap can leave a key with no rows. Only the k-way merge in `merge_index_maps` dropped such a key; the update path wrote it with a zero-row bitmap, and `remap_index_map` did so deliberately, to keep one output row per source key. `BitmapIndex::load` builds `index_map` from the keys column alone, so a written empty key comes back indistinguishable from a live one. Nothing prunes it, and every later merge re-reads and re-emits it, so a column whose values stop occurring -- retired by compaction, deleted, or updated away -- accumulated one permanent directory entry per value it ever held. The resident key directory is what bounds a bitmap index, so the cost lands on the one scarce resource. Skip the entry in `BitmapBatchWriter::emit` instead, the single choke point for every write path, and drop the per-path guards that duplicated it. Absent and present-but-empty are query-equivalent -- both match no rows -- so nothing observable changes for readers, while dropping saves an entry, a bitmap load and a heap step per emptied key in every later merge, and gives every path one rule instead of two. Both filter variants drop emptied fragments, so an `is_empty` set really holds no rows. An all-filtered merge now writes a row-less file, which loads correctly because the file schema carries the value type independently. This also realigns `statistics()`, which counts `index_map` plus the null row, with the file's `num_bitmaps` metadata, which counts entries written. --- rust/lance-index/src/scalar/bitmap.rs | 132 ++++++++++++++++++-------- 1 file changed, 93 insertions(+), 39 deletions(-) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 096c41686fe..15e485d10bc 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -1061,7 +1061,21 @@ impl BitmapBatchWriter { /// Serialize and buffer a single (key, bitmap) pair, flushing the current /// batch to disk if adding it would exceed [`MAX_BUFFERED_BYTES`]. + /// + /// A key whose bitmap is empty is not written at all. This is the single + /// choke point for every write path, so it is the one rule for all of them. pub(crate) async fn emit(&mut self, key: ScalarValue, bitmap: &RowAddrTreeMap) -> Result<()> { + // An old-data filter or a remap can remove every row of a key. Writing + // that key anyway would put it back into `index_map` on load -- which is + // built from the keys column alone and so cannot tell it from a live key + // -- where nothing prunes it, and each later merge would re-read and + // re-emit it. Absent and present-but-empty are query-equivalent, since + // both match no rows. Both filter variants drop emptied fragments, so an + // `is_empty` set really holds no rows. + if bitmap.is_empty() { + return Ok(()); + } + let mut buf = Vec::new(); bitmap.serialize_into(&mut buf).unwrap(); // `key.size()` already covers the `Vec` slot it moves into, @@ -1892,12 +1906,8 @@ pub(crate) async fn build_index_map( /// separately. Nulls live outside `index_map`, in `null_map`, so they are /// remapped separately and emitted first -- a null sorts below every value. /// -/// Emits every key unconditionally, even one whose remapped bitmap comes out -/// empty -- deliberately unlike [`merge_index_maps`], which drops a key its -/// filter empties. Both are query-equivalent, since an absent key and a -/// present-but-empty one both match no rows; this function preserves the key -/// to match the old materialized-map remap path, which always produced one -/// row per source key. +/// A key whose remapped bitmap comes out empty (every one of its rows deleted) +/// is not written, the rule [`BitmapBatchWriter::emit`] applies to every path. pub(crate) async fn remap_index_map( index: &BitmapIndex, mapping: &RowAddrRemap, @@ -1974,9 +1984,8 @@ pub(crate) fn merge_source_entry_count(sources: &[Arc]) -> u64 { /// actually in, since it loads one bitmap per entry. /// /// A key whose merged bitmap comes out empty (every one of its rows retired by -/// `old_data_filter`) is dropped rather than emitted -- deliberately unlike -/// [`remap_index_map`], which preserves such a key. Both are query-equivalent, -/// since an absent key and a present-but-empty one both match no rows. +/// `old_data_filter`) is not written, the rule [`BitmapBatchWriter::emit`] +/// applies to every path. pub(crate) async fn merge_index_maps( sources: &[Arc], old_data_filter: Option<&super::OldIndexDataFilter>, @@ -2000,9 +2009,7 @@ pub(crate) async fn merge_index_maps( .collect(), ); - if let Some(merged_nulls) = old.take_null_bitmap() - && !merged_nulls.is_empty() - { + if let Some(merged_nulls) = old.take_null_bitmap() { 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?; @@ -2018,12 +2025,11 @@ pub(crate) async fn merge_index_maps( let mut last_reported = 0u64; while let Some((key, merged)) = old.take_smallest_if(|_| true).await? { - if !merged.is_empty() { - writer.emit(key, &merged).await?; - } + writer.emit(key, &merged).await?; - // Reported outside the guard above: a key the filter emptied still - // consumed its source entries, and skipping it would stall the count. + // Counted whether or not `emit` wrote the key: one the filter emptied + // still consumed its source entries, and skipping it would stall the + // count. let consumed = old.consumed(); if consumed - last_reported >= report_every && let Some((progress, stage)) = progress @@ -3208,15 +3214,10 @@ 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. + /// Remap must put every address of every source key, nulls included, through + /// the same mapping, and drop a key whose rows were all deleted rather than + /// write it with an empty bitmap -- the one rule `BitmapBatchWriter::emit` + /// applies to every path. #[tokio::test] async fn test_bitmap_remap_matches_materialized_path() { // frag 1 - { 0: null, 1: "a", 2: "b" } @@ -3252,8 +3253,7 @@ mod tests { index.remap(&mapping, dest_store.as_ref()).await.unwrap(); // Read in file order, so this pins the emitted order as well as the - // contents: the null key first, then keys ascending. The old path wrote - // a `HashMap`, in no particular order. + // contents: the null key first, then keys ascending. let frag_3 = |offset: u32| -> Vec { vec![RowAddress::new_from_parts(3, offset).into()] }; let written: Vec<(Option, Vec)> = @@ -3267,17 +3267,24 @@ mod tests { vec![ (None, frag_3(0)), (Some("a".to_string()), frag_3(1)), - // Every row of "b" was deleted, and it still emits a row. - (Some("b".to_string()), Vec::new()), + // Every row of "b" was deleted, so "b" is not written at all. (Some("c".to_string()), vec![addrs[4]]), ] ); - // The emptied key survives the round trip as a key rather than vanishing. + // The emptied key does not come back as a directory entry either. let reloaded = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) .await .unwrap(); - assert_eq!(reloaded.index_map.len(), 3); + assert_eq!(reloaded.index_map.len(), 2); + assert!( + !reloaded + .index_map + .contains_key(&OrderableScalarValue(ScalarValue::Utf8(Some( + "b".to_string() + )))), + "an emptied key must not survive the remap" + ); } #[tokio::test] @@ -3903,19 +3910,19 @@ mod tests { .await .unwrap(); - // A key whose rows were all filtered out may still be materialised - // with an empty bitmap, so only a key outside the vocabulary is a - // fabrication. + // Exactly the keys with at least one surviving row are materialised: + // a key whose rows were all filtered out is not written. for value in VOCAB { let key = OrderableScalarValue(ScalarValue::Utf8(Some(value.to_string()))); let want = expected.get(&Some(value)).cloned().unwrap_or_default(); let got = merged.load_bitmap(&key, None).await.unwrap(); assert_eq!(*got, want, "seed {seed}: row sets differ for {value:?}"); + assert_eq!( + merged.index_map.contains_key(&key), + !want.is_empty(), + "seed {seed}: key {value:?} materialised iff it has rows" + ); } - assert!( - merged.index_map.len() <= VOCAB.len(), - "seed {seed}: merge materialised keys outside the vocabulary" - ); assert_eq!( *merged.null_map, expected.get(&None).cloned().unwrap_or_default(), @@ -3924,6 +3931,53 @@ mod tests { } } + /// A key whose every row the filter removes must not be materialised: `load` + /// builds `index_map` from the keys column alone, so it would come back as a + /// live directory entry that nothing prunes. + #[tokio::test] + async fn test_bitmap_merge_drops_emptied_keys() { + let (_dir, segment) = train_bitmap_segment(&[ + (Some("kept"), addr(0, 0)), + (Some("gone"), addr(5, 0)), + (None, addr(5, 1)), + ]) + .await; + + // Fragment 5 is retired, so "gone" and the null row lose every row. + let filters = vec![Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::from_iter([0u32]), + to_remove: RoaringBitmap::from_iter([5u32]), + })]; + + let (_dest_dir, dest_store) = test_util::index_store(); + BitmapIndex::merge_segments( + &[segment], + value_row_id_stream(&[]), + dest_store.as_ref(), + &filters, + ) + .await + .unwrap(); + let merged = BitmapIndex::load(dest_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + assert_eq!( + merged + .index_map + .keys() + .map(|key| key.0.to_string()) + .collect::>(), + vec!["kept"], + "an emptied key must not be materialised" + ); + assert_eq!(search_addrs(&merged, Some("kept")).await, vec![addr(0, 0)]); + assert!( + merged.null_map.is_empty(), + "an emptied null row must not be materialised" + ); + } + /// A merge must not populate the index cache with the source segments' row /// sets: it reads each one once, and the commit that follows retires those /// segments, so the entries would only evict live ones. From 94ac04a3d593ec1a682577e2909af7755a888cfd Mon Sep 17 00:00:00 2001 From: Adam Cimarosti Date: Wed, 16 Sep 2026 18:23:34 +0200 Subject: [PATCH 7/7] fix(index): carry the last non-null key across a null run in the bitmap ordering check `check_run_order` returned early whenever the previous run was null, so the non-null key before that run was never compared with the one after it. Input such as `["a", NULL, "a"]` passed the check and was written as two `"a"` entries, of which `BitmapIndex::load` keeps only the last, silently losing the first posting. Track the most recent non-null key in `build_index_map` and compare each new non-null key against that rather than against the immediately preceding run. A single null run may still lead or trail; the non-null sequence on either side of it is now checked as one sequence. Two cases cover it: a value reappearing across the null run and a value descending across it. --- rust/lance-index/src/scalar/bitmap.rs | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 15e485d10bc..578cfd69d7c 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -1777,11 +1777,13 @@ pub async fn merge_index_files( /// than merge-joined by value, so a single null run is correct wherever it /// falls, which lets a caller sort nulls first or last. `null_run_closed` says /// whether one has already been flushed, making a second run detectable. +/// `last_non_null` is the most recent non-null key, carried across a null run, +/// so that a value reappearing on the far side of the nulls is still caught. /// /// Both keys come from the same column, so they share a `ScalarValue` variant /// and `OrderableScalarValue`'s `Ord` cannot panic comparing them. fn check_run_order( - previous: &ScalarValue, + last_non_null: Option<&ScalarValue>, next: &ScalarValue, null_run_closed: bool, ) -> Result<()> { @@ -1794,9 +1796,9 @@ fn check_run_order( } return Ok(()); } - if previous.is_null() { + let Some(previous) = last_non_null else { return Ok(()); - } + }; if OrderableScalarValue(next.clone()) <= OrderableScalarValue(previous.clone()) { return Err(Error::invalid_input(format!( "bitmap index: input must be sorted by value, but {next} follows {previous}" @@ -1833,6 +1835,9 @@ pub(crate) async fn build_index_map( // Track whether we emitted a null bitmap (old index stores nulls // separately in null_map, not in index_map). let mut emitted_null = false; + // The most recent non-null key, kept across a null run so the ordering check + // still sees the whole non-null sequence. + let mut last_non_null_key: Option = None; while let Some(batch) = data_source.try_next().await? { let values = batch.column_by_name(VALUE_COLUMN_NAME).expect_ok()?; @@ -1851,7 +1856,7 @@ pub(crate) async fn build_index_map( _ => { // Value changed — flush the previous run. if let Some(prev_key) = current_key.take() { - check_run_order(&prev_key, &key, emitted_null)?; + check_run_order(last_non_null_key.as_ref(), &key, emitted_null)?; let mut prev_bitmap = std::mem::take(&mut current_bitmap); BitmapIndexPlugin::finish_run( prev_key, @@ -1862,6 +1867,9 @@ pub(crate) async fn build_index_map( ) .await?; } + if !key.is_null() { + last_non_null_key = Some(key.clone()); + } current_key = Some(key); current_bitmap = RowAddrTreeMap::default(); current_bitmap.insert(row_id); @@ -4027,6 +4035,10 @@ mod tests { #[rstest] #[case::value_reappears(vec![Some("a"), Some("b"), Some("a")], "a follows b")] #[case::two_null_runs(vec![None, Some("a"), None], "more than one run of nulls")] + // The null run must not reset the ordering: the non-null sequence on either + // side of it is one sequence. + #[case::value_reappears_across_null(vec![Some("a"), None, Some("a")], "a follows a")] + #[case::value_descends_across_null(vec![Some("b"), None, Some("a")], "a follows b")] #[tokio::test] async fn test_bitmap_build_rejects_unsorted_input( #[case] values: Vec>,