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 cfa30ca Add hybrid search support for DataFusion (#473)
cfa30ca is described below
commit cfa30cad84ce122ea6a10d1f269c578e9293618f
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Jul 7 20:22:02 2026 +0800
Add hybrid search support for DataFusion (#473)
---
.../integrations/datafusion/src/hybrid_search.rs | 761 +++++++++++++++++++++
crates/integrations/datafusion/src/lib.rs | 2 +
crates/integrations/datafusion/src/sql_context.rs | 1 +
.../integrations/datafusion/tests/read_tables.rs | 87 +++
crates/paimon/src/lib.rs | 4 +
.../paimon/src/table/full_text_search_builder.rs | 15 +-
crates/paimon/src/table/hybrid_search_builder.rs | 505 ++++++++++++++
crates/paimon/src/table/mod.rs | 11 +
crates/paimon/src/table/vector_search_builder.rs | 40 +-
crates/paimon/src/vector_search.rs | 8 +
docs/src/sql.md | 95 ++-
11 files changed, 1519 insertions(+), 10 deletions(-)
diff --git a/crates/integrations/datafusion/src/hybrid_search.rs
b/crates/integrations/datafusion/src/hybrid_search.rs
new file mode 100644
index 0000000..e608366
--- /dev/null
+++ b/crates/integrations/datafusion/src/hybrid_search.rs
@@ -0,0 +1,761 @@
+// 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.
+
+//! `hybrid_search` table-valued function for DataFusion.
+//!
+//! Spark-compatible shape:
+//! ```sql
+//! SELECT * FROM hybrid_search(
+//! 'table_name',
+//! array(named_struct('field', 'embedding', 'query_vector', array(1.0,
0.0))),
+//! array(named_struct('column', 'content', 'query', 'paimon')),
+//! 10,
+//! 'rrf')
+//! ```
+
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::arrow::array::Array;
+use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
+use datafusion::catalog::{Session, TableFunctionImpl};
+use datafusion::common::{project_schema, ScalarValue};
+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 paimon::catalog::Catalog;
+use paimon::table::{HybridSearchRanker, HybridSearchRoute};
+
+use crate::error::to_datafusion_error;
+use crate::runtime::{await_with_runtime, block_on_with_runtime};
+use crate::table::{PaimonScanBuilder, PaimonTableProvider};
+use crate::table_function_args::{
+ extract_int_literal, extract_string_literal, parse_table_identifier,
+};
+
+const FUNCTION_NAME: &str = "hybrid_search";
+
+pub fn register_hybrid_search(
+ ctx: &SessionContext,
+ catalog: Arc<dyn Catalog>,
+ default_database: &str,
+) {
+ ctx.register_udf(
+ datafusion::functions_nested::make_array::make_array_udf()
+ .as_ref()
+ .clone()
+ .with_aliases(["array"]),
+ );
+ ctx.register_udtf(
+ FUNCTION_NAME,
+ Arc::new(HybridSearchFunction::new(catalog, default_database)),
+ );
+}
+
+pub struct HybridSearchFunction {
+ catalog: Arc<dyn Catalog>,
+ default_database: String,
+}
+
+impl Debug for HybridSearchFunction {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("HybridSearchFunction")
+ .field("default_database", &self.default_database)
+ .finish()
+ }
+}
+
+impl HybridSearchFunction {
+ pub fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
+ Self {
+ catalog,
+ default_database: default_database.to_string(),
+ }
+ }
+}
+
+impl TableFunctionImpl for HybridSearchFunction {
+ fn call(&self, args: &[Expr]) -> DFResult<Arc<dyn TableProvider>> {
+ if args.len() != 4 && args.len() != 5 {
+ return Err(DataFusionError::Plan(
+ "hybrid_search requires 4 or 5 arguments: (table_name,
vector_routes, full_text_routes, limit[, ranker])".to_string(),
+ ));
+ }
+
+ let table_name = extract_string_literal(FUNCTION_NAME, &args[0],
"table_name")?;
+ let limit = extract_int_literal(FUNCTION_NAME, &args[3], "limit")?;
+ if limit <= 0 {
+ return Err(DataFusionError::Plan(
+ "hybrid_search: limit must be positive".to_string(),
+ ));
+ }
+
+ let ranker = if args.len() == 5 {
+ extract_string_literal(FUNCTION_NAME, &args[4], "ranker")?
+ } else {
+ HybridSearchRanker::RRF.to_string()
+ };
+ HybridSearchRanker::parse(&ranker).map_err(to_datafusion_error)?;
+
+ let mut routes = parse_vector_routes(&args[1], limit as usize)?;
+ routes.extend(parse_full_text_routes(&args[2], limit as usize)?);
+
+ let identifier =
+ parse_table_identifier(FUNCTION_NAME, &table_name,
&self.default_database)?;
+ let catalog = Arc::clone(&self.catalog);
+ let table = block_on_with_runtime(
+ async move { catalog.get_table(&identifier).await },
+ "hybrid_search: catalog access thread panicked",
+ )
+ .map_err(to_datafusion_error)?;
+
+ Ok(Arc::new(HybridSearchTableProvider {
+ inner: PaimonTableProvider::try_new(table)?,
+ routes,
+ limit: limit as usize,
+ ranker,
+ }))
+ }
+}
+
+#[derive(Debug)]
+struct HybridSearchTableProvider {
+ inner: PaimonTableProvider,
+ routes: Vec<HybridSearchRoute>,
+ limit: usize,
+ ranker: String,
+}
+
+#[async_trait]
+impl TableProvider for HybridSearchTableProvider {
+ fn schema(&self) -> ArrowSchemaRef {
+ self.inner.schema()
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ _filters: &[Expr],
+ limit: Option<usize>,
+ ) -> DFResult<Arc<dyn ExecutionPlan>> {
+ let table = self.inner.table();
+
+ let row_ranges = await_with_runtime(async {
+ let mut builder = table.new_hybrid_search_builder();
+ for route in self.routes.clone() {
+ builder.add_route(route);
+ }
+ builder
+ .with_limit(self.limit)
+ .with_ranker(&self.ranker)
+ .map_err(to_datafusion_error)?;
+ builder.execute().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)));
+ }
+
+ 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)?;
+
+ PaimonScanBuilder {
+ table,
+ schema: &self.schema(),
+ plan: &plan,
+ scan_trace: None,
+ projection,
+ pushed_predicate: None,
+ limit,
+ target_partitions:
state.config_options().execution.target_partitions,
+ filter_exact: false,
+ }
+ .build()
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> DFResult<Vec<TableProviderFilterPushDown>> {
+ Ok(vec![
+ TableProviderFilterPushDown::Unsupported;
+ filters.len()
+ ])
+ }
+}
+
+fn parse_vector_routes(expr: &Expr, default_limit: usize) ->
DFResult<Vec<HybridSearchRoute>> {
+ if let Some(routes) = extract_literal_array_values(expr, "vector_routes")?
{
+ return routes
+ .iter()
+ .map(|route| parse_vector_route_scalar(route, default_limit))
+ .collect();
+ }
+
+ extract_array_elements(expr, "vector_routes")?
+ .into_iter()
+ .map(|route| parse_vector_route(route, default_limit))
+ .collect()
+}
+
+fn parse_full_text_routes(expr: &Expr, default_limit: usize) ->
DFResult<Vec<HybridSearchRoute>> {
+ if let Some(routes) = extract_literal_array_values(expr,
"full_text_routes")? {
+ return routes
+ .iter()
+ .map(|route| parse_full_text_route_scalar(route, default_limit))
+ .collect();
+ }
+
+ extract_array_elements(expr, "full_text_routes")?
+ .into_iter()
+ .map(|route| parse_full_text_route(route, default_limit))
+ .collect()
+}
+
+fn parse_vector_route(expr: &Expr, default_limit: usize) ->
DFResult<HybridSearchRoute> {
+ let fields = extract_named_struct_fields(expr, "vector route")?;
+ let field_name = optional_field(&fields, &["field", "vector_column"])
+ .ok_or_else(|| {
+ DataFusionError::Plan(
+ "hybrid_search: vector route must define field or
vector_column".to_string(),
+ )
+ })
+ .and_then(|expr| extract_string_literal(FUNCTION_NAME, expr, "vector
route field"))?;
+ let vector = required_field(&fields, "query_vector")
+ .and_then(|expr| extract_float_array(expr, "query_vector"))?;
+ let limit = optional_field(&fields, &["limit"])
+ .map(|expr| extract_positive_usize(expr, "vector route limit"))
+ .transpose()?
+ .unwrap_or(default_limit);
+ let weight = optional_field(&fields, &["weight"])
+ .map(|expr| extract_positive_f32(expr, "weight"))
+ .transpose()?
+ .unwrap_or(1.0);
+ let options = optional_field(&fields, &["options"])
+ .map(extract_options)
+ .transpose()?
+ .unwrap_or_default();
+
+ HybridSearchRoute::vector(field_name, vector, limit, weight, options)
+ .map_err(to_datafusion_error)
+}
+
+fn parse_vector_route_scalar(
+ scalar: &ScalarValue,
+ default_limit: usize,
+) -> DFResult<HybridSearchRoute> {
+ let fields = extract_struct_scalar_fields(scalar, "vector route")?;
+ let field_name = optional_scalar_field(&fields, &["field",
"vector_column"])
+ .ok_or_else(|| {
+ DataFusionError::Plan(
+ "hybrid_search: vector route must define field or
vector_column".to_string(),
+ )
+ })
+ .and_then(|scalar| scalar_to_string(scalar, "vector route field"))?;
+ let vector = required_scalar_field(&fields, "query_vector")
+ .and_then(|scalar| scalar_to_float_array(scalar, "query_vector"))?;
+ let limit = optional_scalar_field(&fields, &["limit"])
+ .map(|scalar| scalar_to_positive_usize(scalar, "vector route limit"))
+ .transpose()?
+ .unwrap_or(default_limit);
+ let weight = optional_scalar_field(&fields, &["weight"])
+ .map(|scalar| scalar_to_positive_f32(scalar, "weight"))
+ .transpose()?
+ .unwrap_or(1.0);
+ let options = optional_scalar_field(&fields, &["options"])
+ .map(scalar_to_options)
+ .transpose()?
+ .unwrap_or_default();
+
+ HybridSearchRoute::vector(field_name, vector, limit, weight, options)
+ .map_err(to_datafusion_error)
+}
+
+fn parse_full_text_route(expr: &Expr, default_limit: usize) ->
DFResult<HybridSearchRoute> {
+ let fields = extract_named_struct_fields(expr, "full-text route")?;
+ let column = required_field(&fields, "column")
+ .and_then(|expr| extract_string_literal(FUNCTION_NAME, expr,
"full-text route column"))?;
+ let query = required_field(&fields, "query")
+ .and_then(|expr| extract_string_literal(FUNCTION_NAME, expr,
"full-text route query"))?;
+ let limit = optional_field(&fields, &["limit"])
+ .map(|expr| extract_positive_usize(expr, "full-text route limit"))
+ .transpose()?
+ .unwrap_or(default_limit);
+ let weight = optional_field(&fields, &["weight"])
+ .map(|expr| extract_positive_f32(expr, "weight"))
+ .transpose()?
+ .unwrap_or(1.0);
+ let options = optional_field(&fields, &["options"])
+ .map(extract_options)
+ .transpose()?
+ .unwrap_or_default();
+
+ HybridSearchRoute::full_text(column, query, limit, weight,
options).map_err(to_datafusion_error)
+}
+
+fn parse_full_text_route_scalar(
+ scalar: &ScalarValue,
+ default_limit: usize,
+) -> DFResult<HybridSearchRoute> {
+ let fields = extract_struct_scalar_fields(scalar, "full-text route")?;
+ let column = required_scalar_field(&fields, "column")
+ .and_then(|scalar| scalar_to_string(scalar, "full-text route
column"))?;
+ let query = required_scalar_field(&fields, "query")
+ .and_then(|scalar| scalar_to_string(scalar, "full-text route query"))?;
+ let limit = optional_scalar_field(&fields, &["limit"])
+ .map(|scalar| scalar_to_positive_usize(scalar, "full-text route
limit"))
+ .transpose()?
+ .unwrap_or(default_limit);
+ let weight = optional_scalar_field(&fields, &["weight"])
+ .map(|scalar| scalar_to_positive_f32(scalar, "weight"))
+ .transpose()?
+ .unwrap_or(1.0);
+ let options = optional_scalar_field(&fields, &["options"])
+ .map(scalar_to_options)
+ .transpose()?
+ .unwrap_or_default();
+
+ HybridSearchRoute::full_text(column, query, limit, weight,
options).map_err(to_datafusion_error)
+}
+
+fn extract_array_elements<'a>(expr: &'a Expr, name: &str) -> DFResult<Vec<&'a
Expr>> {
+ match expr {
+ Expr::ScalarFunction(function)
+ if is_function(function.name(), &["make_array", "array"]) =>
+ {
+ Ok(function.args.iter().collect())
+ }
+ _ => Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be array(...), got: {expr}"
+ ))),
+ }
+}
+
+fn extract_literal_array_values(expr: &Expr, name: &str) ->
DFResult<Option<Vec<ScalarValue>>> {
+ let Expr::Literal(scalar, _) = expr else {
+ return Ok(None);
+ };
+ scalar_array_values(scalar, name).map(Some)
+}
+
+fn scalar_array_values(scalar: &ScalarValue, name: &str) ->
DFResult<Vec<ScalarValue>> {
+ let values = match scalar {
+ ScalarValue::List(array) => array.value(0),
+ ScalarValue::LargeList(array) => array.value(0),
+ ScalarValue::ListView(array) => array.value(0),
+ ScalarValue::LargeListView(array) => array.value(0),
+ ScalarValue::FixedSizeList(array) => array.value(0),
+ _ => {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be an array, got: {scalar}"
+ )));
+ }
+ };
+
+ (0..values.len())
+ .map(|index| ScalarValue::try_from_array(values.as_ref(), index))
+ .collect()
+}
+
+fn extract_named_struct_fields<'a>(
+ expr: &'a Expr,
+ name: &str,
+) -> DFResult<Vec<(String, &'a Expr)>> {
+ let Expr::ScalarFunction(function) = expr else {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be named_struct(...), got: {expr}"
+ )));
+ };
+ if !is_function(function.name(), &["named_struct"]) {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be named_struct(...), got: {expr}"
+ )));
+ }
+ if function.args.len() % 2 != 0 {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must contain key/value pairs"
+ )));
+ }
+
+ let mut fields = Vec::with_capacity(function.args.len() / 2);
+ for pair in function.args.chunks_exact(2) {
+ let key = extract_string_literal(FUNCTION_NAME, &pair[0], "route field
name")?;
+ fields.push((key, &pair[1]));
+ }
+ Ok(fields)
+}
+
+fn extract_struct_scalar_fields(
+ scalar: &ScalarValue,
+ name: &str,
+) -> DFResult<Vec<(String, ScalarValue)>> {
+ let ScalarValue::Struct(array) = scalar else {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be named_struct(...), got: {scalar}"
+ )));
+ };
+ if array.is_null(0) {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} cannot be null"
+ )));
+ }
+
+ array
+ .fields()
+ .iter()
+ .zip(array.columns())
+ .map(|(field, column)| {
+ Ok((
+ field.name().clone(),
+ ScalarValue::try_from_array(column.as_ref(), 0)?,
+ ))
+ })
+ .collect()
+}
+
+fn required_field<'a>(fields: &'a [(String, &'a Expr)], name: &str) ->
DFResult<&'a Expr> {
+ optional_field(fields, &[name])
+ .ok_or_else(|| DataFusionError::Plan(format!("hybrid_search: route
must define {name}")))
+}
+
+fn optional_field<'a>(fields: &'a [(String, &'a Expr)], names: &[&str]) ->
Option<&'a Expr> {
+ fields
+ .iter()
+ .find(|(field_name, _)| names.iter().any(|name| field_name == name))
+ .map(|(_, expr)| *expr)
+}
+
+fn required_scalar_field<'a>(
+ fields: &'a [(String, ScalarValue)],
+ name: &str,
+) -> DFResult<&'a ScalarValue> {
+ optional_scalar_field(fields, &[name])
+ .ok_or_else(|| DataFusionError::Plan(format!("hybrid_search: route
must define {name}")))
+}
+
+fn optional_scalar_field<'a>(
+ fields: &'a [(String, ScalarValue)],
+ names: &[&str],
+) -> Option<&'a ScalarValue> {
+ fields
+ .iter()
+ .find(|(field_name, _)| names.iter().any(|name| field_name == name))
+ .map(|(_, scalar)| scalar)
+}
+
+fn extract_positive_usize(expr: &Expr, name: &str) -> DFResult<usize> {
+ let value = extract_int_literal(FUNCTION_NAME, expr, name)?;
+ if value <= 0 {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be positive"
+ )));
+ }
+ Ok(value as usize)
+}
+
+fn extract_float_array(expr: &Expr, name: &str) -> DFResult<Vec<f32>> {
+ if let Ok(json) = extract_string_literal(FUNCTION_NAME, expr, name) {
+ let vector: Vec<f32> = serde_json::from_str(&json).map_err(|e| {
+ DataFusionError::Plan(format!(
+ "hybrid_search: {name} string must be a JSON array of floats:
{e}"
+ ))
+ })?;
+ if vector.is_empty() {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} cannot be empty"
+ )));
+ }
+ return Ok(vector);
+ }
+
+ let elements = extract_array_elements(expr, name)?;
+ if elements.is_empty() {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} cannot be empty"
+ )));
+ }
+ elements
+ .into_iter()
+ .map(|expr| scalar_to_f32(expr, name))
+ .collect()
+}
+
+fn extract_positive_f32(expr: &Expr, name: &str) -> DFResult<f32> {
+ let value = scalar_to_f32(expr, name)?;
+ if !value.is_finite() || value <= 0.0 {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be finite and positive, got: {value}"
+ )));
+ }
+ Ok(value)
+}
+
+fn scalar_to_f32(expr: &Expr, name: &str) -> DFResult<f32> {
+ let Expr::Literal(scalar, _) = expr else {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be a numeric literal, got: {expr}"
+ )));
+ };
+ match scalar {
+ ScalarValue::Float32(Some(value)) => Ok(*value),
+ ScalarValue::Float64(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int8(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int16(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int32(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int64(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt8(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt16(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt32(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt64(Some(value)) => Ok(*value as f32),
+ ScalarValue::Utf8(Some(value)) => value.parse::<f32>().map_err(|e| {
+ DataFusionError::Plan(format!(
+ "hybrid_search: {name} string must be a float, got '{value}':
{e}"
+ ))
+ }),
+ _ => Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be a numeric literal, got: {expr}"
+ ))),
+ }
+}
+
+fn scalar_to_string(scalar: &ScalarValue, name: &str) -> DFResult<String> {
+ match scalar {
+ ScalarValue::Utf8(Some(value))
+ | ScalarValue::Utf8View(Some(value))
+ | ScalarValue::LargeUtf8(Some(value)) => Ok(value.clone()),
+ _ => Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be a string literal, got: {scalar}"
+ ))),
+ }
+}
+
+fn scalar_to_positive_usize(scalar: &ScalarValue, name: &str) ->
DFResult<usize> {
+ let value = match scalar {
+ ScalarValue::Int8(Some(value)) => *value as i64,
+ ScalarValue::Int16(Some(value)) => *value as i64,
+ ScalarValue::Int32(Some(value)) => *value as i64,
+ ScalarValue::Int64(Some(value)) => *value,
+ ScalarValue::UInt8(Some(value)) => *value as i64,
+ ScalarValue::UInt16(Some(value)) => *value as i64,
+ ScalarValue::UInt32(Some(value)) => *value as i64,
+ ScalarValue::UInt64(Some(value)) => i64::try_from(*value).map_err(|_| {
+ DataFusionError::Plan(format!("hybrid_search: {name} value exceeds
i64 range"))
+ })?,
+ _ => {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be an integer literal, got:
{scalar}"
+ )));
+ }
+ };
+ if value <= 0 {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be positive"
+ )));
+ }
+ Ok(value as usize)
+}
+
+fn scalar_value_to_f32(scalar: &ScalarValue, name: &str) -> DFResult<f32> {
+ match scalar {
+ ScalarValue::Float16(Some(value)) => Ok(value.to_f32()),
+ ScalarValue::Float32(Some(value)) => Ok(*value),
+ ScalarValue::Float64(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int8(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int16(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int32(Some(value)) => Ok(*value as f32),
+ ScalarValue::Int64(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt8(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt16(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt32(Some(value)) => Ok(*value as f32),
+ ScalarValue::UInt64(Some(value)) => Ok(*value as f32),
+ ScalarValue::Utf8(Some(value))
+ | ScalarValue::Utf8View(Some(value))
+ | ScalarValue::LargeUtf8(Some(value)) =>
value.parse::<f32>().map_err(|e| {
+ DataFusionError::Plan(format!(
+ "hybrid_search: {name} string must be a float, got '{value}':
{e}"
+ ))
+ }),
+ _ => Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be a numeric literal, got: {scalar}"
+ ))),
+ }
+}
+
+fn scalar_to_positive_f32(scalar: &ScalarValue, name: &str) -> DFResult<f32> {
+ let value = scalar_value_to_f32(scalar, name)?;
+ if !value.is_finite() || value <= 0.0 {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} must be finite and positive, got: {value}"
+ )));
+ }
+ Ok(value)
+}
+
+fn scalar_to_float_array(scalar: &ScalarValue, name: &str) ->
DFResult<Vec<f32>> {
+ let values = scalar_array_values(scalar, name)?;
+ if values.is_empty() {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: {name} cannot be empty"
+ )));
+ }
+ values
+ .iter()
+ .map(|value| scalar_value_to_f32(value, name))
+ .collect()
+}
+
+fn scalar_to_options(scalar: &ScalarValue) -> DFResult<HashMap<String,
String>> {
+ if matches!(scalar, ScalarValue::Null) {
+ return Ok(HashMap::new());
+ }
+
+ if let Ok(json) = scalar_to_string(scalar, "options") {
+ if json.trim().is_empty() {
+ return Ok(HashMap::new());
+ }
+ return serde_json::from_str(&json).map_err(|e| {
+ DataFusionError::Plan(format!(
+ "hybrid_search: options string must be a JSON object: {e}"
+ ))
+ });
+ }
+
+ let ScalarValue::Map(array) = scalar else {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: options must be map(...), got: {scalar}"
+ )));
+ };
+ if array.is_null(0) {
+ return Ok(HashMap::new());
+ }
+
+ let entries = array.value(0);
+ let keys = entries.column(0);
+ let values = entries.column(1);
+ (0..entries.len())
+ .map(|index| {
+ Ok((
+ scalar_to_string(
+ &ScalarValue::try_from_array(keys.as_ref(), index)?,
+ "options key",
+ )?,
+ scalar_to_string(
+ &ScalarValue::try_from_array(values.as_ref(), index)?,
+ "options value",
+ )?,
+ ))
+ })
+ .collect()
+}
+
+fn extract_options(expr: &Expr) -> DFResult<HashMap<String, String>> {
+ if let Ok(json) = extract_string_literal(FUNCTION_NAME, expr, "options") {
+ if json.trim().is_empty() {
+ return Ok(HashMap::new());
+ }
+ return serde_json::from_str(&json).map_err(|e| {
+ DataFusionError::Plan(format!(
+ "hybrid_search: options string must be a JSON object: {e}"
+ ))
+ });
+ }
+
+ let Expr::ScalarFunction(function) = expr else {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: options must be map(...), got: {expr}"
+ )));
+ };
+ if !is_function(function.name(), &["map", "make_map"]) {
+ return Err(DataFusionError::Plan(format!(
+ "hybrid_search: options must be map(...), got: {expr}"
+ )));
+ }
+ if function.args.is_empty() {
+ return Ok(HashMap::new());
+ }
+
+ if function.args.len() == 2
+ && is_array_expr(&function.args[0])
+ && is_array_expr(&function.args[1])
+ {
+ let keys = extract_array_elements(&function.args[0], "options keys")?;
+ let values = extract_array_elements(&function.args[1], "options
values")?;
+ if keys.len() != values.len() {
+ return Err(DataFusionError::Plan(
+ "hybrid_search: options keys and values must have the same
length".to_string(),
+ ));
+ }
+ return keys
+ .into_iter()
+ .zip(values)
+ .map(|(key, value)| {
+ Ok((
+ extract_string_literal(FUNCTION_NAME, key, "options key")?,
+ extract_string_literal(FUNCTION_NAME, value, "options
value")?,
+ ))
+ })
+ .collect();
+ }
+
+ if function.args.len() % 2 != 0 {
+ return Err(DataFusionError::Plan(
+ "hybrid_search: options map must contain key/value
pairs".to_string(),
+ ));
+ }
+
+ function
+ .args
+ .chunks_exact(2)
+ .map(|pair| {
+ Ok((
+ extract_string_literal(FUNCTION_NAME, &pair[0], "options
key")?,
+ extract_string_literal(FUNCTION_NAME, &pair[1], "options
value")?,
+ ))
+ })
+ .collect()
+}
+
+fn is_array_expr(expr: &Expr) -> bool {
+ matches!(
+ expr,
+ Expr::ScalarFunction(function) if is_function(function.name(),
&["make_array", "array"])
+ )
+}
+
+fn is_function(actual: &str, expected: &[&str]) -> bool {
+ expected
+ .iter()
+ .any(|expected| actual.eq_ignore_ascii_case(expected))
+}
diff --git a/crates/integrations/datafusion/src/lib.rs
b/crates/integrations/datafusion/src/lib.rs
index ec838c0..c4b79b5 100644
--- a/crates/integrations/datafusion/src/lib.rs
+++ b/crates/integrations/datafusion/src/lib.rs
@@ -43,6 +43,7 @@ mod error;
mod filter_pushdown;
#[cfg(feature = "fulltext")]
mod full_text_search;
+mod hybrid_search;
mod lateral_vector_search;
mod merge_into;
mod physical_plan;
@@ -72,6 +73,7 @@ pub use catalog::{PaimonCatalogProvider,
PaimonSchemaProvider};
pub use error::to_datafusion_error;
#[cfg(feature = "fulltext")]
pub use full_text_search::{register_full_text_search, FullTextSearchFunction};
+pub use hybrid_search::{register_hybrid_search, HybridSearchFunction};
pub use physical_plan::PaimonTableScan;
pub use relation_planner::PaimonRelationPlanner;
pub use sql_context::SQLContext;
diff --git a/crates/integrations/datafusion/src/sql_context.rs
b/crates/integrations/datafusion/src/sql_context.rs
index ecfa252..babd650 100644
--- a/crates/integrations/datafusion/src/sql_context.rs
+++ b/crates/integrations/datafusion/src/sql_context.rs
@@ -2582,6 +2582,7 @@ fn register_table_functions(
crate::vector_search::register_vector_search(ctx, Arc::clone(catalog),
default_database);
#[cfg(feature = "fulltext")]
crate::full_text_search::register_full_text_search(ctx,
Arc::clone(catalog), default_database);
+ crate::hybrid_search::register_hybrid_search(ctx, Arc::clone(catalog),
default_database);
}
#[cfg(test)]
diff --git a/crates/integrations/datafusion/tests/read_tables.rs
b/crates/integrations/datafusion/tests/read_tables.rs
index 13c9f12..718a1cd 100644
--- a/crates/integrations/datafusion/tests/read_tables.rs
+++ b/crates/integrations/datafusion/tests/read_tables.rs
@@ -1824,3 +1824,90 @@ mod vector_search_tests {
);
}
}
+
+// ======================= Hybrid Search Tests =======================
+
+mod hybrid_search_tests {
+ use std::sync::Arc;
+
+ use datafusion::arrow::array::Int32Array;
+ use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
+ use paimon_datafusion::SQLContext;
+
+ fn extract_test_warehouse(archive_name: &str) -> (tempfile::TempDir,
String) {
+ let archive_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+ .join("testdata")
+ .join(archive_name);
+ let file = std::fs::File::open(&archive_path)
+ .unwrap_or_else(|e| panic!("Failed to open {}: {e}",
archive_path.display()));
+ let decoder = flate2::read::GzDecoder::new(file);
+ let mut archive = tar::Archive::new(decoder);
+
+ let tmp = tempfile::tempdir().expect("Failed to create temp dir");
+ let db_dir = tmp.path().join("default.db");
+ std::fs::create_dir_all(&db_dir).unwrap();
+ archive.unpack(&db_dir).unwrap();
+
+ let warehouse = format!("file://{}", tmp.path().display());
+ (tmp, warehouse)
+ }
+
+ async fn create_hybrid_search_context() -> (SQLContext, tempfile::TempDir)
{
+ let (tmp, warehouse) =
extract_test_warehouse("test_java_vindex_vector.tar.gz");
+ let mut options = Options::new();
+ options.set(CatalogOptions::WAREHOUSE, warehouse);
+ let catalog = FileSystemCatalog::new(options).expect("Failed to create
catalog");
+ let catalog: Arc<dyn Catalog> = Arc::new(catalog);
+
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog.clone())
+ .await
+ .expect("Failed to register catalog");
+ (ctx, tmp)
+ }
+
+ fn extract_ids(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.sort();
+ ids
+ }
+
+ #[tokio::test]
+ async fn test_hybrid_search_multiple_vector_routes_spark_shape() {
+ let (ctx, _tmp) = create_hybrid_search_context().await;
+ let batches = ctx
+ .sql(
+ "SELECT 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), \
+ named_struct( \
+ 'field', 'embedding', \
+ 'query_vector', array(0.9, 0.1, 0.0, 0.0), \
+ 'limit', 3, \
+ 'weight', 1.0)), \
+ array(), \
+ 3, \
+ 'rrf')",
+ )
+ .await
+ .expect("hybrid_search SQL should parse")
+ .collect()
+ .await
+ .expect("hybrid_search query should execute");
+
+ assert_eq!(extract_ids(&batches), vec![0, 1, 2]);
+ }
+}
diff --git a/crates/paimon/src/lib.rs b/crates/paimon/src/lib.rs
index da34c7d..d9cd290 100644
--- a/crates/paimon/src/lib.rs
+++ b/crates/paimon/src/lib.rs
@@ -51,3 +51,7 @@ pub use table::{
RenamingSnapshotCommit, RowRange, ScanTrace, SnapshotCommit,
SnapshotManager, Table,
TableCommit, TableRead, TableScan, TableUpdate, TableWrite, TagManager,
WriteBuilder,
};
+
+pub use table::{
+ HybridSearchBuilder, HybridSearchRanker, HybridSearchRoute,
HybridSearchRouteKind,
+};
diff --git a/crates/paimon/src/table/full_text_search_builder.rs
b/crates/paimon/src/table/full_text_search_builder.rs
index 0b92674..042a9e8 100644
--- a/crates/paimon/src/table/full_text_search_builder.rs
+++ b/crates/paimon/src/table/full_text_search_builder.rs
@@ -96,6 +96,10 @@ impl<'a> FullTextSearchBuilder<'a> {
///
/// Reference: `FullTextSearchBuilder.executeLocal()`
pub async fn execute(&self) -> crate::Result<Vec<RowRange>> {
+ Ok(self.execute_scored().await?.to_row_ranges())
+ }
+
+ pub async fn execute_scored(&self) -> crate::Result<SearchResult> {
// Fail closed: returns data-derived row ranges outside
`TableScan`/`TableRead`.
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
let text_column =
@@ -123,7 +127,7 @@ impl<'a> FullTextSearchBuilder<'a> {
let snapshot = match snapshot_manager.get_latest_snapshot().await? {
Some(s) => s,
- None => return Ok(Vec::new()),
+ None => return Ok(SearchResult::empty()),
};
let index_entries = match snapshot.index_manifest() {
@@ -168,14 +172,14 @@ async fn evaluate_full_text_search(
evaluation: FullTextSearchEvaluation<'_>,
index_entries: &[IndexManifestEntry],
search: &FullTextSearch,
-) -> crate::Result<Vec<RowRange>> {
+) -> crate::Result<SearchResult> {
let table_path = evaluation.table_path.trim_end_matches('/');
let core_options = CoreOptions::new(evaluation.table_options);
let search_mode = core_options.global_index_search_mode()?;
let field_id = match find_field_id_by_name(evaluation.schema_fields,
&search.field_name) {
Some(id) => id,
- None => return Ok(Vec::new()),
+ None => return Ok(SearchResult::empty()),
};
// Collect tantivy fulltext entries for the target field.
@@ -192,7 +196,7 @@ async fn evaluate_full_text_search(
.collect();
if fulltext_entries.is_empty() && search_mode ==
GlobalIndexSearchMode::Fast {
- return Ok(Vec::new());
+ return Ok(SearchResult::empty());
}
let deleted_row_index = if core_options.data_evolution_enabled() {
@@ -280,8 +284,7 @@ async fn evaluate_full_text_search(
Ok(merged
.without_deleted_row_ranges(deleted_row_index.as_ref())?
- .top_k(search.limit)
- .to_row_ranges())
+ .top_k(search.limit))
}
fn is_tantivy_fulltext_index_file(index_file: &IndexFileMeta) -> bool {
diff --git a/crates/paimon/src/table/hybrid_search_builder.rs
b/crates/paimon/src/table/hybrid_search_builder.rs
new file mode 100644
index 0000000..11b7a02
--- /dev/null
+++ b/crates/paimon/src/table/hybrid_search_builder.rs
@@ -0,0 +1,505 @@
+// 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.
+
+//! Hybrid search builder for combining multiple search routes.
+//!
+//! Reference: `org.apache.paimon.table.source.HybridSearchBuilder`.
+
+use std::collections::HashMap;
+
+use crate::spec::CoreOptions;
+use crate::table::{RowRange, Table};
+use crate::vector_search::SearchResult;
+
+const RRF_K: f32 = 60.0;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum HybridSearchRanker {
+ Rrf,
+ WeightedScore,
+ Mrr,
+}
+
+impl HybridSearchRanker {
+ pub const RRF: &'static str = "rrf";
+ pub const WEIGHTED_SCORE: &'static str = "weighted_score";
+ pub const MRR: &'static str = "mrr";
+
+ pub fn parse(ranker: &str) -> crate::Result<Self> {
+ match ranker.trim().to_ascii_lowercase().as_str() {
+ "" | Self::RRF => Ok(Self::Rrf),
+ Self::WEIGHTED_SCORE => Ok(Self::WeightedScore),
+ Self::MRR => Ok(Self::Mrr),
+ _ => Err(crate::Error::ConfigInvalid {
+ message: format!("Unsupported hybrid ranker: {ranker}"),
+ }),
+ }
+ }
+
+ pub fn as_str(self) -> &'static str {
+ match self {
+ Self::Rrf => Self::RRF,
+ Self::WeightedScore => Self::WEIGHTED_SCORE,
+ Self::Mrr => Self::MRR,
+ }
+ }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum HybridSearchRouteKind {
+ Vector,
+ FullText,
+}
+
+#[derive(Clone, Debug)]
+pub struct HybridSearchRoute {
+ kind: HybridSearchRouteKind,
+ field_name: String,
+ vector: Option<Vec<f32>>,
+ full_text_query: Option<String>,
+ limit: usize,
+ weight: f32,
+ options: HashMap<String, String>,
+}
+
+impl HybridSearchRoute {
+ pub fn vector(
+ field_name: impl Into<String>,
+ vector: Vec<f32>,
+ limit: usize,
+ weight: f32,
+ options: HashMap<String, String>,
+ ) -> crate::Result<Self> {
+ let field_name = field_name.into();
+ Self::validate_common(&field_name, limit, weight)?;
+ if vector.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "Search vector cannot be empty".to_string(),
+ source: None,
+ });
+ }
+ Ok(Self {
+ kind: HybridSearchRouteKind::Vector,
+ field_name,
+ vector: Some(vector),
+ full_text_query: None,
+ limit,
+ weight,
+ options,
+ })
+ }
+
+ pub fn full_text(
+ field_name: impl Into<String>,
+ query: impl Into<String>,
+ limit: usize,
+ weight: f32,
+ options: HashMap<String, String>,
+ ) -> crate::Result<Self> {
+ if !options.is_empty() {
+ return Err(crate::Error::ConfigInvalid {
+ message: "Full-text hybrid route options are not supported
yet".to_string(),
+ });
+ }
+
+ let field_name = field_name.into();
+ let query = query.into();
+ Self::validate_common(&field_name, limit, weight)?;
+ if query.is_empty() {
+ return Err(crate::Error::ConfigInvalid {
+ message: "Full-text route query cannot be empty".to_string(),
+ });
+ }
+
+ Ok(Self {
+ kind: HybridSearchRouteKind::FullText,
+ field_name,
+ vector: None,
+ full_text_query: Some(query),
+ limit,
+ weight,
+ options,
+ })
+ }
+
+ fn validate_common(field_name: &str, limit: usize, weight: f32) ->
crate::Result<()> {
+ if field_name.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "Field name cannot be null or empty".to_string(),
+ source: None,
+ });
+ }
+ if limit == 0 {
+ return Err(crate::Error::ConfigInvalid {
+ message: "Limit must be positive".to_string(),
+ });
+ }
+ if !weight.is_finite() || weight <= 0.0 {
+ return Err(crate::Error::ConfigInvalid {
+ message: format!("Weight must be finite and positive, got:
{weight}"),
+ });
+ }
+ Ok(())
+ }
+
+ pub fn kind(&self) -> HybridSearchRouteKind {
+ self.kind
+ }
+
+ pub fn field_name(&self) -> &str {
+ &self.field_name
+ }
+
+ pub fn vector_value(&self) -> Option<&[f32]> {
+ self.vector.as_deref()
+ }
+
+ pub fn full_text_query(&self) -> Option<&str> {
+ self.full_text_query.as_deref()
+ }
+
+ pub fn limit(&self) -> usize {
+ self.limit
+ }
+
+ pub fn weight(&self) -> f32 {
+ self.weight
+ }
+
+ pub fn options(&self) -> &HashMap<String, String> {
+ &self.options
+ }
+}
+
+pub struct HybridSearchBuilder<'a> {
+ table: &'a Table,
+ routes: Vec<HybridSearchRoute>,
+ limit: Option<usize>,
+ ranker: HybridSearchRanker,
+}
+
+impl<'a> HybridSearchBuilder<'a> {
+ pub(crate) fn new(table: &'a Table) -> Self {
+ Self {
+ table,
+ routes: Vec::new(),
+ limit: None,
+ ranker: HybridSearchRanker::Rrf,
+ }
+ }
+
+ pub fn add_route(&mut self, route: HybridSearchRoute) -> &mut Self {
+ self.routes.push(route);
+ self
+ }
+
+ pub fn add_vector_route(
+ &mut self,
+ field_name: &str,
+ vector: Vec<f32>,
+ limit: usize,
+ weight: f32,
+ options: HashMap<String, String>,
+ ) -> crate::Result<&mut Self> {
+ self.routes.push(HybridSearchRoute::vector(
+ field_name, vector, limit, weight, options,
+ )?);
+ Ok(self)
+ }
+
+ pub fn add_full_text_route(
+ &mut self,
+ field_name: &str,
+ query: &str,
+ limit: usize,
+ weight: f32,
+ options: HashMap<String, String>,
+ ) -> crate::Result<&mut Self> {
+ self.routes.push(HybridSearchRoute::full_text(
+ field_name, query, limit, weight, options,
+ )?);
+ Ok(self)
+ }
+
+ pub fn with_limit(&mut self, limit: usize) -> &mut Self {
+ self.limit = Some(limit);
+ self
+ }
+
+ pub fn with_ranker(&mut self, ranker: &str) -> crate::Result<&mut Self> {
+ self.ranker = HybridSearchRanker::parse(ranker)?;
+ Ok(self)
+ }
+
+ pub fn with_rrf_ranker(&mut self) -> &mut Self {
+ self.ranker = HybridSearchRanker::Rrf;
+ self
+ }
+
+ pub fn with_weighted_score_ranker(&mut self) -> &mut Self {
+ self.ranker = HybridSearchRanker::WeightedScore;
+ self
+ }
+
+ pub fn with_mrr_ranker(&mut self) -> &mut Self {
+ self.ranker = HybridSearchRanker::Mrr;
+ self
+ }
+
+ pub async fn execute(&self) -> crate::Result<Vec<RowRange>> {
+ self.execute_scored().await?.to_row_ranges()
+ }
+
+ pub async fn execute_scored(&self) -> crate::Result<SearchResult> {
+
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
+ let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid {
+ message: "Limit must be set via with_limit()".to_string(),
+ })?;
+ if self.routes.is_empty() {
+ return Err(crate::Error::ConfigInvalid {
+ message: "Routes cannot be empty".to_string(),
+ });
+ }
+
+ let mut route_results = Vec::with_capacity(self.routes.len());
+ for route in &self.routes {
+ let result = match route.kind {
+ HybridSearchRouteKind::Vector => {
+ let mut builder = self.table.new_vector_search_builder();
+ builder
+ .with_vector_column(&route.field_name)
+
.with_query_vector(route.vector.clone().expect("validated vector route"))
+ .with_limit(route.limit)
+ .with_options(route.options.clone());
+ builder.execute_scored().await?
+ }
+ HybridSearchRouteKind::FullText => {
+ execute_full_text_route(self.table, route).await?
+ }
+ };
+ if !result.is_empty() {
+ route_results.push(WeightedRouteResult {
+ result,
+ weight: route.weight,
+ });
+ }
+ }
+
+ Ok(rank_results(self.ranker, &route_results, limit))
+ }
+}
+
+#[cfg(feature = "fulltext")]
+async fn execute_full_text_route(
+ table: &Table,
+ route: &HybridSearchRoute,
+) -> crate::Result<SearchResult> {
+ let mut builder = table.new_full_text_search_builder();
+ builder
+ .with_text_column(&route.field_name)
+ .with_query_text(
+ route
+ .full_text_query
+ .as_deref()
+ .expect("validated full-text route"),
+ )
+ .with_limit(route.limit);
+ let result = builder.execute_scored().await?;
+ Ok(SearchResult::new(result.row_ids, result.scores))
+}
+
+#[cfg(not(feature = "fulltext"))]
+async fn execute_full_text_route(
+ _table: &Table,
+ _route: &HybridSearchRoute,
+) -> crate::Result<SearchResult> {
+ Err(crate::Error::ConfigInvalid {
+ message: "Full-text hybrid routes require the fulltext
feature".to_string(),
+ })
+}
+
+struct WeightedRouteResult {
+ result: SearchResult,
+ weight: f32,
+}
+
+fn rank_results(
+ ranker: HybridSearchRanker,
+ route_results: &[WeightedRouteResult],
+ limit: usize,
+) -> SearchResult {
+ match ranker {
+ HybridSearchRanker::Rrf => rrf(route_results, limit),
+ HybridSearchRanker::WeightedScore => weighted_score(route_results,
limit),
+ HybridSearchRanker::Mrr => mrr(route_results, limit),
+ }
+}
+
+fn rrf(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult {
+ let mut scores = HashMap::new();
+ for route_result in route_results {
+ for (rank, (row_id, _score)) in
ranked_row_ids(&route_result.result).iter().enumerate() {
+ let contribution = route_result.weight / (RRF_K + rank as f32 +
1.0);
+ add_score(&mut scores, *row_id, contribution);
+ }
+ }
+ top_k(scores, limit)
+}
+
+fn mrr(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult {
+ let mut scores = HashMap::new();
+ for route_result in route_results {
+ for (rank, (row_id, _score)) in
ranked_row_ids(&route_result.result).iter().enumerate() {
+ let contribution = route_result.weight / (rank as f32 + 1.0);
+ add_score(&mut scores, *row_id, contribution);
+ }
+ }
+ top_k(scores, limit)
+}
+
+fn weighted_score(route_results: &[WeightedRouteResult], limit: usize) ->
SearchResult {
+ let mut scores = HashMap::new();
+ for route_result in route_results {
+ let ranked = ranked_row_ids(&route_result.result);
+ if ranked.is_empty() {
+ continue;
+ }
+
+ let (mut min, mut max) = (f32::INFINITY, f32::NEG_INFINITY);
+ for (_row_id, score) in &ranked {
+ min = min.min(*score);
+ max = max.max(*score);
+ }
+ let range = max - min;
+
+ for (row_id, score) in ranked {
+ let normalized = if range > 0.0 {
+ (score - min) / range
+ } else {
+ 1.0
+ };
+ add_score(&mut scores, row_id, route_result.weight * normalized);
+ }
+ }
+ top_k(scores, limit)
+}
+
+fn ranked_row_ids(result: &SearchResult) -> Vec<(u64, f32)> {
+ let mut best_scores = HashMap::new();
+ for (&row_id, &score) in result.row_ids.iter().zip(&result.scores) {
+ best_scores
+ .entry(row_id)
+ .and_modify(|old: &mut f32| {
+ if score > *old {
+ *old = score;
+ }
+ })
+ .or_insert(score);
+ }
+
+ let mut ranked: Vec<_> = best_scores.into_iter().collect();
+ ranked.sort_by(|(left_id, left_score), (right_id, right_score)| {
+ right_score
+ .partial_cmp(left_score)
+ .unwrap_or(std::cmp::Ordering::Equal)
+ .then_with(|| left_id.cmp(right_id))
+ });
+ ranked
+}
+
+fn add_score(scores: &mut HashMap<u64, f32>, row_id: u64, score: f32) {
+ scores
+ .entry(row_id)
+ .and_modify(|old_score| *old_score += score)
+ .or_insert(score);
+}
+
+fn top_k(scores: HashMap<u64, f32>, limit: usize) -> SearchResult {
+ if scores.is_empty() || limit == 0 {
+ return SearchResult::empty();
+ }
+
+ let mut entries: Vec<_> = scores.into_iter().collect();
+ entries.sort_by(|(left_id, left_score), (right_id, right_score)| {
+ right_score
+ .partial_cmp(left_score)
+ .unwrap_or(std::cmp::Ordering::Equal)
+ .then_with(|| left_id.cmp(right_id))
+ });
+ entries.truncate(limit);
+
+ let (row_ids, scores): (Vec<_>, Vec<_>) = entries.into_iter().unzip();
+ SearchResult::new(row_ids, scores)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn route_result(row_ids: Vec<u64>, scores: Vec<f32>, weight: f32) ->
WeightedRouteResult {
+ WeightedRouteResult {
+ result: SearchResult::new(row_ids, scores),
+ weight,
+ }
+ }
+
+ #[test]
+ fn test_rrf_prefers_overlap() {
+ let ranked = rank_results(
+ HybridSearchRanker::Rrf,
+ &[
+ route_result(vec![1, 2], vec![0.9, 0.8], 1.0),
+ route_result(vec![2, 3], vec![0.95, 0.1], 1.0),
+ ],
+ 1,
+ );
+
+ assert_eq!(ranked.row_ids, vec![2]);
+ }
+
+ #[test]
+ fn test_weighted_score_min_max_normalizes_per_route() {
+ let ranked = rank_results(
+ HybridSearchRanker::WeightedScore,
+ &[
+ route_result(vec![1, 2, 3], vec![10.0, 5.0, 0.0], 2.0),
+ route_result(vec![1, 2, 3], vec![100.0, 50.0, 0.0], 1.0),
+ ],
+ 3,
+ );
+
+ let scores: HashMap<_, _> =
ranked.row_ids.into_iter().zip(ranked.scores).collect();
+ assert!((scores[&1] - 3.0).abs() < 1e-6);
+ assert!((scores[&2] - 1.5).abs() < 1e-6);
+ assert!((scores[&3] - 0.0).abs() < 1e-6);
+ }
+
+ #[test]
+ fn test_mrr_uses_reciprocal_rank_without_constant() {
+ let ranked = rank_results(
+ HybridSearchRanker::Mrr,
+ &[
+ route_result(vec![1, 2], vec![0.9, 0.8], 1.0),
+ route_result(vec![2, 3], vec![0.95, 0.1], 1.0),
+ ],
+ 2,
+ );
+
+ assert_eq!(ranked.row_ids[0], 2);
+ assert!(ranked.scores[0] > ranked.scores[1]);
+ }
+}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 459e48a..930f9ef 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -38,6 +38,7 @@ mod data_file_writer;
#[cfg(feature = "fulltext")]
mod full_text_search_builder;
pub(crate) mod global_index_scanner;
+mod hybrid_search_builder;
mod kv_file_reader;
mod kv_file_writer;
mod lumina_index_build_builder;
@@ -76,6 +77,9 @@ pub use data_evolution_writer::{DataEvolutionDeleteWriter,
DataEvolutionWriter};
#[cfg(feature = "fulltext")]
pub use full_text_search_builder::FullTextSearchBuilder;
use futures::stream::BoxStream;
+pub use hybrid_search_builder::{
+ HybridSearchBuilder, HybridSearchRanker, HybridSearchRoute,
HybridSearchRouteKind,
+};
pub use lumina_index_build_builder::LuminaIndexBuildBuilder;
pub use read_builder::ReadBuilder;
pub use rest_env::RESTEnv;
@@ -185,6 +189,13 @@ impl Table {
FullTextSearchBuilder::new(self)
}
+ /// Create a hybrid search builder.
+ ///
+ /// Reference:
[HybridSearchBuilderImpl](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/table/source/HybridSearchBuilderImpl.java)
+ pub fn new_hybrid_search_builder(&self) -> HybridSearchBuilder<'_> {
+ HybridSearchBuilder::new(self)
+ }
+
pub fn new_vector_search_builder(&self) -> VectorSearchBuilder<'_> {
VectorSearchBuilder::new(self)
}
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index 33003c4..2384ce9 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -70,6 +70,7 @@ pub struct VectorSearchBuilder<'a> {
vector_column: Option<String>,
query_vector: Option<Vec<f32>>,
limit: Option<usize>,
+ options: HashMap<String, String>,
}
pub struct BatchVectorSearchBuilder<'a> {
@@ -77,6 +78,7 @@ pub struct BatchVectorSearchBuilder<'a> {
vector_column: Option<String>,
query_vectors: Option<Vec<Vec<f32>>>,
limit: Option<usize>,
+ options: HashMap<String, String>,
}
impl<'a> VectorSearchBuilder<'a> {
@@ -86,6 +88,7 @@ impl<'a> VectorSearchBuilder<'a> {
vector_column: None,
query_vector: None,
limit: None,
+ options: HashMap::new(),
}
}
@@ -104,7 +107,16 @@ impl<'a> VectorSearchBuilder<'a> {
self
}
+ pub fn with_options(&mut self, options: HashMap<String, String>) -> &mut
Self {
+ self.options = options;
+ self
+ }
+
pub async fn execute(&self) -> crate::Result<Vec<RowRange>> {
+ self.execute_scored().await?.to_row_ranges()
+ }
+
+ pub async fn execute_scored(&self) -> crate::Result<SearchResult> {
// Fail closed: returns data-derived row ranges outside
`TableScan`/`TableRead`.
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
let vector_column =
@@ -128,11 +140,12 @@ impl<'a> VectorSearchBuilder<'a> {
.with_vector_column(vector_column)
.with_query_vectors(vec![query_vector.clone()])
.with_limit(limit)
+ .with_options(self.options.clone())
.execute()
.await?;
debug_assert_eq!(results.len(), 1);
- results.remove(0).to_row_ranges()
+ Ok(results.remove(0))
}
}
@@ -143,6 +156,7 @@ impl<'a> BatchVectorSearchBuilder<'a> {
vector_column: None,
query_vectors: None,
limit: None,
+ options: HashMap::new(),
}
}
@@ -161,6 +175,11 @@ impl<'a> BatchVectorSearchBuilder<'a> {
self
}
+ pub fn with_options(&mut self, options: HashMap<String, String>) -> &mut
Self {
+ self.options = options;
+ self
+ }
+
pub async fn execute(&self) -> crate::Result<Vec<SearchResult>> {
let vector_column =
self.vector_column
@@ -192,7 +211,10 @@ impl<'a> BatchVectorSearchBuilder<'a> {
let vector_searches = query_vectors
.iter()
- .map(|vector| VectorSearch::new(vector.clone(), limit,
vector_column.to_string()))
+ .map(|vector| {
+ VectorSearch::new(vector.clone(), limit,
vector_column.to_string())
+ .map(|search| search.with_options(self.options.clone()))
+ })
.collect::<crate::Result<Vec<_>>>()?;
let snapshot_manager = SnapshotManager::new(
@@ -282,6 +304,17 @@ async fn evaluate_batch_vector_search(
source: None,
});
}
+ let search_options = vector_searches[0].options.clone();
+ if vector_searches
+ .iter()
+ .any(|vector_search| vector_search.options != search_options)
+ {
+ return Err(crate::Error::DataInvalid {
+ message: "Batch vector search requires all query vectors to use
the same options"
+ .to_string(),
+ source: None,
+ });
+ }
let field_id = match find_field_id_by_name(evaluation.schema_fields,
field_name) {
Some(id) => id,
@@ -346,7 +379,8 @@ async fn evaluate_batch_vector_search(
for vector_search in &mut vector_searches {
vector_search.limit = index_limit;
}
- let options = evaluation.table_options.clone();
+ let mut options = evaluation.table_options.clone();
+ options.extend(search_options.clone());
let input = evaluation.file_io.new_input(&path);
async move {
let input = input?;
diff --git a/crates/paimon/src/vector_search.rs
b/crates/paimon/src/vector_search.rs
index 28b6500..c96edbf 100644
--- a/crates/paimon/src/vector_search.rs
+++ b/crates/paimon/src/vector_search.rs
@@ -22,6 +22,7 @@ pub struct VectorSearch {
pub vector: Vec<f32>,
pub limit: usize,
pub field_name: String,
+ pub options: HashMap<String, String>,
pub include_row_ids: Option<roaring::RoaringTreemap>,
}
@@ -49,10 +50,16 @@ impl VectorSearch {
vector,
limit,
field_name,
+ options: HashMap::new(),
include_row_ids: None,
})
}
+ pub fn with_options(mut self, options: HashMap<String, String>) -> Self {
+ self.options = options;
+ self
+ }
+
pub fn with_include_row_ids(mut self, include_row_ids:
roaring::RoaringTreemap) -> Self {
self.include_row_ids = Some(include_row_ids);
self
@@ -248,6 +255,7 @@ mod tests {
assert_eq!(cloned.vector, vector_search.vector);
assert_eq!(cloned.limit, vector_search.limit);
assert_eq!(cloned.field_name, vector_search.field_name);
+ assert_eq!(cloned.options, vector_search.options);
assert_eq!(cloned.include_row_ids.as_ref(), Some(&include_row_ids));
}
diff --git a/docs/src/sql.md b/docs/src/sql.md
index e983274..70b8099 100644
--- a/docs/src/sql.md
+++ b/docs/src/sql.md
@@ -78,7 +78,7 @@ async fn example() -> Result<(), Box<dyn std::error::Error>> {
}
```
-`SQLContext::new` creates a session context with the Paimon relation planner
pre-registered. Use `register_catalog(...).await` to add one or more Paimon
catalogs; registering a catalog also registers the built-in table-valued
functions (`vector_search`, `full_text_search`) against it. It also manages
session-scoped dynamic options internally for `SET`/`RESET` support.
+`SQLContext::new` creates a session context with the Paimon relation planner
pre-registered. Use `register_catalog(...).await` to add one or more Paimon
catalogs; registering a catalog also registers the built-in table-valued
functions (`vector_search`, `hybrid_search`, and `full_text_search` when the
`fulltext` feature is enabled) against it. It also manages session-scoped
dynamic options internally for `SET`/`RESET` support.
## Data Types
@@ -852,6 +852,99 @@ Vector index behavior is configured via table options
prefixed with `lumina.`:
The Lumina native library must be available at runtime. Set the
`LUMINA_LIB_PATH` environment variable to the path of the shared library, or
place it in the platform default location.
+## Hybrid Search
+
+Paimon supports hybrid search by combining multiple search routes and ranking
the merged results. The `hybrid_search` table-valued function is registered as
a UDTF on the DataFusion session context.
+
+Hybrid search does not require the `fulltext` feature when all routes are
vector routes. Enable `fulltext` only when you include full-text routes.
+
+### Registration
+
+When you use a `SQLContext`, `hybrid_search` is registered automatically for
every catalog you register — no extra setup is needed.
+
+With a raw DataFusion `SessionContext`, register it explicitly:
+
+```rust
+use paimon_datafusion::register_hybrid_search;
+
+register_hybrid_search(&ctx, catalog.clone(), "default");
+```
+
+### Usage
+
+```sql
+SELECT * FROM hybrid_search(
+ 'table_name',
+ vector_routes,
+ full_text_routes,
+ limit,
+ 'ranker'
+)
+```
+
+| Argument | Type | Description |
+|---|---|---|
+| `table_name` | STRING | Table name, fully qualified (`catalog.db.table`) or
short form |
+| `vector_routes` | ARRAY | Vector route definitions; use `array()` when no
vector route is needed |
+| `full_text_routes` | ARRAY | Full-text route definitions; use `array()` for
vector-only hybrid search |
+| `limit` | INT | Maximum number of merged results (top-k) |
+| `ranker` | STRING | Optional ranker: `rrf` (default), `weighted_score`, or
`mrr` |
+
+Route definitions use Spark-compatible `array(named_struct(...))` syntax. A
vector route accepts `field` (or `vector_column`), `query_vector`, optional
`limit`, optional `weight`, and optional `options`:
+
+```sql
+SELECT *
+FROM hybrid_search(
+ 'paimon.my_db.items',
+ array(
+ named_struct(
+ 'field', 'title_embedding',
+ 'query_vector', array(1.0, 0.0, 0.0, 0.0),
+ 'limit', 20,
+ 'weight', 1.0
+ ),
+ named_struct(
+ 'field', 'body_embedding',
+ 'query_vector', array(0.9, 0.1, 0.0, 0.0),
+ 'limit', 20,
+ 'weight', 0.7
+ )
+ ),
+ array(),
+ 10,
+ 'rrf'
+);
+```
+
+A full-text route accepts `column`, `query`, optional `limit`, and optional
`weight`. Full-text routes require the `fulltext` feature:
+
+```sql
+SELECT *
+FROM hybrid_search(
+ 'paimon.my_db.docs',
+ array(
+ named_struct(
+ 'field', 'embedding',
+ 'query_vector', array(1.0, 0.0, 0.0, 0.0),
+ 'limit', 20,
+ 'weight', 1.0
+ )
+ ),
+ array(
+ named_struct(
+ 'column', 'content',
+ 'query', 'paimon search',
+ 'limit', 20,
+ 'weight', 0.8
+ )
+ ),
+ 10,
+ 'weighted_score'
+);
+```
+
+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.
+
## Full-Text Search
Paimon supports full-text search via the Tantivy search engine. The
`full_text_search` table-valued function is registered as a UDTF on the
DataFusion session context.
