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 `$<i+1>`; a column whose name is not a number also
binds `$<name>`. 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

@@ -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

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(

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),
             )
         )
 

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:

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,
                 )
             )
         )

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 ``$<i + 1>`` and, when its name is not a number,
+also ``$<name>``.
+"""
+
+
+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"]

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]

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<Bound<'_, PyAny>>,
+        parameters: Option<PyArrowType<RecordBatch>>,
     ) -> PyResult<Bound<'a, PyAny>> {
         let inner = self_.get_inner()?.clone();
         let default_namespace_path = parse_default_namespace_path(default_namespace_path)?;
         future_into_py(self_.py(), async move {
-            let operation = inner
+            let mut operation = inner
                 .execute_query_async(query)
                 .default_namespace_path(default_namespace_path);
+            if let Some(PyArrowType(parameters)) = parameters {
+                operation = operation.parameters(parameters);
+            }
             operation
                 .execute()
                 .await

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<dyn Database>,
-    query: String,
-    default_namespace_path: Vec<String>,
+    request: ExecuteQueryRequest,
 }
 
 impl ExecuteQueryAsyncBuilder {
     fn new(parent: Arc<dyn Database>, 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<Item = S>,
         S: Into<String>,
     {
-        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 `$<i + 1>`, and a
+    /// column whose name is not a number also binds `$<name>`, 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::<Float32Type, _, _>(
+    ///     [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<crate::sql::Query> {
-        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
     ///

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<String>,
+    /// Values for the statement's placeholders, as a single-row batch.
+    ///
+    /// Column `i` binds `$<i + 1>`, and a column whose name is not a number
+    /// also binds `$<name>`. See
+    /// [`ExecuteQueryAsyncBuilder::parameters`](crate::connection::ExecuteQueryAsyncBuilder::parameters).
+    pub parameters: Option<RecordBatch>,
+}
+
+impl ExecuteQueryRequest {
+    pub fn new(query: impl Into<String>) -> 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<crate::sql::Query> {
         Err(crate::error::Error::NotSupported {
             message: "SQL is not supported by this database".to_string(),

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<S: HttpSend> Database for RemoteDatabase<S> {
         })
     }
 
-    async fn execute_query_async(
-        &self,
-        query: &str,
-        default_namespace_path: &[String],
-    ) -> Result<crate::sql::Query> {
+    async fn execute_query_async(&self, request: ExecuteQueryRequest) -> Result<crate::sql::Query> {
         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<crate::sql::QueryDescription> {

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<Query> {
+        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<Query> {
+        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<QueryDescription> {
         let query = self
             .queries
@@ -530,8 +582,40 @@ impl SqlClientInner {
     }
 }
 
+/// A submitted query as the connection remembers it.
+#[derive(Clone)]
+enum RegisteredQuery {
+    Polled(Arc<RemoteQuery>),
+    Exchanged(Arc<ExchangeQuery>),
+}
+
+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<QueryDescription> {
+        match self {
+            Self::Polled(query) => query.describe().await,
+            Self::Exchanged(query) => query.describe(),
+        }
+    }
+}
+
 struct QueryRegistry {
-    queries: StdMutex<HashMap<Uuid, Arc<RemoteQuery>>>,
+    queries: StdMutex<HashMap<Uuid, RegisteredQuery>>,
 }
 
 impl QueryRegistry {
@@ -541,12 +625,12 @@ impl QueryRegistry {
         }
     }
 
-    fn insert(&self, id: Uuid, query: Arc<RemoteQuery>) {
+    fn insert(&self, id: Uuid, query: RegisteredQuery) {
         self.remove_expired();
         self.queries.lock().unwrap().insert(id, query);
     }
 
-    fn get(&self, id: Uuid) -> Option<Arc<RemoteQuery>> {
+    fn get(&self, id: Uuid) -> Option<RegisteredQuery> {
         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());
     }
 }
 

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<SqlClientInner>,
+    /// The call's rows, until a reader takes them or a cancellation drops them.
+    response: StdMutex<Option<ResultEndpointStream>>,
+    lifecycle: StdMutex<Lifecycle>,
+    cancelled: Notify,
+    terminal_at: OnceLock<Instant>,
+    last_accessed: StdMutex<Instant>,
+}
+
+struct PreparedExchange {
+    schema: SchemaRef,
+    first: Option<RecordBatch>,
+    rest: Option<ResultEndpointStream>,
+}
+
+/// 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<ResultEndpointStream>,
+}
+
+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<SqlClientInner>,
+        query: &str,
+        parameters: RecordBatch,
+        default_namespace_path: &[String],
+    ) -> Result<Self> {
+        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<QueryDescription> {
+        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<PreparedExchange> {
+        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<Self>,
+        mut prepared: PreparedExchange,
+        sender: mpsc::Sender<Result<RecordBatch>>,
+    ) -> 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<Result<RecordBatch>>,
+        batch: RecordBatch,
+    ) -> Result<bool> {
+        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<ExchangeQuery>,
+    result_started: AtomicBool,
+}
+
+impl ExchangeQueryHandle {
+    pub(super) fn new(query: Arc<ExchangeQuery>) -> 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<QueryDescription> {
+        self.query.describe()
+    }
+
+    async fn reader(&self) -> Result<SendableRecordBatchStream> {
+        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(())
+    }
+}

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<std::sync::Mutex<Vec<(String, Option<RecordBatch>)>>>;
+
 #[derive(Clone)]
 struct TestSqlService {
     query_count: Arc<AtomicUsize>,
@@ -77,11 +82,46 @@ struct TestSqlService {
     transient_poll_failures: Arc<AtomicUsize>,
     headers: Arc<std::sync::Mutex<Vec<CapturedHeaders>>>,
     gate_refresh_release: Arc<Notify>,
+    exchanges: ReceivedExchanges,
+    /// How many exchange responses have been dropped, finished or not.
+    exchange_drops: Arc<AtomicUsize>,
     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<AtomicUsize>);
+
+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<arrow_flight::FlightDescriptor>,
     ) -> std::result::Result<Response<PollInfo>, 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<Streaming<FlightData>>,
+        request: Request<Streaming<FlightData>>,
     ) -> std::result::Result<Response<Self::DoExchangeStream>, 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::<CommandStatementQuery>().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<RecordBatch>> =
+            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::<Float32Type, _, _>(
+        [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::<FixedSizeListArray>()
+        .unwrap()
+        .values()
+        .as_any()
+        .downcast_ref::<Float32Array>()
+        .unwrap()
+        .values()
+        .iter()
+        .map(|value| value.to_bits())
+        .collect::<Vec<_>>();
+    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(&registry.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());
 

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<QueryDescription> {
         self.handle.describe().await
     }