diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index f18ada7aa..fb9bada29 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -64,3 +64,6 @@ optional timeoutMs: number; ``` Timeout for query execution in milliseconds + +The timeout also covers reading the results, so consuming them slowly can +exceed it. diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 7e5140928..7f28d8804 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -859,6 +859,13 @@ older server omits structured resource details. A structured `already_exists` response keeps its more specific `AlreadyExistsError` type. Both are subclasses of `HttpError`. +A remote query whose streamed response stops partway raises +`ResponseInterruptedError`, also a subclass of `HttpError`. It carries the number +of batches and rows already returned and, when the server reports it, the table +version the response reflected, so the query can be run again against that +version. It is never retried automatically, because the batches already returned +would be returned again. + SQL response failures raise `SqlError`, retaining the numeric `grpc_code` and `request_id`. Servers that report a recognized `x-lancedb-error-code` select a specific input-error subclass. HTTP failures use the corresponding `HttpError` @@ -926,6 +933,8 @@ generic HTTP/SQL exceptions with their transport status and current request ID. ::: lancedb.remote.errors.ConflictError +::: lancedb.remote.errors.ResponseInterruptedError + ::: lancedb.exceptions.MissingValueError ::: lancedb.exceptions.MissingColumnError diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index cbd927925..bb845276f 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -237,6 +237,9 @@ export interface QueryExecutionOptions { /** * Timeout for query execution in milliseconds + * + * The timeout also covers reading the results, so consuming them slowly can + * exceed it. */ timeoutMs?: number; diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 08ea3b9bc..4a3fea2e2 100644 --- a/python/python/lancedb/query.py +++ b/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`. diff --git a/python/python/lancedb/remote/errors.py b/python/python/lancedb/remote/errors.py index 4dc835e53..ee92ddf53 100644 --- a/python/python/lancedb/remote/errors.py +++ b/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. diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py index d5e92731f..12a193074 100644 --- a/python/python/tests/test_errors.py +++ b/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" diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 25d9db22a..685bc809d 100644 --- a/python/python/tests/test_remote_db.py +++ b/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): diff --git a/python/src/error.rs b/python/src/error.rs index c4f20e884..1ac9f7ac3 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -168,6 +168,28 @@ impl PythonErrorExt for std::result::Result { 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, diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index cc5bd1f4e..7499641fc 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -205,6 +205,48 @@ fn http_timeout_hint(source: &(dyn std::error::Error + 'static)) -> &'static str } } +/// The error that ended a streamed response. +/// +/// It is shared so that a copy of [`Error::ResponseInterrupted`] keeps it: +/// DataFusion hands one error to every partition that reads the failed input. +/// It dereferences to the underlying error, which is also the next error in the +/// source chain. +#[cfg(feature = "remote")] +#[derive(Debug, Clone)] +pub struct InterruptionSource(Arc); + +#[cfg(feature = "remote")] +impl From for InterruptionSource { + fn from(source: BoxError) -> Self { + Self(source.into()) + } +} + +#[cfg(feature = "remote")] +impl std::ops::Deref for InterruptionSource { + type Target = dyn std::error::Error + Send + Sync; + + fn deref(&self) -> &Self::Target { + self.0.as_ref() + } +} + +#[cfg(feature = "remote")] +impl Display for InterruptionSource { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +// Not `Error`: an `Arc` in the source chain hides the error it holds from +// `downcast_ref`, and so would a wrapper that implements `Error`. +#[cfg(feature = "remote")] +impl snafu::AsErrorSource for InterruptionSource { + fn as_error_source(&self) -> &(dyn std::error::Error + 'static) { + self.0.as_ref() + } +} + /// Why a job failed, to whatever precision the backend provides. /// /// A job run in this process carries the error it failed with in [`Self::source`]. @@ -362,6 +404,27 @@ pub enum Error { source: Box, status_code: Option, }, + /// A streamed response stopped before the server signalled that it was complete. + /// + /// The batches already returned are valid, but the result is incomplete. The + /// request is not retried, because re-issuing it would return those batches + /// again. `table_version` is the table version the response reflected, when + /// the server reported one, so a caller can re-issue the request against it. + #[cfg(feature = "remote")] + #[snafu(display( + "Response for request_id={request_id} interrupted after {batches} batches \ + ({rows} rows): {source}{}", + http_timeout_hint(&**source) + ))] + ResponseInterrupted { + request_id: String, + /// Batches returned before the interruption. + batches: u64, + /// Rows in those batches. + rows: u64, + table_version: Option, + source: InterruptionSource, + }, #[snafu(display("Arrow error: {source}"))] Arrow { source: ArrowError }, #[snafu(display("LanceDBError: not supported: {message}"))] @@ -393,6 +456,20 @@ impl From for Error { match source { DataFusionError::ArrowError(source, _) => (*source).into(), DataFusionError::External(source) => Self::from_box_error(source), + // DataFusion adds context to an error, and shares one error between + // the partitions that read a failed input. An error of ours beneath + // either keeps its type; other errors keep the wrapper's message. + DataFusionError::Context(_, inner) if find_error(inner.as_ref()).is_some() => { + (*inner).into() + } + DataFusionError::Shared(shared) => match Arc::try_unwrap(shared) { + Ok(inner) => inner.into(), + Err(shared) => find_error(shared.as_ref()) + .and_then(Self::clone_shared) + .unwrap_or_else(|| Self::External { + source: Box::new(DataFusionError::Shared(shared)), + }), + }, other => Self::External { source: Box::new(other), }, @@ -400,6 +477,12 @@ impl From for Error { } } +/// The first error of ours in an error's source chain. +fn find_error<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a Error> { + std::iter::successors(Some(error), |error| error.source()) + .find_map(|error| error.downcast_ref::()) +} + impl From for Error { fn from(source: lance::Error) -> Self { if has_unsupported_local_filesystem_source(&source) { @@ -509,6 +592,29 @@ impl Error { } } + /// A copy of an error that other partitions of a DataFusion plan still hold. + /// Only an interrupted response can be copied, because it shares its source. + fn clone_shared(&self) -> Option { + #[cfg(feature = "remote")] + if let Self::ResponseInterrupted { + request_id, + batches, + rows, + table_version, + source, + } = self + { + return Some(Self::ResponseInterrupted { + request_id: request_id.clone(), + batches: *batches, + rows: *rows, + table_version: *table_version, + source: source.clone(), + }); + } + None + } + fn from_box_error(mut source: Box) -> Self { source = match source.downcast::() { Ok(e) => match *e { @@ -619,6 +725,67 @@ mod tests { assert!(matches!(Error::from(source), Error::InvalidInput { .. })); } + /// DataFusion shares an error between the partitions that read a failed + /// input, and may add context to it. Whether or not another partition still + /// holds it, an interrupted response keeps its fields and its source. + #[cfg(feature = "remote")] + #[rstest::rstest] + fn shared_interruption_keeps_its_type_and_source( + #[values(false, true)] still_held: bool, + #[values(false, true)] with_context: bool, + ) { + let reset = BoxError::from(std::io::Error::from(std::io::ErrorKind::ConnectionReset)); + let interrupted = Error::ResponseInterrupted { + request_id: "req-1".into(), + batches: 2, + rows: 2048, + table_version: Some(7), + source: reset.into(), + }; + let shared = Arc::new(DataFusionError::External(Box::new(interrupted))); + let held = still_held.then(|| shared.clone()); + let mut error = DataFusionError::Shared(shared); + if with_context { + error = error.context("merging query results"); + } + + let error = Error::from(error); + + assert!( + matches!( + &error, + Error::ResponseInterrupted { + request_id, + batches: 2, + rows: 2048, + table_version: Some(7), + .. + } if request_id == "req-1" + ), + "{error:?}" + ); + let source = std::error::Error::source(&error).unwrap(); + assert_eq!( + source.downcast_ref::().map(|e| e.kind()), + Some(std::io::ErrorKind::ConnectionReset) + ); + drop(held); + } + + #[test] + fn other_shared_errors_keep_their_message() { + let shared = Arc::new(DataFusionError::Execution("boom".into())); + let _held = shared.clone(); + + let error = Error::from(DataFusionError::Shared(shared).context("merging")); + + assert!(matches!(error, Error::External { .. }), "{error:?}"); + assert_eq!( + error.to_string(), + "merging\ncaused by\nExecution error: boom" + ); + } + #[cfg(feature = "remote")] #[test] fn non_timeout_http_error_message_is_unchanged() { diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index ad2dbdcbf..f96e84e96 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -661,6 +661,10 @@ pub struct QueryExecutionOptions { /// By default, this is 1024 pub max_batch_length: u32, /// Max duration to wait for the query to execute before timing out. + /// + /// The timeout also covers reading the results, so a caller that consumes + /// the stream slowly can exceed it. For a remote table it counts from when + /// the query is sent, and the server is given the same timeout. pub timeout: Option, /// How distributed worker metrics should be displayed by /// [`ExecutableQuery::analyze_plan`]. @@ -747,6 +751,11 @@ pub trait ExecutableQuery { /// stream is consumed slowly (this constrains the maximum memory used by a /// single query. /// + /// A remote table decodes batches as they arrive when the server streams + /// query results. If the response stops partway, the stream ends with + /// `Error::ResponseInterrupted`. The batches already returned are valid but + /// incomplete, and the query is not retried. + /// /// For simpler access or row-based access we recommend creating extension traits /// to convert Arrow data into your internal data model. fn execute_with_options( diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a19f4d0d9..6956905c4 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3,6 +3,7 @@ pub mod blobs; pub mod insert; +mod ipc_stream; mod optimize; use self::insert::{RemoteWriteExec, WriteOp}; @@ -61,15 +62,15 @@ use crate::{ }, }; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; -use arrow_ipc::reader::{FileReader, StreamReader}; -use arrow_schema::{ArrowError, DataType, SchemaRef}; +use arrow_ipc::reader::FileReader; +use arrow_schema::{DataType, SchemaRef}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use datafusion_common::DataFusionError; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; use futures::{StreamExt, TryStreamExt}; -use http::header::CONTENT_TYPE; +use http::header::{ACCEPT, CONTENT_TYPE}; use http::{HeaderName, StatusCode}; use lance::arrow::json::{JsonDataType, JsonSchema}; use lance::dataset::refs::TagContents; @@ -91,6 +92,10 @@ const MIN_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-ve const MIN_TIMESTAMP_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-timestamp"); const MIN_READ_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-read-version"); const VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-version"); +/// Asks for stream framing, which is decoded as it arrives. Servers that do not +/// stream query results ignore this and answer with file framing. +const QUERY_ACCEPT: &str = + "application/vnd.apache.arrow.stream, application/vnd.apache.arrow.file;q=0.9"; const METRIC_TYPE_KEY: &str = "metric_type"; const INDEX_TYPE_KEY: &str = "index_type"; const SCHEMA_CACHE_TTL: Duration = Duration::from_secs(30); @@ -1155,25 +1160,30 @@ impl RemoteTable { .map(str::to_owned); let framing = resolve_arrow_ipc_framing(content_type.as_deref(), request_id)?; - // Buffer the whole body. File framing keeps its footer at the end, so /query - // cannot decode incrementally. Stream framing could, via - // arrow_ipc::reader::StreamDecoder, but fetch_blobs concatenates every batch - // before returning, so no caller would see data sooner. - let body = response.bytes().await.err_to_http(request_id.into())?; - type IpcBatchIterator = - Box> + Send>; - let (schema, batches): (SchemaRef, IpcBatchIterator) = match framing { + match framing { ArrowIpcFraming::Stream => { - let reader = StreamReader::try_new(Cursor::new(body), None)?; - (reader.schema(), Box::new(reader)) + let table_version = response + .headers() + .get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()); + ipc_stream::read_ipc_stream( + response.bytes_stream(), + request_id.to_string(), + table_version, + ) + .await } ArrowIpcFraming::File => { + // File framing keeps its footer at the end, so the whole body is + // needed before anything can be decoded. + let body = response.bytes().await.err_to_http(request_id.into())?; let reader = FileReader::try_new(Cursor::new(body), None)?; - (reader.schema(), Box::new(reader)) + let schema = reader.schema(); + let stream = futures::stream::iter(reader).map_err(DataFusionError::from); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } - }; - let stream = futures::stream::iter(batches).map_err(DataFusionError::from); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } } fn apply_query_params( @@ -1544,8 +1554,15 @@ impl RemoteTable { ) -> Result>>> { let mut request = self .client - .post(&format!("/v1/table/{}/query/", self.identifier)); + .post(&format!("/v1/table/{}/query/", self.identifier)) + .header(ACCEPT, QUERY_ACCEPT); + // The timeout bounds the whole query, including reading the response, as + // the server's own deadline for this request does. A deadline too far + // away to represent is never reached. + let deadline = options + .timeout + .and_then(|timeout| Some((tokio::time::Instant::now().checked_add(timeout)?, timeout))); if let Some(timeout) = options.timeout { // Also send to server, so it can abort the query if it takes too long. // (If it doesn't fit into u64, it's not worth sending anyways.) @@ -1569,28 +1586,28 @@ impl RemoteTable { }); let streams = futures::future::try_join_all(futures); - let streams = if let Some(timeout) = options.timeout { - let timeout_future = tokio::time::sleep(timeout); - tokio::pin!(timeout_future); - tokio::pin!(streams); - tokio::select! { - _ = &mut timeout_future => { - return Err(Error::Other { - message: format!("Query timeout after {} ms", timeout.as_millis()), - source: None, - }) - } - result = &mut streams => { - result? - } - } + let streams = if let Some((deadline, timeout)) = deadline { + tokio::time::timeout_at(deadline, streams) + .await + .map_err(|_| Error::Other { + message: format!("Query timeout after {} ms", timeout.as_millis()), + source: None, + })?? } else { streams.await? }; Ok(streams .into_iter() - .map(|stream| restore_projection_order(stream, &query.base().select)) + .map(|stream| { + let stream = restore_projection_order(stream, &query.base().select); + match deadline { + Some((deadline, timeout)) => { + Box::pin(TimeoutStream::with_deadline(stream, timeout, deadline)) + } + None => stream, + } + }) .collect()) } @@ -1714,8 +1731,9 @@ fn schema_from_description(schema: JsonSchema) -> Result { ))) } -/// How a response body frames its Arrow IPC payload. `/query` answers with file framing -/// and `fetch_blobs` with stream framing, so the reader is picked per response. +/// How a response body frames its Arrow IPC payload. `fetch_blobs` answers with stream +/// framing, and `/query` with stream framing only on servers that stream query results, +/// so the reader is picked per response. enum ArrowIpcFraming { File, Stream, @@ -4439,6 +4457,7 @@ mod tests { BinaryArray, Int32Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray, StructArray, record_batch, }; + use arrow_ipc::reader::StreamReader; use arrow_schema::{DataType, Field, Schema}; use chrono::{DateTime, Utc}; use futures::{StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture}; @@ -7093,25 +7112,36 @@ mod tests { assert_eq!(blobs.value(0), b"alpha"); } + // Read as a stream, a file's leading magic declares a message longer than the + // whole body, so the body ends before the stream does. #[rstest] - #[case(ARROW_STREAM_CONTENT_TYPE, write_ipc_file(&one_row_blob_batch("image")))] + #[case(ARROW_STREAM_CONTENT_TYPE, write_ipc_file(&one_row_blob_batch("image")), true)] #[case( ARROW_FILE_CONTENT_TYPE, - write_ipc_stream_uncompressed(&one_row_blob_batch("image")) + write_ipc_stream_uncompressed(&one_row_blob_batch("image")), + false )] #[tokio::test] async fn test_fetch_blobs_fails_when_the_body_contradicts_the_content_type( #[case] content_type: &'static str, #[case] body: Vec, + #[case] ends_early: bool, ) { let table = table_with_fetch_blobs_content_type(Some(content_type), body); let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); - assert!( - matches!(error, Error::Arrow { .. }), - "expected an Arrow decode failure, got {error}" - ); + if ends_early { + assert!( + matches!(error, Error::ResponseInterrupted { batches: 0, .. }), + "expected an incomplete stream, got {error}" + ); + } else { + assert!( + matches!(error, Error::Arrow { .. }), + "expected an Arrow decode failure, got {error}" + ); + } } #[tokio::test] diff --git a/rust/lancedb/src/remote/table/ipc_stream.rs b/rust/lancedb/src/remote/table/ipc_stream.rs new file mode 100644 index 000000000..af31a9b28 --- /dev/null +++ b/rust/lancedb/src/remote/table/ipc_stream.rs @@ -0,0 +1,729 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Incremental decoding of response bodies that use Arrow IPC stream framing. + +use arrow_array::RecordBatch; +use arrow_buffer::Buffer; +use arrow_ipc::reader::StreamDecoder; +use bytes::Bytes; +use datafusion_common::DataFusionError; +use datafusion_physical_plan::SendableRecordBatchStream; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; +use futures::stream::BoxStream; +use futures::{Stream, StreamExt, TryStreamExt}; + +use crate::error::{BoxError, Error, Result}; + +/// Decode an Arrow IPC stream from a response body as its chunks arrive. +/// +/// Only the current chunk and the decoder's state are held, so memory is bounded +/// by the size of a batch rather than the size of the result. This waits for the +/// first batch, or for the end of an empty result, so a body that fails before +/// any data arrives fails the request itself. +/// +/// A complete body ends with the end-of-stream marker; the server aborts the body +/// without it when the query fails after sending headers. A body that ends early, +/// even between messages, is reported as [`Error::ResponseInterrupted`] and never +/// as a short result. +pub(super) async fn read_ipc_stream( + chunks: impl Stream> + Send + 'static, + request_id: String, + table_version: Option, +) -> Result +where + E: Into + 'static, +{ + let mut body = IpcStreamBody { + chunks: chunks.map_err(Into::into).boxed(), + chunk: Buffer::default(), + decoder: StreamDecoder::new(), + finished: false, + request_id, + table_version, + batches: 0, + rows: 0, + }; + let first = body.next_batch().await?; + // The schema is the first message, so only a body holding nothing but the + // end-of-stream marker can complete without one. + let schema = body.decoder.schema().ok_or_else(|| Error::Http { + source: "Arrow IPC stream ended without a schema".into(), + request_id: body.request_id.clone(), + status_code: None, + })?; + let rest = futures::stream::try_unfold(body, |mut body| async move { + Ok(body.next_batch().await?.map(|batch| (batch, body))) + }); + let batches = futures::stream::iter(first.map(Ok)) + .chain(rest) + .map_err(|error: Error| DataFusionError::External(Box::new(error))); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, batches))) +} + +struct IpcStreamBody { + chunks: BoxStream<'static, std::result::Result>, + /// The part of the current chunk that has not been decoded yet. + chunk: Buffer, + decoder: StreamDecoder, + /// Set once the body has ended with the end-of-stream marker. + finished: bool, + request_id: String, + table_version: Option, + batches: u64, + rows: u64, +} + +impl IpcStreamBody { + /// Decode until the next batch is complete. Returns `None` once the body + /// has ended with the end-of-stream marker. + async fn next_batch(&mut self) -> Result> { + while !self.finished { + if self.chunk.is_empty() { + match self.chunks.next().await { + Some(Ok(chunk)) => self.chunk = chunk.into(), + Some(Err(error)) => return Err(self.interrupted(error)), + None if read_end_of_stream(&mut self.decoder) => self.finished = true, + None => { + return Err(self.interrupted( + "the response ended before the Arrow end-of-stream marker".into(), + )); + } + } + continue; + } + // Decoding stops after each batch, keeping the rest of the chunk. + if let Some(batch) = self.decoder.decode(&mut self.chunk)? { + self.batches += 1; + self.rows += batch.num_rows() as u64; + return Ok(Some(batch)); + } + } + Ok(None) + } + + fn interrupted(&self, source: BoxError) -> Error { + Error::ResponseInterrupted { + request_id: self.request_id.clone(), + batches: self.batches, + rows: self.rows, + table_version: self.table_version, + source: source.into(), + } + } +} + +/// Whether the decoder has read the end-of-stream marker. +/// +/// [`StreamDecoder::finish`] also accepts input that stops between messages, +/// because the Arrow format lets a writer end a stream by closing it. A decoder +/// that has read the marker rejects any further input, which is the only way +/// its API tells the two apart. +fn read_end_of_stream(decoder: &mut StreamDecoder) -> bool { + decoder.finish().is_ok() && decoder.decode(&mut Buffer::from_slice_ref([0u8])).is_err() +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use arrow_array::builder::StringDictionaryBuilder; + use arrow_array::types::Int8Type; + use arrow_array::{ArrayRef, Int32Array, StringArray}; + use arrow_ipc::CompressionType; + use arrow_ipc::writer::{DictionaryHandling, FileWriter, IpcWriteOptions, StreamWriter}; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use lance::arrow::json::JsonSchema; + use rstest::rstest; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::sync::mpsc; + + use super::*; + use crate::query::{ExecutableQuery, QueryExecutionOptions}; + use crate::remote::client::{ClientConfig, RetryConfig}; + use crate::remote::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE}; + use crate::{Table, connect}; + + /// Three batches of two rows under a schema with metadata. The second batch's + /// dictionary extends the first one's, which writers can send as a delta, and + /// the third replaces it. + fn batches() -> Vec { + let schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("id", DataType::Int32, false), + Field::new_dictionary("label", DataType::Int8, DataType::Utf8, false), + ], + HashMap::from([("origin".to_string(), "test".to_string())]), + )); + let labels: [(&[&str], [&str; 2]); 3] = [ + (&["a", "b"], ["a", "b"]), + (&["a", "b", "c"], ["c", "a"]), + (&["d"], ["d", "d"]), + ]; + labels + .into_iter() + .enumerate() + .map(|(i, (dictionary, values))| { + let dictionary = StringArray::from(dictionary.to_vec()); + let mut label = + StringDictionaryBuilder::::new_with_dictionary(2, &dictionary) + .unwrap(); + values + .into_iter() + .for_each(|value| label.append_value(value)); + let id = Int32Array::from(vec![2 * i as i32, 2 * i as i32 + 1]); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(id) as ArrayRef, Arc::new(label.finish())], + ) + .unwrap() + }) + .collect() + } + + /// An IPC stream cut where the server cuts it: after the schema, after each + /// batch with its dictionaries, and after the end-of-stream marker. + fn messages( + schema: &Schema, + batches: &[RecordBatch], + options: IpcWriteOptions, + ) -> Vec> { + let mut writer = StreamWriter::try_new_with_options(Vec::new(), schema, options).unwrap(); + let mut messages = vec![std::mem::take(writer.get_mut())]; + for batch in batches { + writer.write(batch).unwrap(); + messages.push(std::mem::take(writer.get_mut())); + } + writer.finish().unwrap(); + messages.push(std::mem::take(writer.get_mut())); + messages + } + + fn chunks(body: &[u8], size: usize) -> Vec> { + body.chunks(size) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect() + } + + /// The schema, the batches yielded, and the error that ended the stream. + async fn read( + chunks: Vec>, + ) -> (Option, Vec, Option) { + let stream = + match read_ipc_stream(futures::stream::iter(chunks), "request".into(), Some(7)).await { + Ok(stream) => stream, + Err(error) => return (None, vec![], Some(error)), + }; + let schema = stream.schema(); + let results = stream.collect::>().await; + let (batches, errors): (Vec<_>, Vec<_>) = results.into_iter().partition(|r| r.is_ok()); + assert!(errors.len() <= 1, "the stream ends at its first error"); + let error = errors + .into_iter() + .next() + .map(|e| Error::from(e.unwrap_err())); + let batches = batches.into_iter().map(|r| r.unwrap()).collect(); + (Some(schema), batches, error) + } + + fn assert_interrupted(error: Option, expected_batches: u64, expected_rows: u64) { + match error { + Some(Error::ResponseInterrupted { + request_id, + batches, + rows, + table_version, + .. + }) => { + assert_eq!(request_id, "request"); + assert_eq!((batches, rows), (expected_batches, expected_rows)); + assert_eq!(table_version, Some(7)); + } + other => panic!("expected an interrupted response, got {other:?}"), + } + } + + // One-byte chunks split every message; a single chunk holds them all. + #[rstest] + #[case(1)] + #[case(7)] + #[case(usize::MAX)] + #[tokio::test] + async fn decodes_any_chunking(#[case] chunk_size: usize) { + let input = batches(); + let body = messages(&input[0].schema(), &input, IpcWriteOptions::default()).concat(); + + let (schema, output, error) = read(chunks(&body, chunk_size)).await; + + assert!(error.is_none(), "{error:?}"); + assert_eq!(schema.unwrap(), input[0].schema()); + assert_eq!(output, input); + } + + #[rstest] + #[case(IpcWriteOptions::default().with_dictionary_handling(DictionaryHandling::Delta))] + #[case(IpcWriteOptions::default().try_with_compression(Some(CompressionType::LZ4_FRAME)).unwrap())] + #[case(IpcWriteOptions::default().try_with_compression(Some(CompressionType::ZSTD)).unwrap())] + #[tokio::test] + async fn decodes_dictionary_deltas_and_compressed_batches(#[case] options: IpcWriteOptions) { + let input = batches(); + let body = messages(&input[0].schema(), &input, options).concat(); + + let (_, output, error) = read(chunks(&body, 5)).await; + + assert!(error.is_none(), "{error:?}"); + assert_eq!(output, input); + } + + /// The server answers an empty result with the schema and the end-of-stream marker. + #[tokio::test] + async fn empty_result_keeps_its_schema() { + let schema = batches()[0].schema(); + let body = messages(&schema, &[], IpcWriteOptions::default()).concat(); + + let (decoded, output, error) = read(chunks(&body, 3)).await; + + assert!(error.is_none(), "{error:?}"); + assert_eq!(decoded.unwrap(), schema); + assert!(output.is_empty()); + } + + #[tokio::test] + async fn body_ending_without_end_of_stream_is_interrupted() { + let input = batches(); + let messages = messages(&input[0].schema(), &input, IpcWriteOptions::default()); + let complete = messages[..messages.len() - 1].concat(); + let after_first_batch = messages[..2].concat().len(); + // Where the body stops, and how many batches it holds by then. + for (end, delivered) in [ + (complete.len(), 3), // between messages, after the last batch + (complete.len() - 3, 2), // inside the last batch + (after_first_batch, 1), // between messages + (messages[0].len(), 0), // after the schema + (messages[0].len() / 2, 0), // inside the schema + (0, 0), // an empty body + ] { + let (_, output, error) = read(chunks(&complete[..end], 4)).await; + assert_eq!(output, input[..delivered], "body cut at {end}"); + assert_interrupted(error, delivered as u64, 2 * delivered as u64); + } + } + + #[tokio::test] + async fn body_error_is_interrupted() { + let input = batches(); + let messages = messages(&input[0].schema(), &input, IpcWriteOptions::default()); + let mut body = chunks(&messages[..2].concat(), usize::MAX); + body.push(Err(std::io::Error::other("connection reset"))); + + let (_, output, error) = read(body).await; + + assert_eq!(output, input[..1]); + let message = error.as_ref().map(ToString::to_string).unwrap_or_default(); + assert!(message.contains("connection reset"), "{message}"); + assert_interrupted(error, 1, 2); + } + + #[tokio::test] + async fn malformed_body_is_an_arrow_error() { + let input = batches(); + let messages = messages(&input[0].schema(), &input, IpcWriteOptions::default()); + // A message whose metadata is not a flatbuffer. + let garbage = [&[0xff; 4][..], &8u32.to_le_bytes(), &[0xab; 8]].concat(); + // Anything after the end-of-stream marker. + let trailing = [messages.concat(), vec![0]].concat(); + for body in [[messages[0].clone(), garbage].concat(), trailing] { + let (_, _, error) = read(chunks(&body, usize::MAX)).await; + assert!(matches!(error, Some(Error::Arrow { .. })), "{error:?}"); + } + } + + /// What the test server does next for a query request. + enum Reply { + /// A complete response with this status and an empty body. + Status(u16), + /// Headers of a chunked body with this content type. + Head(Option<&'static str>), + Chunk(Vec), + /// The last chunk, ending the body cleanly. + End, + /// Close the connection in the middle of the body. + Abort, + } + + struct Server { + table: Table, + replies: mpsc::UnboundedSender, + /// Headers of each query request received. + queries: Arc>>>, + } + + impl Server { + fn reply(&self, replies: impl IntoIterator) { + replies + .into_iter() + .for_each(|reply| self.replies.send(reply).unwrap()); + } + } + + /// Serve table `t` over HTTP. Query requests are answered, in order, by the + /// replies the test sends. + async fn serve() -> Server { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (replies, mut pending) = mpsc::unbounded_channel(); + let queries = Arc::new(Mutex::new(Vec::new())); + let received = queries.clone(); + tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let (path, headers) = read_request(&mut socket).await; + if path.ends_with("/describe/") { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::new_fixed_size_list(DataType::Float32, 2, true), + false, + ), + ]); + let body = serde_json::json!({ + "version": 7, + "schema": JsonSchema::try_from(&schema).unwrap(), + }) + .to_string(); + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + socket.write_all(response.as_bytes()).await.unwrap(); + continue; + } + received.lock().unwrap().push(headers); + while let Some(reply) = pending.recv().await { + let last = !matches!(reply, Reply::Head(_) | Reply::Chunk(_)); + let bytes = match reply { + Reply::Status(status) => format!( + "HTTP/1.1 {status} Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .into_bytes(), + Reply::Head(content_type) => format!( + "HTTP/1.1 200 OK\r\n{}Transfer-Encoding: chunked\r\n\ + x-lancedb-version: 7\r\nConnection: close\r\n\r\n", + content_type + .map(|value| format!("Content-Type: {value}\r\n")) + .unwrap_or_default() + ) + .into_bytes(), + Reply::Chunk(chunk) => { + [format!("{:x}\r\n", chunk.len()).into_bytes(), chunk, b"\r\n".into()] + .concat() + } + Reply::End => b"0\r\n\r\n".to_vec(), + Reply::Abort => break, + }; + // The client may have stopped reading. + if socket.write_all(&bytes).await.is_err() || last { + break; + } + } + } + }); + let db = connect("db://test") + .api_key("key") + .region("us-east-1") + .host_override(&format!("http://{address}")) + .client_config(ClientConfig { + retry_config: RetryConfig { + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }) + .execute() + .await + .unwrap(); + let table = db.open_table("t").execute().await.unwrap(); + Server { + table, + replies, + queries, + } + } + + async fn read_request(socket: &mut TcpStream) -> (String, HashMap) { + let mut request = Vec::new(); + let mut buf = [0; 4096]; + let head_end = loop { + let n = socket.read(&mut buf).await.unwrap(); + assert!(n > 0, "connection closed before the request head"); + request.extend_from_slice(&buf[..n]); + if let Some(end) = request.windows(4).position(|w| w == b"\r\n\r\n") { + break end + 4; + } + }; + let head = String::from_utf8(request[..head_end].to_vec()).unwrap(); + let mut lines = head.lines(); + let path = lines.next().unwrap().split(' ').nth(1).unwrap().to_string(); + let headers: HashMap<_, _> = lines + .filter_map(|line| line.split_once(':')) + .map(|(name, value)| (name.to_ascii_lowercase(), value.trim().to_string())) + .collect(); + let length: usize = headers + .get("content-length") + .map_or(0, |length| length.parse().unwrap()); + while request.len() < head_end + length { + let n = socket.read(&mut buf).await.unwrap(); + assert!(n > 0, "connection closed before the request body"); + request.extend_from_slice(&buf[..n]); + } + (path, headers) + } + + /// The query batches as the server sends them: `id` only. + fn id_messages() -> Vec> { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batches: Vec<_> = (0..3) + .map(|i| { + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![i]))]) + .unwrap() + }) + .collect(); + messages(&schema, &batches, IpcWriteOptions::default()) + } + + fn ids(batch: &RecordBatch) -> Vec { + batch["id"] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[tokio::test] + async fn first_batch_arrives_before_the_rest_of_the_body() { + let server = serve().await; + let messages = id_messages(); + server.reply([ + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages[..2].concat()), + ]); + + // The server sends the rest only after the first batch has been read. + let (results, first) = tokio::time::timeout(Duration::from_secs(10), async { + let mut results = server.table.query().execute().await.unwrap(); + let first = results.next().await.unwrap().unwrap(); + (results, first) + }) + .await + .expect("the first batch must not wait for the whole body"); + assert_eq!(ids(&first), [0]); + server.reply([Reply::Chunk(messages[2..].concat()), Reply::End]); + let rest = results.try_collect::>().await.unwrap(); + assert_eq!(rest.iter().flat_map(ids).collect::>(), [1, 2]); + + let queries = server.queries.lock().unwrap(); + assert_eq!(queries.len(), 1); + assert_eq!(queries[0]["accept"], super::super::QUERY_ACCEPT); + } + + /// A server abort, and a proxy that turns it into a clean end of the body, + /// both fail the query after the batches already read. Neither is retried. + #[rstest] + #[case::abort(Reply::Abort, true)] + #[case::clean_end(Reply::End, false)] + #[tokio::test] + async fn interrupted_body_is_typed_and_not_retried( + #[case] end: Reply, + #[case] transport_error: bool, + ) { + let server = serve().await; + let messages = id_messages(); + server.reply([ + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages[..2].concat()), + end, + ]); + + let mut results = server.table.query().execute().await.unwrap(); + assert_eq!(ids(&results.next().await.unwrap().unwrap()), [0]); + let error = results.next().await.unwrap().unwrap_err(); + assert!(results.next().await.is_none()); + + let Error::ResponseInterrupted { + request_id, + batches, + rows, + table_version, + source, + } = error + else { + panic!("expected an interrupted response, got {error:?}"); + }; + assert!(!request_id.is_empty()); + assert_eq!((batches, rows, table_version), (1, 1, Some(7))); + assert_eq!(source.is::(), transport_error, "{source}"); + assert_eq!(server.queries.lock().unwrap().len(), 1); + } + + /// Vector queries run through a DataFusion plan, which must keep the error's type. + #[tokio::test] + async fn vector_query_keeps_the_interruption_type() { + let server = serve().await; + let messages = id_messages(); + server.reply([ + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages[..2].concat()), + Reply::Abort, + ]); + + let results = server + .table + .query() + .nearest_to(&[0.0f32, 0.0]) + .unwrap() + .execute() + .await + .unwrap(); + let error = results.try_collect::>().await.unwrap_err(); + + assert!( + matches!(error, Error::ResponseInterrupted { batches: 1, .. }), + "{error:?}" + ); + } + + /// A server without multi-vector support gets a request per vector, and the + /// plan merging the responses shares an error between its partitions. + #[tokio::test] + async fn multi_vector_query_keeps_the_interruption() { + let server = serve().await; + let messages = id_messages(); + // Which vector gets which response does not matter. + server.reply([ + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages.concat()), + Reply::End, + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages[..2].concat()), + Reply::Abort, + ]); + + let results = server + .table + .query() + .nearest_to(&[0.0f32, 0.0]) + .unwrap() + .add_query_vector(&[1.0f32, 1.0]) + .unwrap() + .execute() + .await + .unwrap(); + let error = results.try_collect::>().await.unwrap_err(); + + let Error::ResponseInterrupted { + request_id, + batches, + rows, + table_version, + source, + } = error + else { + panic!("expected an interrupted response, got {error:?}"); + }; + assert!(!request_id.is_empty()); + assert_eq!((batches, rows, table_version), (1, 1, Some(7))); + assert!(source.is::(), "{source}"); + assert_eq!(server.queries.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn timeout_covers_reading_the_body() { + let server = serve().await; + let messages = id_messages(); + // The body stalls after the first batch. + server.reply([ + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(messages[..2].concat()), + ]); + + let options = QueryExecutionOptions { + timeout: Some(Duration::from_millis(500)), + ..Default::default() + }; + let mut results = server + .table + .query() + .execute_with_options(options) + .await + .unwrap(); + assert_eq!(ids(&results.next().await.unwrap().unwrap()), [0]); + let error = tokio::time::timeout(Duration::from_secs(10), results.next()) + .await + .expect("the timeout must end a stalled body") + .unwrap() + .unwrap_err(); + + assert!( + !matches!(error, Error::ResponseInterrupted { .. }), + "{error:?}" + ); + assert!( + error.to_string().contains("Query timeout after 500 ms"), + "{error}" + ); + assert_eq!( + server.queries.lock().unwrap()[0]["x-request-timeout-ms"], + "500" + ); + } + + /// Servers that do not stream query results answer with file framing, or + /// with no content type at all. + #[rstest] + #[case(Some(ARROW_FILE_CONTENT_TYPE))] + #[case(None)] + #[tokio::test] + async fn file_framing_is_still_read(#[case] content_type: Option<&'static str>) { + let server = serve().await; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![4, 5]))]) + .unwrap(); + let mut writer = FileWriter::try_new(Vec::new(), &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + server.reply([ + Reply::Head(content_type), + Reply::Chunk(writer.into_inner().unwrap()), + Reply::End, + ]); + + let results = server.table.query().execute().await.unwrap(); + let output = results.try_collect::>().await.unwrap(); + + assert_eq!(output, [batch]); + } + + #[tokio::test] + async fn failure_before_headers_is_retried() { + let server = serve().await; + server.reply([ + Reply::Status(503), + Reply::Head(Some(ARROW_STREAM_CONTENT_TYPE)), + Reply::Chunk(id_messages().concat()), + Reply::End, + ]); + + let results = server.table.query().execute().await.unwrap(); + let output = results.try_collect::>().await.unwrap(); + + assert_eq!(output.iter().flat_map(ids).collect::>(), [0, 1, 2]); + assert_eq!(server.queries.lock().unwrap().len(), 2); + } +} diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 560fe858a..eedff924d 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -680,6 +680,22 @@ impl TimeoutStream { Box::pin(Self::new(inner, timeout)) } + /// Like [`Self::new`], but the stream expires at `deadline` rather than + /// `timeout` after its first poll. `timeout` is only used in the error. + pub fn with_deadline( + inner: SendableRecordBatchStream, + timeout: std::time::Duration, + deadline: tokio::time::Instant, + ) -> Self { + Self { + inner, + state: TimeoutState::Started { + deadline: Box::pin(tokio::time::sleep_until(deadline)), + timeout, + }, + } + } + fn timeout_error(timeout: &std::time::Duration) -> DataFusionError { DataFusionError::Execution(format!("Query timeout after {} ms", timeout.as_millis())) }