feat(sql): bind query parameters over Flight DoExchange

execute_query / execute_query_async accept parameters: a one-row Arrow
batch bound to the statement's $1 / $name placeholders. A statement with
parameters is sent on one DoExchange (command descriptor + parameter
schema, then the row) and its rows come back on the same call, so floats
keep their exact bits and vectors stay compact FixedSizeLists instead of
being spliced into SQL text. Statements without parameters keep the
PollFlightInfo + DoGet path.

Database::execute_query_async now takes an ExecuteQueryRequest.

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,30 @@ 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.
+
         Local connections do not support SQL.
         """
         raise NotImplementedError("SQL is not supported for this connection type")
@@ -3267,6 +3290,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 +3301,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 +3310,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,75 @@ 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) and value.ndim == 1:
+        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
+    if isinstance(parameters, pa.Table):
+        parameters = parameters.combine_chunks()
+        batches = parameters.to_batches()
+        parameters = (
+            batches[0]
+            if len(batches) == 1
+            else pa.RecordBatch.from_pylist([], schema=parameters.schema)
+        )
+    if isinstance(parameters, pa.RecordBatch):
+        batch = parameters
+    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 +156,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