diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 8303a40fd9..a27ed090d9 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -805,8 +805,8 @@ pub struct Grpc { pub internal_grpc_token: String, #[env_config( name = "ZO_GRPC_MAX_MESSAGE_SIZE", - default = 16, - help = "Max grpc message size in MB, default is 16 MB" + default = 32, + help = "Max grpc message size in MB, default is 32 MB" )] pub max_message_size: usize, #[env_config(name = "ZO_GRPC_CONNECT_TIMEOUT", default = 5)] // in seconds diff --git a/src/flight/src/encoder.rs b/src/flight/src/encoder.rs index aec11e3c14..5813806675 100644 --- a/src/flight/src/encoder.rs +++ b/src/flight/src/encoder.rs @@ -13,10 +13,9 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::sync::Arc; - use arrow::{ - array::{Array, ArrayRef, AsArray, RecordBatch, RecordBatchOptions, StringViewBuilder}, + array::RecordBatch, + compute::BatchCoalescer, ipc::{ MessageHeader, writer::{CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions}, @@ -34,8 +33,12 @@ pub struct FlightDataEncoder { options: IpcWriteOptions, data_gen: IpcDataGenerator, dictionary_tracker: DictionaryTracker, - max_flight_data_size: usize, compression_context: CompressionContext, + /// Lazily initialised on the first `encode_batch` call; reused across calls + /// to avoid re-allocating per-column InProgressArray builders every time. + coalescer: Option, + /// Maximum bytes per FlightData message (0 = unlimited). + max_flight_data_size: usize, } impl FlightDataEncoder { @@ -44,8 +47,9 @@ impl FlightDataEncoder { options, data_gen: IpcDataGenerator::default(), dictionary_tracker: DictionaryTracker::new(false), - max_flight_data_size, compression_context: CompressionContext::default(), + coalescer: None, + max_flight_data_size, } } @@ -55,25 +59,57 @@ impl FlightDataEncoder { } /// Encode a RecordBatch to a Vec of FlightData + /// + /// Rows are accumulated inside the coalescer until a full 8192-row chunk is + /// ready. Call [`Self::finish`] when the input stream ends to flush the + /// remaining partial chunk. + /// + /// BatchCoalescer handles both coalescing/splitting and GC: + /// - top-level Utf8View/BinaryView: copies only the string bytes needed per chunk when the + /// source array is sparse (actual_buffer > 2x used) + /// - nested Utf8View inside Struct/List: materialised via concat, so stale backing buffers are + /// released naturally + /// - + /// - pub fn encode_batch(&mut self, batch: RecordBatch) -> Result> { - // refer to: - // https://github.com/openobserve/openobserve/issues/8280 - // https://github.com/apache/datafusion/pull/11587 - let batch = gc_string_view_batch(&batch); - - let mut flight_data = Vec::new(); - for batch in split_batch_for_grpc_response(batch, self.max_flight_data_size) { - let (encoded_dictionaries, encoded_batch) = self.data_gen.encode( - &batch, - &mut self.dictionary_tracker, - &self.options, - &mut self.compression_context, - )?; - - flight_data.extend(encoded_dictionaries.into_iter().map(Into::into)); - flight_data.push(encoded_batch.into()); + if self.coalescer.is_none() { + self.coalescer = Some(BatchCoalescer::new(batch.schema(), 8192)); } + let coalescer = self.coalescer.as_mut().unwrap(); + coalescer.push_batch(batch)?; + + self.drain_completed() + } + + /// Flush any buffered rows and encode them as FlightData. + /// + /// Call this once after all batches have been pushed via [`Self::encode_batch`]. + pub fn finish(&mut self) -> Result> { + let Some(coalescer) = self.coalescer.as_mut() else { + return Ok(vec![]); + }; + coalescer.finish_buffered_batch()?; + self.drain_completed() + } + + fn drain_completed(&mut self) -> Result> { + let Some(coalescer) = self.coalescer.as_mut() else { + return Ok(vec![]); + }; + let mut flight_data = Vec::new(); + while let Some(batch) = coalescer.next_completed_batch() { + for batch in split_batch_for_grpc_response(batch, self.max_flight_data_size) { + let (encoded_dictionaries, encoded_batch) = self.data_gen.encode( + &batch, + &mut self.dictionary_tracker, + &self.options, + &mut self.compression_context, + )?; + flight_data.extend(encoded_dictionaries.into_iter().map(Into::into)); + flight_data.push(encoded_batch.into()); + } + } Ok(flight_data) } @@ -89,6 +125,40 @@ impl FlightDataEncoder { } } +fn split_batch_for_grpc_response( + batch: RecordBatch, + max_flight_data_size: usize, +) -> Vec { + if max_flight_data_size == 0 { + return vec![batch]; + } + + let size = batch + .columns() + .iter() + .map(|col| { + col.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| col.get_buffer_memory_size()) + }) + .sum::(); + + let n_batches = (size / max_flight_data_size + + usize::from(!size.is_multiple_of(max_flight_data_size))) + .max(1); + let rows_per_batch = (batch.num_rows() / n_batches).max(1); + let mut out = Vec::with_capacity(n_batches + 1); + + let mut offset = 0; + while offset < batch.num_rows() { + let length = rows_per_batch.min(batch.num_rows() - offset); + out.push(batch.slice(offset, length)); + offset += length; + } + + out +} + pub fn header_none() -> Bytes { let mut builder: FlatBufferBuilder<'_> = FlatBufferBuilder::new(); @@ -103,122 +173,18 @@ pub fn header_none() -> Bytes { builder.finished_data().to_vec().into() } -fn split_batch_for_grpc_response( - batch: RecordBatch, - max_flight_data_size: usize, -) -> Vec { - let size = batch - .columns() - .iter() - .map(|col| col.get_buffer_memory_size()) - .sum::(); - - let n_batches = (size / max_flight_data_size - + usize::from(!size.is_multiple_of(max_flight_data_size))) - .max(1); - let rows_per_batch = (batch.num_rows() / n_batches).max(1); - let mut out = Vec::with_capacity(n_batches + 1); - - let mut offset = 0; - while offset < batch.num_rows() { - let length = (rows_per_batch).min(batch.num_rows() - offset); - out.push(batch.slice(offset, length)); - - offset += length; - } - - out -} - -/// refer to: https://github.com/apache/datafusion/pull/11587 -/// Heuristically compact `StringViewArray`s to reduce memory usage, if needed -/// -/// Decides when to consolidate the StringView into a new buffer to reduce -/// memory usage and improve string locality for better performance. -/// -/// This differs from `StringViewArray::gc` because: -/// 1. It may not compact the array depending on a heuristic. -/// 2. It uses a precise block size to reduce the number of buffers to track. -/// -/// # Heuristic -/// -/// If the average size of each view is larger than 32 bytes, we compact the array. -/// -/// `StringViewArray` include pointers to buffer that hold the underlying data. -/// One of the great benefits of `StringViewArray` is that many operations -/// (e.g., `filter`) can be done without copying the underlying data. -/// -/// However, after a while (e.g., after `FilterExec` or `HashJoinExec`) the -/// `StringViewArray` may only refer to a small portion of the buffer, -/// significantly increasing memory usage. -fn gc_string_view_batch(batch: &RecordBatch) -> RecordBatch { - let new_columns: Vec = batch - .columns() - .iter() - .map(|c| { - // Try to re-create the `StringViewArray` to prevent holding the underlying buffer too - // long. - let Some(s) = c.as_string_view_opt() else { - return Arc::clone(c); - }; - - // Fast path: if the data buffers are empty, we can return the original array - if s.data_buffers().is_empty() { - return Arc::clone(c); - } - - let ideal_buffer_size: usize = s - .views() - .iter() - .map(|v| { - let len = (*v as u32) as usize; - if len > 12 { len } else { 0 } - }) - .sum(); - - // We don't use get_buffer_memory_size here, because gc is for the contents of the - // data buffers, not views and nulls. - let actual_buffer_size = s.data_buffers().iter().map(|b| b.capacity()).sum::(); - - // Re-creating the array copies data and can be time consuming. - // We only do it if the array is sparse - if actual_buffer_size > (ideal_buffer_size * 2) { - // We set the block size to `ideal_buffer_size` so that the new StringViewArray only - // has one buffer, which accelerate later concat_batches. See https://github.com/apache/arrow-rs/issues/6094 for more details. - let mut builder = StringViewBuilder::with_capacity(s.len()); - if ideal_buffer_size > 0 { - builder = builder.with_fixed_block_size(ideal_buffer_size as u32); - } - - for v in s.iter() { - builder.append_option(v); - } - - let gc_string = builder.finish(); - - debug_assert!(gc_string.data_buffers().len() <= 1); // buffer count can be 0 if the `ideal_buffer_size` is 0 - - Arc::new(gc_string) - } else { - Arc::clone(c) - } - }) - .collect(); - let mut options = RecordBatchOptions::new(); - options = options.with_row_count(Some(batch.num_rows())); - RecordBatch::try_new_with_options(batch.schema(), new_columns, &options) - .expect("Failed to re-create the gc'ed record batch") -} - #[cfg(test)] mod tests { use std::sync::Arc; use arrow::{ - array::{Int32Array, StringArray}, + array::{ + ArrayRef, BinaryArray, Int32Array, ListArray, StringArray, StringViewArray, StructArray, + }, + buffer::OffsetBuffer, ipc::writer::IpcWriteOptions, }; - use arrow_schema::{DataType, Field, Schema}; + use arrow_schema::{DataType, Field, Fields, Schema}; use config::meta::search::ScanStats; use super::*; @@ -266,9 +232,9 @@ mod tests { #[test] fn test_flight_data_encoder_new() { let options = IpcWriteOptions::default(); - let max_size = 8192; - let encoder = FlightDataEncoder::new(options, max_size); - assert_eq!(encoder.max_flight_data_size, max_size); + let encoder = FlightDataEncoder::new(options, 8192); + assert!(encoder.coalescer.is_none()); + assert_eq!(encoder.max_flight_data_size, 8192); } #[test] @@ -290,7 +256,11 @@ mod tests { let mut encoder = FlightDataEncoder::new(options, 8192); let batch = create_test_record_batch(); - let flight_data_vec = encoder.encode_batch(batch).unwrap(); + // encode_batch accumulates; data is only emitted once a full chunk is + // ready or finish() is called at stream end. + let mid = encoder.encode_batch(batch).unwrap(); + let end = encoder.finish().unwrap(); + let flight_data_vec: Vec<_> = mid.into_iter().chain(end).collect(); assert!(!flight_data_vec.is_empty()); for flight_data in flight_data_vec { @@ -334,6 +304,65 @@ mod tests { assert_eq!(message.bodyLength(), 0); } + #[test] + fn test_encode_batch_coalesces() { + let options = IpcWriteOptions::default(); + let mut encoder = FlightDataEncoder::new(options, 0); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + // 10000 rows: BatchCoalescer emits one 8192-row chunk during encode_batch, + // then finish() flushes the remaining 1808 rows → total 2 messages + let ids: Vec = (0..10000).collect(); + let batch = + arrow::array::RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids))]) + .unwrap(); + + let mut flight_data_vec = encoder.encode_batch(batch).unwrap(); + flight_data_vec.extend(encoder.finish().unwrap()); + assert!( + flight_data_vec.len() > 1, + "10000-row batch should be split into multiple FlightData messages" + ); + } + + #[test] + fn test_encode_batch_splits_large_batch() { + let schema = Arc::new(create_test_schema()); + // 100 rows of string data; measure real size and set limit to half of it + let ids: Vec = (0..100).collect(); + let names: Vec> = (0..100) + .map(|i| if i % 10 == 0 { None } else { Some("hello") }) + .collect(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + ], + ) + .unwrap(); + + let batch_size: usize = batch + .columns() + .iter() + .map(|col| { + col.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| col.get_buffer_memory_size()) + }) + .sum(); + let max_flight_data_size = batch_size / 2; + + let options = IpcWriteOptions::default(); + let mut encoder = FlightDataEncoder::new(options, max_flight_data_size); + + let mut flight_data_vec = encoder.encode_batch(batch).unwrap(); + flight_data_vec.extend(encoder.finish().unwrap()); + assert!( + flight_data_vec.len() > 1, + "batch exceeding max_flight_data_size ({max_flight_data_size} bytes) should split" + ); + } + #[test] fn test_split_batch_for_grpc_response_small_batch() { let batch = create_test_record_batch(); @@ -399,7 +428,11 @@ mod tests { let size = batch .columns() .iter() - .map(|col| col.get_buffer_memory_size()) + .map(|col| { + col.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| col.get_buffer_memory_size()) + }) .sum::(); let max_size = size / 2; // Should create 2 batches @@ -409,6 +442,149 @@ mod tests { assert!(batches.len() >= 2); } + #[test] + fn test_split_batch_sliced_array_no_overestimate() { + // Simulate what aggregateExec does: create a large batch then slice a small + // portion of it. get_buffer_memory_size() on the slice returns the full + // backing-buffer size, which caused over-splitting. The fix uses + // get_slice_memory_size() to measure only the used portion. + let schema = Arc::new(create_test_schema()); + let large_id = Arc::new(Int32Array::from((0..1000_i32).collect::>())); + let large_name = Arc::new(StringArray::from( + (0..1000).map(|i| format!("name_{i}")).collect::>(), + )); + let large_batch = RecordBatch::try_new(schema, vec![large_id, large_name]).unwrap(); + + // Slice to 5 rows — backing buffer still holds 1000 rows + let small_batch = large_batch.slice(0, 5); + + // max_size large enough that 5 rows should never need splitting + let slice_size: usize = small_batch + .columns() + .iter() + .map(|col| { + col.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| col.get_buffer_memory_size()) + }) + .sum(); + let max_size = slice_size + 1; + + let batches = split_batch_for_grpc_response(small_batch, max_size); + assert_eq!( + batches.len(), + 1, + "sliced batch should not be over-split due to backing-buffer overestimation" + ); + } + + fn assert_large_splits_slice_does_not(batch: RecordBatch, n_slice: usize) { + let measure = |b: &RecordBatch| -> usize { + b.columns() + .iter() + .map(|col| { + col.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| col.get_buffer_memory_size()) + }) + .sum() + }; + + let small = batch.slice(0, n_slice); + let full_size = measure(&batch); + let slice_size = measure(&small); + assert!( + full_size > slice_size, + "pre-condition: full batch must be larger than slice" + ); + let max_size = (slice_size + full_size) / 2; + + let large_out = split_batch_for_grpc_response(batch.clone(), max_size); + assert!(large_out.len() > 1); + assert_eq!( + large_out.iter().map(|b| b.num_rows()).sum::(), + batch.num_rows() + ); + + let small_out = split_batch_for_grpc_response(small, max_size); + assert_eq!(small_out.len(), 1); + } + + #[test] + fn test_split_batch_utf8() { + let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, true)])); + let strings: Vec> = (0..100).map(|i| Some(format!("word_{i:04}"))).collect(); + let arr: ArrayRef = Arc::new(StringArray::from( + strings.iter().map(|s| s.as_deref()).collect::>(), + )); + let batch = RecordBatch::try_new(schema, vec![arr]).unwrap(); + assert_large_splits_slice_does_not(batch, 5); + } + + #[test] + fn test_split_batch_utf8view() { + let schema = Arc::new(Schema::new(vec![Field::new( + "text", + DataType::Utf8View, + true, + )])); + let strings: Vec> = (0..100) + .map(|i| Some(format!("longer_string_{i:04}"))) + .collect(); + let arr: ArrayRef = Arc::new(StringViewArray::from( + strings.iter().map(|s| s.as_deref()).collect::>(), + )); + let batch = RecordBatch::try_new(schema, vec![arr]).unwrap(); + assert_large_splits_slice_does_not(batch, 5); + } + + #[test] + fn test_split_batch_struct_with_binary() { + let bin_field = Field::new("payload", DataType::Binary, true); + let child_fields = Fields::from(vec![bin_field]); + let struct_field = Field::new("row", DataType::Struct(child_fields.clone()), true); + let schema = Arc::new(Schema::new(vec![struct_field])); + + let bin_data: Vec> = + (0..100).map(|_| Some(b"binary_payload".as_ref())).collect(); + let bin_col: ArrayRef = Arc::new(BinaryArray::from(bin_data)); + let struct_col: ArrayRef = Arc::new(StructArray::new(child_fields, vec![bin_col], None)); + let batch = RecordBatch::try_new(schema, vec![struct_col]).unwrap(); + assert_large_splits_slice_does_not(batch, 5); + } + + #[test] + fn test_split_batch_list_with_utf8view() { + let item_field = Arc::new(Field::new("item", DataType::Utf8View, true)); + let list_field = Field::new("texts", DataType::List(Arc::clone(&item_field)), false); + let schema = Arc::new(Schema::new(vec![list_field])); + + // 100 rows, 2 items each → 200 values + let values: ArrayRef = Arc::new(StringViewArray::from( + (0..200_usize).map(|i| format!("v{i}")).collect::>(), + )); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat(2).take(100)); + let list_col: ArrayRef = Arc::new(ListArray::new(item_field, offsets, values, None)); + let batch = RecordBatch::try_new(schema, vec![list_col]).unwrap(); + assert_large_splits_slice_does_not(batch, 5); + } + + #[test] + fn test_split_batch_list_with_utf8() { + let item_field = Arc::new(Field::new("item", DataType::Utf8, true)); + let list_field = Field::new("words", DataType::List(Arc::clone(&item_field)), false); + let schema = Arc::new(Schema::new(vec![list_field])); + + // 100 rows, 2 words each → 200 values + let values: ArrayRef = Arc::new(StringArray::from( + (0..200_usize).map(|i| format!("w{i}")).collect::>(), + )); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat(2).take(100)); + let list_col: ArrayRef = Arc::new(ListArray::new(item_field, offsets, values, None)); + let batch = RecordBatch::try_new(schema, vec![list_col]).unwrap(); + assert_large_splits_slice_does_not(batch, 5); + } + #[test] fn test_encode_custom_peak_memory() { let options = IpcWriteOptions::default(); diff --git a/src/handler/grpc/flight/stream.rs b/src/handler/grpc/flight/stream.rs index a6e5f43192..6e401973a1 100644 --- a/src/handler/grpc/flight/stream.rs +++ b/src/handler/grpc/flight/stream.rs @@ -132,6 +132,12 @@ impl FlightEncoderStream { Ok(()) } + fn finish_encoder(&mut self) -> Result<(), FlightError> { + let flight_data = self.encoder.finish()?; + self.queue.extend(flight_data); + Ok(()) + } + fn encode_custom_messages(&mut self) -> Result<(), FlightError> { let custom_messages = std::mem::take(&mut self.custom_messages); for message in custom_messages.into_iter() { @@ -144,7 +150,7 @@ impl FlightEncoderStream { Ok(()) } - fn encode_custom_scan_stats(&mut self) -> Result<(), FlightError> { + fn encode_early_emit_messages(&mut self) -> Result<(), FlightError> { let custom_messages = std::mem::take(&mut self.custom_messages); let mut remainder_messages = Vec::new(); for message in custom_messages.into_iter() { @@ -203,17 +209,23 @@ impl Stream for FlightEncoderStream { let batch = ready!(self.inner.poll_next_unpin(cx)); match batch { - None if self.custom_messages.is_empty() => { - self.done = true; - assert!(self.queue.is_empty()); - return Poll::Ready(None); - } None => { + // Flush remaining coalesced batches first, then send custom + // messages (e.g. peak memory). Ordering: all data before stats. + if let Err(e) = self.finish_encoder() { + self.done = true; + self.queue.clear(); + return Poll::Ready(Some(Err(tonic::Status::internal(e.to_string())))); + } + // TODO: maybe we can embed custom messages into the last batch to avoid extra + // FlightData at the end? if let Err(e) = self.encode_custom_messages() { self.done = true; self.queue.clear(); return Poll::Ready(Some(Err(tonic::Status::internal(e.to_string())))); } + self.done = true; + // loop continues to drain queue before returning None } Some(Err(e)) => { self.done = true; @@ -227,9 +239,9 @@ impl Stream for FlightEncoderStream { return Poll::Ready(Some(Err(tonic::Status::internal(e.to_string())))); } Some(Ok(batch)) => { - // before send the first batch, send the scan_stats first + // before send the first batch, send the early emit messages first if self.first_batch && !self.custom_messages.is_empty() { - if let Err(e) = self.encode_custom_scan_stats() { + if let Err(e) = self.encode_early_emit_messages() { self.done = true; self.queue.clear(); return Poll::Ready(Some(Err(tonic::Status::internal(e.to_string()))));