diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index a4699a2efa80..5e5efd78d0e3 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -45,16 +45,26 @@ sha256 = ["dep:sha2"] small_decimals = [] avro_custom_types = ["dep:arrow-select"] +# Enable async APIs +async = ["futures", "tokio"] +# Enable object_store integration +object_store = ["dep:object_store", "async"] + [dependencies] arrow-schema = { workspace = true } arrow-buffer = { workspace = true } arrow-array = { workspace = true } arrow-select = { workspace = true, optional = true } + +object_store = { version = "0.13", default-features = false, optional = true } + +bytes = { version = "1.11.0", default-features = false, features = ["std"] } serde_json = { version = "1.0", default-features = false, features = ["std"] } serde = { version = "1.0.188", features = ["derive"] } flate2 = { version = "1.0", default-features = false, features = [ "rust_backend", ], optional = true } +futures = { version = "0.3", default-features = false, features = ["std"], optional = true } snap = { version = "1.0", default-features = false, optional = true } zstd = { version = "0.13", default-features = false, optional = true } bzip2 = { version = "0.6.0", optional = true } @@ -66,10 +76,11 @@ indexmap = "2.10" rand = "0.9" md5 = { version = "0.8", optional = true } sha2 = { version = "0.10", optional = true } -bytes = "1.11" +tokio = { version = "1.0", optional = true, default-features = false, features = ["macros", "rt", "io-util"] } [dev-dependencies] arrow-data = { workspace = true } +bytes = { version = "1.11.0", default-features = false, features = ["std"] } rand = { version = "0.9.1", default-features = false, features = [ "std", "std_rng", @@ -82,7 +93,9 @@ futures = "0.3.31" async-stream = "0.3.6" apache-avro = "0.21.0" num-bigint = "0.4" +object_store = { version = "0.13", default-features = false, features = ["fs"] } once_cell = "1.21.3" +tokio = { version = "1.0", default-features = false, features = ["macros", "rt-multi-thread", "io-util", "fs"] } [[bench]] name = "avro_reader" diff --git a/arrow-avro/README.md b/arrow-avro/README.md index 85fd76094755..c5776c125b0a 100644 --- a/arrow-avro/README.md +++ b/arrow-avro/README.md @@ -51,7 +51,7 @@ Disable defaults and pick only what you need (see **Feature Flags**): ```toml [dependencies] -arrow-avro = { version = "57.0.0", default-features = false, features = ["deflate", "snappy"] } +arrow-avro = { version = "58.0.0", default-features = false, features = ["deflate", "snappy"] } ``` --- @@ -105,6 +105,36 @@ fn main() -> anyhow::Result<()> { See the crate docs for runnable SOE and Confluent round‑trip examples. +### Async reading from object stores (`object_store` feature) + +```rust,ignore +use std::sync::Arc; +use arrow_avro::reader::{AsyncAvroFileReader, AvroObjectReader}; +use futures::TryStreamExt; +use object_store::ObjectStore; +use object_store::local::LocalFileSystem; +use object_store::path::Path; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let store: Arc = Arc::new(LocalFileSystem::new()); + let path = Path::from("data/example.avro"); + + let meta = store.head(&path).await?; + let reader = AvroObjectReader::new(store, path); + + let stream = AsyncAvroFileReader::builder(reader, meta.size, 1024) + .try_build() + .await?; + + let batches: Vec<_> = stream.try_collect().await?; + for batch in batches { + println!("rows: {}", batch.num_rows()); + } + Ok(()) +} +``` + --- ## Feature Flags (what they do and when to use them) @@ -128,15 +158,22 @@ See the crate docs for runnable SOE and Confluent round‑trip examples. * Only **OCF** uses these codecs (they compress per‑block). They do **not** apply to raw Avro frames used by Confluent wire format or SOE. The crate’s `compression` module is specifically for **OCF blocks**. * `deflate` uses `flate2` with the `rust_backend` (no system zlib required). +### Async & Object Store + +| Feature | Default | What it enables | When to use | +|----------------|--------:|-----------------------------------------------------------------------------|-------------------------------------------------------------------------------| +| `async` | ⬜ | Async APIs for reading Avro via `futures` and `tokio` | Enable for non-blocking async Avro reading with `AsyncAvroFileReader`. | +| `object_store` | ⬜ | Integration with `object_store` crate (implies `async`) | Enable for reading Avro from cloud storage (S3, GCS, Azure Blob, etc.). | + ### Schema fingerprints & custom logical type helpers -| Feature | Default | What it enables | When to use | +| Feature | Default | What it enables | When to use | |-----------------------------|--------:|----------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------| -| `md5` | ⬜ | `md5` dep for optional **MD5** schema fingerprints | If you want to compute MD5 fingerprints of writer schemas (i.e. for custom prefixing/validation). | -| `sha256` | ⬜ | `sha2` dep for optional **SHA‑256** schema fingerprints | If you prefer longer fingerprints; affects max prefix length (i.e. when framing). | +| `md5` | ⬜ | `md5` dep for optional **MD5** schema fingerprints | If you want to compute MD5 fingerprints of writer schemas (i.e. for custom prefixing/validation). | +| `sha256` | ⬜ | `sha2` dep for optional **SHA‑256** schema fingerprints | If you prefer longer fingerprints; affects max prefix length (i.e. when framing). | | `small_decimals` | ⬜ | Extra handling for **small decimal** logical types (`Decimal32` and `Decimal64`) | If your Avro `decimal` values are small and you want more compact Arrow representations. | -| `avro_custom_types` | ⬜ | Annotates Avro values using Arrow specific custom logical types | Enable when you need arrow-avro to reinterpret certain Avro fields as Arrow types that Avro doesn’t natively model. | -| `canonical_extension_types` | ⬜ | Re‑exports Arrow’s canonical extension types support from `arrow-schema` | Enable if your workflow uses Arrow [canonical extension types] and you want `arrow-avro` to respect them. | +| `avro_custom_types` | ⬜ | Annotates Avro values using Arrow specific custom logical types | Enable when you need arrow-avro to reinterpret certain Avro fields as Arrow types that Avro doesn't natively model. | +| `canonical_extension_types` | ⬜ | Re‑exports Arrow's canonical extension types support from `arrow-schema` | Enable if your workflow uses Arrow [canonical extension types] and you want `arrow-avro` to respect them. | [canonical extension types]: https://arrow.apache.org/docs/format/CanonicalExtensions.html @@ -149,17 +186,22 @@ See the crate docs for runnable SOE and Confluent round‑trip examples. * Minimal, fast build (common pipelines): ```toml - arrow-avro = { version = "56", default-features = false, features = ["deflate", "snappy"] } + arrow-avro = { version = "58", default-features = false, features = ["deflate", "snappy"] } ``` * Include Zstandard too (modern data lakes): ```toml - arrow-avro = { version = "56", default-features = false, features = ["deflate", "snappy", "zstd"] } + arrow-avro = { version = "58", default-features = false, features = ["deflate", "snappy", "zstd"] } + ``` +* Async reading from object stores (S3, GCS, etc.): + + ```toml + arrow-avro = { version = "58", features = ["object_store"] } ``` * Fingerprint helpers: ```toml - arrow-avro = { version = "56", features = ["md5", "sha256"] } + arrow-avro = { version = "58", features = ["md5", "sha256"] } ``` --- diff --git a/arrow-avro/src/lib.rs b/arrow-avro/src/lib.rs index eb04ee8fa6fc..da451ea1460e 100644 --- a/arrow-avro/src/lib.rs +++ b/arrow-avro/src/lib.rs @@ -123,16 +123,56 @@ //! # Ok(()) } //! ``` //! +//! ## `async` Reading (`async` feature) +//! +//! The [`reader`] module provides async APIs for reading Avro files when the `async` +//! feature is enabled. +//! +//! [`AsyncAvroFileReader`] implements `Stream>`, +//! allowing efficient async streaming of record batches. When the `object_store` feature +//! is enabled, [`AvroObjectReader`] provides integration with object storage services +//! such as S3 via the [object_store] crate. +//! +//! ```ignore +//! use std::sync::Arc; +//! use arrow_avro::reader::{AsyncAvroFileReader, AvroObjectReader}; +//! use futures::TryStreamExt; +//! use object_store::ObjectStore; +//! use object_store::local::LocalFileSystem; +//! use object_store::path::Path; +//! +//! # async fn example() -> Result<(), Box> { +//! let store: Arc = Arc::new(LocalFileSystem::new()); +//! let path = Path::from("data/example.avro"); +//! let meta = store.head(&path).await?; +//! +//! let reader = AvroObjectReader::new(store, path); +//! let stream = AsyncAvroFileReader::builder(reader, meta.size, 1024) +//! .try_build() +//! .await?; +//! +//! let batches: Vec<_> = stream.try_collect().await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! [object_store]: https://docs.rs/object_store/latest/object_store/ +//! //! --- //! //! ### Modules //! //! - [`reader`]: read Avro (OCF, SOE, Confluent) into Arrow `RecordBatch`es. +//! - With the `async` feature: [`AsyncAvroFileReader`] for async streaming reads. +//! - With the `object_store` feature: [`AvroObjectReader`] for reading from cloud storage. //! - [`writer`]: write Arrow `RecordBatch`es as Avro (OCF, SOE, Confluent, Apicurio). //! - [`schema`]: Avro schema parsing / fingerprints / registries. //! - [`compression`]: codecs used for **OCF block compression** (i.e., Deflate, Snappy, Zstandard, BZip2, and XZ). //! - [`codec`]: internal Avro-Arrow type conversion and row decode/encode plans. //! +//! [`AsyncAvroFileReader`]: reader::AsyncAvroFileReader +//! [`AvroObjectReader`]: reader::AvroObjectReader +//! //! ### Features //! //! **OCF compression (enabled by default)** @@ -142,6 +182,11 @@ //! - `bzip2` — enable BZip2 block compression. //! - `xz` — enable XZ/LZMA block compression. //! +//! **Async & Object Store (opt‑in)** +//! - `async` — enable async APIs for reading Avro (`AsyncAvroFileReader`, `AsyncFileReader` trait). +//! - `object_store` — enable integration with the [`object_store`] crate for reading Avro +//! from cloud storage (S3, GCS, Azure Blob, etc.) via `AvroObjectReader`. Implies `async`. +//! //! **Schema fingerprints & helpers (opt‑in)** //! - `md5` — enable MD5 writer‑schema fingerprints. //! - `sha256` — enable SHA‑256 writer‑schema fingerprints. @@ -156,6 +201,8 @@ //! - OCF compression codecs apply only to **Object Container Files**; they do not affect Avro //! single object encodings. //! +//! [`object_store`]: https://docs.rs/object_store/latest/object_store/ +//! //! [canonical extension types]: https://arrow.apache.org/docs/format/CanonicalExtensions.html //! //! [Apache Arrow]: https://arrow.apache.org/ diff --git a/arrow-avro/src/reader/async_reader/async_file_reader.rs b/arrow-avro/src/reader/async_reader/async_file_reader.rs new file mode 100644 index 000000000000..1257a2f3dd4d --- /dev/null +++ b/arrow-avro/src/reader/async_reader/async_file_reader.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::AvroError; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use std::io::SeekFrom; +use std::ops::Range; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; + +/// The asynchronous interface used by [`super::AsyncAvroFileReader`] to read avro files +/// +/// Notes: +/// +/// 1. There is a default implementation for types that implement [`AsyncRead`] +/// and [`AsyncSeek`], for example [`tokio::fs::File`]. +/// +/// 2. [`super::AvroObjectReader`], available when the `object_store` crate feature +/// is enabled, implements this interface for [`ObjectStore`]. +/// +/// [`ObjectStore`]: object_store::ObjectStore +/// +/// [`tokio::fs::File`]: https://docs.rs/tokio/latest/tokio/fs/struct.File.html +pub trait AsyncFileReader: Send { + /// Retrieve the bytes in `range` + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result>; + + /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result, AvroError>> { + async move { + let mut result = Vec::with_capacity(ranges.len()); + + for range in ranges.into_iter() { + let data = self.get_bytes(range).await?; + result.push(data); + } + + Ok(result) + } + .boxed() + } +} + +/// This allows Box to be used as an AsyncFileReader, +impl AsyncFileReader for Box { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.as_mut().get_bytes(range) + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result, AvroError>> { + self.as_mut().get_byte_ranges(ranges) + } +} + +impl AsyncFileReader for T { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + async move { + self.seek(SeekFrom::Start(range.start)).await?; + + let to_read = range.end - range.start; + let mut buffer = Vec::with_capacity(to_read as usize); + let read = self.take(to_read).read_to_end(&mut buffer).await?; + if read as u64 != to_read { + return Err(AvroError::EOF(format!( + "expected to read {} bytes, got {}", + to_read, read + ))); + } + + Ok(buffer.into()) + } + .boxed() + } +} diff --git a/arrow-avro/src/reader/async_reader/builder.rs b/arrow-avro/src/reader/async_reader/builder.rs new file mode 100644 index 000000000000..0f9a7abf1cd4 --- /dev/null +++ b/arrow-avro/src/reader/async_reader/builder.rs @@ -0,0 +1,256 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::codec::AvroFieldBuilder; +use crate::errors::AvroError; +use crate::reader::async_reader::ReaderState; +use crate::reader::header::{Header, HeaderDecoder}; +use crate::reader::record::RecordDecoder; +use crate::reader::{AsyncAvroFileReader, AsyncFileReader, Decoder}; +use crate::schema::{AvroSchema, FingerprintAlgorithm, SCHEMA_METADATA_KEY}; +use indexmap::IndexMap; +use std::ops::Range; + +const DEFAULT_HEADER_SIZE_HINT: u64 = 16 * 1024; // 16 KB + +/// Builder for an asynchronous Avro file reader. +pub struct ReaderBuilder { + reader: R, + file_size: u64, + batch_size: usize, + range: Option>, + reader_schema: Option, + projection: Option>, + header_size_hint: Option, + utf8_view: bool, + strict_mode: bool, +} + +impl ReaderBuilder { + pub(super) fn new(reader: R, file_size: u64, batch_size: usize) -> Self { + Self { + reader, + file_size, + batch_size, + range: None, + reader_schema: None, + projection: None, + header_size_hint: None, + utf8_view: false, + strict_mode: false, + } + } + + /// Specify a byte range to read from the Avro file. + /// If this is provided, the reader will read all the blocks within the specified range, + /// if the range ends mid-block, it will attempt to fetch the remaining bytes to complete the block, + /// but no further blocks will be read. + /// If this is omitted, the full file will be read. + pub fn with_range(self, range: Range) -> Self { + Self { + range: Some(range), + ..self + } + } + + /// Specify a reader schema to use when reading the Avro file. + /// This can be useful to project specific columns or handle schema evolution. + /// If this is not provided, the schema will be derived from the Arrow schema provided. + pub fn with_reader_schema(self, reader_schema: AvroSchema) -> Self { + Self { + reader_schema: Some(reader_schema), + ..self + } + } + + /// Specify a projection of column indices to read from the Avro file. + /// This can help optimize reading by only fetching the necessary columns. + pub fn with_projection(self, projection: Vec) -> Self { + Self { + projection: Some(projection), + ..self + } + } + + /// Provide a hint for the expected size of the Avro header in bytes. + /// This can help optimize the initial read operation when fetching the header. + pub fn with_header_size_hint(self, hint: u64) -> Self { + Self { + header_size_hint: Some(hint), + ..self + } + } + + /// Enable usage of Utf8View types when reading string data. + pub fn with_utf8_view(self, utf8_view: bool) -> Self { + Self { utf8_view, ..self } + } + + /// Enable strict mode for schema validation and data reading. + pub fn with_strict_mode(self, strict_mode: bool) -> Self { + Self { + strict_mode, + ..self + } + } +} + +impl ReaderBuilder { + async fn read_header(&mut self) -> Result<(Header, u64), AvroError> { + let mut decoder = HeaderDecoder::default(); + let mut position = 0; + loop { + let range_to_fetch = position + ..(position + self.header_size_hint.unwrap_or(DEFAULT_HEADER_SIZE_HINT)) + .min(self.file_size); + + // Maybe EOF after the header, no actual data + if range_to_fetch.is_empty() { + break; + } + + let current_data = self + .reader + .get_bytes(range_to_fetch.clone()) + .await + .map_err(|err| { + AvroError::General(format!( + "Error fetching Avro header from file reader: {err}" + )) + })?; + if current_data.is_empty() { + return Err(AvroError::EOF( + "Unexpected EOF while fetching header data".into(), + )); + } + + let read = current_data.len(); + let decoded = decoder.decode(¤t_data)?; + if decoded != read { + position += decoded as u64; + break; + } + position += read as u64; + } + + decoder + .flush() + .map(|header| (header, position)) + .ok_or_else(|| AvroError::EOF("Unexpected EOF while reading Avro header".into())) + } + + /// 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> { + if self.file_size == 0 { + return Err(AvroError::InvalidArgument("File size cannot be 0".into())); + } + + // Start by reading the header from the beginning of the avro file + // take the writer schema from the header + let (header, header_len) = self.read_header().await?; + let writer_schema = { + let raw = header.get(SCHEMA_METADATA_KEY).ok_or_else(|| { + AvroError::ParseError("No Avro schema present in file header".to_string()) + })?; + let json_string = std::str::from_utf8(raw) + .map_err(|e| { + AvroError::ParseError(format!("Invalid UTF-8 in Avro schema header: {e}")) + })? + .to_string(); + AvroSchema::new(json_string) + }; + + // If projection exists, project the reader schema, + // if no reader schema is provided, parse it from the header(get the raw writer schema), and project that + // this projected schema will be the schema used for reading. + let projected_reader_schema = self + .projection + .as_deref() + .map(|projection| { + let base_schema = if let Some(reader_schema) = &self.reader_schema { + reader_schema + } else { + &writer_schema + }; + base_schema.project(projection) + }) + .transpose()?; + + // Use either the projected reader schema or the original reader schema(if no projection) + // (both optional, at worst no reader schema is provided, in which case we read with the writer schema) + let effective_reader_schema = projected_reader_schema + .as_ref() + .or(self.reader_schema.as_ref()) + .map(|s| s.schema()) + .transpose()?; + + let root = { + let writer_schema = writer_schema.schema()?; + let mut builder = AvroFieldBuilder::new(&writer_schema); + if let Some(reader_schema) = &effective_reader_schema { + builder = builder.with_reader_schema(reader_schema); + } + builder + .with_utf8view(self.utf8_view) + .with_strict_mode(self.strict_mode) + .build() + }?; + + let record_decoder = RecordDecoder::try_new_with_options(root.data_type())?; + let decoder = Decoder::from_parts( + self.batch_size, + record_decoder, + None, + IndexMap::new(), + FingerprintAlgorithm::Rabin, + ); + let range = match self.range { + Some(r) => { + // If this PartitionedFile's range starts at 0, we need to skip the header bytes. + // But then we need to seek back 16 bytes to include the sync marker for the first block, + // as the logic in this reader searches the data for the first sync marker(after which a block starts), + // then reads blocks from the count, size etc. + let start = r.start.max(header_len.checked_sub(16).ok_or(AvroError::ParseError("Avro header length overflow, header was not long enough to contain avro bytes".to_string()))?); + let end = r.end.max(start).min(self.file_size); // Ensure end is not less than start, worst case range is empty + start..end + } + None => 0..self.file_size, + }; + + // 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.compression()?; + let sync_marker = header.sync(); + + Ok(AsyncAvroFileReader::new( + range, + self.file_size, + decoder, + codec, + sync_marker, + reader_state, + )) + } +} diff --git a/arrow-avro/src/reader/async_reader/mod.rs b/arrow-avro/src/reader/async_reader/mod.rs new file mode 100644 index 000000000000..53229f8576eb --- /dev/null +++ b/arrow-avro/src/reader/async_reader/mod.rs @@ -0,0 +1,1764 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::compression::CompressionCodec; +use crate::reader::Decoder; +use crate::reader::block::{BlockDecoder, BlockDecoderState}; +use arrow_array::RecordBatch; +use arrow_schema::ArrowError; +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::{FutureExt, Stream}; +use std::mem; +use std::ops::Range; +use std::pin::Pin; +use std::task::{Context, Poll}; + +mod async_file_reader; +mod builder; + +pub use async_file_reader::AsyncFileReader; +pub use builder::ReaderBuilder; + +#[cfg(feature = "object_store")] +mod store; + +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, + }, + /// 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, + }, + /// 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, +} + +/// An asynchronous Avro file reader that implements `Stream>`. +/// This uses an [`AsyncFileReader`] to fetch data ranges as needed, starting with fetching the header, +/// then reading all the blocks in the provided range where: +/// 1. Reads and decodes data until the header is fully decoded. +/// 2. Searching from `range.start` for the first sync marker, and starting with the following block. +/// (If `range.start` is less than the header length, we start at the header length minus the sync marker bytes) +/// 3. Reading blocks sequentially, decoding them into RecordBatches. +/// 4. If a block is incomplete (due to range ending mid-block), fetching the remaining bytes from the [`AsyncFileReader`]. +/// 5. If no range was originally provided, reads the full file. +/// 6. If the range is 0, file_size is 0, or `range.end` is less than the header length, finish immediately. +/// +/// # Example +/// +/// ``` +/// #[tokio::main(flavor = "current_thread")] +/// async fn main() -> Result<(), Box> { +/// use std::io::Cursor; +/// use std::sync::Arc; +/// use arrow_array::{ArrayRef, Int32Array, RecordBatch}; +/// use arrow_schema::{DataType, Field, Schema}; +/// use arrow_avro::reader::AsyncAvroFileReader; +/// use arrow_avro::writer::AvroWriter; +/// use futures::TryStreamExt; +/// +/// // Build a minimal Arrow schema and batch +/// let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]); +/// let batch = RecordBatch::try_new( +/// Arc::new(schema.clone()), +/// vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], +/// )?; +/// +/// // Write an Avro OCF to memory +/// let buffer: Vec = Vec::new(); +/// let mut writer = AvroWriter::new(buffer, schema)?; +/// writer.write(&batch)?; +/// writer.finish()?; +/// let bytes = writer.into_inner(); +/// +/// // Create an async reader from the in-memory bytes +/// // `tokio::fs::File` also implements `AsyncFileReader` for reading from disk +/// let file_size = bytes.len(); +/// let cursor = Cursor::new(bytes); +/// let reader = AsyncAvroFileReader::builder(cursor, file_size as u64, 1024) +/// .try_build() +/// .await?; +/// +/// // Consume the stream of RecordBatches +/// let batches: Vec = reader.try_collect().await?; +/// assert_eq!(batches.len(), 1); +/// assert_eq!(batches[0].num_rows(), 3); +/// Ok(()) +/// } +/// ``` +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, +} + +impl AsyncAvroFileReader { + /// Returns a builder for a new [`Self`], allowing some optional parameters. + pub fn builder(reader: R, file_size: u64, batch_size: usize) -> ReaderBuilder { + ReaderBuilder::new(reader, file_size, batch_size) + } + + fn new( + range: Range, + file_size: u64, + decoder: Decoder, + codec: Option, + sync_marker: [u8; 16], + reader_state: ReaderState, + ) -> Self { + Self { + range, + file_size, + + decoder, + block_decoder: Default::default(), + codec, + sync_marker, + + reader_state, + finishing_partial_block: false, + } + } + + /// 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 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); + } + } + } + } +} + +// 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) + } +} + +#[cfg(all(test, feature = "object_store"))] +mod tests { + use super::*; + use crate::schema::{AvroSchema, SCHEMA_METADATA_KEY}; + use arrow_array::cast::AsArray; + use arrow_array::types::{Int32Type, Int64Type}; + use arrow_array::*; + use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; + use futures::{StreamExt, TryStreamExt}; + use object_store::local::LocalFileSystem; + use object_store::path::Path; + use object_store::{ObjectStore, ObjectStoreExt}; + use std::collections::HashMap; + use std::sync::Arc; + + fn arrow_test_data(file: &str) -> String { + let base = + std::env::var("ARROW_TEST_DATA").unwrap_or_else(|_| "../testing/data".to_string()); + format!("{}/{}", base, file) + } + + fn get_alltypes_schema() -> SchemaRef { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("bool_col", DataType::Boolean, true), + Field::new("tinyint_col", DataType::Int32, true), + Field::new("smallint_col", DataType::Int32, true), + Field::new("int_col", DataType::Int32, true), + Field::new("bigint_col", DataType::Int64, true), + Field::new("float_col", DataType::Float32, true), + Field::new("double_col", DataType::Float64, true), + Field::new("date_string_col", DataType::Binary, true), + Field::new("string_col", DataType::Binary, true), + Field::new( + "timestamp_col", + DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())), + true, + ), + ]) + .with_metadata(HashMap::from([( + SCHEMA_METADATA_KEY.into(), + r#"{ + "type": "record", + "name": "topLevelRecord", + "fields": [ + { + "name": "id", + "type": [ + "int", + "null" + ] + }, + { + "name": "bool_col", + "type": [ + "boolean", + "null" + ] + }, + { + "name": "tinyint_col", + "type": [ + "int", + "null" + ] + }, + { + "name": "smallint_col", + "type": [ + "int", + "null" + ] + }, + { + "name": "int_col", + "type": [ + "int", + "null" + ] + }, + { + "name": "bigint_col", + "type": [ + "long", + "null" + ] + }, + { + "name": "float_col", + "type": [ + "float", + "null" + ] + }, + { + "name": "double_col", + "type": [ + "double", + "null" + ] + }, + { + "name": "date_string_col", + "type": [ + "bytes", + "null" + ] + }, + { + "name": "string_col", + "type": [ + "bytes", + "null" + ] + }, + { + "name": "timestamp_col", + "type": [ + { + "type": "long", + "logicalType": "timestamp-micros" + }, + "null" + ] + } + ] +} +"# + .into(), + )])); + Arc::new(schema) + } + + fn get_alltypes_with_nulls_schema() -> SchemaRef { + let schema = Schema::new(vec![ + Field::new("string_col", DataType::Binary, true), + Field::new("int_col", DataType::Int32, true), + Field::new("bool_col", DataType::Boolean, true), + Field::new("bigint_col", DataType::Int64, true), + Field::new("float_col", DataType::Float32, true), + Field::new("double_col", DataType::Float64, true), + Field::new("bytes_col", DataType::Binary, true), + ]) + .with_metadata(HashMap::from([( + SCHEMA_METADATA_KEY.into(), + r#"{ + "type": "record", + "name": "topLevelRecord", + "fields": [ + { + "name": "string_col", + "type": [ + "null", + "string" + ], + "default": null + }, + { + "name": "int_col", + "type": [ + "null", + "int" + ], + "default": null + }, + { + "name": "bool_col", + "type": [ + "null", + "boolean" + ], + "default": null + }, + { + "name": "bigint_col", + "type": [ + "null", + "long" + ], + "default": null + }, + { + "name": "float_col", + "type": [ + "null", + "float" + ], + "default": null + }, + { + "name": "double_col", + "type": [ + "null", + "double" + ], + "default": null + }, + { + "name": "bytes_col", + "type": [ + "null", + "bytes" + ], + "default": null + } + ] +}"# + .into(), + )])); + + Arc::new(schema) + } + + fn get_nested_records_schema() -> SchemaRef { + let schema = Schema::new(vec![ + Field::new( + "f1", + DataType::Struct( + vec![ + Field::new("f1_1", DataType::Utf8, false), + Field::new("f1_2", DataType::Int32, false), + Field::new( + "f1_3", + DataType::Struct( + vec![Field::new("f1_3_1", DataType::Float64, false)].into(), + ), + false, + ), + ] + .into(), + ), + false, + ), + Field::new( + "f2", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct( + vec![ + Field::new("f2_1", DataType::Boolean, false), + Field::new("f2_2", DataType::Float32, false), + ] + .into(), + ), + false, + ))), + false, + ), + Field::new( + "f3", + DataType::Struct(vec![Field::new("f3_1", DataType::Utf8, false)].into()), + true, + ), + Field::new( + "f4", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(vec![Field::new("f4_1", DataType::Int64, false)].into()), + true, + ))), + false, + ), + ]) + .with_metadata(HashMap::from([( + SCHEMA_METADATA_KEY.into(), + r#"{ + "type": "record", + "namespace": "ns1", + "name": "record1", + "fields": [ + { + "name": "f1", + "type": { + "type": "record", + "namespace": "ns2", + "name": "record2", + "fields": [ + { + "name": "f1_1", + "type": "string" + }, + { + "name": "f1_2", + "type": "int" + }, + { + "name": "f1_3", + "type": { + "type": "record", + "namespace": "ns3", + "name": "record3", + "fields": [ + { + "name": "f1_3_1", + "type": "double" + } + ] + } + } + ] + } + }, + { + "name": "f2", + "type": { + "type": "array", + "items": { + "type": "record", + "namespace": "ns4", + "name": "record4", + "fields": [ + { + "name": "f2_1", + "type": "boolean" + }, + { + "name": "f2_2", + "type": "float" + } + ] + } + } + }, + { + "name": "f3", + "type": [ + "null", + { + "type": "record", + "namespace": "ns5", + "name": "record5", + "fields": [ + { + "name": "f3_1", + "type": "string" + } + ] + } + ], + "default": null + }, + { + "name": "f4", + "type": { + "type": "array", + "items": [ + "null", + { + "type": "record", + "namespace": "ns6", + "name": "record6", + "fields": [ + { + "name": "f4_1", + "type": "long" + } + ] + } + ] + } + } + ] +} +"# + .into(), + )])); + + Arc::new(schema) + } + + async fn read_async_file( + path: &str, + batch_size: usize, + range: Option>, + schema: Option, + projection: Option>, + ) -> Result, ArrowError> { + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(path).unwrap(); + + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let mut builder = AsyncAvroFileReader::builder(file_reader, file_size, batch_size); + + if let Some(s) = schema { + let reader_schema = AvroSchema::try_from(s.as_ref())?; + builder = builder.with_reader_schema(reader_schema); + } + + if let Some(proj) = projection { + builder = builder.with_projection(proj); + } + + if let Some(range) = range { + builder = builder.with_range(range); + } + + let reader = builder.try_build().await?; + reader.try_collect().await + } + + #[tokio::test] + async fn test_full_file_read() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 11); + + let id_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_array.value(0), 4); + assert_eq!(id_array.value(7), 1); + } + + #[tokio::test] + async fn test_small_batch_size() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 2, None, Some(schema), None) + .await + .unwrap(); + assert_eq!(batches.len(), 4); + + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 11); + } + + #[tokio::test] + async fn test_batch_size_one() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batches.len(), 8); + assert_eq!(batch.num_rows(), 1); + } + + #[tokio::test] + async fn test_batch_size_larger_than_file() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 10000, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + } + + #[tokio::test] + async fn test_empty_range() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let range = 100..100; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + assert_eq!(batches.len(), 0); + } + + #[tokio::test] + async fn test_range_starting_at_zero() { + // Tests that range starting at 0 correctly skips header + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let meta = store.head(&location).await.unwrap(); + + let range = 0..meta.size; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + } + + #[tokio::test] + async fn test_range_after_header() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let meta = store.head(&location).await.unwrap(); + + let range = 100..meta.size; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert!(batch.num_rows() > 0); + } + + #[tokio::test] + async fn test_range_no_sync_marker() { + // Small range unlikely to contain sync marker + let file = arrow_test_data("avro/alltypes_plain.avro"); + let range = 50..150; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + assert_eq!(batches.len(), 0); + } + + #[tokio::test] + async fn test_range_starting_mid_file() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + + let range = 700..768; // Header ends at 675, so this should be mid-block + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + assert_eq!(batches.len(), 0); + } + + #[tokio::test] + async fn test_range_ending_at_file_size() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let meta = store.head(&location).await.unwrap(); + + let range = 200..meta.size; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + } + + #[tokio::test] + async fn test_incomplete_block_requires_fetch() { + // Range ends mid-block, should trigger fetching_rem_block logic + let file = arrow_test_data("avro/alltypes_plain.avro"); + let range = 0..1200; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8) + } + + #[tokio::test] + async fn test_partial_vlq_header_requires_fetch() { + // Range ends mid-VLQ header, triggering the Count|Size partial fetch logic. + let file = arrow_test_data("avro/alltypes_plain.avro"); + let range = 16..676; // Header should end at 675 + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8) + } + + #[cfg(feature = "snappy")] + #[tokio::test] + async fn test_snappy_compressed_with_range() { + { + let file = arrow_test_data("avro/alltypes_plain.snappy.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let meta = store.head(&location).await.unwrap(); + + let range = 200..meta.size; + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert!(batch.num_rows() > 0); + } + } + + #[tokio::test] + async fn test_nulls() { + let file = arrow_test_data("avro/alltypes_nulls_plain.avro"); + let schema = get_alltypes_with_nulls_schema(); + let batches = read_async_file(&file, 1024, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 1); + for col in batch.columns() { + assert!(col.is_null(0)); + } + } + + #[tokio::test] + async fn test_nested_records() { + let file = arrow_test_data("avro/nested_records.avro"); + let schema = get_nested_records_schema(); + let batches = read_async_file(&file, 1024, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert!(batch.num_columns() > 0); + } + + #[tokio::test] + async fn test_stream_produces_multiple_batches() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let schema = get_alltypes_schema(); + let reader_schema = AvroSchema::try_from(schema.as_ref()).unwrap(); + let reader = AsyncAvroFileReader::builder( + file_reader, + file_size, + 2, // Small batch size to force multiple batches + ) + .with_reader_schema(reader_schema) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + + assert!(batches.len() > 1); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 8); + } + + #[tokio::test] + async fn test_stream_early_termination() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let schema = get_alltypes_schema(); + let reader_schema = AvroSchema::try_from(schema.as_ref()).unwrap(); + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1) + .with_reader_schema(reader_schema) + .try_build() + .await + .unwrap(); + + let first_batch = reader.take(1).try_collect::>().await.unwrap(); + + assert_eq!(first_batch.len(), 1); + assert!(first_batch[0].num_rows() > 0); + } + + #[tokio::test] + async fn test_various_batch_sizes() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + + for batch_size in [1, 2, 3, 5, 7, 11, 100] { + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, batch_size, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + // Size should be what was provided, to the limit of the batch in the file + assert_eq!( + batch.num_rows(), + batch_size.min(8), + "Failed with batch_size={}", + batch_size + ); + } + } + + #[tokio::test] + async fn test_range_larger_than_file() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let meta = store.head(&location).await.unwrap(); + + // Range extends beyond file size + let range = 100..(meta.size + 1000); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, Some(range), Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + // Should clamp to file size + assert_eq!(batch.num_rows(), 8); + } + + #[tokio::test] + async fn test_roundtrip_write_then_async_read() { + use crate::writer::AvroWriter; + use arrow_array::{Float64Array, StringArray}; + use std::fs::File; + use std::io::BufWriter; + use tempfile::tempdir; + + // Schema with nullable and non-nullable fields of various types + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("score", DataType::Float64, true), + Field::new("count", DataType::Int64, false), + ])); + + let dir = tempdir().unwrap(); + let file_path = dir.path().join("roundtrip_test.avro"); + + // Write multiple batches with nulls + { + let file = File::create(&file_path).unwrap(); + let writer = BufWriter::new(file); + let mut avro_writer = AvroWriter::new(writer, schema.as_ref().clone()).unwrap(); + + // First batch: 3 rows with some nulls + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![ + Some("alice"), + None, + Some("charlie"), + ])), + Arc::new(Float64Array::from(vec![Some(95.5), Some(87.3), None])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + avro_writer.write(&batch1).unwrap(); + + // Second batch: 2 rows + let batch2 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5])), + Arc::new(StringArray::from(vec![Some("diana"), Some("eve")])), + Arc::new(Float64Array::from(vec![None, Some(88.0)])), + Arc::new(Int64Array::from(vec![40, 50])), + ], + ) + .unwrap(); + avro_writer.write(&batch2).unwrap(); + + avro_writer.finish().unwrap(); + } + + // Read back with small batch size to produce multiple output batches + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file_path).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 2) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + + // Verify we got multiple output batches due to small batch_size + assert!( + batches.len() > 1, + "Expected multiple batches with batch_size=2" + ); + + // Verify total row count + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 5); + + // Concatenate all batches to verify data + let combined = arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(combined.num_rows(), 5); + assert_eq!(combined.num_columns(), 4); + + // Check id column (non-nullable) + let id_array = combined + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_array.values(), &[1, 2, 3, 4, 5]); + + // Check name column (nullable) - verify nulls are preserved + // Avro strings are read as Binary by default + let name_col = combined.column(1); + let name_array = name_col.as_string::(); + assert_eq!(name_array.value(0), "alice"); + assert!(name_col.is_null(1)); // second row has null name + assert_eq!(name_array.value(2), "charlie"); + + // Check score column (nullable) - verify nulls are preserved + let score_array = combined + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!score_array.is_null(0)); + assert!((score_array.value(0) - 95.5).abs() < f64::EPSILON); + assert!(score_array.is_null(2)); // third row has null score + assert!(score_array.is_null(3)); // fourth row has null score + assert!(!score_array.is_null(4)); + assert!((score_array.value(4) - 88.0).abs() < f64::EPSILON); + + // Check count column (non-nullable) + let count_array = combined + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(count_array.values(), &[10, 20, 30, 40, 50]); + } + + #[tokio::test] + async fn test_alltypes_no_schema_no_projection() { + // No reader schema, no projection - uses writer schema from file + let file = arrow_test_data("avro/alltypes_plain.avro"); + let batches = read_async_file(&file, 1024, None, None, None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 11); + assert_eq!(batch.schema().field(0).name(), "id"); + } + + #[tokio::test] + async fn test_alltypes_no_schema_with_projection() { + // No reader schema, with projection - project writer schema + let file = arrow_test_data("avro/alltypes_plain.avro"); + // Project [tinyint_col, id, bigint_col] = indices [2, 0, 5] + let batches = read_async_file(&file, 1024, None, None, Some(vec![2, 0, 5])) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.schema().field(0).name(), "tinyint_col"); + assert_eq!(batch.schema().field(1).name(), "id"); + assert_eq!(batch.schema().field(2).name(), "bigint_col"); + + // Verify data values + let tinyint_col = batch.column(0).as_primitive::(); + assert_eq!(tinyint_col.values(), &[0, 1, 0, 1, 0, 1, 0, 1]); + + let id = batch.column(1).as_primitive::(); + assert_eq!(id.values(), &[4, 5, 6, 7, 2, 3, 0, 1]); + + let bigint_col = batch.column(2).as_primitive::(); + assert_eq!(bigint_col.values(), &[0, 10, 0, 10, 0, 10, 0, 10]); + } + + #[tokio::test] + async fn test_alltypes_with_schema_no_projection() { + // With reader schema, no projection + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + let batches = read_async_file(&file, 1024, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 11); + } + + #[tokio::test] + async fn test_alltypes_with_schema_with_projection() { + // With reader schema, with projection + let file = arrow_test_data("avro/alltypes_plain.avro"); + let schema = get_alltypes_schema(); + // Project [bool_col, id] = indices [1, 0] + let batches = read_async_file(&file, 1024, None, Some(schema), Some(vec![1, 0])) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.schema().field(0).name(), "bool_col"); + assert_eq!(batch.schema().field(1).name(), "id"); + + let bool_col = batch.column(0).as_boolean(); + assert!(bool_col.value(0)); + assert!(!bool_col.value(1)); + + let id = batch.column(1).as_primitive::(); + assert_eq!(id.values(), &[4, 5, 6, 7, 2, 3, 0, 1]); + } + + #[tokio::test] + async fn test_nested_no_schema_no_projection() { + // No reader schema, no projection + let file = arrow_test_data("avro/nested_records.avro"); + let batches = read_async_file(&file, 1024, None, None, None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 4); + assert_eq!(batch.schema().field(0).name(), "f1"); + assert_eq!(batch.schema().field(1).name(), "f2"); + assert_eq!(batch.schema().field(2).name(), "f3"); + assert_eq!(batch.schema().field(3).name(), "f4"); + } + + #[tokio::test] + async fn test_nested_no_schema_with_projection() { + // No reader schema, with projection - reorder nested fields + let file = arrow_test_data("avro/nested_records.avro"); + // Project [f3, f1] = indices [2, 0] + let batches = read_async_file(&file, 1024, None, None, Some(vec![2, 0])) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.schema().field(0).name(), "f3"); + assert_eq!(batch.schema().field(1).name(), "f1"); + } + + #[tokio::test] + async fn test_nested_with_schema_no_projection() { + // With reader schema, no projection + let file = arrow_test_data("avro/nested_records.avro"); + let schema = get_nested_records_schema(); + let batches = read_async_file(&file, 1024, None, Some(schema), None) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 4); + } + + #[tokio::test] + async fn test_nested_with_schema_with_projection() { + // With reader schema, with projection + let file = arrow_test_data("avro/nested_records.avro"); + let schema = get_nested_records_schema(); + // Project [f4, f2, f1] = indices [3, 1, 0] + let batches = read_async_file(&file, 1024, None, Some(schema), Some(vec![3, 1, 0])) + .await + .unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.schema().field(0).name(), "f4"); + assert_eq!(batch.schema().field(1).name(), "f2"); + assert_eq!(batch.schema().field(2).name(), "f1"); + } + + #[tokio::test] + async fn test_projection_error_out_of_bounds() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + // Index 100 is out of bounds for the 11-field schema + let err = read_async_file(&file, 1024, None, None, Some(vec![100])) + .await + .unwrap_err(); + assert!(matches!(err, ArrowError::AvroError(_))); + assert!(err.to_string().contains("out of bounds")); + } + + #[tokio::test] + async fn test_projection_error_duplicate_index() { + let file = arrow_test_data("avro/alltypes_plain.avro"); + // Duplicate index 0 + let err = read_async_file(&file, 1024, None, None, Some(vec![0, 0])) + .await + .unwrap_err(); + assert!(matches!(err, ArrowError::AvroError(_))); + assert!(err.to_string().contains("Duplicate projection index")); + } + + #[tokio::test] + async fn test_with_header_size_hint_small() { + // Use a very small header size hint to force multiple fetches + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let schema = get_alltypes_schema(); + let reader_schema = AvroSchema::try_from(schema.as_ref()).unwrap(); + + // Use a tiny header hint (64 bytes) - header is much larger + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_reader_schema(reader_schema) + .with_header_size_hint(64) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 11); + } + + #[tokio::test] + async fn test_with_header_size_hint_large() { + // Use a larger header size hint than needed + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + let schema = get_alltypes_schema(); + let reader_schema = AvroSchema::try_from(schema.as_ref()).unwrap(); + + // Use a large header hint (64KB) + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_reader_schema(reader_schema) + .with_header_size_hint(64 * 1024) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 8); + assert_eq!(batch.num_columns(), 11); + } + + #[tokio::test] + async fn test_with_utf8_view_enabled() { + // Test that utf8_view produces StringViewArray instead of StringArray + let file = arrow_test_data("avro/nested_records.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_utf8_view(true) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + + // The f1 struct contains f1_1 which is a string field + // With utf8_view enabled, it should be Utf8View type + let f1_col = batch.column(0); + let f1_struct = f1_col.as_struct(); + let f1_1_field = f1_struct.column_by_name("f1_1").unwrap(); + + // Check that the data type is Utf8View + assert_eq!(f1_1_field.data_type(), &DataType::Utf8View); + } + + #[tokio::test] + async fn test_with_utf8_view_disabled() { + // Test that without utf8_view, we get regular Utf8 + let file = arrow_test_data("avro/nested_records.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_utf8_view(false) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + let batch = &batches[0]; + + assert_eq!(batch.num_rows(), 2); + + // The f1 struct contains f1_1 which is a string field + // Without utf8_view, it should be regular Utf8 + let f1_col = batch.column(0); + let f1_struct = f1_col.as_struct(); + let f1_1_field = f1_struct.column_by_name("f1_1").unwrap(); + + assert_eq!(f1_1_field.data_type(), &DataType::Utf8); + } + + #[tokio::test] + async fn test_with_strict_mode_disabled_allows_null_second() { + // Test that with strict_mode disabled, unions of ['T', 'null'] are allowed + // The alltypes_nulls_plain.avro file has unions with null second + let file = arrow_test_data("avro/alltypes_nulls_plain.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + // Without strict mode, this should succeed + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_strict_mode(false) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + } + + #[tokio::test] + async fn test_with_strict_mode_enabled_rejects_null_second() { + // Test that with strict_mode enabled, unions of ['T', 'null'] are rejected + // The alltypes_plain.avro file has unions like ["int", "null"] (null second) + let file = arrow_test_data("avro/alltypes_plain.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + // With strict mode, this should fail because of ['T', 'null'] unions + let result = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_strict_mode(true) + .try_build() + .await; + + match result { + Ok(_) => panic!("Expected error for strict_mode with ['T', 'null'] union"), + Err(err) => { + assert!( + err.to_string().contains("disallowed in strict_mode"), + "Expected strict_mode error, got: {}", + err + ); + } + } + } + + #[tokio::test] + async fn test_with_strict_mode_enabled_valid_schema() { + // Test that strict_mode works with schemas that have proper ['null', 'T'] unions + // The nested_records.avro file has properly ordered unions + let file = arrow_test_data("avro/nested_records.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + // With strict mode, properly ordered unions should still work + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024) + .with_strict_mode(true) + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + } + + #[tokio::test] + async fn test_builder_options_combined() { + // Test combining multiple builder options + let file = arrow_test_data("avro/nested_records.avro"); + let store: Arc = Arc::new(LocalFileSystem::new()); + let location = Path::from_filesystem_path(&file).unwrap(); + let file_size = store.head(&location).await.unwrap().size; + + let file_reader = AvroObjectReader::new(store, location); + + let reader = AsyncAvroFileReader::builder(file_reader, file_size, 2) + .with_header_size_hint(128) + .with_utf8_view(true) + .with_strict_mode(true) + .with_projection(vec![0, 2]) // f1 and f3 + .try_build() + .await + .unwrap(); + + let batches: Vec = reader.try_collect().await.unwrap(); + let batch = &batches[0]; + + // Should have 2 columns (f1 and f3) due to projection + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.schema().field(0).name(), "f1"); + assert_eq!(batch.schema().field(1).name(), "f3"); + + // Verify utf8_view is applied + let f1_col = batch.column(0); + let f1_struct = f1_col.as_struct(); + let f1_1_field = f1_struct.column_by_name("f1_1").unwrap(); + assert_eq!(f1_1_field.data_type(), &DataType::Utf8View); + } +} diff --git a/arrow-avro/src/reader/async_reader/store.rs b/arrow-avro/src/reader/async_reader/store.rs new file mode 100644 index 000000000000..44a4abf1a282 --- /dev/null +++ b/arrow-avro/src/reader/async_reader/store.rs @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::AvroError; +use crate::reader::async_reader::AsyncFileReader; +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::{FutureExt, TryFutureExt}; +use object_store::ObjectStore; +use object_store::ObjectStoreExt; +use object_store::path::Path; +use std::error::Error; +use std::ops::Range; +use std::sync::Arc; +use tokio::runtime::Handle; + +/// An implementation of an AsyncFileReader using the [`ObjectStore`] API. +pub struct AvroObjectReader { + store: Arc, + path: Path, + runtime: Option, +} + +impl AvroObjectReader { + /// Creates a new [`Self`] from a store implementation and file location. + pub fn new(store: Arc, path: Path) -> Self { + Self { + store, + path, + runtime: None, + } + } + + /// Perform IO on the provided tokio runtime + /// + /// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely manner + /// to service IO. Therefore, running IO and CPU-bound tasks, such as avro decoding, + /// on the same tokio runtime can lead to degraded throughput, dropped connections and + /// other issues. For more information see [here]. + /// + /// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/ + pub fn with_runtime(self, handle: Handle) -> Self { + Self { + runtime: Some(handle), + ..self + } + } + + fn spawn(&self, f: F) -> BoxFuture<'_, Result> + where + F: for<'a> FnOnce(&'a Arc, &'a Path) -> BoxFuture<'a, Result> + + Send + + 'static, + O: Send + 'static, + E: Error + Send + 'static, + { + match &self.runtime { + Some(handle) => { + let path = self.path.clone(); + let store = Arc::clone(&self.store); + 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())), + ) + .boxed() + } + None => f(&self.store, &self.path) + .map_err(|e| AvroError::General(e.to_string())) + .boxed(), + } + } +} + +impl AsyncFileReader for AvroObjectReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.spawn(|store, path| async move { store.get_range(path, range).await }.boxed()) + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result, AvroError>> + where + Self: Send, + { + self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed()) + } +} diff --git a/arrow-avro/src/reader/block.rs b/arrow-avro/src/reader/block.rs index 540aedee5227..e20818bf968a 100644 --- a/arrow-avro/src/reader/block.rs +++ b/arrow-avro/src/reader/block.rs @@ -43,7 +43,7 @@ pub struct BlockDecoder { } #[derive(Debug)] -enum BlockDecoderState { +pub(crate) enum BlockDecoderState { Count, Size, Data, @@ -137,3 +137,14 @@ impl BlockDecoder { } } } + +#[cfg(feature = "async")] +impl BlockDecoder { + pub(crate) fn state(&self) -> &BlockDecoderState { + &self.state + } + + pub(crate) fn bytes_remaining(&self) -> usize { + self.bytes_remaining + } +} diff --git a/arrow-avro/src/reader/mod.rs b/arrow-avro/src/reader/mod.rs index 55763aeb6230..a60bc26b49b3 100644 --- a/arrow-avro/src/reader/mod.rs +++ b/arrow-avro/src/reader/mod.rs @@ -499,6 +499,14 @@ mod header; mod record; mod vlq; +#[cfg(feature = "async")] +mod async_reader; + +#[cfg(feature = "object_store")] +pub use async_reader::AvroObjectReader; +#[cfg(feature = "async")] +pub use async_reader::{AsyncAvroFileReader, AsyncFileReader}; + fn is_incomplete_data(err: &AvroError) -> bool { matches!( err, @@ -641,6 +649,25 @@ pub struct Decoder { } impl Decoder { + pub(crate) fn from_parts( + batch_size: usize, + active_decoder: RecordDecoder, + active_fingerprint: Option, + cache: IndexMap, + fingerprint_algorithm: FingerprintAlgorithm, + ) -> Self { + Self { + batch_size, + remaining_capacity: batch_size, + active_fingerprint, + active_decoder, + cache, + fingerprint_algorithm, + pending_schema: None, + awaiting_body: false, + } + } + /// Returns the Arrow schema for the rows decoded by this decoder. /// /// **Note:** With single‑object or Confluent framing, the schema may change @@ -999,25 +1026,6 @@ impl ReaderBuilder { self.make_record_decoder(writer_schema, reader_schema_raw.as_ref()) } - fn make_decoder_with_parts( - &self, - active_decoder: RecordDecoder, - active_fingerprint: Option, - cache: IndexMap, - fingerprint_algorithm: FingerprintAlgorithm, - ) -> Decoder { - Decoder { - batch_size: self.batch_size, - remaining_capacity: self.batch_size, - active_fingerprint, - active_decoder, - cache, - fingerprint_algorithm, - pending_schema: None, - awaiting_body: false, - } - } - fn make_decoder( &self, header: Option<&Header>, @@ -1054,7 +1062,8 @@ impl ReaderBuilder { let effective_reader_schema = projected_reader_schema.as_ref().or(reader_schema); let record_decoder = self.make_record_decoder_from_schemas(&writer_schema, effective_reader_schema)?; - return Ok(self.make_decoder_with_parts( + return Ok(Decoder::from_parts( + self.batch_size, record_decoder, None, IndexMap::new(), @@ -1121,7 +1130,8 @@ impl ReaderBuilder { "Initial fingerprint {start_fingerprint:?} not found in schema store" )) })?; - Ok(self.make_decoder_with_parts( + Ok(Decoder::from_parts( + self.batch_size, active_decoder, Some(start_fingerprint), cache,