From a2964a81e06ed02ee2d17c44074d7d1ae971e3d6 Mon Sep 17 00:00:00 2001 From: Matthew Agostinelli Date: Wed, 26 Nov 2025 11:11:50 -0600 Subject: [PATCH] Fix FlightSQL client to display all batches from multi-partition queries Previously, the FlightSQL client would only display partial results when queries returned multiple batches (common with GROUP BY without ORDER BY due to partitioned execution). This happened because: 1. `run_flightsqls()` only fetched one batch using `match streams.next().await` instead of iterating through all batches 2. `take_record_batches()` only used the first batch when multiple batches existed, ignoring the rest This fix: - Changes `run_flightsqls()` to use `while let` loop to collect all batches from the FlightSQL stream - Updates `take_record_batches()` to concatenate all batches before applying pagination indices, ensuring all rows are visible Known limitations for future consideration: - All batches are collected eagerly into memory before display, which could be problematic for very large result sets - `concat_batches()` creates a temporary copy when concatenating, briefly doubling memory usage during pagination - Unlike the SQL tab which uses lazy batch loading, FlightSQL now loads all data upfront (this matches expected behavior for aggregation queries but may warrant lazy loading for large streaming results in the future) Fixes #346 --- src/tui/execution.rs | 65 ++++++++++++++++++--------------- src/tui/state/tabs/flightsql.rs | 10 +++-- 2 files changed, 42 insertions(+), 33 deletions(-) diff --git a/src/tui/execution.rs b/src/tui/execution.rs index 2cf34b28..08a4d11a 100644 --- a/src/tui/execution.rs +++ b/src/tui/execution.rs @@ -279,37 +279,42 @@ impl TuiExecution { if let Some(streams) = self.flightsql_result_stream.lock().await.as_mut() { - match streams.next().await { - Some((ticket, Ok(batch))) => { - info!("Received batch for {ticket}"); - let duration = start.elapsed(); - let results = ExecutionResultsBatch { - batch, - duration, - query: sql.to_string(), - }; - sender.send( - AppEvent::FlightSQLExecutionResultsNextBatch( - results, - ), - )?; + // Collect all batches from the stream + while let Some((ticket, result)) = + streams.next().await + { + match result { + Ok(batch) => { + info!("Received batch for {ticket}"); + let duration = start.elapsed(); + let results = ExecutionResultsBatch { + batch, + duration, + query: sql.to_string(), + }; + sender.send( + AppEvent::FlightSQLExecutionResultsNextBatch( + results, + ), + )?; + } + Err(e) => { + error!( + "Error executing stream for ticket {ticket}: {:?}", + e + ); + let elapsed = start.elapsed(); + let e = ExecutionError { + query: sql.to_string(), + error: e.to_string(), + duration: elapsed, + }; + sender.send( + AppEvent::FlightSQLExecutionResultsError(e), + )?; + break; + } } - Some((ticket, Err(e))) => { - error!( - "Error executing stream for ticket {ticket}: {:?}", - e - ); - let elapsed = start.elapsed(); - let e = ExecutionError { - query: sql.to_string(), - error: e.to_string(), - duration: elapsed, - }; - sender.send( - AppEvent::FlightSQLExecutionResultsError(e), - )?; - } - None => {} } } } diff --git a/src/tui/state/tabs/flightsql.rs b/src/tui/state/tabs/flightsql.rs index c3389fa1..31a0f2db 100644 --- a/src/tui/state/tabs/flightsql.rs +++ b/src/tui/state/tabs/flightsql.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use color_eyre::Result; use datafusion::arrow::{ array::{Array, RecordBatch, UInt32Array}, - compute::take_record_batch, + compute::{concat_batches, take_record_batch}, datatypes::Schema, error::ArrowError, }; @@ -316,7 +316,11 @@ fn take_record_batches( match batches.len() { 0 => Ok(RecordBatch::new_empty(Arc::new(Schema::empty()))), 1 => take_record_batch(&batches[0], indices), - // For now we just get the first batch - _ => take_record_batch(&batches[0], indices), + _ => { + // Concatenate all batches into a single batch before taking indices + let schema = batches[0].schema(); + let concatenated = concat_batches(&schema, batches)?; + take_record_batch(&concatenated, indices) + } } }