From fb41f90b61376fe4e485817f3241db09cba5f183 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 8 May 2026 06:52:04 +0300 Subject: [PATCH 1/4] refactor: spawn to read avro ObjectStore stream Use the runtime handle provided to the AvroObjectReader to spawn tasks that perform I/O operations on the ObjectStore. --- arrow-avro/Cargo.toml | 3 +- arrow-avro/src/reader/async_reader/store.rs | 96 ++++++++++++++++----- 2 files changed, 76 insertions(+), 23 deletions(-) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index d76f294601f3..6a5e49cf5196 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -76,7 +76,8 @@ indexmap = "2.10" rand = "0.9" md5 = { version = "0.8", optional = true } sha2 = { version = "0.11", optional = true } -tokio = { version = "1.0", optional = true, default-features = false, features = ["macros", "rt", "io-util"] } +tokio = { version = "1.0", optional = true, default-features = false, features = ["macros", "rt", "io-util", "sync"] } +tokio-stream = { version = "0.1", default-features = false } tokio-util = { version = "0.7.18", default-features = false, features = ["io"], optional = true } async-stream = { version = "0.3.6", optional = true } diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index 3cf34d3ec97b..d6fd39ca3dac 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -23,10 +23,15 @@ use futures::stream::BoxStream; use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; use object_store::path::Path; use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; +use tokio::runtime::Handle; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + use std::error::Error; use std::ops::Range; use std::sync::Arc; -use tokio::runtime::Handle; + +const STREAM_BUFFER_SIZE: usize = 8; /// An implementation of an AsyncFileReader using the [`ObjectStore`] API. pub struct AvroObjectReader { @@ -88,6 +93,62 @@ impl AvroObjectReader { .boxed(), } } + + fn spawn_stream( + &self, + f: F, + ) -> BoxFuture<'_, Result>, AvroError>> + where + F: for<'a> FnOnce( + &'a Arc, + &'a Path, + ) + -> BoxFuture<'a, Result>, E>> + + Send + + 'static, + I: Send + 'static, + E: Error + Send + 'static, + { + match &self.runtime { + Some(handle) => { + let path = self.path.clone(); + let store = Arc::clone(&self.store); + async move { + let (sender, receiver) = mpsc::channel(STREAM_BUFFER_SIZE); + let mut stream = handle + .spawn(async move { f(&store, &path).await }) + .map_ok_or_else( + |e| match e.try_into_panic() { + Err(e) => Err(AvroError::External(Box::new(e))), + Ok(p) => std::panic::resume_unwind(p), + }, + |res| res.map_err(|e| AvroError::General(e.to_string())), + ) + .await?; + handle.spawn(async move { + while let Some(item) = stream.next().await { + let send_res = sender.send(item).await; + if send_res.is_err() { + break; + } + } + }); + Ok(ReceiverStream::new(receiver) + .map_err(|e| AvroError::General(e.to_string())) + .boxed()) + } + .boxed() + } + None => f(&self.store, &self.path) + .map_ok(|stream| { + stream + .map_err(|e| AvroError::General(e.to_string())) + .boxed() + }) + .map_err(|e| AvroError::General(e.to_string())) + .boxed(), + } + } } impl AsyncFileReader for AvroObjectReader { @@ -99,27 +160,18 @@ impl AsyncFileReader for AvroObjectReader { &mut self, range: Range, ) -> BoxFuture<'_, Result>, AvroError>> { - // FIXME: can't use self.spawn here because of the signature of the returned stream. - // The signature has to be this way to work with the generic implementation - // for AsyncRead + AsyncSeek types. - async move { - let options = GetOptions { - range: Some(GetRange::Bounded(range)), - ..Default::default() - }; - - let get_result = self - .store - .get_opts(&self.path, options) - .await - .map_err(|e| AvroError::External(Box::new(e)))?; - let stream = get_result - .into_stream() - .map_err(|e| AvroError::External(Box::new(e))) - .boxed(); - Ok(stream) - } - .boxed() + self.spawn_stream(|store, path| { + async move { + let options = GetOptions { + range: Some(GetRange::Bounded(range)), + ..Default::default() + }; + let get_result = store.get_opts(path, options).await?; + let stream = get_result.into_stream(); + Ok(stream) + } + .boxed() + }) } fn get_byte_ranges( From 4410b0b7c532044c54a1433ac021c62733acfbac Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 8 May 2026 15:16:23 +0300 Subject: [PATCH 2/4] doc: stream buffer size magic number --- arrow-avro/src/reader/async_reader/store.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index d6fd39ca3dac..c25bc351d7b6 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -31,6 +31,11 @@ use std::error::Error; use std::ops::Range; use std::sync::Arc; +// Size for the channel buffer between the I/O task driving the object store client +// and the task decoding the received chunks. This should not be too large to +// avoid buffering too much data in memory in case the decoding task is slower +// than the I/O task. A typical data chunk size is 8-64 KiB for HTTP backends +// and 8 MiB for `LocalFileSystem`. const STREAM_BUFFER_SIZE: usize = 8; /// An implementation of an AsyncFileReader using the [`ObjectStore`] API. From 78509670efeaec9dc6f1d87628a5ebdd1b9fc0ae Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 8 May 2026 17:07:43 +0300 Subject: [PATCH 3/4] doc: comment spawn helpers in AvroObjectReader --- arrow-avro/src/reader/async_reader/store.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index c25bc351d7b6..09081cda43e2 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -70,6 +70,11 @@ impl AvroObjectReader { } } + // If the runtime handle is provided, spawns the provided async function + // on the runtime to retrieve the result, and wraps the awaiting for + // the result in a boxed future. + // If no runtime handle is provided, simply invokes the closure + // and adapts the error type in the async result. fn spawn(&self, f: F) -> BoxFuture<'_, Result> where F: for<'a> FnOnce(&'a Arc, &'a Path) -> BoxFuture<'a, Result> @@ -99,6 +104,17 @@ impl AvroObjectReader { } } + // Adaptation of `spawn` for streaming results. If the runtime handle + // is provided, spawns the provided async function on the runtime + // to retrieve the stream. If the stream is successfully established, + // spawns a new task to drive the stream and forward the items to the + // consumer of the returned stream object. + // The two separate tasks are spawned to provide error handling at the + // stream establishment phase and to get internally simpler task states + // to work with. + // If no runtime handle is provided, simply invokes the closure + // and adapts the error type in both the stream establishment result + // and the resulting stream. fn spawn_stream( &self, f: F, From 973dff57e97b75773e690e6f48dc0a350f90a9ba Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Sun, 10 May 2026 12:55:14 +0300 Subject: [PATCH 4/4] perf: minimize buffering in AvroObjectReader --- arrow-avro/src/reader/async_reader/store.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index 09081cda43e2..3d1f76ae993a 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -32,11 +32,12 @@ use std::ops::Range; use std::sync::Arc; // Size for the channel buffer between the I/O task driving the object store client -// and the task decoding the received chunks. This should not be too large to -// avoid buffering too much data in memory in case the decoding task is slower -// than the I/O task. A typical data chunk size is 8-64 KiB for HTTP backends +// and the task decoding the received chunks. This is minimized to +// avoid excessive buffering in case the decoding task is slower +// than the I/O task; the main purpose of the channel is to permit concurrency. +// A typical data chunk size is 8-64 KiB for HTTP backends // and 8 MiB for `LocalFileSystem`. -const STREAM_BUFFER_SIZE: usize = 8; +const STREAM_BUFFER_SIZE: usize = 2; /// An implementation of an AsyncFileReader using the [`ObjectStore`] API. pub struct AvroObjectReader {