From efd19625b410bdeae286f422ee25bf46ffa887c8 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Wed, 16 Sep 2026 12:15:40 +0800 Subject: [PATCH] refactor(servers): simplify Flight coalescing with append-then-check accumulator Remove the lookahead/pending machinery and coalesce with an append-then-check soft budget. Extract the group state into a private BatchAccumulator so the ordinary stream loop reads as a short control flow: append each ready batch, flush when a budget is reached. Budgets are flush thresholds, not memory limits, so a group may exceed a budget by the final appended batch. First-batch direct send and singleton passthrough for a batch that starts a group already at a budget are unchanged. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/servers/src/grpc/flight/stream.rs | 227 +++++++++++++++++++------- 1 file changed, 165 insertions(+), 62 deletions(-) diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index f2a3dd3be67..8fd84336b50 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use arrow_flight::FlightData; use common_error::ext::ErrorExt; use common_grpc::flight::{FlightEncoder, FlightMessage}; -use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::recordbatch::merge_record_batches; +use common_recordbatch::{RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing::{Instrument, info_span}; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{error, info, warn}; @@ -129,6 +129,60 @@ impl Drop for StreamMetrics { } } +/// Coalesces consecutive ready record batches into one outgoing group. +/// +/// The budgets are soft limits (flush thresholds) rather than strict admission +/// caps: `push` appends a batch before reporting whether to flush, so a group may +/// exceed a budget by up to one batch. +struct BatchAccumulator { + batches: Vec, + rows: usize, + bytes: usize, +} + +impl BatchAccumulator { + const MAX_ROWS: usize = 1024; + const MAX_BYTES: usize = 256 * 1024; + const MAX_BATCHES: usize = 16; + + fn new() -> Self { + Self { + batches: Vec::new(), + rows: 0, + bytes: 0, + } + } + + /// Whether a batch fetched to start a new group already reaches a budget, in + /// which case the stream forwards it as a singleton instead of coalescing it. + /// A batch appended inside a group is coalesced even if it exceeds a budget. + fn reaches_budget(rows: usize, bytes: usize) -> bool { + rows >= Self::MAX_ROWS || bytes >= Self::MAX_BYTES + } + + fn is_empty(&self) -> bool { + self.batches.is_empty() + } + + /// Appends `batch` and reports whether the group should be flushed now. + fn push(&mut self, batch: RecordBatch) -> bool { + self.rows += batch.num_rows(); + self.bytes += batch.df_record_batch().get_array_memory_size(); + self.batches.push(batch); + // Append-then-check: the group may exceed a budget by this batch. + self.batches.len() >= Self::MAX_BATCHES || Self::reaches_budget(self.rows, self.bytes) + } + + /// Removes and returns the accumulated batches in arrival order. + /// + /// Merging is left to the caller, which owns the stream schema. + fn drain(&mut self) -> Vec { + self.rows = 0; + self.bytes = 0; + std::mem::take(&mut self.batches) + } +} + #[pin_project(PinnedDrop)] pub struct FlightRecordBatchStream { #[pin] @@ -258,10 +312,6 @@ impl FlightRecordBatchStream { should_send_partial_metrics: bool, can_send_metrics_before_batch: bool, ) { - const MAX_ROWS: usize = 1024; - const MAX_BYTES: usize = 256 * 1024; - const MAX_BATCHES: usize = 16; - let mut metrics = StreamMetrics::new(should_send_partial_metrics); let mut last_metrics_str = None; let recordbatch_schema = recordbatches.schema(); @@ -351,19 +401,13 @@ impl FlightRecordBatchStream { } } else { let mut sent_first_batch = false; - let mut pending = None; + let mut acc = BatchAccumulator::new(); loop { - let (batch_or_err, fetched) = match pending.take() { - Some(recordbatch) => (Ok(recordbatch), false), - None => { - let start = Instant::now(); - let batch_or_err = recordbatches.next().in_current_span().await; - metrics.fetch_content_duration += start.elapsed(); - let Some(batch_or_err) = batch_or_err else { - break; - }; - (batch_or_err, true) - } + let start = Instant::now(); + let batch_or_err = recordbatches.next().in_current_span().await; + metrics.fetch_content_duration += start.elapsed(); + let Some(batch_or_err) = batch_or_err else { + break; }; let recordbatch = match batch_or_err { Ok(recordbatch) => recordbatch, @@ -380,13 +424,14 @@ impl FlightRecordBatchStream { }; let batch_rows = recordbatch.num_rows(); let batch_bytes = recordbatch.df_record_batch().get_array_memory_size(); - if fetched { - metrics.total_rows += batch_rows; - metrics.record_batch_count += 1; - metrics.total_bytes += batch_bytes; - } + metrics.total_rows += batch_rows; + metrics.record_batch_count += 1; + metrics.total_bytes += batch_bytes; - if !sent_first_batch || batch_rows >= MAX_ROWS || batch_bytes >= MAX_BYTES { + // The first batch is forwarded immediately, and a batch starting a + // new group that is already at or over a budget passes through as a + // singleton. A batch appended inside a group is coalesced instead. + if !sent_first_batch || BatchAccumulator::reaches_budget(batch_rows, batch_bytes) { let start = Instant::now(); if let Err(e) = tx .send(Ok(FlightMessage::RecordBatch( @@ -402,36 +447,24 @@ impl FlightRecordBatchStream { continue; } - let mut batches = vec![recordbatch]; - let mut rows = batch_rows; - let mut bytes = batch_bytes; + // Coalesce ready batches until a budget is reached: every ready + // batch is appended first, then the budgets are checked. + debug_assert!(acc.is_empty(), "the previous group must be flushed"); + let mut should_flush = acc.push(recordbatch); let mut eof = false; let mut stream_error = None; - loop { - if batches.len() == MAX_BATCHES || rows >= MAX_ROWS || bytes >= MAX_BYTES { - break; - } + while !should_flush { let start = Instant::now(); let next = poll_fn(|cx| Poll::Ready(recordbatches.as_mut().poll_next(cx))).await; metrics.fetch_content_duration += start.elapsed(); match next { Poll::Ready(Some(Ok(recordbatch))) => { - let batch_rows = recordbatch.num_rows(); - let batch_bytes = recordbatch.df_record_batch().get_array_memory_size(); - if batch_rows > MAX_ROWS - rows || batch_bytes > MAX_BYTES - bytes { - metrics.total_rows += batch_rows; - metrics.record_batch_count += 1; - metrics.total_bytes += batch_bytes; - pending = Some(recordbatch); - break; - } - metrics.total_rows += batch_rows; + metrics.total_rows += recordbatch.num_rows(); metrics.record_batch_count += 1; - metrics.total_bytes += batch_bytes; - rows += batch_rows; - bytes += batch_bytes; - batches.push(recordbatch); + metrics.total_bytes += + recordbatch.df_record_batch().get_array_memory_size(); + should_flush = acc.push(recordbatch); } Poll::Ready(Some(Err(e))) => { stream_error = Some(e); @@ -445,6 +478,7 @@ impl FlightRecordBatchStream { } } + let batches = acc.drain(); let batches = if batches.len() >= 2 { match merge_record_batches(recordbatch_schema.clone(), &batches) { Ok(merged) => vec![merged], @@ -903,7 +937,7 @@ mod test { } #[tokio::test] - async fn test_pending_batch_is_sent_before_upstream_is_polled_again() { + async fn test_over_cap_batch_merges_into_current_aggregate() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "a", ConcreteDataType::int32_datatype(), @@ -940,23 +974,84 @@ mod test { FlightMessage::RecordBatch(_) )); tokio::time::timeout(Duration::from_secs(1), async { - while poll_count.load(std::sync::atomic::Ordering::Relaxed) != 3 { + while poll_count.load(std::sync::atomic::Ordering::Relaxed) < 3 { tokio::task::yield_now().await; } }) .await - .expect("second group must fetch the held batch"); - let held = stream.next().await.unwrap().unwrap(); - assert!(matches!( - decoder.try_decode(&held).unwrap().unwrap(), - FlightMessage::RecordBatch(_) - )); - let pending = stream.next().await.unwrap().unwrap(); - let FlightMessage::RecordBatch(batch) = decoder.try_decode(&pending).unwrap().unwrap() + .expect("the second group must fetch the at-cap batch"); + let aggregated = stream.next().await.unwrap().unwrap(); + let FlightMessage::RecordBatch(batch) = decoder.try_decode(&aggregated).unwrap().unwrap() else { - panic!("expected the held oversized batch"); + panic!("expected the oversized aggregate"); }; - assert_eq!(batch.num_rows(), 1024); + // The at-cap batch is appended before the budget is checked, so the aggregate + // exceeds the row budget by one batch instead of being held back as a + // separate singleton. + assert_eq!(batch.num_rows(), 1025); + } + + #[tokio::test] + async fn test_ready_only_soft_row_budget_flushes_oversized_group() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let messages = flight_messages(Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..600))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..600))), + ]), + poll_count: poll_count.clone(), + })) + .await; + let batches = messages + .iter() + .filter_map(|message| match message.as_ref().unwrap() { + FlightMessage::RecordBatch(batch) => Some(batch.num_rows()), + _ => None, + }) + .collect::>(); + // 600 + 600 exceeds the 1024-row budget, but the second batch is appended + // before the budget is checked, so the group flushes as one 1200-row batch. + assert_eq!(batches, vec![1, 1200]); + // Three fetches, plus the poll that reports end of stream. + assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 4); + } + + #[tokio::test] + async fn test_ready_only_current_at_cap_batch_is_sent_as_singleton() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let messages = flight_messages(Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..1024))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))), + ]), + poll_count: poll_count.clone(), + })) + .await; + let batches = messages + .iter() + .filter_map(|message| match message.as_ref().unwrap() { + FlightMessage::RecordBatch(batch) => Some(batch.num_rows()), + _ => None, + }) + .collect::>(); + // The at-cap batch is current when it is fetched, so it passes through as a + // singleton instead of joining an aggregate. + assert_eq!(batches, vec![1, 1024, 1]); + assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 4); } #[tokio::test] @@ -1018,7 +1113,7 @@ mod test { } #[tokio::test] - async fn test_ready_only_preserves_oversized_and_empty_batch_caps() { + async fn test_ready_only_merges_at_cap_batch_and_coalesces_empty_batches() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "a", ConcreteDataType::int32_datatype(), @@ -1030,8 +1125,11 @@ mod test { ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))), ScriptedItem::Batch(Ok(oversized)), ]; + // 17 trailing empty batches: the 16-batch budget flushes the first 16 as + // one group, leaving one empty batch for a second group. This exercises the + // batch-count bound rather than relying on EOF to flush. items.extend( - (0..16).map(|_| ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone())))), + (0..17).map(|_| ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone())))), ); let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let messages = flight_messages(Box::pin(ScriptedBatchStream { @@ -1047,12 +1145,15 @@ mod test { _ => None, }) .collect::>(); - assert_eq!(batches, vec![1, 1, 1024, 0]); - assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 20); + // The at-cap batch is appended to the current aggregate, so the flushed + // aggregate exceeds the row budget by one batch. The 17 trailing empty + // batches split into two empty groups at the 16-batch budget. + assert_eq!(batches, vec![1, 1025, 0, 0]); + assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 21); } #[tokio::test] - async fn test_ready_only_preserves_byte_oversized_batch() { + async fn test_ready_only_merges_byte_oversized_batch_into_aggregate() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "a", ConcreteDataType::string_datatype(), @@ -1082,7 +1183,9 @@ mod test { _ => None, }) .collect::>(); - assert_eq!(batches, vec![0, 0, 1023]); + // The byte-oversized batch is appended to the leading empty batch, so the + // aggregate exceeds the byte budget by one batch. + assert_eq!(batches, vec![0, 1023]); } #[tokio::test]