fix(datafusion): return vector_search rows in best-first relevance order (#613)

Author: XiaoHongbo-Hope

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<usize>>,
         _filters: &[Expr],
         limit: Option<usize>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
-        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<Vec<TableProviderFilterPushDown>> {
+        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<f32>,
+    /// 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<usize>,
+    projection: Option<Vec<usize>>,
+    output_schema: ArrowSchemaRef,
+    plan_properties: Arc<PlanProperties>,
+}
+
+impl VectorSearchExec {
+    fn new(
+        table: Table,
+        column_name: String,
+        query_vector: Vec<f32>,
+        search_limit: usize,
+        output_limit: Option<usize>,
+        projection: Option<Vec<usize>>,
+        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<RecordBatch> {
+        // 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<RecordBatch> = 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<PlanProperties> {
+        &self.plan_properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> DFResult<Arc<dyn ExecutionPlan>> {
+        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<Vec<TableProviderFilterPushDown>> {
-        Ok(vec![
-            TableProviderFilterPushDown::Unsupported;
-            filters.len()
-        ])
+        partition: usize,
+        _context: Arc<TaskContext>,
+    ) -> DFResult<SendableRecordBatchStream> {
+        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<usize>) -> DFResult<Arc<Statistics>> {
+        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<usize>>,
+) -> DFResult<Vec<DataField>> {
+    let base_fields = datafusion_read_fields(table);
+    let mut read_fields: Vec<DataField> = 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<RecordBatch> {
+    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::<Int64Array>()
+        .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<i64, u32> = 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<u32> = 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<ArrayRef> {
+            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::<DFResult<Vec<_>>>()?;
+
+    // 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)
 }