Skip to content

Commit f03e1bc

Browse files
authored
fix(ipc): handle duplicate projection indices in IPC reader (#9952)
# Which issue does this PR close? - Closes #9950 . # Rationale for this change The current IPC reader does not correctly handle duplicate projection indices. `Schema::project`(in `arrow-schema/src/schema.rs`) and `RecordBatch::project`(in `arrow-array/src/record_batch.rs`) both map each requested index directly, preserve the projection order and allow duplicate indices such as: ```rust id="n4pq0f" vec![1, 1] ``` However, the IPC reader currently uses: ```rust id="gjklyo" projection.iter().position(|p| p == &idx) ``` which only returns the first matching entry. As a result, only one column is decoded even though the projected schema contains multiple fields, leading to schema/column count mismatches when constructing the `RecordBatch`. This also affects reordered duplicate projections such as: ```rust id="jlwmku" vec![2, 0, 2] ``` # What changes are included in this PR? * Updated IPC projection handling in `arrow-ipc/src/reader.rs` to preserve all matching projection entries * Reused the decoded array for duplicate projection indices instead of decoding the same field multiple times * Preserved projection order for reordered duplicate projections # Are these changes tested? Yes. Added `test_projection_duplicate_indices`, which verifies: * duplicate projections (`vec![1, 1]`) * reordered duplicate projections (`vec![2, 0, 2]`) The test compares IPC projection results against `RecordBatch::project`. The test fails before the fix and passes after it. All existing `arrow-ipc` tests also pass `cargo test -p arrow-ipc --lib` # Are there any user-facing changes? No.
1 parent 259cff2 commit f03e1bc

1 file changed

Lines changed: 38 additions & 5 deletions

File tree

arrow-ipc/src/reader.rs

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -560,11 +560,20 @@ impl<'a> RecordBatchDecoder<'a> {
560560
let mut arrays = vec![];
561561
// project fields
562562
for (idx, field) in schema.fields().iter().enumerate() {
563-
// Create array for projected field
564-
if let Some(proj_idx) = projection.iter().position(|p| p == &idx) {
565-
let child = self.create_array(field, &mut variadic_counts)?;
566-
arrays.push((proj_idx, child));
567-
} else {
563+
// A projected field can appear more than once, so collect all matching positions.
564+
let mut child = None;
565+
for (proj_idx, projected_idx) in projection.iter().enumerate() {
566+
if *projected_idx == idx {
567+
if child.is_none() {
568+
child = Some(self.create_array(field, &mut variadic_counts)?);
569+
}
570+
571+
// Reuse the decoded array for duplicate projection entries.
572+
arrays.push((proj_idx, child.as_ref().unwrap().clone()));
573+
}
574+
}
575+
576+
if child.is_none() {
568577
self.skip_field(field, &mut variadic_counts)?;
569578
}
570579
}
@@ -2297,6 +2306,30 @@ mod tests {
22972306
}
22982307
}
22992308

2309+
#[test]
2310+
fn test_projection_duplicate_indices() {
2311+
let schema = create_test_projection_schema();
2312+
let batch = create_test_projection_batch_data(&schema);
2313+
2314+
// Write the batch to IPC
2315+
let mut buf = Vec::new();
2316+
{
2317+
let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2318+
writer.write(&batch).unwrap();
2319+
writer.finish().unwrap();
2320+
}
2321+
2322+
// Verify duplicate([1, 1]) and reordered([2, 0, 2]) projection indices
2323+
for projection in [vec![1, 1], vec![2, 0, 2]] {
2324+
let reader =
2325+
FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2326+
let read_batch = reader.unwrap().next().unwrap().unwrap();
2327+
2328+
let expected_batch = batch.project(&projection).unwrap();
2329+
assert_eq!(read_batch, expected_batch);
2330+
}
2331+
}
2332+
23002333
#[test]
23012334
fn test_arrow_single_float_row() {
23022335
let schema = Schema::new(vec![

0 commit comments

Comments
 (0)