diff --git a/README.md b/README.md index f4dccb2..5b4e167 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,47 @@ Based on the [liblance RFC](https://github.com/lance-format/lance/discussions/60 | [x] | Dataset metadata | `lance_dataset_version()`, `lance_dataset_count_rows()`, `lance_dataset_latest_version()` | | [x] | Filter pushdown | `lance_scanner_set_substrait_filter()` accepts a serialized Substrait `ExtendedExpression`; `lance_scanner_additional_sql_filter()` adds SQL predicates with AND before scanning starts | +## Segment-scoped array label filters + +Ordinary scans can use a `LabelList` index for array membership through the +existing SQL and Substrait filter interfaces. For example, on a `List` +column named `labels`: + +```sql +array_contains(labels, 'red') +array_contains(labels, 'red') AND array_contains(labels, 'blue') +array_contains(labels, 'red') OR array_contains(labels, 'blue') +array_has_all(labels, ['red', 'blue']) +array_has_any(labels, ['red', 'blue']) +``` + +Configure `lance_scanner_set_fragment_ids` and +`lance_scanner_set_scalar_index_segment` with the selected LabelList segment. +Pass a SQL filter to `lance_scanner_new`, or attach a serialized Substrait +`ExtendedExpression` with `lance_scanner_set_substrait_filter`. The Substrait +schema must describe the list field and its element type; replacing that field +with an unsupported-type placeholder cannot express a label predicate. Use +Lance/DataFusion's `array_has` (the canonical name of `array_contains`), +`array_has_all`, or `array_has_any` functions with correctly typed arguments. +Substrait takes precedence over the primary SQL filter; use +`lance_scanner_additional_sql_filter` when an additional condition must be ANDed +with it. + +AND/OR membership expressions can reuse the same selected LabelList segment. +Other columns remain residual filters, evaluated before LIMIT/OFFSET. This API +selects one physical segment; it does not intersect indices on different columns. +Incomplete segment coverage falls back to scanning the entire explicit fragment +scope. Check `scalar_segments_searched` and `scalar_segment_fallbacks` in the +statistics callback to distinguish index acceleration from filter execution. + +Bindings preserve Lance's function semantics, not the semantics of similarly +named functions in another SQL engine. In particular, a NULL search value in +`array_contains` does not match NULL array elements. `array_has_all` tests set +containment, not an ordered contiguous subsequence. The pinned Lance version +also treats an empty all-label query as true even for NULL lists. Integrators +should initially push only non-NULL constant labels with matching element types +and retain conditions whose semantics have not been verified in the calling engine. + ## Distance-bounded vector search After configuring a single-vector nearest-neighbor query, use diff --git a/include/lance/lance.h b/include/lance/lance.h index d29fc0b..ea00436 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -2256,13 +2256,21 @@ int32_t lance_scanner_set_index_segments( * fragment domains and separately include any unindexed data they wish to read. * The segment metadata must identify one key field present in the schema. * - * BTree/Bitmap/LabelList searches use a necessary AND-conjunct of the - * full scanner filter on the selected logical index and require an Exact result. + * BTree/Bitmap/LabelList searches evaluate a candidate expression on the selected + * logical index, including AND, OR, IN and NULL-aware NOT, and require an Exact + * result. LabelList supports Lance array_has/array_contains, array_has_all and + * array_has_any through SQL or Substrait filters with matching list element types. + * These functions retain Lance semantics, including for NULL search values. + * An AND may retain only its supported necessary conditions. OR needs + * candidates for both branches; NOT requires its complete indexed subtree. + * Expressions exceeding 128 nodes or depth 32 use the scoped fallback. + * Filters containing IS [NOT] TRUE/FALSE also use that fallback until the Lance + * planner dependency preserves their NULL semantics under negation. * use_scalar_index=false skips segment search and uses the scoped fallback; * snapshot UUID and fragment validation still applies. * AtMost/AtLeast results fall back to a full filtered scan of fragment_ids. * All predicates are reapplied during candidate reads; other scalar indices - * are disabled. Legacy storage, OR/NOT-only filters, + * are disabled. Legacy storage, expressions without safe scoped candidates, * overlays, fragment reuse, unsupported index types / result domains * and missing coverage use the same domain without an index. No filter also * falls back. LIMIT/OFFSET apply after the complete scanner filter, never to the diff --git a/src/scalar_segment.rs b/src/scalar_segment.rs index 7d60d27..6515fbf 100644 --- a/src/scalar_segment.rs +++ b/src/scalar_segment.rs @@ -8,19 +8,19 @@ use std::collections::HashSet; use std::sync::Arc; use std::time::Instant; +use datafusion::common::tree_node::TreeNode; +use datafusion::logical_expr::Expr; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use lance::Dataset; -use lance::dataset::scanner::{ - ExecutionStatsCallback, ExecutionSummaryCounts, RowAddrMask, Scanner, -}; +use lance::dataset::scanner::{ExecutionStatsCallback, ExecutionSummaryCounts, Scanner}; use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; use lance::io::exec::utils::IndexMetrics; use lance_core::{Error, Result}; use lance_datafusion::planner::Planner; use lance_datafusion::utils::MetricsExt; use lance_index::IndexType; -use lance_index::scalar::SearchResult; -use lance_index::scalar::expression::{PlannerIndexExt, ScalarIndexExpr, ScalarIndexSearch}; +use lance_index::scalar::expression::{PlannerIndexExt, ScalarIndexExpr, ScalarIndexLoader}; +use lance_index::scalar::{MetricsCollector, ScalarIndex}; use uuid::Uuid; pub(crate) struct PreparedScalarSegment { @@ -35,15 +35,83 @@ fn invalid(message: impl Into) -> Error { Error::invalid_input_source(message.into().into()) } -// Only descend through AND: a leaf below OR or NOT need not contain all matches -// of the full expression. The original expression is always reapplied by reader. -fn driver<'a>(expr: &'a ScalarIndexExpr, index_name: &str) -> Option<&'a ScalarIndexSearch> { +// Bound both recursive planning and the concurrent searches in Lance's evaluator. +const MAX_EXPRESSION_NODES: usize = 128; +const MAX_EXPRESSION_DEPTH: usize = 32; + +fn scoped_expression( + expr: &ScalarIndexExpr, + index_name: &str, + column: &str, + require_complete: bool, + depth: usize, + remaining: &mut usize, +) -> std::result::Result, &'static str> { + if depth > MAX_EXPRESSION_DEPTH || *remaining == 0 { + return Err("expression_budget"); + } + *remaining -= 1; + let recurse = |child, complete, remaining: &mut usize| { + scoped_expression(child, index_name, column, complete, depth + 1, remaining) + }; match expr { - ScalarIndexExpr::Query(search) if search.index_name == index_name => Some(search), + ScalarIndexExpr::Query(search) => { + if search.index_name != index_name { + Ok(None) + } else if search.column != column { + Err("field_path") + } else { + Ok(Some(expr.clone())) + } + } ScalarIndexExpr::And(lhs, rhs) => { - driver(lhs, index_name).or_else(|| driver(rhs, index_name)) + let lhs = recurse(lhs, require_complete, remaining)?; + let rhs = recurse(rhs, require_complete, remaining)?; + Ok(match (lhs, rhs) { + (Some(lhs), Some(rhs)) => Some(ScalarIndexExpr::And(Box::new(lhs), Box::new(rhs))), + (lhs, rhs) if !require_complete => lhs.or(rhs), + _ => None, + }) + } + ScalarIndexExpr::Or(lhs, rhs) => { + let lhs = recurse(lhs, require_complete, remaining)?; + let rhs = recurse(rhs, require_complete, remaining)?; + // Both branches must contribute a superset of their matches. Dropping + // an unavailable OR branch would silently exclude valid rows. + Ok(lhs + .zip(rhs) + .map(|(lhs, rhs)| ScalarIndexExpr::Or(Box::new(lhs), Box::new(rhs)))) + } + ScalarIndexExpr::Not(inner) => { + // Negating a pruned AND would turn a safe superset into an unsafe + // subset. Preserve the complete subtree and let Lance track NULLs. + Ok(recurse(inner, true, remaining)?.map(|inner| ScalarIndexExpr::Not(Box::new(inner)))) + } + } +} + +struct SegmentIndexLoader<'a> { + index: Arc, + index_name: &'a str, + column: &'a str, +} + +#[async_trait::async_trait] +impl ScalarIndexLoader for SegmentIndexLoader<'_> { + async fn load_index( + &self, + column: &str, + index_name: &str, + _metrics: &dyn MetricsCollector, + ) -> Result> { + if column != self.column || index_name != self.index_name { + return Err(invalid( + "scalar expression references an index outside the selected segment", + )); } - _ => None, + // Reuse the one physical segment; a dataset loader would search every + // segment of the logical index and violate the caller's restriction. + Ok(self.index.clone()) } } @@ -193,6 +261,18 @@ impl PreparedScalarSegment { let Some(filter) = reader.get_expr_filter()? else { return Ok(Some("no_filter")); }; + // The pinned Lance planner lowers IS TRUE/FALSE to nullable equality. + // Under negation this loses matching NULL rows before the full recheck. + // Inspect the original expression before planning erases that distinction; + // remove this fallback only after adopting lance-format/lance#9568. + if filter.exists(|expr| { + Ok(matches!( + expr, + Expr::IsTrue(_) | Expr::IsFalse(_) | Expr::IsNotTrue(_) | Expr::IsNotFalse(_) + )) + })? { + return Ok(Some("boolean_truth_test")); + } let stored_schema: arrow_schema::Schema = self.dataset.schema().into(); // get_expr_filter validates against the scanner's full filterable // schema, including metadata columns absent from the stored schema. @@ -208,19 +288,25 @@ impl PreparedScalarSegment { let planner = Planner::new(Arc::new(stored_schema)); let index_info = self.dataset.scalar_index_info().await?; let filter_plan = planner.create_filter_plan(filter, &index_info, true)?; - let Some(search) = filter_plan - .index_query - .as_ref() - .and_then(|expr| driver(expr, &index_meta.name)) - else { + let Some(expr) = filter_plan.index_query.as_ref() else { return Ok(Some("no_driver")); }; - if search.column != field.name { - return Ok(Some("field_path")); - } + let mut remaining = MAX_EXPRESSION_NODES; + let expr = match scoped_expression( + expr, + &index_meta.name, + &field.name, + false, + 0, + &mut remaining, + ) { + Ok(Some(expr)) => expr, + Ok(None) => return Ok(Some("no_driver")), + Err(reason) => return Ok(Some(reason)), + }; let index = self .dataset - .open_scalar_index(&search.column, &self.segment_uuid, metrics) + .open_scalar_index(&field.name, &self.segment_uuid, metrics) .await?; // These implementations can return exact candidates. Keep the runtime // Exact check below: a type alone is not a guarantee for every query. @@ -235,7 +321,12 @@ impl PreparedScalarSegment { return Ok(Some("row_id_domain")); } let started = Instant::now(); - let result = index.search(search.query.as_ref(), metrics).await?; + let loader = SegmentIndexLoader { + index, + index_name: &index_meta.name, + column: &field.name, + }; + let result = expr.evaluate(&loader, metrics).await?; stats.all_times.insert( "scalar_segment_search_time".into(), started.elapsed().as_nanos().min(usize::MAX as u128) as usize, @@ -243,18 +334,75 @@ impl PreparedScalarSegment { stats .all_counts .insert("scalar_segments_searched".into(), 1); - let SearchResult::Exact(rows) = result else { + if !result.is_exact() { return Ok(Some("inexact_result")); - }; - stats.all_counts.insert( - "scalar_segment_candidate_rows".into(), - rows.len().unwrap_or(0) as usize, - ); + } + if let Some(rows) = result.upper.max_len() { + stats + .all_counts + .insert("scalar_segment_candidate_rows".into(), rows as usize); + } else { + // A complement mask has no finite cardinality without a row universe. + // Do not report zero candidates for a successful NOT search. + stats + .all_counts + .insert("scalar_segment_candidate_rows_unknown".into(), 1); + } // Do not truncate candidates at LIMIT. The reader evaluates the complete - // filter before applying its existing limit/offset operators. - // The raw selected bitmap can overlap NULL rows; the full filter removes - // those as well. The metric above counts semantic TRUE rows, not mask size. - reader.with_row_addr_prefilter(RowAddrMask::from_allowed(rows.selected_rows().clone())); + // filter before applying its existing limit/offset operators. Its explicit + // fragment domain also bounds complement masks produced by NOT. + reader.with_row_addr_prefilter(result.upper); Ok(None) } } + +#[cfg(test)] +mod tests { + use super::*; + use lance_index::scalar::SargableQuery; + use lance_index::scalar::expression::ScalarIndexSearch; + + fn leaf(index_name: &str) -> ScalarIndexExpr { + ScalarIndexExpr::Query(ScalarIndexSearch { + column: "key".into(), + index_name: index_name.into(), + index_type: "BTree".into(), + query: Arc::new(SargableQuery::IsNull()), + needs_recheck: false, + fragment_bitmap: None, + }) + } + + fn select( + expr: &ScalarIndexExpr, + ) -> std::result::Result, &'static str> { + let mut remaining = MAX_EXPRESSION_NODES; + scoped_expression(expr, "selected", "key", false, 0, &mut remaining) + } + + #[test] + fn never_negate_a_partial_candidate_expression() { + let partial = ScalarIndexExpr::And(Box::new(leaf("selected")), Box::new(leaf("other"))); + assert!(select(&partial).unwrap().is_some()); + let negated = ScalarIndexExpr::Not(Box::new(partial)); + assert!(select(&negated).unwrap().is_none()); + let alternative = ScalarIndexExpr::Or(Box::new(leaf("selected")), Box::new(negated)); + assert!(select(&alternative).unwrap().is_none()); + } + + #[test] + fn expression_budget_bounds_depth_and_concurrent_searches() { + let mut deep = leaf("selected"); + for _ in 0..=MAX_EXPRESSION_DEPTH { + deep = ScalarIndexExpr::Not(Box::new(deep)); + } + assert_eq!(select(&deep).unwrap_err(), "expression_budget"); + let mut wide = leaf("selected"); + for _ in 0..6 { + wide = ScalarIndexExpr::Or(Box::new(wide.clone()), Box::new(wide)); + } + assert!(select(&wide).unwrap().is_some()); + wide = ScalarIndexExpr::Or(Box::new(wide.clone()), Box::new(wide)); + assert_eq!(select(&wide).unwrap_err(), "expression_budget"); + } +} diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 2a00f5e..42a3403 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -15159,9 +15159,35 @@ fn scalar_segment_ids( filter: &str, limit: Option, offset: i64, +) -> (Vec, CapturedScanStatistics) { + scalar_segment_ids_with_filter( + uri, + uuid, + fragments, + ScalarSegmentFilter::Sql(filter), + limit, + offset, + ) +} + +enum ScalarSegmentFilter<'a> { + Sql(&'a str), + Substrait(&'a [u8]), +} + +fn scalar_segment_ids_with_filter( + uri: &str, + uuid: &[u8; 16], + fragments: &[u64], + filter: ScalarSegmentFilter<'_>, + limit: Option, + offset: i64, ) -> (Vec, CapturedScanStatistics) { let uri = c_str(uri); - let filter = c_str(filter); + let sql = match filter { + ScalarSegmentFilter::Sql(sql) => Some(c_str(sql)), + ScalarSegmentFilter::Substrait(_) => None, + }; let id = c_str("id"); let columns = [id.as_ptr(), ptr::null()]; let mut captured = CapturedScanStatistics::default(); @@ -15169,8 +15195,18 @@ fn scalar_segment_ids( unsafe { let ds = lance_dataset_open(uri.as_ptr(), ptr::null(), 0); assert!(!ds.is_null()); - let scanner = lance_scanner_new(ds, columns.as_ptr(), filter.as_ptr()); + let scanner = lance_scanner_new( + ds, + columns.as_ptr(), + sql.as_ref().map_or(ptr::null(), |s| s.as_ptr()), + ); assert!(!scanner.is_null()); + if let ScalarSegmentFilter::Substrait(bytes) = filter { + assert_eq!( + lance_scanner_set_substrait_filter(scanner, bytes.as_ptr(), bytes.len()), + 0 + ); + } assert_eq!( lance_scanner_set_fragment_ids(scanner, fragments.as_ptr(), fragments.len()), 0 @@ -15286,6 +15322,257 @@ fn test_scalar_segment_scope_residual_limit_and_unindexed_fallback() { } } +#[test] +fn test_scalar_segment_boolean_predicates_preserve_scope_and_nulls() { + for kind in [ + lance_index::IndexType::BTree, + lance_index::IndexType::Bitmap, + ] { + for stable in [false, true] { + let (_tmp, uri, uuids) = + create_scalar_segment_fixture_with_options(kind, stable, None, &[&[0, 1], &[2]]); + for (filter, expected) in [ + ("key IN (1, 2)", vec![1, 2, 5, 7]), + ("key = 1 OR key = 2", vec![1, 2, 5, 7]), + ("key IN (0, 1, 2)", vec![1, 2, 3, 5, 6, 7]), + ("key IN (0, 1, 2, 99)", vec![1, 2, 3, 5, 6, 7]), + ("key IN (1, 2) AND id >= 2", vec![2, 5, 7]), + ("(key = 1 OR key = 2) AND key > 1", vec![2, 5]), + ("NOT (key = 1 OR key = 2)", vec![3, 6]), + ("key NOT IN (1, 2)", vec![3, 6]), + ("key IS NULL OR key = 1", vec![0, 1, 4, 7]), + ("key = 98 OR key = 99", vec![]), + ] { + let (ids, stats) = scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0); + assert_eq!(ids, expected, "{kind:?}, stable={stable}, {filter}"); + assert_eq!(stats.calls, 1); + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segments_searched" && *value == 1), + "{kind:?}, stable={stable}, {filter}: {:?}", + stats.metrics + ); + assert!( + !stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segment_fallbacks" && *value != 0), + "{kind:?}, stable={stable}, {filter}: {:?}", + stats.metrics + ); + } + let (ids, _) = scalar_segment_ids(&uri, &uuids[0], &[0], "key IN (1, 2)", None, 0); + assert_eq!( + ids, + vec![1, 2], + "index coverage must not widen the scan scope" + ); + let (ids, _) = scalar_segment_ids( + &uri, + &uuids[0], + &[0, 1], + "key IN (1, 2) AND id >= 2", + Some(1), + 1, + ); + assert_eq!(ids, vec![5], "residual filtering must precede pagination"); + let (ids, stats) = + scalar_segment_ids(&uri, &uuids[0], &[0, 2], "key IN (1, 2)", None, 0); + assert_eq!(ids, vec![1, 2, 10, 11]); + assert!(stats.metrics.iter().any(|(name, _, value)| name + == "scalar_segment_fallback_partial_coverage" + && *value == 1)); + // A missing OR branch cannot be discarded, even if the other branch is indexed. + let (ids, stats) = + scalar_segment_ids(&uri, &uuids[0], &[0, 1], "key IN (1, 2) OR id = 0", None, 0); + assert_eq!(ids, vec![0, 1, 2, 5, 7]); + assert!(stats.metrics.iter().any(|(name, _, value)| name + == "scalar_segment_fallback_no_driver" + && *value == 1)); + for (filter, expected) in [ + ("key IN (1, NULL)", vec![1, 7]), + ("key NOT IN (1, NULL)", vec![]), + ("NOT (key = 1 OR id = 2)", vec![3, 5, 6]), + ] { + let (ids, _) = scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0); + assert_eq!(ids, expected, "{kind:?}, stable={stable}, {filter}"); + } + } + } +} + +#[test] +fn test_scalar_segment_boolean_truth_tests_preserve_null_rows() { + for kind in [ + lance_index::IndexType::BTree, + lance_index::IndexType::Bitmap, + ] { + for stable in [false, true] { + let key = Arc::new(arrow_array::BooleanArray::from(vec![ + None, + Some(true), + Some(false), + Some(false), + None, + Some(true), + Some(false), + Some(true), + None, + Some(false), + Some(true), + Some(false), + ])); + let (_tmp, uri, uuids) = + create_scalar_segment_fixture_from_key(kind, stable, None, &[&[0, 1], &[2]], key); + for (filter, expected) in [ + ("NOT (key IS TRUE)", vec![0, 2, 3, 4, 6]), + ("NOT (key IS FALSE)", vec![0, 1, 4, 5, 7]), + ("key IS NOT TRUE", vec![0, 2, 3, 4, 6]), + ("key IS NOT FALSE", vec![0, 1, 4, 5, 7]), + ("key IS TRUE", vec![1, 5, 7]), + ("key IS FALSE", vec![2, 3, 6]), + ("NOT (key IS TRUE AND id >= 2)", vec![0, 1, 2, 3, 4, 6]), + ("NOT (key IS FALSE OR id = 0)", vec![1, 4, 5, 7]), + ("(key IS TRUE) = false", vec![0, 2, 3, 4, 6]), + ] { + let (ids, stats) = scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0); + assert_eq!(ids, expected, "{kind:?}, stable={stable}, {filter}"); + assert_eq!(stats.calls, 1); + assert!( + stats.metrics.iter().any(|(name, _, value)| name + == "scalar_segment_fallback_boolean_truth_test" + && *value == 1), + "{filter}: {:?}", + stats.metrics + ); + assert!( + !stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segments_searched" && *value != 0) + ); + } + let (ids, _) = + scalar_segment_ids(&uri, &uuids[0], &[0], "NOT (key IS TRUE)", Some(1), 1); + assert_eq!(ids, vec![2], "fallback must preserve scope and pagination"); + let (ids, stats) = + scalar_segment_ids(&uri, &uuids[0], &[0, 1], "NOT (key = true)", None, 0); + assert_eq!(ids, vec![2, 3, 6], "ordinary NOT must still exclude NULLs"); + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segments_searched" && *value == 1) + ); + } + } +} + +#[test] +fn test_scalar_segment_string_in_candidates_and_deleted_rows() { + use lance::index::DatasetIndexExt; + for kind in [ + lance_index::IndexType::BTree, + lance_index::IndexType::Bitmap, + ] { + for stable in [false, true] { + let key = Arc::new(StringArray::from(vec![ + None, + Some("alpha"), + Some("beta"), + Some("gamma"), + None, + Some("beta"), + Some("gamma"), + Some("alpha"), + None, + Some("gamma"), + Some("alpha"), + Some("beta"), + ])); + let (_tmp, uri, uuids) = + create_scalar_segment_fixture_from_key(kind, stable, None, &[&[0, 1], &[2]], key); + let (ids, stats) = scalar_segment_ids( + &uri, + &uuids[0], + &[0, 1], + "key IN ('alpha', 'beta')", + None, + 0, + ); + assert_eq!(ids, vec![1, 2, 5, 7]); + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segment_candidate_rows" && *value == 4) + ); + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "rows_scanned" && *value == 4) + ); + lance_c::runtime::block_on(async { + let mut ds = Dataset::open(&uri).await.unwrap(); + ds.delete("id IN (1, 6)").await.unwrap(); + assert_eq!(ds.load_indices().await.unwrap().len(), 2); + }); + for (filter, expected) in [ + ("key IN ('alpha', 'beta')", vec![2, 5, 7]), + ("NOT (key = 'alpha' OR key = 'beta')", vec![3]), + ] { + let (ids, _) = scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0); + assert_eq!(ids, expected, "{kind:?}, stable={stable}, {filter}"); + } + } + } +} + +#[test] +fn test_scalar_segment_boolean_predicates_do_not_search_other_indices() { + use lance::index::DatasetIndexExt; + use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; + let (_tmp, uri, uuids) = create_scalar_segment_fixture_with_options( + lance_index::IndexType::BTree, + false, + None, + &[&[0, 1], &[2]], + ); + lance_c::runtime::block_on(async { + let mut ds = Dataset::open(&uri).await.unwrap(); + ds.create_index_builder( + &["id"], + lance_index::IndexType::BTree, + &ScalarIndexParams::for_builtin(BuiltinIndexType::BTree), + ) + .name("id_idx".into()) + .await + .unwrap(); + }); + for (filter, expected, accelerated) in [ + ("key = 1 AND id >= 2", vec![7], true), + ( + "(key = 1 AND id >= 2) OR (key = 2 AND id < 6)", + vec![2, 5, 7], + true, + ), + ("key = 1 OR id = 0", vec![0, 1, 7], false), + // Negating a necessary conjunct alone would incorrectly drop id 0 and 1. + ("NOT (key = 1 AND id >= 2)", vec![0, 1, 2, 3, 5, 6], false), + ] { + let (ids, stats) = scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0); + assert_eq!(ids, expected, "{filter}"); + let searched = stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segments_searched" && *value == 1); + assert_eq!(searched, accelerated, "{filter}: {:?}", stats.metrics); + } +} + #[test] fn test_scalar_segment_metadata_residuals_fall_back_within_scope() { for kind in [ @@ -15467,6 +15754,222 @@ fn test_scalar_segment_label_list_exact_candidates() { } } +#[test] +fn test_scalar_segment_string_labels_via_substrait() { + use arrow_array::builder::{ListBuilder, StringBuilder}; + use datafusion::prelude::SessionContext; + use lance_datafusion::{planner::Planner, substrait::encode_substrait}; + use lance_index::IndexType; + + for stable in [false, true] { + let mut labels = ListBuilder::new(StringBuilder::new()); + for row in 0..16 { + match row % 8 { + 0 => labels.append(false), + 1 => labels.append(true), + n => { + if [2, 4, 5, 7].contains(&n) { + labels.values().append_value("red"); + } + if [3, 4, 7].contains(&n) { + labels.values().append_value("blue"); + } + if [5, 6].contains(&n) { + labels.values().append_null(); + } + if n == 7 { + labels.values().append_value("red"); + } + labels.append(true); + } + } + } + let (_tmp, uri, uuids) = create_scalar_segment_fixture_from_key( + IndexType::LabelList, + stable, + None, + &[&[0, 1], &[2]], + Arc::new(labels.finish()), + ); + let schema = lance_c::runtime::block_on(async { + Arc::new(Schema::from(Dataset::open(&uri).await.unwrap().schema())) + }); + let planner = Planner::new(schema.clone()); + let state = SessionContext::new().state(); + let encode = |sql: &str| { + encode_substrait(planner.parse_filter(sql).unwrap(), schema.clone(), &state).unwrap() + }; + for (filter, expected, candidates) in [ + ("array_contains(key, 'red')", vec![2, 4, 5, 7], 4), + ( + "array_contains(key, 'red') AND array_contains(key, 'blue')", + vec![4, 7], + 2, + ), + ( + "array_contains(key, 'red') OR array_contains(key, 'blue')", + vec![2, 3, 4, 5, 7], + 5, + ), + ("array_has_all(key, ['red', 'blue'])", vec![4, 7], 2), + ( + "array_has_any(key, ['red', 'blue'])", + vec![2, 3, 4, 5, 7], + 5, + ), + ("array_contains(key, 'missing')", vec![], 0), + ("array_contains(key, 'red') AND id >= 4", vec![4, 5, 7], 4), + ] { + let bytes = encode(filter); + let (ids, stats) = scalar_segment_ids_with_filter( + &uri, + &uuids[0], + &[0, 1], + ScalarSegmentFilter::Substrait(&bytes), + None, + 0, + ); + assert_eq!(ids, expected, "stable={stable}, {filter}"); + // Candidate counts distinguish indexed label intersections from reading + // one label's matches and applying the second label only as a residual. + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segment_candidate_rows" + && *value == candidates), + "{filter}: {:?}", + stats.metrics + ); + assert_eq!( + ids, + scalar_segment_ids(&uri, &uuids[0], &[0, 1], filter, None, 0).0 + ); + assert!( + stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segments_searched" && *value == 1), + "{filter}: {:?}", + stats.metrics + ); + assert!( + !stats + .metrics + .iter() + .any(|(name, _, value)| name == "scalar_segment_fallbacks" && *value != 0), + "{filter}: {:?}", + stats.metrics + ); + } + // Cross-engine callers must not equate NULL membership or ordered-subsequence + // functions with Lance's membership and set-containment semantics. + for (filter, expected) in [ + ("array_contains(key, CAST(NULL AS STRING))", vec![]), + ("array_has_any(key, [])", vec![]), + // The pinned Lance planner treats an empty all-label query as true, + // including on NULL lists; compare against its unindexed path below. + ("array_has_all(key, [])", vec![0, 1, 2, 3, 4, 5, 6, 7]), + ("array_has_all(key, ['blue', 'red'])", vec![4, 7]), + ("array_contains(key, 'red') OR id = 0", vec![0, 2, 4, 5, 7]), + ] { + let bytes = encode(filter); + let (ids, _) = scalar_segment_ids_with_filter( + &uri, + &uuids[0], + &[0, 1], + ScalarSegmentFilter::Substrait(&bytes), + None, + 0, + ); + let baseline = lance_c::runtime::block_on(async { + use futures::TryStreamExt; + let ds = Dataset::open(&uri).await.unwrap(); + let mut scan = ds.scan(); + scan.project(&["id"]).unwrap(); + scan.use_scalar_index(false); + scan.with_fragments( + ds.get_fragments() + .into_iter() + .take(2) + .map(|f| f.metadata().clone()) + .collect(), + ); + scan.filter(filter).unwrap(); + let batches: Vec = scan + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>() + }); + assert_eq!( + ids, baseline, + "indexed vs unindexed: stable={stable}, {filter}" + ); + assert_eq!(ids, expected, "stable={stable}, {filter}"); + } + let filter = "array_contains(key, 'red') AND array_contains(key, 'blue') AND id >= 4"; + let bytes = encode(filter); + let (ids, _) = scalar_segment_ids_with_filter( + &uri, + &uuids[0], + &[1], + ScalarSegmentFilter::Substrait(&bytes), + Some(1), + 1, + ); + assert_eq!( + ids, + vec![7], + "residual filtering and fragment scope precede pagination" + ); + let (ids, stats) = scalar_segment_ids_with_filter( + &uri, + &uuids[0], + &[0, 3], + ScalarSegmentFilter::Substrait(&bytes), + None, + 0, + ); + assert_eq!( + ids, + vec![12, 15], + "fallback must include matching unindexed fragments" + ); + assert!(stats.metrics.iter().any(|(name, _, value)| name + == "scalar_segment_fallback_partial_coverage" + && *value == 1)); + lance_c::runtime::block_on(async { + let mut ds = Dataset::open(&uri).await.unwrap(); + ds.delete("id = 4").await.unwrap(); + }); + let (ids, _) = scalar_segment_ids_with_filter( + &uri, + &uuids[0], + &[0, 1], + ScalarSegmentFilter::Substrait(&bytes), + None, + 0, + ); + assert_eq!(ids, vec![7]); + } +} + #[test] fn test_scalar_segment_text_indices_still_fall_back() { for kind in [lance_index::IndexType::Fm, lance_index::IndexType::NGram] {