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

The `vector_search` table-valued function materialized its results with a
plain row-range scan, so `SELECT * FROM vector_search(...)` returned the
matching rows in physical file order instead of by relevance. This
contradicts the documented contract ("returns the top-k rows ordered by
relevance score") and means the first returned row is often not the most
similar one, especially at larger `limit`s.

Materialize the scored row-ids (best-first) via `execute_scored`, read the
matching rows (with projection pushdown, plus the internal `_ROW_ID`
column), then realign them to the index's relevance rank via a gather and
drop `_ROW_ID`. A scored row-id that is not materialized fails loud rather
than silently shrinking the top-k. The output schema is unchanged (no score
column exposed). `execute_read` is intentionally not reused: it is
statically `!Send` due to the primary-key-vector path, which the async
`TableProvider::scan` cannot hold.

Add a regression test that builds an ivf-flat vindex table and asserts the
returned order is `[2, 5, 1]` (relevance order), which differs from the
ascending-id file order `[1, 2, 5]` the old code produced, covering both the
projected (`SELECT id`) and full (`SELECT *`) paths.

Author: XiaoHongbo-Hope

crates/integrations/datafusion/src/vector_search.rs

@@ -18,22 +18,32 @@
 use std::fmt::Debug;
 use std::sync::Arc;
 
+use std::collections::HashMap;
+
 use async_trait::async_trait;
+use datafusion::arrow::array::{
+    Array, ArrayRef, Int64Array, RecordBatch, RecordBatchOptions, UInt32Array,
+};
 use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
 use datafusion::catalog::Session;
 use datafusion::catalog::TableFunctionImpl;
 use datafusion::common::project_schema;
+use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::{TableProvider, TableType};
 use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
 use paimon::catalog::Catalog;
+use paimon::spec::{
+    BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME,
+};
 
 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,51 +217,62 @@ impl TableProvider for VectorSearchTableProvider {
 
     async fn scan(
         &self,
-        state: &dyn Session,
+        _state: &dyn Session,
         projection: Option<&Vec<usize>>,
         _filters: &[Expr],
-        limit: Option<usize>,
+        _limit: Option<usize>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
         let table = self.inner.table();
+        let projected_schema = project_schema(&self.schema(), projection)?;
 
-        let row_ranges = await_with_runtime(async {
+        // Best-first row-ids from the index (data-evolution / global-index path;
+        // PK-vector tables are unsupported here, as before).
+        let search_result = await_with_runtime(async {
             let mut builder = 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)
+            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(Arc::new(EmptyExec::new(projected_schema)));
         }
 
-        let mut read_builder = table.new_read_builder();
-        if let Some(limit) = limit {
-            read_builder.with_limit(limit);
-        }
-        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,
-        }
-        .build()
+        // 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(table, projection)?;
+        let row_ranges = search_result.to_row_ranges().map_err(to_datafusion_error)?;
+        let batches = await_with_runtime(async {
+            let mut read_builder = 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` — this is what
+        // honors the documented "top-k rows ordered by relevance score" contract.
+        let ordered = gather_rows_by_rank(&batches, &search_result.row_ids, &projected_schema)?;
+
+        Ok(MemorySourceConfig::try_new_exec(
+            &[vec![ordered]],
+            projected_schema,
+            None,
+        )?)
     }
 
     fn supports_filters_pushdown(
@@ -264,3 +285,104 @@ impl TableProvider for VectorSearchTableProvider {
         ])
     }
 }
+
+/// 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()
+                ))
+            })?;
+            arrow_select::take::take(combined.column(index).as_ref(), &take_indices, None)
+                .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)
+}

crates/integrations/datafusion/tests/read_tables.rs

@@ -1877,6 +1877,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<i32> {
+        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::<Int32Array>())
+                .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 +1977,97 @@ 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"
+        );
+    }
+
     #[tokio::test]
     async fn test_vector_search_branch() {
         let (ctx, catalog, _tmp) = create_java_vindex_vector_search_context().await;