perf: speed up materialization when there is no block list (#1897)

Materialization converts a row id mask into a vector of row ids. The old
approach used a FragIdIter which was too slow. In fact, we probably want
to eventually get rid of FragIdIter entirely (see
#1896). In the meantime, when
there is no block list, we can materialize by filtering the allow list
with the fragment list instead.

This should greatly speed up indexed searches that don't have any vector
search component.

Author: westonpace

python/python/benchmarks/test_search.py

@@ -54,8 +54,18 @@ def create_base_dataset(data_dir: Path) -> lance.LanceDataset:
     if dataset:
         return dataset
 
-    table = create_table(NUM_ROWS, offset=0)
-    dataset = lance.write_dataset(table, tmp_path)
+    rows_remaining = NUM_ROWS
+    offset = 0
+    dataset = None
+    while rows_remaining > 0:
+        next_batch_length = min(rows_remaining, 1024 * 1024)
+        rows_remaining -= next_batch_length
+        table = create_table(next_batch_length, offset)
+        if offset == 0:
+            dataset = lance.write_dataset(table, tmp_path)
+        else:
+            dataset = lance.write_dataset(table, tmp_path, mode="append")
+        offset += next_batch_length
 
     dataset.create_index(
         column="vector",
@@ -66,6 +76,8 @@ def create_base_dataset(data_dir: Path) -> lance.LanceDataset:
         num_bits=8,
     )
 
+    dataset.create_scalar_index("filterable", "BTREE")
+
     return dataset
 
 
@@ -208,14 +220,23 @@ def test_filtered_search(test_dataset, benchmark, selectivity, prefilter, use_in
         "filterable < 5000",
         "filterable > 5000",
     ),
+    ids=[
+        "none",
+        "equality",
+        "not_equality",
+        "in_list_one",
+        "not_in_list_one",
+        "not_equality_and_chain",
+        "not_in_list_three",
+        "less_than_selective",
+        "greater_than_not_selective",
+    ],
 )
-def test_scalar_index_prefilter(datasets: Datasets, benchmark, filter: str):
-    dataset = datasets.clean
-    dataset.create_scalar_index("filterable", "BTREE")
+def test_scalar_index_prefilter(test_dataset, benchmark, filter: str):
     q = pc.random(N_DIMS).cast(pa.float32())
     if filter is None:
         benchmark(
-            dataset.to_table,
+            test_dataset.to_table,
             nearest=dict(
                 column="vector",
                 q=q,
@@ -225,7 +246,7 @@ def test_scalar_index_prefilter(datasets: Datasets, benchmark, filter: str):
         )
     else:
         benchmark(
-            dataset.to_table,
+            test_dataset.to_table,
             nearest=dict(
                 column="vector",
                 q=q,
@@ -235,3 +256,44 @@ def test_scalar_index_prefilter(datasets: Datasets, benchmark, filter: str):
             prefilter=True,
             filter=filter,
         )
+
+
+@pytest.mark.benchmark(group="query_no_vec")
+@pytest.mark.parametrize(
+    "filter",
+    (
+        None,
+        "filterable = 0",
+        "filterable != 0",
+        "filterable IN (0)",
+        "filterable IN (0, 5000, 10000)",
+        "filterable NOT IN (0)",
+        "filterable != 0 AND filterable != 5000 AND filterable != 10000",
+        "filterable NOT IN (0, 5000, 10000)",
+        "filterable < 5000",
+        "filterable > 5000",
+    ),
+    ids=[
+        "none",
+        "equality",
+        "not_equality",
+        "in_list_one",
+        "in_list_three",
+        "not_in_list_one",
+        "not_equality_and_chain",
+        "not_in_list_three",
+        "less_than_selective",
+        "greater_than_not_selective",
+    ],
+)
+def test_scalar_index_search(test_dataset, benchmark, filter: str):
+    if filter is None:
+        benchmark(
+            test_dataset.to_table,
+        )
+    else:
+        benchmark(
+            test_dataset.to_table,
+            prefilter=True,
+            filter=filter,
+        )

rust/lance-core/src/utils/mask.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashSet;
 use std::io::Write;
 use std::{collections::BTreeMap, io::Read};
 
@@ -22,6 +23,8 @@ use roaring::RoaringBitmap;
 
 use crate::Result;
 
+use super::address::RowAddress;
+
 /// A row id mask to select or deselect particular row ids
 ///
 /// If both the allow_list and the block_list are Some then the only selected
@@ -299,6 +302,43 @@ impl RowIdTreeMap {
         self.inner.is_empty()
     }
 
+    /// The number of rows in the map
+    ///
+    /// If there are any "full fragment" items then this is unknown and None is returned
+    pub fn len(&self) -> Option<u64> {
+        self.inner
+            .values()
+            .map(|row_id_selection| match row_id_selection {
+                RowIdSelection::Full => None,
+                RowIdSelection::Partial(indices) => Some(indices.len()),
+            })
+            .try_fold(0_u64, |acc, next| next.map(|next| next + acc))
+    }
+
+    /// An iterator of row ids
+    ///
+    /// If there are any "full fragment" items then this can't be calculated and None
+    /// is returned
+    pub fn row_ids(&self) -> Option<impl Iterator<Item = RowAddress> + '_> {
+        let inner_iters = self
+            .inner
+            .iter()
+            .filter_map(|(frag_id, row_id_selection)| match row_id_selection {
+                RowIdSelection::Full => None,
+                RowIdSelection::Partial(bitmap) => Some(
+                    bitmap
+                        .iter()
+                        .map(|row_offset| RowAddress::new_from_parts(*frag_id, row_offset)),
+                ),
+            })
+            .collect::<Vec<_>>();
+        if inner_iters.len() != self.inner.len() {
+            None
+        } else {
+            Some(inner_iters.into_iter().flatten())
+        }
+    }
+
     /// Add a bitmap for a single fragment
     pub fn insert_bitmap(&mut self, fragment: u32, bitmap: RoaringBitmap) {
         self.inner.insert(fragment, RowIdSelection::Partial(bitmap));
@@ -320,6 +360,12 @@ impl RowIdTreeMap {
         }
     }
 
+    pub fn remove_fragments(&mut self, frag_ids: impl IntoIterator<Item = u32>) {
+        let frag_id_set = frag_ids.into_iter().collect::<HashSet<_>>();
+        self.inner
+            .retain(|frag_id, _| frag_id_set.contains(frag_id));
+    }
+
     /// Compute the serialized size of the set.
     pub fn serialized_size(&self) -> usize {
         // Starts at 4 because of the u32 num_entries

rust/lance/src/io/exec/scalar_index.rs

@@ -33,7 +33,7 @@ use lance_table::format::Fragment;
 use pin_project::pin_project;
 use roaring::RoaringBitmap;
 use snafu::{location, Location};
-use tracing::instrument;
+use tracing::{debug_span, instrument};
 
 use crate::{
     index::{prefilter::PreFilter, DatasetIndexInternalExt},
@@ -257,29 +257,59 @@ impl MaterializeIndexExec {
     ) -> Result<RecordBatch> {
         // TODO: multiple batches, stream without materializing all row ids in memory
         let mask = expr.evaluate(dataset.as_ref());
-        let fragment_bitmap = RoaringBitmap::from_iter(fragments.iter().map(|frag| frag.id as u32));
-        // The user-requested `fragments` is guaranteed to be stricter than the index's fragment
-        // bitmap.  This node only runs on indexed fragments and any fragments that were deleted
-        // when the index was trained will still be deleted when the index is queried.
-        let prefilter = PreFilter::create_deletion_mask(dataset.clone(), fragment_bitmap);
+        let span = debug_span!("create_prefilter");
+        let prefilter = span.in_scope(|| {
+            let fragment_bitmap =
+                RoaringBitmap::from_iter(fragments.iter().map(|frag| frag.id as u32));
+            // The user-requested `fragments` is guaranteed to be stricter than the index's fragment
+            // bitmap.  This node only runs on indexed fragments and any fragments that were deleted
+            // when the index was trained will still be deleted when the index is queried.
+            PreFilter::create_deletion_mask(dataset.clone(), fragment_bitmap)
+        });
         let mask = if let Some(prefilter) = prefilter {
             let (mask, prefilter) = futures::try_join!(mask, prefilter)?;
             mask.also_block((*prefilter).clone())
         } else {
             mask.await?
         };
-        let ids = match (mask.allow_list, mask.block_list) {
+        let span = debug_span!("make_ids");
+        let ids = span.in_scope(|| match (mask.allow_list, mask.block_list) {
             (None, None) => FragIdIter::new(fragments).collect::<Vec<_>>(),
-            (Some(allow_list), None) => FragIdIter::new(fragments)
-                .filter(|row_id| allow_list.contains(*row_id))
-                .collect(),
+            (Some(mut allow_list), None) => {
+                allow_list.remove_fragments(fragments.iter().map(|frag| frag.id as u32));
+                if let Some(allow_list_iter) = allow_list.row_ids() {
+                    allow_list_iter.map(u64::from).collect::<Vec<_>>()
+                } else {
+                    FragIdIter::new(fragments)
+                        .filter(|row_id| allow_list.contains(*row_id))
+                        .collect()
+                }
+            }
             (None, Some(block_list)) => FragIdIter::new(fragments)
                 .filter(|row_id| !block_list.contains(*row_id))
                 .collect(),
-            (Some(allow_list), Some(block_list)) => FragIdIter::new(fragments)
-                .filter(|row_id| !block_list.contains(*row_id) && allow_list.contains(*row_id))
-                .collect(),
-        };
+            (Some(mut allow_list), Some(block_list)) => {
+                allow_list.remove_fragments(fragments.iter().map(|frag| frag.id as u32));
+                if let Some(allow_list_iter) = allow_list.row_ids() {
+                    allow_list_iter
+                        .filter_map(|addr| {
+                            let row_id = u64::from(addr);
+                            if !block_list.contains(row_id) {
+                                Some(row_id)
+                            } else {
+                                None
+                            }
+                        })
+                        .collect::<Vec<_>>()
+                } else {
+                    FragIdIter::new(fragments)
+                        .filter(|row_id| {
+                            !block_list.contains(*row_id) && allow_list.contains(*row_id)
+                        })
+                        .collect()
+                }
+            }
+        });
         let ids = UInt64Array::from(ids);
         Ok(RecordBatch::try_new(
             MATERIALIZE_INDEX_SCHEMA.clone(),