fix(remote): stream query responses instead of buffering them (#4929)

python/python/lancedb/query.py

@@ -1335,6 +1335,9 @@ class LanceQueryBuilder(ABC):
         and also the "_distance" column which is the distance between the query
         vector and the returned vectors.
 
+        This collects all results in memory. To process a large result a batch
+        at a time, use [to_batches][lancedb.query.LanceQueryBuilder.to_batches].
+
         Parameters
         ----------
         timeout: Optional[timedelta]
@@ -1366,6 +1369,11 @@ class LanceQueryBuilder(ABC):
         ``fetch_blobs``. The ``"bytes"`` and ``"descriptions"`` modes hide
         automatic row ids.
 
+        A remote table decodes batches as they arrive when the server streams
+        query results. If the response stops partway, reading raises
+        [ResponseInterruptedError][lancedb.remote.errors.ResponseInterruptedError]
+        after the batches already returned, which are valid but incomplete.
+
         Parameters
         ----------
         batch_size: int
@@ -1373,6 +1381,8 @@ class LanceQueryBuilder(ABC):
         timeout: Optional[timedelta]
             The maximum time to wait for the query to complete.
             If None, wait indefinitely.
+            The timeout also covers reading the batches, so consuming them
+            slowly can exceed it.
         blob_mode: str, default "lazy"
             ``"lazy"`` keeps blob descriptors and auto row ids for manual fetches.
             ``"descriptions"`` omits auto row ids. ``"bytes"`` fetches blob v2
@@ -2303,6 +2313,8 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
         timeout: timedelta, default None
             The maximum time to wait for the query to complete.
             If None, wait indefinitely.
+            The timeout also covers reading the batches, so consuming them
+            slowly can exceed it.
         blob_mode: str, default "lazy"
             Controls blob conversion as described in
             :meth:`LanceQueryBuilder.to_batches`.
@@ -3570,6 +3582,11 @@ class AsyncQueryBase(object):
         """
         Execute the query and return the results as an Apache Arrow RecordBatchReader.
 
+        A remote table decodes batches as they arrive when the server streams
+        query results. If the response stops partway, reading raises
+        [ResponseInterruptedError][lancedb.remote.errors.ResponseInterruptedError]
+        after the batches already returned, which are valid but incomplete.
+
         Parameters
         ----------
 
@@ -3582,6 +3599,8 @@ class AsyncQueryBase(object):
             The maximum time to wait for the query to complete.
             If not specified, no timeout is applied. If the query does not
             complete within the specified time, an error will be raised.
+            The timeout also covers reading the batches, so consuming them
+            slowly can exceed it.
         blob_mode: str, default "lazy"
             ``"lazy"`` keeps blob descriptors and auto row ids for manual fetches.
             ``"descriptions"`` omits auto row ids. ``"bytes"`` fetches blob v2
@@ -5114,6 +5133,8 @@ class BaseQueryBuilder(object):
             The maximum time to wait for the query to complete.
             If not specified, no timeout is applied. If the query does not
             complete within the specified time, an error will be raised.
+            The timeout also covers reading the batches, so consuming them
+            slowly can exceed it.
         blob_mode: str, default "lazy"
             Controls blob conversion as described in
             :meth:`AsyncQueryBase.to_batches`.

python/python/lancedb/remote/errors.py

@@ -286,6 +286,64 @@ class ConflictError(HttpError):
     """
 
 
+class ResponseInterruptedError(HttpError):
+    """A streamed query response stopped before the server finished sending it.
+
+    The batches already returned are valid, but the result is incomplete. The
+    query is not retried automatically, because a new attempt would return those
+    rows a second time. To read the same data in a new attempt, check out
+    ``table_version`` first.
+
+    Attributes
+    ----------
+    batches: int
+        The number of batches returned before the interruption.
+    rows: int
+        The number of rows in those batches.
+    table_version: int or None
+        The table version the response reflected, if the server reported it.
+
+    Examples
+    --------
+    >>> error = ResponseInterruptedError("interrupted", "req-1", 2, 2048, 7)
+    >>> isinstance(error, HttpError)
+    True
+    >>> (error.rows, error.table_version)
+    (2048, 7)
+    """
+
+    # HttpError.__new__ takes HttpError's arguments, to choose a subclass.
+    def __new__(cls, *args, **kwargs):
+        return RuntimeError.__new__(cls)
+
+    def __init__(
+        self,
+        message: str,
+        request_id: str,
+        batches: int,
+        rows: int,
+        table_version: Optional[int] = None,
+        is_transport_error: bool = False,
+    ):
+        super().__init__(message, request_id, is_transport_error=is_transport_error)
+        self.batches = batches
+        self.rows = rows
+        self.table_version = table_version
+
+    def __reduce__(self) -> tuple[type, tuple]:
+        return (
+            self.__class__,
+            (
+                str(self),
+                self.request_id,
+                self.batches,
+                self.rows,
+                self.table_version,
+                self.is_transport_error,
+            ),
+        )
+
+
 class TableNotFoundError(HttpError, ValueError, _NamespaceTableNotFoundError):
     """A remote table was not found.
 

python/python/tests/test_errors.py

@@ -15,6 +15,7 @@ from lancedb.remote.errors import (
     HttpError,
     LanceDBClientError,
     NamespaceNotFoundError,
+    ResponseInterruptedError,
     RetryError,
     SecretNotFoundError,
     TableNotFoundError,
@@ -73,6 +74,17 @@ def test_pickle_http_transport_error():
     assert restored.request_id == "req-789"
 
 
+def test_pickle_response_interrupted_error():
+    err = ResponseInterruptedError("interrupted", "req-1", 2, 2048, 7, True)
+    restored = pickle.loads(pickle.dumps(err))
+    assert type(restored) is ResponseInterruptedError
+    assert str(restored) == "interrupted"
+    assert restored.request_id == "req-1"
+    assert (restored.batches, restored.rows, restored.table_version) == (2, 2048, 7)
+    assert restored.is_transport_error is True
+    assert restored.status_code is None
+
+
 @pytest.mark.parametrize("error_code", [None, "conflict", "future_conflict_code"])
 def test_http_conflict_error_preserves_reason_and_metadata(error_code):
     reason = "Conflict: branch 'exp' cannot be deleted, main depends on its files"

python/python/tests/test_remote_db.py

@@ -48,6 +48,7 @@ from lancedb.remote.errors import (
     ConflictError,
     HttpError,
     NamespaceNotFoundError,
+    ResponseInterruptedError,
     RetryError,
 )
 import pytest
@@ -3398,6 +3399,100 @@ def test_query_sync_minimal():
         assert data == expected
 
 
+def _arrow_stream(batches):
+    sink = pa.BufferOutputStream()
+    with pa.ipc.new_stream(sink, batches[0].schema) as writer:
+        for batch in batches:
+            writer.write_batch(batch)
+    return sink.getvalue().to_pybytes()
+
+
+@contextlib.contextmanager
+def streamed_query_table(write_body):
+    """Yield a remote table and the headers of each query request it sends.
+
+    Queries get a 200 response with stream framing, whose body ``write_body``
+    writes after the other headers.
+    """
+    queries = []
+
+    def handler(request):
+        if request.path == "/v1/table/test/describe/":
+            id_field = {"name": "id", "type": {"type": "int32"}, "nullable": True}
+            schema = {"fields": [id_field]}
+            request.send_response(200)
+            request.send_header("Content-Type", "application/json")
+            request.end_headers()
+            request.wfile.write(json.dumps({"version": 7, "schema": schema}).encode())
+        elif request.path == "/v1/table/test/query/":
+            queries.append(request.headers)
+            request.rfile.read(int(request.headers["Content-Length"]))
+            request.send_response(200)
+            request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
+            request.send_header("x-lancedb-version", "7")
+            write_body(request)
+        else:
+            request.send_response(404)
+            request.end_headers()
+
+    with mock_lancedb_connection(handler) as db:
+        yield db.open_table("test"), queries
+
+
+def test_query_sync_to_batches_yields_before_the_body_is_complete():
+    batches = [pa.record_batch({"id": pa.array([i], pa.int32())}) for i in range(3)]
+    body = _arrow_stream(batches)
+    # The schema and first batch, without the 8-byte end-of-stream marker.
+    head = _arrow_stream(batches[:1])[:-8]
+    first_batch_read = threading.Event()
+    waited_for_client = []
+
+    def write_body(request):
+        request.end_headers()
+        request.wfile.write(head)
+        waited_for_client.append(first_batch_read.wait(10))
+        request.wfile.write(body[len(head) :])
+
+    with streamed_query_table(write_body) as (table, queries):
+        reader = table.search().to_batches()
+        assert reader.read_next_batch()["id"].to_pylist() == [0]
+        first_batch_read.set()
+        assert reader.read_all()["id"].to_pylist() == [1, 2]
+
+    # The wait times out if the client needs the whole body for the first batch.
+    assert waited_for_client == [True]
+    assert queries[0]["Accept"] == (
+        "application/vnd.apache.arrow.stream, application/vnd.apache.arrow.file;q=0.9"
+    )
+
+
+@pytest.mark.parametrize("abort", [False, True], ids=["clean_end", "abort"])
+def test_query_sync_interrupted_response_raises(abort):
+    batches = [pa.record_batch({"id": pa.array([i], pa.int32())}) for i in range(3)]
+    head = _arrow_stream(batches[:1])[:-8]
+
+    def write_body(request):
+        if abort:
+            # Closing before the promised length is a transport error. Otherwise
+            # the body ends cleanly, but without the end-of-stream marker.
+            request.send_header("Content-Length", str(len(head) + 100))
+        request.end_headers()
+        request.wfile.write(head)
+
+    with streamed_query_table(write_body) as (table, queries):
+        reader = table.search().to_batches()
+        assert reader.read_next_batch()["id"].to_pylist() == [0]
+        with pytest.raises(ResponseInterruptedError) as exc_info:
+            reader.read_next_batch()
+
+    error = exc_info.value
+    assert isinstance(error, HttpError)
+    assert (error.batches, error.rows, error.table_version) == (1, 1, 7)
+    assert error.request_id == queries[0]["x-request-id"]
+    assert error.is_transport_error is abort
+    assert len(queries) == 1
+
+
 @pytest.mark.parametrize("query_type", ["auto", "vector", "hybrid"])
 @pytest.mark.parametrize("vector_column_name", [None, "embedding"])
 def test_query_sync_text_embedding(query_type, vector_column_name):

python/src/error.rs

@@ -168,6 +168,28 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
                     err.setattr(intern!(py, "__cause__"), cause_err)?;
                     Err(PyErr::from_value(err))
                 }),
+                LanceError::ResponseInterrupted {
+                    request_id,
+                    batches,
+                    rows,
+                    table_version,
+                    source,
+                } => Python::attach(|py| {
+                    let error = py
+                        .import(intern!(py, "lancedb.remote.errors"))?
+                        .getattr(intern!(py, "ResponseInterruptedError"))?
+                        .call1((
+                            err.to_string(),
+                            request_id,
+                            *batches,
+                            *rows,
+                            *table_version,
+                            is_transport_error(&**source),
+                        ))?;
+                    let cause = http_from_rust_error(py, &**source, request_id, None)?;
+                    error.setattr(intern!(py, "__cause__"), cause)?;
+                    Err(PyErr::from_value(error))
+                }),
                 LanceError::Sql {
                     message,
                     request_id,