From 3866a45f33ae88165a4c1b77490ddd3d26708823 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Mon, 13 Jul 2026 17:26:33 +0800 Subject: [PATCH] feat(datafusion): expose hybrid search scores DataFusion hybrid search discarded fused scores, preventing queries from selecting the Spark-compatible search-score metadata column. Expose __paimon_search_score and materialize it by matching scan rows through _ROW_ID while preserving projection, empty-result, and row-tracking-only behavior. Signed-off-by: QuakeWang --- .../datafusion/src/hybrid_search.rs | 168 +++++++++-- .../datafusion/src/physical_plan/mod.rs | 2 + .../src/physical_plan/search_score.rs | 273 ++++++++++++++++++ .../integrations/datafusion/src/table/mod.rs | 8 + .../datafusion/tests/read_tables.rs | 243 +++++++++++++++- docs/src/sql.md | 21 +- 6 files changed, 691 insertions(+), 24 deletions(-) create mode 100644 crates/integrations/datafusion/src/physical_plan/search_score.rs diff --git a/crates/integrations/datafusion/src/hybrid_search.rs b/crates/integrations/datafusion/src/hybrid_search.rs index f5c0e7a3..b37e4caa 100644 --- a/crates/integrations/datafusion/src/hybrid_search.rs +++ b/crates/integrations/datafusion/src/hybrid_search.rs @@ -33,7 +33,9 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::array::Array; -use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema, SchemaRef as ArrowSchemaRef, +}; use datafusion::catalog::{Session, TableFunctionImpl}; use datafusion::common::{project_schema, ScalarValue}; use datafusion::datasource::{TableProvider, TableType}; @@ -43,17 +45,22 @@ use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use paimon::catalog::Catalog; -use paimon::table::{HybridSearchRanker, HybridSearchRoute}; +use paimon::spec::{ + BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, +}; +use paimon::table::{HybridSearchRanker, HybridSearchRoute, Table}; use crate::error::to_datafusion_error; +use crate::physical_plan::{SearchScoreExec, SearchScoreOutputColumn}; use crate::runtime::{await_with_runtime, block_on_with_runtime}; -use crate::table::{PaimonScanBuilder, PaimonTableProvider}; +use crate::table::{datafusion_read_fields, PaimonScanBuilder, PaimonTableProvider}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; use crate::table_loader::load_data_table_for_read; const FUNCTION_NAME: &str = "hybrid_search"; +const SEARCH_SCORE_COLUMN: &str = "__paimon_search_score"; pub fn register_hybrid_search( ctx: &SessionContext, @@ -128,27 +135,71 @@ impl TableFunctionImpl for HybridSearchFunction { "hybrid_search: catalog access thread panicked", )?; - Ok(Arc::new(HybridSearchTableProvider { - inner: PaimonTableProvider::try_new(table)?, + Ok(Arc::new(HybridSearchTableProvider::try_new( + PaimonTableProvider::try_new(table)?, routes, - limit: limit as usize, + limit as usize, ranker, - })) + )?)) } } #[derive(Debug)] struct HybridSearchTableProvider { inner: PaimonTableProvider, + schema: ArrowSchemaRef, routes: Vec, limit: usize, ranker: String, } +impl HybridSearchTableProvider { + fn try_new( + inner: PaimonTableProvider, + routes: Vec, + limit: usize, + ranker: String, + ) -> DFResult { + let inner_schema = inner.schema(); + if inner_schema + .fields() + .iter() + .any(|field| field.name() == SEARCH_SCORE_COLUMN) + { + return Err(DataFusionError::Plan(format!( + "hybrid_search: table already contains reserved column {SEARCH_SCORE_COLUMN}" + ))); + } + + let mut fields = inner_schema + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect::>(); + fields.push(Field::new( + SEARCH_SCORE_COLUMN, + ArrowDataType::Float32, + true, + )); + let schema = Arc::new(Schema::new_with_metadata( + fields, + inner_schema.metadata().clone(), + )); + + Ok(Self { + inner, + schema, + routes, + limit, + ranker, + }) + } +} + #[async_trait] impl TableProvider for HybridSearchTableProvider { fn schema(&self) -> ArrowSchemaRef { - self.inner.schema() + self.schema.clone() } fn table_type(&self) -> TableType { @@ -160,11 +211,11 @@ impl TableProvider for HybridSearchTableProvider { state: &dyn Session, projection: Option<&Vec>, _filters: &[Expr], - limit: Option, + _limit: Option, ) -> DFResult> { let table = self.inner.table(); - let row_ranges = await_with_runtime(async { + let search_result = await_with_runtime(async { let mut builder = table.new_hybrid_search_builder(); for route in self.routes.clone() { builder.add_route(route); @@ -173,37 +224,91 @@ impl TableProvider for HybridSearchTableProvider { .with_limit(self.limit) .with_ranker(&self.ranker) .map_err(to_datafusion_error)?; - builder.execute().await.map_err(to_datafusion_error) + builder.execute_scored().await.map_err(to_datafusion_error) }) .await?; + let row_ranges = search_result.to_row_ranges().map_err(to_datafusion_error)?; + if row_ranges.is_empty() { - let schema = project_schema(&self.schema(), projection)?; + let schema = project_schema(&self.schema, projection)?; return Ok(Arc::new(EmptyExec::new(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); + // Raw row-range scans can include non-matching rows from selected files, + // so the outer limit must remain above the row-ID filter. + let scan = table + .new_read_builder() + .new_scan() + .with_row_ranges(row_ranges); let plan = await_with_runtime(scan.plan()) .await .map_err(to_datafusion_error)?; - PaimonScanBuilder { + let inner_schema = self.inner.schema(); + let score_index = inner_schema.fields().len(); + let input_read_fields = search_read_fields(table)?; + let input_schema = paimon::arrow::build_target_arrow_schema(&input_read_fields) + .map_err(to_datafusion_error)?; + let row_id_table_index = input_read_fields + .iter() + .position(|field| field.name() == ROW_ID_FIELD_NAME) + .expect("search read fields contain _ROW_ID"); + let projected_indices = projection + .cloned() + .unwrap_or_else(|| (0..self.schema.fields().len()).collect()); + let mut input_projection = projected_indices + .iter() + .copied() + .filter(|index| *index != score_index) + .collect::>(); + if !input_projection.contains(&row_id_table_index) { + input_projection.push(row_id_table_index); + } + let row_id_input_index = input_projection + .iter() + .position(|index| *index == row_id_table_index) + .expect("row ID was added to the input projection"); + let output_columns = projected_indices + .iter() + .map(|index| { + if *index == score_index { + SearchScoreOutputColumn::Score + } else { + let input_index = input_projection + .iter() + .position(|input_index| input_index == index) + .expect("projected table column exists in the input projection"); + SearchScoreOutputColumn::Input(input_index) + } + }) + .collect(); + let output_schema = project_schema(&self.schema, projection)?; + let input = PaimonScanBuilder { table, - schema: &self.schema(), + schema: &input_schema, plan: &plan, scan_trace: None, - projection, + projection: Some(&input_projection), pushed_predicate: None, - limit, + limit: None, target_partitions: state.config_options().execution.target_partitions, filter_exact: false, case_sensitive: true, } - .build() + .build_with_read_fields(input_read_fields)?; + let scores = search_result + .row_ids + .into_iter() + .zip(search_result.scores) + .collect(); + Ok(Arc::new(SearchScoreExec::new( + input, + output_schema, + row_id_input_index, + output_columns, + Arc::new(scores), + ))) } fn supports_filters_pushdown( @@ -217,6 +322,25 @@ impl TableProvider for HybridSearchTableProvider { } } +fn search_read_fields(table: &Table) -> DFResult> { + let mut fields = datafusion_read_fields(table); + if fields.iter().any(|field| field.name() == ROW_ID_FIELD_NAME) { + return Ok(fields); + } + if !CoreOptions::new(table.schema().options()).row_tracking_enabled() { + return Err(DataFusionError::Plan( + "hybrid_search: cannot materialize search results because _ROW_ID is not available" + .to_string(), + )); + } + fields.push(DataField::new( + ROW_ID_FIELD_ID, + ROW_ID_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::with_nullable(true)), + )); + Ok(fields) +} + fn parse_vector_routes(expr: &Expr, default_limit: usize) -> DFResult> { if let Some(routes) = extract_literal_array_values(expr, "vector_routes")? { return routes diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index 2fa35bf1..2d190503 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/datafusion/src/physical_plan/mod.rs @@ -16,7 +16,9 @@ // under the License. pub(crate) mod scan; +mod search_score; pub(crate) mod sink; pub use scan::PaimonTableScan; +pub(crate) use search_score::{SearchScoreExec, SearchScoreOutputColumn}; pub use sink::PaimonDataSink; diff --git a/crates/integrations/datafusion/src/physical_plan/search_score.rs b/crates/integrations/datafusion/src/physical_plan/search_score.rs new file mode 100644 index 00000000..39d01ced --- /dev/null +++ b/crates/integrations/datafusion/src/physical_plan/search_score.rs @@ -0,0 +1,273 @@ +// 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::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, Float32Array, Int64Array, RecordBatch, RecordBatchOptions, UInt32Array, +}; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::common::stats::Precision; +use datafusion::common::{internal_err, DataFusionError, Statistics}; +use datafusion::error::Result as DFResult; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::EquivalenceProperties; +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 futures::StreamExt; + +#[derive(Clone, Copy, Debug)] +pub(crate) enum SearchScoreOutputColumn { + Input(usize), + Score, +} + +/// Filters Paimon rows by search result row ID and optionally exposes their scores. +#[derive(Debug, Clone)] +pub(crate) struct SearchScoreExec { + input: Arc, + output_schema: ArrowSchemaRef, + row_id_index: usize, + output_columns: Arc<[SearchScoreOutputColumn]>, + scores: Arc>, + plan_properties: Arc, +} + +impl SearchScoreExec { + pub(crate) fn new( + input: Arc, + output_schema: ArrowSchemaRef, + row_id_index: usize, + output_columns: Vec, + scores: Arc>, + ) -> Self { + let partition_count = input.output_partitioning().partition_count(); + let plan_properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + input, + output_schema, + row_id_index, + output_columns: output_columns.into(), + scores, + plan_properties, + } + } + + fn filter_and_project(&self, batch: RecordBatch) -> DFResult { + let row_ids = batch + .column(self.row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal( + "_ROW_ID must be Int64 when materializing hybrid search results".to_string(), + ) + })?; + + let mut input_indices = Vec::with_capacity(batch.num_rows()); + let mut scores = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return internal_err!("_ROW_ID cannot be null in hybrid search results"); + } + let row_id = u64::try_from(row_ids.value(row)).map_err(|_| { + DataFusionError::Internal(format!( + "negative _ROW_ID {} in hybrid search results", + row_ids.value(row) + )) + })?; + if let Some(score) = self.scores.get(&row_id) { + input_indices.push(row as u32); + scores.push(*score); + } + } + let matched_row_count = input_indices.len(); + let scores: ArrayRef = Arc::new(Float32Array::from(scores)); + // Raw row-range scans may conservatively return non-matching rows from a + // selected file. The scored row IDs are the authoritative result set. + let input_indices = + (input_indices.len() != batch.num_rows()).then(|| UInt32Array::from(input_indices)); + + let columns = self + .output_columns + .iter() + .map(|column| -> DFResult { + match column { + SearchScoreOutputColumn::Input(index) => match &input_indices { + Some(indices) => Ok(arrow_select::take::take( + batch.column(*index).as_ref(), + indices, + None, + )?), + None => Ok(Arc::clone(batch.column(*index))), + }, + SearchScoreOutputColumn::Score => Ok(Arc::clone(&scores)), + } + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(matched_row_count)); + RecordBatch::try_new_with_options(self.output_schema.clone(), columns, &options) + .map_err(DataFusionError::from) + } +} + +impl DisplayAs for SearchScoreExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "SearchScoreExec") + } +} + +impl ExecutionPlan for SearchScoreExec { + fn name(&self) -> &str { + "SearchScoreExec" + } + + 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!("SearchScoreExec expects one child"); + } + Ok(Arc::new(Self::new( + children.remove(0), + Arc::clone(&self.output_schema), + self.row_id_index, + self.output_columns.to_vec(), + Arc::clone(&self.scores), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DFResult { + let input = self.input.execute(partition, context)?; + let exec = self.clone(); + let stream = input.map(move |batch| batch.and_then(|batch| exec.filter_and_project(batch))); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.output_schema.clone(), + Box::pin(stream), + ))) + } + + fn partition_statistics(&self, _partition: Option) -> DFResult> { + Ok(Arc::new(Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: Statistics::unknown_column(&self.output_schema), + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::empty::EmptyExec; + + #[test] + fn test_filter_and_project_by_row_id_filters_non_matches_and_reorders_columns() { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("_ROW_ID", DataType::Int64, true), + ])); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("__paimon_search_score", DataType::Float32, true), + Field::new("id", DataType::Int32, false), + ])); + let exec = SearchScoreExec::new( + Arc::new(EmptyExec::new(input_schema.clone())), + output_schema, + 1, + vec![ + SearchScoreOutputColumn::Score, + SearchScoreOutputColumn::Input(0), + ], + Arc::new(HashMap::from([(7, 0.25), (3, 0.75)])), + ); + let batch = RecordBatch::try_new( + input_schema, + vec![ + Arc::new(Int32Array::from(vec![70, 40, 30])), + Arc::new(Int64Array::from(vec![7, 4, 3])), + ], + ) + .expect("input batch"); + + let output = exec.filter_and_project(batch).expect("filter and project"); + let scores = output + .column(0) + .as_any() + .downcast_ref::() + .expect("score column"); + let ids = output + .column(1) + .as_any() + .downcast_ref::() + .expect("id column"); + assert_eq!(scores.values(), &[0.25, 0.75]); + assert_eq!(ids.values(), &[70, 30]); + } + + #[test] + fn test_filter_and_project_preserves_row_count_without_output_columns() { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "_ROW_ID", + DataType::Int64, + true, + )])); + let output_schema = Arc::new(Schema::empty()); + let exec = SearchScoreExec::new( + Arc::new(EmptyExec::new(input_schema.clone())), + output_schema, + 0, + Vec::new(), + Arc::new(HashMap::from([(7, 0.25), (3, 0.75)])), + ); + let batch = RecordBatch::try_new( + input_schema, + vec![Arc::new(Int64Array::from(vec![7, 4, 3]))], + ) + .expect("input batch"); + + let output = exec.filter_and_project(batch).expect("filter and project"); + assert_eq!(output.num_columns(), 0); + assert_eq!(output.num_rows(), 2); + } +} diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index cc152f3b..c84a48e0 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -292,6 +292,14 @@ impl PaimonScanBuilder<'_> { /// Build a [`PaimonTableScan`] from the configured parameters. pub(crate) fn build(self) -> DFResult> { let read_fields = datafusion_read_fields(self.table); + self.build_with_read_fields(read_fields) + } + + /// Build using an internal read schema which may contain hidden system fields. + pub(crate) fn build_with_read_fields( + self, + read_fields: Vec, + ) -> DFResult> { let (projected_schema, read_type) = if let Some(indices) = self.projection { let fields: Vec = indices .iter() diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index df3ed0f1..5ee87099 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -2149,7 +2149,9 @@ mod vector_search_tests { mod hybrid_search_tests { use std::sync::Arc; - use datafusion::arrow::array::Int32Array; + #[cfg(feature = "fulltext")] + use datafusion::arrow::array::{Array, Int64Array}; + use datafusion::arrow::array::{Float32Array, Int32Array}; use paimon::catalog::Identifier; use paimon::table::BranchManager; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; @@ -2202,6 +2204,245 @@ mod hybrid_search_tests { ids } + fn extract_scored_ids( + batches: &[datafusion::arrow::record_batch::RecordBatch], + ) -> Vec<(i32, f32)> { + let mut rows = Vec::new(); + for batch in batches { + let id_array = batch + .column_by_name("id") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("Expected Int32Array for id"); + let score_array = batch + .column_by_name("__paimon_search_score") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("Expected Float32Array for __paimon_search_score"); + for row in 0..batch.num_rows() { + rows.push((id_array.value(row), score_array.value(row))); + } + } + rows + } + + fn extract_scores(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec { + let mut scores = Vec::new(); + for batch in batches { + assert_eq!(batch.num_columns(), 1, "only the score column was selected"); + let score_array = batch + .column_by_name("__paimon_search_score") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("Expected Float32Array for __paimon_search_score"); + scores.extend((0..batch.num_rows()).map(|row| score_array.value(row))); + } + scores + } + + #[tokio::test] + async fn test_hybrid_search_exposes_ranked_scores() { + let (ctx, _catalog, _tmp) = create_hybrid_search_context().await; + let batches = ctx + .sql( + "SELECT __paimon_search_score, id FROM hybrid_search( \ + 'paimon.default.test_java_vindex_vector', \ + array(named_struct( \ + 'field', 'embedding', \ + 'query_vector', array(1.0, 0.0, 0.0, 0.0), \ + 'limit', 3, \ + 'weight', 1.0)), \ + array(), \ + 3, \ + 'rrf') \ + ORDER BY __paimon_search_score DESC", + ) + .await + .expect("hybrid_search score SQL should parse") + .collect() + .await + .expect("hybrid_search score query should execute"); + + let rows = extract_scored_ids(&batches); + assert_eq!(rows.len(), 3); + for ((id, score), (expected_id, expected_score)) in + rows.into_iter() + .zip([(0, 1.0 / 61.0), (1, 1.0 / 62.0), (2, 1.0 / 63.0)]) + { + assert_eq!(id, expected_id); + assert!( + (score - expected_score).abs() < 1e-6, + "score for id {id}: expected {expected_score}, got {score}" + ); + } + } + + #[tokio::test] + async fn test_hybrid_search_score_only_and_empty_results() { + let (ctx, _catalog, _tmp) = create_hybrid_search_context().await; + let score_batches = ctx + .sql( + "SELECT __paimon_search_score FROM hybrid_search( \ + 'paimon.default.test_java_vindex_vector', \ + array(named_struct( \ + 'field', 'embedding', \ + 'query_vector', array(1.0, 0.0, 0.0, 0.0), \ + 'limit', 3, \ + 'weight', 1.0)), \ + array(), \ + 3, \ + 'rrf') \ + ORDER BY __paimon_search_score DESC", + ) + .await + .expect("score-only hybrid_search SQL should parse") + .collect() + .await + .expect("score-only hybrid_search query should execute"); + + let scores = extract_scores(&score_batches); + assert_eq!(scores.len(), 3); + for (score, expected) in scores.into_iter().zip([1.0 / 61.0, 1.0 / 62.0, 1.0 / 63.0]) { + assert!((score - expected).abs() < 1e-6); + } + + let empty_batches = ctx + .sql( + "SELECT __paimon_search_score FROM hybrid_search( \ + 'paimon.default.test_java_vindex_vector', \ + array(named_struct( \ + 'field', 'missing_embedding', \ + 'query_vector', array(1.0), \ + 'limit', 3, \ + 'weight', 1.0)), \ + array(), \ + 3, \ + 'rrf')", + ) + .await + .expect("empty hybrid_search score SQL should parse") + .collect() + .await + .expect("empty hybrid_search score query should execute"); + assert_eq!( + empty_batches + .iter() + .map(|batch| batch.num_rows()) + .sum::(), + 0 + ); + } + + #[cfg(feature = "fulltext")] + #[tokio::test] + async fn test_hybrid_search_score_with_row_tracking_only_fulltext() { + let (_tmp, ctx) = super::common::setup_sql_context().await; + super::common::exec( + &ctx, + "CREATE TABLE paimon.test_db.hybrid_raw_fulltext (id INT, content STRING) \ + WITH ( \ + 'row-tracking.enabled' = 'true', \ + 'global-index.search-mode' = 'full')", + ) + .await; + super::common::exec( + &ctx, + "INSERT INTO paimon.test_db.hybrid_raw_fulltext VALUES \ + (1, 'paimon search'), (2, 'other'), (3, 'paimon table')", + ) + .await; + + const SEARCH: &str = "hybrid_search( \ + 'paimon.test_db.hybrid_raw_fulltext', \ + array(), \ + array(named_struct( \ + 'column', 'content', \ + 'query', 'paimon', \ + 'limit', 10, \ + 'weight', 1.0)), \ + 10, \ + 'rrf')"; + let explicit_sql = format!("SELECT id, __paimon_search_score FROM {SEARCH} ORDER BY id"); + let explicit_batches = ctx + .sql(&explicit_sql) + .await + .expect("row-tracking-only score SQL should parse") + .collect() + .await + .expect("row-tracking-only score query should execute"); + + let mut ids = Vec::new(); + for batch in &explicit_batches { + let id_array = batch + .column_by_name("id") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("id column"); + let score_array = batch + .column_by_name("__paimon_search_score") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("score column"); + assert_eq!(score_array.null_count(), 0); + for row in 0..batch.num_rows() { + ids.push(id_array.value(row)); + assert!(score_array.value(row) > 0.0); + } + } + assert_eq!(ids, vec![1, 3]); + + let limited_sql = format!("SELECT id, __paimon_search_score FROM {SEARCH} LIMIT 2"); + let limited_batches = ctx + .sql(&limited_sql) + .await + .expect("row-tracking-only limited SQL should parse") + .collect() + .await + .expect("row-tracking-only limited query should execute"); + assert_eq!(extract_scored_ids(&limited_batches).len(), 2); + + let projected_sql = format!("SELECT id FROM {SEARCH} ORDER BY id"); + let projected_batches = ctx + .sql(&projected_sql) + .await + .expect("row-tracking-only projected SQL should parse") + .collect() + .await + .expect("row-tracking-only projected query should execute"); + assert_eq!(extract_ids(&projected_batches), vec![1, 3]); + + let count_sql = format!("SELECT COUNT(*) AS matches FROM {SEARCH}"); + let count_batches = ctx + .sql(&count_sql) + .await + .expect("row-tracking-only aggregate SQL should parse") + .collect() + .await + .expect("row-tracking-only aggregate query should execute"); + let matches = count_batches + .first() + .and_then(|batch| batch.column_by_name("matches")) + .and_then(|column| column.as_any().downcast_ref::()) + .expect("matches column"); + assert_eq!(matches.value(0), 2); + + let star_sql = format!("SELECT * FROM {SEARCH}"); + let star_batches = ctx + .sql(&star_sql) + .await + .expect("row-tracking-only SELECT * SQL should parse") + .collect() + .await + .expect("row-tracking-only SELECT * query should execute"); + assert_eq!( + star_batches + .iter() + .map(|batch| batch.num_rows()) + .sum::(), + 2 + ); + for batch in &star_batches { + assert_eq!(batch.num_columns(), 3); + assert!(batch.column_by_name("__paimon_search_score").is_some()); + assert!(batch.column_by_name("_ROW_ID").is_none()); + } + } + #[tokio::test] async fn test_hybrid_search_multiple_vector_routes_spark_shape() { let (ctx, _catalog, _tmp) = create_hybrid_search_context().await; diff --git a/docs/src/sql.md b/docs/src/sql.md index 587603e4..0207062c 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1391,7 +1391,26 @@ FROM hybrid_search( ); ``` -The function searches each route independently, merges route results with the selected ranker, and returns the top-k matching rows from the target table. The current DataFusion table function returns table rows only; it does not expose a metadata score column. +The function searches each route independently, merges route results with the selected ranker, and returns the top-k matching rows from the target table. Its output also includes a nullable `FLOAT` metadata column named `__paimon_search_score`, which contains the final score produced by the selected ranker: + +```sql +SELECT id, __paimon_search_score +FROM hybrid_search( + 'paimon.my_db.docs', + array( + named_struct( + 'field', 'embedding', + 'query_vector', array(1.0, 0.0, 0.0, 0.0) + ) + ), + array(), + 10, + 'rrf' +) +ORDER BY __paimon_search_score DESC; +``` + +The metadata column is part of the DataFusion table function schema, so `SELECT *` includes it. Use an explicit `ORDER BY __paimon_search_score DESC` when result ranking order matters; scan output order is not an ordering guarantee. ## Full-Text Search