Merge main into gatekeeper/fix-2325-1

This commit is contained in:
Gatefixer
2026-10-16 18:58:03 +00:00
593 changed files with 85326 additions and 17533 deletions
+280
View File
@@ -0,0 +1,280 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
use lancedb::authz::*;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::types::{PyAnyMethods, PyBytes, PyString};
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use std::sync::Arc;
#[pyclass]
pub struct AuthorizationClient {
pub(crate) inner: Arc<dyn Authorization>,
}
#[pymethods]
impl AuthorizationClient {
fn list_acls<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
let request: ListAccessControlEntriesRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let result = inner.list_acls(request).await.infer_error()?;
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
})
}
fn add_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
let request: AccessControlEntryRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.add_acl(request).await.infer_error()
})
}
fn delete_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
let request: AccessControlEntryRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.delete_acl(request).await.infer_error()
})
}
fn list_role_bindings<'py>(
self_: PyRef<'py, Self>,
request: &str,
) -> PyResult<Bound<'py, PyAny>> {
let request: ListRoleBindingsRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
request.validate().infer_error()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let result = inner.list_role_bindings(request).await.infer_error()?;
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
})
}
fn add_role_binding<'py>(
self_: PyRef<'py, Self>,
request: &str,
) -> PyResult<Bound<'py, PyAny>> {
let request: RoleBindingRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
request.validate().infer_error()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.add_role_binding(request).await.infer_error()
})
}
fn delete_role_binding<'py>(
self_: PyRef<'py, Self>,
request: &str,
) -> PyResult<Bound<'py, PyAny>> {
let request: RoleBindingRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
request.validate().infer_error()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.delete_role_binding(request).await.infer_error()
})
}
fn get_principal<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
let request: GetPrincipalRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
request.validate().infer_error()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let result = inner.get_principal(request).await.infer_error()?;
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
})
}
fn get_group<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
let request: GetGroupRequest = serde_json::from_str(request)
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
request.validate().infer_error()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let result = inner.get_group(request).await.infer_error()?;
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
})
}
}
/// The shared privilege parser and canonical wire names.
#[pyclass(frozen)]
pub struct AuthzPrivilege {
inner: Privilege,
}
#[pymethods]
impl AuthzPrivilege {
#[new]
fn new(value: &str) -> PyResult<Self> {
Ok(Self {
inner: value.parse().map_err(PyValueError::new_err)?,
})
}
#[staticmethod]
fn known() -> Vec<String> {
Privilege::KNOWN.iter().map(ToString::to_string).collect()
}
fn __str__(&self) -> String {
self.inner.to_string()
}
}
#[pyclass(frozen)]
pub struct AuthzSubject {
inner: Subject,
}
#[pymethods]
impl AuthzSubject {
#[new]
fn new(kind: &str, value: &Bound<'_, PyAny>) -> PyResult<Self> {
let value: String = value
.extract()
.map_err(|_| PyValueError::new_err("Subject value must be a nonempty string"))?;
Ok(Self {
inner: Subject::new(kind.parse().infer_error()?, value).infer_error()?,
})
}
#[staticmethod]
fn from_canonical(value: &str) -> PyResult<Self> {
Ok(Self {
inner: Subject::from_canonical(value).infer_error()?,
})
}
#[getter]
fn kind(&self) -> &str {
self.inner.kind().as_str()
}
#[getter]
fn value(&self) -> &str {
self.inner.value()
}
#[getter]
fn display_value(&self) -> &str {
self.inner.display_value()
}
#[pyo3(signature = (context=None))]
fn to_wire(&self, context: Option<&str>) -> PyResult<String> {
match context {
None => {}
Some("principal") => self.inner.validate_principal().infer_error()?,
Some("group") => self.inner.validate_group().infer_error()?,
Some("role_member") => self.inner.validate_role_member().infer_error()?,
Some(_) => return Err(PyValueError::new_err("Invalid subject context")),
}
Ok(self.inner.to_wire())
}
fn __repr__(&self) -> String {
format!("{:?}", self.inner)
}
fn __str__(&self) -> String {
self.inner.to_string()
}
}
#[pyclass(frozen)]
pub struct AuthzObject {
inner: Object,
}
impl AuthzObject {
fn validated(inner: Object) -> PyResult<Self> {
inner.validate().map_err(PyValueError::new_err)?;
Ok(Self { inner })
}
}
fn namespace_path(value: &Bound<'_, PyAny>) -> PyResult<NamespacePath> {
if value.is_instance_of::<PyString>() || value.is_instance_of::<PyBytes>() {
return Err(PyValueError::new_err(
"namespace must be a sequence of components",
));
}
let components = value
.extract()
.map_err(|_| PyValueError::new_err("namespace must be a sequence of string components"))?;
NamespacePath::new(components).map_err(PyValueError::new_err)
}
#[pymethods]
impl AuthzObject {
#[new]
fn new(value: &str) -> PyResult<Self> {
Ok(Self {
inner: value.parse().infer_error()?,
})
}
#[staticmethod]
fn system() -> Self {
Self {
inner: Object::System,
}
}
#[staticmethod]
fn database(name: &str) -> PyResult<Self> {
Self::validated(Object::Database(DatabaseObject::new(name.into())))
}
#[staticmethod]
fn namespace(database: &str, namespace: &Bound<'_, PyAny>) -> PyResult<Self> {
Self::validated(Object::Namespace(NamespaceObject {
database: database.into(),
namespace: namespace_path(namespace)?,
}))
}
#[staticmethod]
fn table(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
Self::validated(Object::Table(TableObject {
database: database.into(),
namespace: namespace_path(namespace)?,
table: name.into(),
}))
}
#[staticmethod]
fn view(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
Self::validated(Object::View(ViewObject {
database: database.into(),
namespace: namespace_path(namespace)?,
view: name.into(),
}))
}
#[staticmethod]
fn secret(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
Self::validated(Object::Secret(SecretObject {
database: database.into(),
namespace: namespace_path(namespace)?,
secret: name.into(),
}))
}
#[staticmethod]
fn function(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
Self::validated(Object::Function(FunctionObject {
database: database.into(),
namespace: namespace_path(namespace)?,
function: name.into(),
}))
}
fn __str__(&self) -> String {
self.inner.to_string()
}
fn __repr__(&self) -> String {
format!("{:?}", self.inner)
}
}
+52 -15
View File
@@ -1,12 +1,13 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::{Arc, Mutex};
use std::time::Duration;
use lancedb::catalog::{
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
use pyo3::exceptions::PyValueError;
use futures::{StreamExt, future::poll_fn};
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
use crate::connection::{Connection, PyClientConfig};
@@ -20,6 +21,13 @@ pub struct Catalog {
#[pymethods]
impl Catalog {
#[getter]
fn authz(&self) -> PyResult<crate::authz::AuthorizationClient> {
Ok(crate::authz::AuthorizationClient {
inner: self.inner.authz().infer_error()?,
})
}
#[getter]
fn uri(&self) -> &str {
self.inner.uri()
@@ -67,19 +75,48 @@ impl Catalog {
})
}
#[pyo3(signature = (*, limit=None, page_token=None))]
fn list_databases<'py>(
self_: PyRef<'py, Self>,
limit: Option<u32>,
page_token: Option<String>,
) -> PyResult<Bound<'py, PyAny>> {
#[pyo3(signature = (*, page_token=None, page_limit=None))]
fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
DatabaseNames {
inner: Arc::new(Mutex::new(
self.inner.list_databases(page_token, page_limit),
)),
next_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
}
#[pyclass]
pub struct DatabaseNames {
inner: Arc<Mutex<lancedb::catalog::DatabaseNames>>,
next_lock: Arc<tokio::sync::Mutex<()>>,
}
#[pymethods]
impl DatabaseNames {
fn num_page_results(&self, py: Python<'_>) -> usize {
py.detach(|| self.inner.lock().unwrap().num_page_results())
}
fn page_token(&self, py: Python<'_>) -> Option<String> {
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
}
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
self_
}
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let mut request = ListDatabasesRequest::default();
request.limit = limit;
request.page_token = page_token;
let response = inner.list_databases(request).await.infer_error()?;
Ok((response.databases, response.page_token))
// Serialize advances while allowing synchronous state inspection
// between polls, including while a REST request is pending.
let _guard = next_lock.lock().await;
poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
.await
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
.infer_error()
})
}
}
+86 -67
View File
@@ -14,9 +14,10 @@ use crate::{
table::Table,
};
use arrow::{
array::RecordBatch,
datatypes::Schema,
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, ToPyArrow},
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
};
use lancedb::{
connection::Connection as LanceConnection,
@@ -148,18 +149,22 @@ impl Connection {
self.get_inner().map(|inner| inner.uri().to_string())
}
#[pyo3(signature = (query, *, default_namespace_path=None))]
#[pyo3(signature = (query, *, default_namespace_path=None, parameters=None))]
pub fn execute_query_async<'a>(
self_: PyRef<'a, Self>,
query: String,
default_namespace_path: Option<Bound<'_, PyAny>>,
parameters: Option<PyArrowType<RecordBatch>>,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.get_inner()?.clone();
let default_namespace_path = parse_default_namespace_path(default_namespace_path)?;
future_into_py(self_.py(), async move {
let operation = inner
let mut operation = inner
.execute_query_async(query)
.default_namespace_path(default_namespace_path);
if let Some(PyArrowType(parameters)) = parameters {
operation = operation.parameters(parameters);
}
operation
.execute()
.await
@@ -407,8 +412,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -417,17 +421,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -443,8 +440,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
@@ -453,17 +449,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -480,12 +469,16 @@ impl Connection {
})
}
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let views = inner.list_materialized_views().await.infer_error()?;
Ok(views)
})
#[pyo3(signature = (*, page_token=None, page_limit=None))]
pub fn list_materialized_views(
&self,
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::NameListing> {
Ok(crate::listing::NameListing::new(
self.get_inner()?
.list_materialized_views(crate::listing::options(page_token, page_limit)),
))
}
#[pyo3(signature = (name, namespace_path=None))]
@@ -754,31 +747,37 @@ impl Connection {
})
}
#[pyo3(signature = (request_json, namespace_path=None))]
pub fn create_function_async(
self_: PyRef<'_, Self>,
request_json: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let request = lancedb::function::FunctionRegistrationRequest::from_json(&request_json)
.infer_error()?;
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.create_function_async(request)
.create_function_async(request, &namespace_path)
.await
.infer_error()
.map(crate::job::Job::new_typed)
})
}
#[pyo3(signature = (name, version, namespace_path=None))]
pub fn get_function(
self_: PyRef<'_, Self>,
name: String,
version: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.get_function(name, version)
.get_function(name, version, &namespace_path)
.await
.infer_error()?
.to_canonical_json()
@@ -786,39 +785,50 @@ impl Connection {
})
}
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner
.list_functions()
.await
.infer_error()?
.into_iter()
.map(|function| function.to_canonical_json().infer_error())
.collect::<PyResult<Vec<_>>>()
})
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
pub fn list_functions(
&self,
namespace_path: Option<Vec<String>>,
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::FunctionListing> {
Ok(crate::listing::FunctionListing::new(
self.get_inner()?.list_functions(
&namespace_path.unwrap_or_default(),
crate::listing::options(page_token, page_limit),
),
))
}
#[pyo3(signature = (name, version, namespace_path=None))]
pub fn drop_function(
self_: PyRef<'_, Self>,
name: String,
version: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner.drop_function(name, version).await.infer_error()
inner
.drop_function(name, version, &namespace_path)
.await
.infer_error()
})
}
#[pyo3(signature = (name, version, namespace_path=None))]
pub fn drop_function_async(
self_: PyRef<'_, Self>,
name: String,
version: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.drop_function_async(name, version)
.drop_function_async(name, version, &namespace_path)
.await
.infer_error()
.map(|(dropped, job)| (dropped, crate::job::Job::new(job)))
@@ -859,16 +869,19 @@ impl Connection {
})
}
#[pyo3(signature = (namespace_path=None))]
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
pub fn list_secrets(
self_: PyRef<'_, Self>,
&self,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner.list_secrets(&namespace_path).await.infer_error()
})
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::NameListing> {
Ok(crate::listing::NameListing::new(
self.get_inner()?.list_secrets(
&namespace_path.unwrap_or_default(),
crate::listing::options(page_token, page_limit),
),
))
}
#[pyo3(signature = (name, namespace_path=None))]
@@ -969,27 +982,31 @@ impl Connection {
})
}
#[pyo3(signature = (namespace_path=None))]
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
pub fn list_views(
self_: PyRef<'_, Self>,
&self,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner.list_views(&namespace_path).await.infer_error()
})
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::NameListing> {
Ok(crate::listing::NameListing::new(
self.get_inner()?.list_views(
&namespace_path.unwrap_or_default(),
crate::listing::options(page_token, page_limit),
),
))
}
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let jobs = inner.list_jobs().await.infer_error()?;
Ok(jobs
.into_iter()
.map(crate::job::JobInfo::from)
.collect::<Vec<_>>())
})
#[pyo3(signature = (*, page_token=None, page_limit=None))]
pub fn list_jobs(
&self,
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::JobListing> {
Ok(crate::listing::JobListing::new(
self.get_inner()?
.list_jobs(crate::listing::options(page_token, page_limit)),
))
}
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
@@ -1203,6 +1220,7 @@ pub struct PyClientConfig {
user_agent: String,
retry_config: Option<PyClientRetryConfig>,
timeout_config: Option<PyClientTimeoutConfig>,
blob_request_concurrency: Option<usize>,
extra_headers: Option<HashMap<String, String>>,
id_delimiter: Option<String>,
tls_config: Option<PyClientTlsConfig>,
@@ -1288,6 +1306,7 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
user_agent: value.user_agent,
retry_config: value.retry_config.map(Into::into).unwrap_or_default(),
timeout_config: value.timeout_config.map(Into::into).unwrap_or_default(),
blob_request_concurrency: value.blob_request_concurrency,
extra_headers: value.extra_headers.unwrap_or_default(),
id_delimiter: value.id_delimiter,
tls_config: value.tls_config.map(Into::into),
+8 -4
View File
@@ -54,8 +54,6 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
.with_position(params.with_position)
.lower_case(params.lower_case)
.max_token_length(params.max_token_length)
.remove_stop_words(params.remove_stop_words)
.stem(params.stem)
.ascii_folding(params.ascii_folding)
.ngram_min_length(params.ngram_min_length)
.ngram_max_length(params.ngram_max_length)
@@ -65,6 +63,12 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
DocumentGranularity::try_from(params.document_granularity.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))?,
);
if let Some(stem) = params.stem {
inner_opts = inner_opts.stem(stem);
}
if let Some(remove_stop_words) = params.remove_stop_words {
inner_opts = inner_opts.remove_stop_words(remove_stop_words);
}
if let Some(memory_limit) = params.memory_limit {
inner_opts = inner_opts.memory_limit_mb(memory_limit);
}
@@ -215,8 +219,8 @@ struct FtsParams {
language: String,
max_token_length: Option<usize>,
lower_case: bool,
stem: bool,
remove_stop_words: bool,
stem: Option<bool>,
remove_stop_words: Option<bool>,
custom_stop_words: Option<Vec<String>>,
ascii_folding: bool,
ngram_min_length: u32,
+13 -2
View File
@@ -16,12 +16,13 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery};
use session::Session;
use table::{
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, LsmWriteSpec, MergeResult,
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, MemWalWriteSpec, MergeResult,
PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult, Table,
UpdateFieldMetadataResult, UpdateResult,
};
pub mod arrow;
pub mod authz;
pub mod catalog;
pub mod connection;
pub mod error;
@@ -29,6 +30,7 @@ pub mod expr;
pub mod header;
pub mod index;
pub mod job;
mod listing;
pub mod namespace;
pub mod oauth;
pub mod otel;
@@ -64,6 +66,15 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
env_logger::init_from_env(env);
m.add_class::<Connection>()?;
m.add_class::<catalog::Catalog>()?;
m.add_class::<catalog::DatabaseNames>()?;
m.add_class::<listing::NameListing>()?;
m.add_class::<listing::FunctionListing>()?;
m.add_class::<listing::JobListing>()?;
m.add_class::<authz::AuthorizationClient>()?;
m.add_class::<authz::AuthzSubject>()?;
m.add_class::<authz::AuthzObject>()?;
m.add_class::<authz::AuthzPrivilege>()?;
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
m.add_class::<Session>()?;
m.add_class::<Table>()?;
@@ -93,7 +104,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<UpdateFieldMetadataResult>()?;
m.add_class::<AddResult>()?;
m.add_class::<MergeResult>()?;
m.add_class::<LsmWriteSpec>()?;
m.add_class::<MemWalWriteSpec>()?;
m.add_class::<DeleteResult>()?;
m.add_class::<DropColumnsResult>()?;
m.add_class::<UpdateResult>()?;
+131
View File
@@ -0,0 +1,131 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
use futures::{StreamExt, future::poll_fn};
use lancedb::listing::{Listing, ListingOptions};
use pyo3::exceptions::PyStopAsyncIteration;
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
use std::sync::{Arc, Mutex};
pub fn options(page_token: Option<String>, page_limit: Option<u32>) -> ListingOptions {
let mut options = ListingOptions::default();
options.page_token = page_token;
options.page_limit = page_limit;
options
}
#[pyclass]
pub struct NameListing {
inner: Arc<Mutex<Listing<String>>>,
next_lock: Arc<tokio::sync::Mutex<()>>,
}
impl NameListing {
pub fn new(inner: Listing<String>) -> Self {
Self {
inner: Arc::new(Mutex::new(inner)),
next_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
}
#[pymethods]
impl NameListing {
fn num_page_results(&self, py: Python<'_>) -> usize {
py.detach(|| self.inner.lock().unwrap().num_page_results())
}
fn page_token(&self, py: Python<'_>) -> Option<String> {
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
}
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
self_
}
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let _guard = next_lock.lock().await;
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
.await
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
.infer_error()?;
Ok(value)
})
}
}
#[pyclass]
pub struct FunctionListing {
inner: Arc<Mutex<Listing<lancedb::function::FunctionVersion>>>,
next_lock: Arc<tokio::sync::Mutex<()>>,
}
impl FunctionListing {
pub fn new(inner: Listing<lancedb::function::FunctionVersion>) -> Self {
Self {
inner: Arc::new(Mutex::new(inner)),
next_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
}
#[pymethods]
impl FunctionListing {
fn num_page_results(&self, py: Python<'_>) -> usize {
py.detach(|| self.inner.lock().unwrap().num_page_results())
}
fn page_token(&self, py: Python<'_>) -> Option<String> {
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
}
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
self_
}
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let _guard = next_lock.lock().await;
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
.await
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
.infer_error()?;
value.to_canonical_json().infer_error()
})
}
}
#[pyclass]
pub struct JobListing {
inner: Arc<Mutex<Listing<lancedb::database::JobInfo>>>,
next_lock: Arc<tokio::sync::Mutex<()>>,
}
impl JobListing {
pub fn new(inner: Listing<lancedb::database::JobInfo>) -> Self {
Self {
inner: Arc::new(Mutex::new(inner)),
next_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
}
#[pymethods]
impl JobListing {
fn num_page_results(&self, py: Python<'_>) -> usize {
py.detach(|| self.inner.lock().unwrap().num_page_results())
}
fn page_token(&self, py: Python<'_>) -> Option<String> {
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
}
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
self_
}
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let _guard = next_lock.lock().await;
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
.await
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
.infer_error()?;
Ok(crate::job::JobInfo::from(value))
})
}
}
+58 -28
View File
@@ -71,7 +71,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
match class_name(&ob)?.as_str() {
"MatchQuery" => {
let query = ob.getattr("query")?.extract()?;
let column = ob.getattr("column")?.extract()?;
let column = ob.getattr("column")?.extract::<Option<String>>()?;
let boost = ob.getattr("boost")?.extract()?;
let fuzziness = ob.getattr("fuzziness")?.extract()?;
let max_expansions = ob.getattr("max_expansions")?.extract()?;
@@ -88,7 +88,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
let mut query =
MatchQuery::new(query)
.with_column(Some(column))
.with_column(column)
.with_boost(boost)
.with_fuzziness(fuzziness)
.with_max_expansions(max_expansions)
@@ -180,6 +180,9 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
.expect("Failed to import namespace");
match self.0 {
FtsQuery::CombinedFields(_) => Err(PyValueError::new_err(
"Combined-fields FTS queries are not yet supported in Python",
)),
FtsQuery::Match(query) => {
let kwargs = PyDict::new(py);
kwargs.set_item("boost", query.boost)?;
@@ -196,7 +199,7 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
}
namespace
.getattr(intern!(py, "MatchQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
.call((query.terms, query.column), Some(&kwargs))
}
FtsQuery::Phrase(query) => {
let kwargs = PyDict::new(py);
@@ -330,13 +333,13 @@ pub struct PyQueryRequest {
pub select_source_columns: Option<HashMap<String, String>>,
pub fast_search: Option<bool>,
pub with_row_id: Option<bool>,
pub use_lsm: Option<bool>,
pub use_memwal: Option<bool>,
pub column: Option<String>,
pub query_vector: Option<PyQueryVectors>,
pub nprobes: Option<usize>,
pub minimum_nprobes: Option<usize>,
// None means user did not set it and default should be used (currently 20)
// Some(0) means user set it to None and there is no limit
pub maximum_nprobes: Option<usize>,
pub search_effort: Option<f64>,
pub lower_bound: Option<f32>,
pub upper_bound: Option<f32>,
pub ef: Option<usize>,
@@ -363,11 +366,13 @@ impl From<AnyQuery> for PyQueryRequest {
select: PySelect(query_request.select),
fast_search: Some(query_request.fast_search),
with_row_id: Some(query_request.with_row_id),
use_lsm: query_request.use_lsm,
use_memwal: query_request.use_memwal,
column: None,
query_vector: None,
nprobes: None,
minimum_nprobes: None,
maximum_nprobes: None,
search_effort: None,
lower_bound: None,
upper_bound: None,
ef: None,
@@ -390,14 +395,13 @@ impl From<AnyQuery> for PyQueryRequest {
select: PySelect(vector_query.base.select),
fast_search: Some(vector_query.base.fast_search),
with_row_id: Some(vector_query.base.with_row_id),
use_lsm: vector_query.base.use_lsm,
use_memwal: vector_query.base.use_memwal,
column: vector_query.column,
query_vector: Some(PyQueryVectors(vector_query.query_vector)),
minimum_nprobes: Some(vector_query.minimum_nprobes),
maximum_nprobes: match vector_query.maximum_nprobes {
None => Some(0),
Some(value) => Some(value),
},
nprobes: vector_query.nprobes,
minimum_nprobes: vector_query.minimum_nprobes,
maximum_nprobes: vector_query.maximum_nprobes,
search_effort: vector_query.search_effort,
lower_bound: vector_query.lower_bound,
upper_bound: vector_query.upper_bound,
ef: vector_query.ef,
@@ -536,8 +540,8 @@ impl Query {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -702,8 +706,8 @@ impl TakeQuery {
self.inner = self.inner.clone().with_row_id();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
#[pyo3(signature = ())]
@@ -815,8 +819,8 @@ impl FTSQuery {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -966,8 +970,8 @@ impl VectorQuery {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -992,8 +996,18 @@ impl VectorQuery {
self.inner = self.inner.clone().refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?;
Ok(())
}
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
self.inner = self
.inner
.clone()
.search_effort(search_effort)
.infer_error()?;
Ok(())
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
@@ -1164,9 +1178,9 @@ impl HybridQuery {
self.inner_fts.postfilter();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner_vec.use_lsm(enable);
self.inner_fts.use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner_vec.use_memwal(enable);
self.inner_fts.use_memwal(enable);
}
pub fn add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> {
@@ -1185,8 +1199,20 @@ impl HybridQuery {
self.inner_vec.refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner_vec.nprobes(nprobe);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner_vec.nprobes(nprobe)
}
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
self.inner_vec.search_effort(search_effort)
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
self.inner_vec.minimum_nprobes(minimum_nprobes)
}
pub fn maximum_nprobes(&mut self, maximum_nprobes: u32) -> PyResult<()> {
self.inner_vec.maximum_nprobes(maximum_nprobes)
}
pub fn ef(&mut self, ef: u32) {
@@ -1233,6 +1259,10 @@ impl HybridQuery {
req.query_vector = vec_req.query_vector;
req.column = vec_req.column;
req.distance_type = vec_req.distance_type;
req.nprobes = vec_req.nprobes;
req.minimum_nprobes = vec_req.minimum_nprobes;
req.maximum_nprobes = vec_req.maximum_nprobes;
req.search_effort = vec_req.search_effort;
req.ef = vec_req.ef;
req.refine_factor = vec_req.refine_factor;
req.lower_bound = vec_req.lower_bound;
+175 -112
View File
@@ -41,16 +41,16 @@ use pyo3::{
mod scannable;
/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
/// Convert `MemWalStats` to a Python dict, preserving the per-table-shard list.
///
/// Deliberately not flattened to a table-level summary: a table is N
/// buckets on one node, and the per-bucket detail is the reason the
/// endpoint exists — flattening hides the single hot bucket someone opened
/// table shards on one node, and the per-shard detail is the reason the
/// endpoint exists — flattening hides the single hot table shard someone opened
/// it to find.
fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
fn memwal_stats_to_py(py: Python<'_>, stats: &lancedb::table::MemWalStats) -> PyResult<Py<PyDict>> {
let out = PyDict::new(py);
let buckets = PyList::empty(py);
for b in &stats.buckets {
let table_shards = PyList::empty(py);
for b in &stats.table_shards {
let e = PyDict::new(py);
e.set_item("shard_id", &b.shard_id)?;
e.set_item("status", &b.status)?;
@@ -66,15 +66,15 @@ fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult
b.wal_entry_position_last_seen,
)?;
let generations = PyList::empty(py);
for g in &b.generations {
let sstables = PyList::empty(py);
for g in &b.sstables {
let ge = PyDict::new(py);
ge.set_item("generation", g.generation)?;
ge.set_item("bytes", g.bytes)?;
ge.set_item("rows", g.rows)?;
generations.append(ge)?;
sstables.append(ge)?;
}
e.set_item("generations", generations)?;
e.set_item("sstables", sstables)?;
e.set_item("compacting", b.compacting)?;
e.set_item(
@@ -96,9 +96,9 @@ fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult
})
.transpose()?,
)?;
buckets.append(e)?;
table_shards.append(e)?;
}
out.set_item("buckets", buckets)?;
out.set_item("table_shards", table_shards)?;
Ok(out.unbind())
}
@@ -399,20 +399,20 @@ fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
/// MemWAL supports, resolved on install.
/// table has, including ones created later.
#[pyclass(module = "lancedb._lancedb", from_py_object)]
#[derive(Clone, Debug)]
pub struct LsmWriteSpec {
inner: lancedb::table::LsmWriteSpec,
pub struct MemWalWriteSpec {
inner: lancedb::table::MemWalWriteSpec,
}
#[pymethods]
impl LsmWriteSpec {
impl MemWalWriteSpec {
/// Hash-bucket sharding by the unenforced primary key column.
#[staticmethod]
pub fn bucket(column: String, num_buckets: u32) -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::bucket(column, num_buckets),
inner: lancedb::table::MemWalWriteSpec::bucket(column, num_buckets),
}
}
@@ -426,7 +426,7 @@ impl LsmWriteSpec {
#[staticmethod]
pub fn identity(column: String) -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::identity(column),
inner: lancedb::table::MemWalWriteSpec::identity(column),
}
}
@@ -435,13 +435,13 @@ impl LsmWriteSpec {
#[staticmethod]
pub fn unsharded() -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::unsharded(),
inner: lancedb::table::MemWalWriteSpec::unsharded(),
}
}
/// Set which indexes the MemWAL maintains. `None` (the default)
/// resolves every supported index on install; a list is verbatim,
/// and an empty list maintains nothing.
/// Set which indexes the MemWAL maintains. `None` (the default) is
/// every index the table has, including ones created later; a list is
/// verbatim, and an empty list maintains nothing.
#[pyo3(signature = (indexes))]
pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
Self {
@@ -459,33 +459,33 @@ impl LsmWriteSpec {
pub fn __repr__(&self) -> String {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket {
lancedb::table::MemWalWriteSpec::Bucket {
column,
num_buckets,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
column,
num_buckets,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Identity {
lancedb::table::MemWalWriteSpec::Identity {
column,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
column,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Unsharded {
lancedb::table::MemWalWriteSpec::Unsharded {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
@@ -497,9 +497,9 @@ impl LsmWriteSpec {
#[getter]
pub fn spec_type(&self) -> &'static str {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { .. } => "bucket",
lancedb::table::LsmWriteSpec::Identity { .. } => "identity",
lancedb::table::LsmWriteSpec::Unsharded { .. } => "unsharded",
lancedb::table::MemWalWriteSpec::Bucket { .. } => "bucket",
lancedb::table::MemWalWriteSpec::Identity { .. } => "identity",
lancedb::table::MemWalWriteSpec::Unsharded { .. } => "unsharded",
}
}
@@ -507,9 +507,9 @@ impl LsmWriteSpec {
#[getter]
pub fn column(&self) -> Option<String> {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { column, .. }
| lancedb::table::LsmWriteSpec::Identity { column, .. } => Some(column.clone()),
lancedb::table::LsmWriteSpec::Unsharded { .. } => None,
lancedb::table::MemWalWriteSpec::Bucket { column, .. }
| lancedb::table::MemWalWriteSpec::Identity { column, .. } => Some(column.clone()),
lancedb::table::MemWalWriteSpec::Unsharded { .. } => None,
}
}
@@ -517,7 +517,7 @@ impl LsmWriteSpec {
#[getter]
pub fn num_buckets(&self) -> Option<u32> {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
lancedb::table::MemWalWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
_ => None,
}
}
@@ -535,14 +535,14 @@ impl LsmWriteSpec {
}
}
impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
fn from(spec: LsmWriteSpec) -> Self {
impl From<MemWalWriteSpec> for lancedb::table::MemWalWriteSpec {
fn from(spec: MemWalWriteSpec) -> Self {
spec.inner
}
}
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
impl From<lancedb::table::MemWalWriteSpec> for MemWalWriteSpec {
fn from(inner: lancedb::table::MemWalWriteSpec) -> Self {
Self { inner }
}
}
@@ -724,8 +724,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
fn from(result: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match result.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
@@ -831,6 +829,14 @@ impl PyBlobFile {
}
}
fn blob_file_error(operation: &str, error: lance_core::Error) -> pyo3::PyErr {
let message = format!("blob {operation} failed: {error}");
match error {
lance_core::Error::InvalidInput { .. } => PyValueError::new_err(message),
_ => PyRuntimeError::new_err(message),
}
}
#[pymethods]
impl PyBlobFile {
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
@@ -839,7 +845,7 @@ impl PyBlobFile {
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
.map_err(|e| blob_file_error("read", e))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
}
@@ -847,10 +853,7 @@ impl PyBlobFile {
self_.ensure_open()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let bytes = inner
.read()
.await
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
let bytes = inner.read().await.map_err(|e| blob_file_error("read", e))?;
Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind()))
})
}
@@ -873,7 +876,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.close().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
.map_err(|e| blob_file_error("close", e))
}
fn is_closed(self_: PyRef<'_, Self>) -> bool {
@@ -886,7 +889,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.seek(position).await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
.map_err(|e| blob_file_error("seek", e))
}
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
@@ -895,7 +898,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.tell().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
.map_err(|e| blob_file_error("tell", e))
}
fn size(self_: PyRef<'_, Self>) -> u64 {
@@ -912,17 +915,39 @@ impl PyBlobFile {
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
.map_err(|e| blob_file_error("read_range", e))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
}
/// Read blob-local byte ranges without moving the cursor.
fn read_ranges(self_: PyRef<'_, Self>, ranges: Vec<(u64, u64)>) -> PyResult<Vec<Py<PyBytes>>> {
let ranges = ranges
.into_iter()
.map(|(offset, length)| {
offset
.checked_add(length)
.map(|end| offset..end)
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))
})
.collect::<PyResult<Vec<_>>>()?;
let inner = self_.inner.clone();
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read_ranges(&ranges).await }))
.map_err(|e| blob_file_error("read_ranges", e))?;
Ok(bytes
.iter()
.map(|bytes| PyBytes::new(py, bytes.as_ref()).unbind())
.collect())
}
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
self_.ensure_open()?;
let inner = self_.inner.clone();
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read_up_to(length).await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
.map_err(|e| blob_file_error("read_up_to", e))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
}
}
@@ -957,8 +982,8 @@ impl From<LanceDbFtsToken> for FtsToken {
language = "English".to_string(),
max_token_length = Some(40),
lower_case = true,
stem = true,
remove_stop_words = true,
stem = None,
remove_stop_words = None,
custom_stop_words = None,
ascii_folding = true,
ngram_min_length = 3,
@@ -972,15 +997,15 @@ pub fn tokenize(
language: String,
max_token_length: Option<u32>,
lower_case: bool,
stem: bool,
remove_stop_words: bool,
stem: Option<bool>,
remove_stop_words: Option<bool>,
custom_stop_words: Option<Vec<String>>,
ascii_folding: bool,
ngram_min_length: u32,
ngram_max_length: u32,
prefix_only: bool,
) -> PyResult<Vec<FtsToken>> {
let params = FtsIndexBuilder::default()
let mut params = FtsIndexBuilder::default()
.base_tokenizer(base_tokenizer)
.language(&language)
.map_err(|_| {
@@ -991,13 +1016,17 @@ pub fn tokenize(
})?
.max_token_length(max_token_length.map(|value| value as usize))
.lower_case(lower_case)
.stem(stem)
.remove_stop_words(remove_stop_words)
.ascii_folding(ascii_folding)
.ngram_min_length(ngram_min_length)
.ngram_max_length(ngram_max_length)
.ngram_prefix_only(prefix_only)
.custom_stop_words(custom_stop_words);
if let Some(stem) = stem {
params = params.stem(stem);
}
if let Some(remove_stop_words) = remove_stop_words {
params = params.remove_stop_words(remove_stop_words);
}
let tokens = lancedb_tokenize(&query, &params).infer_error()?;
Ok(tokens.into_iter().map(FtsToken::from).collect())
}
@@ -1047,6 +1076,11 @@ impl Table {
self.inner.is_some()
}
/// Whether this table has a local Lance dataset that the client may open.
pub fn _is_native(&self) -> PyResult<bool> {
Ok(self.inner_ref()?.as_native().is_some())
}
/// Closes the table, releasing any resources associated with it.
pub fn close(&mut self) {
self.inner.take();
@@ -1400,6 +1434,7 @@ impl Table {
let dict = PyDict::new(py);
dict.set_item("total_bytes", stats.total_bytes)?;
dict.set_item("num_rows", stats.num_rows)?;
dict.set_item("num_deleted_rows", stats.num_deleted_rows)?;
dict.set_item("num_indices", stats.num_indices)?;
let fragment_stats = PyDict::new(py);
@@ -1586,14 +1621,30 @@ impl Table {
pub fn fetch_blob_ranges(
self_: PyRef<'_, Self>,
column: String,
requests: Vec<(u64, u64, u64)>,
requests: Vec<Vec<i128>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
let requests = requests
.into_iter()
.map(|values| {
let [row_id, offset, length] = <[i128; 3]>::try_from(values).map_err(|_| {
PyValueError::new_err("each blob range request must have three values")
})?;
let to_u64 = |value, name| {
u64::try_from(value).map_err(|_| {
PyValueError::new_err(format!(
"blob range {name} must be a non-negative 64-bit integer"
))
})
};
Ok(BlobRangeRequest::new(
to_u64(row_id, "row id")?,
to_u64(offset, "offset")?,
to_u64(length, "length")?,
))
})
.collect::<PyResult<Vec<_>>>()?;
future_into_py(self_.py(), async move {
let requests = requests
.into_iter()
.map(|(row_id, offset, length)| BlobRangeRequest::new(row_id, offset, length))
.collect::<Vec<_>>();
let blobs: LargeBinaryArray = inner
.fetch_blob_ranges(column, requests)
.await
@@ -1719,8 +1770,8 @@ impl Table {
if let Some(use_index) = parameters.use_index {
builder.use_index(use_index);
}
if let Some(use_lsm) = parameters.use_lsm {
builder.use_lsm(use_lsm);
if let Some(use_memwal) = parameters.use_memwal {
builder.use_memwal(use_memwal);
}
if let Some(validate_single_shard) = parameters.validate_single_shard {
builder.validate_single_shard(validate_single_shard);
@@ -1745,29 +1796,29 @@ impl Table {
})
}
pub fn set_lsm_write_spec<'a>(
pub fn enable_memwal<'a>(
self_: PyRef<'a, Self>,
spec: LsmWriteSpec,
spec: MemWalWriteSpec,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.inner_ref()?.clone();
let native_spec = lancedb::table::LsmWriteSpec::from(spec);
let native_spec = lancedb::table::MemWalWriteSpec::from(spec);
future_into_py(self_.py(), async move {
inner.set_lsm_write_spec(native_spec).await.infer_error()
inner.enable_memwal(native_spec).await.infer_error()
})
}
pub fn unset_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn disable_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.unset_lsm_write_spec().await.infer_error()
inner.disable_memwal().await.infer_error()
})
}
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn get_memwal_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let spec = inner.get_lsm_write_spec().await.infer_error()?;
Ok(spec.map(LsmWriteSpec::from))
let spec = inner.get_memwal_spec().await.infer_error()?;
Ok(spec.map(MemWalWriteSpec::from))
})
}
@@ -1786,55 +1837,54 @@ impl Table {
Ok(())
}
/// Converge the table's LSM write path into its base table.
/// Converge the table's MemWAL write path into its base table.
///
/// Best-effort: with writes flowing, new rows may land after the last
/// pass. Errors if the table stops making progress.
pub fn checkpoint_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn checkpoint_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.checkpoint_lsm().await.infer_error()
inner.checkpoint_memwal().await.infer_error()
})
}
/// Seal every bucket's active memtable into L0.
pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(
self_.py(),
async move { inner.flush_lsm().await.infer_error() },
)
}
/// Trigger a background L0 → base pass per bucket. Returns once the
/// passes are dispatched, not once they finish — watch `get_lsm_stats`.
pub fn compact_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
/// Freeze every table shard's active memtable into an SSTable.
pub fn flush_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_lsm().await.infer_error()
inner.flush_memwal().await.infer_error()
})
}
/// Live LSM state, or `None` when the LSM write path is not enabled.
#[pyo3(signature = (include_generation_rows=false))]
pub fn get_lsm_stats(
/// Trigger a background SSTable compaction pass per table shard. Returns once the
/// passes are dispatched, not once they finish — watch `get_memwal_stats`.
pub fn compact_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_memwal().await.infer_error()
})
}
/// Live MemWAL state, or `None` when the MemWAL write path is not enabled.
#[pyo3(signature = (include_sstable_rows=false))]
pub fn get_memwal_stats(
self_: PyRef<'_, Self>,
include_generation_rows: bool,
include_sstable_rows: bool,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let stats = inner
.get_lsm_stats(include_generation_rows)
.get_memwal_stats(include_sstable_rows)
.await
.infer_error()?;
Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
Python::attach(|py| stats.map(|s| memwal_stats_to_py(py, &s)).transpose())
})
}
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn close_memwal_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.close_lsm_writers().await.infer_error()
inner.close_memwal_writers().await.infer_error()
})
}
@@ -1843,7 +1893,10 @@ impl Table {
future_into_py(self_.py(), async move {
inner
.as_native()
.ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))?
.ok_or_else(|| lancedb::Error::NotSupported {
message: "uses_v2_manifest_paths is not supported for remote tables.".into(),
})
.infer_error()?
.uses_v2_manifest_paths()
.await
.infer_error()
@@ -1855,7 +1908,10 @@ impl Table {
future_into_py(self_.py(), async move {
inner
.as_native()
.ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))?
.ok_or_else(|| lancedb::Error::NotSupported {
message: "migrate_manifest_paths_v2 is not supported for remote tables.".into(),
})
.infer_error()?
.migrate_manifest_paths_v2()
.await
.infer_error()
@@ -1951,10 +2007,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1962,7 +2017,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1971,10 +2026,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view_async(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1982,7 +2036,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1997,11 +2051,9 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
view.definition().to_json().map_err(|err| {
PyRuntimeError::new_err(format!(
"failed to serialize materialized-view definition: {err}"
))
})
Ok(lancedb::materialized_view::definition_metadata_from_sql(
view.definition_sql(),
))
})
}
@@ -2031,6 +2083,17 @@ impl Table {
let alterations = alterations
.iter()
.map(|alteration| {
for key in alteration.keys().iter() {
let key = key.extract::<String>()?;
if !matches!(
key.as_str(),
"path" | "rename" | "name" | "nullable" | "data_type"
) {
return Err(PyValueError::new_err(format!(
"Unknown column alteration key '{key}'"
)));
}
}
let path = alteration
.get_item("path")?
.ok_or_else(|| PyValueError::new_err("Missing path"))?
@@ -2153,7 +2216,7 @@ pub struct MergeInsertParams {
when_not_matched_by_source_condition_expr: Option<PyExpr>,
timeout: Option<std::time::Duration>,
use_index: Option<bool>,
use_lsm: Option<bool>,
use_memwal: Option<bool>,
validate_single_shard: Option<bool>,
}