From e47f6da4cd1a8dc6ad0827b0e4aa98b9a41aa424 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Tue, 31 Mar 2026 13:29:14 +0300 Subject: [PATCH 01/11] feat: async stream read --- arrow-avro/Cargo.toml | 3 +- .../reader/async_reader/async_file_reader.rs | 34 +- arrow-avro/src/reader/async_reader/builder.rs | 18 +- arrow-avro/src/reader/async_reader/mod.rs | 855 +++++++++--------- arrow-avro/src/reader/async_reader/store.rs | 39 +- 5 files changed, 515 insertions(+), 434 deletions(-) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index f46ef7e7b999..681304200449 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -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"] } +async-stream = "0.3.6" [dev-dependencies] arrow-data = { workspace = true } @@ -90,7 +92,6 @@ criterion = { workspace = true, default-features = false } tempfile = "3.3" arrow = { workspace = true, features = ["prettyprint"] } futures = "0.3.31" -async-stream = "0.3.6" apache-avro = "0.21.0" num-bigint = "0.4" object_store = { workspace = true, features = ["fs"] } diff --git a/arrow-avro/src/reader/async_reader/async_file_reader.rs b/arrow-avro/src/reader/async_reader/async_file_reader.rs index 1257a2f3dd4d..18facd536f82 100644 --- a/arrow-avro/src/reader/async_reader/async_file_reader.rs +++ b/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) -> BoxFuture<'_, Result>; + /// Retrieve a range of bytes as a stream + fn get_stream( + &mut self, + range: Range, + ) -> BoxFuture<'_, Result>, 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 to be used as an AsyncFileReader, impl AsyncFileReader for Box { + fn get_stream( + &mut self, + range: Range, + ) -> BoxFuture<'_, Result>, AvroError>> { + self.as_mut().get_stream(range) + } + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { self.as_mut().get_bytes(range) } @@ -74,6 +89,23 @@ impl AsyncFileReader for Box { } impl AsyncFileReader for T { + fn get_stream( + &mut self, + range: Range, + ) -> BoxFuture<'_, Result>, 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(ArrowError::from) + .boxed(); + + Ok(stream) + } + .boxed() + } + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { async move { self.seek(SeekFrom::Start(range.start)).await?; diff --git a/arrow-avro/src/reader/async_reader/builder.rs b/arrow-avro/src/reader/async_reader/builder.rs index d3cca70425de..ddead347d279 100644 --- a/arrow-avro/src/reader/async_reader/builder.rs +++ b/arrow-avro/src/reader/async_reader/builder.rs @@ -183,7 +183,7 @@ where .ok_or_else(|| AvroError::EOF("Unexpected EOF while reading Avro header".into())) } -impl ReaderBuilder { +impl ReaderBuilder { /// 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, AvroError> { @@ -273,24 +273,18 @@ impl ReaderBuilder { // 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 codec = header_info.compression()?; - let sync_marker = header_info.sync(); + let finished = range.start == range.end || self.header_len == range.end; + let codec = self.header.compression()?; + let sync_marker = self.header.sync(); Ok(AsyncAvroFileReader::new( + self.inner.reader, range, self.file_size, decoder, codec, sync_marker, - reader_state, + finished, )) } } diff --git a/arrow-avro/src/reader/async_reader/mod.rs b/arrow-avro/src/reader/async_reader/mod.rs index c034411edb03..f467b364553c 100644 --- a/arrow-avro/src/reader/async_reader/mod.rs +++ b/arrow-avro/src/reader/async_reader/mod.rs @@ -25,9 +25,9 @@ use crate::reader::Decoder; use crate::reader::block::{BlockDecoder, BlockDecoderState}; use arrow_array::RecordBatch; use arrow_schema::{ArrowError, SchemaRef}; -use bytes::Bytes; -use futures::future::BoxFuture; -use futures::{FutureExt, Stream}; +use bytes::{Bytes, BytesMut}; +use futures::future::{BoxFuture, BoxStream}; +use futures::{FutureExt, Stream, StreamExt}; use std::mem; use std::ops::Range; use std::pin::Pin; @@ -46,39 +46,214 @@ use crate::errors::AvroError; #[cfg(feature = "object_store")] pub use store::AvroObjectReader; -enum FetchNextBehaviour { - /// Initial read: scan for sync marker, then move to decoding blocks - ReadSyncMarker, - /// Parse VLQ header bytes one at a time until Data state, then continue decoding - DecodeVLQHeader, - /// Continue decoding the current block with the fetched data - ContinueDecoding, -} - -enum ReaderState { - /// Intermediate state to fix ownership issues - InvalidState, - /// Initial state, fetch initial range - Idle { reader: R }, - /// Fetching data from the reader - FetchingData { - future: BoxFuture<'static, Result<(R, Bytes), AvroError>>, - next_behaviour: FetchNextBehaviour, +/// State of the block-reading loop inside [`make_stream`]. +/// +/// The active `stream` lives alongside `reader` as a local variable in the async block; +/// the compiler's async state machine (via `Pin`) handles that self-referential borrow +/// automatically — no `unsafe` required. +enum Phase { + ReadSyncMarker { + accumulated: BytesMut, + }, + DecodingBlock { + buffered: Bytes, }, - /// Decode a block in a loop until completion - DecodingBlock { data: Bytes, reader: R }, - /// Output batches from a decoded block ReadingBatches { - data: Bytes, block_data: Bytes, - remaining_in_block: usize, - reader: R, + remaining: usize, + /// Bytes left over in the stream chunk after the completed block. + next_buffered: Bytes, }, - /// Successfully finished reading file contents; drain any remaining buffered records - /// from the decoder into (possibly partial) output batches. - Flushing, - /// Done, flush decoder and return - Finished, +} + +/// The non-reader fields of [`AvroReaderState`], split out so that +/// `remaining_block_range` can borrow them while `reader` is mutably borrowed by the active stream. +struct AvroReaderMeta { + range: Range, + file_size: u64, + decoder: Decoder, + codec: Option, + sync_marker: [u8; 16], +} + +impl AvroReaderMeta { + /// Calculate the byte range needed to complete the current block. + /// Only valid when `block_decoder` is in `Data` or `Sync` state. + fn remaining_block_range( + &self, + block_decoder: &BlockDecoder, + ) -> Result, ArrowError> { + let remaining = block_decoder.bytes_remaining() as u64 + + match block_decoder.state() { + BlockDecoderState::Data => 16, // Include sync marker + BlockDecoderState::Sync => 0, + state => { + return Err(ArrowError::ParseError(format!( + "remaining_block_range called in unexpected state: {state:?}" + ))); + } + }; + let fetch_end = self.range.end + remaining; + if fetch_end > self.file_size { + return Err(ArrowError::ParseError(format!( + "Avro block requires more bytes than what exists in the file: \ + fetch_end {fetch_end}, remaining {remaining}, file_size {}", + self.file_size + ))); + } + Ok(self.range.end..fetch_end) + } +} + +/// Owns all the state needed to drive the Avro block-reading loop. +struct AvroReaderState { + reader: R, + meta: AvroReaderMeta, +} + +/// Drives all reading and decoding as a natural async generator. +/// +/// `state.reader` is owned (moved in). `stream` is a local borrow of `state.reader` inside the +/// `async_stream::try_stream!` block; the compiler's generated async state machine stores +/// both and manages the lifetime relationship safely via `Pin`. +fn make_stream( + mut state: AvroReaderState, +) -> impl Stream> + Send + 'static { + async_stream::try_stream! { + if state.meta.range.start >= state.meta.range.end { + Err(ArrowError::InvalidArgumentError(format!( + "Invalid range specified for Avro file: start {} >= end {}, file_size: {}", + state.meta.range.start, state.meta.range.end, state.meta.file_size + )))?; + } + + let mut block_decoder = BlockDecoder::default(); + let mut finishing_partial_block = false; + + // Obtain the initial stream and begin scanning for the first sync marker. + let mut stream = state.reader.get_stream(state.meta.range.clone()).await?; + let mut phase = Phase::ReadSyncMarker { accumulated: BytesMut::new() }; + + loop { + phase = match phase { + Phase::ReadSyncMarker { mut accumulated } => { + match stream.next().await { + Some(Ok(chunk)) => { + accumulated.extend_from_slice(&chunk); + match accumulated.windows(16).position(|w| w == state.meta.sync_marker) { + Some(pos) => Phase::DecodingBlock { + buffered: accumulated.freeze().slice(pos + 16..), + }, + None => { + // Only the trailing 15 bytes matter: the marker could + // straddle a chunk boundary by at most (sync_marker.len() - 1) bytes. + // Copy them to the front to avoid reallocation on the next extend. + let discard = accumulated.len().saturating_sub(state.meta.sync_marker.len() - 1); + accumulated.copy_within(discard.., 0); + accumulated.truncate(accumulated.len() - discard); + Phase::ReadSyncMarker { accumulated } + } + } + } + Some(Err(e)) => Err(e)?, + None => break, // file has no blocks after the header + } + } + + Phase::DecodingBlock { mut buffered } => { + let consumed = block_decoder.decode(&buffered)?; + buffered = buffered.slice(consumed..); + + if let Some(block) = block_decoder.flush() { + // A complete block is ready — decompress and start yielding batches. + let block_data = Bytes::from_owner(match &state.meta.codec { + Some(c) => c.decompress(&block.data)?, + None => block.data, + }); + Phase::ReadingBatches { + block_data, + remaining: block.count, + next_buffered: buffered, + } + } else if !buffered.is_empty() { + Err(ArrowError::ParseError( + "BlockDecoder failed to make progress decoding Avro block".into(), + ))? + } else { + // Block incomplete and buffer empty — pull next chunk from stream. + match stream.next().await { + Some(Ok(chunk)) => Phase::DecodingBlock { buffered: chunk }, + Some(Err(e)) => Err(e)?, + None => { + // Stream exhausted — compute the next range to fetch. + let range_to_fetch = match block_decoder.state() { + BlockDecoderState::Count | BlockDecoderState::Size => + { + // VLQ header is at most 20 bytes; fetch just enough to + // complete it so block_decoder can parse count + size. + const MAX_VLQ_HEADER_SIZE: u64 = 20; + let fetch_end = + (state.meta.range.end + MAX_VLQ_HEADER_SIZE).min(state.meta.file_size); + if fetch_end == state.meta.range.end || finishing_partial_block { + // Nothing more to fetch, or partial fetch already done. + break; + } + let r = state.meta.range.end..fetch_end; + state.meta.range.end = fetch_end; + r + }, + BlockDecoderState::Data | BlockDecoderState::Sync | BlockDecoderState::Finished => { + // Mid-block (Data/Sync): fetch exactly the bytes needed + // to finish this block (data remainder + sync marker). + if finishing_partial_block { + Err(ArrowError::ParseError( + "Unexpected EOF while reading last Avro block".into(), + ))?; + } + finishing_partial_block = true; + state.meta.remaining_block_range(&block_decoder)? + }, + }; + + // Drop the exhausted stream before reborrowing `state.reader`. + drop(stream); + stream = state.reader.get_stream(range_to_fetch).await?; + Phase::DecodingBlock { buffered: Bytes::new() } + } + } + } + } + + Phase::ReadingBatches { mut block_data, mut remaining, next_buffered } => { + let (consumed, decoded) = state.meta.decoder.decode_block(&block_data, remaining)?; + remaining -= decoded; + block_data = block_data.slice(consumed..); + + let next_phase = if remaining == 0 { + Phase::DecodingBlock { buffered: next_buffered } + } else { + Phase::ReadingBatches { block_data, remaining, next_buffered } + }; + + if state.meta.decoder.batch_is_full() { + match state.meta.decoder.flush()? { + Some(batch) => yield batch, + None => Err(ArrowError::ParseError( + "Decoder reported a full batch, but flush returned None".into(), + ))?, + } + } + + next_phase + } + }; + } + + // Flush any records that did not fill a complete batch. + while let Some(batch) = state.meta.decoder.flush()? { + yield batch; + } + } } /// An asynchronous Avro file reader that implements `Stream>`. @@ -135,19 +310,8 @@ enum ReaderState { /// } /// ``` pub struct AsyncAvroFileReader { - // Members required to fetch data - range: Range, - file_size: u64, - - // Members required to actually decode and read data - decoder: Decoder, - block_decoder: BlockDecoder, - codec: Option, - sync_marker: [u8; 16], - - // Members keeping the current state of the reader - reader_state: ReaderState, - finishing_partial_block: bool, + inner: BoxStream<'static, Result>, + _phantom: PhantomData, } impl AsyncAvroFileReader { @@ -155,391 +319,45 @@ impl AsyncAvroFileReader { pub fn builder(reader: R, file_size: u64, batch_size: usize) -> ReaderBuilder { ReaderBuilder::new(reader, file_size, batch_size) } +} +impl AsyncAvroFileReader { fn new( + reader: R, range: Range, file_size: u64, decoder: Decoder, codec: Option, sync_marker: [u8; 16], - reader_state: ReaderState, + finished: bool, ) -> Self { - Self { - range, - file_size, - - decoder, - block_decoder: Default::default(), - codec, - sync_marker, - - reader_state, - finishing_partial_block: false, - } - } - - /// Returns the Arrow schema for batches produced by this reader. - /// - /// The schema is determined by the writer schema in the file and the reader schema provided to the builder. - pub fn schema(&self) -> SchemaRef { - self.decoder.schema() - } - - /// Calculate the byte range needed to complete the current block. - /// Only valid when block_decoder is in Data or Sync state. - /// Returns the range to fetch, or an error if EOF would be reached. - fn remaining_block_range(&self) -> Result, AvroError> { - let remaining = self.block_decoder.bytes_remaining() as u64 - + match self.block_decoder.state() { - BlockDecoderState::Data => 16, // Include sync marker - BlockDecoderState::Sync => 0, - state => { - return Err(AvroError::General(format!( - "remaining_block_range called in unexpected state: {state:?}" - ))); - } + let inner: BoxStream<'static, Result> = if finished { + Box::pin(futures::stream::empty()) + } else { + let state = AvroReaderState { + reader, + meta: AvroReaderMeta { + range, + file_size, + decoder, + codec, + sync_marker, + }, }; - - let fetch_end = self.range.end + remaining; - if fetch_end > self.file_size { - return Err(AvroError::EOF( - "Avro block requires more bytes than what exists in the file".into(), - )); - } - - Ok(self.range.end..fetch_end) - } - - /// Terminate the stream after returning this error once. - #[inline] - fn finish_with_error( - &mut self, - error: AvroError, - ) -> Poll>> { - self.reader_state = ReaderState::Finished; - Poll::Ready(Some(Err(error))) - } - - #[inline] - fn start_flushing(&mut self) { - self.reader_state = ReaderState::Flushing; - } - - /// Drain any remaining buffered records from the decoder. - #[inline] - fn poll_flush(&mut self) -> Poll>> { - match self.decoder.flush() { - Ok(Some(batch)) => { - self.reader_state = ReaderState::Flushing; - Poll::Ready(Some(Ok(batch))) - } - Ok(None) => { - self.reader_state = ReaderState::Finished; - Poll::Ready(None) - } - Err(e) => self.finish_with_error(e), - } - } -} - -impl AsyncAvroFileReader { - // The forbid question mark thing shouldn't apply here, as it is within the future, - // so exported this to a separate function. - async fn fetch_bytes(mut reader: R, range: Range) -> Result<(R, Bytes), AvroError> { - let data = reader.get_bytes(range).await?; - Ok((reader, data)) - } - - #[forbid(clippy::question_mark_used)] - fn read_next(&mut self, cx: &mut Context<'_>) -> Poll>> { - loop { - match mem::replace(&mut self.reader_state, ReaderState::InvalidState) { - ReaderState::Idle { reader } => { - let range = self.range.clone(); - if range.start >= range.end { - return self.finish_with_error(AvroError::InvalidArgument(format!( - "Invalid range specified for Avro file: start {} >= end {}, file_size: {}", - range.start, range.end, self.file_size - ))); - } - - let future = Self::fetch_bytes(reader, range).boxed(); - self.reader_state = ReaderState::FetchingData { - future, - next_behaviour: FetchNextBehaviour::ReadSyncMarker, - }; - } - ReaderState::FetchingData { - mut future, - next_behaviour, - } => { - let (reader, data_chunk) = match future.poll_unpin(cx) { - Poll::Ready(Ok(data)) => data, - Poll::Ready(Err(e)) => return self.finish_with_error(e), - Poll::Pending => { - self.reader_state = ReaderState::FetchingData { - future, - next_behaviour, - }; - return Poll::Pending; - } - }; - - match next_behaviour { - FetchNextBehaviour::ReadSyncMarker => { - let sync_marker_pos = data_chunk - .windows(16) - .position(|slice| slice == self.sync_marker); - let block_start = match sync_marker_pos { - Some(pos) => pos + 16, // Move past the sync marker - None => { - // Sync marker not found, valid if we arbitrarily split the file at its end. - self.reader_state = ReaderState::Finished; - return Poll::Ready(None); - } - }; - - self.reader_state = ReaderState::DecodingBlock { - reader, - data: data_chunk.slice(block_start..), - }; - } - FetchNextBehaviour::DecodeVLQHeader => { - let mut data = data_chunk; - - // Feed bytes one at a time until we reach Data state (VLQ header complete) - while !matches!(self.block_decoder.state(), BlockDecoderState::Data) { - if data.is_empty() { - return self.finish_with_error(AvroError::EOF( - "Unexpected EOF while reading Avro block header".into(), - )); - } - let consumed = match self.block_decoder.decode(&data[..1]) { - Ok(consumed) => consumed, - Err(e) => return self.finish_with_error(e), - }; - if consumed == 0 { - return self.finish_with_error(AvroError::General( - "BlockDecoder failed to consume byte during VLQ header parsing" - .into(), - )); - } - data = data.slice(consumed..); - } - - // Now we know the block size. Slice remaining data to what we need. - let bytes_remaining = self.block_decoder.bytes_remaining(); - let data_to_use = data.slice(..data.len().min(bytes_remaining)); - let consumed = match self.block_decoder.decode(&data_to_use) { - Ok(consumed) => consumed, - Err(e) => return self.finish_with_error(e), - }; - if consumed != data_to_use.len() { - return self.finish_with_error(AvroError::General( - "BlockDecoder failed to consume all bytes after VLQ header parsing" - .into(), - )); - } - - // May need more data to finish the block. - let range_to_fetch = match self.remaining_block_range() { - Ok(range) if range.is_empty() => { - // All bytes fetched, move to decoding block directly - self.reader_state = ReaderState::DecodingBlock { - reader, - data: Bytes::new(), - }; - continue; - } - Ok(range) => range, - Err(e) => return self.finish_with_error(e), - }; - - let future = Self::fetch_bytes(reader, range_to_fetch).boxed(); - self.reader_state = ReaderState::FetchingData { - future, - next_behaviour: FetchNextBehaviour::ContinueDecoding, - }; - continue; - } - FetchNextBehaviour::ContinueDecoding => { - self.reader_state = ReaderState::DecodingBlock { - reader, - data: data_chunk, - }; - } - } - } - ReaderState::InvalidState => { - return self.finish_with_error(AvroError::General( - "AsyncAvroFileReader in invalid state".into(), - )); - } - ReaderState::DecodingBlock { reader, mut data } => { - // Try to decode another block from the buffered reader. - let consumed = match self.block_decoder.decode(&data) { - Ok(consumed) => consumed, - Err(e) => return self.finish_with_error(e), - }; - data = data.slice(consumed..); - - // If we reached the end of the block, flush it, and move to read batches. - if let Some(block) = self.block_decoder.flush() { - // Successfully decoded a block. - let block_count = block.count; - let block_data = Bytes::from_owner(if let Some(ref codec) = self.codec { - match codec.decompress(&block.data) { - Ok(decompressed) => decompressed, - Err(e) => return self.finish_with_error(e), - } - } else { - block.data - }); - - // Since we have an active block, move to reading batches - self.reader_state = ReaderState::ReadingBatches { - reader, - data, - block_data, - remaining_in_block: block_count, - }; - continue; - } - - // data should always be consumed unless Finished, if it wasn't, something went wrong - if !data.is_empty() { - return self.finish_with_error(AvroError::General( - "BlockDecoder failed to make progress decoding Avro block".into(), - )); - } - - if matches!(self.block_decoder.state(), BlockDecoderState::Finished) { - // We've already flushed, so if no batch was produced, we are simply done. - self.finishing_partial_block = false; - self.start_flushing(); - continue; - } - - // If we've tried the following stage before, and still can't decode, - // this means the file is truncated or corrupted. - if self.finishing_partial_block { - return self.finish_with_error(AvroError::EOF( - "Unexpected EOF while reading last Avro block".into(), - )); - } - - // Avro splitting case: block is incomplete, we need to: - // 1. Parse the length so we know how much to read - // 2. Fetch more data from the reader - // 3. Create a new block data from the remaining slice and the newly fetched data - // 4. Continue decoding until end of block - self.finishing_partial_block = true; - - // Mid-block, but we don't know how many bytes are missing yet - if matches!( - self.block_decoder.state(), - BlockDecoderState::Count | BlockDecoderState::Size - ) { - // Max VLQ header is 20 bytes (10 bytes each for count and size). - // Fetch just enough to complete it. - const MAX_VLQ_HEADER_SIZE: u64 = 20; - let fetch_end = (self.range.end + MAX_VLQ_HEADER_SIZE).min(self.file_size); - - // If there is nothing more to fetch, error out - if fetch_end == self.range.end { - return self.finish_with_error(AvroError::EOF( - "Unexpected EOF while reading Avro block header".into(), - )); - } - - let range_to_fetch = self.range.end..fetch_end; - self.range.end = fetch_end; // Track that we've fetched these bytes - - let future = Self::fetch_bytes(reader, range_to_fetch).boxed(); - self.reader_state = ReaderState::FetchingData { - future, - next_behaviour: FetchNextBehaviour::DecodeVLQHeader, - }; - continue; - } - - // Otherwise, we're mid-block but know how many bytes are remaining to fetch. - let range_to_fetch = match self.remaining_block_range() { - Ok(range) => range, - Err(e) => return self.finish_with_error(e), - }; - - let future = Self::fetch_bytes(reader, range_to_fetch).boxed(); - self.reader_state = ReaderState::FetchingData { - future, - next_behaviour: FetchNextBehaviour::ContinueDecoding, - }; - continue; - } - ReaderState::ReadingBatches { - reader, - data, - mut block_data, - mut remaining_in_block, - } => { - let (consumed, records_decoded) = - match self.decoder.decode_block(&block_data, remaining_in_block) { - Ok((consumed, records_decoded)) => (consumed, records_decoded), - Err(e) => return self.finish_with_error(e), - }; - - remaining_in_block -= records_decoded; - - if remaining_in_block == 0 { - if data.is_empty() { - // No more data to read, drain remaining buffered records - self.start_flushing(); - } else { - // Finished this block, move to decode next block in the next iteration - self.reader_state = ReaderState::DecodingBlock { reader, data }; - } - } else { - // Still more records to decode in this block, slice the already-read data and stay in this state - block_data = block_data.slice(consumed..); - self.reader_state = ReaderState::ReadingBatches { - reader, - data, - block_data, - remaining_in_block, - }; - } - - // We have a full batch ready, emit it - // (This is not mutually exclusive with the block being finished, so the state change is valid) - if self.decoder.batch_is_full() { - return match self.decoder.flush() { - Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), - Ok(None) => self.finish_with_error(AvroError::General( - "Decoder reported a full batch, but flush returned None".into(), - )), - Err(e) => self.finish_with_error(e), - }; - } - } - ReaderState::Flushing => { - return self.poll_flush(); - } - ReaderState::Finished => { - // Terminal: once finished (including after an error), always yield None - self.reader_state = ReaderState::Finished; - return Poll::Ready(None); - } - } + Box::pin(make_stream(state)) + }; + Self { + inner, + _phantom: PhantomData, } } } -// To maintain compatibility with the expected stream results in the ecosystem, this returns ArrowError. impl Stream for AsyncAvroFileReader { type Item = Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.read_next(cx).map_err(Into::into) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::get_mut(self).inner.as_mut().poll_next(cx) } } @@ -554,7 +372,9 @@ mod tests { use arrow_array::types::{Int32Type, Int64Type}; use arrow_array::*; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; - use futures::{StreamExt, TryStreamExt}; + use futures::future::BoxFuture; + use futures::{FutureExt, StreamExt, TryStreamExt}; + use object_store::ObjectStore; use object_store::local::LocalFileSystem; use object_store::path::Path; use object_store::{ObjectStore, ObjectStoreExt}; @@ -1964,4 +1784,209 @@ mod tests { let f1_1_field = f1_struct.column_by_name("f1_1").unwrap(); assert_eq!(f1_1_field.data_type(), &DataType::Utf8View); } + + /// A test [`AsyncFileReader`] that returns the byte stream as small fixed-size chunks, + /// simulating a real object store that streams HTTP responses in multiple pieces. + /// + /// This exercises the multi-chunk streaming path: when a block completes mid-stream + /// the remaining chunks must not be discarded. + struct ChunkedReader { + data: bytes::Bytes, + chunk_size: usize, + } + + impl ChunkedReader { + fn new(data: bytes::Bytes, chunk_size: usize) -> Self { + Self { data, chunk_size } + } + } + + impl AsyncFileReader for ChunkedReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + let slice = self.data.slice(range.start as usize..range.end as usize); + async move { Ok(slice) }.boxed() + } + + fn get_stream( + &mut self, + range: Range, + ) -> BoxFuture<'_, Result>, AvroError>> { + let range_data = self.data.slice(range.start as usize..range.end as usize); + let chunk_size = self.chunk_size; + + let mut chunks: Vec> = Vec::new(); + let mut pos = 0; + let len = range_data.len(); + while pos < len { + let end = (pos + chunk_size).min(len); + chunks.push(Ok(range_data.slice(pos..end))); + pos = end; + } + + // The stream is built from owned `Bytes` clones so it is in fact 'static, + // but we declare '_ to match the trait signature. + let stream: BoxStream<'_, _> = futures::stream::iter(chunks).boxed(); + std::future::ready(Ok(stream)).boxed() + } + } + + /// Regression test for the bug where multi-chunk streaming discards bytes for subsequent blocks. + /// + /// When an Avro block completes mid-stream the remaining stream chunks contain bytes for the + /// next block. The old code recovered the reader (discarding those chunks) and then tried to + /// fetch from `range.end`, which is past EOF when the range covers the full file — producing: + /// "Avro block requires more bytes than what exists in the file" + #[tokio::test] + async fn test_multi_chunk_stream_multiple_blocks() { + use crate::writer::AvroWriter; + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + + // Write 3 separate batches → 3 separate Avro blocks in the file. + let mut buffer: Vec = Vec::new(); + let mut writer = AvroWriter::new(&mut buffer, schema.as_ref().clone()).unwrap(); + for i in 0..3_i32 { + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![i])) as _], + ) + .unwrap(); + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + + let bytes = bytes::Bytes::from(buffer); + let file_size = bytes.len() as u64; + + // Use a tiny chunk size so every block boundary falls in the middle of a stream. + let reader = ChunkedReader::new(bytes, 32); + let batches: Vec = AsyncAvroFileReader::builder(reader, file_size, 1024) + .try_build() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); + + let id_values: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(id_values, vec![0, 1, 2]); + } + + /// Regression test: after a partial-block fetch (`finishing_partial_block = true`), + /// the reader must stop — no further blocks beyond the assigned range should be read. + /// + /// Uses 4 blocks and a range ending mid-block-2. Blocks 1 and 2 must be returned; + /// blocks 3 and 4 must be skipped (they belong to the next Spark task's range). + #[tokio::test] + async fn test_range_ends_mid_block_no_further_blocks_read() { + use crate::writer::AvroWriter; + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + const ROWS_PER_BLOCK: i32 = 500; + const NUM_BLOCKS: i32 = 4; + + // Write 4 separate batches → 4 separate Avro blocks. + let mut buffer: Vec = Vec::new(); + let mut writer = AvroWriter::new(&mut buffer, schema.as_ref().clone()).unwrap(); + for block in 0..NUM_BLOCKS { + let start = block * ROWS_PER_BLOCK; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from( + (start..start + ROWS_PER_BLOCK).collect::>(), + )) as _], + ) + .unwrap(); + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + + let bytes = bytes::Bytes::from(buffer); + let file_size = bytes.len() as u64; + + // file_size/2 falls in the middle of block 2 (blocks 1+2 occupy roughly half the + // file after the small header; each block is ~1 KB of zigzag-encoded ints). + // Block 2 needs a partial-block fetch; blocks 3 and 4 must NOT be read. + let range = 0..(file_size / 2); + + let reader = ChunkedReader::new(bytes.clone(), 32); + let batches: Vec = AsyncAvroFileReader::builder(reader, file_size, 1024) + .with_range(range) + .try_build() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + // Exactly blocks 1 and 2 should be returned (1000 rows total). + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + (ROWS_PER_BLOCK * 2) as usize, + "expected blocks 1+2 only; got {total_rows} rows (file_size={file_size})" + ); + } + + /// Verifies that when a range ends mid-block, the reader completes that block via a + /// partial-block fetch and returns exactly its rows — the following block is not read. + #[tokio::test] + async fn test_range_ending_mid_block_returns_correct_rows() { + use crate::writer::AvroWriter; + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + const ROWS_PER_BLOCK: i32 = 500; + + // Two blocks: block 1 will be completed via a partial-block fetch; block 2 must + // not be touched (it belongs to the next Spark range). + let mut buffer: Vec = Vec::new(); + let mut writer = AvroWriter::new(&mut buffer, schema.as_ref().clone()).unwrap(); + for block in 0..2_i32 { + let start = block * ROWS_PER_BLOCK; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from( + (start..start + ROWS_PER_BLOCK).collect::>(), + )) as _], + ) + .unwrap(); + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + + let bytes = bytes::Bytes::from(buffer); + let file_size = bytes.len() as u64; + + // file_size/4 lands in the middle of block 1 (the header is small relative to + // the ~1 KB blocks), so a partial-block fetch is required to complete it. + let range = 0..(file_size / 4); + + let reader = ChunkedReader::new(bytes.clone(), 32); + let batches: Vec = AsyncAvroFileReader::builder(reader, file_size, 1024) + .with_range(range) + .try_build() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, ROWS_PER_BLOCK as usize, + "expected block 1 only; got {total_rows} rows (file_size={file_size})" + ); + } } diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index 44a4abf1a282..585987dfcab8 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -19,10 +19,12 @@ use crate::errors::AvroError; use crate::reader::async_reader::AsyncFileReader; use bytes::Bytes; use futures::future::BoxFuture; -use futures::{FutureExt, TryFutureExt}; +use futures::stream::BoxStream; +use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; use object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path; +use object_store::{GetOptions, GetRange, ObjectStore}; use std::error::Error; use std::ops::Range; use std::sync::Arc; @@ -62,12 +64,12 @@ impl AvroObjectReader { fn spawn(&self, f: F) -> BoxFuture<'_, Result> where - F: for<'a> FnOnce(&'a Arc, &'a Path) -> BoxFuture<'a, Result> - + Send - + 'static, + F: FnOnce(Arc, Path) -> BoxFuture<'static, Result> + Send + 'static, O: Send + 'static, E: Error + Send + 'static, { + let path = self.path.clone(); + let store = Arc::clone(&self.store); match &self.runtime { Some(handle) => { let path = self.path.clone(); @@ -95,6 +97,33 @@ impl AsyncFileReader for AvroObjectReader { self.spawn(|store, path| async move { store.get_range(path, range).await }.boxed()) } + fn get_stream( + &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| ArrowError::from_external_error(Box::new(e)))?; + let stream = get_result + .into_stream() + .map_err(|e| ArrowError::from_external_error(Box::new(e))) + .boxed(); + Ok(stream) + } + .boxed() + } + fn get_byte_ranges( &mut self, ranges: Vec>, @@ -102,6 +131,6 @@ impl AsyncFileReader for AvroObjectReader { where Self: Send, { - self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed()) + self.spawn(|store, path| async move { store.get_ranges(&path, &ranges).await }.boxed()) } } From 76e2bcf97e408ec316594111a57706059b6433be Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 1 Apr 2026 17:16:34 +0300 Subject: [PATCH 02/11] fix: feature-gate the async imports --- arrow-avro/Cargo.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index 681304200449..0302b6d6582e 100644 --- a/arrow-avro/Cargo.toml +++ b/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,8 +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"] } -async-stream = "0.3.6" +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 } From d7febee00cff184b2c131df0b7dbd8ef9f5036c9 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 1 Apr 2026 17:31:20 +0300 Subject: [PATCH 03/11] fix: compilation failures with object-store --- .../reader/async_reader/async_file_reader.rs | 2 +- arrow-avro/src/reader/async_reader/builder.rs | 9 ++++---- arrow-avro/src/reader/async_reader/mod.rs | 22 +++++++++++++------ arrow-avro/src/reader/async_reader/store.rs | 16 ++++++-------- 4 files changed, 27 insertions(+), 22 deletions(-) diff --git a/arrow-avro/src/reader/async_reader/async_file_reader.rs b/arrow-avro/src/reader/async_reader/async_file_reader.rs index 18facd536f82..e13d0ce00687 100644 --- a/arrow-avro/src/reader/async_reader/async_file_reader.rs +++ b/arrow-avro/src/reader/async_reader/async_file_reader.rs @@ -98,7 +98,7 @@ impl AsyncFileReader for T { let ranged_reader = self.take(range.end - range.start); let stream = ReaderStream::new(ranged_reader) - .map_err(ArrowError::from) + .map_err(AvroError::from) .boxed(); Ok(stream) diff --git a/arrow-avro/src/reader/async_reader/builder.rs b/arrow-avro/src/reader/async_reader/builder.rs index ddead347d279..f85c71a752f7 100644 --- a/arrow-avro/src/reader/async_reader/builder.rs +++ b/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}; @@ -273,12 +272,12 @@ impl ReaderBuilder { // 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 finished = range.start == range.end || self.header_len == range.end; - let codec = self.header.compression()?; - let sync_marker = self.header.sync(); + 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.inner.reader, + self.reader, range, self.file_size, decoder, diff --git a/arrow-avro/src/reader/async_reader/mod.rs b/arrow-avro/src/reader/async_reader/mod.rs index f467b364553c..98dca498f22e 100644 --- a/arrow-avro/src/reader/async_reader/mod.rs +++ b/arrow-avro/src/reader/async_reader/mod.rs @@ -26,9 +26,8 @@ use crate::reader::block::{BlockDecoder, BlockDecoderState}; use arrow_array::RecordBatch; use arrow_schema::{ArrowError, SchemaRef}; use bytes::{Bytes, BytesMut}; -use futures::future::{BoxFuture, BoxStream}; -use futures::{FutureExt, Stream, StreamExt}; -use std::mem; +use futures::{Stream, StreamExt, stream::BoxStream}; +use std::marker::PhantomData; use std::ops::Range; use std::pin::Pin; use std::task::{Context, Poll}; @@ -42,7 +41,6 @@ pub use builder::{ReaderBuilder, read_header_info}; #[cfg(feature = "object_store")] mod store; -use crate::errors::AvroError; #[cfg(feature = "object_store")] pub use store::AvroObjectReader; @@ -311,6 +309,7 @@ fn make_stream( /// ``` pub struct AsyncAvroFileReader { inner: BoxStream<'static, Result>, + schema: SchemaRef, _phantom: PhantomData, } @@ -319,6 +318,13 @@ impl AsyncAvroFileReader { pub fn builder(reader: R, file_size: u64, batch_size: usize) -> ReaderBuilder { ReaderBuilder::new(reader, file_size, batch_size) } + + /// Returns the Arrow schema for batches produced by this reader. + /// + /// The schema is determined by the writer schema in the file and the reader schema provided to the builder. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } } impl AsyncAvroFileReader { @@ -331,6 +337,7 @@ impl AsyncAvroFileReader { sync_marker: [u8; 16], finished: bool, ) -> Self { + let schema = decoder.schema(); let inner: BoxStream<'static, Result> = if finished { Box::pin(futures::stream::empty()) } else { @@ -348,6 +355,7 @@ impl AsyncAvroFileReader { }; Self { inner, + schema, _phantom: PhantomData, } } @@ -365,6 +373,7 @@ impl Stream for AsyncAvroFileReader { mod tests { use super::*; use crate::codec::Tz; + use crate::errors::AvroError; use crate::schema::{ AVRO_NAME_METADATA_KEY, AVRO_NAMESPACE_METADATA_KEY, AvroSchema, SCHEMA_METADATA_KEY, }; @@ -374,7 +383,6 @@ mod tests { use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; - use object_store::ObjectStore; use object_store::local::LocalFileSystem; use object_store::path::Path; use object_store::{ObjectStore, ObjectStoreExt}; @@ -1802,7 +1810,7 @@ mod tests { } impl AsyncFileReader for ChunkedReader { - fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { let slice = self.data.slice(range.start as usize..range.end as usize); async move { Ok(slice) }.boxed() } @@ -1814,7 +1822,7 @@ mod tests { let range_data = self.data.slice(range.start as usize..range.end as usize); let chunk_size = self.chunk_size; - let mut chunks: Vec> = Vec::new(); + let mut chunks: Vec> = Vec::new(); let mut pos = 0; let len = range_data.len(); while pos < len { diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs index 585987dfcab8..3cf34d3ec97b 100644 --- a/arrow-avro/src/reader/async_reader/store.rs +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -21,10 +21,8 @@ use bytes::Bytes; use futures::future::BoxFuture; use futures::stream::BoxStream; use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; -use object_store::ObjectStore; -use object_store::ObjectStoreExt; use object_store::path::Path; -use object_store::{GetOptions, GetRange, ObjectStore}; +use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; use std::error::Error; use std::ops::Range; use std::sync::Arc; @@ -64,12 +62,12 @@ impl AvroObjectReader { fn spawn(&self, f: F) -> BoxFuture<'_, Result> where - F: FnOnce(Arc, Path) -> BoxFuture<'static, Result> + Send + 'static, + F: for<'a> FnOnce(&'a Arc, &'a Path) -> BoxFuture<'a, Result> + + Send + + 'static, O: Send + 'static, E: Error + Send + 'static, { - let path = self.path.clone(); - let store = Arc::clone(&self.store); match &self.runtime { Some(handle) => { let path = self.path.clone(); @@ -114,10 +112,10 @@ impl AsyncFileReader for AvroObjectReader { .store .get_opts(&self.path, options) .await - .map_err(|e| ArrowError::from_external_error(Box::new(e)))?; + .map_err(|e| AvroError::External(Box::new(e)))?; let stream = get_result .into_stream() - .map_err(|e| ArrowError::from_external_error(Box::new(e))) + .map_err(|e| AvroError::External(Box::new(e))) .boxed(); Ok(stream) } @@ -131,6 +129,6 @@ impl AsyncFileReader for AvroObjectReader { where Self: Send, { - self.spawn(|store, path| async move { store.get_ranges(&path, &ranges).await }.boxed()) + self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed()) } } From 5081cde6f1b911ec3394939072e34750a22adcd7 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 15 Apr 2026 10:06:25 +0300 Subject: [PATCH 04/11] fix: add back async-stream to dev-dependencies --- arrow-avro/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index 0302b6d6582e..d76f294601f3 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -92,6 +92,7 @@ criterion = { workspace = true, default-features = false } tempfile = "3.3" arrow = { workspace = true, features = ["prettyprint"] } futures = "0.3.31" +async-stream = "0.3.6" apache-avro = "0.21.0" num-bigint = "0.4" object_store = { workspace = true, features = ["fs"] } From fb41f90b61376fe4e485817f3241db09cba5f183 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 8 May 2026 06:52:04 +0300 Subject: [PATCH 05/11] 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 06/11] 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 07/11] 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 08/11] 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 { From 4cbc8ba09d9c28815bec5c9ba88c595386898433 Mon Sep 17 00:00:00 2001 From: Ariel Miculas-Trif Date: Mon, 11 May 2026 14:08:16 +0300 Subject: [PATCH 09/11] fix: mark tokio-stream optional Co-authored-by: Mikhail Zabaluev --- arrow-avro/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index 6a5e49cf5196..14c5865231bc 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -46,7 +46,7 @@ small_decimals = [] avro_custom_types = ["dep:arrow-select"] # Enable async APIs -async = ["futures", "tokio", "tokio-util", "async-stream"] +async = ["futures", "tokio", "tokio-stream", "tokio-util", "async-stream"] # Enable object_store integration object_store = ["dep:object_store", "async"] From a02bfacfd1ad16c69c29c80c83ea8edccd03cd22 Mon Sep 17 00:00:00 2001 From: Ariel Miculas-Trif Date: Mon, 11 May 2026 14:08:24 +0300 Subject: [PATCH 10/11] fix: mark tokio-stream optional Co-authored-by: Mikhail Zabaluev --- arrow-avro/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index 14c5865231bc..768081db7cf9 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -77,7 +77,7 @@ 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", "sync"] } -tokio-stream = { version = "0.1", default-features = false } +tokio-stream = { version = "0.1", optional = true, default-features = false } tokio-util = { version = "0.7.18", default-features = false, features = ["io"], optional = true } async-stream = { version = "0.3.6", optional = true } From d741ad467d07f670325f6d1350333b8e50d6bc4f Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Wed, 22 Jul 2026 04:00:06 +0300 Subject: [PATCH 11/11] chore: commit Cargo.lock changes --- Cargo.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 74d0759d5f9d..9dfbdcfd4f0f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -262,6 +262,8 @@ dependencies = [ "strum_macros 0.28.0", "tempfile", "tokio", + "tokio-stream", + "tokio-util", "uuid", "zstd", ]