diff --git a/.bumpversion.toml b/.bumpversion.toml index 4cac71e4e51..0b6016416b1 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -95,6 +95,11 @@ filename = "Cargo.toml" search = 'lance-namespace-impls = {{ version = "={current_version}"' replace = 'lance-namespace-impls = {{ version = "={new_version}"' +[[tool.bumpversion.files]] +filename = "Cargo.toml" +search = 'lance-select = {{ version = "={current_version}"' +replace = 'lance-select = {{ version = "={new_version}"' + [[tool.bumpversion.files]] filename = "Cargo.toml" search = 'lance-table = {{ version = "={current_version}"' diff --git a/Cargo.lock b/Cargo.lock index 6fb213092ed..1b902bd3e4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4378,6 +4378,7 @@ dependencies = [ "lance-io", "lance-linalg", "lance-namespace", + "lance-select", "lance-table", "lance-test-macros", "lance-testing", @@ -4487,7 +4488,6 @@ dependencies = [ "async-trait", "byteorder", "bytes", - "criterion", "datafusion-common", "datafusion-sql", "deepsize", @@ -4740,6 +4740,7 @@ dependencies = [ "lance-geo", "lance-io", "lance-linalg", + "lance-select", "lance-table", "lance-testing", "lance-tokenizer", @@ -4924,6 +4925,23 @@ dependencies = [ "url", ] +[[package]] +name = "lance-select" +version = "7.1.0-beta.1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "byteorder", + "bytes", + "criterion", + "deepsize", + "itertools 0.13.0", + "lance-core", + "proptest", + "roaring", + "rstest", +] + [[package]] name = "lance-table" version = "7.1.0-beta.1" @@ -4947,6 +4965,7 @@ dependencies = [ "lance-datagen", "lance-file", "lance-io", + "lance-select", "log", "object_store", "pprof", diff --git a/Cargo.toml b/Cargo.toml index fdb44b1342b..7bd2192e770 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ "rust/lance-namespace", "rust/lance-namespace-impls", "rust/lance-namespace-datafusion", + "rust/lance-select", "rust/lance-tokenizer", "rust/lance-table", "rust/lance-test-macros", @@ -70,6 +71,7 @@ lance-namespace = { version = "=7.1.0-beta.1", path = "./rust/lance-namespace" } lance-namespace-impls = { version = "=7.1.0-beta.1", path = "./rust/lance-namespace-impls" } lance-namespace-datafusion = { version = "=7.0.0-beta.9", path = "./rust/lance-namespace-datafusion" } lance-namespace-reqwest-client = "0.7.7" +lance-select = { version = "=7.1.0-beta.1", path = "./rust/lance-select" } lance-tokenizer = { version = "=7.1.0-beta.1", path = "./rust/lance-tokenizer" } lance-table = { version = "=7.1.0-beta.1", path = "./rust/lance-table" } lance-test-macros = { version = "=7.1.0-beta.1", path = "./rust/lance-test-macros" } diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 2d2ae4b24c1..264ea4294b7 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -3660,6 +3660,7 @@ dependencies = [ "lance-io", "lance-linalg", "lance-namespace", + "lance-select", "lance-table", "lance-tokenizer", "log", @@ -3921,6 +3922,7 @@ dependencies = [ "lance-geo", "lance-io", "lance-linalg", + "lance-select", "lance-table", "lance-tokenizer", "libm", @@ -4096,6 +4098,20 @@ dependencies = [ "url", ] +[[package]] +name = "lance-select" +version = "7.1.0-beta.1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "byteorder", + "bytes", + "deepsize", + "itertools 0.13.0", + "lance-core", + "roaring", +] + [[package]] name = "lance-table" version = "7.1.0-beta.1" @@ -4115,6 +4131,7 @@ dependencies = [ "lance-core", "lance-file", "lance-io", + "lance-select", "log", "object_store", "prost", diff --git a/python/Cargo.lock b/python/Cargo.lock index ae5c73e68c5..399f18afee3 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4020,6 +4020,7 @@ dependencies = [ "lance-io", "lance-linalg", "lance-namespace", + "lance-select", "lance-table", "lance-tokenizer", "log", @@ -4282,6 +4283,7 @@ dependencies = [ "lance-geo", "lance-io", "lance-linalg", + "lance-select", "lance-table", "lance-tokenizer", "libm", @@ -4421,6 +4423,20 @@ dependencies = [ "url", ] +[[package]] +name = "lance-select" +version = "7.1.0-beta.1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "byteorder", + "bytes", + "deepsize", + "itertools 0.13.0", + "lance-core", + "roaring", +] + [[package]] name = "lance-table" version = "7.1.0-beta.1" @@ -4442,6 +4458,7 @@ dependencies = [ "lance-core", "lance-file", "lance-io", + "lance-select", "log", "object_store", "prost", diff --git a/rust/lance-core/Cargo.toml b/rust/lance-core/Cargo.toml index ccab121b43f..9dff4b001a4 100644 --- a/rust/lance-core/Cargo.toml +++ b/rust/lance-core/Cargo.toml @@ -48,16 +48,11 @@ log.workspace = true libc = { version = "0.2" } [dev-dependencies] -criterion.workspace = true proptest.workspace = true rstest.workspace = true [features] datafusion = ["dep:datafusion-common", "dep:datafusion-sql"] -[[bench]] -name = "row_addr_mask" -harness = false - [lints] workspace = true diff --git a/rust/lance-core/src/utils.rs b/rust/lance-core/src/utils.rs index 565036311f9..a7ac74a5b27 100644 --- a/rust/lance-core/src/utils.rs +++ b/rust/lance-core/src/utils.rs @@ -11,7 +11,6 @@ pub mod cpu; pub mod deletion; pub mod futures; pub mod hash; -pub mod mask; pub mod parse; pub mod path; pub mod tempfile; diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index af5e0204320..8bde58fce7e 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -45,6 +45,7 @@ lance-file.workspace = true lance-geo = { workspace = true, optional = true } lance-io.workspace = true lance-linalg.workspace = true +lance-select.workspace = true lance-tokenizer.workspace = true lance-table.workspace = true libm.workspace = true diff --git a/rust/lance-index/src/frag_reuse.rs b/rust/lance-index/src/frag_reuse.rs index 309df62d4c3..d145108d3c0 100644 --- a/rust/lance-index/src/frag_reuse.rs +++ b/rust/lance-index/src/frag_reuse.rs @@ -8,8 +8,8 @@ use arrow_array::{Array, ArrayRef, PrimitiveArray, RecordBatch, UInt64Array}; use async_trait::async_trait; use deepsize::{Context, DeepSizeOf}; use itertools::Itertools; -use lance_core::utils::mask::RowAddrTreeMap; use lance_core::{Error, Result}; +use lance_select::RowAddrTreeMap; use lance_table::format::pb::fragment_reuse_index_details::InlineContent; use lance_table::format::{ExternalFile, Fragment, pb}; use roaring::{RoaringBitmap, RoaringTreemap}; diff --git a/rust/lance-index/src/prefilter.rs b/rust/lance-index/src/prefilter.rs index b431ac2fc24..6671bac0a81 100644 --- a/rust/lance-index/src/prefilter.rs +++ b/rust/lance-index/src/prefilter.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use async_trait::async_trait; use lance_core::Result; -use lance_core::utils::mask::RowAddrMask; +use lance_select::RowAddrMask; /// A trait to be implemented by anything supplying a prefilter row addr mask /// diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index c89d75c11db..bc0e39c206f 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -21,9 +21,9 @@ use datafusion_expr::Expr; use datafusion_expr::expr::ScalarFunction; use deepsize::DeepSizeOf; use inverted::query::{FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, fill_fts_query_column}; -use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use lance_core::{Error, Result}; use lance_io::stream::{RecordBatchStream, RecordBatchStreamAdapter}; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use roaring::RoaringBitmap; use serde::Serialize; diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 76d387f92b7..eb0276dcb9f 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -23,17 +23,14 @@ use lance_arrow::ipc::{ read_ipc_stream_single_at, read_len_prefixed_bytes_at, write_ipc_stream, write_len_prefixed_bytes, }; -use lance_core::utils::mask::RowSetOps; use lance_core::{ Error, ROW_ID, Result, cache::{CacheCodec, CacheCodecImpl, CacheKey, LanceCache, WeakLanceCache}, error::LanceOptionExt, - utils::{ - mask::{NullableRowAddrSet, RowAddrTreeMap}, - tokio::get_num_compute_intensive_cpus, - }, + utils::tokio::get_num_compute_intensive_cpus, }; use lance_io::object_store::ObjectStore; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use object_store::path::Path; use roaring::RoaringBitmap; use serde::{Deserialize, Serialize}; @@ -1761,9 +1758,9 @@ mod tests { } use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::stream; - use lance_core::utils::mask::RowSetOps; use lance_core::utils::{address::RowAddress, tempfile::TempObjDir}; use lance_io::object_store::ObjectStore; + use lance_select::RowSetOps; use std::collections::HashMap; fn assert_state_roundtrips(state: &BitmapIndexState) { @@ -1947,8 +1944,8 @@ mod tests { use arrow_schema::DataType; use datafusion_common::ScalarValue; use lance_core::cache::LanceCache; - use lance_core::utils::mask::RowAddrTreeMap; use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; use std::collections::HashMap; use std::sync::Arc; diff --git a/rust/lance-index/src/scalar/bloomfilter.rs b/rust/lance-index/src/scalar/bloomfilter.rs index b861340b9ee..95807c47d3f 100644 --- a/rust/lance-index/src/scalar/bloomfilter.rs +++ b/rust/lance-index/src/scalar/bloomfilter.rs @@ -1154,12 +1154,9 @@ mod tests { use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_common::ScalarValue; use futures::{StreamExt, stream}; - use lance_core::{ - ROW_ADDR, - cache::LanceCache, - utils::{mask::RowAddrTreeMap, tempfile::TempObjDir}, - }; + use lance_core::{ROW_ADDR, cache::LanceCache, utils::tempfile::TempObjDir}; use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; use crate::scalar::{ BloomFilterQuery, ScalarIndex, SearchResult, diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index eba6f1c6205..6de490b9572 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -50,7 +50,6 @@ use lance_core::{ cache::{CacheCodec, CacheCodecImpl, CacheKey, LanceCache, WeakLanceCache}, error::LanceOptionExt, utils::{ - mask::NullableRowAddrSet, tokio::get_num_compute_intensive_cpus, tracing::{IO_TYPE_LOAD_SCALAR_PART, TRACE_IO_EVENTS}, }, @@ -60,6 +59,7 @@ use lance_datafusion::{ exec::{LanceExecutionOptions, OneShotExec, execute_plan}, }; use lance_io::object_store::ObjectStore; +use lance_select::NullableRowAddrSet; use log::{debug, warn}; use object_store::{Error as ObjectStoreError, path::Path}; use rangemap::RangeInclusiveMap; @@ -2987,12 +2987,12 @@ mod tests { use deepsize::DeepSizeOf; use futures::TryStreamExt; use futures::stream; - use lance_core::utils::mask::RowSetOps; + use lance_core::cache::LanceCache; use lance_core::utils::tempfile::TempObjDir; - use lance_core::{cache::LanceCache, utils::mask::RowAddrTreeMap}; use lance_datafusion::{chunker::break_stream, datagen::DatafusionDatagenExt}; use lance_datagen::{ArrayGeneratorExt, BatchCount, RowCount, array, gen_batch}; use lance_io::object_store::ObjectStore; + use lance_select::{RowAddrTreeMap, RowSetOps}; use object_store::path::Path; use crate::metrics::LocalMetricsCollector; diff --git a/rust/lance-index/src/scalar/btree/flat.rs b/rust/lance-index/src/scalar/btree/flat.rs index 10f0b1ad339..4240753772b 100644 --- a/rust/lance-index/src/scalar/btree/flat.rs +++ b/rust/lance-index/src/scalar/btree/flat.rs @@ -18,7 +18,7 @@ use lance_arrow::ipc::{read_ipc_stream_single_at, read_len_prefixed_bytes_at, wr use lance_core::Result; use lance_core::cache::CacheCodecImpl; use lance_core::utils::address::RowAddress; -use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use roaring::RoaringBitmap; use tracing::instrument; diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index d8f61d7dabc..00ac3ebec7d 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -24,11 +24,9 @@ use super::{ }; #[cfg(feature = "geo")] use super::{GeoQuery, RelationQuery}; -use lance_core::{ - Error, Result, - utils::mask::{NullableRowAddrMask, RowAddrMask}, -}; +use lance_core::{Error, Result}; use lance_datafusion::{expr::safe_coerce_scalar, planner::Planner}; +use lance_select::{NullableRowAddrMask, RowAddrMask}; use roaring::RoaringBitmap; use tracing::instrument; @@ -1308,12 +1306,12 @@ pub static INDEX_EXPR_RESULT_SCHEMA: LazyLock = LazyLock::new(|| { ])) }); -#[derive(Debug)] -enum NullableIndexExprResult { - Exact(NullableRowAddrMask), - AtMost(NullableRowAddrMask), - AtLeast(NullableRowAddrMask), -} +// `IndexExprResult` and `NullableIndexExprResult` themselves live in the +// `lance-select` crate so that benchmarks and downstream consumers can +// depend on the mask substrate without pulling in all of `lance-index`. +// The wire-format helpers below stay here, where the `INDEX_EXPR_RESULT_SCHEMA` +// constant they reference is defined. +pub use lance_select::{IndexExprResult, NullableIndexExprResult}; impl From for NullableIndexExprResult { fn from(result: SearchResult) -> Self { @@ -1325,132 +1323,50 @@ impl From for NullableIndexExprResult { } } -impl std::ops::BitAnd for NullableIndexExprResult { - type Output = Self; - - fn bitand(self, rhs: Self) -> Self { - match (self, rhs) { - (Self::Exact(lhs), Self::Exact(rhs)) => Self::Exact(lhs & rhs), - (Self::Exact(lhs), Self::AtMost(rhs)) | (Self::AtMost(lhs), Self::Exact(rhs)) => { - Self::AtMost(lhs & rhs) - } - (Self::Exact(exact), Self::AtLeast(_)) | (Self::AtLeast(_), Self::Exact(exact)) => { - // We could do better here, elements in both lhs and rhs are known - // to be true and don't require a recheck. We only need to recheck - // elements in lhs that are not in rhs - Self::AtMost(exact) - } - (Self::AtMost(lhs), Self::AtMost(rhs)) => Self::AtMost(lhs & rhs), - (Self::AtLeast(lhs), Self::AtLeast(rhs)) => Self::AtLeast(lhs & rhs), - (Self::AtMost(most), Self::AtLeast(_)) | (Self::AtLeast(_), Self::AtMost(most)) => { - Self::AtMost(most) - } - } - } -} - -impl std::ops::BitOr for NullableIndexExprResult { - type Output = Self; - - fn bitor(self, rhs: Self) -> Self { - match (self, rhs) { - (Self::Exact(lhs), Self::Exact(rhs)) => Self::Exact(lhs | rhs), - (Self::Exact(lhs), Self::AtMost(rhs)) | (Self::AtMost(rhs), Self::Exact(lhs)) => { - // We could do better here, elements in lhs are known to be true - // and don't require a recheck. We only need to recheck elements - // in rhs that are not in lhs - Self::AtMost(lhs | rhs) - } - (Self::Exact(lhs), Self::AtLeast(rhs)) | (Self::AtLeast(rhs), Self::Exact(lhs)) => { - Self::AtLeast(lhs | rhs) - } - (Self::AtMost(lhs), Self::AtMost(rhs)) => Self::AtMost(lhs | rhs), - (Self::AtLeast(lhs), Self::AtLeast(rhs)) => Self::AtLeast(lhs | rhs), - (Self::AtMost(_), Self::AtLeast(least)) | (Self::AtLeast(least), Self::AtMost(_)) => { - Self::AtLeast(least) - } - } - } -} - -impl NullableIndexExprResult { - pub fn drop_nulls(self) -> IndexExprResult { - match self { - Self::Exact(mask) => IndexExprResult::Exact(mask.drop_nulls()), - Self::AtMost(mask) => IndexExprResult::AtMost(mask.drop_nulls()), - Self::AtLeast(mask) => IndexExprResult::AtLeast(mask.drop_nulls()), - } - } -} - -#[derive(Debug)] -pub enum IndexExprResult { - // The answer is exactly the rows in the allow list minus the rows in the block list - Exact(RowAddrMask), - // The answer is at most the rows in the allow list minus the rows in the block list - // Some of the rows in the allow list may not be in the result and will need to be filtered - // by a recheck. Every row in the block list is definitely not in the result. - AtMost(RowAddrMask), - // The answer is at least the rows in the allow list minus the rows in the block list - // Some of the rows in the block list might be in the result. Every row in the allow list is - // definitely in the result. - AtLeast(RowAddrMask), -} - -impl IndexExprResult { - pub fn row_addr_mask(&self) -> &RowAddrMask { - match self { - Self::Exact(mask) => mask, - Self::AtMost(mask) => mask, - Self::AtLeast(mask) => mask, - } - } - - pub fn discriminant(&self) -> u32 { - match self { - Self::Exact(_) => 0, - Self::AtMost(_) => 1, - Self::AtLeast(_) => 2, - } - } - - pub fn from_parts(mask: RowAddrMask, discriminant: u32) -> Result { - match discriminant { - 0 => Ok(Self::Exact(mask)), - 1 => Ok(Self::AtMost(mask)), - 2 => Ok(Self::AtLeast(mask)), - _ => Err(Error::invalid_input_source( - format!("Invalid IndexExprResult discriminant: {}", discriminant).into(), - )), - } - } - - #[instrument(skip_all)] - pub fn serialize_to_arrow( - &self, - fragments_covered_by_result: &RoaringBitmap, - ) -> Result { - let row_addr_mask = self.row_addr_mask(); - let row_addr_mask_arr = row_addr_mask.into_arrow()?; - let discriminant = self.discriminant(); - let discriminant_arr = - Arc::new(UInt32Array::from(vec![discriminant, discriminant])) as Arc; - let mut fragments_covered_builder = BinaryBuilder::new(); - let fragments_covered_bytes_len = fragments_covered_by_result.serialized_size(); - let mut fragments_covered_bytes = Vec::with_capacity(fragments_covered_bytes_len); - fragments_covered_by_result.serialize_into(&mut fragments_covered_bytes)?; - fragments_covered_builder.append_value(fragments_covered_bytes); - fragments_covered_builder.append_null(); - let fragments_covered_arr = Arc::new(fragments_covered_builder.finish()) as Arc; - Ok(RecordBatch::try_new( - INDEX_EXPR_RESULT_SCHEMA.clone(), - vec![ - Arc::new(row_addr_mask_arr), - Arc::new(discriminant_arr), - Arc::new(fragments_covered_arr), - ], - )?) - } +/// Parse an `IndexExprResult` from its serialized `(mask, discriminant)` +/// representation. Counterpart to [`serialize_index_expr_result`]. +pub fn index_expr_result_from_parts( + mask: RowAddrMask, + discriminant: u32, +) -> Result { + match discriminant { + 0 => Ok(IndexExprResult::Exact(mask)), + 1 => Ok(IndexExprResult::AtMost(mask)), + 2 => Ok(IndexExprResult::AtLeast(mask)), + _ => Err(Error::invalid_input_source( + format!("Invalid IndexExprResult discriminant: {}", discriminant).into(), + )), + } +} + +/// Serialize an `IndexExprResult` plus its applicable-fragments bitmap +/// into the `INDEX_EXPR_RESULT_SCHEMA` record-batch layout used to hand +/// scalar-index results to the read planner. +#[instrument(skip_all)] +pub fn serialize_index_expr_result( + result: &IndexExprResult, + fragments_covered_by_result: &RoaringBitmap, +) -> Result { + let row_addr_mask = result.row_addr_mask(); + let row_addr_mask_arr = row_addr_mask.into_arrow()?; + let discriminant = result.discriminant(); + let discriminant_arr = + Arc::new(UInt32Array::from(vec![discriminant, discriminant])) as Arc; + let mut fragments_covered_builder = BinaryBuilder::new(); + let fragments_covered_bytes_len = fragments_covered_by_result.serialized_size(); + let mut fragments_covered_bytes = Vec::with_capacity(fragments_covered_bytes_len); + fragments_covered_by_result.serialize_into(&mut fragments_covered_bytes)?; + fragments_covered_builder.append_value(fragments_covered_bytes); + fragments_covered_builder.append_null(); + let fragments_covered_arr = Arc::new(fragments_covered_builder.finish()) as Arc; + Ok(RecordBatch::try_new( + INDEX_EXPR_RESULT_SCHEMA.clone(), + vec![ + Arc::new(row_addr_mask_arr), + Arc::new(discriminant_arr), + Arc::new(fragments_covered_arr), + ], + )?) } impl ScalarIndexExpr { @@ -1469,16 +1385,7 @@ impl ScalarIndexExpr { match self { Self::Not(inner) => { let result = inner.evaluate_impl(index_loader, metrics).await?; - // Flip certainty: NOT(AtMost) → AtLeast, NOT(AtLeast) → AtMost - Ok(match result { - NullableIndexExprResult::Exact(mask) => NullableIndexExprResult::Exact(!mask), - NullableIndexExprResult::AtMost(mask) => { - NullableIndexExprResult::AtLeast(!mask) - } - NullableIndexExprResult::AtLeast(mask) => { - NullableIndexExprResult::AtMost(!mask) - } - }) + Ok(!result) } Self::And(lhs, rhs) => { let lhs_result = lhs.evaluate_impl(index_loader, metrics); @@ -2663,7 +2570,7 @@ mod tests { #[tokio::test] async fn test_not_flips_certainty() { - use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap}; + use lance_select::{NullableRowAddrSet, RowAddrTreeMap}; // Test that NOT flips certainty for inexact index results // This tests the implementation in evaluate_impl for Self::Not @@ -2709,7 +2616,7 @@ mod tests { #[tokio::test] async fn test_and_or_preserve_certainty() { - use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap}; + use lance_select::{NullableRowAddrSet, RowAddrTreeMap}; // Test that AND/OR correctly propagate certainty let make_at_most = || { diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 7d89bed85d8..3d82bfe3818 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -23,10 +23,10 @@ use lance_arrow::json::JSON_EXT_NAME; use lance_arrow::{ARROW_EXT_NAME_KEY, iter_str_array}; use lance_core::cache::LanceCache; use lance_core::error::LanceOptionExt; -use lance_core::utils::mask::RowSetOps; use lance_core::utils::tokio::{IO_CORE_RESERVATION, get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; use lance_io::object_store::ObjectStore; +use lance_select::RowSetOps; use object_store::path::Path; use roaring::RoaringBitmap; use smallvec::SmallVec; @@ -3093,7 +3093,7 @@ mod tests { let stream2 = Box::pin(stream2); // Use RowIds filter instead of Fragments — should not affect deleted_fragments - let mut valid_ids = lance_core::utils::mask::RowAddrTreeMap::new(); + let mut valid_ids = lance_select::RowAddrTreeMap::new(); valid_ids.insert(0); let old_data_filter = Some(crate::scalar::OldIndexDataFilter::RowIds(valid_ids)); diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index d9c675b40ec..d08dacd26e7 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -42,10 +42,10 @@ use itertools::Itertools; use lance_arrow::{RecordBatchExt, iter_str_array}; use lance_core::cache::{CacheCodec, CacheKey, LanceCache, WeakLanceCache}; use lance_core::error::{DataFusionResult, LanceOptionExt}; -use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::utils::tracing::{IO_TYPE_LOAD_SCALAR_PART, TRACE_IO_EVENTS}; use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_select::{RowAddrMask, RowAddrTreeMap}; use roaring::RoaringBitmap; use std::sync::LazyLock; use tokio::task::spawn_blocking; diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index b06c75c0021..02dce14dfeb 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -12,7 +12,7 @@ use arrow_array::Array; use itertools::Itertools; use lance_core::Result; use lance_core::utils::address::RowAddress; -use lance_core::utils::mask::RowAddrMask; +use lance_select::RowAddrMask; use crate::metrics::MetricsCollector; diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index 1efb62bd566..d0055e201ac 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -22,8 +22,8 @@ use futures::{StreamExt, TryStream, TryStreamExt, stream::BoxStream}; use lance_arrow::ipc::{read_len_prefixed_bytes_at, write_len_prefixed_bytes}; use lance_core::cache::{CacheCodec, CacheCodecImpl, CacheKey, LanceCache}; use lance_core::error::LanceOptionExt; -use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use lance_core::{Error, ROW_ID, Result}; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use roaring::RoaringBitmap; use tracing::instrument; diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index efc5554d848..9fcd888876d 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -416,9 +416,9 @@ mod tests { use datafusion_common::ScalarValue; use futures::FutureExt; use lance_core::ROW_ID; - use lance_core::utils::mask::{RowAddrTreeMap, RowSetOps}; use lance_core::utils::tempfile::TempDir; use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch}; + use lance_select::{RowAddrTreeMap, RowSetOps}; fn test_store(tempdir: &TempDir) -> Arc { let test_path = tempdir.obj_path(); diff --git a/rust/lance-index/src/scalar/ngram.rs b/rust/lance-index/src/scalar/ngram.rs index df526bebcc6..5f7ce57da4a 100644 --- a/rust/lance-index/src/scalar/ngram.rs +++ b/rust/lance-index/src/scalar/ngram.rs @@ -38,9 +38,9 @@ use lance_core::utils::address::RowAddress; use lance_core::utils::tempfile::TempDir; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::utils::tracing::{IO_TYPE_LOAD_SCALAR_PART, TRACE_IO_EVENTS}; -use lance_core::{Error, utils::mask::RowAddrTreeMap}; -use lance_core::{ROW_ID, Result}; +use lance_core::{Error, ROW_ID, Result}; use lance_io::object_store::ObjectStore; +use lance_select::RowAddrTreeMap; use lance_tokenizer::{ AlphaNumOnlyFilter, AsciiFoldingFilter, LowerCaser, NgramTokenizer, RawTokenizer, TextAnalyzer, }; @@ -1334,13 +1334,10 @@ mod tests { use datafusion_common::DataFusionError; use futures::{TryStreamExt, stream}; use itertools::Itertools; - use lance_core::{ - ROW_ID, - cache::LanceCache, - utils::{mask::RowAddrTreeMap, tempfile::TempDir}, - }; + use lance_core::{ROW_ID, cache::LanceCache, utils::tempfile::TempDir}; use lance_datagen::{BatchCount, ByteCount, RowCount}; use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; use lance_tokenizer::TextAnalyzer; use crate::scalar::{ diff --git a/rust/lance-index/src/scalar/rtree.rs b/rust/lance-index/src/scalar/rtree.rs index adc365e53d2..246c41ba006 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/rust/lance-index/src/scalar/rtree.rs @@ -34,12 +34,12 @@ use geoarrow_schema::{Dimension, RectType}; use lance_arrow::RecordBatchExt; use lance_core::cache::{CacheKey, LanceCache, WeakLanceCache}; use lance_core::utils::address::RowAddress; -use lance_core::utils::mask::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use lance_core::utils::tempfile::TempDir; use lance_core::{Error, ROW_ID, Result}; use lance_datafusion::chunker::chunk_concat_stream; pub use lance_geo::bbox::{BoundingBox, bounding_box, total_bounds}; use lance_io::object_store::ObjectStore; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use roaring::RoaringBitmap; use serde::{Deserialize, Serialize}; use sort::hilbert_sort::HilbertSorter; diff --git a/rust/lance-index/src/scalar/zoned.rs b/rust/lance-index/src/scalar/zoned.rs index b5d90ab7624..7ceed851bae 100644 --- a/rust/lance-index/src/scalar/zoned.rs +++ b/rust/lance-index/src/scalar/zoned.rs @@ -12,9 +12,9 @@ use datafusion::execution::SendableRecordBatchStream; use futures::TryStreamExt; use lance_core::error::Error; use lance_core::utils::address::RowAddress; -use lance_core::utils::mask::RowAddrTreeMap; use lance_core::{ROW_ADDR, Result}; use lance_datafusion::chunker::chunk_concat_stream; +use lance_select::RowAddrTreeMap; // // Example: Suppose we have two fragments, each with 4 rows. diff --git a/rust/lance-index/src/scalar/zonemap.rs b/rust/lance-index/src/scalar/zonemap.rs index 45ded3b0db5..8ba08d3a255 100644 --- a/rust/lance-index/src/scalar/zonemap.rs +++ b/rust/lance-index/src/scalar/zonemap.rs @@ -967,17 +967,16 @@ mod tests { use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_common::ScalarValue; use futures::{StreamExt, TryStreamExt, stream}; - use lance_core::utils::mask::NullableRowAddrSet; use lance_core::utils::tempfile::TempObjDir; use lance_core::{ ROW_ADDR, cache::{LanceCache, WeakLanceCache}, - utils::mask::RowAddrTreeMap, }; use lance_datafusion::datagen::DatafusionDatagenExt; use lance_datagen::ArrayGeneratorExt; use lance_datagen::{BatchCount, RowCount, array}; use lance_io::object_store::ObjectStore; + use lance_select::{NullableRowAddrSet, RowAddrTreeMap}; use crate::scalar::{ SargableQuery, ScalarIndex, SearchResult, diff --git a/rust/lance-index/src/vector/flat/index.rs b/rust/lance-index/src/vector/flat/index.rs index 382601b88e7..9f7d3016206 100644 --- a/rust/lance-index/src/vector/flat/index.rs +++ b/rust/lance-index/src/vector/flat/index.rs @@ -510,7 +510,7 @@ mod tests { use arrow_array::FixedSizeListArray; use async_trait::async_trait; use lance_arrow::FixedSizeListArrayExt; - use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; + use lance_select::{RowAddrMask, RowAddrTreeMap}; use crate::metrics::NoOpMetricsCollector; use crate::prefilter::NoFilter; diff --git a/rust/lance-select/Cargo.toml b/rust/lance-select/Cargo.toml new file mode 100644 index 00000000000..c7deb6f8894 --- /dev/null +++ b/rust/lance-select/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "lance-select" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +readme.workspace = true +keywords.workspace = true +categories.workspace = true +rust-version.workspace = true +description = "Row-selection masks and index-result algebra for Lance" + +[dependencies] +arrow-array = { workspace = true } +arrow-buffer = { workspace = true } +byteorder = { workspace = true } +bytes = { workspace = true } +deepsize = { workspace = true } +itertools = { workspace = true } +lance-core = { workspace = true } +roaring = { workspace = true } + +[dev-dependencies] +criterion = { workspace = true } +proptest = { workspace = true } +rstest = { workspace = true } + +[[bench]] +name = "index_expr_result" +harness = false + +[[bench]] +name = "row_addr_mask" +harness = false + +[lints] +workspace = true diff --git a/rust/lance-select/benches/index_expr_result.rs b/rust/lance-select/benches/index_expr_result.rs new file mode 100644 index 00000000000..6e37e7217ad --- /dev/null +++ b/rust/lance-select/benches/index_expr_result.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Benchmarks for the `NullableIndexExprResult` boolean algebra +//! (`Not` / `BitAnd` / `BitOr`). +//! +//! Captures a baseline of the current 3-variant `Exact`/`AtMost`/`AtLeast` +//! representation before we convert it to a 2-mask `{lower, upper}` form. +//! After the conversion every binary op will do two `NullableRowAddrMask` +//! operations instead of zero or one, so we expect a 2x slowdown on +//! same-variant micro-ops in exchange for richer results (Refined-style +//! interval bounds). +//! +//! Each variant is constructed from a fragment with a single contiguous +//! run of selected rows (the shape produced by zone-map / bloom-filter +//! IsNull searches and by `mask_to_offset_ranges` on contiguous segments). +//! We sweep N across 10K..10M to expose mask-size scaling. + +use std::hint::black_box; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use lance_select::{ + NullableIndexExprResult, NullableRowAddrMask, NullableRowAddrSet, RowAddrTreeMap, +}; + +const ROW_COUNTS: &[u64] = &[10_000, 100_000, 1_000_000, 10_000_000]; + +/// Build a `NullableRowAddrMask::AllowList` covering `0..n` in fragment 0 +/// with no NULL bits set. This is the shape an index search produces +/// when it identifies a contiguous run as a match. +fn allow_run(n: u64) -> NullableRowAddrMask { + let mut tree = RowAddrTreeMap::new(); + tree.insert_range(0..n); + NullableRowAddrMask::AllowList(NullableRowAddrSet::new(tree, RowAddrTreeMap::new())) +} + +/// Build an overlapping mask: covers `n/4..3n/4` (the middle half of an +/// equivalent `allow_run(n)`), so `&` / `|` against `allow_run(n)` +/// produces a non-trivial output that still has roaring-internal runs. +fn allow_middle(n: u64) -> NullableRowAddrMask { + let mut tree = RowAddrTreeMap::new(); + tree.insert_range((n / 4)..(3 * n / 4)); + NullableRowAddrMask::AllowList(NullableRowAddrSet::new(tree, RowAddrTreeMap::new())) +} + +// --- NOT --------------------------------------------------------------- + +fn bench_not(c: &mut Criterion) { + let mut group = c.benchmark_group("not"); + for &n in ROW_COUNTS { + group.throughput(Throughput::Elements(n)); + for label in ["Exact", "AtMost", "AtLeast"] { + let id = BenchmarkId::new(label, n); + let make = move || match label { + "Exact" => NullableIndexExprResult::Exact(allow_run(n)), + "AtMost" => NullableIndexExprResult::AtMost(allow_run(n)), + "AtLeast" => NullableIndexExprResult::AtLeast(allow_run(n)), + _ => unreachable!(), + }; + group.bench_function(id, |b| { + b.iter_batched(make, |r| black_box(!r), criterion::BatchSize::SmallInput); + }); + } + } + group.finish(); +} + +// --- AND / OR ---------------------------------------------------------- + +type PairFn = Box (NullableIndexExprResult, NullableIndexExprResult)>; + +/// Helper that builds `(lhs, rhs)` for every combination we want to bench. +/// Returns `(label, build_fn)` pairs. We test the same-variant cases +/// (Exact/Exact, AtMost/AtMost, AtLeast/AtLeast) plus the cross-variant +/// cases that today drop information (Exact/AtMost, Exact/AtLeast, +/// AtMost/AtLeast). +fn pair_cases(n: u64) -> Vec<(&'static str, PairFn)> { + use NullableIndexExprResult::*; + let lhs = move || allow_run(n); + let rhs = move || allow_middle(n); + vec![ + ( + "Exact_Exact", + Box::new(move || (Exact(lhs()), Exact(rhs()))), + ), + ( + "AtMost_AtMost", + Box::new(move || (AtMost(lhs()), AtMost(rhs()))), + ), + ( + "AtLeast_AtLeast", + Box::new(move || (AtLeast(lhs()), AtLeast(rhs()))), + ), + ( + "Exact_AtMost", + Box::new(move || (Exact(lhs()), AtMost(rhs()))), + ), + ( + "Exact_AtLeast", + Box::new(move || (Exact(lhs()), AtLeast(rhs()))), + ), + ( + "AtMost_AtLeast", + Box::new(move || (AtMost(lhs()), AtLeast(rhs()))), + ), + ] +} + +fn bench_and(c: &mut Criterion) { + let mut group = c.benchmark_group("and"); + for &n in ROW_COUNTS { + group.throughput(Throughput::Elements(n)); + for (label, make) in pair_cases(n) { + let id = BenchmarkId::new(label, n); + group.bench_function(id, |b| { + b.iter_batched( + &make, + |(l, r)| black_box(l & r), + criterion::BatchSize::SmallInput, + ); + }); + } + } + group.finish(); +} + +fn bench_or(c: &mut Criterion) { + let mut group = c.benchmark_group("or"); + for &n in ROW_COUNTS { + group.throughput(Throughput::Elements(n)); + for (label, make) in pair_cases(n) { + let id = BenchmarkId::new(label, n); + group.bench_function(id, |b| { + b.iter_batched( + &make, + |(l, r)| black_box(l | r), + criterion::BatchSize::SmallInput, + ); + }); + } + } + group.finish(); +} + +criterion_group!(benches, bench_not, bench_and, bench_or); +criterion_main!(benches); diff --git a/rust/lance-core/benches/row_addr_mask.rs b/rust/lance-select/benches/row_addr_mask.rs similarity index 98% rename from rust/lance-core/benches/row_addr_mask.rs rename to rust/lance-select/benches/row_addr_mask.rs index c1f09484b69..c5d6c44dd85 100644 --- a/rust/lance-core/benches/row_addr_mask.rs +++ b/rust/lance-select/benches/row_addr_mask.rs @@ -14,12 +14,12 @@ //! ranges fixed at 1. A range-aware representation should make these //! near-constant time; today they are linear in N. //! -//! Run with `cargo bench -p lance-core --bench row_addr_mask`. +//! Run with `cargo bench -p lance-select --bench row_addr_mask`. use std::ops::Range; use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; -use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; +use lance_select::{RowAddrMask, RowAddrTreeMap}; /// Row counts we sweep across. Chosen to cover the realistic range of /// matches a zonemap produces for an `IS NULL`-like predicate on a single diff --git a/rust/lance-select/src/lib.rs b/rust/lance-select/src/lib.rs new file mode 100644 index 00000000000..ed13a94707d --- /dev/null +++ b/rust/lance-select/src/lib.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Row-selection primitives shared across Lance. +//! +//! This crate contains: +//! +//! * [`mask`] — `RowAddrMask` / `NullableRowAddrMask` and their underlying +//! set types. These describe which rows survive a filter and are produced +//! by scalar-index searches, prefilters, and the read planner. +//! * [`result`] — `IndexExprResult` / `NullableIndexExprResult`: the +//! certainty-tagged wrappers around a `RowAddrMask` returned by a scalar- +//! index expression evaluation, plus their boolean algebra +//! (`Not`/`BitAnd`/`BitOr`). +//! +//! These types were extracted from `lance-core` and `lance-index` so that +//! consumers (benchmarks, downstream filtering code) can depend on the +//! mask substrate without pulling in either of those larger crates. + +pub mod mask; +pub mod result; + +pub use mask::{ + NullableRowAddrMask, NullableRowAddrSet, RowAddrMask, RowAddrSelection, RowAddrTreeMap, + RowIdMask, RowIdSet, RowSetOps, bitmap_to_ranges, ranges_to_bitmap, +}; +pub use result::{IndexExprResult, NullableIndexExprResult}; diff --git a/rust/lance-core/src/utils/mask.rs b/rust/lance-select/src/mask.rs similarity index 99% rename from rust/lance-core/src/utils/mask.rs rename to rust/lance-select/src/mask.rs index b904a0e3748..86a475866c9 100644 --- a/rust/lance-core/src/utils/mask.rs +++ b/rust/lance-select/src/mask.rs @@ -13,10 +13,9 @@ use deepsize::DeepSizeOf; use itertools::Itertools; use roaring::{MultiOps, RoaringBitmap, RoaringTreemap}; -use crate::cache::CacheCodecImpl; -use crate::{Error, Result}; - -use super::address::RowAddress; +use lance_core::cache::CacheCodecImpl; +use lance_core::utils::address::RowAddress; +use lance_core::{Error, Result}; mod nullable; @@ -460,7 +459,7 @@ impl RowAddrTreeMap { /// Returns true if the value was not already in the set. /// /// ```rust - /// use lance_core::utils::mask::{RowAddrTreeMap, RowSetOps}; + /// use lance_select::{RowAddrTreeMap, RowSetOps}; /// /// let mut set = RowAddrTreeMap::new(); /// assert_eq!(set.insert(10), true); diff --git a/rust/lance-core/src/utils/mask/nullable.rs b/rust/lance-select/src/mask/nullable.rs similarity index 100% rename from rust/lance-core/src/utils/mask/nullable.rs rename to rust/lance-select/src/mask/nullable.rs diff --git a/rust/lance-select/src/result.rs b/rust/lance-select/src/result.rs new file mode 100644 index 00000000000..759c139a0fd --- /dev/null +++ b/rust/lance-select/src/result.rs @@ -0,0 +1,151 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Certainty-tagged wrappers around a row-address mask returned by a +//! scalar-index expression evaluation. +//! +//! These types model the three possible degrees of knowledge an index +//! search can return: +//! +//! * [`Exact`] — the mask is the precise answer; no recheck needed. +//! * [`AtMost`] — the mask is a *superset* of the true answer; the rows +//! inside the mask must be rechecked against the predicate. +//! * [`AtLeast`] — the mask is a *subset* of the true answer; the rows +//! outside the mask must be rechecked against the predicate. +//! +//! The boolean algebra (`Not`/`BitAnd`/`BitOr`) is implemented on both +//! [`NullableIndexExprResult`] (the form during evaluation, carrying SQL +//! three-valued logic via [`NullableRowAddrMask`]) and +//! [`IndexExprResult`] (the form consumed by the read planner, after +//! `drop_nulls` collapses NULL rows into FALSE). +//! +//! [`Exact`]: IndexExprResult::Exact +//! [`AtMost`]: IndexExprResult::AtMost +//! [`AtLeast`]: IndexExprResult::AtLeast + +use crate::mask::{NullableRowAddrMask, RowAddrMask}; + +/// Result of an index search before NULL rows are dropped. Carries +/// three-valued-logic information via [`NullableRowAddrMask`]. +#[derive(Debug)] +pub enum NullableIndexExprResult { + Exact(NullableRowAddrMask), + AtMost(NullableRowAddrMask), + AtLeast(NullableRowAddrMask), +} + +impl std::ops::Not for NullableIndexExprResult { + type Output = Self; + + fn not(self) -> Self { + // Flip certainty: NOT(AtMost) → AtLeast, NOT(AtLeast) → AtMost. + // NULL info is preserved by `NullableRowAddrMask::not` (it flips + // AllowList ↔ BlockList without touching the `nulls` field), which + // is the 3VL-correct negation: TRUE↔FALSE swap, NULL stays NULL. + match self { + Self::Exact(mask) => Self::Exact(!mask), + Self::AtMost(mask) => Self::AtLeast(!mask), + Self::AtLeast(mask) => Self::AtMost(!mask), + } + } +} + +impl std::ops::BitAnd for NullableIndexExprResult { + type Output = Self; + + fn bitand(self, rhs: Self) -> Self { + match (self, rhs) { + (Self::Exact(lhs), Self::Exact(rhs)) => Self::Exact(lhs & rhs), + (Self::Exact(lhs), Self::AtMost(rhs)) | (Self::AtMost(lhs), Self::Exact(rhs)) => { + Self::AtMost(lhs & rhs) + } + (Self::Exact(exact), Self::AtLeast(_)) | (Self::AtLeast(_), Self::Exact(exact)) => { + // We could do better here, elements in both lhs and rhs are known + // to be true and don't require a recheck. We only need to recheck + // elements in lhs that are not in rhs + Self::AtMost(exact) + } + (Self::AtMost(lhs), Self::AtMost(rhs)) => Self::AtMost(lhs & rhs), + (Self::AtLeast(lhs), Self::AtLeast(rhs)) => Self::AtLeast(lhs & rhs), + (Self::AtMost(most), Self::AtLeast(_)) | (Self::AtLeast(_), Self::AtMost(most)) => { + Self::AtMost(most) + } + } + } +} + +impl std::ops::BitOr for NullableIndexExprResult { + type Output = Self; + + fn bitor(self, rhs: Self) -> Self { + match (self, rhs) { + (Self::Exact(lhs), Self::Exact(rhs)) => Self::Exact(lhs | rhs), + (Self::Exact(lhs), Self::AtMost(rhs)) | (Self::AtMost(rhs), Self::Exact(lhs)) => { + // We could do better here, elements in lhs are known to be true + // and don't require a recheck. We only need to recheck elements + // in rhs that are not in lhs + Self::AtMost(lhs | rhs) + } + (Self::Exact(lhs), Self::AtLeast(rhs)) | (Self::AtLeast(rhs), Self::Exact(lhs)) => { + Self::AtLeast(lhs | rhs) + } + (Self::AtMost(lhs), Self::AtMost(rhs)) => Self::AtMost(lhs | rhs), + (Self::AtLeast(lhs), Self::AtLeast(rhs)) => Self::AtLeast(lhs | rhs), + (Self::AtMost(_), Self::AtLeast(least)) | (Self::AtLeast(least), Self::AtMost(_)) => { + Self::AtLeast(least) + } + } + } +} + +impl NullableIndexExprResult { + /// Project NULL rows out of the result. + /// + /// Under a `WHERE` clause, NULL is treated as FALSE — so `drop_nulls` + /// removes them from `AllowList`s (NULL rows are not selected) and + /// folds them into `BlockList`s (NULL rows are still blocked). + pub fn drop_nulls(self) -> IndexExprResult { + match self { + Self::Exact(mask) => IndexExprResult::Exact(mask.drop_nulls()), + Self::AtMost(mask) => IndexExprResult::AtMost(mask.drop_nulls()), + Self::AtLeast(mask) => IndexExprResult::AtLeast(mask.drop_nulls()), + } + } +} + +/// Result of an index search after NULL rows have been dropped. This is +/// what the read planner consumes. +#[derive(Debug)] +pub enum IndexExprResult { + /// The answer is exactly the rows in the allow list minus the rows + /// in the block list. + Exact(RowAddrMask), + /// The answer is at most the rows in the allow list minus the rows + /// in the block list. Some of the rows in the allow list may not be + /// in the result and will need to be filtered by a recheck. Every + /// row in the block list is definitely not in the result. + AtMost(RowAddrMask), + /// The answer is at least the rows in the allow list minus the rows + /// in the block list. Some of the rows in the block list might be in + /// the result. Every row in the allow list is definitely in the + /// result. + AtLeast(RowAddrMask), +} + +impl IndexExprResult { + pub fn row_addr_mask(&self) -> &RowAddrMask { + match self { + Self::Exact(mask) => mask, + Self::AtMost(mask) => mask, + Self::AtLeast(mask) => mask, + } + } + + pub fn discriminant(&self) -> u32 { + match self { + Self::Exact(_) => 0, + Self::AtMost(_) => 1, + Self::AtLeast(_) => 2, + } + } +} diff --git a/rust/lance-table/Cargo.toml b/rust/lance-table/Cargo.toml index c9cfc828f70..5848480212c 100644 --- a/rust/lance-table/Cargo.toml +++ b/rust/lance-table/Cargo.toml @@ -15,6 +15,7 @@ rust-version.workspace = true lance-arrow.workspace = true lance-core.workspace = true lance-file.workspace = true +lance-select.workspace = true lance-io.workspace = true arrow.workspace = true arrow-array.workspace = true diff --git a/rust/lance-table/src/rowids.rs b/rust/lance-table/src/rowids.rs index 69a2d635e76..0b56be84f56 100644 --- a/rust/lance-table/src/rowids.rs +++ b/rust/lance-table/src/rowids.rs @@ -26,15 +26,12 @@ use deepsize::DeepSizeOf; // These are the public API. pub use index::FragmentRowIdIndex; pub use index::RowIdIndex; -use lance_core::{ - Error, Result, - utils::mask::{RowAddrMask, RowAddrTreeMap}, -}; +use lance_core::{Error, Result}; use lance_io::ReadBatchParams; +use lance_select::{RowAddrMask, RowAddrTreeMap, RowSetOps}; pub use serde::{read_row_ids, write_row_ids}; use crate::utils::LanceIteratorExtension; -use lance_core::utils::mask::RowSetOps; use segment::U64Segment; use tracing::instrument; diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index ca2bfdeaf91..fa4de229cc9 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -29,6 +29,7 @@ lance-io = { workspace = true } lance-linalg = { workspace = true } lance-index = { workspace = true } lance-namespace = { workspace = true } +lance-select = { workspace = true } lance-tokenizer = { workspace = true } lance-table = { workspace = true } arrow-arith = { workspace = true } diff --git a/rust/lance/benches/scalar_index.rs b/rust/lance/benches/scalar_index.rs index 918b08a78a0..49b94c0207e 100644 --- a/rust/lance/benches/scalar_index.rs +++ b/rust/lance/benches/scalar_index.rs @@ -12,7 +12,6 @@ use datafusion::{physical_plan::SendableRecordBatchStream, scalar::ScalarValue}; use futures::{FutureExt, TryStreamExt}; use lance::{Dataset, io::ObjectStore}; use lance_core::cache::LanceCache; -use lance_core::utils::mask::RowSetOps; use lance_core::utils::tempfile::TempStrDir; use lance_datafusion::utils::reader_to_stream; use lance_datagen::{BatchCount, RowCount, array, gen_batch}; @@ -23,6 +22,7 @@ use lance_index::scalar::{ registry::ScalarIndexPlugin, }; use lance_index::{metrics::NoOpMetricsCollector, scalar::btree::BTreeIndexPlugin}; +use lance_select::RowSetOps; #[cfg(target_os = "linux")] use pprof::criterion::{Output, PProfProfiler}; diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index f2c16b87d0b..547707affcb 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -53,7 +53,6 @@ use lance_core::datatypes::{ }; use lance_core::error::LanceOptionExt; use lance_core::utils::address::RowAddress; -use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ROW_ADDR, ROW_ID, ROW_OFFSET}; use lance_datafusion::aggregate::Aggregate; @@ -66,7 +65,9 @@ use lance_file::reader::FileReaderOptions; use lance_index::IndexCriteria; use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::expression::ScalarIndexExpr; -use lance_index::scalar::expression::{INDEX_EXPR_RESULT_SCHEMA, IndexExprResult, PlannerIndexExt}; +use lance_index::scalar::expression::{ + INDEX_EXPR_RESULT_SCHEMA, IndexExprResult, PlannerIndexExt, serialize_index_expr_result, +}; use lance_index::scalar::inverted::query::{ FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, PhraseQuery, fill_fts_query_column, }; @@ -75,6 +76,7 @@ use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, DIST_COL, Query}; use lance_index::{metrics::NoOpMetricsCollector, scalar::inverted::FTS_SCHEMA}; use lance_io::stream::RecordBatchStream; use lance_linalg::distance::MetricType; +use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::{Fragment, IndexMetadata}; use roaring::RoaringBitmap; use tracing::{Span, info_span, instrument}; @@ -2834,7 +2836,7 @@ impl Scanner { let row_addr_mask = RowAddrMask::from_allowed(row_addrs); let index_result = IndexExprResult::Exact(row_addr_mask); let fragments_covered = self.dataset.fragment_bitmap.as_ref().clone(); - let batch = index_result.serialize_to_arrow(&fragments_covered)?; + let batch = serialize_index_expr_result(&index_result, &fragments_covered)?; let stream = futures::stream::once(async move { Ok(batch) }); let stream = Box::pin(RecordBatchStreamAdapter::new( INDEX_EXPR_RESULT_SCHEMA.clone(), diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 45b78b48bcb..95f4484c099 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -4,9 +4,9 @@ use std::collections::HashMap; use std::sync::Arc; -use lance_core::utils::mask::RowAddrTreeMap; use lance_file::version::LanceFileVersion; use lance_io::object_store::{ObjectStore, ObjectStoreParams}; +use lance_select::RowAddrTreeMap; use lance_table::{ format::{DataStorageFormat, is_detached_version}, io::commit::{CommitConfig, CommitHandler, ManifestNamingScheme}, diff --git a/rust/lance/src/dataset/write/delete.rs b/rust/lance/src/dataset/write/delete.rs index 2c43b49ac73..a063d28ad7b 100644 --- a/rust/lance/src/dataset/write/delete.rs +++ b/rust/lance/src/dataset/write/delete.rs @@ -11,8 +11,8 @@ use crate::{ use datafusion::logical_expr::Expr; use datafusion::scalar::ScalarValue; use futures::{StreamExt, TryStreamExt}; -use lance_core::utils::mask::RowAddrTreeMap; use lance_core::{Error, ROW_ID, Result}; +use lance_select::RowAddrTreeMap; use lance_table::format::Fragment; use roaring::RoaringTreemap; use std::collections::BTreeMap; diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index c88edef262e..e9432fbad27 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -101,7 +101,7 @@ use lance_core::{ Error, ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD, Result, datatypes::{OnMissing, OnTypeMismatch, SchemaCompareOptions}, error::{InvalidInputSnafu, box_error}, - utils::{futures::Capacity, mask::RowAddrTreeMap, tokio::get_num_compute_intensive_cpus}, + utils::{futures::Capacity, tokio::get_num_compute_intensive_cpus}, }; use lance_datafusion::{ chunker::chunk_stream, @@ -115,6 +115,7 @@ use lance_datafusion::{ use lance_file::version::LanceFileVersion; use lance_index::IndexCriteria; use lance_index::mem_wal::MergedGeneration; +use lance_select::RowAddrTreeMap; use lance_table::format::{Fragment, IndexMetadata, RowIdMeta}; use log::info; use roaring::RoaringTreemap; diff --git a/rust/lance/src/dataset/write/update.rs b/rust/lance/src/dataset/write/update.rs index 67a4f5bbd6d..fba7b2342da 100644 --- a/rust/lance/src/dataset/write/update.rs +++ b/rust/lance/src/dataset/write/update.rs @@ -25,10 +25,10 @@ use datafusion::scalar::ScalarValue; use futures::StreamExt; use lance_arrow::RecordBatchExt; use lance_core::error::{InvalidInputSnafu, box_error}; -use lance_core::utils::mask::RowAddrTreeMap; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ROW_ADDR_FIELD, ROW_ID_FIELD, ROW_OFFSET_FIELD}; use lance_datafusion::expr::safe_coerce_scalar; +use lance_select::RowAddrTreeMap; use lance_table::format::{Fragment, RowIdMeta}; use roaring::RoaringTreemap; use snafu::ResultExt; diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 66c0dc84337..16b17752ef4 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -4,10 +4,7 @@ use std::sync::Arc; use futures::{FutureExt, TryStreamExt}; -use lance_core::{ - Error, Result, - utils::mask::{RowAddrTreeMap, RowSetOps}, -}; +use lance_core::{Error, Result}; use lance_index::{ INDEX_FILE_NAME, IndexType, metrics::NoOpMetricsCollector, @@ -17,6 +14,7 @@ use lance_index::{ CreatedIndex, OldIndexDataFilter, inverted::InvertedIndex, lance_format::LanceIndexStore, }, }; +use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::{Fragment, IndexMetadata, list_index_files_with_sizes}; use roaring::RoaringBitmap; use uuid::Uuid; diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 64df23e0fc9..842252f9a45 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -1373,8 +1373,8 @@ mod tests { #[tokio::test] async fn test_distributed_build_bitmap() { use datafusion::common::ScalarValue; - use lance_core::utils::mask::RowSetOps; use lance_index::scalar::{SargableQuery, SearchResult, bitmap::BITMAP_LOOKUP_NAME}; + use lance_select::RowSetOps; let tmpdir = TempStrDir::default(); let dataset_uri = format!("file://{}", tmpdir.as_str()); diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 95184e63f15..ebb4c600e65 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -19,8 +19,8 @@ use futures::TryStreamExt; use futures::future::BoxFuture; use futures::stream; use lance_core::utils::deletion::DeletionVector; -use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; use lance_core::utils::tokio::spawn_cpu; +use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::Fragment; use lance_table::format::IndexMetadata; use lance_table::rowids::RowIdSequence; @@ -379,7 +379,7 @@ impl PreFilter for DatasetPreFilter { #[cfg(test)] mod test { - use lance_core::utils::mask::RowSetOps; + use lance_select::RowSetOps; use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; use crate::dataset::WriteParams; diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 2cc9dd5cbe4..162a36a0c97 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -9,11 +9,11 @@ use std::sync::Arc; use async_trait::async_trait; use deepsize::{Context, DeepSizeOf}; use futures::future::try_join_all; -use lance_core::utils::mask::NullableRowAddrSet; use lance_core::{Error, Result}; use lance_index::metrics::MetricsCollector; use lance_index::scalar::{AnyQuery, CreatedIndex, ScalarIndex, SearchResult, UpdateCriteria}; use lance_index::{Index, IndexType}; +use lance_select::NullableRowAddrSet; use lance_table::format::IndexMetadata; use roaring::RoaringBitmap; use serde_json::json; diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index 05779fb96fa..1f1a89d80c8 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -645,10 +645,10 @@ mod tests { use lance_linalg::kernels::normalize_fsl; use crate::index::vector::ivf::build_ivf_model; - use lance_core::utils::mask::RowAddrMask; use lance_index::metrics::NoOpMetricsCollector; use lance_index::vector::DEFAULT_QUERY_PARALLELISM; use lance_index::vector::ivf::IvfBuildParams; + use lance_select::RowAddrMask; use lance_testing::datagen::{ generate_random_array_with_range, generate_random_array_with_seed, }; diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 67e57ed3320..d0a2934552c 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -26,10 +26,10 @@ use std::time::Instant; use conflict_resolver::TransactionRebase; use lance_core::utils::backoff::{Backoff, SlotBackoff}; -use lance_core::utils::mask::RowAddrTreeMap; use lance_file::version::LanceFileVersion; use lance_index::metrics::NoOpMetricsCollector; use lance_io::utils::CachedFileSize; +use lance_select::RowAddrTreeMap; use lance_table::format::{ DETACHED_VERSION_MASK, DataStorageFormat, DeletionFile, Fragment, IndexMetadata, Manifest, WriterVersion, is_detached_version, list_index_files_with_sizes, pb, diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 2e0a04be312..b242cd5b3dd 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -10,13 +10,10 @@ use crate::{ dataset::transaction::{Operation, Transaction}, }; use futures::{StreamExt, TryStreamExt}; -use lance_core::utils::mask::RowSetOps; -use lance_core::{ - Error, Result, - utils::{deletion::DeletionVector, mask::RowAddrTreeMap}, -}; +use lance_core::{Error, Result, utils::deletion::DeletionVector}; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{MEM_WAL_INDEX_NAME, MergedGeneration}; +use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::IndexMetadata; use lance_table::{format::Fragment, io::deletion::write_deletion_file}; use std::{ diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 73ddd52a4aa..5a06408090b 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -32,9 +32,6 @@ use lance_arrow::RecordBatchExt; use lance_core::datatypes::OnMissing; use lance_core::utils::deletion::DeletionVector; use lance_core::utils::futures::FinallyStreamExt; -use lance_core::utils::mask::{ - RowAddrMask, RowAddrSelection, RowAddrTreeMap, bitmap_to_ranges, ranges_to_bitmap, -}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{Error, Result, datatypes::Projection}; use lance_datafusion::planner::Planner; @@ -43,8 +40,11 @@ use lance_datafusion::utils::{ ROWS_SCANNED_METRIC, TASK_WAIT_TIME_METRIC, }; use lance_file::reader::FileReaderOptions; -use lance_index::scalar::expression::{FilterPlan, IndexExprResult}; +use lance_index::scalar::expression::{FilterPlan, IndexExprResult, index_expr_result_from_parts}; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; +use lance_select::{ + RowAddrMask, RowAddrSelection, RowAddrTreeMap, bitmap_to_ranges, ranges_to_bitmap, +}; use lance_table::format::Fragment; use lance_table::rowids::RowIdSequence; use lance_table::utils::stream::ReadBatchFut; @@ -101,7 +101,7 @@ impl EvaluatedIndex { } let row_addr_mask = RowAddrMask::from_arrow(batch.column(0).as_binary())?; let match_type = batch.column(1).as_primitive::().values()[0]; - let index_result = IndexExprResult::from_parts(row_addr_mask, match_type)?; + let index_result = index_expr_result_from_parts(row_addr_mask, match_type)?; let applicable_fragments = batch.column(2).as_binary::(); let applicable_fragments = RoaringBitmap::deserialize_from(applicable_fragments.value(0))?; diff --git a/rust/lance/src/io/exec/filtered_read_proto.rs b/rust/lance/src/io/exec/filtered_read_proto.rs index c373a691f04..4eb329506aa 100644 --- a/rust/lance/src/io/exec/filtered_read_proto.rs +++ b/rust/lance/src/io/exec/filtered_read_proto.rs @@ -20,10 +20,10 @@ use datafusion::execution::SessionState; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; use lance_core::datatypes::{BlobHandling, Projection}; -use lance_core::utils::mask::RowAddrTreeMap; use lance_core::{Error, Result}; use lance_datafusion::pb; use lance_datafusion::substrait::{encode_substrait, parse_substrait, prune_schema_for_substrait}; +use lance_select::RowAddrTreeMap; use lance_table::format::Fragment; use crate::Dataset; @@ -487,8 +487,8 @@ mod tests { use arrow_schema::{DataType, Field}; use datafusion::prelude::SessionContext; use lance_core::datatypes::OnMissing; - use lance_core::utils::mask::RowAddrTreeMap; use lance_datagen::{array, gen_batch}; + use lance_select::RowAddrTreeMap; use roaring::RoaringBitmap; use std::collections::HashMap; use std::collections::HashSet; diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index b15ed55d004..f6fd1d0a197 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -29,14 +29,7 @@ use datafusion::{ }; use datafusion_physical_expr::EquivalenceProperties; use futures::{StreamExt, TryFutureExt, TryStreamExt, stream::BoxStream}; -use lance_core::utils::mask::RowSetOps; -use lance_core::{ - Error, ROW_ID_FIELD, Result, - utils::{ - address::RowAddress, - mask::{RowAddrMask, RowAddrTreeMap}, - }, -}; +use lance_core::{Error, ROW_ID_FIELD, Result, utils::address::RowAddress}; use lance_datafusion::{ chunker::break_stream, utils::{ @@ -49,10 +42,11 @@ use lance_index::{ SargableQuery, ScalarIndex, expression::{ INDEX_EXPR_RESULT_SCHEMA, IndexExprResult, ScalarIndexExpr, ScalarIndexLoader, - ScalarIndexSearch, + ScalarIndexSearch, serialize_index_expr_result, }, }, }; +use lance_select::{RowAddrMask, RowAddrTreeMap, RowSetOps}; use lance_table::format::Fragment; use roaring::RoaringBitmap; use tracing::{debug_span, instrument}; @@ -159,7 +153,7 @@ impl ScalarIndexExec { { let ser_time = plan_metrics.new_time(SCALAR_INDEX_SER_TIME_METRIC, 0); let _timer = ser_time.timer(); - query_result.serialize_to_arrow(&fragments_covered_by_result) + serialize_index_expr_result(&query_result, &fragments_covered_by_result) } } } diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index 1c9c1d9d476..004b1a95b5f 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -28,9 +28,9 @@ use futures::stream::FuturesUnordered; use futures::{Stream, StreamExt, TryStreamExt}; use lance_core::error::{CloneableResult, Error}; use lance_core::utils::futures::{Capacity, SharedStreamExt}; -use lance_core::utils::mask::{RowAddrMask, RowAddrTreeMap}; use lance_core::{ROW_ID, Result}; use lance_index::prefilter::FilterLoader; +use lance_select::{RowAddrMask, RowAddrTreeMap}; use std::future::Future; use crate::Dataset; diff --git a/rust/lance/src/session/caches.rs b/rust/lance/src/session/caches.rs index f8cdd520dad..eab758418f7 100644 --- a/rust/lance/src/session/caches.rs +++ b/rust/lance/src/session/caches.rs @@ -15,8 +15,9 @@ use std::{borrow::Cow, ops::Deref}; use deepsize::{Context, DeepSizeOf}; use lance_core::{ cache::{CacheKey, LanceCache}, - utils::{deletion::DeletionVector, mask::RowAddrMask}, + utils::deletion::DeletionVector, }; +use lance_select::RowAddrMask; use lance_table::{ format::{DeletionFile, Manifest}, rowids::{RowIdIndex, RowIdSequence},