From 4487b1108c65d9efcd2f69d1a73642d097b5c971 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 8 Oct 2026 18:23:46 -0700 Subject: [PATCH] feat(sql): bind query parameters over Flight DoExchange (#4345) `execute_query` / `execute_query_async` can now bind values to `$1` / `$name` placeholders instead of splicing them into the SQL text. The values travel as Arrow, so a float32 keeps its exact bits and type and a vector stays a compact `FixedSizeList` (3 KB for 768 dims rather than ~9 KB of decimal literals parsed back as Float64). ```rust db.execute_query_async("SELECT id FROM docs ORDER BY distance(vector, $vector) LIMIT $k") .parameters(batch) // one-row RecordBatch .execute() .await?; ``` ```python db.execute_query(sql, parameters={"vector": np.float32_vector, "k": 10}) ``` - A statement with parameters is sent on one `DoExchange`: the first message carries the `CommandStatementQuery` descriptor and the parameter schema, the second the one-row batch, and the result rows stream back on the same call. Statements without parameters keep `PollFlightInfo` + `DoGet` unchanged. - Column `i` binds `$`; a column whose name is not a number also binds `$`. Python accepts a list (positional), a dict (named), or a one-row `RecordBatch` / `Table`; 1-D numpy / pyarrow arrays become a `FixedSizeList` of their element type. - A parameterized query runs on its call: `cancel()`, dropping the handle before reading, or dropping the reader ends the call and stops the server-side statement. It is still registered for `describe_query`. - `Database::execute_query_async` now takes an `ExecuteQueryRequest` options struct. Server side: lancedb/sophon#8256. --- docs/src/python/python.md | 11 + python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/db.py | 37 +- python/python/lancedb/namespace.py | 6 +- python/python/lancedb/remote/db.py | 6 +- python/python/lancedb/sql.py | 80 ++++- python/python/tests/test_sql.py | 113 +++++- python/src/connection.rs | 11 +- rust/lancedb/src/connection.rs | 71 +++- rust/lancedb/src/database.rs | 30 +- rust/lancedb/src/remote/db.rs | 12 +- rust/lancedb/src/remote/sql.rs | 94 ++++- rust/lancedb/src/remote/sql/exchange.rs | 435 ++++++++++++++++++++++++ rust/lancedb/src/remote/sql_test.rs | 414 ++++++++++++++++++++-- rust/lancedb/src/sql.rs | 13 +- 15 files changed, 1267 insertions(+), 67 deletions(-) create mode 100644 rust/lancedb/src/remote/sql/exchange.rs diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 863edbe4a..3776283dc 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -86,6 +86,15 @@ print(query.describe().status) for batch in query.reader(): print(batch.num_rows) +# Values bind to $1 / $name placeholders and travel as Arrow, so a float32 +# stays a float32 and a vector stays a compact fixed-size list: +import numpy as np + +reader = db.execute_query( + "SELECT id FROM docs ORDER BY distance(vector, $vector) LIMIT $k", + parameters={"vector": np.random.rand(768).astype(np.float32), "k": 10}, +) + # The async connection exposes the same lifecycle without blocking: # async_db = await lancedb.connect_async( # "db://analytics", @@ -201,6 +210,8 @@ listing a storage directory. ::: lancedb.sql.QueryDescription +::: lancedb.sql.QueryParameters + ## Materialized Views (Synchronous) ::: lancedb.materialized_view.MaterializedView diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f7477e891..e4679382f 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -217,6 +217,7 @@ class Connection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[pa.RecordBatch] = None, ) -> SqlQuery: ... async def describe_query(self, query_id: UUID) -> QueryDescription: ... async def create_table( diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index f890c5a14..aff86e9bf 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -53,7 +53,7 @@ from .functions import FunctionVersion, UdfDefinition from .job import AsyncJob, Job, _typed_job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery -from .sql import QueryDescription +from .sql import QueryDescription, QueryParameters, to_parameter_batch from .materialized_view import ( AsyncMaterializedView, MaterializedView, @@ -1178,6 +1178,7 @@ class DBConnection(EnforceOverrides): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> pa.RecordBatchReader: """Execute SQL and return a blocking Arrow reader. @@ -1188,6 +1189,7 @@ class DBConnection(EnforceOverrides): return self.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ).reader() def execute_query_async( @@ -1195,9 +1197,32 @@ class DBConnection(EnforceOverrides): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> SqlQuery: """Start executing SQL and return its query handle. + ``parameters`` binds values to the statement's ``$1`` / ``$name`` + placeholders. They travel as Arrow rather than as SQL text, so a + float keeps its exact value and type and a vector stays a compact + fixed-size list. A list binds by position, a dict binds by name, and + a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds both + ways. A 1-D numpy or pyarrow array becomes a fixed-size list of its + element type, the shape of a vector column: + + .. code-block:: python + + db.execute_query( + "SELECT id FROM docs ORDER BY distance(vector, $1) LIMIT $2", + parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10], + ) + + Parameters are supported in queries, not in DDL or DML. A + parameterized query runs on the call that returns its rows: + cancelling it, or dropping its handle before reading, stops it on the + server. Its status follows that stream as well: ``describe()`` + reports ``running`` until the reader has received every row, so open + ``reader()`` directly rather than waiting for ``finished``. + Local connections do not support SQL. """ raise NotImplementedError("SQL is not supported for this connection type") @@ -3267,6 +3292,7 @@ class AsyncConnection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncRecordBatchReader: """Execute SQL and return an asynchronous Arrow reader. @@ -3277,6 +3303,7 @@ class AsyncConnection(object): submitted = await self.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) return await submitted.reader() @@ -3285,17 +3312,21 @@ class AsyncConnection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncSqlQuery: """Start executing SQL and return its query handle. The database from ``connect_async`` is used for unqualified database - references. The namespace defaults to ``["public"]``. Local - connections raise ``NotImplementedError``. + references. The namespace defaults to ``["public"]``. ``parameters`` + binds values to ``$1`` / ``$name`` placeholders; see + [DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async]. + Local connections raise ``NotImplementedError``. """ return AsyncSqlQuery( await self._inner.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=to_parameter_batch(parameters), ) ) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a3571f64b..0fe77bfef 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -53,7 +53,7 @@ from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job from lancedb.listing import AsyncListing, Listing from lancedb.sql import AsyncQuery as AsyncSqlQuery -from lancedb.sql import QueryDescription +from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( LanceNamespace, connect as namespace_connect, @@ -1584,11 +1584,13 @@ class AsyncLanceNamespaceDBConnection: query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncRecordBatchReader: """Execute SQL when supported by the underlying connection.""" return await self._inner.execute_query( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) async def execute_query_async( @@ -1596,6 +1598,7 @@ class AsyncLanceNamespaceDBConnection: query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncSqlQuery: """Start executing SQL when supported by the underlying connection. @@ -1604,6 +1607,7 @@ class AsyncLanceNamespaceDBConnection: return await self._inner.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) async def describe_query(self, query_id: UUID) -> QueryDescription: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 2688c3efc..a0a6ef6c1 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -41,7 +41,7 @@ from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job from ..sql import Query as SqlQuery -from ..sql import QueryDescription +from ..sql import QueryDescription, QueryParameters from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg from ..secrets import EnvVarSecret, SecretInfo from ..view import ViewDescription @@ -1090,18 +1090,22 @@ class RemoteDBConnection(DBConnection): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> SqlQuery: """Start executing SQL through this remote connection. Unqualified tables use this connection's database and the ``["public"]`` namespace by default. Fully qualified table names may reference other databases available to the same deployment. + ``parameters`` binds values to ``$1`` / ``$name`` placeholders; see + [DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async]. """ return SqlQuery( LOOP.run( self._conn.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) ) ) diff --git a/python/python/lancedb/sql.py b/python/python/lancedb/sql.py index 41bbb5328..05b82b0d3 100644 --- a/python/python/lancedb/sql.py +++ b/python/python/lancedb/sql.py @@ -3,8 +3,10 @@ """Handles to SQL queries running on a remote database.""" +from typing import Any, Mapping, Optional, Sequence, Union from uuid import UUID +import numpy as np import pyarrow as pa from lancedb.background_loop import LOOP @@ -14,6 +16,82 @@ from .arrow import AsyncRecordBatchReader QueryDescription = _lancedb.QueryDescription +QueryParameters = Union[pa.RecordBatch, pa.Table, Sequence[Any], Mapping[str, Any]] +"""Values for a statement's ``$1`` / ``$name`` placeholders. + +A sequence binds by position (``$1``, ``$2``, ...), a mapping binds by name +(``$name``), and a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds +both ways: column ``i`` is ``$`` and, when its name is not a number, +also ``$``. +""" + + +def _parameter_array(value: Any) -> pa.Array: + """One parameter as a one-element array of its own type.""" + if isinstance(value, pa.Scalar): + return pa.array([value.as_py()], type=value.type) + # A 1-D array is a vector: keep its element type and make it the + # fixed-size list a vector column stores, rather than a list of Python + # floats that would come back as float64. + if isinstance(value, pa.ChunkedArray): + value = value.combine_chunks() + if isinstance(value, np.ndarray): + if value.ndim == 0: + return pa.array(value.reshape(1)) + if value.ndim != 1: + raise ValueError( + "a numpy query parameter must be a scalar or a 1-D vector, " + f"got {value.ndim} dimensions" + ) + value = pa.array(value) + if isinstance(value, pa.Array): + return pa.FixedSizeListArray.from_arrays(value, len(value)) + return pa.array([value]) + + +def to_parameter_batch( + parameters: Optional[QueryParameters], +) -> Optional[pa.RecordBatch]: + """Turn ``parameters`` into the one-row batch a parameterized query sends. + + Scalars keep the type pyarrow infers for them (numpy scalars keep their + dtype), a ``pyarrow.Scalar`` keeps its own type, and a 1-D numpy or + pyarrow array becomes a ``FixedSizeList`` of its element type, which is + what a vector column stores. + """ + if parameters is None: + return None + batch: pa.RecordBatch + if isinstance(parameters, pa.RecordBatch): + batch = parameters + elif isinstance(parameters, pa.Table): + batches = parameters.combine_chunks().to_batches() + batch = ( + batches[0] + if len(batches) == 1 + else pa.RecordBatch.from_pylist([], schema=parameters.schema) + ) + elif isinstance(parameters, Mapping): + batch = pa.RecordBatch.from_arrays( + [_parameter_array(value) for value in parameters.values()], + names=[str(name) for name in parameters.keys()], + ) + elif isinstance(parameters, Sequence) and not isinstance(parameters, (str, bytes)): + batch = pa.RecordBatch.from_arrays( + [_parameter_array(value) for value in parameters], + names=[str(position + 1) for position in range(len(parameters))], + ) + else: + raise TypeError( + "query parameters must be a sequence, a mapping, or a one-row " + f"pyarrow RecordBatch or Table, not {type(parameters).__name__}" + ) + if batch.num_rows != 1: + raise ValueError( + f"query parameters must be exactly one row, got {batch.num_rows}" + ) + return batch + class AsyncQuery: """A handle to a submitted SQL query on an asynchronous connection.""" @@ -85,4 +163,4 @@ class Query: LOOP.run(self._inner.cancel()) -__all__ = ["AsyncQuery", "Query", "QueryDescription"] +__all__ = ["AsyncQuery", "Query", "QueryDescription", "QueryParameters"] diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index eeadb3a33..e6f5fd3c7 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import threading from uuid import UUID +import numpy as np import pytest import pyarrow as pa @@ -11,7 +13,7 @@ from lancedb import _lancedb from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection from lancedb.remote.db import RemoteDBConnection -from lancedb.sql import AsyncQuery, Query +from lancedb.sql import AsyncQuery, Query, to_parameter_batch NIL_QUERY_ID = UUID(int=0) @@ -24,12 +26,20 @@ class FakeNativeQuery: class FakeNativeConnection: - async def execute_query_async(self, query, *, default_namespace_path=None): + def __init__(self): + self.parameters = [] + + async def execute_query_async( + self, query, *, default_namespace_path=None, parameters=None + ): + self.parameters.append(parameters) return FakeNativeQuery() class FakeAsyncConnection: - async def execute_query_async(self, query, *, default_namespace_path=None): + async def execute_query_async( + self, query, *, default_namespace_path=None, parameters=None + ): return AsyncQuery(FakeNativeQuery()) @@ -160,3 +170,100 @@ def test_execute_query_async_rejects_invalid_namespace_components( remote_connection().execute_query_async( "SELECT 1", default_namespace_path=default_namespace_path ) + + +def test_positional_parameters_keep_their_types(): + vector = np.array([0.1, 0.2, 0.3], dtype=np.float32) + batch = to_parameter_batch( + [42, 1.5, "a", np.float32(0.25), vector, pa.scalar(7, pa.int16())] + ) + assert batch.num_rows == 1 + assert batch.schema.names == ["1", "2", "3", "4", "5", "6"] + assert batch.schema.types == [ + pa.int64(), + pa.float64(), + pa.string(), + pa.float32(), + pa.list_(pa.float32(), 3), + pa.int16(), + ] + sent = batch.column(4).flatten().to_numpy() + assert sent.dtype == np.float32 + assert sent.tobytes() == vector.tobytes() + + +def test_named_parameters_and_batches(): + half = pa.array(np.array([1.0, 2.0], dtype=np.float16)) + batch = to_parameter_batch({"vector": half, "k": 3}) + assert batch.schema.names == ["vector", "k"] + assert batch.schema.field("vector").type == pa.list_(pa.float16(), 2) + + row = pa.record_batch({"id": [1]}) + assert to_parameter_batch(row) is row + assert to_parameter_batch(pa.table({"id": [1]})).equals(row) + assert to_parameter_batch(None) is None + + +@pytest.mark.parametrize( + "parameters", + [pa.record_batch({"id": [1, 2]}), pa.table({"id": pa.array([], pa.int64())})], +) +def test_parameters_are_exactly_one_row(parameters): + with pytest.raises(ValueError, match="exactly one row"): + to_parameter_batch(parameters) + + +def test_parameters_must_be_a_collection(): + with pytest.raises(TypeError, match="query parameters"): + to_parameter_batch("SELECT 1") + + +@pytest.mark.asyncio +async def test_async_execute_query_sends_parameters_as_a_batch(): + native = FakeNativeConnection() + connection = AsyncConnection(native) + await connection.execute_query("SELECT $1", parameters=[1]) + await connection.execute_query_async("SELECT 1") + sent, unparameterized = native.parameters + assert sent.equals(pa.record_batch({"1": [1]})) + assert unparameterized is None + + +def test_parameters_reach_the_server_as_the_values_that_were_sent(): + """A round trip through the bindings against a Flight server that answers + each exchange with the parameter row it received.""" + flight = pytest.importorskip("pyarrow.flight") + + class EchoServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + self.commands = [] + self.lock = threading.Lock() + + def do_exchange(self, context, descriptor, reader, writer): + parameters = reader.read_all() + with self.lock: + self.commands.append(descriptor.command) + writer.begin(parameters.schema) + writer.write_table(parameters) + writer.close() + + vector = np.array([0.1, 1e-45, -0.0], dtype=np.float32) + with EchoServer() as server: + db = lancedb.connect( + "db://analytics", + api_key="test-key", + host_override="http://localhost:10024", + sql_host_override=f"grpc://127.0.0.1:{server.port}", + ) + echoed = db.execute_query( + "SELECT echo", parameters={"vector": vector, "k": 10} + ).read_all() + (command,) = server.commands + + assert b"CommandStatementQuery" in command + assert b"SELECT echo" in command + assert echoed.schema.field("vector").type == pa.list_(pa.float32(), 3) + received = echoed.column("vector").combine_chunks().flatten().to_numpy() + assert received.view(np.uint32).tolist() == vector.view(np.uint32).tolist() + assert echoed.column("k").to_pylist() == [10] diff --git a/python/src/connection.rs b/python/src/connection.rs index 71ba70e0e..ec0768eb2 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -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>, + parameters: Option>, ) -> PyResult> { 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 diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index b299cf4ff..c6138bcdd 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -23,8 +23,8 @@ use crate::connection::create_table::CreateTableBuilder; use crate::data::scannable::Scannable; use crate::database::listing::ListingDatabase; use crate::database::{ - CloneTableRequest, Database, DatabaseOptions, JobInfo, OpenTableRequest, PauseJobStatus, - ReadConsistency, ResumeJobStatus, TableNamesRequest, + CloneTableRequest, Database, DatabaseOptions, ExecuteQueryRequest, JobInfo, OpenTableRequest, + PauseJobStatus, ReadConsistency, ResumeJobStatus, TableNamesRequest, }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -335,16 +335,14 @@ pub struct CloneTableBuilder { /// Builder for asynchronously executing a SQL statement on a remote database. pub struct ExecuteQueryAsyncBuilder { parent: Arc, - query: String, - default_namespace_path: Vec, + request: ExecuteQueryRequest, } impl ExecuteQueryAsyncBuilder { fn new(parent: Arc, query: String) -> Self { Self { parent, - query, - default_namespace_path: vec!["public".to_string()], + request: ExecuteQueryRequest::new(query), } } @@ -357,15 +355,63 @@ impl ExecuteQueryAsyncBuilder { I: IntoIterator, S: Into, { - self.default_namespace_path = path.into_iter().map(Into::into).collect(); + self.request.default_namespace_path = path.into_iter().map(Into::into).collect(); + self + } + + /// Bind values to the statement's placeholders. + /// + /// `parameters` is a single-row batch. Column `i` binds `$`, and a + /// column whose name is not a number also binds `$`, so + /// `WHERE id = $1` and `WHERE id = $id` both work. The values travel as + /// Arrow rather than as SQL text, so a float keeps its exact bits and type + /// and a vector stays a compact `FixedSizeList`. + /// + /// Every placeholder needs a value and every parameter must be used by a + /// placeholder. Parameters are supported in queries (`SELECT`, and + /// `EXPLAIN` of one), not in DDL or DML, nor in table-function arguments. + /// Do not follow a named placeholder directly with `$`: `$a$` starts a + /// dollar-quoted string. + /// + /// A parameterized statement runs on the call that returns its rows rather + /// than detached from it: cancelling the query, or dropping its handle + /// before reading it, stops the statement on the server. Its status follows + /// that stream too -- it reports `Running` until the reader has received + /// every row -- so open the reader directly rather than waiting for + /// `Finished`. + /// + /// # Example + /// + /// ```no_run + /// # use std::sync::Arc; + /// # use arrow_array::{ArrayRef, FixedSizeListArray, Int64Array, RecordBatch}; + /// # use arrow_array::types::Float32Type; + /// # async fn query(db: &lancedb::Connection) -> lancedb::Result<()> { + /// let vector = FixedSizeListArray::from_iter_primitive::( + /// [Some([0.1_f32, 0.2, 0.3].map(Some))], + /// 3, + /// ); + /// let parameters = RecordBatch::try_from_iter([ + /// ("vector", Arc::new(vector) as ArrayRef), + /// ("k", Arc::new(Int64Array::from(vec![10])) as ArrayRef), + /// ])?; + /// let query = db + /// .execute_query_async("SELECT id FROM docs ORDER BY distance(vector, $vector) LIMIT $k") + /// .parameters(parameters) + /// .execute() + /// .await?; + /// let rows = query.reader().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn parameters(mut self, parameters: RecordBatch) -> Self { + self.request.parameters = Some(parameters); self } /// Start the statement and return its asynchronous query handle. pub async fn execute(self) -> Result { - self.parent - .execute_query_async(&self.query, &self.default_namespace_path) - .await + self.parent.execute_query_async(self.request).await } } @@ -456,8 +502,9 @@ impl Connection { /// /// The query can reference tables in other databases with SQL dot notation. /// Use [`ExecuteQueryAsyncBuilder::default_namespace_path`] to avoid qualifying - /// tables in the default namespace. Local connections return - /// [`Error::NotSupported`]. + /// tables in the default namespace, and [`ExecuteQueryAsyncBuilder::parameters`] + /// to bind Arrow values to `$1` / `$name` placeholders instead of writing them + /// into the SQL text. Local connections return [`Error::NotSupported`]. /// /// # Example /// diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index ad4e8118d..05386e0b7 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use arrow_array::RecordBatch; use lance::dataset::ReadParams; use lance_namespace::LanceNamespace; use lance_namespace::models::{ @@ -191,6 +192,32 @@ impl CloneTableRequest { } } +/// A request to start a SQL statement on a remote database. +#[derive(Clone, Debug)] +pub struct ExecuteQueryRequest { + /// The SQL text. + pub query: String, + /// The namespace unqualified table names resolve in. Empty is `public`, + /// the SQL name of the root namespace. + pub default_namespace_path: Vec, + /// Values for the statement's placeholders, as a single-row batch. + /// + /// Column `i` binds `$`, and a column whose name is not a number + /// also binds `$`. See + /// [`ExecuteQueryAsyncBuilder::parameters`](crate::connection::ExecuteQueryAsyncBuilder::parameters). + pub parameters: Option, +} + +impl ExecuteQueryRequest { + pub fn new(query: impl Into) -> Self { + Self { + query: query.into(), + default_namespace_path: vec!["public".to_string()], + parameters: None, + } + } +} + /// How long until a change is reflected from one Table instance to another /// /// Tables are always internally consistent. If a write method is called on @@ -555,8 +582,7 @@ pub trait Database: /// Start executing a SQL statement on a remote database. async fn execute_query_async( &self, - _query: &str, - _default_namespace_path: &[String], + _request: ExecuteQueryRequest, ) -> Result { Err(crate::error::Error::NotSupported { message: "SQL is not supported by this database".to_string(), diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index c44f7558e..57513c703 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -20,8 +20,8 @@ use lance_namespace::models::{ use crate::Error; use crate::database::{ - CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, JobInfo, - OpenTableRequest, ReadConsistency, TableNamesRequest, + CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, + ExecuteQueryRequest, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; use crate::function::{ @@ -1329,18 +1329,14 @@ impl Database for RemoteDatabase { }) } - async fn execute_query_async( - &self, - query: &str, - default_namespace_path: &[String], - ) -> Result { + async fn execute_query_async(&self, request: ExecuteQueryRequest) -> Result { let client = self .sql_client .as_ref() .ok_or_else(|| Error::NotSupported { message: "SQL is unavailable for this remote database client".to_string(), })?; - client.submit(query, default_namespace_path).await + client.execute(request).await } async fn describe_query(&self, query_id: uuid::Uuid) -> Result { diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index e8ddcf6ff..a58a9afad 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -27,12 +27,17 @@ use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity use uuid::Uuid; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; +use crate::database::ExecuteQueryRequest; use crate::error::{Error, Result}; use crate::job::Job; use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; +mod exchange; + +use exchange::{ExchangeQuery, ExchangeQueryHandle}; + const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(120); @@ -185,6 +190,23 @@ impl SqlClient { } } + pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { + match request.parameters { + Some(parameters) => { + self.submit_with_parameters( + &request.query, + parameters, + &request.default_namespace_path, + ) + .await + } + None => { + self.submit(&request.query, &request.default_namespace_path) + .await + } + } + } + pub(super) async fn submit( &self, query: &str, @@ -206,7 +228,8 @@ impl SqlClient { default_namespace_path.to_vec(), poll_info, )?); - self.queries.insert(query_id, query.clone()); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) }) .await @@ -248,6 +271,35 @@ impl SqlClient { }) } + /// Values cannot ride on a poll, so a statement that has them is sent on + /// its own exchange instead. + pub(super) async fn submit_with_parameters( + &self, + query: &str, + parameters: RecordBatch, + default_namespace_path: &[String], + ) -> Result { + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(default_namespace_path)?; + let query_id = Uuid::now_v7(); + let query = Arc::new( + ExchangeQuery::open( + query_id, + self.inner.clone(), + query, + parameters, + default_namespace_path, + ) + .await?, + ); + self.queries + .insert(query_id, RegisteredQuery::Exchanged(query.clone())); + Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) + }) + .await + } + pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries @@ -530,8 +582,40 @@ impl SqlClientInner { } } +/// A submitted query as the connection remembers it. +#[derive(Clone)] +enum RegisteredQuery { + Polled(Arc), + Exchanged(Arc), +} + +impl RegisteredQuery { + fn touch(&self) { + match self { + Self::Polled(query) => query.touch(), + Self::Exchanged(query) => query.touch(), + } + } + + /// Whether the registry may forget this query. It is abandoned when the + /// registry holds the only reference to it. + fn registry_expired(&self) -> bool { + match self { + Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + } + } + + async fn describe(&self) -> Result { + match self { + Self::Polled(query) => query.describe().await, + Self::Exchanged(query) => query.describe(), + } + } +} + struct QueryRegistry { - queries: StdMutex>>, + queries: StdMutex>, } impl QueryRegistry { @@ -541,12 +625,12 @@ impl QueryRegistry { } } - fn insert(&self, id: Uuid, query: Arc) { + fn insert(&self, id: Uuid, query: RegisteredQuery) { self.remove_expired(); self.queries.lock().unwrap().insert(id, query); } - fn get(&self, id: Uuid) -> Option> { + fn get(&self, id: Uuid) -> Option { self.remove_expired(); let query = self.queries.lock().unwrap().get(&id).cloned(); if let Some(query) = &query { @@ -559,7 +643,7 @@ impl QueryRegistry { self.queries .lock() .unwrap() - .retain(|_, query| !query.registry_expired(Arc::strong_count(query) == 1)); + .retain(|_, query| !query.registry_expired()); } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/exchange.rs new file mode 100644 index 000000000..523aaf5d7 --- /dev/null +++ b/rust/lancedb/src/remote/sql/exchange.rs @@ -0,0 +1,435 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! SQL statements sent with parameters. +//! +//! A statement without parameters is polled: `PollFlightInfo` starts it +//! detached from the connection and `DoGet` reads its endpoints. Neither call +//! carries Arrow data, so a statement's values would have to be spelled into +//! its text. A statement with parameters is sent on one `DoExchange` instead: +//! the command and the parameter schema, then one row, and the result rows +//! come back on the same call. The call is the statement's lifetime, so +//! cancelling it -- or dropping the handle before reading -- stops the +//! statement on the server. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex as StdMutex, OnceLock}; +use std::time::Instant; + +use arrow_array::RecordBatch; +use arrow_flight::FlightDescriptor; +use arrow_flight::encode::FlightDataEncoderBuilder; +use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_schema::SchemaRef; +use prost::Message; +use tokio::sync::{Notify, mpsc}; +use uuid::Uuid; + +use super::{ + ABANDONED_QUERY_RETENTION, DEFAULT_READ_TIMEOUT, ResultEndpointStream, ResultStartGuard, + SqlClientInner, TERMINAL_QUERY_RETENTION, resolve_timeout, sql_error, with_overall_timeout, +}; +use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; +use crate::error::{Error, Result}; +use crate::sql::{QueryDescription, QueryHandle, QueryStatus}; + +#[derive(Clone, Debug, PartialEq, Eq)] +enum Lifecycle { + Running, + /// Every row has been received. + Finished, + Cancelled, + Failed(String), +} + +/// A failed statement, keeping the message `describe` reports it with -- not +/// `Runtime error: …` again when `describe` wraps it in a runtime error. +fn failure(error: &Error) -> Lifecycle { + Lifecycle::Failed(match error { + Error::Runtime { message } => message.clone(), + other => other.to_string(), + }) +} + +/// A statement running on its own `DoExchange` call. +pub(super) struct ExchangeQuery { + id: Uuid, + client: Arc, + /// The call's rows, until a reader takes them or a cancellation drops them. + response: StdMutex>, + lifecycle: StdMutex, + cancelled: Notify, + terminal_at: OnceLock, + last_accessed: StdMutex, +} + +struct PreparedExchange { + schema: SchemaRef, + first: Option, + rest: Option, +} + +/// The call's rows while a reader waits for their first batch. +/// +/// A reader whose future is dropped mid-wait -- its caller gave up, or its +/// task was cancelled -- puts them back, so the next reader finds the query as +/// this one did. A query cancelled meanwhile has its rows dropped instead, +/// which ends the call. +struct Borrowed<'a> { + query: &'a ExchangeQuery, + rows: Option, +} + +impl Borrowed<'_> { + fn get(&mut self) -> &mut ResultEndpointStream { + self.rows.as_mut().expect("rows are borrowed until kept") + } + + fn keep(mut self) -> ResultEndpointStream { + self.rows.take().expect("rows are borrowed until kept") + } +} + +impl Drop for Borrowed<'_> { + fn drop(&mut self) { + let Some(rows) = self.rows.take() else { + return; + }; + // Checked under the response lock: `cancel` settles first and takes + // the response after, so rows put back here are never missed by it. + let mut response = self.query.response.lock().unwrap(); + if *self.query.lifecycle.lock().unwrap() == Lifecycle::Running { + *response = Some(rows); + } + } +} + +impl ExchangeQuery { + /// Send the statement and its parameters, returning once the server has + /// planned it: a statement it refuses is refused here, as a polled one is + /// on its first poll. + pub(super) async fn open( + id: Uuid, + client: Arc, + query: &str, + parameters: RecordBatch, + default_namespace_path: &[String], + ) -> Result { + if parameters.num_rows() != 1 { + return Err(Error::InvalidInput { + message: format!( + "SQL query parameters must be exactly one row, got {}", + parameters.num_rows() + ), + }); + } + let request_id = Uuid::new_v4().to_string(); + let read_timeout = resolve_timeout( + client.client_config.timeout_config.read_timeout, + "LANCE_CLIENT_READ_TIMEOUT", + Some(DEFAULT_READ_TIMEOUT), + )? + .unwrap(); + let descriptor = FlightDescriptor::new_cmd( + CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(), + ); + // The descriptor rides on the schema message, and the row follows it. + let request = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(descriptor)) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])); + let mut flight = client + .client_with_headers(default_namespace_path, &request_id) + .await?; + let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) + .await + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|err| sql_error(&request_id, err))?; + Ok(Self { + id, + client, + response: StdMutex::new(Some(ResultEndpointStream { + stream, + request_id, + read_timeout, + })), + lifecycle: StdMutex::new(Lifecycle::Running), + cancelled: Notify::new(), + terminal_at: OnceLock::new(), + last_accessed: StdMutex::new(Instant::now()), + }) + } + + pub(super) fn touch(&self) { + *self.last_accessed.lock().unwrap() = Instant::now(); + } + + pub(super) fn registry_expired(&self, abandoned: bool) -> bool { + if let Some(finished) = self.terminal_at.get() { + return finished.elapsed() >= TERMINAL_QUERY_RETENTION; + } + abandoned && self.last_accessed.lock().unwrap().elapsed() >= ABANDONED_QUERY_RETENTION + } + + pub(super) fn describe(&self) -> Result { + self.touch(); + let status = match &*self.lifecycle.lock().unwrap() { + Lifecycle::Running => QueryStatus::Running, + Lifecycle::Finished => QueryStatus::Finished, + Lifecycle::Cancelled => QueryStatus::Cancelled, + Lifecycle::Failed(message) => { + return Err(Error::Runtime { + message: message.clone(), + }); + } + }; + Ok(QueryDescription { + id: self.id, + status, + progress: (status == QueryStatus::Finished).then_some(1.0), + expires_at: None, + }) + } + + /// Stop the statement. Idempotent, and a no-op once it has ended. + pub(super) fn cancel(&self) { + self.touch(); + let settled = self.settle(Lifecycle::Cancelled); + // Unread rows are dropped here, which cancels the call; a query that + // already ended has none. A reader's stream is dropped by its task + // when it sees the cancellation. + self.response.lock().unwrap().take(); + if settled { + self.cancelled.notify_waiters(); + } + } + + /// Move a running statement to `state`. `true` if this call did. + fn settle(&self, state: Lifecycle) -> bool { + let mut lifecycle = self.lifecycle.lock().unwrap(); + if *lifecycle != Lifecycle::Running { + return false; + } + *lifecycle = state; + drop(lifecycle); + let _ = self.terminal_at.set(Instant::now()); + true + } + + fn is_cancelled(&self) -> bool { + *self.lifecycle.lock().unwrap() == Lifecycle::Cancelled + } + + fn cancelled_error(&self) -> Error { + Error::JobCancelled { + job_id: Some(self.id.to_string()), + } + } + + async fn wait_for_cancellation(&self) { + loop { + let cancelled = self.cancelled.notified(); + if self.is_cancelled() { + return; + } + cancelled.await; + } + } + + /// Take the rows and wait for the first batch, which carries the schema. + async fn prepare(&self) -> Result { + let Some(rows) = self.response.lock().unwrap().take() else { + // Taken by an earlier attempt that failed, or dropped by a cancel. + return Err(match &*self.lifecycle.lock().unwrap() { + Lifecycle::Failed(message) => Error::Runtime { + message: message.clone(), + }, + _ => self.cancelled_error(), + }); + }; + let mut rows = Borrowed { + query: self, + rows: Some(rows), + }; + let first = tokio::select! { + biased; + _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), + result = rows.get().next_batch() => result, + }; + // Read or failed, the rows are this attempt's now: a failed read is + // not worth putting back. + let rows = rows.keep(); + let first = first?; + let schema = first + .as_ref() + .map(RecordBatch::schema) + .or_else(|| rows.stream.schema().cloned()) + .ok_or_else(|| Error::Runtime { + message: "SQL result stream did not include a schema".to_string(), + })?; + if first.is_none() { + self.settle(Lifecycle::Finished); + } + Ok(PreparedExchange { + schema, + rest: first.is_some().then_some(rows), + first, + }) + } + + async fn run( + self: Arc, + mut prepared: PreparedExchange, + sender: mpsc::Sender>, + ) -> Result<()> { + if let Some(batch) = prepared.first.take() + && !self.send(&sender, batch).await? + { + return Ok(()); + } + let Some(mut rows) = prepared.rest.take() else { + return Ok(()); + }; + loop { + let batch = tokio::select! { + biased; + // The reader was dropped: the rows go with this task, which + // ends the call. + _ = sender.closed() => { + self.settle(Lifecycle::Cancelled); + return Ok(()); + } + _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), + result = rows.next_batch() => result, + }; + match batch { + Ok(Some(batch)) => { + if !self.send(&sender, batch).await? { + return Ok(()); + } + } + Ok(None) => { + self.settle(Lifecycle::Finished); + return Ok(()); + } + Err(error) => { + self.settle(failure(&error)); + return Err(error); + } + } + } + } + + async fn send( + &self, + sender: &mpsc::Sender>, + batch: RecordBatch, + ) -> Result { + tokio::select! { + biased; + _ = self.wait_for_cancellation() => Err(self.cancelled_error()), + result = sender.send(Ok(batch)) => { + if result.is_err() { + self.settle(Lifecycle::Cancelled); + } + Ok(result.is_ok()) + } + } + } +} + +pub(super) struct ExchangeQueryHandle { + query: Arc, + result_started: AtomicBool, +} + +impl ExchangeQueryHandle { + pub(super) fn new(query: Arc) -> Self { + Self { + query, + result_started: AtomicBool::new(false), + } + } +} + +impl Drop for ExchangeQueryHandle { + // Rows nobody took are rows nobody will read. Keeping the call open would + // hold the statement running on the server, stalled on flow control, + // until the registry forgot it. + fn drop(&mut self) { + if self.query.response.lock().unwrap().is_some() { + self.query.cancel(); + } + } +} + +#[async_trait::async_trait] +impl QueryHandle for ExchangeQueryHandle { + fn id(&self) -> Uuid { + self.query.touch(); + self.query.id + } + + async fn describe(&self) -> Result { + self.query.describe() + } + + async fn reader(&self) -> Result { + let timeout = self.query.client.overall_timeout()?; + if self + .result_started + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + return Err(Error::Runtime { + message: "SQL query results can only be consumed once".to_string(), + }); + } + let result_start = ResultStartGuard::new(&self.result_started); + self.query.touch(); + let started = Instant::now(); + let prepared = + match with_overall_timeout(timeout, "SQL query result", self.query.prepare()).await { + Ok(prepared) => prepared, + Err(error) => { + // A failed read or the timeout ends the query, and with it + // the call, including rows a timed-out wait put back. + self.query.settle(failure(&error)); + self.query.response.lock().unwrap().take(); + return Err(error); + } + }; + let remaining_timeout = timeout.map(|timeout| timeout.saturating_sub(started.elapsed())); + let schema = prepared.schema.clone(); + let (sender, receiver) = mpsc::channel(2); + let error_sender = sender.clone(); + let query = self.query.clone(); + tokio::spawn(async move { + let result = with_overall_timeout( + remaining_timeout, + "SQL query result", + query.clone().run(prepared, sender), + ) + .await; + if let Err(error) = result { + query.settle(failure(&error)); + let _ = error_sender.send(Err(error)).await; + } + }); + let stream = futures::stream::unfold(receiver, |mut receiver| async move { + receiver.recv().await.map(|item| (item, receiver)) + }); + result_start.commit(); + Ok(Box::pin(SimpleRecordBatchStream::new(stream, schema))) + } + + async fn cancel(&self) -> Result<()> { + self.query.cancel(); + Ok(()) + } +} diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 0d7efc9c4..79ca339db 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -4,7 +4,8 @@ use std::sync::atomic::AtomicUsize; use arrow_array::builder::StringDictionaryBuilder; -use arrow_array::{Array, Int64Array, StringArray, types::Int32Type}; +use arrow_array::types::{Float32Type, Int32Type}; +use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array, StringArray}; use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::flight_service_server::{FlightService, FlightServiceServer}; use arrow_flight::sql::{Any, CommandStatementQuery}; @@ -15,6 +16,7 @@ use arrow_flight::{ use arrow_schema::{DataType, Field, Schema}; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; +use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; use super::*; @@ -63,6 +65,9 @@ struct CapturedHeaders { database_prefix: String, } +/// Every exchange's statement and the parameter row it arrived with. +type ReceivedExchanges = Arc)>>>; + #[derive(Clone)] struct TestSqlService { query_count: Arc, @@ -77,11 +82,46 @@ struct TestSqlService { transient_poll_failures: Arc, headers: Arc>>, gate_refresh_release: Arc, + exchanges: ReceivedExchanges, + /// How many exchange responses have been dropped, finished or not. + exchange_drops: Arc, result: RecordBatch, large_result: RecordBatch, dictionary_result: RecordBatch, } +/// Counts a response stream being dropped, which is how a server sees a call +/// end -- by the client cancelling it or by the rows running out. +struct DropSignal(Arc); + +impl Drop for DropSignal { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +impl TestSqlService { + fn capture_headers(&self, metadata: &MetadataMap) { + let header = |name| { + metadata + .get(name) + .and_then(|value| value.to_str().ok()) + .unwrap() + .to_string() + }; + self.headers.lock().unwrap().push(CapturedHeaders { + database: header("database"), + namespace_path: header("namespace-path"), + request_id: header("x-request-id"), + api_key: header("x-api-key"), + database_prefix: metadata + .get("x-lancedb-database-prefix") + .map(|value| value.to_str().unwrap().to_string()) + .unwrap_or_default(), + }); + } +} + impl Default for TestSqlService { fn default() -> Self { let schema = Arc::new(Schema::new(vec![Field::new( @@ -126,6 +166,8 @@ impl Default for TestSqlService { transient_poll_failures: Arc::new(AtomicUsize::new(0)), headers: Arc::new(std::sync::Mutex::new(Vec::new())), gate_refresh_release: Arc::new(Notify::new()), + exchanges: Arc::new(std::sync::Mutex::new(Vec::new())), + exchange_drops: Arc::new(AtomicUsize::new(0)), result, large_result, dictionary_result, @@ -168,24 +210,7 @@ impl FlightService for TestSqlService { &self, request: Request, ) -> std::result::Result, Status> { - let metadata = request.metadata(); - let header = |name| { - metadata - .get(name) - .and_then(|value| value.to_str().ok()) - .unwrap() - .to_string() - }; - self.headers.lock().unwrap().push(CapturedHeaders { - database: header("database"), - namespace_path: header("namespace-path"), - request_id: header("x-request-id"), - api_key: header("x-api-key"), - database_prefix: metadata - .get("x-lancedb-database-prefix") - .map(|value| value.to_str().unwrap().to_string()) - .unwrap_or_default(), - }); + self.capture_headers(request.metadata()); let command = Any::decode(request.get_ref().cmd.as_ref()) .ok() @@ -394,11 +419,336 @@ impl FlightService for TestSqlService { Err(Status::unimplemented("list_actions")) } + /// Answers with the parameter row it was sent, so a test can compare what + /// came back with what went out. async fn do_exchange( &self, - _request: Request>, + request: Request>, ) -> std::result::Result, Status> { - Err(Status::unimplemented("do_exchange")) + self.capture_headers(request.metadata()); + let mut messages = request.into_inner(); + let first = messages + .message() + .await? + .ok_or_else(|| Status::invalid_argument("empty exchange"))?; + let query = first + .flight_descriptor + .as_ref() + .and_then(|descriptor| Any::decode(descriptor.cmd.as_ref()).ok()) + .and_then(|any| any.unpack::().ok().flatten()) + .map(|command| command.query) + .ok_or_else(|| Status::invalid_argument("exchange without a statement"))?; + let data = futures::stream::once(async move { Ok(first) }) + .chain(messages) + .map_err(|status| FlightError::Tonic(Box::new(status))); + let parameters = FlightRecordBatchStream::new_from_flight_data(data) + .try_next() + .await + .map_err(|error| Status::invalid_argument(error.to_string()))?; + self.exchanges + .lock() + .unwrap() + .push((query.clone(), parameters.clone())); + if query == "SELECT exchange refused" { + return Err(Status::invalid_argument( + "parameter 1 is not referenced by the statement", + )); + } + if query == "SELECT exchange stall" { + return Ok(Response::new(futures::stream::pending().boxed())); + } + let rows = + parameters.ok_or_else(|| Status::invalid_argument("exchange without parameters"))?; + let schema = rows.schema(); + let rest: BoxStream<'static, arrow_flight::error::Result> = + if query == "SELECT exchange slow" { + futures::stream::pending().boxed() + } else { + futures::stream::empty().boxed() + }; + let delay = if query == "SELECT exchange delayed" { + Duration::from_millis(300) + } else { + Duration::ZERO + }; + let dropped = DropSignal(self.exchange_drops.clone()); + let stream = FlightDataEncoderBuilder::new() + .with_schema(schema) + .build( + futures::stream::once(async move { + tokio::time::sleep(delay).await; + Ok(rows) + }) + .chain(rest), + ) + .map_err(Status::from) + .map(move |message| { + let _dropped = &dropped; + message + }); + Ok(Response::new(Box::pin(stream))) + } +} + +fn parameter_row() -> RecordBatch { + let vector = FixedSizeListArray::from_iter_primitive::( + [Some([0.1_f32, 1e-45, -0.0].map(Some))], + 3, + ); + RecordBatch::try_from_iter([ + ("vector", Arc::new(vector) as ArrayRef), + ("k", Arc::new(Int64Array::from(vec![10])) as ArrayRef), + ]) + .unwrap() +} + +fn parameterized(query: &str, parameters: RecordBatch) -> ExecuteQueryRequest { + ExecuteQueryRequest { + parameters: Some(parameters), + ..ExecuteQueryRequest::new(query) + } +} + +async fn start_sql_service( + service: TestSqlService, + client_config: ClientConfig, +) -> (SqlClient, tokio::sync::oneshot::Sender<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::try_unfold(listener, |listener| async { + let (socket, _) = listener.accept().await?; + Ok::<_, std::io::Error>(Some((socket, listener))) + }); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + tokio::spawn( + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }), + ); + let client = SqlClient::new( + "analytics".to_string(), + None, + "test-key".to_string(), + None, + Some(format!("grpc://{address}")), + client_config, + ); + (client, shutdown_tx) +} + +async fn wait_for(count: &AtomicUsize, expected: usize, what: &str) { + tokio::time::timeout(Duration::from_secs(5), async { + while count.load(Ordering::SeqCst) < expected { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{what}")); +} + +/// The statement and its row go out on one exchange -- no poll, no ticket -- +/// and the row arrives as exactly the bytes that were sent. +#[tokio::test] +async fn parameters_travel_on_one_exchange() { + let service = TestSqlService::default(); + let exchanges = service.exchanges.clone(); + let query_count = service.query_count.clone(); + let do_get_count = service.do_get_count.clone(); + let headers = service.headers.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let parameters = parameter_row(); + let query = client + .execute(ExecuteQueryRequest { + default_namespace_path: vec!["events".to_string(), "raw".to_string()], + ..parameterized("SELECT echo", parameters.clone()) + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + let rows = collect_result(&query).await.unwrap(); + assert_eq!(rows, vec![parameters.clone()]); + let echoed = rows[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .map(|value| value.to_bits()) + .collect::>(); + assert_eq!( + echoed, + [0.1_f32, 1e-45, -0.0].map(f32::to_bits), + "every float keeps its exact bits, the subnormal and the sign of zero included" + ); + + let finished = query.describe().await.unwrap(); + assert_eq!(finished.status, QueryStatus::Finished); + assert_eq!(finished.progress, Some(1.0)); + assert_eq!( + client.describe(query.id()).await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + + assert_eq!( + *exchanges.lock().unwrap(), + vec![("SELECT echo".to_string(), Some(parameters))] + ); + assert_eq!(query_count.load(Ordering::SeqCst), 0); + assert_eq!(do_get_count.load(Ordering::SeqCst), 0); + let headers = headers.lock().unwrap(); + assert_eq!(headers.len(), 1); + assert_eq!(headers[0].database, "analytics"); + assert_eq!(headers[0].namespace_path, "events$raw"); + assert_eq!(headers[0].api_key, "test-key"); +} + +/// The server plans before it answers, so a refusal is the submission's. +#[tokio::test] +async fn a_refused_statement_fails_its_submission() { + let (client, _shutdown) = + start_sql_service(TestSqlService::default(), ClientConfig::default()).await; + let refused = client + .execute(parameterized("SELECT exchange refused", parameter_row())) + .await + .unwrap_err(); + assert!( + refused + .to_string() + .contains("not referenced by the statement"), + "{refused}" + ); + + let two_rows = + RecordBatch::try_from_iter([("k", Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef)]) + .unwrap(); + assert!(matches!( + client.execute(parameterized("SELECT echo", two_rows)).await, + Err(Error::InvalidInput { .. }) + )); +} + +/// The call is the statement's lifetime: cancelling a read, or dropping a +/// handle nobody read, ends the call and so the work behind it. +#[tokio::test] +async fn cancelling_or_abandoning_an_exchange_ends_the_call() { + let service = TestSqlService::default(); + let drops = service.exchange_drops.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let reading = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let mut rows = reading.reader().await.unwrap(); + assert_eq!(rows.try_next().await.unwrap().unwrap(), parameter_row()); + reading.cancel().await.unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), rows.try_next()) + .await + .expect("cancellation must wake the reader"), + Err(Error::JobCancelled { .. }) + )); + assert_eq!( + reading.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + wait_for(&drops, 1, "cancelling must end the call on the server").await; + reading.cancel().await.unwrap(); + + let unread = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let id = unread.id(); + drop(unread); + wait_for(&drops, 2, "dropping an unread handle must end the call").await; + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Cancelled + ); + + let dropped_reader = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let rows = dropped_reader.reader().await.unwrap(); + drop(rows); + wait_for(&drops, 3, "dropping a reader must end the call").await; +} + +/// A reader abandoned before the first batch -- its task cancelled -- must +/// leave the rows for the next reader, as a polled query does, rather than +/// take the call down with it. +#[tokio::test] +async fn an_abandoned_reader_leaves_the_rows_for_the_next_one() { + let service = TestSqlService::default(); + let drops = service.exchange_drops.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let delayed = Arc::new( + client + .execute(parameterized("SELECT exchange delayed", parameter_row())) + .await + .unwrap(), + ); + let abandoned = { + let delayed = delayed.clone(); + tokio::spawn(async move { delayed.reader().await.map(|_| ()) }) + }; + tokio::time::sleep(Duration::from_millis(50)).await; + abandoned.abort(); + assert!(abandoned.await.unwrap_err().is_cancelled()); + assert_eq!( + delayed.describe().await.unwrap().status, + QueryStatus::Running + ); + assert_eq!(drops.load(Ordering::SeqCst), 0, "the call must survive"); + + assert_eq!( + collect_result(&delayed).await.unwrap(), + vec![parameter_row()] + ); + assert_eq!( + delayed.describe().await.unwrap().status, + QueryStatus::Finished + ); +} + +/// A reader that gives up before the first batch has already taken the rows, +/// so the query must say it failed rather than report `Running` for a call +/// that is gone. +#[tokio::test] +async fn a_reader_that_times_out_reports_the_failure() { + let mut client_config = ClientConfig::default(); + client_config.timeout_config.timeout = Some(Duration::from_millis(200)); + let (client, _shutdown) = start_sql_service(TestSqlService::default(), client_config).await; + + let stalled = client + .execute(parameterized("SELECT exchange stall", parameter_row())) + .await + .unwrap(); + assert_overall_timeout(stalled.reader().await, "result"); + let Err(retried) = stalled.reader().await else { + panic!("the rows are gone, so a second reader cannot succeed"); + }; + for error in [ + stalled.describe().await.unwrap_err(), + client.describe(stalled.id()).await.unwrap_err(), + retried, + ] { + assert!( + matches!(&error, Error::Runtime { message } if message == "SQL query result timed out"), + "{error}" + ); } } @@ -430,7 +780,10 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { let database = RemoteDatabase::for_catalog("https://catalog.example", Some(name), &options).unwrap(); let query = database - .execute_query_async("SELECT 42", &[]) + .execute_query_async(ExecuteQueryRequest { + default_namespace_path: vec![], + ..ExecuteQueryRequest::new("SELECT 42") + }) .await .unwrap(); assert_eq!( @@ -813,7 +1166,11 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { .submit("SELECT slow", &["public".to_string()]) .await .unwrap(); - let tracked_dropped_reader = client.queries.get(dropped_reader.id()).unwrap(); + let Some(RegisteredQuery::Polled(tracked_dropped_reader)) = + client.queries.get(dropped_reader.id()) + else { + panic!("a submitted statement is registered as a polled query"); + }; let continuation_count_before = first_continuation_count.load(Ordering::SeqCst); let dropped_result_stream = dropped_reader.reader().await.unwrap(); tokio::time::timeout(Duration::from_millis(100), async { @@ -1107,8 +1464,11 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { ) .unwrap(), ); - registry.insert(id, query.clone()); - assert!(Arc::ptr_eq(®istry.get(id).unwrap(), &query)); + registry.insert(id, RegisteredQuery::Polled(query.clone())); + let Some(RegisteredQuery::Polled(found)) = registry.get(id) else { + panic!("the registry returns the query it was given"); + }; + assert!(Arc::ptr_eq(&found, &query)); } let expired_id = Uuid::now_v7(); @@ -1125,7 +1485,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { ) .unwrap(), ); - registry.insert(expired_id, expired_query); + registry.insert(expired_id, RegisteredQuery::Polled(expired_query)); assert!(registry.get(expired_id).is_none()); let stale_id = Uuid::now_v7(); @@ -1142,7 +1502,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { .unwrap(), ); *stale_query.last_accessed.lock().unwrap() = Instant::now() - ABANDONED_QUERY_RETENTION; - registry.insert(stale_id, stale_query.clone()); + registry.insert(stale_id, RegisteredQuery::Polled(stale_query.clone())); drop(stale_query); assert!(registry.get(stale_id).is_none()); diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index 7c040c51b..c2ef13ba5 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -17,6 +17,10 @@ pub enum QueryStatus { /// The server is still executing the query. Running, /// The server has made the complete result available. + /// + /// A query submitted with parameters streams its rows on the call that + /// runs it, so it reports `Finished` only once its reader has received + /// the last batch. Read it rather than waiting for this status. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -60,7 +64,9 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query. +/// Dropping it does not cancel the server-side query, unless the query was +/// submitted with parameters: those run on the call that returns their rows, +/// so dropping the handle before reading, or dropping the reader, stops them. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -88,6 +94,11 @@ impl Query { } /// Get a point-in-time description of the query. + /// + /// For a query submitted with parameters the status follows its result + /// stream: it stays [`QueryStatus::Running`] until the reader has received + /// every row, so open [`Self::reader`] directly instead of polling for + /// [`QueryStatus::Finished`] first. pub async fn describe(&self) -> Result { self.handle.describe().await }