This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 885dd29c fix(datafusion): return vector_search rows in best-first
relevance order (#613)
885dd29c is described below
commit 885dd29ceee509024f75674761fe362842ced57c
Author: XiaoHongbo <[email protected]>
AuthorDate: Mon Jul 27 09:59:53 2026 +0800
fix(datafusion): return vector_search rows in best-first relevance order
(#613)
---
.../integrations/datafusion/src/vector_search.rs | 351 ++++++++++++++++++---
.../integrations/datafusion/tests/read_tables.rs | 232 +++++++++++++-
2 files changed, 541 insertions(+), 42 deletions(-)
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<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)
}
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<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 +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;
