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..3d1f76ae993a 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -23,10 +23,21 @@ 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; + +// Size for the channel buffer between the I/O task driving the object store client +// 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 = 2; /// An implementation of an AsyncFileReader using the [`ObjectStore`] API. pub struct AvroObjectReader { @@ -60,6 +71,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> @@ -88,6 +104,73 @@ impl AvroObjectReader { .boxed(), } } + + // 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, + ) -> 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 +182,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(