diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index a0a9bf54..e78d10fa 100644 --- a/crates/integrations/datafusion/Cargo.toml +++ b/crates/integrations/datafusion/Cargo.toml @@ -33,6 +33,7 @@ mosaic = ["paimon/mosaic"] vortex = ["paimon/vortex"] [dependencies] +arrow-select = { workspace = true } async-trait = "0.1" chrono = "0.4" constant_time_eq = { workspace = true } diff --git a/crates/integrations/datafusion/src/lateral_vector_search.rs b/crates/integrations/datafusion/src/lateral_vector_search.rs new file mode 100644 index 00000000..91a202a0 --- /dev/null +++ b/crates/integrations/datafusion/src/lateral_vector_search.rs @@ -0,0 +1,698 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::array::{ + new_empty_array, Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array, ListArray, + RecordBatch, UInt32Array, +}; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::catalog::default_table_source::source_as_provider; +use datafusion::common::stats::Precision; +use datafusion::common::tree_node::Transformed; +use datafusion::common::{ + internal_err, DFSchemaRef, DataFusionError, JoinType, Result as DFResult, Statistics, +}; +use datafusion::datasource::TableProvider; +use datafusion::execution::context::{QueryPlanner, SessionState}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, TableScan, UserDefinedLogicalNode}; +use datafusion::optimizer::{ApplyOrder, Optimizer, OptimizerConfig, OptimizerRule}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, +}; +use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; +use datafusion::prelude::SessionConfig; +use futures::{StreamExt, TryStreamExt}; +use paimon::spec::ROW_ID_FIELD_NAME; +use paimon::table::{RowRange, Table}; +use paimon::vector_search::SearchResult; + +use crate::error::to_datafusion_error; +use crate::vector_search::LateralVectorSearchTableProvider; + +#[derive(Debug)] +pub(crate) struct PaimonQueryPlanner; + +impl PaimonQueryPlanner { + pub(crate) fn new() -> Self { + Self + } +} + +#[async_trait] +impl QueryPlanner for PaimonQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session_state: &SessionState, + ) -> DFResult> { + let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + LateralVectorSearchExtensionPlanner, + )]); + planner + .create_physical_plan(logical_plan, session_state) + .await + } +} + +#[derive(Debug)] +pub(crate) struct RewriteLateralVectorSearch; + +impl RewriteLateralVectorSearch { + pub(crate) fn new() -> Self { + Self + } +} + +pub(crate) fn optimizer_rules() -> Vec> { + let mut rules: Vec> = + vec![Arc::new(RewriteLateralVectorSearch::new())]; + rules.extend(Optimizer::default().rules); + rules +} + +impl OptimizerRule for RewriteLateralVectorSearch { + fn name(&self) -> &str { + "rewrite_lateral_vector_search" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::BottomUp) + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> DFResult> { + let LogicalPlan::Join(join) = plan else { + return Ok(Transformed::no(plan)); + }; + + if join.join_type != JoinType::Inner || !join.on.is_empty() || join.filter.is_some() { + return Ok(Transformed::no(LogicalPlan::Join(join))); + } + + let Some(spec) = find_lateral_vector_search_provider(&join.right)? else { + return Ok(Transformed::no(LogicalPlan::Join(join))); + }; + + let node = LateralVectorSearchNode::new( + Arc::clone(&join.left), + spec.target_table, + spec.target_schema, + spec.target_column, + spec.query_vector_expr, + spec.limit, + Arc::clone(&join.schema), + ); + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(node), + }))) + } +} + +fn find_lateral_vector_search_provider( + plan: &LogicalPlan, +) -> DFResult> { + match plan { + LogicalPlan::TableScan(TableScan { source, .. }) => { + let provider = source_as_provider(source)?; + let Some(provider) = provider + .as_any() + .downcast_ref::() + else { + return Ok(None); + }; + Ok(Some(LateralVectorSearchSpec { + target_table: provider.inner().table().clone(), + target_schema: provider.inner().schema(), + target_column: provider.column_name().to_string(), + query_vector_expr: provider.query_vector_expr().clone(), + limit: provider.limit(), + })) + } + LogicalPlan::Subquery(subquery) => find_lateral_vector_search_provider(&subquery.subquery), + LogicalPlan::SubqueryAlias(alias) => find_lateral_vector_search_provider(&alias.input), + _ => Ok(None), + } +} + +struct LateralVectorSearchSpec { + target_table: Table, + target_schema: ArrowSchemaRef, + target_column: String, + query_vector_expr: Expr, + limit: usize, +} + +#[derive(Debug, Clone)] +pub(crate) struct LateralVectorSearchNode { + input: Arc, + target_table: Table, + target_schema: ArrowSchemaRef, + target_column: String, + query_vector_expr: Expr, + limit: usize, + schema: DFSchemaRef, +} + +impl LateralVectorSearchNode { + fn new( + input: Arc, + target_table: Table, + target_schema: ArrowSchemaRef, + target_column: String, + query_vector_expr: Expr, + limit: usize, + schema: DFSchemaRef, + ) -> Self { + Self { + input, + target_table, + target_schema, + target_column, + query_vector_expr, + limit, + schema, + } + } + + fn target_table(&self) -> &Table { + &self.target_table + } + + fn target_schema(&self) -> &ArrowSchemaRef { + &self.target_schema + } + + fn target_column(&self) -> &str { + &self.target_column + } + + fn query_vector_expr(&self) -> &Expr { + &self.query_vector_expr + } + + fn limit(&self) -> usize { + self.limit + } +} + +impl UserDefinedLogicalNode for LateralVectorSearchNode { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "LateralVectorSearch" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn check_invariants(&self, _check: datafusion::logical_expr::InvariantLevel) -> DFResult<()> { + Ok(()) + } + + fn expressions(&self) -> Vec { + vec![self.query_vector_expr.clone()] + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "LateralVectorSearch: column={}, limit={}", + self.target_column, self.limit + ) + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + inputs: Vec, + ) -> DFResult> { + if exprs.len() != 1 || inputs.len() != 1 { + return internal_err!("LateralVectorSearch expects one expression and one input"); + } + Ok(Arc::new(Self { + input: Arc::new(inputs.into_iter().next().unwrap()), + target_table: self.target_table.clone(), + target_schema: Arc::clone(&self.target_schema), + target_column: self.target_column.clone(), + query_vector_expr: exprs.into_iter().next().unwrap(), + limit: self.limit, + schema: Arc::clone(&self.schema), + })) + } + + fn dyn_hash(&self, mut state: &mut dyn Hasher) { + self.name().hash(&mut state); + self.input.hash(&mut state); + self.target_table.location().hash(&mut state); + self.target_column.hash(&mut state); + self.query_vector_expr.hash(&mut state); + self.limit.hash(&mut state); + } + + fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool { + other.as_any().downcast_ref::().is_some_and(|other| { + self.input == other.input + && self.target_table.location() == other.target_table.location() + && self.target_column == other.target_column + && self.query_vector_expr == other.query_vector_expr + && self.limit == other.limit + }) + } + + fn dyn_ord(&self, other: &dyn UserDefinedLogicalNode) -> Option { + let other = other.as_any().downcast_ref::()?; + if self.dyn_eq(other) { + Some(Ordering::Equal) + } else { + Some(format!("{self:?}").cmp(&format!("{other:?}"))) + } + } +} + +#[derive(Debug)] +struct LateralVectorSearchExtensionPlanner; + +#[async_trait] +impl ExtensionPlanner for LateralVectorSearchExtensionPlanner { + async fn plan_extension( + &self, + planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + session_state: &SessionState, + ) -> DFResult>> { + let Some(node) = node.as_any().downcast_ref::() else { + return Ok(None); + }; + if logical_inputs.len() != 1 || physical_inputs.len() != 1 { + return internal_err!("LateralVectorSearch physical planning expects one input"); + } + + let query_vector_expr = planner.create_physical_expr( + node.query_vector_expr(), + logical_inputs[0].schema(), + session_state, + )?; + Ok(Some(Arc::new(LateralVectorSearchExec::new( + Arc::clone(&physical_inputs[0]), + node.target_table().clone(), + Arc::clone(node.target_schema()), + node.target_column().to_string(), + query_vector_expr, + node.limit(), + Arc::new(node.schema().as_arrow().clone()), + )))) + } +} + +#[derive(Debug, Clone)] +struct LateralVectorSearchExec { + input: Arc, + target_table: Table, + target_schema: ArrowSchemaRef, + target_column: String, + query_vector_expr: Arc, + limit: usize, + output_schema: ArrowSchemaRef, + plan_properties: Arc, +} + +impl LateralVectorSearchExec { + fn new( + input: Arc, + target_table: Table, + target_schema: ArrowSchemaRef, + target_column: String, + query_vector_expr: Arc, + limit: usize, + output_schema: ArrowSchemaRef, + ) -> Self { + let partition_count = input.output_partitioning().partition_count(); + let plan_properties = Arc::new(PlanProperties::new( + datafusion::physical_expr::EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + input, + target_table, + target_schema, + target_column, + query_vector_expr, + limit, + output_schema, + plan_properties, + } + } + + async fn process_batch(&self, batch: RecordBatch) -> DFResult { + if batch.num_rows() == 0 { + return empty_batch(self.output_schema.clone()); + } + + let vector_array = self + .query_vector_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?; + let (query_vectors, left_query_rows) = collect_query_vectors(&vector_array)?; + if query_vectors.is_empty() { + return empty_batch(self.output_schema.clone()); + } + + let mut builder = self.target_table.new_batch_vector_search_builder(); + let results = builder + .with_vector_column(&self.target_column) + .with_query_vectors(query_vectors) + .with_limit(self.limit) + .execute() + .await + .map_err(to_datafusion_error)?; + + let (target_batch, target_row_id_to_index) = + read_target_rows(&self.target_table, &self.target_schema, &results).await?; + + let mut left_indices = Vec::new(); + let mut right_indices = Vec::new(); + for (query_index, result) in results.iter().enumerate() { + let left_row = left_query_rows[query_index] as u32; + for row_id in &result.row_ids { + if let Some(&right_row) = target_row_id_to_index.get(row_id) { + left_indices.push(left_row); + right_indices.push(right_row); + } + } + } + + if left_indices.is_empty() { + return empty_batch(self.output_schema.clone()); + } + + let left_indices = UInt32Array::from(left_indices); + let right_indices = UInt32Array::from(right_indices); + let mut columns = Vec::with_capacity(batch.num_columns() + target_batch.num_columns()); + for column in batch.columns() { + columns.push(arrow_select::take::take( + column.as_ref(), + &left_indices, + None, + )?); + } + for column in target_batch.columns() { + columns.push(arrow_select::take::take( + column.as_ref(), + &right_indices, + None, + )?); + } + + RecordBatch::try_new(self.output_schema.clone(), columns).map_err(DataFusionError::from) + } +} + +impl DisplayAs for LateralVectorSearchExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "LateralVectorSearchExec: column={}, limit={}", + self.target_column, self.limit + ) + } +} + +impl ExecutionPlan for LateralVectorSearchExec { + fn name(&self) -> &str { + "LateralVectorSearchExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + &self.plan_properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DFResult> { + if children.len() != 1 { + return internal_err!("LateralVectorSearchExec expects one child"); + } + Ok(Arc::new(Self::new( + children.remove(0), + self.target_table.clone(), + Arc::clone(&self.target_schema), + self.target_column.clone(), + Arc::clone(&self.query_vector_expr), + self.limit, + Arc::clone(&self.output_schema), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DFResult { + let input = self.input.execute(partition, context)?; + let exec = self.clone(); + let stream = input.then(move |batch| { + let exec = exec.clone(); + async move { + let batch = batch?; + exec.process_batch(batch).await + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.output_schema.clone(), + Box::pin(stream), + ))) + } + + fn partition_statistics(&self, _partition: Option) -> DFResult { + Ok(Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: Statistics::unknown_column(&self.output_schema), + }) + } +} + +fn collect_query_vectors(array: &ArrayRef) -> DFResult<(Vec>, Vec)> { + enum VectorLayout<'a> { + List(&'a ListArray), + Fixed(&'a FixedSizeListArray), + } + let layout = if let Some(array) = array.as_any().downcast_ref::() { + VectorLayout::List(array) + } else if let Some(array) = array.as_any().downcast_ref::() { + VectorLayout::Fixed(array) + } else { + return Err(DataFusionError::Plan( + "lateral vector_search query vector must be List or FixedSizeList" + .to_string(), + )); + }; + let values = match layout { + VectorLayout::List(array) => array.values(), + VectorLayout::Fixed(array) => array.values(), + } + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan( + "lateral vector_search query vector elements must be Float32".to_string(), + ) + })?; + + let row_count = match layout { + VectorLayout::List(array) => array.len(), + VectorLayout::Fixed(array) => array.len(), + }; + let mut vectors = Vec::new(); + let mut rows = Vec::new(); + for row in 0..row_count { + let is_null = match layout { + VectorLayout::List(array) => array.is_null(row), + VectorLayout::Fixed(array) => array.is_null(row), + }; + if is_null { + continue; + } + + let (start, end) = match layout { + VectorLayout::List(array) => { + let offsets = array.value_offsets(); + (offsets[row] as usize, offsets[row + 1] as usize) + } + VectorLayout::Fixed(array) => { + let len = array.value_length() as usize; + (row * len, (row + 1) * len) + } + }; + let mut vector = Vec::with_capacity(end - start); + for value_index in start..end { + if values.is_null(value_index) { + return Err(DataFusionError::Plan( + "lateral vector_search query vector cannot contain null elements".to_string(), + )); + } + vector.push(values.value(value_index)); + } + vectors.push(vector); + rows.push(row); + } + Ok((vectors, rows)) +} + +async fn read_target_rows( + table: &Table, + target_schema: &ArrowSchemaRef, + results: &[SearchResult], +) -> DFResult<(RecordBatch, HashMap)> { + let mut row_ids = results + .iter() + .flat_map(|result| result.row_ids.iter().copied()) + .collect::>(); + row_ids.sort_unstable(); + row_ids.dedup(); + if row_ids.is_empty() { + return Ok((empty_batch(target_schema.clone())?, HashMap::new())); + } + + let row_ranges = row_ranges_from_row_ids(&row_ids)?; + let mut projection = target_schema + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect::>(); + if !projection.iter().any(|column| column == ROW_ID_FIELD_NAME) { + projection.push(ROW_ID_FIELD_NAME.to_string()); + } + let projection_refs = projection.iter().map(String::as_str).collect::>(); + + let mut read_builder = table.new_read_builder(); + read_builder + .with_projection(&projection_refs) + .with_row_ranges(row_ranges); + let plan = read_builder + .new_scan() + .plan() + .await + .map_err(to_datafusion_error)?; + if plan.splits().is_empty() { + return Ok((empty_batch(target_schema.clone())?, HashMap::new())); + } + + let read = read_builder.new_read().map_err(to_datafusion_error)?; + let mut stream = read.to_arrow(plan.splits()).map_err(to_datafusion_error)?; + let mut batches = Vec::new(); + while let Some(batch) = stream.try_next().await.map_err(to_datafusion_error)? { + batches.push(batch); + } + if batches.is_empty() { + return Ok((empty_batch(target_schema.clone())?, HashMap::new())); + } + + let read_schema = batches[0].schema(); + let batch = arrow_select::concat::concat_batches(&read_schema, &batches) + .map_err(DataFusionError::from)?; + let row_id_index = batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(DataFusionError::from)?; + let row_id_array = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal( + "_ROW_ID must be Int64 in vector search target read".to_string(), + ) + })?; + + let mut row_id_to_index = HashMap::new(); + for row in 0..batch.num_rows() { + if row_id_array.is_null(row) { + continue; + } + let row_id = u64::try_from(row_id_array.value(row)).map_err(|_| { + DataFusionError::Internal(format!( + "negative _ROW_ID {} in vector search target read", + row_id_array.value(row) + )) + })?; + row_id_to_index.insert(row_id, row as u32); + } + + let target_columns = (0..target_schema.fields().len()) + .map(|index| Arc::clone(batch.column(index))) + .collect::>(); + let target_batch = RecordBatch::try_new(target_schema.clone(), target_columns) + .map_err(DataFusionError::from)?; + Ok((target_batch, row_id_to_index)) +} + +fn row_ranges_from_row_ids(row_ids: &[u64]) -> DFResult> { + let scores = vec![0.0; row_ids.len()]; + SearchResult::new(row_ids.to_vec(), scores) + .to_row_ranges() + .map_err(to_datafusion_error) +} + +fn empty_batch(schema: ArrowSchemaRef) -> DFResult { + let columns = schema + .fields() + .iter() + .map(|field| new_empty_array(field.data_type())) + .collect::>(); + RecordBatch::try_new(schema, columns).map_err(DataFusionError::from) +} + +pub(crate) fn session_config() -> SessionConfig { + SessionConfig::new().with_information_schema(true) +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index f11cfce4..de6e9328 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -43,6 +43,7 @@ mod error; mod filter_pushdown; #[cfg(feature = "fulltext")] mod full_text_search; +mod lateral_vector_search; mod merge_into; mod physical_plan; mod procedures; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 515cead1..8ccdc501 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -46,7 +46,8 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::TableReference; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; -use datafusion::prelude::{DataFrame, SessionConfig, SessionContext}; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::sqlparser::ast::{ AlterTableOperation, ColumnDef, CreateTable, CreateTableOptions, CreateView, Delete, Expr as SqlExpr, FromTable, Insert, Merge, ObjectName, ObjectType, RenameTableNameKind, Reset, @@ -93,12 +94,18 @@ impl Default for SQLContext { impl SQLContext { /// Creates a new empty SQL context. pub fn new() -> Self { - let ctx = - SessionContext::new_with_config(SessionConfig::new().with_information_schema(true)); - ctx.register_relation_planner(Arc::new( - crate::relation_planner::PaimonRelationPlanner::new(), - )) - .expect("failed to register relation planner"); + let state = SessionStateBuilder::new() + .with_config(crate::lateral_vector_search::session_config()) + .with_default_features() + .with_relation_planners(vec![Arc::new( + crate::relation_planner::PaimonRelationPlanner::new(), + )]) + .with_optimizer_rules(crate::lateral_vector_search::optimizer_rules()) + .with_query_planner(Arc::new( + crate::lateral_vector_search::PaimonQueryPlanner::new(), + )) + .build(); + let ctx = SessionContext::new_with_state(state); Self { ctx, catalogs: HashMap::new(), diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index 5b53d146..536e3795 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -25,7 +25,7 @@ use datafusion::catalog::Session; use datafusion::catalog::TableFunctionImpl; use datafusion::common::project_schema; use datafusion::datasource::{TableProvider, TableType}; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::ExecutionPlan; @@ -84,29 +84,14 @@ impl TableFunctionImpl for VectorSearchFunction { let table_name = extract_string_literal(FUNCTION_NAME, &args[0], "table_name")?; let column_name = extract_string_literal(FUNCTION_NAME, &args[1], "column_name")?; - let query_vector_json = - extract_string_literal(FUNCTION_NAME, &args[2], "query_vector_json")?; let limit = extract_int_literal(FUNCTION_NAME, &args[3], "limit")?; if limit <= 0 { - return Err(datafusion::error::DataFusionError::Plan( + return Err(DataFusionError::Plan( "vector_search: limit must be positive".to_string(), )); } - let query_vector: Vec = serde_json::from_str(&query_vector_json).map_err(|e| { - datafusion::error::DataFusionError::Plan(format!( - "vector_search: query_vector_json must be a JSON array of floats, got '{}': {}", - query_vector_json, e - )) - })?; - - if query_vector.is_empty() { - return Err(datafusion::error::DataFusionError::Plan( - "vector_search: query vector cannot be empty".to_string(), - )); - } - let identifier = parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?; @@ -118,6 +103,32 @@ impl TableFunctionImpl for VectorSearchFunction { .map_err(to_datafusion_error)?; let inner = PaimonTableProvider::try_new(table)?; + let query_vector_json = + match extract_string_literal(FUNCTION_NAME, &args[2], "query_vector_json") { + Ok(value) => value, + Err(_) if matches!(args[2], Expr::Column(_)) => { + return Ok(Arc::new(LateralVectorSearchTableProvider { + inner, + column_name, + query_vector_expr: args[2].clone(), + limit: limit as usize, + })); + } + Err(err) => return Err(err), + }; + + let query_vector: Vec = serde_json::from_str(&query_vector_json).map_err(|e| { + DataFusionError::Plan(format!( + "vector_search: query_vector_json must be a JSON array of floats, got '{}': {}", + query_vector_json, e + )) + })?; + + if query_vector.is_empty() { + return Err(DataFusionError::Plan( + "vector_search: query vector cannot be empty".to_string(), + )); + } Ok(Arc::new(VectorSearchTableProvider { inner, @@ -128,6 +139,59 @@ impl TableFunctionImpl for VectorSearchFunction { } } +#[derive(Debug)] +pub(crate) struct LateralVectorSearchTableProvider { + inner: PaimonTableProvider, + column_name: String, + query_vector_expr: Expr, + limit: usize, +} + +impl LateralVectorSearchTableProvider { + pub(crate) fn inner(&self) -> &PaimonTableProvider { + &self.inner + } + + pub(crate) fn column_name(&self) -> &str { + &self.column_name + } + + pub(crate) fn query_vector_expr(&self) -> &Expr { + &self.query_vector_expr + } + + pub(crate) fn limit(&self) -> usize { + self.limit + } +} + +#[async_trait] +impl TableProvider for LateralVectorSearchTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> ArrowSchemaRef { + self.inner.schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> DFResult> { + Err(DataFusionError::Plan( + "lateral vector_search must be planned through a lateral join".to_string(), + )) + } +} + #[derive(Debug)] struct VectorSearchTableProvider { inner: PaimonTableProvider, diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index cc4c781a..1d723804 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1343,6 +1343,7 @@ mod vector_search_tests { 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::{Catalog, CatalogOptions, FileSystemCatalog, Options}; @@ -1467,6 +1468,27 @@ mod vector_search_tests { ids } + fn extract_query_result_ids( + batches: &[datafusion::arrow::record_batch::RecordBatch], + ) -> Vec<(i32, i32)> { + let mut rows = Vec::new(); + for batch in batches { + let query_id_array = batch + .column_by_name("query_id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("Expected Int32Array for query_id"); + let result_id_array = batch + .column_by_name("result_id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("Expected Int32Array for result_id"); + for i in 0..batch.num_rows() { + rows.push((query_id_array.value(i), result_id_array.value(i))); + } + } + rows.sort(); + rows + } + fn extract_index_rows( batches: &[datafusion::arrow::record_batch::RecordBatch], ) -> Vec<(String, i64, i64, i64, String)> { @@ -1586,6 +1608,35 @@ mod vector_search_tests { assert_eq!(ids, vec![0, 1, 2]); } + #[tokio::test] + async fn test_vector_search_lateral_join_uses_query_vectors() { + let (ctx, _tmp) = create_java_vindex_vector_search_context().await; + let query_batch = build_vector_batch( + vec![10, 20], + vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]], + ); + let query_table = MemTable::try_new(query_batch.schema(), vec![vec![query_batch]]) + .expect("Failed to create query vector table"); + ctx.register_temp_table("paimon.default.queries", Arc::new(query_table)) + .expect("Failed to register query vector table"); + + let batches = ctx + .sql( + "SELECT q.id AS query_id, r.id AS result_id \ + FROM paimon.default.queries q \ + CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 2) AS r \ + ORDER BY query_id, result_id", + ) + .await + .expect("lateral vector_search SQL should parse") + .collect() + .await + .expect("lateral vector_search query should execute"); + + let rows = extract_query_result_ids(&batches); + assert_eq!(rows, vec![(10, 0), (10, 1), (20, 1), (20, 2)]); + } + // Manual run with a local Lumina native library: // LUMINA_LIB_PATH=/path/to/liblumina_py.so cargo test -p paimon-datafusion \ // vector_search_tests::test_lumina_build_then_vector_search_query \