starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-04 16:09:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: make split_batch_for_grpc_response based on the real memory size (#12466)

Browse Source 
- close https://github.com/openobserve/o2-enterprise/issues/1865

## TODO
- [x] make split_batch_for_grpc_response based on the real memory sizew
- [x] coalesce and gc view before send to leader

```
SELECT trace_id, count(*) as cnt, max(_timestamp) as ts FROM "k8s_logs" group by trace_id order by cnt desc limit 10

utf8view enable:
main: 30s (15k batches send to leader)
this branch: 15s (1.5k batches send to leader)

utf8view disable:
main: 13s (42k batches send to leader)
this branch: 12s (1.5k batches send to leader)
```

one found is the AggregateExec's time_calculating_group_ids, utf8 is
faster than utf8view
This commit is contained in:
Huaijin
2026-06-05 09:24:41 +00:00
committed by GitHub
parent 17a713fe4c
commit 7fbbcc4b02
3 changed files with 333 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/src/config.rs
+2 -2
View File
@@ -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
src/flight/src/encoder.rs
+311 -135
View File
@@ -13,10 +13,9 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<BatchCoalescer>,
/// 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
/// - <https://github.com/openobserve/openobserve/issues/8280>
/// - <https://github.com/apache/datafusion/pull/11587>
pub fn encode_batch(&mut self, batch: RecordBatch) -> Result<Vec<FlightData>> {
// 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<Vec<FlightData>> {
let Some(coalescer) = self.coalescer.as_mut() else {
return Ok(vec![]);
};
coalescer.finish_buffered_batch()?;
self.drain_completed()
}
fn drain_completed(&mut self) -> Result<Vec<FlightData>> {
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<RecordBatch> {
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::<usize>();
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<RecordBatch> {
let size = batch
.columns()
.iter()
.map(|col| col.get_buffer_memory_size())
.sum::<usize>();
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<ArrayRef> = 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::<usize>();
// 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<i32> = (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<i32> = (0..100).collect();
let names: Vec<Option<&str>> = (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::<usize>();
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::<Vec<_>>()));
let large_name = Arc::new(StringArray::from(
(0..1000).map(|i| format!("name_{i}")).collect::<Vec<_>>(),
));
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::<usize>(),
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<Option<String>> = (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::<Vec<_>>(),
));
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<Option<String>> = (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::<Vec<_>>(),
));
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<Option<&[u8]>> =
(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::<Vec<_>>(),
));
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::<Vec<_>>(),
));
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();
src/handler/grpc/flight/stream.rs
+20 -8
View File
@@ -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()))));