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.

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