refactor: spawn to read avro ObjectStore stream

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 }
 

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<F, O, E>(&self, f: F) -> BoxFuture<'_, Result<O, AvroError>>
     where
         F: for<'a> FnOnce(&'a Arc<dyn ObjectStore>, &'a Path) -> BoxFuture<'a, Result<O, E>>
@@ -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<F, I, E>(
+        &self,
+        f: F,
+    ) -> BoxFuture<'_, Result<BoxStream<'_, Result<I, AvroError>>, AvroError>>
+    where
+        F: for<'a> FnOnce(
+                &'a Arc<dyn ObjectStore>,
+                &'a Path,
+            )
+                -> BoxFuture<'a, Result<BoxStream<'static, Result<I, E>>, 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<u64>,
     ) -> BoxFuture<'_, Result<BoxStream<'_, Result<Bytes, AvroError>>, 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(