This is an automated email from the ASF dual-hosted git repository.
jerry-024 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 8c4c7943 datafusion: cast lateral vector search rows to provider
schema (#781)
8c4c7943 is described below
commit 8c4c79434d823ab12967fe72764db0c746b53265
Author: shyjsarah <[email protected]>
AuthorDate: Wed Sep 2 14:47:13 2026 +0800
datafusion: cast lateral vector search rows to provider schema (#781)
---
.../datafusion/src/lateral_vector_search.rs | 22 ++++++++++--
.../integrations/datafusion/tests/read_tables.rs | 42 ++++++++++++++++++++--
2 files changed, 58 insertions(+), 6 deletions(-)
diff --git a/crates/integrations/datafusion/src/lateral_vector_search.rs
b/crates/integrations/datafusion/src/lateral_vector_search.rs
index 8e88c538..2500d8b5 100644
--- a/crates/integrations/datafusion/src/lateral_vector_search.rs
+++ b/crates/integrations/datafusion/src/lateral_vector_search.rs
@@ -31,6 +31,7 @@ use datafusion::arrow::array::{
new_empty_array, Array, ArrayRef, FixedSizeListArray, Float32Array,
Int64Array, ListArray,
RecordBatch, UInt32Array,
};
+use datafusion::arrow::compute::cast;
use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datafusion::catalog::default_table_source::source_as_provider;
use datafusion::common::stats::Precision;
@@ -1066,9 +1067,24 @@ async fn read_target_rows(
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::<Vec<_>>();
+ // Paimon reads use storage Arrow types (for example `Utf8`), while the
+ // DataFusion provider schema may expose view types such as `Utf8View`.
+ let target_columns = target_schema
+ .fields()
+ .iter()
+ .map(|field| -> DFResult<ArrayRef> {
+ let index = batch
+ .schema()
+ .index_of(field.name())
+ .map_err(DataFusionError::from)?;
+ let column = batch.column(index);
+ if column.data_type() == field.data_type() {
+ Ok(Arc::clone(column))
+ } else {
+ cast(column.as_ref(),
field.data_type()).map_err(DataFusionError::from)
+ }
+ })
+ .collect::<DFResult<Vec<_>>>()?;
let target_batch = RecordBatch::try_new(target_schema.clone(),
target_columns)
.map_err(DataFusionError::from)?;
Ok((target_batch, row_id_to_index))
diff --git a/crates/integrations/datafusion/tests/read_tables.rs
b/crates/integrations/datafusion/tests/read_tables.rs
index 2226b18c..e2d9189b 100644
--- a/crates/integrations/datafusion/tests/read_tables.rs
+++ b/crates/integrations/datafusion/tests/read_tables.rs
@@ -2082,9 +2082,8 @@ mod vector_search_tests {
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`).
+ /// Regression: direct and lateral vector search results containing a
string column
+ /// must align the Paimon read's `Utf8` arrays with the provider's
`Utf8View` schema.
#[tokio::test]
async fn test_vector_search_casts_string_column() {
let (ctx, catalog, _tmp) = create_empty_vector_search_context().await;
@@ -2189,6 +2188,43 @@ mod vector_search_tests {
let mut ids = extract_ids_in_order(&batches);
ids.sort();
assert_eq!(ids, vec![1, 2]);
+
+ let query_batch = build_vector_batch(vec![10], vec![vec![0.0, 1.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.string_queries",
Arc::new(query_table))
+ .expect("Failed to register query vector table");
+
+ let lateral_batches = ctx
+ .sql(
+ "SELECT r.id, r.name \
+ FROM paimon.default.string_queries q \
+ CROSS JOIN LATERAL vector_search( \
+ 'paimon.default.vindex_string_e2e', \
+ 'embedding', \
+ q.embedding, \
+ 2 \
+ ) AS r",
+ )
+ .await
+ .expect("lateral vector search SQL should parse")
+ .collect()
+ .await
+ .expect("lateral vector search with a string column should
execute");
+ let mut lateral_ids = extract_ids_in_order(&lateral_batches);
+ lateral_ids.sort();
+ assert_eq!(lateral_ids, vec![1, 2]);
+ let mut lateral_names = lateral_batches
+ .iter()
+ .flat_map(|batch| {
+ let names = batch.column_by_name("name").expect("Expected name
column");
+ (0..batch.num_rows())
+ .map(|row| string_value(names.as_ref(), row).to_string())
+ .collect::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
+ lateral_names.sort();
+ assert_eq!(lateral_names, vec!["one", "two"]);
}
#[tokio::test]
