Merge branch 'main' into gatekeeper/fix-2325-1

python/src/connection.rs

@@ -609,7 +609,7 @@ impl Connection {
                 .create_function_async(request)
                 .await
                 .infer_error()
-                .map(crate::job::FunctionJob::new)
+                .map(crate::job::Job::new_typed)
         })
     }
 

python/src/index.rs

@@ -8,7 +8,7 @@ use lancedb::index::vector::{
 };
 use lancedb::index::{
     Index as LanceDbIndex,
-    scalar::{BTreeIndexBuilder, FmIndexBuilder, FtsIndexBuilder},
+    scalar::{BTreeIndexBuilder, DocumentGranularity, FmIndexBuilder, FtsIndexBuilder},
 };
 use pyo3::IntoPyObject;
 use pyo3::types::PyStringMethods;
@@ -60,7 +60,11 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
                     .ngram_min_length(params.ngram_min_length)
                     .ngram_max_length(params.ngram_max_length)
                     .ngram_prefix_only(params.prefix_only)
-                    .custom_stop_words(params.custom_stop_words);
+                    .custom_stop_words(params.custom_stop_words)
+                    .document_granularity(
+                        DocumentGranularity::try_from(params.document_granularity.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))?,
+                    );
                 if let Some(memory_limit) = params.memory_limit {
                     inner_opts = inner_opts.memory_limit_mb(memory_limit);
                 }
@@ -221,6 +225,7 @@ struct FtsParams {
     block_size: usize,
     memory_limit: Option<u64>,
     num_workers: Option<usize>,
+    document_granularity: String,
 }
 
 #[derive(FromPyObject)]
@@ -481,6 +486,7 @@ mod tests {
     block_size = 128
     memory_limit = 2048
     num_workers = 7
+    document_granularity = 'row'
 
 config = FTS()",
                 None,

python/src/job.rs

@@ -5,72 +5,33 @@ use std::sync::Arc;
 
 use crate::runtime::future_into_py;
 use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
+use serde::Serialize;
 
 use crate::error::PythonErrorExt;
 
 #[pyclass]
 pub struct Job {
-    inner: Arc<lancedb::Job>,
-}
-
-/// Python bridge for a typed remote Function registration job.
-///
-/// The public Python layer decodes the canonical JSON returned by `wait`
-/// into its immutable `FunctionVersion` model.
-#[pyclass]
-pub struct FunctionJob {
-    inner: Arc<lancedb::Job<lancedb::function::FunctionVersion>>,
-}
-
-impl FunctionJob {
-    pub(crate) fn new(inner: lancedb::Job<lancedb::function::FunctionVersion>) -> Self {
-        Self {
-            inner: Arc::new(inner),
-        }
-    }
+    inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
 }
 
 impl Job {
     pub(crate) fn new(inner: lancedb::Job) -> Self {
         Self {
-            inner: Arc::new(inner),
+            inner: Arc::new(inner.map(|()| Ok(None))),
         }
     }
-}
 
-#[pymethods]
-impl FunctionJob {
-    #[getter]
-    pub fn id(&self) -> Option<String> {
-        self.inner.id().map(str::to_string)
-    }
-
-    pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.inner.clone();
-        future_into_py(
-            self_.py(),
-            async move { inner.status().await.infer_error() },
-        )
-    }
-
-    pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.inner.clone();
-        future_into_py(self_.py(), async move {
-            inner
-                .wait()
-                .await
-                .infer_error()?
-                .to_canonical_json()
-                .infer_error()
-        })
-    }
-
-    pub fn cancel(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.inner.clone();
-        future_into_py(self_.py(), async move {
-            inner.cancel().await.infer_error()?;
-            Ok(())
-        })
+    pub(crate) fn new_typed<T>(inner: lancedb::Job<T>) -> Self
+    where
+        T: Clone + Serialize + Send + Sync + 'static,
+    {
+        Self {
+            inner: Arc::new(inner.map(|result| {
+                serde_json::to_string(&result)
+                    .map(Some)
+                    .map_err(|error| format!("failed to serialize typed job result: {error}"))
+            })),
+        }
     }
 }
 
@@ -92,8 +53,10 @@ impl Job {
     pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
         let inner = self_.inner.clone();
         future_into_py(self_.py(), async move {
-            inner.wait().await.infer_error()?;
-            Ok(None::<()>)
+            let result = inner.wait().await.infer_error()?;
+            result
+                .map_err(|message| lancedb::Error::Runtime { message })
+                .infer_error()
         })
     }
 

python/src/lib.rs

@@ -47,7 +47,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<Session>()?;
     m.add_class::<Table>()?;
     m.add_class::<crate::job::Job>()?;
-    m.add_class::<crate::job::FunctionJob>()?;
     m.add_class::<crate::job::JobInfo>()?;
     m.add_class::<crate::job::JobDescription>()?;
     m.add_class::<crate::job::JobFailureInfo>()?;

python/src/query.rs

@@ -16,8 +16,8 @@ use arrow::pyarrow::FromPyArrow;
 use arrow::pyarrow::IntoPyArrow;
 use arrow::pyarrow::ToPyArrow;
 use lancedb::index::scalar::{
-    BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
-    Operator, PhraseQuery,
+    BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery,
+    MultiMatchQuery, Occur, Operator, PhraseQuery,
 };
 use lancedb::query::AnalyzePlanDistributedMetrics;
 use lancedb::query::QueryBase;
@@ -76,8 +76,16 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
                 let max_expansions = ob.getattr("max_expansions")?.extract()?;
                 let operator = ob.getattr("operator")?.extract::<String>()?;
                 let prefix_length = ob.getattr("prefix_length")?.extract()?;
+                let document_granularity = ob
+                    .getattr("document_granularity")?
+                    .extract::<Option<String>>()?
+                    .map(|value| {
+                        DocumentGranularity::try_from(value.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))
+                    })
+                    .transpose()?;
 
-                Ok(Self(
+                let mut query =
                     MatchQuery::new(query)
                         .with_column(Some(column))
                         .with_boost(boost)
@@ -86,21 +94,32 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
                         .with_operator(Operator::try_from(operator.as_str()).map_err(|e| {
                             PyValueError::new_err(format!("Invalid operator: {}", e))
                         })?)
-                        .with_prefix_length(prefix_length)
-                        .into(),
-                ))
+                        .with_prefix_length(prefix_length);
+                if let Some(document_granularity) = document_granularity {
+                    query = query.with_document_granularity(document_granularity);
+                }
+                Ok(Self(query.into()))
             }
             "PhraseQuery" => {
                 let query = ob.getattr("query")?.extract()?;
                 let column = ob.getattr("column")?.extract()?;
                 let slop = ob.getattr("slop")?.extract()?;
+                let document_granularity = ob
+                    .getattr("document_granularity")?
+                    .extract::<Option<String>>()?
+                    .map(|value| {
+                        DocumentGranularity::try_from(value.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))
+                    })
+                    .transpose()?;
 
-                Ok(Self(
-                    PhraseQuery::new(query)
-                        .with_column(Some(column))
-                        .with_slop(slop)
-                        .into(),
-                ))
+                let mut query = PhraseQuery::new(query)
+                    .with_column(Some(column))
+                    .with_slop(slop);
+                if let Some(document_granularity) = document_granularity {
+                    query = query.with_document_granularity(document_granularity);
+                }
+                Ok(Self(query.into()))
             }
             "BoostQuery" => {
                 let positive: Self = ob.getattr("positive")?.extract()?;
@@ -167,6 +186,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
                 kwargs.set_item("max_expansions", query.max_expansions)?;
                 kwargs.set_item::<_, &str>("operator", query.operator.into())?;
                 kwargs.set_item("prefix_length", query.prefix_length)?;
+                if let Some(document_granularity) = query.document_granularity {
+                    let value = match document_granularity {
+                        DocumentGranularity::Row => "row",
+                        DocumentGranularity::ListElement => "list_element",
+                    };
+                    kwargs.set_item("document_granularity", value)?;
+                }
                 namespace
                     .getattr(intern!(py, "MatchQuery"))?
                     .call((query.terms, query.column.unwrap()), Some(&kwargs))
@@ -174,6 +200,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
             FtsQuery::Phrase(query) => {
                 let kwargs = PyDict::new(py);
                 kwargs.set_item("slop", query.slop)?;
+                if let Some(document_granularity) = query.document_granularity {
+                    let value = match document_granularity {
+                        DocumentGranularity::Row => "row",
+                        DocumentGranularity::ListElement => "list_element",
+                    };
+                    kwargs.set_item("document_granularity", value)?;
+                }
                 namespace
                     .getattr(intern!(py, "PhraseQuery"))?
                     .call((query.terms, query.column.unwrap()), Some(&kwargs))

python/src/table.rs

@@ -903,15 +903,19 @@ impl Table {
         })
     }
 
-    #[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
+    #[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
     pub fn add<'a>(
         self_: PyRef<'a, Self>,
         data: PyScannable,
         mode: String,
         progress: Option<Py<PyAny>>,
         write_parallelism: Option<usize>,
+        allow_external_blob_outside_bases: bool,
     ) -> PyResult<Bound<'a, PyAny>> {
-        let mut op = self_.inner_ref()?.add(data);
+        let mut op = self_
+            .inner_ref()?
+            .add(data)
+            .allow_external_blob_outside_bases(allow_external_blob_outside_bases);
         if mode == "append" {
             op = op.mode(AddDataMode::Append);
         } else if mode == "overwrite" {
@@ -1744,7 +1748,7 @@ impl Table {
         let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {
             let job = inner.refresh_column_async(column).await.infer_error()?;
-            Ok(crate::job::Job::new(job))
+            Ok(crate::job::Job::new_typed(job))
         })
     }