scalar index segment

Author: zhangstar333

docs/scalar-segment-scans.md

@@ -0,0 +1,78 @@
+# Scalar index segment scans
+
+An ordinary scanner can use one physical BTree/Bitmap segment to generate
+candidates, then read those candidates with the complete scanner filter. This
+does not run a global search of the other segments of the logical index. It does
+not subdivide a physical segment or make its own index search incremental.
+
+## Configuring a task
+
+Open a fixed dataset version. Select the physical index UUID from that version's
+metadata and pass the task's complete fragment domain explicitly:
+
+```c
+LanceScanner *scanner = lance_scanner_new(dataset, columns, full_filter_sql);
+/* Check every return value in production. */
+lance_scanner_set_fragment_ids(scanner, fragment_ids, fragment_count);
+lance_scanner_set_scalar_index_segment(scanner, segment_uuid_16_bytes);
+lance_scanner_set_limit(scanner, 20000);
+/* The scanner-owning thread calls lance_scanner_next as usual. */
+```
+
+SQL, Substrait and additional SQL filters keep their existing precedence and AND
+composition. The caller does not supply a separate driver predicate: Lance-C
+uses the typed filter planner and selects a necessary indexed leaf belonging to
+the requested logical index. It only descends through AND, never through OR or
+NOT. It then searches the selected UUID and applies the complete filter while
+reading candidates with automatic scalar-index planning disabled.
+
+Each task's fragment IDs define its result domain, including on fallback. A
+distributed planner must assign disjoint domains whose union covers the intended
+scan. Unindexed fragments need their own tasks, or an explicit domain including
+them (which causes that task to use fallback). Merely listing indexed segments
+does not include appended, unindexed data automatically.
+
+An unknown UUID, absent fragment or invalid option combination is an error. A
+known segment with incomplete/unknown coverage, no suitable driver, unsupported
+index type, nested key, overlays, fragment reuse, non-exact results or unsupported
+row-ID domain falls back to a non-indexed scan of the entire explicit domain.
+I/O and corruption errors are propagated, not converted to empty results or
+successful fallback.
+
+The first implementation supports live-row ordinary scans and cannot be combined
+with vector/FTS queries. Physical row-address
+results on stable-row-ID datasets currently fall back; results already expressed
+in the correct row-ID domain use the candidate path. Deletes and all remaining
+predicates are handled by the ordinary reader. No candidate-count limit is
+applied: LIMIT/OFFSET remain after the scanner's complete filter.
+
+If the host has additional predicates outside Lance, do not set a local limit
+before those predicates. Never divide the global limit by the number of tasks.
+Global OFFSET belongs to the coordinator, not independently to each task.
+
+## Stopping after the host limit
+
+This mode uses the existing scanner/stream lifecycle. Once the host has enough
+rows, it stops requesting further batches and closes the scanner after any active
+call has returned. Do not call `lance_scanner_close` concurrently with `next`.
+Exported Arrow streams remain owned by the caller and must also be released after
+their active consumers have finished.
+
+A host stop flag does not interrupt an in-progress `lance_scanner_next`: current
+index evaluation or I/O may finish before the host observes stop and closes the
+stream. No separate cancellation signal or thread is introduced. The host remains
+responsible for enforcing the global LIMIT across concurrent tasks.
+
+## Memory and statistics
+
+Candidate masks stay in Rust, and record batches are streamed. Each active task
+can still hold a complete segment's candidate set; scanner I/O buffer size does
+not cap that allocation. Control task concurrency and physical segment size.
+
+Successful exhaustion merges segment-search metrics into the existing statistics
+callback exactly once. New metrics include `scalar_segments_requested`,
+`scalar_segments_searched`, `scalar_segment_candidate_rows`,
+`scalar_segment_prepare_time`, `scalar_segment_search_time`, and
+`scalar_segment_fallback_*` reasons. `prepare_time` includes search time. Early
+release, cancellation and errors retain the existing callback contract: final
+statistics are not guaranteed. Metrics do not establish global task concurrency.

include/lance/lance.h

@@ -1763,6 +1763,30 @@ int32_t lance_scanner_set_index_segments(
     size_t len
 );
 
+/**
+ * Accelerate an ordinary scalar-filtered scan with one physical index segment.
+ * segment_uuid points to 16 UUID bytes in RFC 4122 order; NULL clears the setting.
+ * Must be configured before scanning. Requires explicit nonempty fragment_ids,
+ * which define BOTH the read and fallback domain, independently of the segment.
+ * Missing snapshot UUIDs / fragment IDs are errors. Extra segment coverage is
+ * excluded by fragment_ids; incomplete coverage falls back to a full filtered
+ * scan of those fragment_ids. Callers distributing work must assign disjoint
+ * fragment domains and separately include any unindexed data they wish to read.
+ *
+ * BTree/Bitmap searches use a necessary AND-conjunct of the full scanner filter
+ * on the selected logical index. All predicates are reapplied during candidate
+ * reads; other scalar indices are disabled. OR/NOT-only filters, 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
+ * unfiltered candidate set. Vector/FTS queries are rejected.
+ *
+ * UUID bytes are copied. Metadata and final option compatibility are validated
+ * when creating the stream. Index corruption or I/O failures remain errors.
+ */
+int32_t lance_scanner_set_scalar_index_segment(
+    LanceScanner* scanner, const uint8_t* segment_uuid);
+
 /* ─── Full-text search (Phase 2) ─── */
 
 /**

include/lance/lance.hpp

@@ -1276,6 +1276,20 @@ class Scanner {
         return *this;
     }
 
+    /// Restrict scalar candidate generation to one segment; fragment_ids is
+    /// required and defines the complete read/fallback domain. See lance.h.
+    Scanner& scalar_index_segment(const std::array<uint8_t, 16>& segment_uuid) {
+        if (lance_scanner_set_scalar_index_segment(handle_.get(), segment_uuid.data()) != 0)
+            check_error();
+        return *this;
+    }
+
+    Scanner& clear_scalar_index_segment() {
+        if (lance_scanner_set_scalar_index_segment(handle_.get(), nullptr) != 0)
+            check_error();
+        return *this;
+    }
+
     /// Restrict scan to specific fragment IDs.
     Scanner& fragment_ids(const uint64_t* ids, size_t len) {
         if (lance_scanner_set_fragment_ids(handle_.get(), ids, len) != 0)

src/lib.rs

@@ -39,6 +39,7 @@ mod index_segment;
 mod merge_insert;
 mod restore;
 pub mod runtime;
+mod scalar_segment;
 mod scanner;
 mod session;
 pub mod stream_guard;

src/scalar_segment.rs

@@ -0,0 +1,224 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+//! Segment-scoped candidate generation for ordinary scans. The explicit fragment
+//! list is the read domain, including on fallback; a segment is only an accelerator.
+
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::time::Instant;
+
+use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
+use lance::Dataset;
+use lance::dataset::scanner::{
+    ExecutionStatsCallback, ExecutionSummaryCounts, RowAddrMask, 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 uuid::Uuid;
+
+pub(crate) struct PreparedScalarSegment {
+    pub dataset: Arc<Dataset>,
+    pub segment_uuid: Uuid,
+    pub fragment_ids: Vec<u64>,
+    pub callback: Option<ExecutionStatsCallback>,
+}
+
+fn invalid(message: impl Into<String>) -> 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> {
+    match expr {
+        ScalarIndexExpr::Query(search) if search.index_name == index_name => Some(search),
+        ScalarIndexExpr::And(lhs, rhs) => {
+            driver(lhs, index_name).or_else(|| driver(rhs, index_name))
+        }
+        _ => None,
+    }
+}
+
+impl PreparedScalarSegment {
+    pub async fn configure(self, mut reader: Scanner) -> Result<Scanner> {
+        // Never let either candidate reads or fallback re-enter a global index search.
+        reader.use_scalar_index(false);
+        let mut stats = ExecutionSummaryCounts::default();
+        stats
+            .all_counts
+            .insert("scalar_segments_requested".into(), 1);
+        let plan_metrics = ExecutionPlanMetricsSet::new();
+        let metrics = IndexMetrics::new(&plan_metrics, 0);
+        let started = Instant::now();
+        let reason = self
+            .configure_candidates(&mut reader, &metrics, &mut stats)
+            .await?;
+        metrics.flush_io();
+        stats.all_times.insert(
+            "scalar_segment_prepare_time".into(),
+            started.elapsed().as_nanos().min(usize::MAX as u128) as usize,
+        );
+        if let Some(reason) = reason {
+            stats
+                .all_counts
+                .insert("scalar_segment_fallbacks".into(), 1);
+            stats
+                .all_counts
+                .insert(format!("scalar_segment_fallback_{reason}"), 1);
+        }
+        for (name, count) in plan_metrics.clone_inner().iter_counts() {
+            let name = name.as_ref();
+            match name {
+                "iops" => stats.iops += count.value(),
+                "requests" => stats.requests += count.value(),
+                "bytes_read" => stats.bytes_read += count.value(),
+                "indices_loaded" => stats.indices_loaded += count.value(),
+                "parts_loaded" => stats.parts_loaded += count.value(),
+                "index_comparisons" => stats.index_comparisons += count.value(),
+                _ => *stats.all_counts.entry(name.to_string()).or_default() += count.value(),
+            }
+        }
+        if let Some(callback) = self.callback {
+            // Preserve the callback's once-per-successfully-exhausted-stream contract.
+            // Candidate work is not part of the underlying reader's plan metrics.
+            reader.scan_stats_callback(Arc::new(move |read| {
+                let mut combined = read.clone();
+                combined.iops += stats.iops;
+                combined.requests += stats.requests;
+                combined.bytes_read += stats.bytes_read;
+                combined.indices_loaded += stats.indices_loaded;
+                combined.parts_loaded += stats.parts_loaded;
+                combined.index_comparisons += stats.index_comparisons;
+                for (name, value) in &stats.all_counts {
+                    *combined.all_counts.entry(name.clone()).or_default() += value;
+                }
+                for (name, value) in &stats.all_times {
+                    *combined.all_times.entry(name.clone()).or_default() += value;
+                }
+                callback(&combined);
+            }));
+        }
+        Ok(reader)
+    }
+
+    async fn configure_candidates(
+        &self,
+        reader: &mut Scanner,
+        metrics: &IndexMetrics,
+        stats: &mut ExecutionSummaryCounts,
+    ) -> Result<Option<&'static str>> {
+        let fragments = self.dataset.get_fragments();
+        let visible: HashSet<u64> = fragments.iter().map(|f| f.id() as u64).collect();
+        if self.fragment_ids.iter().any(|id| !visible.contains(id)) {
+            return Err(invalid(
+                "scalar segment fragment_ids contains a fragment absent from the dataset snapshot",
+            ));
+        }
+        let indices = self.dataset.load_indices().await?;
+        let index_meta = indices
+            .iter()
+            .find(|i| i.uuid == self.segment_uuid)
+            .ok_or_else(|| {
+                invalid(format!(
+                    "scalar index segment {} is absent from the dataset snapshot",
+                    self.segment_uuid
+                ))
+            })?;
+        let field_id = index_meta
+            .keyed_field()
+            .ok_or_else(|| invalid("scalar segment must index a single key field"))?;
+        let field =
+            self.dataset.schema().field_by_id(field_id).ok_or_else(|| {
+                invalid("scalar segment key field is absent from the dataset schema")
+            })?;
+        // Keep V1 to flat scalar fields. A dotted name is not sufficient to prove
+        // the field path of an evolved or nested schema.
+        if !self
+            .dataset
+            .schema()
+            .fields
+            .iter()
+            .any(|f| f.id == field.id)
+        {
+            return Ok(Some("nested_field"));
+        }
+        let scope: HashSet<u64> = self.fragment_ids.iter().copied().collect();
+        let Some(coverage) = index_meta.fragment_bitmap.as_ref() else {
+            return Ok(Some("unknown_coverage"));
+        };
+        if self
+            .fragment_ids
+            .iter()
+            .any(|id| u32::try_from(*id).map_or(true, |id| !coverage.contains(id)))
+        {
+            // Scan the ENTIRE explicit read domain, not just the covered part.
+            return Ok(Some("partial_coverage"));
+        }
+        if fragments
+            .iter()
+            .filter(|f| scope.contains(&(f.id() as u64)))
+            .any(|f| !f.metadata().overlays.is_empty() || f.metadata().physical_rows.is_none())
+        {
+            return Ok(Some("fragment_state"));
+        }
+        // Fragment reuse can change the domain of an old segment. Until its
+        // coverage mapping is handled here, preserve correctness with a scoped scan.
+        if self.dataset.frag_reuse_index_uuid().await.is_some() {
+            return Ok(Some("fragment_reuse"));
+        }
+        let Some(filter) = reader.get_expr_filter()? else {
+            return Ok(Some("no_filter"));
+        };
+        let planner = Planner::new(Arc::new(self.dataset.schema().into()));
+        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 {
+            return Ok(Some("no_driver"));
+        };
+        if search.column != field.name {
+            return Ok(Some("field_path"));
+        }
+        let index = self
+            .dataset
+            .open_scalar_index(&search.column, &self.segment_uuid, metrics)
+            .await?;
+        if !matches!(index.index_type(), IndexType::BTree | IndexType::Bitmap) {
+            return Ok(Some("index_type"));
+        }
+        // External masks use _rowid, not necessarily physical row addresses.
+        if index.results_are_row_addresses() && self.dataset.manifest.uses_stable_row_ids() {
+            return Ok(Some("row_id_domain"));
+        }
+        let started = Instant::now();
+        let result = index.search(search.query.as_ref(), metrics).await?;
+        stats.all_times.insert(
+            "scalar_segment_search_time".into(),
+            started.elapsed().as_nanos().min(usize::MAX as u128) as usize,
+        );
+        stats
+            .all_counts
+            .insert("scalar_segments_searched".into(), 1);
+        let SearchResult::Exact(rows) = result else {
+            return Ok(Some("inexact_result"));
+        };
+        stats.all_counts.insert(
+            "scalar_segment_candidate_rows".into(),
+            rows.len().unwrap_or(0) as usize,
+        );
+        // Do not truncate candidates at LIMIT. The reader evaluates the complete
+        // filter before applying its existing limit/offset operators.
+        reader.with_row_addr_prefilter(RowAddrMask::from_allowed(rows.selected_rows().clone()));
+        Ok(None)
+    }
+}