rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-30 01:48:19 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
ab74aa620cef4a84fc2f01abf1c7dc0f17cb160c
lancedb/python/src/query.rs
T
lancedb-gatefixer[bot] 25645d82d4 feat(python): accept expressions in update filters (#3876) 
## Summary
- allow Python sync, async, and remote table updates to accept type-safe
`Expr` filters
- serialize expression filters before invoking the existing update
implementation
- cover numeric-looking text and apostrophe-containing text in sync and
async regression tests

## Root cause
`Table.update` was the remaining Python write path that required callers
to construct a raw SQL predicate. Dynamic text interpolated without SQL
literal encoding could therefore be parsed as an integer, float, or
unterminated string instead of Utf8. The expression API already encodes
literals safely for query and delete filters.

## Validation
- `cd python && .venv/bin/pytest
python/tests/test_table.py::test_update_async
python/tests/test_table.py::test_update_expr_filter_literals_async
python/tests/test_table.py::test_update
python/tests/test_table.py::test_update_expr_filter_literals -q`
- `cd python && .venv/bin/pytest python/tests/test_expr.py -q`
- `cd python && .venv/bin/ruff format --check .`
- `cd python && .venv/bin/ruff check .`

Fixes #1869

<!-- lance-gatekeeper-fix:v1 agent=01f1e7b69c65e8b6d3b3c1e1a7918179
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 20:28:21 +08:00

1252 lines
44 KiB
Rust
Raw Blame History

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use crate::expr::PyExpr;
use crate::runtime::future_into_py;
use crate::util::parse_distance_type;
use crate::{arrow::RecordBatchStream, util::PyLanceDB};
use crate::{error::PythonErrorExt, index::class_name};
use arrow::array::Array;
use arrow::array::ArrayData;
use arrow::array::make_array;
use arrow::pyarrow::FromPyArrow;
use arrow::pyarrow::IntoPyArrow;
use arrow::pyarrow::ToPyArrow;
use lancedb::index::scalar::{
BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery,
MultiMatchQuery, Occur, Operator, PhraseQuery,
};
use lancedb::query::AnalyzePlanDistributedMetrics;
use lancedb::query::QueryBase;
use lancedb::query::QueryExecutionOptions;
use lancedb::query::QueryFilter;
use lancedb::query::{
ColumnOrdering, ExecutableQuery, Query as LanceDbQuery, Select, TakeQuery as LanceDbTakeQuery,
VectorQuery as LanceDbVectorQuery,
};
use lancedb::table::AnyQuery;
use pyo3::Bound;
use pyo3::IntoPyObject;
use pyo3::PyAny;
use pyo3::PyRef;
use pyo3::PyResult;
use pyo3::Python;
use pyo3::prelude::{PyAnyMethods, PyDictMethods};
use pyo3::pyfunction;
use pyo3::pymethods;
use pyo3::types::PyList;
use pyo3::types::{PyDict, PyString};
use pyo3::{Borrowed, FromPyObject, exceptions::PyRuntimeError};
use pyo3::{PyErr, pyclass};
use pyo3::{exceptions::PyValueError, intern};
fn analyze_plan_options(distributed_metrics: Option<&str>) -> PyResult<QueryExecutionOptions> {
let analyze_plan_distributed_metrics = match distributed_metrics.unwrap_or("aggregate") {
"aggregate" => AnalyzePlanDistributedMetrics::Aggregate,
"per_worker" => AnalyzePlanDistributedMetrics::PerWorker,
"full" => AnalyzePlanDistributedMetrics::Full,
mode => {
return Err(PyValueError::new_err(format!(
"Invalid distributed_metrics value '{}'. Expected one of: \
'aggregate', 'per_worker', 'full'",
mode
)));
}
};
let mut options = QueryExecutionOptions::default();
options.analyze_plan_distributed_metrics = analyze_plan_distributed_metrics;
Ok(options)
}
impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
type Error = PyErr;
fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let ob = ob.to_owned();
match class_name(&ob)?.as_str() {
"MatchQuery" => {
let query = ob.getattr("query")?.extract()?;
let column = ob.getattr("column")?.extract()?;
let boost = ob.getattr("boost")?.extract()?;
let fuzziness = ob.getattr("fuzziness")?.extract()?;
let max_expansions = ob.getattr("max_expansions")?.extract()?;
let operator = ob.getattr("operator")?.extract::<String>()?;
let prefix_length = ob.getattr("prefix_length")?.extract()?;
let document_granularity = ob
.getattr("document_granularity")?
.extract::<Option<String>>()?
.map(|value| {
DocumentGranularity::try_from(value.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))
})
.transpose()?;
let mut query =
MatchQuery::new(query)
.with_column(Some(column))
.with_boost(boost)
.with_fuzziness(fuzziness)
.with_max_expansions(max_expansions)
.with_operator(Operator::try_from(operator.as_str()).map_err(|e| {
PyValueError::new_err(format!("Invalid operator: {}", e))
})?)
.with_prefix_length(prefix_length);
if let Some(document_granularity) = document_granularity {
query = query.with_document_granularity(document_granularity);
}
Ok(Self(query.into()))
}
"PhraseQuery" => {
let query = ob.getattr("query")?.extract()?;
let column = ob.getattr("column")?.extract()?;
let slop = ob.getattr("slop")?.extract()?;
let document_granularity = ob
.getattr("document_granularity")?
.extract::<Option<String>>()?
.map(|value| {
DocumentGranularity::try_from(value.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))
})
.transpose()?;
let mut query = PhraseQuery::new(query)
.with_column(Some(column))
.with_slop(slop);
if let Some(document_granularity) = document_granularity {
query = query.with_document_granularity(document_granularity);
}
Ok(Self(query.into()))
}
"BoostQuery" => {
let positive: Self = ob.getattr("positive")?.extract()?;
let negative: Self = ob.getattr("negative")?.extract()?;
let negative_boost = ob.getattr("negative_boost")?.extract()?;
Ok(Self(
BoostQuery::new(positive.0, negative.0, negative_boost).into(),
))
}
"MultiMatchQuery" => {
let query = ob.getattr("query")?.extract()?;
let columns = ob.getattr("columns")?.extract()?;
let boosts: Option<Vec<f32>> = ob.getattr("boosts")?.extract()?;
let operator: String = ob.getattr("operator")?.extract()?;
let q = MultiMatchQuery::try_new(query, columns)
.map_err(|e| PyValueError::new_err(format!("Invalid query: {}", e)))?;
let q = if let Some(boosts) = boosts {
q.try_with_boosts(boosts)
.map_err(|e| PyValueError::new_err(format!("Invalid boosts: {}", e)))?
} else {
q
};
let op = Operator::try_from(operator.as_str())
.map_err(|e| PyValueError::new_err(format!("Invalid operator: {}", e)))?;
Ok(Self(q.with_operator(op).into()))
}
"BooleanQuery" => {
let queries: Vec<(String, Self)> = ob.getattr("queries")?.extract()?;
let mut sub_queries = Vec::with_capacity(queries.len());
for (occur, q) in queries {
let occur = Occur::try_from(occur.as_str())
.map_err(|e| PyValueError::new_err(e.to_string()))?;
sub_queries.push((occur, q.0));
}
Ok(Self(BooleanQuery::new(sub_queries).into()))
}
name => Err(PyValueError::new_err(format!(
"Unsupported FTS query type: {}",
name
))),
}
}
}
impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;
type Error = PyErr;
fn into_pyobject(self, py: pyo3::Python<'py>) -> PyResult<Self::Output> {
let namespace = py
.import(intern!(py, "lancedb"))
.and_then(|m| m.getattr(intern!(py, "query")))
.expect("Failed to import namespace");
match self.0 {
FtsQuery::Match(query) => {
let kwargs = PyDict::new(py);
kwargs.set_item("boost", query.boost)?;
kwargs.set_item("fuzziness", query.fuzziness)?;
kwargs.set_item("max_expansions", query.max_expansions)?;
kwargs.set_item::<_, &str>("operator", query.operator.into())?;
kwargs.set_item("prefix_length", query.prefix_length)?;
if let Some(document_granularity) = query.document_granularity {
let value = match document_granularity {
DocumentGranularity::Row => "row",
DocumentGranularity::ListElement => "list_element",
};
kwargs.set_item("document_granularity", value)?;
}
namespace
.getattr(intern!(py, "MatchQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
}
FtsQuery::Phrase(query) => {
let kwargs = PyDict::new(py);
kwargs.set_item("slop", query.slop)?;
if let Some(document_granularity) = query.document_granularity {
let value = match document_granularity {
DocumentGranularity::Row => "row",
DocumentGranularity::ListElement => "list_element",
};
kwargs.set_item("document_granularity", value)?;
}
namespace
.getattr(intern!(py, "PhraseQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
}
FtsQuery::Boost(query) => {
let positive = Self(query.positive.as_ref().clone()).into_pyobject(py)?;
let negative = Self(query.negative.as_ref().clone()).into_pyobject(py)?;
let kwargs = PyDict::new(py);
kwargs.set_item("negative_boost", query.negative_boost)?;
namespace
.getattr(intern!(py, "BoostQuery"))?
.call((positive, negative), Some(&kwargs))
}
FtsQuery::MultiMatch(query) => {
let first = &query.match_queries[0];
let (columns, boosts): (Vec<_>, Vec<_>) = query
.match_queries
.iter()
.map(|q| (q.column.as_ref().unwrap().clone(), q.boost))
.unzip();
let kwargs = PyDict::new(py);
kwargs.set_item("boosts", boosts)?;
kwargs.set_item::<_, &str>("operator", first.operator.into())?;
namespace
.getattr(intern!(py, "MultiMatchQuery"))?
.call((first.terms.clone(), columns), Some(&kwargs))
}
FtsQuery::Boolean(query) => {
let mut queries: Vec<(&str, Bound<'py, PyAny>)> = Vec::with_capacity(
query.should.len() + query.must.len() + query.must_not.len(),
);
for q in query.should {
queries.push((Occur::Should.into(), Self(q).into_pyobject(py)?));
}
for q in query.must {
queries.push((Occur::Must.into(), Self(q).into_pyobject(py)?));
}
for q in query.must_not {
queries.push((Occur::MustNot.into(), Self(q).into_pyobject(py)?));
}
namespace
.getattr(intern!(py, "BooleanQuery"))?
.call1((queries,))
}
}
}
}
// Python representation of query vector(s)
#[derive(Clone)]
pub struct PyQueryVectors(Vec<Arc<dyn Array>>);
#[derive(Clone, FromPyObject)]
#[pyo3(from_item_all)]
pub struct PyColumnOrdering {
pub column_name: String,
pub ascending: bool,
pub nulls_first: bool,
}
impl From<ColumnOrdering> for PyColumnOrdering {
fn from(ordering: ColumnOrdering) -> Self {
Self {
column_name: ordering.column_name,
ascending: ordering.ascending,
nulls_first: ordering.nulls_first,
}
}
}
impl From<PyColumnOrdering> for ColumnOrdering {
fn from(ordering: PyColumnOrdering) -> Self {
Self {
column_name: ordering.column_name,
ascending: ordering.ascending,
nulls_first: ordering.nulls_first,
}
}
}
impl<'py> IntoPyObject<'py> for PyColumnOrdering {
type Target = PyDict;
type Output = Bound<'py, Self::Target>;
type Error = PyErr;
fn into_pyobject(self, py: pyo3::Python<'py>) -> PyResult<Self::Output> {
let dict = PyDict::new(py);
dict.set_item("column_name", self.column_name)?;
dict.set_item("ascending", self.ascending)?;
dict.set_item("nulls_first", self.nulls_first)?;
Ok(dict)
}
}
impl<'py> IntoPyObject<'py> for PyQueryVectors {
type Target = PyList;
type Output = Bound<'py, Self::Target>;
type Error = PyErr;
fn into_pyobject(self, py: pyo3::Python<'py>) -> PyResult<Self::Output> {
let py_objs = self
.0
.into_iter()
.map(|v| v.to_data().into_pyarrow(py))
.collect::<Result<Vec<_>, _>>()?;
PyList::new(py, py_objs)
}
}
// Python representation of a query
#[pyclass(get_all)]
pub struct PyQueryRequest {
pub limit: Option<usize>,
pub offset: Option<usize>,
pub filter: Option<PyQueryFilter>,
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
pub select: PySelect,
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 column: Option<String>,
pub query_vector: Option<PyQueryVectors>,
pub minimum_nprobes: Option<usize>,
// None means user did not set it and default shoud be used (currenty 20)
// Some(0) means user set it to None and there is no limit
pub maximum_nprobes: Option<usize>,
pub lower_bound: Option<f32>,
pub upper_bound: Option<f32>,
pub ef: Option<usize>,
pub refine_factor: Option<u32>,
pub distance_type: Option<String>,
pub bypass_vector_index: Option<bool>,
pub postfilter: Option<bool>,
pub norm: Option<String>,
pub order_by: Option<Vec<PyColumnOrdering>>,
}
impl From<AnyQuery> for PyQueryRequest {
fn from(query: AnyQuery) -> Self {
match query {
AnyQuery::Query(query_request) => Self {
limit: query_request.limit,
offset: query_request.offset,
filter: query_request.filter.map(PyQueryFilter),
full_text_search: query_request
.full_text_search
.map(|fts| PyLanceDB(fts.query)),
select_source_columns: PySelect::source_columns(&query_request.select),
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,
column: None,
query_vector: None,
minimum_nprobes: None,
maximum_nprobes: None,
lower_bound: None,
upper_bound: None,
ef: None,
refine_factor: None,
distance_type: None,
bypass_vector_index: None,
postfilter: None,
norm: None,
order_by: query_request
.order_by
.map(|order_by| order_by.into_iter().map(PyColumnOrdering::from).collect()),
},
AnyQuery::VectorQuery(vector_query) => Self {
limit: vector_query.base.limit,
offset: vector_query.base.offset,
filter: vector_query.base.filter.map(PyQueryFilter),
full_text_search: None,
select_source_columns: PySelect::source_columns(&vector_query.base.select),
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,
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),
},
lower_bound: vector_query.lower_bound,
upper_bound: vector_query.upper_bound,
ef: vector_query.ef,
refine_factor: vector_query.refine_factor,
distance_type: vector_query.distance_type.map(|d| d.to_string()),
bypass_vector_index: Some(!vector_query.use_index),
postfilter: Some(!vector_query.base.prefilter),
norm: vector_query.base.norm.map(|n| n.to_string()),
order_by: vector_query
.base
.order_by
.map(|order_by| order_by.into_iter().map(PyColumnOrdering::from).collect()),
},
}
}
}
// Python representation of query selection
#[derive(Clone)]
pub struct PySelect(Select);
impl PySelect {
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
match select {
Select::Expr(pairs) => Some(
pairs
.iter()
.filter_map(|(output, expr)| match expr {
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
Some((output.clone(), column.name.clone()))
}
_ => None,
})
.collect(),
),
_ => None,
}
}
}
impl<'py> IntoPyObject<'py> for PySelect {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;
type Error = PyErr;
fn into_pyobject(self, py: pyo3::Python<'py>) -> PyResult<Self::Output> {
match self.0 {
Select::All => Ok(py.None().into_bound(py).into_any()),
Select::Columns(columns) => Ok(columns.into_pyobject(py)?.into_any()),
Select::Dynamic(columns) => Ok(columns.into_pyobject(py)?.into_any()),
Select::Expr(pairs) => {
// Serialize DataFusion Expr -> SQL string so Python sees the same
// format as Select::Dynamic: a list of (name, sql_string) tuples.
let sql_pairs: PyResult<Vec<(String, String)>> = pairs
.into_iter()
.map(|(name, expr)| {
lancedb::expr::expr_to_sql_string(&expr)
.map(|sql| (name, sql))
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
.collect();
Ok(sql_pairs?.into_pyobject(py)?.into_any())
}
}
}
}
// Python representation of query filter
#[derive(Clone)]
pub struct PyQueryFilter(QueryFilter);
impl<'py> IntoPyObject<'py> for PyQueryFilter {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;
type Error = PyErr;
fn into_pyobject(self, py: pyo3::Python<'py>) -> PyResult<Self::Output> {
match self.0 {
QueryFilter::Datafusion(expr) => {
// Serialize the DataFusion expression to a SQL string so that
// callers (e.g. remote tables) see the same format as Sql.
let sql = lancedb::expr::expr_to_sql_string(&expr)
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
Ok(sql.into_pyobject(py)?.into_any())
}
QueryFilter::Sql(sql) => Ok(sql.into_pyobject(py)?.into_any()),
QueryFilter::Substrait(substrait) => Ok(substrait.into_pyobject(py)?.into_any()),
}
}
}
#[pyclass]
pub struct Query {
inner: LanceDbQuery,
}
impl Query {
pub fn new(query: LanceDbQuery) -> Self {
Self { inner: query }
}
}
#[pymethods]
impl Query {
pub fn r#where(&mut self, predicate: String) {
self.inner = self.inner.clone().only_if(predicate);
}
pub fn where_expr(&mut self, expr: PyExpr) {
self.inner = self.inner.clone().only_if_expr(expr.0);
}
pub fn select(&mut self, columns: Vec<(String, String)>) {
self.inner = self.inner.clone().select(Select::dynamic(&columns));
}
pub fn select_expr(&mut self, columns: Vec<(String, PyExpr)>) {
let pairs: Vec<(String, lancedb::expr::DfExpr)> =
columns.into_iter().map(|(name, e)| (name, e.0)).collect();
self.inner = self.inner.clone().select(Select::Expr(pairs));
}
pub fn select_columns(&mut self, columns: Vec<String>) {
self.inner = self.inner.clone().select(Select::columns(&columns));
}
pub fn limit(&mut self, limit: u32) {
self.inner = self.inner.clone().limit(limit as usize);
}
pub fn offset(&mut self, offset: u32) {
self.inner = self.inner.clone().offset(offset as usize);
}
pub fn fast_search(&mut self) {
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 with_row_id(&mut self) {
self.inner = self.inner.clone().with_row_id();
}
pub fn postfilter(&mut self) {
self.inner = self.inner.clone().postfilter();
}
pub fn nearest_to(&mut self, vector: Bound<'_, PyAny>) -> PyResult<VectorQuery> {
let data: ArrayData = ArrayData::from_pyarrow_bound(&vector)?;
let array = make_array(data);
let inner = self.inner.clone().nearest_to(array).infer_error()?;
Ok(VectorQuery { inner })
}
pub fn nearest_to_text(&mut self, query: Bound<'_, PyDict>) -> PyResult<FTSQuery> {
let fts_query = query
.get_item("query")?
.ok_or(PyErr::new::<PyRuntimeError, _>(
"Query text is required for nearest_to_text",
))?;
let query = if let Ok(query_text) = fts_query.cast::<PyString>() {
let mut query_text = query_text.to_string();
let columns = query
.get_item("columns")?
.map(|columns| columns.extract::<Vec<String>>())
.transpose()?;
let is_phrase =
query_text.len() >= 2 && query_text.starts_with('"') && query_text.ends_with('"');
let is_multi_match = columns.as_ref().map(|cols| cols.len() > 1).unwrap_or(false);
if is_phrase {
// Remove the surrounding quotes for phrase queries
query_text = query_text[1..query_text.len() - 1].to_string();
}
let query: FtsQuery = match (is_phrase, is_multi_match) {
(false, _) => MatchQuery::new(query_text).into(),
(true, false) => PhraseQuery::new(query_text).into(),
(true, true) => {
return Err(PyValueError::new_err(
"Phrase queries cannot be used with multiple columns.",
));
}
};
let mut query = FullTextSearchQuery::new_query(query);
match columns {
Some(cols) if !cols.is_empty() => {
query = query.with_columns(&cols).map_err(|e| {
PyValueError::new_err(format!(
"Failed to set full text search columns: {}",
e
))
})?;
}
_ => {}
}
query
} else {
let query = fts_query.extract::<PyLanceDB<FtsQuery>>()?;
FullTextSearchQuery::new_query(query.0)
};
Ok(FTSQuery {
inner: self.inner.clone(),
fts_query: query,
})
}
pub fn order_by(&mut self, ordering: Option<Vec<PyColumnOrdering>>) -> PyResult<()> {
let ordering =
ordering.map(|ordering| ordering.into_iter().map(ColumnOrdering::from).collect());
self.inner = self.inner.clone().order_by(ordering);
Ok(())
}
#[pyo3(signature = ())]
pub fn output_schema(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let schema = inner.output_schema().await.infer_error()?;
Python::attach(|py| schema.to_pyarrow(py).map(|obj| obj.unbind()))
})
}
#[pyo3(signature = (max_batch_length=None, timeout=None))]
pub fn execute(
self_: PyRef<'_, Self>,
max_batch_length: Option<u32>,
timeout: Option<Duration>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let mut opts = QueryExecutionOptions::default();
if let Some(max_batch_length) = max_batch_length {
opts.max_batch_length = max_batch_length;
}
if let Some(timeout) = timeout {
opts.timeout = Some(timeout);
}
let inner_stream = inner.execute_with_options(opts).await.infer_error()?;
Ok(RecordBatchStream::new(inner_stream))
})
}
pub fn explain_plan(self_: PyRef<'_, Self>, verbose: bool) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.explain_plan(verbose)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
#[pyo3(signature = (distributed_metrics=None))]
pub fn analyze_plan(
self_: PyRef<'_, Self>,
distributed_metrics: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let options = analyze_plan_options(distributed_metrics.as_deref())?;
future_into_py(self_.py(), async move {
inner
.analyze_plan_with_options(options)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
pub fn to_query_request(&self) -> PyQueryRequest {
PyQueryRequest::from(AnyQuery::Query(self.inner.clone().into_request()))
}
}
#[pyclass]
pub struct TakeQuery {
inner: LanceDbTakeQuery,
}
impl TakeQuery {
pub fn new(query: LanceDbTakeQuery) -> Self {
Self { inner: query }
}
}
#[pymethods]
impl TakeQuery {
pub fn select(&mut self, columns: Vec<(String, String)>) {
self.inner = self.inner.clone().select(Select::dynamic(&columns));
}
pub fn select_columns(&mut self, columns: Vec<String>) {
self.inner = self.inner.clone().select(Select::columns(&columns));
}
pub fn with_row_id(&mut self) {
self.inner = self.inner.clone().with_row_id();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
}
#[pyo3(signature = ())]
pub fn output_schema(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let schema = inner.output_schema().await.infer_error()?;
Python::attach(|py| schema.to_pyarrow(py).map(|obj| obj.unbind()))
})
}
#[pyo3(signature = (max_batch_length=None, timeout=None))]
pub fn execute(
self_: PyRef<'_, Self>,
max_batch_length: Option<u32>,
timeout: Option<Duration>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let mut opts = QueryExecutionOptions::default();
if let Some(max_batch_length) = max_batch_length {
opts.max_batch_length = max_batch_length;
}
if let Some(timeout) = timeout {
opts.timeout = Some(timeout);
}
let inner_stream = inner.execute_with_options(opts).await.infer_error()?;
Ok(RecordBatchStream::new(inner_stream))
})
}
pub fn explain_plan(self_: PyRef<'_, Self>, verbose: bool) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.explain_plan(verbose)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
#[pyo3(signature = (distributed_metrics=None))]
pub fn analyze_plan(
self_: PyRef<'_, Self>,
distributed_metrics: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let options = analyze_plan_options(distributed_metrics.as_deref())?;
future_into_py(self_.py(), async move {
inner
.analyze_plan_with_options(options)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
pub fn to_query_request(&self) -> PyQueryRequest {
PyQueryRequest::from(AnyQuery::Query(self.inner.clone().into_request()))
}
}
#[pyclass(from_py_object)]
#[derive(Clone)]
pub struct FTSQuery {
inner: LanceDbQuery,
fts_query: FullTextSearchQuery,
}
#[pymethods]
impl FTSQuery {
pub fn r#where(&mut self, predicate: String) {
self.inner = self.inner.clone().only_if(predicate);
}
pub fn where_expr(&mut self, expr: PyExpr) {
self.inner = self.inner.clone().only_if_expr(expr.0);
}
pub fn select(&mut self, columns: Vec<(String, String)>) {
self.inner = self.inner.clone().select(Select::dynamic(&columns));
}
pub fn select_expr(&mut self, columns: Vec<(String, PyExpr)>) {
let pairs: Vec<(String, lancedb::expr::DfExpr)> =
columns.into_iter().map(|(name, e)| (name, e.0)).collect();
self.inner = self.inner.clone().select(Select::Expr(pairs));
}
pub fn select_columns(&mut self, columns: Vec<String>) {
self.inner = self.inner.clone().select(Select::columns(&columns));
}
pub fn limit(&mut self, limit: u32) {
self.inner = self.inner.clone().limit(limit as usize);
}
pub fn offset(&mut self, offset: u32) {
self.inner = self.inner.clone().offset(offset as usize);
}
pub fn order_by(&mut self, ordering: Option<Vec<PyColumnOrdering>>) -> PyResult<()> {
let ordering =
ordering.map(|ordering| ordering.into_iter().map(ColumnOrdering::from).collect());
self.inner = self.inner.clone().order_by(ordering);
Ok(())
}
pub fn fast_search(&mut self) {
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 with_row_id(&mut self) {
self.inner = self.inner.clone().with_row_id();
}
pub fn postfilter(&mut self) {
self.inner = self.inner.clone().postfilter();
}
#[pyo3(signature = ())]
pub fn output_schema(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let schema = inner.output_schema().await.infer_error()?;
Python::attach(|py| schema.to_pyarrow(py).map(|obj| obj.unbind()))
})
}
#[pyo3(signature = (max_batch_length=None, timeout=None))]
pub fn execute(
self_: PyRef<'_, Self>,
max_batch_length: Option<u32>,
timeout: Option<Duration>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_
.inner
.clone()
.full_text_search(self_.fts_query.clone());
future_into_py(self_.py(), async move {
let mut opts = QueryExecutionOptions::default();
if let Some(max_batch_length) = max_batch_length {
opts.max_batch_length = max_batch_length;
}
if let Some(timeout) = timeout {
opts.timeout = Some(timeout);
}
let inner_stream = inner.execute_with_options(opts).await.infer_error()?;
Ok(RecordBatchStream::new(inner_stream))
})
}
pub fn nearest_to(&mut self, vector: Bound<'_, PyAny>) -> PyResult<HybridQuery> {
let vector_query = Query::new(self.inner.clone()).nearest_to(vector)?;
Ok(HybridQuery {
inner_fts: self.clone(),
inner_vec: vector_query,
})
}
pub fn explain_plan(self_: PyRef<'_, Self>, verbose: bool) -> PyResult<Bound<'_, PyAny>> {
let inner = self_
.inner
.clone()
.full_text_search(self_.fts_query.clone());
future_into_py(self_.py(), async move {
inner
.explain_plan(verbose)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
#[pyo3(signature = (distributed_metrics=None))]
pub fn analyze_plan(
self_: PyRef<'_, Self>,
distributed_metrics: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_
.inner
.clone()
.full_text_search(self_.fts_query.clone());
let options = analyze_plan_options(distributed_metrics.as_deref())?;
future_into_py(self_.py(), async move {
inner
.analyze_plan_with_options(options)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
pub fn get_query(&self) -> String {
self.fts_query.query.query().clone()
}
pub fn to_query_request(&self) -> PyQueryRequest {
let mut req = self.inner.clone().into_request();
req.full_text_search = Some(self.fts_query.clone());
PyQueryRequest::from(AnyQuery::Query(req))
}
}
#[pyclass(from_py_object)]
#[derive(Clone)]
pub struct VectorQuery {
inner: LanceDbVectorQuery,
}
#[pymethods]
impl VectorQuery {
pub fn r#where(&mut self, predicate: String) {
self.inner = self.inner.clone().only_if(predicate);
}
pub fn where_expr(&mut self, expr: PyExpr) {
self.inner = self.inner.clone().only_if_expr(expr.0);
}
pub fn add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> {
let data: ArrayData = ArrayData::from_pyarrow_bound(&vector)?;
let array = make_array(data);
self.inner = self.inner.clone().add_query_vector(array).infer_error()?;
Ok(())
}
pub fn select(&mut self, columns: Vec<(String, String)>) {
self.inner = self.inner.clone().select(Select::dynamic(&columns));
}
pub fn select_expr(&mut self, columns: Vec<(String, PyExpr)>) {
let pairs: Vec<(String, lancedb::expr::DfExpr)> =
columns.into_iter().map(|(name, e)| (name, e.0)).collect();
self.inner = self.inner.clone().select(Select::Expr(pairs));
}
pub fn select_columns(&mut self, columns: Vec<String>) {
self.inner = self.inner.clone().select(Select::columns(&columns));
}
pub fn limit(&mut self, limit: u32) {
self.inner = self.inner.clone().limit(limit as usize);
}
pub fn offset(&mut self, offset: u32) {
self.inner = self.inner.clone().offset(offset as usize);
}
pub fn order_by(&mut self, ordering: Option<Vec<PyColumnOrdering>>) -> PyResult<()> {
let ordering =
ordering.map(|ordering| ordering.into_iter().map(ColumnOrdering::from).collect());
self.inner = self.inner.clone().order_by(ordering);
Ok(())
}
pub fn fast_search(&mut self) {
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 with_row_id(&mut self) {
self.inner = self.inner.clone().with_row_id();
}
pub fn column(&mut self, column: String) {
self.inner = self.inner.clone().column(&column);
}
pub fn distance_type(&mut self, distance_type: String) -> PyResult<()> {
let distance_type = parse_distance_type(distance_type)?;
self.inner = self.inner.clone().distance_type(distance_type);
Ok(())
}
pub fn postfilter(&mut self) {
self.inner = self.inner.clone().postfilter();
}
pub fn refine_factor(&mut self, refine_factor: u32) {
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 minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
self.inner = self
.inner
.clone()
.minimum_nprobes(minimum_nprobes as usize)
.infer_error()?;
Ok(())
}
pub fn maximum_nprobes(&mut self, maximum_nprobes: u32) -> PyResult<()> {
let maximum_nprobes = if maximum_nprobes == 0 {
None
} else {
Some(maximum_nprobes as usize)
};
self.inner = self
.inner
.clone()
.maximum_nprobes(maximum_nprobes)
.infer_error()?;
Ok(())
}
#[pyo3(signature = (lower_bound=None, upper_bound=None))]
pub fn distance_range(&mut self, lower_bound: Option<f32>, upper_bound: Option<f32>) {
self.inner = self.inner.clone().distance_range(lower_bound, upper_bound);
}
pub fn ef(&mut self, ef: u32) {
self.inner = self.inner.clone().ef(ef as usize);
}
pub fn bypass_vector_index(&mut self) {
self.inner = self.inner.clone().bypass_vector_index()
}
#[pyo3(signature = ())]
pub fn output_schema(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let schema = inner.output_schema().await.infer_error()?;
Python::attach(|py| schema.to_pyarrow(py).map(|obj| obj.unbind()))
})
}
#[pyo3(signature = (max_batch_length=None, timeout=None))]
pub fn execute(
self_: PyRef<'_, Self>,
max_batch_length: Option<u32>,
timeout: Option<Duration>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let mut opts = QueryExecutionOptions::default();
if let Some(max_batch_length) = max_batch_length {
opts.max_batch_length = max_batch_length;
}
if let Some(timeout) = timeout {
opts.timeout = Some(timeout);
}
let inner_stream = inner.execute_with_options(opts).await.infer_error()?;
Ok(RecordBatchStream::new(inner_stream))
})
}
pub fn explain_plan(self_: PyRef<'_, Self>, verbose: bool) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.explain_plan(verbose)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
#[pyo3(signature = (distributed_metrics=None))]
pub fn analyze_plan(
self_: PyRef<'_, Self>,
distributed_metrics: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let options = analyze_plan_options(distributed_metrics.as_deref())?;
future_into_py(self_.py(), async move {
inner
.analyze_plan_with_options(options)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
})
}
pub fn nearest_to_text(&mut self, query: Bound<'_, PyDict>) -> PyResult<HybridQuery> {
let base_query = self.inner.clone().into_plain();
let fts_query = Query::new(base_query).nearest_to_text(query)?;
Ok(HybridQuery {
inner_vec: self.clone(),
inner_fts: fts_query,
})
}
pub fn to_query_request(&self) -> PyQueryRequest {
PyQueryRequest::from(AnyQuery::VectorQuery(self.inner.clone().into_request()))
}
}
#[pyclass]
pub struct HybridQuery {
inner_vec: VectorQuery,
inner_fts: FTSQuery,
}
#[pymethods]
impl HybridQuery {
pub fn r#where(&mut self, predicate: String) {
self.inner_vec.r#where(predicate.clone());
self.inner_fts.r#where(predicate);
}
pub fn where_expr(&mut self, expr: PyExpr) {
self.inner_vec.where_expr(expr.clone());
self.inner_fts.where_expr(expr);
}
pub fn select(&mut self, columns: Vec<(String, String)>) {
self.inner_vec.select(columns.clone());
self.inner_fts.select(columns);
}
pub fn select_expr(&mut self, columns: Vec<(String, PyExpr)>) {
self.inner_vec.select_expr(columns.clone());
self.inner_fts.select_expr(columns);
}
pub fn select_columns(&mut self, columns: Vec<String>) {
self.inner_vec.select_columns(columns.clone());
self.inner_fts.select_columns(columns);
}
pub fn limit(&mut self, limit: u32) {
self.inner_vec.limit(limit);
self.inner_fts.limit(limit);
}
pub fn offset(&mut self, offset: u32) {
self.inner_vec.offset(offset);
self.inner_fts.offset(offset);
}
pub fn order_by(&mut self, ordering: Option<Vec<PyColumnOrdering>>) -> PyResult<()> {
self.inner_vec.order_by(ordering.clone())?;
self.inner_fts.order_by(ordering)?;
Ok(())
}
pub fn fast_search(&mut self) {
self.inner_vec.fast_search();
self.inner_fts.fast_search();
}
pub fn with_row_id(&mut self) {
self.inner_fts.with_row_id();
self.inner_vec.with_row_id();
}
pub fn postfilter(&mut self) {
self.inner_vec.postfilter();
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 add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> {
self.inner_vec.add_query_vector(vector)
}
pub fn column(&mut self, column: String) {
self.inner_vec.column(column);
}
pub fn distance_type(&mut self, distance_type: String) -> PyResult<()> {
self.inner_vec.distance_type(distance_type)
}
pub fn refine_factor(&mut self, refine_factor: u32) {
self.inner_vec.refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner_vec.nprobes(nprobe);
}
pub fn ef(&mut self, ef: u32) {
self.inner_vec.ef(ef);
}
pub fn bypass_vector_index(&mut self) {
self.inner_vec.bypass_vector_index();
}
#[pyo3(signature = (lower_bound=None, upper_bound=None))]
pub fn distance_range(&mut self, lower_bound: Option<f32>, upper_bound: Option<f32>) {
self.inner_vec.distance_range(lower_bound, upper_bound);
}
pub fn to_vector_query(&mut self) -> PyResult<VectorQuery> {
Ok(VectorQuery {
inner: self.inner_vec.inner.clone(),
})
}
pub fn to_fts_query(&mut self) -> PyResult<FTSQuery> {
Ok(FTSQuery {
inner: self.inner_fts.inner.clone(),
fts_query: self.inner_fts.fts_query.clone(),
})
}
pub fn get_limit(&mut self) -> Option<u32> {
self.inner_fts
.inner
.current_request()
.limit
.map(|i| i as u32)
}
pub fn get_with_row_id(&mut self) -> bool {
self.inner_fts.inner.current_request().with_row_id
}
pub fn to_query_request(&self) -> PyQueryRequest {
let mut req = self.inner_fts.to_query_request();
let vec_req = self.inner_vec.to_query_request();
req.query_vector = vec_req.query_vector;
req.column = vec_req.column;
req.distance_type = vec_req.distance_type;
req.ef = vec_req.ef;
req.refine_factor = vec_req.refine_factor;
req.lower_bound = vec_req.lower_bound;
req.upper_bound = vec_req.upper_bound;
req
}
}
/// Convert a Python FTS query to JSON string
#[pyfunction]
pub fn fts_query_to_json(query_obj: &Bound<'_, PyAny>) -> PyResult<String> {
let wrapped: PyLanceDB<FtsQuery> = query_obj.extract()?;
lancedb::table::datafusion::udtf::fts::to_json(&wrapped.0).map_err(|e| {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
"Failed to serialize FTS query to JSON: {}",
e
))
})
}
Reference in New Issue View Git Blame Copy Permalink