diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 8f5b12feced..578cfd69d7c 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}; @@ -510,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"])) @@ -527,16 +542,104 @@ 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 { &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 +961,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?; @@ -956,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, @@ -1316,6 +1435,142 @@ 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::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 +/// 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` 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 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))); + } + } + + // 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 +1590,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 +1606,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. @@ -1518,6 +1764,49 @@ 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. +/// `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( + last_non_null: Option<&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(()); + } + 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}" + ))); + } + Ok(()) +} + /// Build a bitmap index map from value-sorted `(value, row_id)` input, emitting /// one key at a time into `writer`. /// @@ -1525,28 +1814,20 @@ 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. 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. +/// 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_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; @@ -1554,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()?; @@ -1572,35 +1856,20 @@ 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(last_non_null_key.as_ref(), &key, emitted_null)?; let mut prev_bitmap = std::mem::take(&mut current_bitmap); BitmapIndexPlugin::finish_run( prev_key, &mut prev_bitmap, - old_index, - &mut old_keys, + &mut old, &mut emitted_null, writer, - old_data_filter, ) .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); @@ -1615,32 +1884,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?; } @@ -1655,12 +1914,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, @@ -1723,22 +1978,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. @@ -1747,9 +1992,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>, @@ -1763,18 +2007,22 @@ 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() { 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 +2032,13 @@ 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(); + while let Some((key, merged)) = old.take_smallest_if(|_| true).await? { + writer.emit(key, &merged).await?; - // 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); - if !merged.is_empty() { - writer.emit(next_key.0.clone(), &merged).await?; - } - - // Reported outside the guard above: a key the filter emptied still - // consumed its source entries, and skipping it would stall the count. + // 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 { @@ -1846,64 +2051,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 +2351,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 @@ -2179,6 +2376,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) { @@ -3024,15 +3222,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" } @@ -3068,8 +3261,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)> = @@ -3083,17 +3275,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] @@ -3239,8 +3438,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 +3485,597 @@ 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); + } + + /// 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(); + + // 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_eq!( + *merged.null_map, + expected.get(&None).cloned().unwrap_or_default(), + "seed {seed}: null row sets differ" + ); + } + } + + /// 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. + #[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. + #[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>, + #[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. @@ -3356,7 +4139,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 +4215,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..10dba71a58a 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}; @@ -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 @@ -570,8 +568,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 +681,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..27b5b728334 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,437 @@ 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 { + count_rows_where(dataset, &format!("cat = '{cat}'")).await + } + + async fn count_rows_where(dataset: &Dataset, predicate: &str) -> usize { + dataset + .scan() + .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. + #[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(); + // 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)] { + assert_eq!( + count_cat(&dataset, cat).await, + expected, + "bitmap merge returned stale rows for cat = {cat}" + ); + } + } } 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}; 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,