Support primary-key full-text results through scan and read

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:

bindings/python/src/context.rs

@@ -801,7 +801,6 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     this.add_class::<crate::full_text_search::PyFullTextScan>()?;
     this.add_class::<crate::full_text_search::PyFullTextScanPlan>()?;
     this.add_class::<crate::full_text_search::PyFullTextRead>()?;
-    this.add_class::<crate::full_text_search::PyFullTextSearchResult>()?;
     this.add_class::<crate::hybrid_search::PyHybridSearchBuilder>()?;
     this.add_class::<crate::hybrid_search::PyHybridSearchResult>()?;
     this.add_class::<crate::vector_search::PyVectorSearchBuilder>()?;
@@ -810,8 +809,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     this.add_class::<crate::vector_search::PyVectorScanPlan>()?;
     this.add_class::<crate::vector_search::PyVectorRead>()?;
     this.add_class::<crate::vector_search::PyBatchVectorRead>()?;
-    this.add_class::<crate::vector_search::PySearchResult>()?;
-    this.add_class::<crate::vector_search::PySearchResultReadBuilder>()?;
+    this.add_class::<crate::search_result::PySearchResult>()?;
+    this.add_class::<crate::search_result::PySearchResultReadBuilder>()?;
     this.add_class::<crate::read::PyReadBuilder>()?;
     this.add_class::<crate::read::PyTableScan>()?;
     this.add_class::<crate::read::PyStreamTableScan>()?;

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<PyFullTextSearchResult> {
+    fn execute_local(&self, py: Python<'_>) -> PyResult<PySearchResult> {
         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<PyFullTextSearchResult> {
+    fn read_plan(&self, py: Python<'_>, plan: &PyFullTextScanPlan) -> PyResult<PySearchResult> {
         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<u64, f32> {
-        self.inner
-            .row_ids
-            .iter()
-            .copied()
-            .zip(self.inner.scores.iter().copied())
-            .collect()
-    }
-}

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;

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<PyBytes>, i32, String, i64, f32);
+
+#[pymethods]
+impl PySearchResult {
+    fn snapshot_id(&self) -> Option<i64> {
+        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<HashMap<u64, f32>> {
+        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<Vec<PySearchPosition>> {
+        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<Vec<Py<PyBytes>>> {
+        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<Vec<String>>,
+}
+
+#[pymethods]
+impl PySearchResultReadBuilder {
+    fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec<String>) -> PyRefMut<'_, Self> {
+        slf.projection = Some(columns);
+        slf
+    }
+    fn read(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
+        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::<Vec<_>>().await
+                })
+            })
+            .map_err(to_py_err)?;
+        batches
+            .iter()
+            .map(|batch| batch.to_pyarrow(py).map(Bound::unbind))
+            .collect()
+    }
+}

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<PyBytes>, i32, String, i64, f32);
-
-#[pymethods]
-impl PySearchResult {
-    fn snapshot_id(&self) -> Option<i64> {
-        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<HashMap<u64, f32>> {
-        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<Vec<PySearchPosition>> {
-        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<Vec<Py<PyBytes>>> {
-        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<Vec<String>>,
-}
-
-#[pymethods]
-impl PySearchResultReadBuilder {
-    fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec<String>) -> PyRefMut<'_, Self> {
-        slf.projection = Some(columns);
-        slf
-    }
-    fn read(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
-        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::<Vec<_>>().await
-                })
-            })
-            .map_err(to_py_err)?;
-        batches
-            .iter()
-            .map(|batch| batch.to_pyarrow(py).map(Bound::unbind))
-            .collect()
-    }
-}