diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 6fe172d9715..3b57438d700 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -206,6 +206,10 @@ harness = false name = "count_pushdown" harness = false +[[bench]] +name = "logical_scan_planner" +harness = false + [[bench]] name = "vector_index" harness = false diff --git a/rust/lance/benches/logical_scan_planner.rs b/rust/lance/benches/logical_scan_planner.rs new file mode 100644 index 00000000000..c4d4c2e617c --- /dev/null +++ b/rust/lance/benches/logical_scan_planner.rs @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Benchmarks comparing the imperative scan planner against the logical-plan prototype. +//! +//! Both paths are called directly rather than through `LANCE_LOGICAL_SCAN_PLANNER`, so the two +//! appear in one process and criterion can put them side by side. +//! +//! Three questions, which is why there are three groups: +//! +//! * `plan/` — is going through a logical plan, a rule loop and a physical planner more expensive +//! than hand-building the exec tree? This is the cost the prototype adds, measured alone. +//! * `scan/` and `search/` — does the resulting plan execute at the same speed? Planning is a +//! fixed cost per query; execution is what a real workload pays. +//! +//! ```text +//! cargo bench -p lance --bench logical_scan_planner +//! ``` + +use std::sync::Arc; + +use arrow_array::Float32Array; +use arrow_array::types::{Float32Type, Int32Type}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::dataset::WriteParams; +use lance::dataset::scanner::Scanner; +use lance::index::DatasetIndexExt; +use lance::index::vector::VectorIndexParams; +use lance_core::utils::tempfile::TempStrDir; +use lance_datafusion::exec::{LanceExecutionOptions, execute_plan}; +use lance_datagen::{BatchCount, ByteCount, Dimension, RowCount, array, gen_batch}; +use lance_index::IndexType; +use lance_index::scalar::ScalarIndexParams; +use lance_linalg::distance::DistanceType; +#[cfg(target_os = "linux")] +use lance_testing::pprof::{Output, PProfProfiler}; + +const DIM: u32 = 64; +const ROWS_PER_FRAGMENT: u64 = 25_000; +const NUM_FRAGMENTS: u32 = 8; +const TOTAL_ROWS: i32 = (ROWS_PER_FRAGMENT as u32 * NUM_FRAGMENTS) as i32; + +/// On-disk rather than `memory://`: the prototype's plans differ from the imperative ones in how +/// much they read and when, which an in-memory store would flatten out. +struct Fixture { + _datadir: TempStrDir, + dataset: Arc, +} + +impl Fixture { + async fn open() -> Self { + let datadir = TempStrDir::default(); + let reader = gen_batch() + .col("i", array::step::()) + .col("s", array::rand_utf8(ByteCount::from(32), false)) + .col("vec", array::rand_vec::(Dimension::from(DIM))) + .into_reader_rows( + RowCount::from(ROWS_PER_FRAGMENT), + BatchCount::from(NUM_FRAGMENTS), + ); + let mut dataset = Dataset::write( + reader, + datadir.as_str(), + Some(WriteParams { + max_rows_per_file: ROWS_PER_FRAGMENT as usize, + ..Default::default() + }), + ) + .await + .unwrap(); + + // A scalar index on the filter column so the planning groups exercise the scalar-index + // rules, which are the prototype's most expensive rewrite. + dataset + .create_index( + &["i"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + dataset + .create_index( + &["vec"], + IndexType::Vector, + None, + &VectorIndexParams::ivf_pq(16, 8, 8, DistanceType::L2, 20), + true, + ) + .await + .unwrap(); + + Self { + _datadir: datadir, + dataset: Arc::new(dataset), + } + } +} + +fn query_vector() -> Float32Array { + Float32Array::from((0..DIM).map(|v| v as f32).collect::>()) +} + +/// One benchmarked query, as a name and a way to build the scanner for it. +/// +/// The planner under test is chosen by the caller rather than baked in here, so both paths are +/// guaranteed to plan the identical query. +type Shape = (&'static str, fn(&Dataset) -> Scanner); + +/// The query shapes both groups run, named so criterion's output reads as a comparison. +fn shapes() -> Vec { + vec![ + ("full_scan", |dataset| dataset.scan()), + ("filtered_scan", |dataset| { + let mut scan = dataset.scan(); + scan.project(&["s"]).unwrap(); + scan.filter(&format!("i < {}", TOTAL_ROWS / 100)).unwrap(); + scan + }), + ("filtered_scan_with_limit", |dataset| { + let mut scan = dataset.scan(); + scan.project(&["s"]).unwrap(); + scan.filter(&format!("i < {}", TOTAL_ROWS / 2)).unwrap(); + scan.limit(Some(100), None).unwrap(); + scan + }), + ("ann", |dataset| { + let mut scan = dataset.scan(); + scan.project(&["s"]).unwrap(); + scan.nearest("vec", &query_vector(), 10).unwrap(); + scan + }), + ("ann_prefiltered", |dataset| { + let mut scan = dataset.scan(); + scan.project(&["s"]).unwrap(); + scan.prefilter(true); + scan.filter(&format!("i < {}", TOTAL_ROWS / 10)).unwrap(); + scan.nearest("vec", &query_vector(), 10).unwrap(); + scan + }), + ] +} + +/// Whether a shape's cost is dominated by reading rows or by the index search, which is the only +/// reason to separate the two execution groups. +fn is_search(shape: &str) -> bool { + shape.starts_with("ann") +} + +async fn plan( + scanner: &Scanner, + use_logical: bool, +) -> Arc { + if use_logical { + scanner.create_plan_logical().await.unwrap() + } else { + scanner.create_plan().await.unwrap() + } +} + +async fn plan_and_execute(scanner: &Scanner, use_logical: bool) -> usize { + let plan = plan(scanner, use_logical).await; + execute_plan(plan, LanceExecutionOptions::default()) + .unwrap() + .try_fold(0, |rows, batch| async move { Ok(rows + batch.num_rows()) }) + .await + .unwrap() +} + +fn bench_planning(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let fixture = rt.block_on(Fixture::open()); + let dataset = &fixture.dataset; + + let mut group = c.benchmark_group("plan"); + for (shape, configure) in shapes() { + let scanner = configure(dataset); + // Plan once per path before measuring: both read index metadata on their first call and + // cache it on the dataset, and that one-time cost would otherwise land in whichever path + // criterion warmed up first. + for use_logical in [false, true] { + rt.block_on(plan(&scanner, use_logical)); + let path = if use_logical { "logical" } else { "imperative" }; + group.bench_function(BenchmarkId::new(path, shape), |b| { + b.iter(|| rt.block_on(plan(&scanner, use_logical))) + }); + } + } + group.finish(); +} + +fn bench_execution(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let fixture = rt.block_on(Fixture::open()); + let dataset = &fixture.dataset; + + for (name, want_search) in [("scan", false), ("search", true)] { + let mut group = c.benchmark_group(name); + for (shape, configure) in shapes() { + if is_search(shape) != want_search { + continue; + } + let scanner = configure(dataset); + for use_logical in [false, true] { + let path = if use_logical { "logical" } else { "imperative" }; + // Assert the two paths agree on row count before timing them. A path that returns + // fewer rows would otherwise look like a speedup. + let rows = rt.block_on(plan_and_execute(&scanner, use_logical)); + assert!(rows > 0, "{path}/{shape} returned no rows"); + group.bench_function(BenchmarkId::new(path, shape), |b| { + b.iter(|| rt.block_on(plan_and_execute(&scanner, use_logical))) + }); + } + } + group.finish(); + } +} + +#[cfg(target_os = "linux")] +criterion_group!( + name = benches; + config = Criterion::default().significance_level(0.1).sample_size(10) + .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); + targets = bench_planning, bench_execution); + +#[cfg(not(target_os = "linux"))] +criterion_group!( + name = benches; + config = Criterion::default().significance_level(0.1).sample_size(10); + targets = bench_planning, bench_execution); + +criterion_main!(benches); diff --git a/rust/lance/src/datafusion.rs b/rust/lance/src/datafusion.rs index e8749e86c8f..0747156467e 100644 --- a/rust/lance/src/datafusion.rs +++ b/rust/lance/src/datafusion.rs @@ -6,4 +6,5 @@ pub(crate) mod dataframe; pub(crate) mod logical_plan; -pub use dataframe::LanceTableProvider; +pub use crate::dataset::scanner::logical::dataframe::{LanceContextExt, LanceDataFrameExt}; +pub use dataframe::{LanceTableProvider, SessionContextExt}; diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 223a848167a..79a6207ce99 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2799,6 +2799,19 @@ impl Scanner { if logical::is_enabled() { return logical::create_plan(self).await; } + self.create_plan_imperative().await + } + + /// Plan this scan through the logical planner, whatever `LANCE_LOGICAL_SCAN_PLANNER` says. + /// + /// Exists for the planning benchmark, which needs both paths in one process so criterion can + /// put them side by side. Goes away with the imperative path. + #[doc(hidden)] + pub async fn create_plan_logical(&self) -> Result> { + logical::create_plan(self).await + } + + async fn create_plan_imperative(&self) -> Result> { self.validate_options()?; let full_text_query = match &self.full_text_query { diff --git a/rust/lance/src/dataset/scanner/logical/builder.rs b/rust/lance/src/dataset/scanner/logical/builder.rs index a0f70537c22..ea8fb09016d 100644 --- a/rust/lance/src/dataset/scanner/logical/builder.rs +++ b/rust/lance/src/dataset/scanner/logical/builder.rs @@ -20,7 +20,7 @@ use datafusion::logical_expr::{ use datafusion::prelude::col; use lance_core::{ - ROW_ADDR, ROW_ID, + ROW_ADDR, ROW_ID, ROW_OFFSET, datatypes::{OnMissing, Projection}, }; use lance_index::scalar::inverted::{DOC_INDEX_COL, SCORE_COL}; @@ -28,6 +28,7 @@ use lance_index::vector::DIST_COL; use super::fts; use super::prepare::PreparedQueries; +use super::row_offset::RowOffsetNode; use super::source::{LanceScanSource, ScanSourceOptions}; use super::{LanceTakeNode, TakeSettings, VectorAccessPath, VectorRerankNode, VectorSearchNode}; use crate::dataset::scanner::{ColumnOrdering, MaterializationStyle}; @@ -49,8 +50,18 @@ pub fn build(scanner: &Scanner, prepared: &PreparedQueries) -> Result Result>, + /// Fragment row counts and deletion vectors, loaded only when the plan asks for `_rowoffset`. + /// `AddRowOffsetExec::try_new` is the one physical constructor that does I/O; this is how it + /// stops being one. + row_offsets: Option, } impl ScanPlanningContext { @@ -197,8 +203,14 @@ impl ScanPlanningContext { pub async fn collect(plan: &LogicalPlan) -> Result { let mut leaf = None; let mut searches: Vec<(String, Option)> = Vec::new(); + let mut needs_row_offsets = false; let mut takes: Vec = Vec::new(); plan.apply(|node| { + if let LogicalPlan::Extension(extension) = node + && extension.node.as_any().is::() + { + needs_row_offsets = true; + } takes.extend(take_operations(node)); if let LogicalPlan::Extension(extension) = node && let Some(search) = extension.node.as_any().downcast_ref::() @@ -260,6 +272,10 @@ impl ScanPlanningContext { let scalar_indices = Arc::new(dataset.scalar_index_info().await?); let index_staleness = prefetch_index_staleness(&dataset, &fragments).await?; + let row_offsets = match needs_row_offsets { + true => Some(RowOffsetMap::load(dataset.clone()).await?), + false => None, + }; let take_rows = resolve_takes(&dataset, takes).await?; Ok(Self { @@ -274,10 +290,15 @@ impl ScanPlanningContext { }, scalar_indices, index_staleness, + row_offsets, take_rows, }) } + pub fn row_offsets(&self) -> Option<&RowOffsetMap> { + self.row_offsets.as_ref() + } + /// The rows a take-shaped predicate selects, if resolving it needed I/O and stage 2 did it. pub fn take_rows(&self, take: &TakeOperation) -> Option<&Arc> { self.take_rows.get(&take_key(take)) diff --git a/rust/lance/src/dataset/scanner/logical/dataframe.rs b/rust/lance/src/dataset/scanner/logical/dataframe.rs new file mode 100644 index 00000000000..d386f40be71 --- /dev/null +++ b/rust/lance/src/dataset/scanner/logical/dataframe.rs @@ -0,0 +1,273 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance search as DataFusion `DataFrame` operators. +//! +//! The scanner exposes one fixed query shape: filter, search, sort, limit, project, in that order. +//! Everything the logical path added to make that shape planable — a scan leaf that is a real +//! `TableProvider`, search nodes that are real logical nodes, and a lowering stage that reads only +//! the plan — also makes the shape unnecessary. A `DataFrame` can put a search anywhere, and join, +//! aggregate, or window the result with anything DataFusion can express. +//! +//! ``` +//! # use std::sync::Arc; +//! # use datafusion::prelude::{SessionContext, col, lit}; +//! # use lance::Result; +//! # use lance::dataset::Dataset; +//! # use lance::datafusion::{LanceContextExt, LanceDataFrameExt}; +//! # use lance_index::scalar::FullTextSearchQuery; +//! # use lance_index::vector::Query; +//! # async fn hybrid(dataset: Arc, text: FullTextSearchQuery, vector: Query) -> Result<()> { +//! let ctx = SessionContext::new(); +//! let plan = ctx +//! .read_lance_dataset(dataset)? +//! .filter(col("category").eq(lit("news")))? +//! .full_text_search(text) +//! .await? +//! .nearest(vector)? +//! .limit(0, Some(10))? +//! .lance_plan() +//! .await?; +//! # let _ = plan; +//! # Ok(()) +//! # } +//! ``` +//! +//! The filter is pushed into the Lance scan leaf, the text search runs against the inverted index, +//! and the vector search re-scores its matches — one plan, planned once. Note that this is *not* +//! gated by [`is_enabled`](super::is_enabled): there is no imperative equivalent to fall back to. + +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::TableProvider; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::logical_expr::{LogicalPlan, Projection as DfProjection}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::{DataFrame, SessionContext, col}; +use lance_core::ROW_ID; +use lance_core::datatypes::{OnMissing, Projection}; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::scalar::inverted::{DOC_INDEX_COL, SCORE_COL}; +use lance_index::vector::{DIST_COL, Query}; + +use super::builder::{extension, source_options}; +use super::fts::{self, FtsCompoundNode, FtsLeafNode}; +use super::source::LanceScanSource; +use super::{LanceTakeNode, TakeSettings, VectorAccessPath, VectorSearchNode, with_lance_source}; +use crate::dataset::Dataset; +use crate::{Error, Result}; + +/// Read a Lance dataset as a `DataFrame` that the Lance scan planner will lower. +pub trait LanceContextExt { + /// A `DataFrame` over `dataset`, backed by the logical path's own scan leaf. + /// + /// Unlike [`SessionContextExt::read_lance`](crate::datafusion::SessionContextExt::read_lance), + /// this does not go through `Dataset::scan()` — the returned frame's leaf is the same + /// `TableProvider` the scanner's logical plans are built on, so filters and projections land in + /// the scan itself and [`LanceDataFrameExt`]'s operators can be stacked on top. + fn read_lance_dataset(&self, dataset: Arc) -> Result; +} + +impl LanceContextExt for SessionContext { + fn read_lance_dataset(&self, dataset: Arc) -> Result { + let defaults = dataset.scan(); + let source = LanceScanSource::new(dataset, source_options(&defaults))?; + Ok(self.read_table(Arc::new(source) as Arc)?) + } +} + +/// Lance search operators for a `DataFrame` whose leaf is a Lance dataset. +#[async_trait] +pub trait LanceDataFrameExt: Sized { + /// Nearest-neighbour search over this frame's rows, ordered by `_distance`. + /// + /// The result carries the frame's existing columns plus `_distance`; the columns are re-read by + /// row id, so a search never has to carry them through. + /// + /// Whether the search uses a vector index is left to the planner — *unless* this frame is + /// already the result of a search, in which case its rows are the search space and the scoring + /// is exact. That is the same rule the scanner applies to a vector search over a full-text + /// filter. + fn nearest(self, query: Query) -> Result; + + /// Full-text search over this frame's rows, ordered by descending `_score`. + /// + /// Async because a query that names neither a column nor a document granularity is completed + /// from the dataset's inverted indices, which is I/O. + async fn full_text_search(self, query: FullTextSearchQuery) -> Result; + + /// Lower this frame through the Lance scan planner. + /// + /// Use this instead of `DataFrame::create_physical_plan`: the Lance nodes need Lance's own + /// analyzer, optimizer, and physical rules, and the index metadata they read is prefetched + /// here. The frame's session config and runtime are kept. + async fn lance_plan(self) -> Result>; +} + +#[async_trait] +impl LanceDataFrameExt for DataFrame { + fn nearest(self, query: Query) -> Result { + let (state, plan) = self.into_parts(); + let dataset = lance_dataset(&plan)?; + let plan = carrying_row_ids(plan)?; + + let search = if is_search_result(&plan) { + // These rows are the search space, so they are scored exactly — and scoring them means + // reading their vectors, which a search's output does not carry. + let candidates = fts::take_column( + plan.clone(), + &dataset, + &query.column, + &TakeSettings::default(), + )?; + VectorSearchNode::try_new(candidates, dataset.clone(), query)? + .with_resolution(VectorAccessPath::Flat) + } else { + VectorSearchNode::try_new(plan.clone(), dataset.clone(), query)? + }; + let searched = with_take(extension(search), &dataset, &plan)?; + + Ok(Self::new(state, searched).sort(vec![col(DIST_COL).sort(true, false)])?) + } + + async fn full_text_search(self, query: FullTextSearchQuery) -> Result { + let (state, plan) = self.into_parts(); + let dataset = lance_dataset(&plan)?; + let plan = carrying_row_ids(plan)?; + + let resolved = dataset + .scan() + .resolve_full_text_search_query(&query) + .await?; + let mut searched = fts::build_source(plan.clone(), &dataset, &resolved, None)?; + // A list-element query scores each matching element, so one row can come back several + // times. Rows are what a frame's later operators expect. + if searched + .schema() + .has_column_with_unqualified_name(DOC_INDEX_COL) + { + searched = fts::dedupe_rows(searched)?; + } + let searched = with_take(searched, &dataset, &plan)?; + + Ok(Self::new(state, searched).sort(vec![col(SCORE_COL).sort(false, false)])?) + } + + async fn lance_plan(self) -> Result> { + let (state, plan) = self.into_parts(); + // The frame's own state has DataFusion's physical rules, not Lance's, and lowering a Lance + // node depends on them — `EnforceDistribution` in particular decides how a search fans out. + let state = SessionStateBuilder::new_from_existing(state) + .with_physical_optimizer_rules(super::physical_optimizer_rules()) + .build(); + super::lower(plan, Arc::new(state)).await + } +} + +/// The dataset a frame reads, recovered from its scan leaf. +fn lance_dataset(plan: &LogicalPlan) -> Result> { + let mut dataset = None; + plan.apply(|node| { + dataset = with_lance_source(node, |source| source.dataset().clone()); + Ok(match dataset { + Some(_) => TreeNodeRecursion::Stop, + None => TreeNodeRecursion::Continue, + }) + })?; + dataset.ok_or_else(|| { + Error::invalid_input( + "a Lance search can only be added to a DataFrame read with read_lance_dataset" + .to_string(), + ) + }) +} + +/// The same plan, emitting `_rowid`. +/// +/// A search identifies its results by row id — that is how a prefilter reaches the index, and how +/// the take above the search reads the columns back. A user's own projection has no reason to keep +/// the column, so it is put back here rather than made a rule of the API. +fn carrying_row_ids(plan: LogicalPlan) -> Result { + if plan.schema().has_column_with_unqualified_name(ROW_ID) { + return Ok(plan); + } + let plan = plan + .transform_up(|node| { + let LogicalPlan::Projection(projection) = &node else { + return Ok(Transformed::no(node)); + }; + if !projection + .input + .schema() + .has_column_with_unqualified_name(ROW_ID) + { + return Ok(Transformed::no(node)); + } + let mut exprs = projection.expr.clone(); + exprs.push(col(ROW_ID)); + Ok(Transformed::yes(LogicalPlan::Projection( + DfProjection::try_new(exprs, projection.input.clone())?, + ))) + })? + .data; + + if !plan.schema().has_column_with_unqualified_name(ROW_ID) { + return Err(Error::invalid_input(format!( + "a Lance search needs {ROW_ID}, and nothing below this point in the plan produces it", + ))); + } + Ok(plan) +} + +/// Whether this plan already scored its rows, and so is a candidate set rather than a table. +fn is_search_result(plan: &LogicalPlan) -> bool { + let mut found = false; + let _ = plan.apply(|node| { + let LogicalPlan::Extension(extension) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + let node = extension.node.as_any(); + found = node.is::() + || node.is::() + || node.is::(); + Ok(match found { + true => TreeNodeRecursion::Stop, + false => TreeNodeRecursion::Continue, + }) + }); + found +} + +/// Re-read the columns `before` carried, which a search's output does not. +fn with_take( + searched: LogicalPlan, + dataset: &Arc, + before: &LogicalPlan, +) -> Result { + let columns = before + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let mut projection = + Projection::empty(dataset.clone() as Arc) + .union_columns(&columns, OnMissing::Ignore)?; + // A row's version columns come from its position within its fragment, so only an ordered scan + // can produce them and a take never can. A frame with no projection of its own asks for every + // column the leaf advertises, these among them, so they are dropped here rather than left for + // the read to be asked for something it cannot return. + projection.with_row_last_updated_at_version = false; + projection.with_row_created_at_version = false; + if LanceTakeNode::is_noop(&searched, &projection)? { + return Ok(searched); + } + Ok(extension(LanceTakeNode::try_new( + searched, + dataset.clone(), + projection, + TakeSettings::default(), + )?)) +} diff --git a/rust/lance/src/dataset/scanner/logical/mod.rs b/rust/lance/src/dataset/scanner/logical/mod.rs index 15d4aa4b285..0324cb54e95 100644 --- a/rust/lance/src/dataset/scanner/logical/mod.rs +++ b/rust/lance/src/dataset/scanner/logical/mod.rs @@ -87,6 +87,7 @@ //! coverage splitting a search across indexed and unindexed fragments //! scan_index recording on each scan how it finds its rows (index query, or a take) //! planner stage 4: dispatch to each node's lowering +//! row_offset the `_rowoffset` column, whose node is the one that needs a prefetch //! take/ late materialization //! vector/ node, rerank, rules, planner <- five entry points //! fts/ node, rules, planner, prefetch <- the same five @@ -101,9 +102,11 @@ pub(super) mod builder; pub(super) mod context; pub(super) mod coverage; +pub mod dataframe; pub(super) mod fts; pub(super) mod planner; pub(super) mod prepare; +pub(super) mod row_offset; pub(super) mod rules; pub(super) mod scan_index; pub(super) mod source; @@ -113,6 +116,7 @@ mod tests; pub(super) mod vector; pub use coverage::*; +pub use row_offset::*; pub use rules::*; pub use scan_index::*; pub use take::*; @@ -134,7 +138,6 @@ use datafusion::physical_optimizer::join_selection::JoinSelection; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; -use lance_core::ROW_OFFSET; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_datafusion::exec::{StrictBatchSizeExec, get_session_context}; @@ -319,6 +322,7 @@ fn analyzer_rules(context: &Arc) -> Vec Result<()> { nearest_query_count: _, } = scanner; - if reads_row_offset(scanner)? { - // `_rowoffset` is computed above the scan rather than read from it, so producing it needs a - // node this path does not have yet. - return Err(Error::not_supported_source( - "The logical scan planner cannot produce _rowoffset yet".into(), - )); - } if *include_deleted_rows && (nearest.is_some() || full_text_query.is_some()) { // The imperative path rejects these in `vector_search_source`/`fts_search_source` for the // same reason: a search returns row ids, and a deleted row does not have one. @@ -492,15 +489,3 @@ fn ensure_supported(scanner: &Scanner) -> Result<()> { } Ok(()) } - -fn reads_row_offset(scanner: &Scanner) -> Result { - let names_it = |expr: &datafusion::logical_expr::Expr| { - expr.column_refs().iter().any(|col| col.name == ROW_OFFSET) - }; - Ok(scanner - .projection_plan - .requested_output_expr - .iter() - .any(|output| names_it(&output.expr)) - || scanner.get_expr_filter()?.as_ref().is_some_and(names_it)) -} diff --git a/rust/lance/src/dataset/scanner/logical/planner.rs b/rust/lance/src/dataset/scanner/logical/planner.rs index 62f2d743da6..be2808a2b46 100644 --- a/rust/lance/src/dataset/scanner/logical/planner.rs +++ b/rust/lance/src/dataset/scanner/logical/planner.rs @@ -17,6 +17,7 @@ use datafusion::physical_plan::expressions; use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; use datafusion_physical_expr::PhysicalSortExpr; +use super::row_offset::{RowOffsetNode, plan_row_offset}; use super::{LanceTakeNode, PrefilterSourceKind, VectorRerankNode, VectorSearchNode}; use super::{plan_flat_knn, plan_take, plan_vector_search}; use crate::Result; @@ -62,6 +63,9 @@ impl ExtensionPlanner for LanceExtensionPlanner { if let Some(take) = node.as_any().downcast_ref::() { return Ok(Some(plan_take(take, input)?)); } + if let Some(offsets) = node.as_any().downcast_ref::() { + return Ok(Some(plan_row_offset(offsets, input)?)); + } Ok(None) } } diff --git a/rust/lance/src/dataset/scanner/logical/row_offset.rs b/rust/lance/src/dataset/scanner/logical/row_offset.rs new file mode 100644 index 00000000000..1fed0288fef --- /dev/null +++ b/rust/lance/src/dataset/scanner/logical/row_offset.rs @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The `_rowoffset` column: node, resolution rule, and lowering. +//! +//! A row's offset is its position in the dataset once deletions are accounted for, so computing one +//! needs every earlier fragment's row count and deletion vector. That is the only I/O any physical +//! node's constructor does, which is why the load happens in stage 2 and a rule hands the result to +//! the node here. + +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow_schema::Schema as ArrowSchema; +use datafusion::common::tree_node::Transformed; +use datafusion::common::{DFSchema, DFSchemaRef, plan_err}; +use datafusion::config::ConfigOptions; +use datafusion::logical_expr::{ + Expr, Extension, InvariantLevel, LogicalPlan, UserDefinedLogicalNodeCore, +}; +use datafusion::optimizer::AnalyzerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::expressions; +use datafusion::physical_plan::projection::ProjectionExec; +use lance_core::{ROW_ADDR, ROW_OFFSET, ROW_OFFSET_FIELD}; + +use super::context::ScanPlanningContext; +use super::rules::analyze_bottom_up; +use crate::Result; +use crate::io::exec::{AddRowOffsetExec, RowOffsetMap}; + +/// Append `_rowoffset` to the input's columns. +#[derive(Clone)] +pub struct RowOffsetNode { + input: LogicalPlan, + /// The per-fragment state the computation needs. Filled in by [`ResolveRowOffsets`]; a node + /// that still has `None` here cannot be lowered. + offsets: Option, + schema: DFSchemaRef, +} + +impl RowOffsetNode { + pub fn try_new(input: LogicalPlan) -> Result { + if !input.schema().has_column_with_unqualified_name(ROW_ADDR) { + return Err(crate::Error::internal(format!( + "a {ROW_OFFSET} column is computed from {ROW_ADDR}, which its input does not have" + ))); + } + let mut fields = input.schema().as_arrow().fields().to_vec(); + fields.push(Arc::new(ROW_OFFSET_FIELD.clone())); + let schema = Arc::new(DFSchema::try_from(ArrowSchema::new_with_metadata( + fields, + input.schema().as_arrow().metadata().clone(), + ))?); + Ok(Self { + input, + offsets: None, + schema, + }) + } + + pub fn offsets(&self) -> Option<&RowOffsetMap> { + self.offsets.as_ref() + } + + fn with_offsets(mut self, offsets: RowOffsetMap) -> Self { + self.offsets = Some(offsets); + self + } +} + +impl fmt::Debug for RowOffsetNode { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.fmt_for_explain(f) + } +} + +impl PartialEq for RowOffsetNode { + fn eq(&self, other: &Self) -> bool { + self.input == other.input && self.offsets.is_some() == other.offsets.is_some() + } +} + +impl Eq for RowOffsetNode {} + +impl Hash for RowOffsetNode { + fn hash(&self, state: &mut H) { + self.input.hash(state); + self.offsets.is_some().hash(state); + } +} + +impl PartialOrd for RowOffsetNode { + fn partial_cmp(&self, other: &Self) -> Option { + self.input.partial_cmp(&other.input) + } +} + +impl UserDefinedLogicalNodeCore for RowOffsetNode { + fn name(&self) -> &str { + "RowOffset" + } + + fn check_invariants(&self, check: InvariantLevel) -> datafusion::common::Result<()> { + if matches!(check, InvariantLevel::Executable) && self.offsets.is_none() { + return plan_err!("{ROW_OFFSET} reached execution with no fragment offsets loaded"); + } + Ok(()) + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "RowOffset") + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + mut inputs: Vec, + ) -> datafusion::common::Result { + if !exprs.is_empty() { + return Err(datafusion::common::DataFusionError::Internal( + "RowOffset takes no expressions".into(), + )); + } + if inputs.len() != 1 { + return Err(datafusion::common::DataFusionError::Internal(format!( + "RowOffset takes exactly one input, got {}", + inputs.len() + ))); + } + Ok(Self { + input: inputs.remove(0), + offsets: self.offsets.clone(), + schema: self.schema.clone(), + }) + } + + /// Every input column carries through, and `_rowaddr` is read even when nothing above wants it. + fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { + Some(vec![(0..self.input.schema().fields().len()).collect()]) + } +} + +/// Hand each `_rowoffset` node the fragment offsets stage 2 loaded for it. +#[derive(Debug)] +pub struct ResolveRowOffsets { + context: Arc, +} + +impl ResolveRowOffsets { + pub fn new(context: Arc) -> Self { + Self { context } + } +} + +impl AnalyzerRule for ResolveRowOffsets { + fn name(&self) -> &str { + "resolve_row_offsets" + } + + fn analyze( + &self, + plan: LogicalPlan, + _config: &ConfigOptions, + ) -> datafusion::common::Result { + analyze_bottom_up(plan, |node| { + let LogicalPlan::Extension(extension) = &node else { + return Ok(Transformed::no(node)); + }; + let Some(offsets) = extension + .node + .as_any() + .downcast_ref::() + .filter(|row_offset| row_offset.offsets().is_none()) + else { + return Ok(Transformed::no(node)); + }; + let Some(loaded) = self.context.row_offsets() else { + return plan_err!( + "{ROW_OFFSET} was requested but stage 2 did not load its offsets" + ); + }; + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(offsets.clone().with_offsets(loaded.clone())), + }))) + }) + } +} + +pub fn plan_row_offset( + node: &RowOffsetNode, + input: Arc, +) -> Result> { + let offsets = node.offsets().ok_or_else(|| { + crate::Error::internal(format!( + "{ROW_OFFSET} lowering ran before its offsets loaded" + )) + })?; + let with_offsets = Arc::new(AddRowOffsetExec::try_new_from_map(input, offsets)?); + + // The read emits the system columns in its own order, which need not be the order the node + // declared. Restating the declared order here is what keeps the two schemas equal, which + // DataFusion checks at every extension boundary. + let schema = with_offsets.schema(); + let columns = node + .schema() + .fields() + .iter() + .map(|field| { + let name = field.name(); + Ok((expressions::col(name, schema.as_ref())?, name.clone())) + }) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(columns, with_offsets)?)) +} diff --git a/rust/lance/src/dataset/scanner/logical/scan_index.rs b/rust/lance/src/dataset/scanner/logical/scan_index.rs index b000674c9c1..609d5a45656 100644 --- a/rust/lance/src/dataset/scanner/logical/scan_index.rs +++ b/rust/lance/src/dataset/scanner/logical/scan_index.rs @@ -9,16 +9,16 @@ use std::sync::Arc; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::datasource::{provider_as_source, source_as_provider}; -use datafusion::logical_expr::LogicalPlan; use datafusion::logical_expr::expr_rewriter::unnormalize_cols; use datafusion::logical_expr::utils::conjunction; +use datafusion::logical_expr::{Extension, Filter, LogicalPlan}; use datafusion::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule}; use lance_select::mask::RowAddrTreeMap; use roaring::RoaringBitmap; -use super::PrefilterSourceKind; use super::context::{OverlayStaleness, ScanPlanningContext}; use super::source::{LanceScanSource, ScanRestriction}; +use super::{PrefilterSourceKind, RowOffsetNode}; use crate::dataset::scanner::TakeOperation; /// Derive each Lance scan's scalar index query and record it on the scan's source. @@ -255,6 +255,10 @@ impl OptimizerRule for ResolveTake { ) -> datafusion::common::Result> { match &plan { LogicalPlan::TableScan(_) => self.restrict_scan(plan), + // A `_rowoffset` predicate never reaches the scan: the column is computed by the node + // above it, and `PushDownFilter` does not push through an extension. Matching that pair + // is what lets those predicates become takes too. + LogicalPlan::Filter(_) => self.restrict_below_row_offsets(plan), _ => Ok(Transformed::no(plan)), } } @@ -304,4 +308,50 @@ impl ResolveTake { source.restricted_to(&ScanRestriction::Rows(rows.clone())) })?)) } + + /// [`Self::restrict_scan`] for a predicate stranded above a [`RowOffsetNode`]. + fn restrict_below_row_offsets( + &self, + plan: LogicalPlan, + ) -> datafusion::common::Result> { + let LogicalPlan::Filter(filter) = &plan else { + return Ok(Transformed::no(plan)); + }; + let LogicalPlan::Extension(offsets) = filter.input.as_ref() else { + return Ok(Transformed::no(plan)); + }; + if offsets + .node + .as_any() + .downcast_ref::() + .is_none() + { + return Ok(Transformed::no(plan)); + } + let [scan] = offsets.node.inputs()[..] else { + return Ok(Transformed::no(plan)); + }; + let restrictable = with_lance_source(scan, |source| source.options().rows.is_none()); + if restrictable != Some(true) { + return Ok(Transformed::no(plan)); + } + let Some((take, remainder)) = TakeOperation::try_from_expr(&filter.predicate) else { + return Ok(Transformed::no(plan)); + }; + let Some(rows) = self.rows_for(&take) else { + return Ok(Transformed::no(plan)); + }; + let restricted = map_lance_scan(scan, |source| { + source.restricted_to(&ScanRestriction::Rows(rows.clone())) + })?; + let offsets = LogicalPlan::Extension(Extension { + node: offsets + .node + .with_exprs_and_inputs(vec![], vec![restricted])?, + }); + Ok(Transformed::yes(match remainder { + Some(remainder) => LogicalPlan::Filter(Filter::try_new(remainder, Arc::new(offsets))?), + None => offsets, + })) + } } diff --git a/rust/lance/src/dataset/scanner/logical/tests/dataframe.rs b/rust/lance/src/dataset/scanner/logical/tests/dataframe.rs new file mode 100644 index 00000000000..606382716e6 --- /dev/null +++ b/rust/lance/src/dataset/scanner/logical/tests/dataframe.rs @@ -0,0 +1,200 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance search as DataFrame operators. +//! +//! The scanner's fixed query shape is not the oracle here — the point of these is queries the +//! scanner cannot express — so they assert directly on the rows instead. + +use std::sync::Arc; + +use arrow::datatypes::{Float32Type, Int32Type}; +use arrow_array::cast::AsArray; +use datafusion::prelude::{SessionContext, col, lit}; +use lance_index::vector::Query; + +use super::fts::{fts_dataset, match_query}; +use super::harness::*; +use crate::datafusion::{LanceContextExt, LanceDataFrameExt}; +use crate::dataset::Dataset; + +fn vector_query(dataset: &Dataset, k: usize) -> Query { + let mut scan = dataset.scan(); + scan.nearest("vec", &query_vector(), k).unwrap(); + scan.nearest.clone().expect("nearest sets the query") +} + +/// A filter through the DataFrame API reaches the same rows as the scanner's own. +#[tokio::test] +async fn test_dataframe_filter_matches_the_scanner() { + let dataset = Arc::new(test_dataset().await); + + let ctx = SessionContext::new(); + let plan = ctx + .read_lance_dataset(dataset.clone()) + .unwrap() + .select(vec![col("i")]) + .unwrap() + .filter(col("i").gt(lit(10)).and(col("i").lt(lit(20)))) + .unwrap() + .lance_plan() + .await + .unwrap(); + + let batch = run(plan).await.unwrap(); + let values = batch["i"].as_primitive::().values().to_vec(); + assert_eq!(values, (11..20).collect::>()); +} + +/// A vector search over a frame that is a plain scan is free to use the index. +#[tokio::test] +async fn test_dataframe_vector_search_uses_the_index() { + let dataset = Arc::new(indexed_vector_dataset().await); + + let ctx = SessionContext::new(); + let plan = ctx + .read_lance_dataset(dataset.clone()) + .unwrap() + .select(vec![col("i")]) + .unwrap() + .nearest(vector_query(&dataset, 10)) + .unwrap() + .lance_plan() + .await + .unwrap(); + + let display = datafusion::physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!(display.contains("ANNSubIndex"), "{display}"); + + let batch = run(plan).await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let distances = batch[lance_index::vector::DIST_COL] + .as_primitive::() + .values() + .to_vec(); + assert!(distances.windows(2).all(|pair| pair[0] <= pair[1])); + assert!(batch.column_by_name("i").is_some(), "take did not run"); +} + +/// A filter below the search is a prefilter — the scan leaf applies it, and the search only ever +/// sees the surviving rows. +#[tokio::test] +async fn test_dataframe_filter_below_a_search_is_a_prefilter() { + let dataset = Arc::new(indexed_vector_dataset().await); + + let ctx = SessionContext::new(); + let plan = ctx + .read_lance_dataset(dataset.clone()) + .unwrap() + .select(vec![col("i")]) + .unwrap() + .filter(col("i").lt(lit(100))) + .unwrap() + .nearest(vector_query(&dataset, 5)) + .unwrap() + .lance_plan() + .await + .unwrap(); + + let batch = run(plan).await.unwrap(); + assert_eq!(batch.num_rows(), 5); + assert!( + batch["i"] + .as_primitive::() + .values() + .iter() + .all(|value| *value < 100), + "{batch:?}" + ); +} + +/// The headline: text search then vector search, in one plan. +/// +/// The vector search scores the text matches exactly rather than consulting an index, because its +/// input is already a candidate set. The scanner expresses this only as a `query_filter`; here it +/// is just two operators stacked. +#[tokio::test] +async fn test_dataframe_hybrid_text_then_vector() { + let dataset = Arc::new(fts_dataset().await); + + let ctx = SessionContext::new(); + let plan = ctx + .read_lance_dataset(dataset.clone()) + .unwrap() + .select(vec![col("i"), col("s")]) + .unwrap() + .full_text_search(match_query("hello")) + .await + .unwrap() + .nearest(vector_query(&dataset, 5)) + .unwrap() + .lance_plan() + .await + .unwrap(); + + let display = datafusion::physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!(display.contains("MatchQuery"), "{display}"); + assert!( + !display.contains("ANNSubIndex"), + "a search over candidates must be exact: {display}" + ); + + let batch = run(plan).await.unwrap(); + assert_eq!(batch.num_rows(), 5); + for text in batch["s"].as_string::().iter() { + assert!( + text.expect("s is not nullable").contains("hello"), + "{batch:?}" + ); + } +} + +/// A search needs a Lance dataset under it, and says so rather than panicking on the downcast. +#[tokio::test] +async fn test_search_needs_a_lance_frame() { + let dataset = Arc::new(indexed_vector_dataset().await); + let query = vector_query(&dataset, 5); + + let ctx = SessionContext::new(); + let frame = ctx.read_empty().unwrap(); + let err = frame.nearest(query).expect_err("not a Lance frame"); + assert!(err.to_string().contains("read_lance_dataset"), "{err}"); +} + +/// Aggregating a frame's search results. +/// +/// A frame can put an aggregate straight above a search, where the scanner always has a take in +/// between, so this is the first place DataFusion compares the search node's declared schema against +/// the plan it lowered to. +#[tokio::test] +async fn test_dataframe_aggregate_over_a_search() { + use datafusion::functions_aggregate::expr_fn::count; + + let dataset = Arc::new(vector_dataset().await); + + let ctx = SessionContext::new(); + let plan = ctx + .read_lance_dataset(dataset.clone()) + .unwrap() + .nearest(vector_query(&dataset, 10)) + .unwrap() + .aggregate(vec![], vec![count(lit(1))]) + .unwrap() + .lance_plan() + .await + .unwrap(); + + let batch = run(plan).await.unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!( + batch + .column(0) + .as_primitive::() + .value(0), + 10 + ); +} diff --git a/rust/lance/src/dataset/scanner/logical/tests/mod.rs b/rust/lance/src/dataset/scanner/logical/tests/mod.rs index 63102ef5cfe..ddca85da11d 100644 --- a/rust/lance/src/dataset/scanner/logical/tests/mod.rs +++ b/rust/lance/src/dataset/scanner/logical/tests/mod.rs @@ -6,6 +6,7 @@ //! Most of these are equivalence tests: they build the same query through both the imperative and //! the logical path and compare the rows. See [`harness`] for that oracle. +mod dataframe; mod fts; mod harness; mod planner; diff --git a/rust/lance/src/dataset/scanner/logical/tests/scan.rs b/rust/lance/src/dataset/scanner/logical/tests/scan.rs index 8273b1115d9..510e96a4524 100644 --- a/rust/lance/src/dataset/scanner/logical/tests/scan.rs +++ b/rust/lance/src/dataset/scanner/logical/tests/scan.rs @@ -501,3 +501,45 @@ async fn test_a_row_address_take() { .await .unwrap(); } + +/// `_rowoffset` counts live rows from the start of the dataset, so a deleted row shifts it. +#[tokio::test] +#[rstest::rstest] +#[case::projected(false)] +#[case::taken(true)] +async fn test_row_offsets(#[case] as_take: bool) { + use arrow::datatypes::UInt64Type; + use arrow_array::cast::AsArray; + + let mut dataset = test_dataset().await; + dataset.delete("i = 3").await.unwrap(); + + let scan_config = config(move |scan: &mut crate::dataset::Scanner| { + scan.project(&["i", "_rowoffset"])?; + if as_take { + scan.filter("_rowoffset IN (5, 10, 20)")?; + } + Ok(scan) + }); + let batch = scan_rows(&dataset, scan_config).await.unwrap(); + + // Offsets count live rows, so deleting `i = 3` shifts every row after it down by one. + let offsets = batch[lance_core::ROW_OFFSET] + .as_primitive::() + .values(); + let ids = batch["i"] + .as_primitive::() + .values(); + let expected_id = |offset: u64| if offset < 3 { offset } else { offset + 1 } as i32; + for (offset, id) in offsets.iter().zip(ids) { + assert_eq!( + *id, + expected_id(*offset), + "offset {offset} names the wrong row" + ); + } + match as_take { + true => assert_eq!(offsets, &[5, 10, 20]), + false => assert_eq!(offsets, &(0..199).collect::>()), + } +}