perf: various btree performance improvements

rust/lance-arrow/Cargo.toml

@@ -18,6 +18,7 @@ arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
 arrow-data = { workspace = true }
 arrow-cast = { workspace = true }
+arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-select = { workspace = true }
 bytes = { workspace = true }

rust/lance-arrow/src/lib.rs

@@ -18,7 +18,7 @@ use arrow_array::{
 };
 use arrow_buffer::MutableBuffer;
 use arrow_data::ArrayDataBuilder;
-use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema};
+use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema, SortOptions};
 use arrow_select::{interleave::interleave, take::take};
 use rand::prelude::*;
 
@@ -604,6 +604,9 @@ pub trait RecordBatchExt {
 
     /// Create a new RecordBatch with compacted memory after slicing.
     fn shrink_to_fit(&self) -> Result<RecordBatch>;
+
+    /// Helper method to sort the RecordBatch by a column
+    fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<RecordBatch>;
 }
 
 impl RecordBatchExt for RecordBatch {
@@ -778,6 +781,18 @@ impl RecordBatchExt for RecordBatch {
         // Deep copy the sliced record batch, instead of whole batch
         crate::deepcopy::deep_copy_batch_sliced(self)
     }
+
+    fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<Self> {
+        if column >= self.num_columns() {
+            return Err(ArrowError::InvalidArgumentError(format!(
+                "Column index out of bounds: {}",
+                column
+            )));
+        }
+        let column = self.column(column);
+        let sorted = arrow_ord::sort::sort_to_indices(column, options, None)?;
+        self.take(&sorted)
+    }
 }
 
 fn project(struct_array: &StructArray, fields: &Fields) -> Result<StructArray> {

rust/lance-core/Cargo.toml

@@ -24,6 +24,7 @@ datafusion-common = { workspace = true, optional = true }
 datafusion-sql = { workspace = true, optional = true }
 deepsize.workspace = true
 futures.workspace = true
+itertools.workspace = true
 libc.workspace = true
 mock_instant.workspace = true
 moka.workspace = true

rust/lance-core/src/error.rs

@@ -184,7 +184,7 @@ impl<T> LanceOptionExt<T> for Option<T> {
     }
 }
 
-trait ToSnafuLocation {
+pub trait ToSnafuLocation {
     fn to_snafu_location(&'static self) -> snafu::Location;
 }
 

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

@@ -10,9 +10,11 @@ use arrow_array::{Array, BinaryArray, GenericBinaryArray};
 use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer};
 use byteorder::{ReadBytesExt, WriteBytesExt};
 use deepsize::DeepSizeOf;
+use itertools::Itertools;
 use roaring::{MultiOps, RoaringBitmap, RoaringTreemap};
 
-use crate::Result;
+use crate::error::ToSnafuLocation;
+use crate::{Error, Result};
 
 use super::address::RowAddress;
 
@@ -595,6 +597,30 @@ impl RowAddrTreeMap {
                 }),
             })
     }
+
+    #[track_caller]
+    pub fn from_sorted_iter(iter: impl IntoIterator<Item = u64>) -> Result<Self> {
+        let mut iter = iter.into_iter().peekable();
+        let mut inner = BTreeMap::new();
+
+        while let Some(row_id) = iter.peek() {
+            let fragment_id = (row_id >> 32) as u32;
+            let next_bitmap_iter = iter
+                .peeking_take_while(|row_id| (row_id >> 32) as u32 == fragment_id)
+                .map(|row_id| row_id as u32);
+            let Ok(bitmap) = RoaringBitmap::from_sorted_iter(next_bitmap_iter) else {
+                return Err(Error::Internal {
+                    message: "RowAddrTreeMap::from_sorted_iter called with non-sorted input"
+                        .to_string(),
+                    // Use the caller location since we aren't the one that got it out of order
+                    location: std::panic::Location::caller().to_snafu_location(),
+                });
+            };
+            inner.insert(fragment_id, RowAddrSelection::Partial(bitmap));
+        }
+
+        Ok(Self { inner })
+    }
 }
 
 impl std::ops::BitOr<Self> for RowAddrTreeMap {
@@ -1504,6 +1530,17 @@ mod tests {
             prop_assert_eq!(expected, left);
         }
 
+        #[test]
+        fn test_from_sorted_iter(
+            mut rows in proptest::collection::vec(0..u64::MAX, 0..1000)
+        ) {
+            rows.sort();
+            let num_rows = rows.len();
+            let mask = RowAddrTreeMap::from_sorted_iter(rows).unwrap();
+            prop_assert_eq!(mask.len(), Some(num_rows as u64));
+        }
+
+
     }
 
     #[test]

rust/lance-index/benches/btree.rs

@@ -19,15 +19,13 @@ use std::{
     time::Duration,
 };
 
-use arrow_schema::DataType;
 use common::{LOW_CARDINALITY_COUNT, TOTAL_ROWS};
 use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion_common::ScalarValue;
 use lance_core::cache::LanceCache;
 use lance_index::metrics::NoOpMetricsCollector;
 use lance_index::pbold;
 use lance_index::scalar::btree::{train_btree_index, BTreeIndexPlugin, DEFAULT_BTREE_BATCH_SIZE};
-use lance_index::scalar::flat::FlatIndexMetadata;
 use lance_index::scalar::lance_format::LanceIndexStore;
 use lance_index::scalar::registry::ScalarIndexPlugin;
 use lance_index::scalar::{SargableQuery, ScalarIndex};
@@ -107,17 +105,10 @@ async fn create_int_unique_index(
     use_cache: bool,
 ) -> Arc<dyn ScalarIndex> {
     let stream = common::generate_int_unique_stream();
-    let sub_index = FlatIndexMetadata::new(DataType::Int64);
-
-    train_btree_index(
-        stream,
-        &sub_index,
-        store.as_ref(),
-        DEFAULT_BTREE_BATCH_SIZE,
-        None,
-    )
-    .await
-    .unwrap();
+
+    train_btree_index(stream, store.as_ref(), DEFAULT_BTREE_BATCH_SIZE, None)
+        .await
+        .unwrap();
 
     let cache = get_cache(use_cache, "int_unique");
     let details = prost_types::Any::from_msg(&pbold::BTreeIndexDetails::default()).unwrap();
@@ -135,17 +126,10 @@ async fn create_int_low_card_index(
     use_cache: bool,
 ) -> Arc<dyn ScalarIndex> {
     let stream = common::generate_int_low_cardinality_stream();
-    let sub_index = FlatIndexMetadata::new(DataType::Int64);
-
-    train_btree_index(
-        stream,
-        &sub_index,
-        store.as_ref(),
-        DEFAULT_BTREE_BATCH_SIZE,
-        None,
-    )
-    .await
-    .unwrap();
+
+    train_btree_index(stream, store.as_ref(), DEFAULT_BTREE_BATCH_SIZE, None)
+        .await
+        .unwrap();
 
     let cache = get_cache(use_cache, "int_low_card");
     let details = prost_types::Any::from_msg(&pbold::BTreeIndexDetails::default()).unwrap();
@@ -163,17 +147,10 @@ async fn create_string_unique_index(
     use_cache: bool,
 ) -> Arc<dyn ScalarIndex> {
     let stream = common::generate_string_unique_stream();
-    let sub_index = FlatIndexMetadata::new(DataType::Utf8);
-
-    train_btree_index(
-        stream,
-        &sub_index,
-        store.as_ref(),
-        DEFAULT_BTREE_BATCH_SIZE,
-        None,
-    )
-    .await
-    .unwrap();
+
+    train_btree_index(stream, store.as_ref(), DEFAULT_BTREE_BATCH_SIZE, None)
+        .await
+        .unwrap();
 
     let cache = get_cache(use_cache, "string_unique");
     let details = prost_types::Any::from_msg(&pbold::BTreeIndexDetails::default()).unwrap();
@@ -191,17 +168,10 @@ async fn create_string_low_card_index(
     use_cache: bool,
 ) -> Arc<dyn ScalarIndex> {
     let stream = common::generate_string_low_cardinality_stream();
-    let sub_index = FlatIndexMetadata::new(DataType::Utf8);
-
-    train_btree_index(
-        stream,
-        &sub_index,
-        store.as_ref(),
-        DEFAULT_BTREE_BATCH_SIZE,
-        None,
-    )
-    .await
-    .unwrap();
+
+    train_btree_index(stream, store.as_ref(), DEFAULT_BTREE_BATCH_SIZE, None)
+        .await
+        .unwrap();
 
     let cache = get_cache(use_cache, "string_low_card");
     let details = prost_types::Any::from_msg(&pbold::BTreeIndexDetails::default()).unwrap();

rust/lance-index/src/scalar.rs

@@ -32,7 +32,6 @@ pub mod bitmap;
 pub mod bloomfilter;
 pub mod btree;
 pub mod expression;
-pub mod flat;
 pub mod inverted;
 pub mod json;
 pub mod label_list;