Skip to content

feat: expose arrow schema on async avro reader #9534

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 4 commits into from
Mar 11, 2026
Merged

feat: expose arrow schema on async avro reader #9534

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a59b6f2
feat: expose Arrow schema from async avro reader
mzabaluev Mar 10, 2026
10b81dc
test: verify arrow schema on async avro reader
mzabaluev Mar 10, 2026
b11ca64
test: fix up expected nested records schema
mzabaluev Mar 10, 2026
4b39bef
test(arrow-avro): clean up async reader tests
mzabaluev Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 140 additions & 21 deletions arrow-avro/src/reader/async_reader/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use crate::compression::CompressionCodec;
use crate::reader::Decoder;
use crate::reader::block::{BlockDecoder, BlockDecoderState};
use arrow_array::RecordBatch;
use arrow_schema::ArrowError;
use arrow_schema::{ArrowError, SchemaRef};
use bytes::Bytes;
use futures::future::BoxFuture;
use futures::{FutureExt, Stream};
Expand Down Expand Up @@ -173,6 +173,13 @@ impl<R> AsyncAvroFileReader<R> {
}
}

/// 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.
Expand Down Expand Up @@ -534,7 +541,9 @@ impl<R: AsyncFileReader + Unpin + 'static> Stream for AsyncAvroFileReader<R> {
#[cfg(all(test, feature = "object_store"))]
mod tests {
use super::*;
use crate::schema::{AvroSchema, SCHEMA_METADATA_KEY};
use crate::schema::{
AVRO_NAME_METADATA_KEY, AVRO_NAMESPACE_METADATA_KEY, AvroSchema, SCHEMA_METADATA_KEY,
};
use arrow_array::cast::AsArray;
use arrow_array::types::{Int32Type, Int64Type};
use arrow_array::*;
Expand Down Expand Up @@ -758,39 +767,63 @@ mod tests {
vec![Field::new("f1_3_1", DataType::Float64, false)].into(),
),
false,
),
)
.with_metadata(HashMap::from([
(AVRO_NAMESPACE_METADATA_KEY.to_owned(), "ns3".to_owned()),
(AVRO_NAME_METADATA_KEY.to_owned(), "record3".to_owned()),
])),
]
.into(),
),
false,
),
)
.with_metadata(HashMap::from([
(AVRO_NAMESPACE_METADATA_KEY.to_owned(), "ns2".to_owned()),
(AVRO_NAME_METADATA_KEY.to_owned(), "record2".to_owned()),
])),
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,
))),
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,
)
.with_metadata(HashMap::from([
(AVRO_NAMESPACE_METADATA_KEY.to_owned(), "ns4".to_owned()),
(AVRO_NAME_METADATA_KEY.to_owned(), "record4".to_owned()),
])),
)),
false,
),
Field::new(
"f3",
DataType::Struct(vec![Field::new("f3_1", DataType::Utf8, false)].into()),
true,
),
)
.with_metadata(HashMap::from([
(AVRO_NAMESPACE_METADATA_KEY.to_owned(), "ns5".to_owned()),
(AVRO_NAME_METADATA_KEY.to_owned(), "record5".to_owned()),
])),
Field::new(
"f4",
DataType::List(Arc::new(Field::new(
"item",
DataType::Struct(vec![Field::new("f4_1", DataType::Int64, false)].into()),
true,
))),
DataType::List(Arc::new(
Field::new(
"item",
DataType::Struct(vec![Field::new("f4_1", DataType::Int64, false)].into()),
true,
)
.with_metadata(HashMap::from([
(AVRO_NAMESPACE_METADATA_KEY.to_owned(), "ns6".to_owned()),
(AVRO_NAME_METADATA_KEY.to_owned(), "record6".to_owned()),
])),
)),
false,
),
])
Expand Down Expand Up @@ -1538,6 +1571,92 @@ mod tests {
assert!(err.to_string().contains("Duplicate projection index"));
}

#[tokio::test]
async fn test_arrow_schema_from_reader_no_reader_schema() {
let file = arrow_test_data("avro/alltypes_plain.avro");
let store: Arc<dyn ObjectStore> = 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);

@alamb alamb Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also reduce some of the duplication in this test so that it is easier to understand what is actually being tested and what is different between the tests?

@mzabaluev mzabaluev Mar 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have added clarifying comments explaining the purpose and differences of code in each of the added cases. Hope this helps.

let expected_schema = get_alltypes_schema()
.as_ref()
.clone()
.with_metadata(Default::default());

// Build reader without providing reader schema - should use writer schema from file
let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024)
.try_build()
.await
.unwrap();

assert_eq!(reader.schema().as_ref(), &expected_schema);

let batches: Vec<RecordBatch> = reader.try_collect().await.unwrap();
let batch = &batches[0];

assert_eq!(batch.schema().as_ref(), &expected_schema);
}

#[tokio::test]
async fn test_arrow_schema_from_reader_with_reader_schema() {
let file = arrow_test_data("avro/alltypes_plain.avro");
let store: Arc<dyn ObjectStore> = 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()
.project(&[0, 1, 7])
.unwrap()
.with_metadata(Default::default());
let reader_schema = AvroSchema::try_from(&schema).unwrap();
let expected_schema = schema.clone();

// Build reader with provided reader schema - must apply the projection
let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024)
.with_reader_schema(reader_schema)
.try_build()
.await
.unwrap();

assert_eq!(reader.schema().as_ref(), &expected_schema);

let batches: Vec<RecordBatch> = reader.try_collect().await.unwrap();
let batch = &batches[0];

assert_eq!(batch.schema().as_ref(), &expected_schema);
}

#[tokio::test]
async fn test_arrow_schema_from_reader_nested_records() {
let file = arrow_test_data("avro/nested_records.avro");
let store: Arc<dyn ObjectStore> = 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);

// The schema produced by the reader should match the expected schema,
// attaching Avro type name metadata to fields of record and list types.
let expected_schema = get_nested_records_schema()
.as_ref()
.clone()
.with_metadata(Default::default());

let reader = AsyncAvroFileReader::builder(file_reader, file_size, 1024)
.try_build()
.await
.unwrap();

assert_eq!(reader.schema().as_ref(), &expected_schema);

let batches: Vec<RecordBatch> = reader.try_collect().await.unwrap();
let batch = &batches[0];

assert_eq!(batch.schema().as_ref(), &expected_schema);
}

#[tokio::test]
async fn test_with_header_size_hint_small() {
// Use a very small header size hint to force multiple fetches
Expand Down
Loading