This is an automated email from the ASF dual-hosted git repository.
Dandandan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new f3055a415 perf(core): adopt transport blocks instead of copying them
in BlockDataStream (#2312)
f3055a415 is described below
commit f3055a415b2f3fe49090622c160e9e04487cc31b
Author: Daniël Heres <[email protected]>
AuthorDate: Sun Aug 16 09:48:38 2026 +0200
perf(core): adopt transport blocks instead of copying them in
BlockDataStream (#2312)
`BlockDataStream` concatenated every incoming block into a freshly
allocated buffer, even when nothing was pending. `Buffer::from(Bytes)`
adopts the transport allocation rather than copying it, and
`StreamDecoder::decode` drains `state_buffer` completely before the
stream asks for another block, so in the steady state the block can be
taken as-is.
Replace `combine_buffers` with `append_block`, which adopts the incoming
block when nothing is pending and falls back to concatenating only for a
partial message straddling a block boundary — the case the schema
accumulation loop in `try_new` relies on.
65 MiB payload at the server's 8 MiB block size: 69.7 ms -> 22.7 ms
(-67%). At 1 MiB blocks: 41.8 ms -> 34.9 ms (-17%).
Co-authored-by: Claude Opus 5 <[email protected]>
---
ballista/core/src/client.rs | 50 ++++++++++++++++++++++++++++++++++++++-------
1 file changed, 43 insertions(+), 7 deletions(-)
diff --git a/ballista/core/src/client.rs b/ballista/core/src/client.rs
index 267ef8300..ab50deffa 100644
--- a/ballista/core/src/client.rs
+++ b/ballista/core/src/client.rs
@@ -482,8 +482,7 @@ impl<S: Stream<Item = Result<prost::bytes::Bytes>> + Unpin>
BlockDataStream<S> {
match ipc_stream.next().await {
Some(Ok(blob)) => {
- state_buffer =
- Self::combine_buffers(&state_buffer,
&Buffer::from(blob));
+ state_buffer = Self::append_block(state_buffer, blob);
match try_schema_from_ipc_buffer(state_buffer.as_slice()) {
Ok(schema) => {
@@ -517,10 +516,21 @@ impl<S: Stream<Item = Result<prost::bytes::Bytes>> +
Unpin> BlockDataStream<S> {
}
impl<S: Stream<Item = Result<prost::bytes::Bytes>> + Unpin> BlockDataStream<S>
{
- fn combine_buffers(first: &Buffer, second: &Buffer) -> Buffer {
- let mut combined = MutableBuffer::new(first.len() + second.len());
- combined.extend_from_slice(first.as_slice());
- combined.extend_from_slice(second.as_slice());
+ /// Appends a transport block to the bytes still waiting to be decoded.
+ ///
+ /// `Buffer::from(Bytes)` adopts the transport allocation instead of
copying
+ /// it, so when nothing is pending — which is the steady state, since
+ /// [`StreamDecoder::decode`] drains `state_buffer` completely before the
+ /// stream asks for another block — the block is taken as-is. Only a
partial
+ /// message straddling a block boundary needs the concatenating path.
+ fn append_block(pending: Buffer, blob: prost::bytes::Bytes) -> Buffer {
+ let incoming = Buffer::from(blob);
+ if pending.is_empty() {
+ return incoming;
+ }
+ let mut combined = MutableBuffer::new(pending.len() + incoming.len());
+ combined.extend_from_slice(pending.as_slice());
+ combined.extend_from_slice(incoming.as_slice());
combined.into()
}
@@ -533,7 +543,8 @@ impl<S: Stream<Item = Result<prost::bytes::Bytes>> + Unpin>
BlockDataStream<S> {
//TODO: do we want to limit maximum buffer size here as well?
//
self.transmitted += blob.len();
- self.state_buffer = Self::combine_buffers(&self.state_buffer,
&Buffer::from(blob))
+ let pending = std::mem::take(&mut self.state_buffer);
+ self.state_buffer = Self::append_block(pending, blob);
}
}
@@ -699,6 +710,31 @@ mod tests {
assert_eq!(batches, result.unwrap())
}
+ #[tokio::test]
+ async fn should_process_multi_block_payload() {
+ // Realistic transport shape: a payload spanning several whole blocks.
+ // Once the decoder has drained the previous block, the next one is
+ // adopted rather than copied; block sizes that leave a partial schema
+ // message still exercise the concatenating path in `try_new`.
+ let batches = generate_batches();
+ let ipc_blob = generate_ipc_stream(&batches);
+
+ for block_size in [8usize, 64, 512] {
+ let stream = futures::stream::iter(ipc_blob.clone())
+ .chunks(block_size)
+ .map(|b| Ok(Bytes::from(b)));
+
+ let result: datafusion::error::Result<Vec<RecordBatch>> =
+ BlockDataStream::try_new(stream)
+ .await
+ .unwrap()
+ .try_collect()
+ .await;
+
+ assert_eq!(batches, result.unwrap(), "block_size={block_size}");
+ }
+ }
+
#[tokio::test]
async fn should_process_concatenated_streams() {
let batches = generate_batches();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
