diff --git a/python/python/lancedb/sql.py b/python/python/lancedb/sql.py index 7d1ca44e6..05b82b0d3 100644 --- a/python/python/lancedb/sql.py +++ b/python/python/lancedb/sql.py @@ -35,7 +35,14 @@ def _parameter_array(value: Any) -> pa.Array: # floats that would come back as float64. if isinstance(value, pa.ChunkedArray): value = value.combine_chunks() - if isinstance(value, np.ndarray) and value.ndim == 1: + if isinstance(value, np.ndarray): + if value.ndim == 0: + return pa.array(value.reshape(1)) + if value.ndim != 1: + raise ValueError( + "a numpy query parameter must be a scalar or a 1-D vector, " + f"got {value.ndim} dimensions" + ) value = pa.array(value) if isinstance(value, pa.Array): return pa.FixedSizeListArray.from_arrays(value, len(value)) @@ -54,16 +61,16 @@ def to_parameter_batch( """ if parameters is None: return None - if isinstance(parameters, pa.Table): - parameters = parameters.combine_chunks() - batches = parameters.to_batches() - parameters = ( + batch: pa.RecordBatch + if isinstance(parameters, pa.RecordBatch): + batch = parameters + elif isinstance(parameters, pa.Table): + batches = parameters.combine_chunks().to_batches() + batch = ( batches[0] if len(batches) == 1 else pa.RecordBatch.from_pylist([], schema=parameters.schema) ) - if isinstance(parameters, pa.RecordBatch): - batch = parameters elif isinstance(parameters, Mapping): batch = pa.RecordBatch.from_arrays( [_parameter_array(value) for value in parameters.values()], diff --git a/rust/lancedb/src/remote/sql_exchange.rs b/rust/lancedb/src/remote/sql_exchange.rs index d5c45a39c..062223584 100644 --- a/rust/lancedb/src/remote/sql_exchange.rs +++ b/rust/lancedb/src/remote/sql_exchange.rs @@ -42,6 +42,15 @@ enum Lifecycle { Failed(String), } +/// A failed statement, keeping the message `describe` reports it with -- not +/// `Runtime error: …` again when `describe` wraps it in a runtime error. +fn failure(error: &Error) -> Lifecycle { + Lifecycle::Failed(match error { + Error::Runtime { message } => message.clone(), + other => other.to_string(), + }) +} + /// A statement running on its own `DoExchange` call. pub(super) struct ExchangeQuery { id: Uuid, @@ -209,14 +218,7 @@ impl ExchangeQuery { let first = tokio::select! { biased; _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), - result = rows.next_batch() => result, - }; - let first = match first { - Ok(first) => first, - Err(error) => { - self.settle(Lifecycle::Failed(error.to_string())); - return Err(error); - } + result = rows.next_batch() => result?, }; let schema = first .as_ref() @@ -271,7 +273,7 @@ impl ExchangeQuery { return Ok(()); } Err(error) => { - self.settle(Lifecycle::Failed(error.to_string())); + self.settle(failure(&error)); return Err(error); } } @@ -347,7 +349,15 @@ impl QueryHandle for ExchangeQueryHandle { self.query.touch(); let started = Instant::now(); let prepared = - with_overall_timeout(timeout, "SQL query result", self.query.prepare()).await?; + match with_overall_timeout(timeout, "SQL query result", self.query.prepare()).await { + Ok(prepared) => prepared, + Err(error) => { + // The rows were taken before the wait, so whatever ended + // it -- a failed read or the timeout -- the call is gone. + self.query.settle(failure(&error)); + return Err(error); + } + }; let remaining_timeout = timeout.map(|timeout| timeout.saturating_sub(started.elapsed())); let schema = prepared.schema.clone(); let (sender, receiver) = mpsc::channel(2); @@ -361,7 +371,7 @@ impl QueryHandle for ExchangeQueryHandle { ) .await; if let Err(error) = result { - query.settle(Lifecycle::Failed(error.to_string())); + query.settle(failure(&error)); let _ = error_sender.send(Err(error)).await; } }); diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 1b3364280..d2b24fd85 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -65,6 +65,9 @@ struct CapturedHeaders { database_prefix: String, } +/// Every exchange's statement and the parameter row it arrived with. +type ReceivedExchanges = Arc)>>>; + #[derive(Clone)] struct TestSqlService { query_count: Arc, @@ -79,8 +82,7 @@ struct TestSqlService { transient_poll_failures: Arc, headers: Arc>>, gate_refresh_release: Arc, - /// Every exchange's statement and the parameter row it arrived with. - exchanges: Arc)>>>, + exchanges: ReceivedExchanges, /// How many exchange responses have been dropped, finished or not. exchange_drops: Arc, result: RecordBatch, @@ -452,6 +454,9 @@ impl FlightService for TestSqlService { "parameter 1 is not referenced by the statement", )); } + if query == "SELECT exchange stall" { + return Ok(Response::new(futures::stream::pending().boxed())); + } let rows = parameters.ok_or_else(|| Status::invalid_argument("exchange without parameters"))?; let schema = rows.schema(); @@ -495,6 +500,7 @@ fn parameterized(query: &str, parameters: RecordBatch) -> ExecuteQueryRequest { async fn start_sql_service( service: TestSqlService, + client_config: ClientConfig, ) -> (SqlClient, tokio::sync::oneshot::Sender<()>) { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); @@ -516,7 +522,7 @@ async fn start_sql_service( "test-key".to_string(), None, Some(format!("grpc://{address}")), - ClientConfig::default(), + client_config, ); (client, shutdown_tx) } @@ -540,7 +546,7 @@ async fn parameters_travel_on_one_exchange() { let query_count = service.query_count.clone(); let do_get_count = service.do_get_count.clone(); let headers = service.headers.clone(); - let (client, _shutdown) = start_sql_service(service).await; + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; let parameters = parameter_row(); let query = client @@ -597,7 +603,8 @@ async fn parameters_travel_on_one_exchange() { /// The server plans before it answers, so a refusal is the submission's. #[tokio::test] async fn a_refused_statement_fails_its_submission() { - let (client, _shutdown) = start_sql_service(TestSqlService::default()).await; + let (client, _shutdown) = + start_sql_service(TestSqlService::default(), ClientConfig::default()).await; let refused = client .execute(parameterized("SELECT exchange refused", parameter_row())) .await @@ -624,7 +631,7 @@ async fn a_refused_statement_fails_its_submission() { async fn cancelling_or_abandoning_an_exchange_ends_the_call() { let service = TestSqlService::default(); let drops = service.exchange_drops.clone(); - let (client, _shutdown) = start_sql_service(service).await; + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; let reading = client .execute(parameterized("SELECT exchange slow", parameter_row())) @@ -667,6 +674,35 @@ async fn cancelling_or_abandoning_an_exchange_ends_the_call() { wait_for(&drops, 3, "dropping a reader must end the call").await; } +/// A reader that gives up before the first batch has already taken the rows, +/// so the query must say it failed rather than report `Running` for a call +/// that is gone. +#[tokio::test] +async fn a_reader_that_times_out_reports_the_failure() { + let mut client_config = ClientConfig::default(); + client_config.timeout_config.timeout = Some(Duration::from_millis(200)); + let (client, _shutdown) = start_sql_service(TestSqlService::default(), client_config).await; + + let stalled = client + .execute(parameterized("SELECT exchange stall", parameter_row())) + .await + .unwrap(); + assert_overall_timeout(stalled.reader().await, "result"); + let Err(retried) = stalled.reader().await else { + panic!("the rows are gone, so a second reader cannot succeed"); + }; + for error in [ + stalled.describe().await.unwrap_err(), + client.describe(stalled.id()).await.unwrap_err(), + retried, + ] { + assert!( + matches!(&error, Error::Runtime { message } if message == "SQL query result timed out"), + "{error}" + ); + } +} + #[tokio::test] async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();