This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 2cc31af529 GH-10676: Fix IPC StreamReader projected schema (#10677)
2cc31af529 is described below
commit 2cc31af5297492b41ddf4b8f9f0d33ff15d9a400
Author: anchor <[email protected]>
AuthorDate: Sat Aug 15 22:21:34 2026 +0800
GH-10676: Fix IPC StreamReader projected schema (#10677)
# Which issue does this PR close?
- Closes #10676.
# Rationale for this change
With a projection set, `StreamReader` yields batches with the projected
schema, but `StreamReader::schema()` and its
`RecordBatchReader::schema()` impl return the full stream schema. Same
as #10382 (csv) and #10613 (ipc file reader); #10627 fixed `FileReader`,
`StreamReader` was the remaining case.
# What changes are included in this PR?
- Keep the projected schema `StreamReader::try_new` already computes as
a `SchemaRef` and return it from `StreamReader::schema()`; the
`RecordBatchReader` impl delegates to it, as `FileReader` does.
- Decoding is untouched: batch and dictionary decoding still use the
full stream schema, and with no projection the reported schema is
unchanged.
- Add a regression test mirroring
`test_file_reader_projected_schema_matches_batch_schema`.
# Are these changes tested?
Yes, the new test fails before this change (14 column stream schema vs 3
column batch schema) and passes after. I ran:
- `cargo test -p arrow-ipc` (143 passed) and `--all-features` (148
passed)
- `cargo clippy -p arrow-ipc --all-targets --all-features -- -D
warnings`
- `cargo fmt --all -- --check`
I did not run the full workspace suite; outside `arrow-ipc` every
`StreamReader::try_new` call site passes `None`, which is unaffected.
# Are there any user-facing changes?
Yes. A projected IPC `StreamReader` now reports the same schema as the
batches it produces. No public API signature changes.
AI assistance: an AI agent found this by auditing the remaining
`RecordBatchReader` impls after #10613 and drafted the fix and its test;
I reviewed every line and ran the reproduction and the checks above
locally.
Co-authored-by: codeAnqiang-ma
<[email protected]>
---
arrow-ipc/src/reader.rs | 35 +++++++++++++++++++++++++++++------
1 file changed, 29 insertions(+), 6 deletions(-)
diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs
index 99e7c89c9c..c0d99c370e 100644
--- a/arrow-ipc/src/reader.rs
+++ b/arrow-ipc/src/reader.rs
@@ -1541,8 +1541,8 @@ pub struct StreamReader<R> {
/// This value is set to `true` the first time the reader's `next()`
returns `None`.
finished: bool,
- /// Optional projection
- projection: Option<(Vec<usize>, Schema)>,
+ /// Optional projection: column indices and the resulting projected schema
+ projection: Option<(Vec<usize>, SchemaRef)>,
/// Should validation be skipped when reading data? Defaults to false.
///
@@ -1612,7 +1612,7 @@ impl<R: Read> StreamReader<R> {
let projection = match projection {
Some(projection_indices) => {
- let schema = schema.project(&projection_indices)?;
+ let schema = Arc::new(schema.project(&projection_indices)?);
Some((projection_indices, schema))
}
_ => None,
@@ -1628,9 +1628,12 @@ impl<R: Read> StreamReader<R> {
})
}
- /// Return the schema of the stream
+ /// Return the schema of the record batches produced by this reader
pub fn schema(&self) -> SchemaRef {
- self.schema.clone()
+ match &self.projection {
+ Some((_, projected_schema)) => projected_schema.clone(),
+ None => self.schema.clone(),
+ }
}
/// Check if the stream is finished
@@ -1786,7 +1789,7 @@ impl<R: Read> Iterator for StreamReader<R> {
impl<R: Read> RecordBatchReader for StreamReader<R> {
fn schema(&self) -> SchemaRef {
- self.schema.clone()
+ self.schema()
}
}
@@ -2378,6 +2381,26 @@ mod tests {
assert_eq!(reader_schema, read_batch.schema());
}
+ #[test]
+ fn test_stream_reader_projected_schema_matches_batch_schema() {
+ let schema = create_test_projection_schema();
+ let batch = create_test_projection_batch_data(&schema);
+
+ let mut buf = Vec::new();
+ {
+ let mut writer = crate::writer::StreamWriter::try_new(&mut buf,
&schema).unwrap();
+ writer.write(&batch).unwrap();
+ writer.finish().unwrap();
+ }
+
+ let projection = vec![3, 2, 1];
+ let mut reader = StreamReader::try_new(Cursor::new(buf),
Some(projection)).unwrap();
+ let reader_schema = RecordBatchReader::schema(&reader);
+ let read_batch = reader.next().unwrap().unwrap();
+
+ assert_eq!(reader_schema, read_batch.schema());
+ }
+
#[test]
fn test_file_reader_rejects_invalid_projection() {
let schema = create_test_projection_schema();
