feat: async stream read

arrow-avro/Cargo.toml

@@ -46,7 +46,7 @@ small_decimals = []
 avro_custom_types = ["dep:arrow-select"]
 
 # Enable async APIs
-async = ["futures", "tokio"]
+async = ["futures", "tokio", "tokio-util", "async-stream"]
 # Enable object_store integration
 object_store = ["dep:object_store", "async"]
 
@@ -77,6 +77,8 @@ 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-util = { version = "0.7.18", default-features = false, features = ["io"], optional = true }
+async-stream = { version = "0.3.6", optional = true }
 
 [dev-dependencies]
 arrow-data = { workspace = true }

arrow-avro/src/reader/async_reader/async_file_reader.rs

@@ -17,11 +17,13 @@
 
 use crate::errors::AvroError;
 use bytes::Bytes;
-use futures::FutureExt;
 use futures::future::BoxFuture;
+use futures::stream::BoxStream;
+use futures::{FutureExt, StreamExt, TryStreamExt};
 use std::io::SeekFrom;
 use std::ops::Range;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
+use tokio_util::io::ReaderStream;
 
 /// The asynchronous interface used by [`super::AsyncAvroFileReader`] to read avro files
 ///
@@ -40,6 +42,12 @@ pub trait AsyncFileReader: Send {
     /// Retrieve the bytes in `range`
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>>;
 
+    /// Retrieve a range of bytes as a stream
+    fn get_stream(
+        &mut self,
+        range: Range<u64>,
+    ) -> BoxFuture<'_, Result<BoxStream<'_, Result<Bytes, AvroError>>, AvroError>>;
+
     /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially
     fn get_byte_ranges(
         &mut self,
@@ -61,6 +69,13 @@ pub trait AsyncFileReader: Send {
 
 /// This allows Box<dyn AsyncFileReader + '_> to be used as an AsyncFileReader,
 impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
+    fn get_stream(
+        &mut self,
+        range: Range<u64>,
+    ) -> BoxFuture<'_, Result<BoxStream<'_, Result<Bytes, AvroError>>, AvroError>> {
+        self.as_mut().get_stream(range)
+    }
+
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
         self.as_mut().get_bytes(range)
     }
@@ -74,6 +89,23 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
 }
 
 impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
+    fn get_stream(
+        &mut self,
+        range: Range<u64>,
+    ) -> BoxFuture<'_, Result<BoxStream<'_, Result<Bytes, AvroError>>, AvroError>> {
+        async move {
+            self.seek(SeekFrom::Start(range.start)).await?;
+
+            let ranged_reader = self.take(range.end - range.start);
+            let stream = ReaderStream::new(ranged_reader)
+                .map_err(AvroError::from)
+                .boxed();
+
+            Ok(stream)
+        }
+        .boxed()
+    }
+
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
         async move {
             self.seek(SeekFrom::Start(range.start)).await?;

arrow-avro/src/reader/async_reader/builder.rs

@@ -17,7 +17,6 @@
 
 use crate::codec::{AvroFieldBuilder, Tz};
 use crate::errors::AvroError;
-use crate::reader::async_reader::ReaderState;
 use crate::reader::header::{Header, HeaderDecoder, HeaderInfo};
 use crate::reader::record::RecordDecoder;
 use crate::reader::{AsyncAvroFileReader, AsyncFileReader, Decoder};
@@ -183,7 +182,7 @@ where
         .ok_or_else(|| AvroError::EOF("Unexpected EOF while reading Avro header".into()))
 }
 
-impl<R: AsyncFileReader> ReaderBuilder<R> {
+impl<R: AsyncFileReader + Unpin + Send + 'static> ReaderBuilder<R> {
     /// Build the asynchronous Avro reader with the provided parameters.
     /// This reads the header first to initialize the reader state.
     pub async fn try_build(mut self) -> Result<AsyncAvroFileReader<R>, AvroError> {
@@ -273,24 +272,18 @@ impl<R: AsyncFileReader> ReaderBuilder<R> {
 
         // Determine if there is actually data to fetch, note that we subtract the header len from range.start,
         // so we need to check if range.end == header_len to see if there's no data after the header
-        let reader_state = if range.start == range.end || header_len == range.end {
-            ReaderState::Finished
-        } else {
-            ReaderState::Idle {
-                reader: self.reader,
-            }
-        };
-
+        let finished = range.start == range.end || header_len == range.end;
         let codec = header_info.compression()?;
         let sync_marker = header_info.sync();
 
         Ok(AsyncAvroFileReader::new(
+            self.reader,
             range,
             self.file_size,
             decoder,
             codec,
             sync_marker,
-            reader_state,
+            finished,
         ))
     }
 }