fix(sql): settle a parameterized query whose reader fails before its first batch

A reader that times out or fails before the first batch has already taken the call's rows; the query now reports that failure instead of Running. Also type-checks the numpy parameter conversion and keeps describe from double-prefixing runtime errors.

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()],

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;
             }
         });

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<std::sync::Mutex<Vec<(String, Option<RecordBatch>)>>>;
+
 #[derive(Clone)]
 struct TestSqlService {
     query_count: Arc<AtomicUsize>,
@@ -79,8 +82,7 @@ struct TestSqlService {
     transient_poll_failures: Arc<AtomicUsize>,
     headers: Arc<std::sync::Mutex<Vec<CapturedHeaders>>>,
     gate_refresh_release: Arc<Notify>,
-    /// Every exchange's statement and the parameter row it arrived with.
-    exchanges: Arc<std::sync::Mutex<Vec<(String, Option<RecordBatch>)>>>,
+    exchanges: ReceivedExchanges,
     /// How many exchange responses have been dropped, finished or not.
     exchange_drops: Arc<AtomicUsize>,
     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();