diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 7b605d686..1edff10eb 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -806,7 +806,7 @@ class FullTextSearchBuilder: def with_partition_filter(self, predicate: Dict[str, Any]) -> FullTextSearchBuilder: ... def new_full_text_scan(self) -> FullTextScan: ... def new_full_text_read(self) -> FullTextRead: ... - def execute_local(self) -> FullTextSearchResult: ... + def execute_local(self) -> SearchResult: ... class FullTextScan: def scan(self) -> FullTextScanPlan: ... @@ -815,11 +815,7 @@ class FullTextScanPlan: def snapshot_id(self) -> Optional[int]: ... class FullTextRead: - def read_plan(self, plan: FullTextScanPlan) -> FullTextSearchResult: ... - -class FullTextSearchResult: - def __len__(self) -> int: ... - def row_ids(self) -> Dict[int, float]: ... + def read_plan(self, plan: FullTextScanPlan) -> SearchResult: ... class ShardTableUpdator: def arrow_reader(self) -> RecordBatchReader: diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 372775131..1c29b4564 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -801,7 +801,6 @@ 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::()?; @@ -810,8 +809,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::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/full_text_search.rs b/bindings/python/src/full_text_search.rs index b73c0fb6c..ca364d1a9 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -19,12 +19,12 @@ use crate::error::to_py_err; use crate::predicate::dict_to_table_predicate; +use crate::search_result::PySearchResult; use paimon::spec::Predicate; use paimon::table::{FullTextRead, FullTextScan, FullTextScanPlan, Table}; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; use pyo3::types::PyDict; -use std::collections::HashMap; use std::sync::Arc; #[derive(Default)] @@ -123,10 +123,10 @@ impl PyFullTextSearchBuilder { .map_err(to_py_err)?, }) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + fn execute_local(&self, py: Python<'_>) -> PyResult { let builder = self.config.builder(&self.table).map_err(to_py_err)?; py.detach(|| runtime().block_on(builder.execute_scored())) - .map(|inner| PyFullTextSearchResult { inner }) + .map(PySearchResult::new) .map_err(to_py_err) } } @@ -161,32 +161,9 @@ pub struct PyFullTextRead { } #[pymethods] impl PyFullTextRead { - fn read_plan( - &self, - py: Python<'_>, - plan: &PyFullTextScanPlan, - ) -> PyResult { + fn read_plan(&self, py: Python<'_>, plan: &PyFullTextScanPlan) -> PyResult { py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(|inner| PyFullTextSearchResult { inner }) + .map(PySearchResult::new) .map_err(to_py_err) } } - -#[pyclass(name = "FullTextSearchResult", module = "pypaimon_rust.datafusion")] -pub struct PyFullTextSearchResult { - inner: paimon::full_text::SearchResult, -} -#[pymethods] -impl PyFullTextSearchResult { - 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 a15c7600d..c55aa5bdf 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -31,6 +31,7 @@ mod oss_cpp; mod predicate; mod read; mod schema; +mod search_result; mod shard_update; mod table; mod udf; diff --git a/bindings/python/src/search_result.rs b/bindings/python/src/search_result.rs new file mode 100644 index 000000000..2e889e30f --- /dev/null +++ b/bindings/python/src/search_result.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::error::to_py_err; +use arrow::pyarrow::ToPyArrow; +use futures::TryStreamExt; +use paimon::table::SearchResult; +use paimon_datafusion::runtime::runtime; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use std::collections::HashMap; + +#[pyclass(name = "SearchResult", module = "pypaimon_rust.datafusion")] +pub struct PySearchResult { + inner: SearchResult, +} + +impl PySearchResult { + pub(crate) fn new(inner: SearchResult) -> Self { + Self { inner } + } +} + +type PySearchPosition = (Py, i32, String, i64, f32); + +#[pymethods] +impl PySearchResult { + fn snapshot_id(&self) -> Option { + self.inner.snapshot_id() + } + fn new_read_builder(&self) -> PySearchResultReadBuilder { + PySearchResultReadBuilder { + result: self.inner.clone(), + projection: None, + } + } + fn __len__(&self) -> usize { + self.inner.len() + } + fn row_ids(&self) -> PyResult> { + let hits = self.inner.row_ids().map_err(to_py_err)?; + Ok(hits + .row_ids + .iter() + .copied() + .zip(hits.scores.iter().copied()) + .collect()) + } + fn positions(&self, py: Python<'_>) -> PyResult> { + Ok(self + .inner + .positions() + .map_err(to_py_err)? + .iter() + .map(|position| { + ( + PyBytes::new(py, &position.partition.to_serialized_bytes()).unbind(), + position.bucket, + position.data_file_name.clone(), + position.row_position, + position.score, + ) + }) + .collect()) + } + fn splits(&self, py: Python<'_>) -> PyResult>> { + Ok(self + .inner + .serialize_primary_key_splits() + .map_err(to_py_err)? + .iter() + .map(|bytes| PyBytes::new(py, bytes).unbind()) + .collect()) + } +} + +#[pyclass(name = "SearchResultReadBuilder", module = "pypaimon_rust.datafusion")] +pub struct PySearchResultReadBuilder { + result: SearchResult, + projection: Option>, +} + +#[pymethods] +impl PySearchResultReadBuilder { + fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec) -> PyRefMut<'_, Self> { + slf.projection = Some(columns); + slf + } + fn read(&self, py: Python<'_>) -> PyResult>> { + let batches = py + .detach(|| { + runtime().block_on(async { + let mut builder = self.result.new_read_builder(); + if let Some(projection) = &self.projection { + let columns: Vec<&str> = projection.iter().map(String::as_str).collect(); + builder.with_projection(&columns); + } + builder.read().await?.try_collect::>().await + }) + }) + .map_err(to_py_err)?; + batches + .iter() + .map(|batch| batch.to_pyarrow(py).map(Bound::unbind)) + .collect() + } +} diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index 49fe05b61..540c579d1 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -18,17 +18,15 @@ use std::collections::HashMap; use std::sync::Arc; -use arrow::pyarrow::ToPyArrow; -use futures::TryStreamExt; use paimon::spec::Predicate; use paimon::table::{BatchVectorRead, Table, VectorRead, VectorScan, VectorScanPlan}; -use paimon::vector_search::SearchResult; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict}; +use pyo3::types::PyDict; use crate::error::to_py_err; use crate::predicate::dict_to_table_predicate; +use crate::search_result::PySearchResult; #[derive(Clone, Default)] struct SearchConfig { @@ -332,99 +330,3 @@ impl PyBatchVectorRead { .map_err(to_py_err) } } - -#[pyclass(name = "SearchResult", module = "pypaimon_rust.datafusion")] -pub struct PySearchResult { - inner: SearchResult, -} - -impl PySearchResult { - fn new(inner: SearchResult) -> Self { - Self { inner } - } -} - -type PySearchPosition = (Py, i32, String, i64, f32); - -#[pymethods] -impl PySearchResult { - fn snapshot_id(&self) -> Option { - self.inner.snapshot_id() - } - fn new_read_builder(&self) -> PySearchResultReadBuilder { - PySearchResultReadBuilder { - result: self.inner.clone(), - projection: None, - } - } - fn __len__(&self) -> usize { - self.inner.len() - } - fn row_ids(&self) -> PyResult> { - let hits = self.inner.row_ids().map_err(to_py_err)?; - Ok(hits - .row_ids - .iter() - .copied() - .zip(hits.scores.iter().copied()) - .collect()) - } - fn positions(&self, py: Python<'_>) -> PyResult> { - Ok(self - .inner - .positions() - .map_err(to_py_err)? - .iter() - .map(|position| { - ( - PyBytes::new(py, &position.partition.to_serialized_bytes()).unbind(), - position.bucket, - position.data_file_name.clone(), - position.row_position, - position.score, - ) - }) - .collect()) - } - fn splits(&self, py: Python<'_>) -> PyResult>> { - Ok(self - .inner - .serialize_primary_key_splits() - .map_err(to_py_err)? - .iter() - .map(|bytes| PyBytes::new(py, bytes).unbind()) - .collect()) - } -} - -#[pyclass(name = "SearchResultReadBuilder", module = "pypaimon_rust.datafusion")] -pub struct PySearchResultReadBuilder { - result: SearchResult, - projection: Option>, -} - -#[pymethods] -impl PySearchResultReadBuilder { - fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec) -> PyRefMut<'_, Self> { - slf.projection = Some(columns); - slf - } - fn read(&self, py: Python<'_>) -> PyResult>> { - let batches = py - .detach(|| { - runtime().block_on(async { - let mut builder = self.result.new_read_builder(); - if let Some(projection) = &self.projection { - let columns: Vec<&str> = projection.iter().map(String::as_str).collect(); - builder.with_projection(&columns); - } - builder.read().await?.try_collect::>().await - }) - }) - .map_err(to_py_err)?; - batches - .iter() - .map(|batch| batch.to_pyarrow(py).map(Bound::unbind)) - .collect() - } -} diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 8255768ca..11d85bfbd 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -38,7 +38,7 @@ use crate::table::{ find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, }; use arrow_array::{Array, Int64Array, LargeStringArray, RecordBatch, StringArray}; -use futures::{stream, StreamExt, TryStreamExt}; +use futures::{StreamExt, TryStreamExt}; use paimon_ftindex_core::io::PosWriter; use paimon_ftindex_core::{FullTextIndexConfig, FullTextIndexWriter}; use roaring::RoaringTreemap; @@ -119,11 +119,15 @@ impl<'a> FullTextSearchBuilder<'a> { /// the search against full-text indexes. /// /// Reference: `FullTextSearchBuilder.executeLocal()` + /// For primary-key file positions, use [`Self::execute_scored`]. pub async fn execute(&self) -> crate::Result> { - Ok(self.execute_scored().await?.to_row_ranges()) + self.execute_scored().await?.row_ids()?.to_row_ranges() } - pub async fn execute_scored(&self) -> crate::Result { + /// Search one snapshot without materializing user columns. Results carry + /// global row IDs or primary-key file positions and can be read later with + /// [`crate::table::SearchResult::new_read_builder`]. + pub async fn execute_scored(&self) -> crate::Result { let read = self.new_read()?; read.read(self.new_scan()?.scan().await?).await } @@ -178,15 +182,6 @@ impl<'a> FullTextSearchBuilder<'a> { .ok_or_else(|| crate::Error::ConfigInvalid { message: "Text column must be set via with_text_column()".to_string(), })?; - let query_text = self - .query_text - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query text must be set via with_query_text()".to_string(), - })?; - let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { - message: "Limit must be set via with_limit()".to_string(), - })?; if !resolves_to_pk_full_text_path(&core, text_column) { return Err(crate::Error::Unsupported { @@ -197,74 +192,7 @@ impl<'a> FullTextSearchBuilder<'a> { }); } - // FAST-only: reject FULL/DETAIL loud rather than silently degrading. - if core.full_text_index_search_mode()? != GlobalIndexSearchMode::Fast { - return Err(crate::Error::DataInvalid { - message: "primary-key full-text search supports only the FAST global-index search \ - mode" - .to_string(), - source: None, - }); - } - - // Reject a non-positive limit at construction, before the empty-plan fast - // path — an empty table must not mask an invalid limit (mirrors Java - // `PrimaryKeyFullTextRead`, and matches `search_bucket`'s own guard). - if limit == 0 { - return Err(crate::Error::ConfigInvalid { - message: "Limit must be positive".to_string(), - }); - } - - // Resolve the queried column's schema field id for the scan/field-id guard. - let field_id = find_field_id_by_name(self.table.schema().fields(), text_column) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("full-text search column '{text_column}' does not exist"), - source: None, - })?; - - if let Some(filter) = &self.filter { - let (_, data) = super::bucket_filter::split_partition_and_data_predicates( - filter.clone(), - self.table.schema().fields(), - self.table.schema().partition_keys(), - ); - if !data.is_empty() { - return Err(crate::Error::Unsupported { - message: "Primary-key full-text search does not support non-partition filters" - .into(), - }); - } - } - let plan = PrimaryKeyFullTextScan::new(self.table, field_id, self.filter.clone()) - .plan() - .await?; - if plan.splits.is_empty() { - return Ok(Box::pin(stream::empty())); - } - - // A predicate-free materialization reader projecting the user table - // columns (mirrors `table_read.rs::new_data_file_reader` with an empty - // predicate list). The PK read appends the score column itself. - let materialize_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - self.table.schema().fields().to_vec(), - Vec::new(), - ) - .with_table_options(self.table.schema().options().clone()); - let read = PrimaryKeyFullTextRead::new( - self.table.file_io().clone(), - materialize_reader, - self.table.location().trim_end_matches('/').to_string(), - self.table - .schema() - .core_options() - .index_file_in_data_file_dir(), - ); - read.read(&plan, query_text, limit).await + self.execute_scored().await?.new_read_builder().read().await } } @@ -1414,10 +1342,28 @@ mod tests { ) } + #[tokio::test] + async fn pk_full_text_scan_and_read_accept_an_empty_pk_table() { + let table = pk_full_text_table("pk_ft_plan_read_empty", &[]); + let mut builder = table.new_full_text_search_builder(); + builder + .with_query("body", r#"{"match":{"query":"alpha"}}"#) + .with_limit(10); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + assert_eq!(plan.snapshot_id(), None); + assert!(builder + .new_read() + .unwrap() + .read(plan) + .await + .unwrap() + .is_empty()); + } + // (d) The PK full-text path produces physical positions, not global row-ids: - // `execute` / `execute_scored` must fail loud and point callers at execute_read. + // Only the ranges-only `execute` API rejects the physical-position result. #[tokio::test] - async fn pk_full_text_execute_fails_loud_use_execute_read() { + async fn pk_full_text_ranges_reject_physical_positions() { let table = pk_full_text_table("pk_ft_execute_loud", &[]); let err = table .new_full_text_search_builder() @@ -1428,26 +1374,26 @@ mod tests { .await .unwrap_err(); assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("execute_read")), - "PK full-text execute must fail loud pointing at execute_read, got: {err}" + matches!(&err, crate::Error::Unsupported { message } if message.contains("physical file positions")), + "the ranges-only API must reject primary-key file positions, got: {err}" ); } #[tokio::test] - async fn pk_full_text_execute_scored_fails_loud_use_execute_read() { + async fn pk_full_text_scored_empty_result_keeps_its_address_space() { let table = pk_full_text_table("pk_ft_scored_loud", &[]); - let err = table + let result = table .new_full_text_search_builder() .with_text_column("body") .with_query_text(r#"{"match":{"query":"alpha"}}"#) .with_limit(10) .execute_scored() .await - .unwrap_err(); - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("global row ids")), - "PK full-text execute_scored must fail loud, got: {err}" - ); + .unwrap(); + assert_eq!(result.snapshot_id(), None); + assert!(result.positions().unwrap().is_empty()); + assert!(result.row_ids().is_err()); + assert!(result.serialize_primary_key_splits().unwrap().is_empty()); } // (c) FAST-only: FULL / DETAIL global-index search modes must fail loud on the @@ -1521,7 +1467,7 @@ mod tests { // Dispatch: with data-evolution ENABLED, a configured pk-full-text column must // NOT resolve to the PK path (mirrors Java `primaryKeyFullTextDefinition`), so - // execute_scored takes the append/DE path instead of the PK fail-loud guard. + // execute_scored returns global row IDs rather than physical PK positions. #[tokio::test] async fn data_evolution_enabled_does_not_take_pk_path() { let de_on = HashMap::from([ diff --git a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs index b8948262c..5992e8dcd 100644 --- a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs @@ -135,12 +135,101 @@ fn query(table: &Table) -> FullTextSearchBuilder<'_> { builder } -fn ids(result: SearchResult) -> Vec { - let mut ids = result.row_ids; +fn ids(result: crate::table::SearchResult) -> Vec { + let mut ids = result.into_row_ids().unwrap().row_ids; ids.sort_unstable(); ids } +#[tokio::test] +async fn full_text_result_reads_projection_and_scores_from_the_captured_snapshot() { + let table = new_table(false); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + let result = query(&table).execute_scored().await.unwrap(); + let snapshot = result.snapshot_id().unwrap(); + let expected = result.row_ids().unwrap(); + assert!(result.positions().is_err()); + append(&table, 8, 4).await; + assert!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap() + .unwrap() + > snapshot + ); + + let batches: Vec = result + .new_read_builder() + .with_projection(&["label", "id"]) + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut user_ids = Vec::new(); + let mut scores = Vec::new(); + for batch in batches { + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + ["label", "id", "__paimon_search_score"] + ); + user_ids.extend_from_slice( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + scores.extend_from_slice( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + assert_eq!( + user_ids, + expected + .row_ids + .iter() + .map(|id| *id as i32) + .collect::>() + ); + assert_eq!(scores, expected.scores); + let batches: Vec = result + .new_read_builder() + .with_projection(&[]) + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected.len() + ); + assert!(batches.iter().all(|batch| batch.num_columns() == 1)); + assert!(result + .new_read_builder() + .with_projection(&["missing"]) + .read() + .await + .is_err()); +} + #[tokio::test] async fn scalar_modes_and_refinement_follow_java_before_top_k() { let table = new_table(false); @@ -293,7 +382,7 @@ async fn empty_string_filters_preserve_java_search_refinement() { } } -async fn user_ids(table: &Table, result: SearchResult) -> Vec { +async fn user_ids(table: &Table, result: crate::table::SearchResult) -> Vec { let mut builder = table.new_read_builder(); builder.with_projection(&["id", ROW_ID_FIELD_NAME]).unwrap(); let plan = builder.new_scan().plan().await.unwrap(); @@ -322,6 +411,8 @@ async fn user_ids(table: &Table, result: SearchResult) -> Vec { } } let mut values = result + .row_ids() + .unwrap() .row_ids .iter() .map(|id| mapping[id]) @@ -461,6 +552,8 @@ async fn raw_filter_keeps_bm25_statistics_of_the_whole_corpus() { .await .unwrap(); assert_eq!(ids(filtered.clone()), vec![1, 5, 7, 9, 11]); + let unfiltered = unfiltered.into_row_ids().unwrap(); + let filtered = filtered.into_row_ids().unwrap(); let scores: HashMap<_, _> = unfiltered .row_ids .into_iter() diff --git a/crates/paimon/src/table/full_text_search_builder/read.rs b/crates/paimon/src/table/full_text_search_builder/read.rs index 3a8b21694..934affe9f 100644 --- a/crates/paimon/src/table/full_text_search_builder/read.rs +++ b/crates/paimon/src/table/full_text_search_builder/read.rs @@ -43,7 +43,32 @@ impl FullTextRead { .ok_or_else(|| crate::Error::ConfigInvalid { message: "Query text must be set via with_query() or with_query_text()".into(), })?; - let query = if builder.query_is_dsl { + if context.primary_key_field_id.is_some() { + if builder + .table + .schema() + .core_options() + .full_text_index_search_mode()? + != GlobalIndexSearchMode::Fast + { + return Err(crate::Error::DataInvalid { message: "primary-key full-text search supports only the FAST global-index search mode".into(), source: None }); + } + if let Some(filter) = &context.filter { + let (_, data) = super::super::bucket_filter::split_partition_and_data_predicates( + filter.clone(), + builder.table.schema().fields(), + builder.table.schema().partition_keys(), + ); + if !data.is_empty() { + return Err(crate::Error::Unsupported { + message: + "Primary-key full-text search does not support non-partition filters" + .into(), + }); + } + } + } + let query = if builder.query_is_dsl || context.primary_key_field_id.is_some() { query.to_string() } else { normalize_query_text(query, &context.column)? @@ -56,7 +81,7 @@ impl FullTextRead { }) } - pub async fn read(&self, plan: FullTextScanPlan) -> crate::Result { + pub async fn read(&self, plan: FullTextScanPlan) -> crate::Result { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; CoreOptions::new(plan.table.schema().options()).ensure_read_authorized()?; if self.context != plan.context { @@ -65,12 +90,44 @@ impl FullTextRead { source: None, }); } - if plan.snapshot_id.is_none() { - return Ok(SearchResult::empty()); + let table = std::sync::Arc::new(plan.table.clone()); + if self.context.primary_key_field_id.is_some() { + let empty = super::super::pk_full_text_scan::PrimaryKeyFullTextScanPlan { + snapshot_id: 0, + splits: Vec::new(), + }; + let source = plan.primary_key_plan.as_ref().unwrap_or(&empty); + let materialize = DataFileReader::new( + plan.table.file_io().clone(), + plan.table.schema_manager().clone(), + plan.table.schema().id(), + plan.table.schema().fields().to_vec(), + plan.table.schema().fields().to_vec(), + Vec::new(), + ) + .with_table_options(plan.table.schema().options().clone()); + let read = PrimaryKeyFullTextRead::new( + plan.table.file_io().clone(), + materialize, + plan.table.location().trim_end_matches('/').to_string(), + plan.table + .schema() + .core_options() + .index_file_in_data_file_dir(), + ); + let result = read + .search_route(source, &self.search.query_text, self.search.limit) + .await?; + return crate::table::SearchResult::from_full_text_primary_key( + table, + result.snapshot_id, + result.candidates, + result.splits, + ); } let mut search = self.search.clone(); search.include_row_ids = self.context.include_row_ids.clone(); - evaluate_full_text_search( + let hits = evaluate_full_text_search( FullTextSearchEvaluation { table: Some(&plan.table), file_io: plan.table.file_io(), @@ -84,6 +141,10 @@ impl FullTextRead { &plan.entries, &search, ) - .await + .await?; + Ok(crate::table::SearchResult::from_full_text_row_ids( + table, + crate::vector_search::ScoredRowIds::new(hits.row_ids, hits.scores), + )) } } diff --git a/crates/paimon/src/table/full_text_search_builder/scan.rs b/crates/paimon/src/table/full_text_search_builder/scan.rs index 68e94f221..96e494716 100644 --- a/crates/paimon/src/table/full_text_search_builder/scan.rs +++ b/crates/paimon/src/table/full_text_search_builder/scan.rs @@ -29,6 +29,7 @@ pub(super) struct PlanContext { pub(super) column: String, pub(super) filter: Option, pub(super) include_row_ids: Option, + pub(super) primary_key_field_id: Option, } impl PlanContext { @@ -57,8 +58,15 @@ impl PlanContext { message: format!("Full-text column '{column}' must be a character string"), }); } - if resolves_to_pk_full_text_path(&builder.table.schema().core_options(), column) { - return Err(crate::Error::DataInvalid { message: "primary-key full-text search does not produce global row ids; use the materialized read (execute_read) instead".into(), source: None }); + let primary_key_field_id = + resolves_to_pk_full_text_path(&builder.table.schema().core_options(), column) + .then_some(field.id()); + if primary_key_field_id.is_some() && builder.include_row_ids.is_some() { + return Err(crate::Error::Unsupported { + message: + "primary-key full-text search uses physical file positions, not global row IDs" + .into(), + }); } Ok(Self { location: builder.table.location().trim_end_matches('/').into(), @@ -66,11 +74,12 @@ impl PlanContext { column: column.into(), filter: builder.filter.clone(), include_row_ids: builder.include_row_ids.clone(), + primary_key_field_id, }) } } -/// One selected snapshot and its partition-pruned index entries. A plan remains +/// One selected snapshot and its partition-pruned indexes or PK buckets. A plan remains /// usable after its scan is dropped, and does not re-resolve the latest snapshot. #[derive(Clone)] pub struct FullTextScanPlan { @@ -81,6 +90,8 @@ pub struct FullTextScanPlan { pub(super) snapshot_id: Option, pub(super) partition_filter: Option, pub(super) row_filter: Option, + pub(super) primary_key_plan: + Option, } impl FullTextScanPlan { @@ -114,6 +125,7 @@ impl FullTextScan { snapshot_id: None, partition_filter: None, row_filter: None, + primary_key_plan: None, }; let mut partition_pruner = None; if let Some(filter) = &self.context.filter { @@ -165,6 +177,14 @@ impl FullTextScan { // Pin the data without changing the schema already bound by the query. // A schema-only rename does not create a new data snapshot. plan.table = self.table.copy_with_pinned_snapshot(Some(&snapshot)); + if let Some(field_id) = self.context.primary_key_field_id { + plan.primary_key_plan = Some( + PrimaryKeyFullTextScan::new(&plan.table, field_id, self.context.filter.clone()) + .plan() + .await?, + ); + return Ok(plan); + } if let Some(name) = snapshot.index_manifest() { let path = self.table.snapshot_manager().manifest_path(name); plan.entries = IndexManifest::read(self.table.file_io(), &path).await?; diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index d6edb496a..602a3337b 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -1043,7 +1043,7 @@ async fn execute_full_text_route( ) -> 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)) + result.into_row_ids() } #[cfg(not(feature = "fulltext"))] diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 92d5b7c9f..f7caeb626 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -138,6 +138,7 @@ mod row_position_selection; mod row_sidecar; mod scan_trace; pub(crate) mod schema_manager; +pub(crate) mod search_result; mod shard_table_updater; pub(crate) mod snapshot_commit; mod snapshot_manager; @@ -170,7 +171,6 @@ mod vector_read; mod vector_scan; mod vector_search_builder; mod vector_search_common; -pub(crate) mod vector_search_result; #[cfg(test)] mod vector_search_test_utils; mod vindex_index_build_builder; @@ -222,6 +222,7 @@ pub use read_builder::ReadBuilder; pub use rest_env::RESTEnv; pub use scan_trace::ScanTrace; pub use schema_manager::SchemaManager; +pub use search_result::{PrimaryKeySearchPosition, SearchResult, SearchResultReadBuilder}; pub use shard_table_updater::ShardTableUpdater; pub use snapshot_commit::{RESTSnapshotCommit, RenamingSnapshotCommit, SnapshotCommit}; pub use snapshot_manager::SnapshotManager; diff --git a/crates/paimon/src/table/pk_full_text_read.rs b/crates/paimon/src/table/pk_full_text_read.rs index 18376cf57..a5523a302 100644 --- a/crates/paimon/src/table/pk_full_text_read.rs +++ b/crates/paimon/src/table/pk_full_text_read.rs @@ -16,28 +16,21 @@ // under the License. //! Primary-key full-text search read foundation: the score-based candidate type, -//! its cross-bucket score-descending Top-K, and the FAST-only materialized read -//! orchestration (plan the buckets, search each, fuse by score, materialize the -//! winning physical rows best-score-first with a `__paimon_search_score` column). +//! its cross-bucket score-descending Top-K, and the FAST-only search over +//! physical positions. SearchResult materializes the selected rows later. use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; -use arrow_array::RecordBatch; -use futures::{stream, TryStreamExt}; - use crate::deletion_vector::DeletionVector; use crate::io::FileIO; use crate::spec::BinaryRow; use crate::table::data_file_reader::DataFileReader; use crate::table::pk_full_text_bucket_search::search_bucket; use crate::table::pk_full_text_scan::{PrimaryKeyFullTextScanPlan, PrimaryKeyFullTextSearchSplit}; -use crate::table::pk_vector_indexed_split_read::{PkVectorIndexedSplit, PkVectorIndexedSplitRead}; +use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit; use crate::table::source::DataSplitBuilder; -use crate::table::vector_search_common::{ - collect_ranked_rows, reorder_and_strip_position, RankedRow, -}; -use crate::table::{ArrowRecordBatchStream, RowRange}; +use crate::table::RowRange; fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { @@ -255,24 +248,19 @@ pub(crate) fn build_full_text_indexed_splits( /// [`PrimaryKeyFullTextRead::search_route`]. pub(crate) struct PrimaryKeyFullTextRouteResult<'a> { pub(crate) candidates: Vec, - // Read by the hybrid route consumer (fuses routes before materializing); the - // materialized read reaches candidates directly and holds the plan itself. + // Both shared results and hybrid routes retain the selected snapshot. pub(crate) snapshot_id: i64, pub(crate) splits: &'a [PrimaryKeyFullTextSearchSplit], } -/// FAST-only primary-key full-text materialized read. Given a planned set of +/// FAST-only primary-key full-text search. Given a planned set of /// per-bucket search splits, it searches each bucket through the full-text archive -/// reader, fuses the hits cross-bucket by score, materializes the winning physical -/// rows, and re-orders them best-score-first with a `__paimon_search_score` column -/// (the internal `_PKEY_VECTOR_POSITION` column is stripped). Mirrors Java -/// `PrimaryKeyFullTextRead` feeding its result splits into an ordinary table read: -/// the search decides which rows, the reader decides which columns. +/// reader and fuses physical positions by score. SearchResult owns the selections +/// and their later materialization, as Java PrimaryKeyScoredResult does. /// /// The planning + FAST-mode guard live in the caller (`FullTextSearchBuilder`); -/// this type is the Table-free search+materialize core (it holds only the -/// predicate-free materialization reader, the `FileIO`, and the table path for the -/// archive directory), so it can be driven end-to-end from a hand-built plan. +/// this type searches a prepared plan. Its data reader loads deletion vectors; +/// projected user columns are read later by SearchResult. pub(crate) struct PrimaryKeyFullTextRead { file_io: FileIO, materialize_reader: DataFileReader, @@ -299,8 +287,8 @@ impl PrimaryKeyFullTextRead { /// best-score-first Top-`limit`, WITHOUT materializing any rows. Returns the /// surviving candidates together with the route source context (the plan's /// resolved snapshot id and its per-bucket source splits) so a caller can fuse - /// them with another route before materializing. [`read`](Self::read) is - /// layered on top of this. `query` is passed verbatim to the archive reader + /// them with another route before materializing through SearchResult. + /// `query` is passed verbatim to the archive reader /// (spec D2). pub(crate) async fn search_route<'p>( &self, @@ -362,71 +350,6 @@ impl PrimaryKeyFullTextRead { splits: &plan.splits, }) } - - /// Search every planned bucket, fuse the hits by score into a global Top-`limit`, - /// materialize the winning rows, and emit them best-score-first with the - /// unified score column. An empty plan or an empty result yields an empty - /// stream. `query` is passed verbatim to the archive reader (spec D2). - pub(crate) async fn read( - &self, - plan: &PrimaryKeyFullTextScanPlan, - query: &str, - limit: usize, - ) -> crate::Result { - let survivors = self.search_route(plan, query, limit).await?.candidates; - if survivors.is_empty() { - return Ok(Box::pin(stream::empty())); - } - - // Rank each survivor by its best-first position, keyed by its FULL physical - // position `(partition bytes, bucket, file, row position)` so the physical - // materialization order can be reduced back to best-first (gap 8). Keying on - // just file+position would collide across partitions/buckets. - let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); - for (rank, c) in survivors.iter().enumerate() { - rank_of.insert( - ( - c.partition.to_serialized_bytes(), - c.bucket, - c.data_file_name.clone(), - c.row_position, - ), - rank, - ); - } - - let indexed_splits = build_full_text_indexed_splits(survivors, &plan.splits)?; - - // Materialize every indexed split, retaining each batch and, per row, the - // (rank, batch_index, row_index) tuple so we can reorder to best-first. - // Top-K is small, so full in-memory collection is acceptable. - let mut batches: Vec = Vec::new(); - let mut ranked: Vec = Vec::new(); - for indexed in indexed_splits { - let partition_bytes = indexed.split.partition().to_serialized_bytes(); - let bucket = indexed.split.bucket(); - let file_name = indexed.split.data_files()[0].file_name.clone(); - let mut stream = - PkVectorIndexedSplitRead::new(self.materialize_reader.clone()).read(&indexed)?; - while let Some(batch) = stream.try_next().await? { - let batch_index = batches.len(); - collect_ranked_rows( - &batch, - batch_index, - &partition_bytes, - bucket, - &file_name, - &rank_of, - &mut ranked, - )?; - batches.push(batch); - } - } - - // Reorder to best-first and drop the internal position column. - let output = reorder_and_strip_position(&batches, ranked)?; - Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) - } } #[cfg(test)] @@ -566,8 +489,10 @@ mod read_tests { use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; use crate::table::schema_manager::SchemaManager; use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile}; + use crate::table::ArrowRecordBatchStream; use arrow_array::{Array, Float32Array, Int32Array, RecordBatch}; use bytes::Bytes; + use futures::TryStreamExt; use paimon_ftindex_core::io::PosWriter; use paimon_ftindex_core::{FullTextIndexConfig, FullTextIndexWriter}; use paimon_mosaic_core::spec::COMPRESSION_NONE; @@ -575,6 +500,37 @@ mod read_tests { use std::io; use std::sync::Arc; + // Exercise the same shared result materialization as the public builder. + impl PrimaryKeyFullTextRead { + async fn read( + &self, + plan: &PrimaryKeyFullTextScanPlan, + query: &str, + limit: usize, + ) -> crate::Result { + let schema = crate::spec::Schema::builder() + .column("id", DataType::Int(IntType::new())) + .build()?; + let table = crate::table::Table::new( + self.file_io.clone(), + crate::catalog::Identifier::new("default", "text"), + self.table_path.clone(), + crate::spec::TableSchema::new(1, &schema), + None, + ); + let route = self.search_route(plan, query, limit).await?; + crate::table::SearchResult::from_full_text_primary_key( + Arc::new(table), + route.snapshot_id, + route.candidates, + route.splits, + )? + .new_read_builder() + .read() + .await + } + } + const FILE_SOURCE_COMPACT: i32 = 1; // ---- helpers ---- diff --git a/crates/paimon/src/table/pk_full_text_scan.rs b/crates/paimon/src/table/pk_full_text_scan.rs index 350b17dc2..679d20c25 100644 --- a/crates/paimon/src/table/pk_full_text_scan.rs +++ b/crates/paimon/src/table/pk_full_text_scan.rs @@ -43,6 +43,7 @@ fn data_invalid(message: impl Into) -> crate::Error { /// bucket-split invariants (unique data files; each payload covers ≥1 active /// source; no active file double-covered; uncovered files active/non-duplicate/ /// not-covered; and covered ∪ uncovered == every active data file in the split). +#[derive(Clone)] pub(crate) struct PrimaryKeyFullTextSearchSplit { pub data_split: DataSplit, pub current_payloads: Vec, @@ -248,6 +249,7 @@ impl BucketAccumulator { } /// The per-bucket search splits produced by planning. +#[derive(Clone)] pub(crate) struct PrimaryKeyFullTextScanPlan { // The snapshot the plan resolved during planning (pinned before the index // manifest is read); the read guards every split against it before searching. diff --git a/crates/paimon/src/table/vector_search_result.rs b/crates/paimon/src/table/search_result.rs similarity index 79% rename from crates/paimon/src/table/vector_search_result.rs rename to crates/paimon/src/table/search_result.rs index 01610637d..4c5861e0d 100644 --- a/crates/paimon/src/table/vector_search_result.rs +++ b/crates/paimon/src/table/search_result.rs @@ -45,7 +45,9 @@ pub struct PrimaryKeySearchPosition { #[derive(Debug, Clone)] enum SearchHits { DataEvolution { - vector_column: String, + // Vector results retain their target-column validation at read time; + // full-text columns have already been validated by FullTextRead. + vector_column: Option, hits: ScoredRowIds, }, PrimaryKey { @@ -54,9 +56,9 @@ enum SearchHits { }, } -/// Ranked vector-search hits together with the source snapshot needed to read them. +/// Ranked search hits together with the source snapshot needed to read them. /// -/// Both single and batch searches return this type. DE hits use global row IDs; +/// Single/batch vector and full-text searches return this type. DE hits use global row IDs; /// PK hits use physical file positions, as in Java's `PrimaryKeyScoredResult`. /// Searching does not materialize projected user columns. Use /// [`new_read_builder`](Self::new_read_builder) when rows are needed. @@ -78,12 +80,51 @@ impl SearchResult { table, snapshot_id, hits: SearchHits::DataEvolution { - vector_column, + vector_column: Some(vector_column), + hits, + }, + } + } + + #[cfg(feature = "fulltext")] + pub(super) fn from_full_text_row_ids(table: Arc, hits: ScoredRowIds) -> Self { + let snapshot_id = table.travel_snapshot().map(|snapshot| snapshot.id()); + Self { + table, + snapshot_id, + hits: SearchHits::DataEvolution { + vector_column: None, hits, }, } } + #[cfg(feature = "fulltext")] + pub(super) fn from_full_text_primary_key( + table: Arc
, + snapshot_id: i64, + candidates: Vec, + source_splits: &[super::pk_full_text_scan::PrimaryKeyFullTextSearchSplit], + ) -> crate::Result { + let positions = candidates + .iter() + .map(|c| PrimaryKeySearchPosition { + partition: c.partition.clone(), + bucket: c.bucket, + data_file_name: c.data_file_name.clone(), + row_position: c.row_position, + score: c.score, + }) + .collect(); + let splits = + super::pk_full_text_read::build_full_text_indexed_splits(candidates, source_splits)?; + Ok(Self { + table, + snapshot_id: (snapshot_id != 0).then_some(snapshot_id), + hits: SearchHits::PrimaryKey { positions, splits }, + }) + } + pub(super) fn from_primary_key( table: Arc
, snapshot_id: i64, @@ -200,7 +241,7 @@ fn no_global_row_ids() -> crate::Error { /// Reads selected user columns and `__paimon_search_score`, in relevance order. /// Internal global row IDs and PK positions are not output. Reading a result -/// never reruns vector search or replans its primary-key source files. +/// never reruns search or replans its primary-key source files. #[derive(Debug, Clone)] pub struct SearchResultReadBuilder<'a> { result: &'a SearchResult, @@ -223,8 +264,12 @@ impl SearchResultReadBuilder<'_> { vector_column, hits, } => { - let read_type = - DeVectorRead::read_type(table, vector_column, self.projection.as_deref())?; + let read_type = match vector_column { + Some(column) => { + DeVectorRead::read_type(table, column, self.projection.as_deref())? + } + None => resolve_materialize_read_type(table, self.projection.as_deref())?, + }; materialize_row_ids(table, hits, read_type).await } SearchHits::PrimaryKey { positions, splits } => { diff --git a/crates/paimon/src/table/vector_search_result/tests.rs b/crates/paimon/src/table/search_result/tests.rs similarity index 100% rename from crates/paimon/src/table/vector_search_result/tests.rs rename to crates/paimon/src/table/search_result/tests.rs diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index 0599b3f49..1ebc3f716 100644 --- a/crates/paimon/src/vector_search.rs +++ b/crates/paimon/src/vector_search.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -pub use crate::table::vector_search_result::{ +pub use crate::table::search_result::{ PrimaryKeySearchPosition, SearchResult, SearchResultReadBuilder, };