diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index 897b385f..cdfed832 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -15,25 +15,42 @@ // specific language governing permissions and limitations // under the License. -use std::fmt::Debug; +use std::collections::HashMap; +use std::fmt::{self, Debug}; use std::sync::Arc; use async_trait::async_trait; +use datafusion::arrow::array::{ + Array, ArrayRef, Int64Array, RecordBatch, RecordBatchOptions, UInt32Array, +}; +use datafusion::arrow::compute::cast; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::catalog::Session; use datafusion::catalog::TableFunctionImpl; -use datafusion::common::project_schema; +use datafusion::common::stats::Precision; +use datafusion::common::{internal_err, project_schema, Statistics}; use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, +}; use datafusion::prelude::SessionContext; +use futures::{stream, TryStreamExt}; use paimon::catalog::Catalog; +use paimon::spec::{ + BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, +}; +use paimon::table::Table; use crate::error::to_datafusion_error; use crate::runtime::{await_with_runtime, block_on_with_runtime}; -use crate::table::{PaimonScanBuilder, PaimonTableProvider}; +use crate::table::{datafusion_read_fields, PaimonTableProvider}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; @@ -207,60 +224,314 @@ impl TableProvider for VectorSearchTableProvider { async fn scan( &self, - state: &dyn Session, + _state: &dyn Session, projection: Option<&Vec>, _filters: &[Expr], limit: Option, ) -> DFResult> { - let table = self.inner.table(); + let projected_schema = project_schema(&self.schema(), projection)?; + + // An outer `LIMIT 0` needs no rows. + if limit == Some(0) { + return Ok(Arc::new(EmptyExec::new(projected_schema))); + } - let row_ranges = await_with_runtime(async { - let mut builder = table.new_vector_search_builder(); + // The search runs with the table function's own top-k (`self.limit`) so the ANN + // recall/search width is unchanged; the outer DataFusion `limit` only truncates + // the already-ranked result before any rows are read (so a large top-k with a + // small outer LIMIT doesn't read/materialize everything). All of this — search, + // read and rank-order gather — runs at execution time in the exec's stream, so + // planning / EXPLAIN stays cheap and the work is driven by the TaskContext. + Ok(Arc::new(VectorSearchExec::new( + self.inner.table().clone(), + self.column_name.clone(), + self.query_vector.clone(), + self.limit, + limit, + projection.cloned(), + projected_schema, + ))) + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DFResult> { + Ok(vec![ + TableProviderFilterPushDown::Unsupported; + filters.len() + ]) + } +} + +/// Execution-time plan for `vector_search`: runs the ANN search, reads the matching +/// rows, and gathers them into best-first relevance order when its stream is polled, +/// so planning (and `EXPLAIN`) stays cheap and the work runs under DataFusion's +/// `TaskContext`. +#[derive(Debug, Clone)] +struct VectorSearchExec { + table: Table, + column_name: String, + query_vector: Vec, + /// The table function's own top-k — drives the ANN search width; never reduced. + search_limit: usize, + /// The outer DataFusion `LIMIT`, applied by truncating the ranked result. + output_limit: Option, + projection: Option>, + output_schema: ArrowSchemaRef, + plan_properties: Arc, +} + +impl VectorSearchExec { + fn new( + table: Table, + column_name: String, + query_vector: Vec, + search_limit: usize, + output_limit: Option, + projection: Option>, + output_schema: ArrowSchemaRef, + ) -> Self { + let plan_properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + table, + column_name, + query_vector, + search_limit, + output_limit, + projection, + output_schema, + plan_properties, + } + } + + async fn compute_batch(&self) -> DFResult { + // Best-first row-ids from the index, searched at the full top-k so the ANN + // recall is unchanged (data-evolution / global-index path; PK-vector tables are + // unsupported here, as before). + let mut search_result = await_with_runtime(async { + let mut builder = self.table.new_vector_search_builder(); builder .with_vector_column(&self.column_name) .with_query_vector(self.query_vector.clone()) - .with_limit(self.limit); - builder.execute().await.map_err(to_datafusion_error) + .with_limit(self.search_limit); + builder.execute_scored().await.map_err(to_datafusion_error) }) .await?; - if row_ranges.is_empty() { - let schema = project_schema(&self.schema(), projection)?; - return Ok(Arc::new(EmptyExec::new(schema))); + if search_result.is_empty() { + return Ok(RecordBatch::new_empty(self.output_schema.clone())); } - let mut read_builder = table.new_read_builder(); - if let Some(limit) = limit { - read_builder.with_limit(limit); + // Apply the outer LIMIT by truncating the ranked result *before* reading, so a + // large top-k with a small outer LIMIT only reads/materializes the rows it can + // return (the search itself is unaffected). + if let Some(n) = self.output_limit { + if search_result.row_ids.len() > n { + search_result.row_ids.truncate(n); + search_result.scores.truncate(n); + } } - let scan = read_builder.new_scan().with_row_ranges(row_ranges); - let plan = await_with_runtime(scan.plan()) - .await - .map_err(to_datafusion_error)?; - let target = state.config_options().execution.target_partitions; - PaimonScanBuilder { - table, - schema: &self.schema(), - plan: &plan, - scan_trace: None, - projection, - pushed_predicate: None, - limit, - target_partitions: target, - filter_exact: false, - case_sensitive: true, + // Read the projected columns (+ internal `_ROW_ID`); the row-range scan yields + // file order, realigned to relevance rank below. + let read_fields = projected_read_fields(&self.table, self.projection.as_ref())?; + let row_ranges = search_result.to_row_ranges().map_err(to_datafusion_error)?; + let batches = await_with_runtime(async { + let mut read_builder = self.table.new_read_builder(); + read_builder + .with_read_type(read_fields) + .with_row_ranges(row_ranges); + let scan = read_builder.new_scan(); + let plan = scan.plan().await.map_err(to_datafusion_error)?; + let table_read = read_builder.new_read().map_err(to_datafusion_error)?; + let mut stream = table_read + .to_arrow(plan.splits()) + .map_err(to_datafusion_error)?; + let mut batches: Vec = Vec::new(); + while let Some(batch) = stream.try_next().await.map_err(to_datafusion_error)? { + batches.push(batch); + } + Ok::<_, DataFusionError>(batches) + }) + .await?; + + // Realign file-order rows to best-first rank and drop `_ROW_ID`. + gather_rows_by_rank(&batches, &search_result.row_ids, &self.output_schema) + } +} + +impl DisplayAs for VectorSearchExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "VectorSearchExec: column={}, search_limit={}, output_limit={:?}", + self.column_name, self.search_limit, self.output_limit + ) + } +} + +impl ExecutionPlan for VectorSearchExec { + fn name(&self) -> &str { + "VectorSearchExec" + } + + fn properties(&self) -> &Arc { + &self.plan_properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DFResult> { + if !children.is_empty() { + return internal_err!("VectorSearchExec is a leaf and takes no children"); } - .build() + Ok(self) } - fn supports_filters_pushdown( + fn execute( &self, - filters: &[&Expr], - ) -> DFResult> { - Ok(vec![ - TableProviderFilterPushDown::Unsupported; - filters.len() - ]) + partition: usize, + _context: Arc, + ) -> DFResult { + if partition != 0 { + return internal_err!( + "VectorSearchExec has a single partition, got partition {partition}" + ); + } + let exec = self.clone(); + let stream = stream::once(async move { exec.compute_batch().await }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.output_schema.clone(), + Box::pin(stream), + ))) + } + + fn partition_statistics(&self, _partition: Option) -> DFResult> { + Ok(Arc::new(Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: Statistics::unknown_column(&self.output_schema), + })) + } +} + +/// Projected user columns (+ internal `_ROW_ID`, needed to realign rows to rank). +/// Errors if the table has no row tracking, since results then can't be ordered. +fn projected_read_fields( + table: &paimon::table::Table, + projection: Option<&Vec>, +) -> DFResult> { + let base_fields = datafusion_read_fields(table); + let mut read_fields: Vec = match projection { + Some(indices) => indices.iter().map(|&i| base_fields[i].clone()).collect(), + None => base_fields, + }; + if !read_fields + .iter() + .any(|field| field.name() == ROW_ID_FIELD_NAME) + { + if !CoreOptions::new(table.schema().options()).row_tracking_enabled() { + return Err(DataFusionError::Plan( + "vector_search: cannot order results by relevance because _ROW_ID is not available" + .to_string(), + )); + } + read_fields.push(DataField::new( + ROW_ID_FIELD_ID, + ROW_ID_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::with_nullable(true)), + )); + } + Ok(read_fields) +} + +/// Gather the file-order `batches` into `ranked_row_ids` order (rank == slice index), +/// producing `output_schema` (which excludes `_ROW_ID`). A permutation driven by the +/// index's existing ranking, not a re-sort. +fn gather_rows_by_rank( + batches: &[RecordBatch], + ranked_row_ids: &[u64], + output_schema: &ArrowSchemaRef, +) -> DFResult { + let input_schema = batches.first().map(|batch| batch.schema()).ok_or_else(|| { + DataFusionError::Internal("vector_search: no rows materialized".to_string()) + })?; + let combined = arrow_select::concat::concat_batches(&input_schema, batches) + .map_err(DataFusionError::from)?; + + let row_id_index = combined.schema().index_of(ROW_ID_FIELD_NAME).map_err(|_| { + DataFusionError::Internal(format!( + "vector_search: materialized rows are missing the {ROW_ID_FIELD_NAME} column" + )) + })?; + let row_ids = combined + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal(format!("vector_search: {ROW_ID_FIELD_NAME} must be Int64")) + })?; + + // Map global row id -> physical position in the materialized batch. + let mut position_of: HashMap = HashMap::with_capacity(combined.num_rows()); + for row in 0..combined.num_rows() { + if !row_ids.is_null(row) { + position_of.insert(row_ids.value(row), row as u32); + } + } + + // Emit in rank order; extra scanned rows are ignored, but a missing scored id + // fails loud rather than silently shrinking the top-k. + let mut take_indices: Vec = Vec::with_capacity(ranked_row_ids.len()); + for &row_id in ranked_row_ids { + let position = position_of.get(&(row_id as i64)).ok_or_else(|| { + DataFusionError::Internal(format!( + "vector_search: scored row id {row_id} was not materialized; \ + cannot return the requested top-k" + )) + })?; + take_indices.push(*position); } + let take_indices = UInt32Array::from(take_indices); + let row_count = take_indices.len(); + + let columns = output_schema + .fields() + .iter() + .map(|field| -> DFResult { + let index = combined.schema().index_of(field.name()).map_err(|_| { + DataFusionError::Internal(format!( + "vector_search: materialized rows are missing expected column '{}'", + field.name() + )) + })?; + let taken = + arrow_select::take::take(combined.column(index).as_ref(), &take_indices, None) + .map_err(DataFusionError::from)?; + // The Paimon read keeps its own arrow types (e.g. `Utf8`), but the provider + // schema may differ (e.g. DataFusion's `Utf8View`); cast to match, as the + // normal scan path does via `to_datafusion_batch`. + if taken.data_type() == field.data_type() { + Ok(taken) + } else { + cast(taken.as_ref(), field.data_type()).map_err(DataFusionError::from) + } + }) + .collect::>>()?; + + // Preserve the row count for a zero-column projection (e.g. `COUNT(*)`). + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + RecordBatch::try_new_with_options(Arc::clone(output_schema), columns, &options) + .map_err(DataFusionError::from) } diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index b4637b5b..e4fc4721 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1719,14 +1719,16 @@ mod fulltext_tests { mod vector_search_tests { use std::sync::Arc; - use datafusion::arrow::array::{ArrayRef, Float32Builder, Int32Array, ListBuilder}; + use datafusion::arrow::array::{ + ArrayRef, Float32Builder, Int32Array, ListBuilder, StringArray, + }; use datafusion::arrow::datatypes::{ DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, }; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; use paimon::catalog::Identifier; - use paimon::spec::{ArrayType, DataType, FloatType, IntType, Schema}; + use paimon::spec::{ArrayType, DataType, FloatType, IntType, Schema, VarCharType}; use paimon::table::BranchManager; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::{register_vector_search, SQLContext}; @@ -1877,6 +1879,22 @@ mod vector_search_tests { ids } + /// Like [`extract_ids`] but preserves the row order returned by the query (used to + /// assert relevance ordering). + fn extract_ids_in_order(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec { + let mut ids = Vec::new(); + for batch in batches { + let id_array = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("Expected Int32Array for id"); + for i in 0..batch.num_rows() { + ids.push(id_array.value(i)); + } + } + ids + } + fn extract_query_result_ids( batches: &[datafusion::arrow::record_batch::RecordBatch], ) -> Vec<(i32, i32)> { @@ -1961,6 +1979,216 @@ mod vector_search_tests { assert!(ids.contains(&0), "exact match [1,0,0,0] should be in top 3"); } + /// `vector_search` must return rows in relevance order, not file order. Querying + /// id 2's vector ranks `[2, 5, 1]` by L2 distance, which differs from the + /// ascending-id file order `[1, 2, 5]` the pre-fix scan produced. + #[tokio::test] + async fn test_vector_search_orders_by_relevance() { + let (ctx, catalog, _tmp) = create_empty_vector_search_context().await; + let identifier = Identifier::new("default", "vindex_order_e2e"); + catalog + .create_table(&identifier, build_vindex_table_schema(), false) + .await + .expect("Failed to create table"); + let table = catalog + .get_table(&identifier) + .await + .expect("Failed to load table"); + + let write_builder = table + .new_write_builder() + .with_commit_user("test-user") + .expect("Failed to configure write builder"); + let mut table_write = write_builder + .new_write() + .expect("Failed to create table write"); + table_write + .write_arrow_batch(&build_vector_batch( + vec![0, 1, 2, 3, 4, 5], + vec![ + vec![1.0, 0.0], + vec![0.9, 0.1], + vec![0.0, 1.0], + vec![-1.0, 0.0], + vec![0.0, -1.0], + vec![0.7, 0.3], + ], + )) + .await + .expect("Failed to write vector batch"); + let messages = table_write + .prepare_commit() + .await + .expect("Failed to prepare commit"); + write_builder + .new_commit() + .commit(messages) + .await + .expect("Failed to commit vector data"); + + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.vindex_order_e2e', \ + index_column => 'embedding', \ + index_type => 'ivf-flat', \ + options => 'ivf-flat.dimension=2,ivf-flat.nlist=1,ivf-flat.distance.metric=l2')", + ) + .await + .expect("vindex index build SQL should parse") + .collect() + .await + .expect("vindex index build SQL should execute"); + + let batches = ctx + .sql("SELECT id FROM vector_search('paimon.default.vindex_order_e2e', 'embedding', '[0.0, 1.0]', 3)") + .await + .expect("SQL should parse") + .collect() + .await + .expect("query should execute"); + + let ordered = extract_ids_in_order(&batches); + assert_eq!( + ordered, + vec![2, 5, 1], + "vector_search must return rows best-first by relevance, not in file order" + ); + + // The full-projection path (`SELECT *`, i.e. projection = None) must stay + // ordered too, and round-trip the other columns. + let star_batches = ctx + .sql("SELECT * FROM vector_search('paimon.default.vindex_order_e2e', 'embedding', '[0.0, 1.0]', 3)") + .await + .expect("SQL should parse") + .collect() + .await + .expect("query should execute"); + assert_eq!( + extract_ids_in_order(&star_batches), + vec![2, 5, 1], + "SELECT * must also be relevance-ordered" + ); + + // An outer LIMIT truncates the ranked result: top-2 of [2, 5, 1] is [2, 5]. + let limited = ctx + .sql("SELECT id FROM vector_search('paimon.default.vindex_order_e2e', 'embedding', '[0.0, 1.0]', 3) LIMIT 2") + .await + .expect("SQL should parse") + .collect() + .await + .expect("query should execute"); + assert_eq!(extract_ids_in_order(&limited), vec![2, 5]); + } + + /// Regression: a result containing a string column must not fail on the provider's + /// `Utf8View` schema vs the Paimon read's `Utf8` — the gather casts to the output + /// schema (the plain scan path did this via `to_datafusion_batch`). + #[tokio::test] + async fn test_vector_search_casts_string_column() { + let (ctx, catalog, _tmp) = create_empty_vector_search_context().await; + let identifier = Identifier::new("default", "vindex_string_e2e"); + + let mut options = std::collections::HashMap::new(); + options.insert("row-tracking.enabled".to_string(), "true".to_string()); + options.insert("data-evolution.enabled".to_string(), "true".to_string()); + options.insert("global-index.enabled".to_string(), "true".to_string()); + options.insert( + "global-index.row-count-per-shard".to_string(), + "3".to_string(), + ); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(64).unwrap())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) + .build() + .expect("Failed to build table schema"); + catalog + .create_table(&identifier, schema, false) + .await + .expect("Failed to create table"); + let table = catalog + .get_table(&identifier) + .await + .expect("Failed to load table"); + + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + for vector in [vec![1.0, 0.0], vec![0.0, 1.0], vec![0.7, 0.3]] { + for value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef, + Arc::new(StringArray::from(vec!["zero", "one", "two"])) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .expect("Failed to build batch"); + + let write_builder = table + .new_write_builder() + .with_commit_user("test-user") + .expect("Failed to configure write builder"); + let mut table_write = write_builder + .new_write() + .expect("Failed to create table write"); + table_write + .write_arrow_batch(&batch) + .await + .expect("Failed to write batch"); + let messages = table_write + .prepare_commit() + .await + .expect("Failed to prepare commit"); + write_builder + .new_commit() + .commit(messages) + .await + .expect("Failed to commit data"); + + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.vindex_string_e2e', \ + index_column => 'embedding', \ + index_type => 'ivf-flat', \ + options => 'ivf-flat.dimension=2,ivf-flat.nlist=1,ivf-flat.distance.metric=l2')", + ) + .await + .expect("vindex index build SQL should parse") + .collect() + .await + .expect("vindex index build SQL should execute"); + + // The point of this test is that the string `name` column round-trips without + // a Utf8/Utf8View schema-cast error (guarded by the successful collect below). + // Relevance ordering is covered by test_vector_search_orders_by_relevance, so + // here we only assert the hit set, order-independently. + let batches = ctx + .sql("SELECT id, name FROM vector_search('paimon.default.vindex_string_e2e', 'embedding', '[0.0, 1.0]', 2)") + .await + .expect("SQL should parse") + .collect() + .await + .expect("query with a string column should execute"); + let mut ids = extract_ids_in_order(&batches); + ids.sort(); + assert_eq!(ids, vec![1, 2]); + } + #[tokio::test] async fn test_vector_search_branch() { let (ctx, catalog, _tmp) = create_java_vindex_vector_search_context().await;