mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 04:02:27 +00:00
Merge main into gatekeeper/fix-2325-1
This commit is contained in:
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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>()?;
|
||||
|
||||
@@ -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
@@ -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
@@ -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, ¶ms).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>,
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user