yongster opened a new issue, #10908:
URL: https://github.com/apache/arrow-rs/issues/10908
### Describe the bug
`take` is supposed to leave a null in the output when the index at that slot
is null.
That is what happens for primitive arrays, dictionaries, lists, and
**sparse** unions.
Two arms do not:
1. **Dense `UnionArray`** — a null index becomes the first row of the
selected child (usually the value at offset 0).
2. **`RunEndEncoded`** — a null index becomes the first run's value.
Sparse union on the same data is already correct, so this is not “union take
is broken” in general. The dense path rebuilds child offsets itself and drops
the index nulls. The REE path looks up physical runs from `indices.values()`,
which still holds `0` in a null slot.
Existing tests for both arms only use non-null indices, so they do not catch
this.
### To Reproduce
On current `main`:
```rust
use std::sync::Arc;
use arrow_array::builder::PrimitiveRunBuilder;
use arrow_array::cast::AsArray;
use arrow_array::types::Int32Type;
use arrow_array::{Array, Int32Array, UnionArray, UInt32Array};
use arrow_buffer::ScalarBuffer;
use arrow_schema::{DataType, Field, UnionFields};
use arrow_select::take::take;
fn main() {
let indices = UInt32Array::from(vec![Some(0), None, Some(2)]);
// --- Dense union [1, 2, 3] ---
let fields: UnionFields = [(0, Arc::new(Field::new("i", DataType::Int32,
true)))]
.into_iter()
.collect();
let dense = UnionArray::try_new(
fields.clone(),
ScalarBuffer::from(vec![0_i8, 0, 0]),
Some(ScalarBuffer::from(vec![0_i32, 1, 2])),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let dense_out = take(&dense, &indices, None).unwrap();
println!("dense union take: {dense_out:?}");
// actual: [1, 1, 3] (null slot became the first value)
// expected: [1, null, 3]
// Sparse union, same values: already correct
let sparse = UnionArray::try_new(
fields,
ScalarBuffer::from(vec![0_i8, 0, 0]),
None,
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let sparse_out = take(&sparse, &indices, None).unwrap();
println!("sparse union take: {sparse_out:?}");
// [1, null, 3]
// --- REE [10, 10, 99, 1, 1, 1] ---
let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
for v in [10, 10, 99, 1, 1, 1] {
b.append_value(v);
}
let ree = b.finish();
let ree_out = take(&ree, &indices, None).unwrap();
println!("REE take: {ree_out:?}");
// actual: [10, 10, 99] (null slot became the first run)
// expected: [10, null, 99]
// Primitive, same values: already correct
let prim = Int32Array::from(vec![10, 10, 99, 1, 1, 1]);
let prim_out = take(&prim, &indices, None).unwrap();
println!("primitive take: {prim_out:?}");
// [10, null, 99]
}
```
### Expected behavior
A null in `indices` should produce a null in the result, same as primitive /
dictionary / list / sparse union.
- Dense union `[1, 2, 3]` with indices `[0, null, 2]` → `[1, null, 3]`, not
`[1, 1, 3]`
- REE `[10, 10, 99, …]` with indices `[0, null, 2]` → `[10, null, 99]`, not
`[10, 10, 99]`
### Additional context
Dense union (`arrow-select/src/take.rs`): `take_native` fills a null slot
with `0` (default `type_id` / `offset`). The next `PrimitiveArray::try_new(...,
None)` drops the index null bitmap. The fake offset `0` is then used to `take`
from the child, so the first child value is copied into that slot.
REE (`take_run` in the same file):
```rust
let physical_indices =
run_array.get_physical_indices(logical_indices.values())?;
```
`values()` is the raw integer buffer. A null `UInt32` slot is typically `0`,
so the lookup uses logical index 0 (the first run). `indices.nulls()` is never
read, and the output run array has no hole for that slot.
I reproduced this locally on current `main`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
