diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 9279e8940..15ca03ce6 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -222,6 +222,7 @@ class Table: def new_read_builder(self, options: Optional[Dict[str, str]] = None) -> ReadBuilder: ... def new_stream_scan(self) -> StreamTableScan: ... def new_full_text_search_builder(self) -> FullTextSearchBuilder: ... + def new_hybrid_search_builder(self) -> HybridSearchBuilder: ... def new_vector_search_builder(self) -> "VectorSearchBuilder": ... def new_batch_vector_search_builder(self) -> "BatchVectorSearchBuilder": ... def new_global_index_build_builder(self) -> "GlobalIndexBuildBuilder": ... @@ -615,6 +616,26 @@ class SearchResultReadBuilder: """Read user columns and scores in relevance order from the search snapshot.""" ... +class HybridSearchBuilder: + def with_limit(self, limit: int) -> HybridSearchBuilder: ... + def with_snapshot(self, snapshot_json: Optional[str]) -> HybridSearchBuilder: + """Reuse resolved Java Snapshot JSON; None pins an empty read view.""" + ... + def with_ranker(self, ranker: Optional[str]) -> HybridSearchBuilder: ... + def with_rrf_ranker(self) -> HybridSearchBuilder: ... + def with_weighted_score_ranker(self) -> HybridSearchBuilder: ... + def with_filter(self, predicate: Dict[str, Any]) -> HybridSearchBuilder: ... + def with_partition_filter(self, predicate: Dict[str, Any]) -> HybridSearchBuilder: ... + def add_vector_route(self, field_name: str, vector: List[float], limit: int, + weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... + def add_full_text_route(self, field_name: str, query: str, limit: int, + weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... + def execute_local(self) -> HybridSearchResult: ... + +class HybridSearchResult: + def __len__(self) -> int: ... + def row_ids(self) -> Dict[int, float]: ... + class FullTextSearchBuilder: def with_query(self, field_name: str, query: str) -> FullTextSearchBuilder: ... def with_limit(self, limit: int) -> FullTextSearchBuilder: ... diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 4c6637a45..ba8c3d87a 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -732,6 +732,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> this.add_class::()?; this.add_class::()?; this.add_class::()?; + this.add_class::()?; + this.add_class::()?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs new file mode 100644 index 000000000..749bdd93f --- /dev/null +++ b/bindings/python/src/hybrid_search.rs @@ -0,0 +1,206 @@ +// 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 configuration forwarding; snapshot, routes and fusion belong to core. + +use crate::error::to_py_err; +use crate::predicate::dict_to_table_predicate; +use paimon::spec::{Predicate, Snapshot}; +use paimon::table::{HybridSearchRanker, HybridSearchRoute, Table}; +use paimon::vector_search::ScoredRowIds; +use paimon_datafusion::runtime::runtime; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use std::collections::HashMap; +use std::sync::Arc; + +#[pyclass(name = "HybridSearchBuilder", module = "pypaimon_rust.datafusion")] +pub struct PyHybridSearchBuilder { + table: Arc, + routes: Vec, + filters: Vec, + limit: Option, + ranker: HybridSearchRanker, + pinned_snapshot: Option>, +} + +impl PyHybridSearchBuilder { + pub(crate) fn new(table: Arc
) -> Self { + Self { + table, + routes: Vec::new(), + filters: Vec::new(), + limit: None, + ranker: HybridSearchRanker::Rrf, + pinned_snapshot: None, + } + } + + fn builder(&self) -> paimon::table::HybridSearchBuilder<'_> { + let mut builder = self.table.new_hybrid_search_builder(); + for route in &self.routes { + builder.add_route(route.clone()); + } + if let Some(limit) = self.limit { + builder.with_limit(limit); + } + // The stored ranker has already been validated by the core parser. + builder + .with_ranker(self.ranker.as_str()) + .expect("validated ranker"); + for filter in &self.filters { + builder.with_filter(filter.clone()); + } + if let Some(snapshot) = &self.pinned_snapshot { + builder.with_snapshot(snapshot.as_ref()); + } + builder + } +} + +#[pymethods] +impl PyHybridSearchBuilder { + #[pyo3(signature = (field_name, vector, limit, weight=1.0, options=None))] + fn add_vector_route( + mut slf: PyRefMut<'_, Self>, + field_name: String, + vector: Vec, + limit: usize, + weight: f32, + options: Option>, + ) -> PyResult> { + let route = HybridSearchRoute::vector( + field_name, + vector, + limit, + weight, + options.unwrap_or_default(), + ) + .map_err(to_py_err)?; + slf.routes.push(route); + Ok(slf) + } + + #[pyo3(signature = (field_name, query, limit, weight=1.0, options=None))] + fn add_full_text_route( + mut slf: PyRefMut<'_, Self>, + field_name: String, + query: String, + limit: usize, + weight: f32, + options: Option>, + ) -> PyResult> { + let route = HybridSearchRoute::full_text( + field_name, + query, + limit, + weight, + options.unwrap_or_default(), + ) + .map_err(to_py_err)?; + slf.routes.push(route); + Ok(slf) + } + + fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyRefMut<'_, Self> { + slf.limit = Some(limit); + slf + } + + #[pyo3(signature = (snapshot_json))] + fn with_snapshot<'py>( + mut slf: PyRefMut<'py, Self>, + snapshot_json: Option<&str>, + ) -> PyResult> { + let snapshot = snapshot_json + .map(serde_json::from_str) + .transpose() + .map_err(|error| { + pyo3::exceptions::PyValueError::new_err(format!("Invalid snapshot JSON: {error}")) + })?; + slf.pinned_snapshot = Some(snapshot); + Ok(slf) + } + + #[pyo3(signature = (ranker))] + fn with_ranker<'py>( + mut slf: PyRefMut<'py, Self>, + ranker: Option<&str>, + ) -> PyResult> { + slf.ranker = HybridSearchRanker::parse(ranker.unwrap_or_default()).map_err(to_py_err)?; + Ok(slf) + } + + fn with_rrf_ranker(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> { + slf.ranker = HybridSearchRanker::Rrf; + slf + } + + fn with_weighted_score_ranker(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> { + slf.ranker = HybridSearchRanker::WeightedScore; + slf + } + + fn with_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let filter = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.filters.push(filter); + Ok(slf) + } + + fn with_partition_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let filter = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.builder() + .with_partition_filter(filter.clone()) + .map_err(to_py_err)?; + slf.filters.push(filter); + Ok(slf) + } + + fn execute_local(&self, py: Python<'_>) -> PyResult { + let builder = self.builder(); + py.detach(|| runtime().block_on(builder.execute_scored())) + .map(|inner| PyHybridSearchResult { inner }) + .map_err(to_py_err) + } +} + +#[pyclass(name = "HybridSearchResult", module = "pypaimon_rust.datafusion")] +pub struct PyHybridSearchResult { + inner: ScoredRowIds, +} + +#[pymethods] +impl PyHybridSearchResult { + fn __len__(&self) -> usize { + self.inner.len() + } + + fn row_ids(&self) -> HashMap { + self.inner + .row_ids + .iter() + .copied() + .zip(self.inner.scores.iter().copied()) + .collect() + } +} diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e4a5c0270..87a3362c5 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -22,6 +22,7 @@ mod blob_uri_reader; mod context; mod error; mod full_text_search; +mod hybrid_search; mod index_build; mod merge_into; mod oss_cpp; diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs index e0daf6ea4..2c33c4b0a 100644 --- a/bindings/python/src/table.rs +++ b/bindings/python/src/table.rs @@ -202,6 +202,10 @@ impl PyTable { crate::full_text_search::PyFullTextSearchBuilder::new(Arc::clone(&self.inner)) } + fn new_hybrid_search_builder(&self) -> crate::hybrid_search::PyHybridSearchBuilder { + crate::hybrid_search::PyHybridSearchBuilder::new(Arc::clone(&self.inner)) + } + fn new_vector_search_builder(&self) -> crate::vector_search::PyVectorSearchBuilder { crate::vector_search::PyVectorSearchBuilder::new(Arc::clone(&self.inner)) } diff --git a/crates/integrations/datafusion/src/hybrid_search.rs b/crates/integrations/datafusion/src/hybrid_search.rs index 1912c5f58..76a2889a7 100644 --- a/crates/integrations/datafusion/src/hybrid_search.rs +++ b/crates/integrations/datafusion/src/hybrid_search.rs @@ -22,7 +22,7 @@ //! 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')), +//! array(named_struct('column', 'content', 'query', '{"match":{"query":"paimon"}}')), //! 10, //! 'rrf') //! ``` @@ -53,7 +53,9 @@ 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::{datafusion_read_fields, PaimonScanBuilder, PaimonTableProvider}; +use crate::table::{ + datafusion_arrow_schema, datafusion_read_fields, PaimonScanBuilder, PaimonTableProvider, +}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; @@ -248,8 +250,7 @@ impl TableProvider for HybridSearchTableProvider { 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 input_schema = datafusion_arrow_schema(&input_read_fields, true)?; let row_id_table_index = input_read_fields .iter() .position(|field| field.name() == ROW_ID_FIELD_NAME) diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index f867b07d8..caa00fa9d 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -2751,7 +2751,7 @@ mod hybrid_search_tests { use std::sync::Arc; #[cfg(feature = "fulltext")] - use datafusion::arrow::array::{Array, Int64Array}; + use datafusion::arrow::array::{Array, Int64Array, StringViewArray}; use datafusion::arrow::array::{Float32Array, Int32Array}; use paimon::catalog::Identifier; use paimon::table::BranchManager; @@ -2904,7 +2904,7 @@ mod hybrid_search_tests { assert!((score - expected).abs() < 1e-6); } - let empty_batches = ctx + let missing_error = ctx .sql( "SELECT __paimon_search_score FROM hybrid_search( \ 'paimon.default.test_java_vindex_vector', \ @@ -2918,10 +2918,29 @@ mod hybrid_search_tests { 'rrf')", ) .await - .expect("empty hybrid_search score SQL should parse") + .expect("invalid-column hybrid_search SQL should parse") + .collect() + .await + .expect_err("an unknown vector column must not silently return empty"); + assert!(missing_error.to_string().contains("does not exist")); + super::common::exec( + &ctx, + "CREATE TABLE paimon.default.hybrid_empty (id INT, embedding ARRAY) \ + WITH ('row-tracking.enabled' = 'true')", + ) + .await; + let empty_batches = ctx + .sql( + "SELECT __paimon_search_score FROM hybrid_search( \ + 'paimon.default.hybrid_empty', \ + array(named_struct('field', 'embedding', 'query_vector', array(1.0), 'limit', 3)), \ + array(), 3, 'rrf')", + ) + .await + .unwrap() .collect() .await - .expect("empty hybrid_search score query should execute"); + .unwrap(); assert_eq!( empty_batches .iter() @@ -2933,19 +2952,22 @@ mod hybrid_search_tests { #[cfg(feature = "fulltext")] #[tokio::test] - async fn test_hybrid_search_score_with_row_tracking_only_fulltext() { + async fn test_hybrid_search_score_with_fulltext_definition() { 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', \ + 'bucket' = '-1', \ + 'data-evolution.enabled' = 'true', \ + 'global-index.enabled' = 'true', \ 'global-index.search-mode' = 'full')", ) .await; super::common::exec( &ctx, - "INSERT INTO paimon.test_db.hybrid_raw_fulltext VALUES \ + "INSERT INTO paimon.test_db.hybrid_raw_fulltext (id, content) VALUES \ (1, 'paimon search'), (2, 'other'), (3, 'paimon table')", ) .await; @@ -2955,11 +2977,26 @@ mod hybrid_search_tests { array(), \ array(named_struct( \ 'column', 'content', \ - 'query', 'paimon', \ + 'query', '{\"match\":{\"query\":\"paimon\"}}', \ 'limit', 10, \ 'weight', 1.0)), \ 10, \ 'rrf')"; + // Java returns empty without a full-text definition, even in FULL mode. + let before = ctx + .sql(&format!("SELECT id FROM {SEARCH}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert!(before.iter().all(|batch| batch.num_rows() == 0)); + super::common::exec( + &ctx, + "CALL sys.create_global_index(table => 'test_db.hybrid_raw_fulltext', \ + index_column => 'content', index_type => 'full-text')", + ) + .await; let explicit_sql = format!("SELECT id, __paimon_search_score FROM {SEARCH} ORDER BY id"); let explicit_batches = ctx .sql(&explicit_sql) @@ -3038,10 +3075,25 @@ mod hybrid_search_tests { 2 ); for batch in &star_batches { - assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_columns(), 4); assert!(batch.column_by_name("__paimon_search_score").is_some()); - assert!(batch.column_by_name("_ROW_ID").is_none()); + assert!(batch.column_by_name("_ROW_ID").is_some()); } + let mut text = star_batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("content") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|value| value.unwrap().to_string()) + }) + .collect::>(); + text.sort(); + assert_eq!(text, vec!["paimon search", "paimon table"]); } #[tokio::test] diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index f6be569f1..e8838214a 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -22,9 +22,11 @@ use std::collections::{BTreeMap, HashMap}; use arrow_array::RecordBatch; -use futures::{stream, TryStreamExt}; +use futures::{stream, StreamExt, TryStreamExt}; -use crate::spec::{CoreOptions, SCAN_SNAPSHOT_ID_OPTION}; +#[cfg(all(test, feature = "fulltext"))] +use crate::spec::SCAN_SNAPSHOT_ID_OPTION; +use crate::spec::{CoreOptions, Predicate, Snapshot}; use crate::table::data_file_reader::DataFileReader; use crate::table::pk_search_position::PrimaryKeySearchPosition; use crate::table::pk_search_ranker::{self, Ranking}; @@ -46,6 +48,10 @@ use crate::table::pk_full_text_read::{build_full_text_indexed_splits, PrimaryKey use crate::table::pk_full_text_scan::PrimaryKeyFullTextScan; const RRF_K: f32 = 60.0; +const MAX_ROUTE_WORKERS: usize = 4; + +#[cfg(all(test, feature = "fulltext"))] +mod parity_tests; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum HybridSearchRanker { @@ -210,6 +216,8 @@ pub struct HybridSearchBuilder<'a> { routes: Vec, limit: Option, ranker: HybridSearchRanker, + filter: Option, + pinned_snapshot: Option>, } impl<'a> HybridSearchBuilder<'a> { @@ -219,6 +227,8 @@ impl<'a> HybridSearchBuilder<'a> { routes: Vec::new(), limit: None, ranker: HybridSearchRanker::Rrf, + filter: None, + pinned_snapshot: None, } } @@ -260,6 +270,26 @@ impl<'a> HybridSearchBuilder<'a> { self } + /// Reuse a read view already resolved by the caller, preserving the current + /// query schema. `None` pins an empty view even if a later commit appears. + pub fn with_snapshot(&mut self, snapshot: Option<&Snapshot>) -> &mut Self { + self.pinned_snapshot = Some(snapshot.cloned()); + self + } + + /// Apply the same pre-filter to every route before its own Top-K. + /// Partition conjuncts are pruned by each route's scan. + pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { + super::vector_search_common::add_filter(&mut self.filter, filter); + self + } + + /// Add a partition-only predicate, using table field indices like other builders. + pub fn with_partition_filter(&mut self, filter: Predicate) -> crate::Result<&mut Self> { + super::partition_filter::validate_partition_filter(self.table, &filter)?; + Ok(self.with_filter(filter)) + } + pub fn with_ranker(&mut self, ranker: &str) -> crate::Result<&mut Self> { self.ranker = HybridSearchRanker::parse(ranker)?; Ok(self) @@ -287,14 +317,7 @@ impl<'a> HybridSearchBuilder<'a> { pub async fn execute_scored(&self) -> crate::Result { let core = CoreOptions::new(self.table.schema().options()); core.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 limit = self.validate_search()?; // A primary-key hybrid fuses PHYSICAL positions, not global row ids, so a // scored/row-range result is unsupported on it: fail loud and direct callers @@ -313,30 +336,109 @@ impl<'a> HybridSearchBuilder<'a> { HybridAddressSpace::Global => {} } - let mut route_results = Vec::with_capacity(self.routes.len()); + self.validate_global_routes()?; + // Resolve once even for an empty table or a retained tag. Re-resolving + // selectors independently in the routes can observe different versions. + let Some(table) = self.resolve_pinned_route_table().await? else { + return Ok(ScoredRowIds::empty()); + }; + self.execute_global_routes(&table, limit).await + } + + fn validate_search(&self) -> crate::Result { + if self.routes.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "Routes cannot be empty".into(), + }); + } + self.limit + .filter(|limit| *limit > 0) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Limit must be positive, set via with_limit()".into(), + }) + } + + fn vector_route_builder<'t>( + &self, + table: &'t Table, + route: &HybridSearchRoute, + ) -> super::VectorSearchBuilder<'t> { + let mut builder = 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()); + if let Some(filter) = &self.filter { + builder.with_filter(filter.clone()); + } + builder + } + + fn validate_global_routes(&self) -> crate::Result<()> { for route in &self.routes { - let result = match route.kind { + 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().await?.into_row_ids()? + if !self + .table + .schema() + .fields() + .iter() + .any(|field| field.name() == route.field_name) + { + return Err(crate::Error::ConfigInvalid { + message: format!("Vector column '{}' does not exist", route.field_name), + }); + } + let builder = self.vector_route_builder(self.table, route); + builder.new_scan()?; + builder.new_read()?; } HybridSearchRouteKind::FullText => { - execute_full_text_route(self.table, route).await? + #[cfg(feature = "fulltext")] + { + let builder = + full_text_route_builder(self.table, route, self.filter.as_ref()); + builder.new_scan()?; + builder.new_read()?; + } + #[cfg(not(feature = "fulltext"))] + return Err(crate::Error::ConfigInvalid { + message: "Full-text hybrid routes require the fulltext feature".into(), + }); } - }; - if !result.is_empty() { - route_results.push(WeightedRouteResult { - result, - weight: route.weight, - }); } } + Ok(()) + } + async fn execute_global_routes( + &self, + table: &Table, + limit: usize, + ) -> crate::Result { + // Bound concurrent route work while retaining input order for float + // accumulation and keeping each route's weight and limit attached to it. + let route_results = stream::iter(self.routes.clone()) + .map(|route| async move { + let result = match route.kind { + HybridSearchRouteKind::Vector => self + .vector_route_builder(table, &route) + .execute() + .await? + .into_row_ids()?, + HybridSearchRouteKind::FullText => { + execute_full_text_route(table, &route, self.filter.as_ref()).await? + } + }; + Ok::<_, crate::Error>(WeightedRouteResult { + result, + weight: route.weight, + }) + }) + .buffered(MAX_ROUTE_WORKERS) + .try_collect::>() + .await?; Ok(rank_results(self.ranker, &route_results, limit)) } @@ -351,14 +453,7 @@ impl<'a> HybridSearchBuilder<'a> { pub async fn execute_read(&self) -> crate::Result { let core = CoreOptions::new(self.table.schema().options()); core.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 limit = self.validate_search()?; match self.classify_routes(&core)? { HybridAddressSpace::PrimaryKey => { @@ -418,27 +513,27 @@ impl<'a> HybridSearchBuilder<'a> { // fused score. Reject it up front — before any route runs and even when the // fused result is empty — reusing the primary-key vector read's guard. ensure_no_reserved_read_columns(self.table.schema().fields())?; + self.validate_primary_key_routes(core)?; // Snapshot pinning: resolve ONE snapshot for the whole primary-key hybrid // read and plan every route against it, mirroring Java // `HybridSearchBuilderImpl.routeBuilders()` (which resolves the snapshot // once and injects it into every route builder). Only the "read latest" - // case is racy — a concurrent commit landing between the two route plans - // would otherwise pick different snapshots — because every time-travel - // selector resolves deterministically to the same snapshot on each plan. - // So pin latest once via `scan.snapshot-id` and leave the already-pinned - // (time-travel) paths to resolve their own fixed snapshot. - let pinned_table = self.resolve_pinned_route_table().await?; - let route_table: &Table = pinned_table.as_ref().unwrap_or(self.table); + // case is racy: a concurrent commit landing between the two route plans + // would otherwise pick different snapshots. Preserve the resolved snapshot + // itself so retained tags also work after the original snapshot expires. + let Some(route_table) = self.resolve_pinned_route_table().await? else { + return Ok(Box::pin(stream::empty())); + }; // Per-route search: positions (converted from candidates) + single-file // source splits + the snapshot each route's plan pinned. let mut routes: Vec = Vec::with_capacity(self.routes.len()); for route in &self.routes { let pk_route = match route.kind { - HybridSearchRouteKind::Vector => self.pk_vector_route(route_table, route).await?, + HybridSearchRouteKind::Vector => self.pk_vector_route(&route_table, route).await?, HybridSearchRouteKind::FullText => { - self.pk_full_text_route(route_table, core, route).await? + self.pk_full_text_route(&route_table, core, route).await? } }; routes.push(pk_route); @@ -522,42 +617,23 @@ impl<'a> HybridSearchBuilder<'a> { Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) } - /// Resolve the ONE snapshot every primary-key route must plan against, as an + /// Resolve the ONE snapshot every route must plan against, as an /// optional pinned table copy. Mirrors Java /// `HybridSearchBuilderImpl.routeBuilders()`, which resolves a single snapshot /// up front and injects it into every route builder. /// - /// A time-travel selector (`scan.version` / `scan.timestamp-millis` / - /// `scan.snapshot-id` / `scan.tag-name`) already resolves deterministically to - /// the same snapshot on every route plan, so no extra pinning is needed — - /// return `None` and let each route resolve it. Only the default "read latest" - /// path is racy: a concurrent commit landing between two route plans would let - /// them resolve different snapshots. For that path resolve the latest snapshot - /// id ONCE and return a table copy pinned to it via `scan.snapshot-id`, so both - /// routes plan the same version. A table with no snapshot at all also returns - /// `None` (nothing to pin; every route plans empty). + /// Cache the resolved snapshot rather than just its id: retained tags must + /// remain readable even when the original snapshot file has expired. + /// `None` means the one resolved view is empty. async fn resolve_pinned_route_table(&self) -> crate::Result> { - let core = CoreOptions::new(self.table.schema().options()); - // Already targeting a fixed snapshot (resolved travel copy or a selector - // that resolves deterministically): every route agrees without pinning. - if self.table.has_resolved_travel_snapshot() || core.has_time_travel_selector() { - return Ok(None); + if let Some(snapshot) = &self.pinned_snapshot { + return Ok(snapshot + .as_ref() + .map(|snapshot| self.table.copy_with_pinned_snapshot(snapshot))); } - // Read-latest: pin the current latest snapshot once so a concurrent commit - // cannot split the routes across versions. - let Some(latest) = self - .table - .snapshot_manager() - .get_latest_snapshot_id() + Ok(super::time_travel::resolve_snapshot(self.table) .await? - else { - return Ok(None); - }; - let pinned = self.table.copy_with_options(HashMap::from([( - SCAN_SNAPSHOT_ID_OPTION.to_string(), - latest.to_string(), - )])); - Ok(Some(pinned)) + .map(|snapshot| self.table.copy_with_pinned_snapshot(&snapshot))) } /// Consume the vector search's scored positions, retaining its source files @@ -567,13 +643,7 @@ impl<'a> HybridSearchBuilder<'a> { table: &Table, route: &HybridSearchRoute, ) -> crate::Result { - let vector = route.vector.as_deref().expect("validated vector route"); - let mut builder = table.new_vector_search_builder(); - builder - .with_vector_column(&route.field_name) - .with_query_vector(vector.to_vec()) - .with_limit(route.limit) - .with_options(route.options.clone()); + let builder = self.vector_route_builder(table, route); let result = builder.execute().await?; let positions = result .positions()? @@ -593,13 +663,29 @@ impl<'a> HybridSearchBuilder<'a> { }) } + fn validate_primary_key_routes(&self, core: &CoreOptions<'_>) -> crate::Result<()> { + for route in &self.routes { + match route.kind { + HybridSearchRouteKind::Vector => { + let builder = self.vector_route_builder(self.table, route); + builder.new_scan()?; + builder.new_read()?; + } + HybridSearchRouteKind::FullText => { + self.pk_full_text_field(self.table, core, route)?; + } + } + } + Ok(()) + } + #[cfg(feature = "fulltext")] - async fn pk_full_text_route( + fn pk_full_text_field( &self, table: &Table, core: &CoreOptions<'_>, route: &HybridSearchRoute, - ) -> crate::Result { + ) -> crate::Result { // FAST-only: the primary-key full-text read searches compaction-visible // payloads and rejects FULL/DETAIL loud rather than silently degrading, // mirroring `full_text_search_builder::execute_read` and Java @@ -613,10 +699,6 @@ impl<'a> HybridSearchBuilder<'a> { }); } - let query = route - .full_text_query - .as_deref() - .expect("validated full-text route"); let field_id = find_field_id_by_name(table.schema().fields(), &route.field_name) .ok_or_else(|| crate::Error::DataInvalid { message: format!( @@ -625,7 +707,47 @@ impl<'a> HybridSearchBuilder<'a> { ), source: None, })?; - let plan = PrimaryKeyFullTextScan::new(table, field_id, None) + if let Some(filter) = &self.filter { + let mapping = crate::spec::field_idx_to_partition_idx( + table.schema().fields(), + table.schema().partition_keys(), + ); + if !filter.references_only_mapped_fields(&mapping) { + return Err(crate::Error::Unsupported { + message: + "Primary-key full-text search does not support non-partition filters yet." + .into(), + }); + } + } + Ok(field_id) + } + + #[cfg(not(feature = "fulltext"))] + fn pk_full_text_field( + &self, + _table: &Table, + _core: &CoreOptions<'_>, + _route: &HybridSearchRoute, + ) -> crate::Result { + Err(crate::Error::ConfigInvalid { + message: "primary-key full-text hybrid routes require the fulltext feature".into(), + }) + } + + #[cfg(feature = "fulltext")] + async fn pk_full_text_route( + &self, + table: &Table, + core: &CoreOptions<'_>, + route: &HybridSearchRoute, + ) -> crate::Result { + let field_id = self.pk_full_text_field(table, core, route)?; + let query = route + .full_text_query + .as_deref() + .expect("validated full-text route"); + let plan = PrimaryKeyFullTextScan::new(table, field_id, self.filter.clone()) .plan() .await?; let materialize_reader = DataFileReader::new( @@ -892,20 +1014,34 @@ fn build_hybrid_indexed_splits( } #[cfg(feature = "fulltext")] -async fn execute_full_text_route( - table: &Table, +fn full_text_route_builder<'t>( + table: &'t Table, route: &HybridSearchRoute, -) -> crate::Result { + filter: Option<&Predicate>, +) -> super::FullTextSearchBuilder<'t> { let mut builder = table.new_full_text_search_builder(); builder - .with_text_column(&route.field_name) - .with_query_text( + .with_query( + &route.field_name, route .full_text_query .as_deref() .expect("validated full-text route"), ) .with_limit(route.limit); + if let Some(filter) = filter { + builder.with_filter(filter.clone()); + } + builder +} + +#[cfg(feature = "fulltext")] +async fn execute_full_text_route( + table: &Table, + route: &HybridSearchRoute, + filter: Option<&Predicate>, +) -> crate::Result { + let builder = full_text_route_builder(table, route, filter); let result = builder.execute_scored().await?; Ok(ScoredRowIds::new(result.row_ids, result.scores)) } @@ -914,6 +1050,7 @@ async fn execute_full_text_route( async fn execute_full_text_route( _table: &Table, _route: &HybridSearchRoute, + _filter: Option<&Predicate>, ) -> crate::Result { Err(crate::Error::ConfigInvalid { message: "Full-text hybrid routes require the fulltext feature".to_string(), @@ -1541,6 +1678,18 @@ mod pk_hybrid_tests { "score column must be present" ); } + builder.with_snapshot(None); + let empty: Vec = builder + .execute_read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert!( + empty.is_empty(), + "a pinned empty view must not observe the existing rows" + ); } #[tokio::test] @@ -1861,6 +2010,13 @@ mod pk_hybrid_tests { format!("{err:?}").contains("FAST"), "PK hybrid full-text route under FULL mode must fail loud, got: {err:?}" ); + builder.with_snapshot(None); + let err = builder + .execute_read() + .await + .err() + .expect("empty views must still validate the mode"); + assert!(err.to_string().contains("FAST")); } // (g) Snapshot pinning: the read-latest path resolves ONE snapshot up front and diff --git a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs new file mode 100644 index 000000000..df2f8e255 --- /dev/null +++ b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs @@ -0,0 +1,393 @@ +// 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. + +//! Global hybrid execution against real indexes, matching Java route semantics. + +use super::*; +use crate::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::{ + ArrayType, DataType, Datum, FloatType, IntType, PredicateBuilder, Schema, TableSchema, + VarCharType, +}; +use crate::table::{TableCommit, TableWrite}; +use arrow_array::types::Float32Type; +use arrow_array::{ArrayRef, Int32Array, ListArray, StringArray}; +use arrow_schema::{DataType as ArrowType, Field, Schema as ArrowSchema}; +use std::sync::Arc; + +fn table() -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("text", DataType::VarChar(VarCharType::string_type())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .column("label", DataType::VarChar(VarCharType::string_type())) + .column("pt", DataType::Int(IntType::new())) + .partition_keys(["pt"]) + .options(HashMap::from([ + ("bucket".into(), "-1".into()), + ("row-tracking.enabled".into(), "true".into()), + ("data-evolution.enabled".into(), "true".into()), + ("global-index.enabled".into(), "true".into()), + ("scalar-index.search-mode".into(), "full".into()), + ("vector-index.search-mode".into(), "full".into()), + ("full-text-index.search-mode".into(), "full".into()), + ])) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "hybrid"), + format!("memory:/hybrid-parity-{}", uuid::Uuid::new_v4()), + TableSchema::new(0, &schema), + None, + ) +} + +async fn append(table: &Table, start: i32) { + let columns: Vec = vec![ + Arc::new(Int32Array::from(vec![ + start, + start + 1, + start + 2, + start + 3, + ])), + Arc::new(StringArray::from(vec![ + Some("paimon"), + Some("paimon paimon paimon"), + Some("paimon lake"), + None, + ])), + Arc::new(ListArray::from_iter_primitive::( + (start..start + 4).map(|id| Some(vec![Some(id as f32), Some(0.0)])), + )), + Arc::new(StringArray::from(vec![ + Some("drop"), + Some("keep"), + Some("keep"), + None, + ])), + Arc::new(Int32Array::from(vec![start / 4; 4])), + ]; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowType::Int32, true), + Field::new("text", ArrowType::Utf8, true), + Field::new("embedding", columns[2].data_type().clone(), true), + Field::new("label", ArrowType::Utf8, true), + Field::new("pt", ArrowType::Int32, true), + ])); + let mut writer = TableWrite::new(table, "hybrid".into()).unwrap(); + writer + .write_arrow_batch(&RecordBatch::try_new(schema, columns).unwrap()) + .await + .unwrap(); + TableCommit::new(table.clone(), "hybrid".into()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); +} + +async fn index(table: &Table, kind: &str, column: &str) { + let mut builder = table.new_global_index_build_builder(); + builder.with_index_type(kind).with_index_column(column); + if kind == "ivf-flat" { + builder.with_options(HashMap::from([ + ("ivf-flat.dimension".into(), "2".into()), + ("ivf-flat.nlist".into(), "1".into()), + ("ivf-flat.distance.metric".into(), "l2".into()), + ])); + } + builder.execute().await.unwrap(); +} + +#[tokio::test] +async fn full_text_route_uses_java_structured_query() { + let table = table(); + append(&table, 0).await; + index(&table, "full-text", "text").await; + let mut builder = table.new_hybrid_search_builder(); + builder + .add_full_text_route( + "text", + r#"{"match":{"query":"paimon lake","operator":"And"}}"#, + 4, + 1.0, + HashMap::new(), + ) + .unwrap(); + builder.with_limit(4); + assert_eq!(builder.execute_scored().await.unwrap().row_ids, vec![2]); +} + +#[tokio::test] +async fn full_text_route_rejects_plain_text_as_java_does() { + let table = table(); + append(&table, 0).await; + index(&table, "full-text", "text").await; + let mut builder = table.new_hybrid_search_builder(); + builder + .add_full_text_route("text", "paimon", 4, 1.0, HashMap::new()) + .unwrap(); + builder.with_limit(4); + assert!(builder.execute_scored().await.is_err()); +} + +fn mixed_query(table: &Table, route_limit: usize) -> HybridSearchBuilder<'_> { + let mut builder = table.new_hybrid_search_builder(); + builder + .add_vector_route( + "embedding", + vec![0.0, 0.0], + route_limit, + 1.0, + HashMap::new(), + ) + .unwrap() + .add_full_text_route( + "text", + r#"{"match":{"query":"paimon"}}"#, + route_limit, + 1.0, + HashMap::new(), + ) + .unwrap() + .with_limit(16); + builder +} + +#[tokio::test] +async fn both_routes_filter_before_top_k_and_accumulate_predicates() { + let table = table(); + append(&table, 0).await; + append(&table, 4).await; + index(&table, "ivf-flat", "embedding").await; + index(&table, "full-text", "text").await; + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = mixed_query(&table, 1); + builder + .with_filter(pb.equal("label", Datum::String("keep".into())).unwrap()) + .with_filter(pb.equal("id", Datum::Int(6)).unwrap()) + .with_partition_filter(pb.equal("pt", Datum::Int(1)).unwrap()) + .unwrap(); + let result = builder.execute_scored().await.unwrap(); + assert_eq!(result.row_ids, vec![6]); + assert!((result.scores[0] - 2.0 / 61.0).abs() < 1e-6); + builder + .with_partition_filter(pb.equal("pt", Datum::Int(0)).unwrap()) + .unwrap(); + assert!(builder.execute_scored().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn scalar_coverage_and_refinement_are_shared_by_both_routes() { + let table = table(); + append(&table, 0).await; + index(&table, "ivf-flat", "embedding").await; + index(&table, "full-text", "text").await; + let pb = PredicateBuilder::new(table.schema().fields()); + let keep = pb.equal("label", Datum::String("keep".into())).unwrap(); + for mode in ["fast", "full", "detail"] { + let table = table.copy_with_options(HashMap::from([( + "scalar-index.search-mode".into(), + mode.into(), + )])); + let mut ids = mixed_query(&table, 4) + .with_filter(keep.clone()) + .execute_scored() + .await + .unwrap() + .row_ids; + ids.sort_unstable(); + // Vector FAST reads indexed ranges from data when no scalar index can + // evaluate the filter; full-text FAST excludes those ranges instead. + assert_eq!(ids, vec![1, 2], "scalar mode {mode}"); + let mut text = table.new_full_text_search_builder(); + text.with_query("text", r#"{"match":{"query":"paimon"}}"#) + .with_filter(keep.clone()) + .with_limit(4); + assert_eq!( + text.execute_scored().await.unwrap().len(), + if mode == "fast" { 0 } else { 2 } + ); + } + index(&table, "bitmap", "label").await; + let candidate = pb.contains("label", Datum::String("eep".into())).unwrap(); + for refine in [false, true] { + let table = table.copy_with_options(HashMap::from([ + ( + "global-index.filter.refine-from-data".into(), + refine.to_string(), + ), + ("scalar-index.search-mode".into(), "fast".into()), + ])); + let mut ids = mixed_query(&table, 4) + .with_filter(candidate.clone()) + .execute_scored() + .await + .unwrap() + .row_ids; + ids.sort_unstable(); + assert_eq!(ids, if refine { vec![1, 2] } else { vec![] }); + } +} + +#[tokio::test] +async fn concurrent_routes_preserve_weights_limits_and_java_fusion() { + let table = table(); + append(&table, 0).await; + index(&table, "ivf-flat", "embedding").await; + for (ranker, ids, scores) in [ + ( + "rrf", + vec![3, 2, 0, 1], + vec![2.0 / 61.0, 2.0 / 62.0, 1.0 / 61.0, 1.0 / 62.0], + ), + ("mrr", vec![3, 0, 2, 1], vec![2.0, 1.0, 1.0, 0.5]), + ("weighted_score", vec![3, 0, 1, 2], vec![2.0, 1.0, 0.0, 0.0]), + ] { + let mut builder = table.new_hybrid_search_builder(); + builder + .add_vector_route("embedding", vec![0.0, 0.0], 2, 1.0, HashMap::new()) + .unwrap() + .add_vector_route("embedding", vec![3.0, 0.0], 2, 2.0, HashMap::new()) + .unwrap() + .with_limit(4) + .with_ranker(ranker) + .unwrap(); + let result = builder.execute_scored().await.unwrap(); + assert_eq!(result.row_ids, ids, "ranker {ranker}"); + for (actual, expected) in result.scores.iter().zip(scores) { + assert!( + (actual - expected).abs() < 1e-6, + "ranker {ranker}: {actual} != {expected}" + ); + } + builder.with_limit(1); + assert_eq!(builder.execute_scored().await.unwrap().row_ids, vec![3]); + } +} + +#[tokio::test] +async fn retained_tag_keeps_both_routes_on_the_expired_snapshot() { + let table = table(); + append(&table, 0).await; + index(&table, "ivf-flat", "embedding").await; + index(&table, "full-text", "text").await; + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + table + .tag_manager() + .create("hybrid-before", &snapshot) + .await + .unwrap(); + append(&table, 4).await; + table + .file_io() + .delete_file(&table.snapshot_manager().snapshot_path(snapshot.id())) + .await + .unwrap(); + let tagged = table.copy_with_options(HashMap::from([( + "scan.tag-name".into(), + "hybrid-before".into(), + )])); + let mut old = mixed_query(&tagged, 16) + .execute_scored() + .await + .unwrap() + .row_ids; + old.sort_unstable(); + assert_eq!(old, vec![0, 1, 2, 3]); + assert_eq!( + mixed_query(&table, 16) + .execute_scored() + .await + .unwrap() + .len(), + 8 + ); +} + +#[tokio::test] +async fn caller_pinned_empty_and_snapshot_views_do_not_read_latest() { + let table = table(); + append(&table, 0).await; + index(&table, "ivf-flat", "embedding").await; + index(&table, "full-text", "text").await; + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + append(&table, 4).await; + assert!(mixed_query(&table, 16) + .with_snapshot(None) + .execute_scored() + .await + .unwrap() + .is_empty()); + table + .file_io() + .delete_file(&table.snapshot_manager().snapshot_path(snapshot.id())) + .await + .unwrap(); + let mut old = mixed_query(&table, 16) + .with_snapshot(Some(&snapshot)) + .execute_scored() + .await + .unwrap() + .row_ids; + old.sort_unstable(); + assert_eq!(old, vec![0, 1, 2, 3]); +} + +#[tokio::test] +async fn empty_tables_still_validate_route_and_partition_configuration() { + let table = table(); + assert!(mixed_query(&table, 4) + .execute_scored() + .await + .unwrap() + .is_empty()); + let pb = PredicateBuilder::new(table.schema().fields()); + assert!(mixed_query(&table, 4) + .with_partition_filter(pb.equal("id", Datum::Int(1)).unwrap()) + .is_err()); + let mut invalid = table.new_hybrid_search_builder(); + invalid + .add_full_text_route("text", "plain text", 4, 1.0, HashMap::new()) + .unwrap() + .with_limit(4); + // Java does not parse the DSL until a route actually evaluates a corpus. + assert!(invalid.execute_scored().await.unwrap().is_empty()); + append(&table, 0).await; + index(&table, "full-text", "text").await; + assert!(invalid.execute_scored().await.is_err()); + let mut missing = table.new_hybrid_search_builder(); + missing + .add_vector_route("missing", vec![0.0, 0.0], 4, 1.0, HashMap::new()) + .unwrap() + .with_limit(4); + assert!(missing.execute_scored().await.is_err()); +}