mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-12 08:12:28 +00:00
Jack Ye
e639b1b650
## Summary
Add SQL execution to remote LanceDB connections. On the standard
synchronous connection, `execute_query` waits for the initial result
stream and returns its Arrow reader. `execute_query_async` is called
without Python `await` and immediately returns a query handle for status
inspection, streaming, or cancellation. Local databases report that SQL
is not supported.
The transport and query lifecycle live in Rust. Python exposes
native-backed synchronous and asynchronous connection methods and query
wrappers; it does not use PyArrow's Flight client.
## User experience
The standard synchronous connection supports both direct reads and
background query execution:
```python
db = lancedb.connect(
"db://analytics",
api_key="ldb_...",
sql_host_override="grpc+tls://sql.example.com:10026",
)
# Direct execution waits only until the initial result stream is available.
# Later batches continue streaming as the query progresses.
reader = db.execute_query(
"SELECT * FROM events",
default_namespace_path=["production"],
)
for batch in reader:
print(batch.num_rows)
# Background execution returns a query handle immediately. Despite the
# `_async` suffix, no Python `await` is needed on a synchronous connection.
query = db.execute_query_async("SELECT * FROM events")
print(query.id)
description = db.describe_query(query.id)
print(description.status)
print(description.progress)
print(description.expires_at)
# Start reading as soon as the service advertises partial results. The reader
# continues polling and yields newly available record batches until the query
# and all result endpoints are complete.
reader = query.reader()
for batch in reader:
print(batch.num_rows)
# Or cancel a different still-running query. Its status becomes "cancelling"
# while the server is still working, then "cancelled" once confirmed.
cancelled_query = db.execute_query_async("SELECT * FROM large_events")
cancelled_query.cancel()
```
The less commonly used asynchronous connection exposes the same
operations as coroutines:
```python
async_db = await lancedb.connect_async(
"db://analytics",
api_key="ldb_...",
sql_host_override="grpc+tls://sql.example.com:10026",
)
query = await async_db.execute_query_async("SELECT * FROM events")
async for batch in await query.reader():
print(batch.num_rows)
```
The UUIDv7 query id is scoped to the connection that submitted it. The
connection retains lightweight shared query state used by
`query.describe()` and `db.describe_query(query.id)`; the id does not
encode SQL or a Flight continuation token and is not a cross-connection
resume token. Abandoned state has bounded retention, and terminal state
remains available briefly.
Unqualified table names use the connected database and the `public`
namespace by default. `default_namespace_path` accepts a list such as
`["production", "events"]`. SQL can still use qualified names to
reference other databases and namespaces available to the deployment.
## Design
- Uses Arrow Flight `PollFlightInfo` for submission and long polling,
`DoGet` for results, and `CancelFlightInfo` for cancellation. Each
`PollInfo.info` is treated as the cumulative set of currently available
endpoints, so advertised tickets are consumed once and batches can be
delivered before execution is complete.
- Serializes result completion and cancellation into one lifecycle. A
server-accepted request reports `cancelling` and wakes blocked
status/result work; a later retry can confirm `cancelled`. Result
retrieval is rejected after cancellation is accepted, while cancellation
after a result was already delivered is a no-op.
- Assigns a time-ordered UUIDv7 connection-scoped query id and retains
only shared evolving lifecycle state, keeping SQL, Flight continuation
tokens, and Arrow result data out of public ids and the registry.
- Leaves admission control to the server while honoring server
expiration and a local fallback retention window for abandoned entries.
- Retains terminal ids for five minutes so they remain available for
connection-level description.
- Keeps one lazily initialized SQL client on each remote database
connection and attaches fresh authentication, routing, namespace, and
request metadata to every operation.
- Applies the configured overall timeout to each execution, description,
reader, and cancellation operation. A result reader carries one absolute
deadline from `reader()` through the end of streaming; connect and read
timeouts continue to bound their individual phases.
- Returns a bounded, backpressured, single-consumer Arrow stream rather
than collecting the full result in memory. Dropping the reader stops
downloading but does not implicitly cancel the server query.
- Preserves typed schemas for empty result sets through the stream
schema.
- Accepts Flight result messages up to 1 GiB so a valid row containing a
large blob, string, or vector is not rejected by tonic's 4 MiB default
receive limit.
- Supports the Python client first while keeping the authoritative
implementation in the Rust core.
100 lines
3.5 KiB
Rust
100 lines
3.5 KiB
Rust
// SPDX-License-Identifier: Apache-2.0
|
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
use arrow::RecordBatchStream;
|
|
use connection::{Connection, connect, connect_namespace, connect_namespace_client};
|
|
use env_logger::Env;
|
|
use expr::{PyExpr, expr_col, expr_func, expr_lit};
|
|
use index::IndexConfig;
|
|
use permutation::{PyAsyncPermutationBuilder, PyPermutationReader};
|
|
use pyo3::{
|
|
Bound, PyResult, Python, pymodule,
|
|
types::{PyModule, PyModuleMethods},
|
|
wrap_pyfunction,
|
|
};
|
|
use query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
|
use session::Session;
|
|
use table::{
|
|
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
|
LsmWriteSpec, MergeResult, PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult,
|
|
Table, UpdateFieldMetadataResult, UpdateResult,
|
|
};
|
|
|
|
pub mod arrow;
|
|
pub mod connection;
|
|
pub mod error;
|
|
pub mod expr;
|
|
pub mod header;
|
|
pub mod index;
|
|
pub mod job;
|
|
pub mod namespace;
|
|
pub mod oauth;
|
|
pub mod otel;
|
|
pub mod permutation;
|
|
pub mod query;
|
|
pub mod runtime;
|
|
pub mod session;
|
|
pub mod sql;
|
|
pub mod table;
|
|
pub mod util;
|
|
|
|
#[pymodule]
|
|
pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|
let env = Env::new()
|
|
.filter_or("LANCEDB_LOG", "warn")
|
|
.write_style("LANCEDB_LOG_STYLE");
|
|
env_logger::init_from_env(env);
|
|
m.add_class::<Connection>()?;
|
|
m.add_class::<Session>()?;
|
|
m.add_class::<Table>()?;
|
|
m.add_class::<crate::job::Job>()?;
|
|
m.add_class::<crate::job::JobInfo>()?;
|
|
m.add_class::<crate::job::JobDescription>()?;
|
|
m.add_class::<crate::job::JobFailureInfo>()?;
|
|
m.add_class::<crate::sql::Query>()?;
|
|
m.add_class::<crate::sql::QueryDescription>()?;
|
|
m.add_class::<PyBlobFile>()?;
|
|
m.add_class::<IndexConfig>()?;
|
|
m.add_class::<Query>()?;
|
|
m.add_class::<FTSQuery>()?;
|
|
m.add_class::<HybridQuery>()?;
|
|
m.add_class::<VectorQuery>()?;
|
|
m.add_class::<RecordBatchStream>()?;
|
|
m.add_class::<AddColumnsResult>()?;
|
|
m.add_class::<RefreshColumnResult>()?;
|
|
m.add_class::<RefreshMaterializedViewResult>()?;
|
|
m.add_class::<AlterColumnsResult>()?;
|
|
m.add_class::<UpdateFieldMetadataResult>()?;
|
|
m.add_class::<AddResult>()?;
|
|
m.add_class::<MergeResult>()?;
|
|
m.add_class::<LsmWriteSpec>()?;
|
|
m.add_class::<DeleteResult>()?;
|
|
m.add_class::<DropColumnsResult>()?;
|
|
m.add_class::<UpdateResult>()?;
|
|
m.add_class::<FtsToken>()?;
|
|
m.add_class::<PyAsyncPermutationBuilder>()?;
|
|
m.add_class::<PyPermutationReader>()?;
|
|
m.add_class::<PyExpr>()?;
|
|
// OpenTelemetry metrics bridge
|
|
m.add_class::<otel::PyMetricPoint>()?;
|
|
m.add_class::<otel::PyMetricDescription>()?;
|
|
m.add_function(wrap_pyfunction!(
|
|
otel::register_lancedb_metrics_recorder,
|
|
m
|
|
)?)?;
|
|
m.add_function(wrap_pyfunction!(otel::lancedb_metrics_catalog, m)?)?;
|
|
m.add_function(wrap_pyfunction!(otel::snapshot_lancedb_metrics, m)?)?;
|
|
m.add_function(wrap_pyfunction!(connect, m)?)?;
|
|
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
|
m.add_function(wrap_pyfunction!(connect_namespace_client, m)?)?;
|
|
m.add_function(wrap_pyfunction!(table::tokenize, m)?)?;
|
|
m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?;
|
|
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
|
m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?;
|
|
m.add_function(wrap_pyfunction!(expr_col, m)?)?;
|
|
m.add_function(wrap_pyfunction!(expr_lit, m)?)?;
|
|
m.add_function(wrap_pyfunction!(expr_func, m)?)?;
|
|
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
|
|
Ok(())
|
|
}
|