diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 5d1393924..41d3e16d7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -221,10 +221,19 @@ The connected database and `default_namespace_path=["public"]` are used for unqualified tables. Fully qualified references can still query other databases and namespaces available to the same deployment. `execute_query` returns a reader as soon as its initial result stream is available. `execute_query_async` -returns a query handle immediately; use it to inspect progress, open a reader, +returns a query handle once submission is accepted; use it to inspect progress, open a reader, or cancel the query. The SQL client is initialized by the first query and retained for the lifetime of the remote connection. Query ids are random, -connection-scoped references rather than encoded SQL or durable resume tokens: +connection-scoped references rather than encoded SQL or durable resume tokens. + +Both APIs start with one Flight call. The server can stream results on that +call or return a polling continuation that the SDK follows automatically. +This behavior is the same with and without parameters. Open the reader directly; +a directly streamed query needs its reader to make progress. Dropping an unread +handle or its reader stops direct streaming. Deferred execution continues until +explicitly cancelled. + +For example: ```python import lancedb diff --git a/docs/web/basics/sql.mdx b/docs/web/basics/sql.mdx index 1c4ad46cd..c89bd7a20 100644 --- a/docs/web/basics/sql.mdx +++ b/docs/web/basics/sql.mdx @@ -82,6 +82,8 @@ In application code, load credentials from environment variables or a secrets manager. The [connection guide](/basics/connect#enterprise) shows that setup. `execute_query()` returns a blocking Arrow reader when results start arriving. +The server can stream rows on the first call or return result tickets that the +SDK follows automatically, with the same behavior for parameterized queries. The example uses `read_all()` to collect the small result in memory. For larger results, consume the reader in batches. To submit a query and check its status separately, use `execute_query_async()`. See [Queries](/basics/query) for query diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 74acc28ba..8e53eac04 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1253,12 +1253,13 @@ class DBConnection(EnforceOverrides): 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``. + Parameters are supported in queries, not in DDL or DML. With or + without parameters, the server chooses whether to stream rows on the + initial call or return a polling continuation. The reader follows + either path transparently. Dropping an unread handle or its reader + stops a directly streamed query; deferred execution requires explicit + cancellation. Open ``reader()`` directly: a streamed query needs its + reader to make progress and reports ``running`` until the stream ends. Local connections do not support SQL. """ diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 482acf89a..26afef5e8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -372,12 +372,10 @@ impl ExecuteQueryAsyncBuilder { /// Do not follow a named placeholder directly with `$`: `$a$` starts a /// dollar-quoted string. /// - /// A parameterized statement runs on the call that returns its rows rather - /// than detached from it: cancelling the query, or dropping its handle - /// before reading it, stops the statement on the server. Its status follows - /// that stream too -- it reports `Running` until the reader has received - /// every row -- so open the reader directly rather than waiting for - /// `Finished`. + /// With or without parameters, the server may stream rows on the initial + /// call or return a polling continuation. The reader follows either path + /// transparently. Open it directly instead of waiting for `Finished`: + /// directly streamed queries need their reader to make progress. /// /// # Example /// diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index a58a9afad..d8f04119f 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -34,9 +34,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; -mod exchange; +mod stream; -use exchange::{ExchangeQuery, ExchangeQueryHandle}; +use stream::{StreamingQuery, StreamingQueryHandle, Submission}; const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; @@ -191,20 +191,39 @@ impl SqlClient { } pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { - match request.parameters { - Some(parameters) => { - self.submit_with_parameters( - &request.query, - parameters, - &request.default_namespace_path, - ) - .await + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(&request.default_namespace_path)?; + let query_id = Uuid::now_v7(); + let submission = StreamingQuery::open( + self.inner.clone(), + &request.query, + request.parameters, + &request.default_namespace_path, + ) + .await?; + match submission { + Submission::Stream(response) => { + let query = + Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response)); + self.queries + .insert(query_id, RegisteredQuery::Streamed(query.clone())); + Ok(Query::new(Arc::new(StreamingQueryHandle::new(query)))) + } + Submission::Deferred(poll_info) => { + let query = Arc::new(RemoteQuery::new( + query_id, + self.inner.clone(), + request.default_namespace_path, + poll_info, + )?); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); + Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) + } } - None => { - self.submit(&request.query, &request.default_namespace_path) - .await - } - } + }) + .await } pub(super) async fn submit( @@ -271,35 +290,6 @@ impl SqlClient { }) } - /// Values cannot ride on a poll, so a statement that has them is sent on - /// its own exchange instead. - pub(super) async fn submit_with_parameters( - &self, - query: &str, - parameters: RecordBatch, - default_namespace_path: &[String], - ) -> Result { - let timeout = self.inner.overall_timeout()?; - with_overall_timeout(timeout, "SQL query submission", async { - validate_namespace_path(default_namespace_path)?; - let query_id = Uuid::now_v7(); - let query = Arc::new( - ExchangeQuery::open( - query_id, - self.inner.clone(), - query, - parameters, - default_namespace_path, - ) - .await?, - ); - self.queries - .insert(query_id, RegisteredQuery::Exchanged(query.clone())); - Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) - }) - .await - } - pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries @@ -586,14 +576,14 @@ impl SqlClientInner { #[derive(Clone)] enum RegisteredQuery { Polled(Arc), - Exchanged(Arc), + Streamed(Arc), } impl RegisteredQuery { fn touch(&self) { match self { Self::Polled(query) => query.touch(), - Self::Exchanged(query) => query.touch(), + Self::Streamed(query) => query.touch(), } } @@ -602,14 +592,14 @@ impl RegisteredQuery { fn registry_expired(&self) -> bool { match self { Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), - Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Streamed(query) => query.registry_expired(Arc::strong_count(query) == 1), } } async fn describe(&self) -> Result { match self { Self::Polled(query) => query.describe().await, - Self::Exchanged(query) => query.describe(), + Self::Streamed(query) => query.describe(), } } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/stream.rs similarity index 69% rename from rust/lancedb/src/remote/sql/exchange.rs rename to rust/lancedb/src/remote/sql/stream.rs index 523aaf5d7..ed85c6c92 100644 --- a/rust/lancedb/src/remote/sql/exchange.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -1,26 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! SQL statements sent with parameters. -//! -//! A statement without parameters is polled: `PollFlightInfo` starts it -//! detached from the connection and `DoGet` reads its endpoints. Neither call -//! carries Arrow data, so a statement's values would have to be spelled into -//! its text. A statement with parameters is sent on one `DoExchange` instead: -//! the command and the parameter schema, then one row, and the result rows -//! come back on the same call. The call is the statement's lifetime, so -//! cancelling it -- or dropping the handle before reading -- stops the -//! statement on the server. +//! SQL submitted on one call, with either streamed rows or a polling handoff. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::Instant; use arrow_array::RecordBatch; -use arrow_flight::FlightDescriptor; +use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::encode::FlightDataEncoderBuilder; -use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_flight::sql::{Any, CommandStatementQuery, ProstMessageExt}; +use arrow_flight::{FlightDescriptor, PollInfo, Ticket}; use arrow_schema::SchemaRef; +use futures::TryStreamExt; use prost::Message; use tokio::sync::{Notify, mpsc}; use uuid::Uuid; @@ -51,8 +44,8 @@ fn failure(error: &Error) -> Lifecycle { }) } -/// A statement running on its own `DoExchange` call. -pub(super) struct ExchangeQuery { +/// A statement streaming on its initial `DoGet` or `DoExchange` call. +pub(super) struct StreamingQuery { id: Uuid, client: Arc, /// The call's rows, until a reader takes them or a cancellation drops them. @@ -63,7 +56,7 @@ pub(super) struct ExchangeQuery { last_accessed: StdMutex, } -struct PreparedExchange { +struct PreparedStream { schema: SchemaRef, first: Option, rest: Option, @@ -76,7 +69,7 @@ struct PreparedExchange { /// this one did. A query cancelled meanwhile has its rows dropped instead, /// which ends the call. struct Borrowed<'a> { - query: &'a ExchangeQuery, + query: &'a StreamingQuery, rows: Option, } @@ -104,18 +97,37 @@ impl Drop for Borrowed<'_> { } } -impl ExchangeQuery { - /// Send the statement and its parameters, returning once the server has - /// planned it: a statement it refuses is refused here, as a polled one is - /// on its first poll. +pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol"; +pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1"; +const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse"; + +#[derive(Clone, PartialEq, Message)] +pub(super) struct QueryResponse { + #[prost(message, optional, tag = "1")] + pub poll_info: Option, +} + +#[derive(Clone, PartialEq, Message)] +struct CommandTicket { + #[prost(bytes = "vec", tag = "2")] + command: Vec, +} + +pub(super) enum Submission { + Stream(ResultEndpointStream), + Deferred(PollInfo), +} + +impl StreamingQuery { pub(super) async fn open( - id: Uuid, client: Arc, query: &str, - parameters: RecordBatch, + parameters: Option, default_namespace_path: &[String], - ) -> Result { - if parameters.num_rows() != 1 { + ) -> Result { + if let Some(parameters) = ¶meters + && parameters.num_rows() != 1 + { return Err(Error::InvalidInput { message: format!( "SQL query parameters must be exactly one row, got {}", @@ -130,39 +142,102 @@ impl ExchangeQuery { Some(DEFAULT_READ_TIMEOUT), )? .unwrap(); - let descriptor = FlightDescriptor::new_cmd( - CommandStatementQuery { - query: query.to_string(), - transaction_id: None, + let command = CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(); + let connection = client.connection(&request_id).await?; + let headers = client.headers(default_namespace_path, &request_id).await?; + let mut metadata = super::client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + let mut flight = connection.client.clone(); + let response = if let Some(parameters) = parameters { + let messages = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(FlightDescriptor::new_cmd(command))) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])) + .try_collect::>() + .await + .map_err(|error| sql_error(&request_id, error))?; + let mut request = tonic::Request::new(futures::stream::iter(messages)); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_exchange(request)).await + } else { + let mut request = + tonic::Request::new(Ticket::new(CommandTicket { command }.encode_to_vec())); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_get(request)).await + } + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|error| sql_error(&request_id, error))?; + let negotiated = response + .metadata() + .get(PROTOCOL_HEADER) + .is_some_and(|value| value == PROTOCOL_VERSION); + let mut response_stream = response.into_inner(); + if negotiated { + let first = tokio::time::timeout(read_timeout, response_stream.message()) + .await + .map_err(|_| sql_error(&request_id, "SQL query handoff timed out"))? + .map_err(|error| sql_error(&request_id, error))? + .ok_or_else(|| { + sql_error(&request_id, "SQL response omitted its execution decision") + })?; + if !first.data_header.is_empty() || !first.data_body.is_empty() { + return Err(sql_error( + &request_id, + "SQL response sent rows before its execution decision", + )); } - .as_any() - .encode_to_vec(), + let envelope = + Any::decode(first.app_metadata).map_err(|error| sql_error(&request_id, error))?; + if envelope.type_url != RESPONSE_TYPE_URL { + return Err(sql_error( + &request_id, + "SQL response returned an unknown execution decision", + )); + } + let decision = QueryResponse::decode(envelope.value) + .map_err(|error| sql_error(&request_id, error))?; + if let Some(info) = decision.poll_info { + if info.info.is_none() && info.flight_descriptor.is_none() { + return Err(sql_error( + &request_id, + "SQL handoff omitted result and continuation information", + )); + } + return Ok(Submission::Deferred(info)); + } + } + let stream = FlightRecordBatchStream::new_from_flight_data( + response_stream + .map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))), ); - // The descriptor rides on the schema message, and the row follows it. - let request = FlightDataEncoderBuilder::new() - .with_flight_descriptor(Some(descriptor)) - .with_schema(parameters.schema()) - .build(futures::stream::iter([Ok(parameters)])); - let mut flight = client - .client_with_headers(default_namespace_path, &request_id) - .await?; - let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) - .await - .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? - .map_err(|err| sql_error(&request_id, err))?; - Ok(Self { + Ok(Submission::Stream(ResultEndpointStream { + stream, + request_id, + read_timeout, + })) + } + + pub(super) fn new( + id: Uuid, + client: Arc, + response: ResultEndpointStream, + ) -> Self { + Self { id, client, - response: StdMutex::new(Some(ResultEndpointStream { - stream, - request_id, - read_timeout, - })), + response: StdMutex::new(Some(response)), lifecycle: StdMutex::new(Lifecycle::Running), cancelled: Notify::new(), terminal_at: OnceLock::new(), last_accessed: StdMutex::new(Instant::now()), - }) + } } pub(super) fn touch(&self) { @@ -242,7 +317,7 @@ impl ExchangeQuery { } /// Take the rows and wait for the first batch, which carries the schema. - async fn prepare(&self) -> Result { + async fn prepare(&self) -> Result { let Some(rows) = self.response.lock().unwrap().take() else { // Taken by an earlier attempt that failed, or dropped by a cancel. return Err(match &*self.lifecycle.lock().unwrap() { @@ -275,7 +350,7 @@ impl ExchangeQuery { if first.is_none() { self.settle(Lifecycle::Finished); } - Ok(PreparedExchange { + Ok(PreparedStream { schema, rest: first.is_some().then_some(rows), first, @@ -284,7 +359,7 @@ impl ExchangeQuery { async fn run( self: Arc, - mut prepared: PreparedExchange, + mut prepared: PreparedStream, sender: mpsc::Sender>, ) -> Result<()> { if let Some(batch) = prepared.first.take() @@ -343,13 +418,13 @@ impl ExchangeQuery { } } -pub(super) struct ExchangeQueryHandle { - query: Arc, +pub(super) struct StreamingQueryHandle { + query: Arc, result_started: AtomicBool, } -impl ExchangeQueryHandle { - pub(super) fn new(query: Arc) -> Self { +impl StreamingQueryHandle { + pub(super) fn new(query: Arc) -> Self { Self { query, result_started: AtomicBool::new(false), @@ -357,7 +432,7 @@ impl ExchangeQueryHandle { } } -impl Drop for ExchangeQueryHandle { +impl Drop for StreamingQueryHandle { // Rows nobody took are rows nobody will read. Keeping the call open would // hold the statement running on the server, stalled on flow control, // until the registry forgot it. @@ -369,7 +444,7 @@ impl Drop for ExchangeQueryHandle { } #[async_trait::async_trait] -impl QueryHandle for ExchangeQueryHandle { +impl QueryHandle for StreamingQueryHandle { fn id(&self) -> Uuid { self.query.touch(); self.query.id diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 79ca339db..4837baab7 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -19,6 +19,7 @@ use futures::{StreamExt, TryStreamExt}; use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; +use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; use crate::remote::RemoteCatalogOptions; @@ -70,6 +71,9 @@ type ReceivedExchanges = Arc)> #[derive(Clone)] struct TestSqlService { + adaptive: bool, + defer: bool, + invalid_decision: bool, query_count: Arc, do_get_count: Arc, cancel_count: Arc, @@ -120,6 +124,51 @@ impl TestSqlService { .unwrap_or_default(), }); } + fn adaptive_response( + &self, + query: &str, + rows: BoxStream<'static, std::result::Result>, + ) -> Response>> { + if !self.adaptive { + return Response::new(rows); + } + let poll_info = + self.defer.then(|| PollInfo { + info: Some(FlightInfo::new().with_endpoint( + FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())), + )), + flight_descriptor: Some(FlightDescriptor::new_cmd(format!("poll:1:{query}"))), + progress: Some(0.25), + expiration_time: None, + }); + let first = FlightData { + app_metadata: if self.invalid_decision { + b"invalid protobuf".to_vec().into() + } else { + Any { + type_url: "type.googleapis.com/lancedb.sql.QueryResponse".to_string(), + value: QueryResponse { poll_info }.encode_to_vec().into(), + } + .encode_to_vec() + .into() + }, + ..Default::default() + }; + let rest = if self.defer { + futures::stream::empty().boxed() + } else { + rows + }; + let mut response = Response::new( + futures::stream::once(async { Ok(first) }) + .chain(rest) + .boxed(), + ); + response + .metadata_mut() + .insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + response + } } impl Default for TestSqlService { @@ -154,6 +203,9 @@ impl Default for TestSqlService { let dictionary_result = RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap(); Self { + adaptive: false, + defer: false, + invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), @@ -308,7 +360,22 @@ impl FlightService for TestSqlService { request: Request, ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); - let ticket = request.get_ref().ticket.as_ref(); + let initial = Any::decode(request.get_ref().ticket.clone()) + .ok() + .and_then(|envelope| Any::decode(envelope.value).ok()) + .and_then(|command| command.unpack::().ok().flatten()); + if initial.is_some() { + self.capture_headers(request.metadata()); + assert_eq!( + request.metadata().get(PROTOCOL_HEADER).unwrap(), + PROTOCOL_VERSION + ); + } + let query = initial + .as_ref() + .map(|command| command.query.as_bytes()) + .unwrap_or(request.get_ref().ticket.as_ref()); + let ticket = query; let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; @@ -337,7 +404,11 @@ impl FlightService for TestSqlService { encoder = encoder.with_max_flight_data_size(8 * 1024 * 1024); } let stream = encoder.build(input).map_err(Status::from); - Ok(Response::new(Box::pin(stream))) + if let Some(initial) = initial { + Ok(self.adaptive_response(&initial.query, stream.boxed())) + } else { + Ok(Response::new(stream.boxed())) + } } async fn do_put( @@ -486,7 +557,7 @@ impl FlightService for TestSqlService { let _dropped = &dropped; message }); - Ok(Response::new(Box::pin(stream))) + Ok(self.adaptive_response(&query, stream.boxed())) } } @@ -1609,3 +1680,138 @@ fn validates_metadata_with_header_map() { assert!(insert_header(&mut headers, "bad header", "value").is_err()); assert!(insert_header(&mut headers, "valid-header", "bad\nvalue").is_err()); } + +#[tokio::test] +async fn adaptive_streams_start_on_one_rpc_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = parameters.clone().unwrap_or_else(|| service.result.clone()); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT echo") + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(gets.load(Ordering::SeqCst), usize::from(!with_parameters)); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_share_polling_and_survive_handle_drop() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = service.result.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT adaptive") + }) + .await + .unwrap(); + let id = query.id(); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + drop(query); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "the original SQL must not be resubmitted through polling" + ); + assert_eq!( + gets.load(Ordering::SeqCst), + 1 + usize::from(!with_parameters) + ); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_use_server_cancellation_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let cancels = service.cancel_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT cancel") + }) + .await + .unwrap(); + query.cancel().await.unwrap(); + assert_eq!(cancels.load(Ordering::SeqCst), 1); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + assert!(query.reader().await.is_err()); + } +} + +#[tokio::test] +async fn malformed_adaptive_decision_never_resubmits_the_statement() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + invalid_decision: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + assert!( + client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT invalid") + }) + .await + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + } +} diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index c2ef13ba5..fee2d5d79 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -18,9 +18,9 @@ pub enum QueryStatus { Running, /// The server has made the complete result available. /// - /// A query submitted with parameters streams its rows on the call that - /// runs it, so it reports `Finished` only once its reader has received - /// the last batch. Read it rather than waiting for this status. + /// Directly streamed queries report `Finished` when their stream reaches + /// its end. Deferred queries report the server's execution state. Open the + /// reader directly instead of waiting for this status before reading. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -64,9 +64,10 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query, unless the query was -/// submitted with parameters: those run on the call that returns their rows, -/// so dropping the handle before reading, or dropping the reader, stops them. +/// The server may stream results on the initial call or defer execution and +/// return result tickets. Dropping an unread handle or its reader stops a +/// directly streamed query; deferred execution continues until explicitly +/// cancelled. This applies equally with and without parameters. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -95,10 +96,9 @@ impl Query { /// Get a point-in-time description of the query. /// - /// For a query submitted with parameters the status follows its result - /// stream: it stays [`QueryStatus::Running`] until the reader has received - /// every row, so open [`Self::reader`] directly instead of polling for - /// [`QueryStatus::Finished`] first. + /// Directly streamed queries stay [`QueryStatus::Running`] until the result + /// stream ends. Deferred queries expose the server's execution state. Open + /// [`Self::reader`] directly instead of waiting for [`QueryStatus::Finished`]. pub async fn describe(&self) -> Result { self.handle.describe().await }