Jefffrey commented on code in PR #10909:
URL: https://github.com/apache/arrow-rs/pull/10909#discussion_r3889376536
##########
arrow-select/src/take.rs:
##########
@@ -384,6 +392,40 @@ fn take_impl<IndexType: ArrowPrimitiveType, const CHECKED:
bool>(
}
}
+/// Takes union type ids, substituting a valid type id for null take indices.
+///
+/// Union arrays do not have a top-level null bitmap. A null is represented by
selecting an
+/// arbitrary valid child type id with a null value in that child. In
particular, a null index
+/// cannot fall back to type id `0`, as unions are not required to have such a
child.
+fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
+ fields: &UnionFields,
+ type_ids: &ScalarBuffer<i8>,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ScalarBuffer<i8>, ArrowError> {
+ let taken_type_ids = take_native(type_ids, indices);
+ if indices.null_count() == 0 {
+ return Ok(taken_type_ids);
+ }
+
+ let null_type_id = fields
+ .iter()
+ .next()
+ .map(|(type_id, _)| type_id)
+ .ok_or_else(|| ArrowError::ComputeError("Cannot take from an empty
union".into()))?;
Review Comment:
```suggestion
.ok_or_else(|| ArrowError::ComputeError("Cannot take from a union
with zero fields when indices contains nulls".into()))?;
```
this is an interesting edge case, nice spot; could we add a unit test for
this? for both this case, union with 0 fields and indices has a null, and when
union with 0 fields but indices has no null (should still be valid, albeit
degenerate case)
##########
arrow-select/src/take.rs:
##########
@@ -384,6 +392,40 @@ fn take_impl<IndexType: ArrowPrimitiveType, const CHECKED:
bool>(
}
}
+/// Takes union type ids, substituting a valid type id for null take indices.
+///
+/// Union arrays do not have a top-level null bitmap. A null is represented by
selecting an
+/// arbitrary valid child type id with a null value in that child. In
particular, a null index
+/// cannot fall back to type id `0`, as unions are not required to have such a
child.
+fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
+ fields: &UnionFields,
+ type_ids: &ScalarBuffer<i8>,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ScalarBuffer<i8>, ArrowError> {
+ let taken_type_ids = take_native(type_ids, indices);
+ if indices.null_count() == 0 {
+ return Ok(taken_type_ids);
+ }
+
+ let null_type_id = fields
Review Comment:
actually i thought of another potential edge case: the field must be
nullable as well 🤔
we can probably do this in a followup, as perhaps it might be a wider issue
##########
arrow-select/src/take.rs:
##########
@@ -2787,6 +2910,22 @@ mod tests {
assert_eq!(take_out_values.values(), &[2, 1]);
}
+ #[test]
+ fn test_take_runs_null_indices() {
Review Comment:
oh theres another separate issue i just thought of: if the input run array
has non-null for its values, but we take with indices that has nulls, then
we're creating a run array with non-null values (in its field definition) but
actually has nulls...
i think we can deal with it in a followup, as i suspect we might have
similar issues with other kernels (interleave in particular) and i dont want to
scope creep this PR too much
for reference this is an example unit test
```rust
#[test]
fn test123() {
let array = unsafe {
RunArray::<Int16Type>::new_unchecked(
DataType::RunEndEncoded(
Field::new("run_ends", DataType::Int16, false).into(),
Field::new("values", DataType::Int32, false).into(),
),
RunEndBuffer::new(vec![1, 2].into(), 0, 2),
Arc::new(Int32Array::from(vec![1, 2])),
)
};
dbg!(&array);
let indices = Int32Array::from(vec![Some(0), None, Some(1)]);
let taken = take(&array, &indices, None).unwrap();
dbg!(&taken);
println!("{}", array.data_type());
println!("{}", taken.data_type());
}
```
and the output:
```sh
[arrow-select/src/take.rs:2941:9] &array = RunArray {run_ends: [1, 2],
values: PrimitiveArray<Int32>
[
1,
2,
]}
[arrow-select/src/take.rs:2944:9] &taken = RunArray {run_ends: [1, 2, 3],
values: PrimitiveArray<Int32>
[
1,
null,
2,
]}
RunEndEncoded("run_ends": non-null Int16, "values": non-null Int32)
RunEndEncoded("run_ends": non-null Int16, "values": non-null Int32)
```
- so input & output array has identical datatype, down to the nullability of
its values; however can see the actual values array has nulls
##########
arrow-select/src/take.rs:
##########
@@ -1127,6 +1181,75 @@ fn take_run<T: RunEndIndexType, I: ArrowPrimitiveType>(
)
}
+/// Physical run index for each logical take slot.
+///
+/// `None` means the logical index is null. Only valid indices are passed to
+/// [`RunArray::get_physical_indices`]; a null slot's backing integer is
ignored
+/// and may be out of range.
+fn physical_indices_for_take<T: RunEndIndexType, I: ArrowPrimitiveType>(
+ run_array: &RunArray<T>,
+ logical_indices: &PrimitiveArray<I>,
+) -> Result<Vec<Option<usize>>, ArrowError> {
+ if logical_indices.null_count() == 0 {
+ return Ok(run_array
+ .get_physical_indices(logical_indices.values())?
+ .into_iter()
+ .map(Some)
+ .collect());
+ }
+
+ let valid_logical: Vec<_> = logical_indices.iter().flatten().collect();
+
+ let valid_physical = if valid_logical.is_empty() {
+ Vec::new()
+ } else {
+ run_array.get_physical_indices(&valid_logical)?
+ };
+
+ let mut valid_physical = valid_physical.into_iter();
+ Ok(logical_indices
+ .iter()
+ .map(|index| index.map(|_| valid_physical.next().unwrap()))
+ .collect())
+}
+
+fn is_new_run_take(
+ values: &dyn Array,
+ prev_idx: Option<usize>,
+ cur_idx: Option<usize>,
+ values_cmp: &arrow_cmp::DynComparator,
+) -> bool {
+ let prev_valid = prev_idx.is_some_and(|idx| values.is_valid(idx));
+ let cur_valid = cur_idx.is_some_and(|idx| values.is_valid(idx));
+ match (prev_valid, cur_valid) {
+ (false, false) => false,
+ (true, true) => {
+ let prev = prev_idx.unwrap();
+ let cur = cur_idx.unwrap();
+ prev != cur && values_cmp(cur, prev).is_ne()
+ }
+ _ => true,
+ }
+}
+
+fn push_run_take_index<I: ArrowPrimitiveType>(
+ take_value_indices: &mut Vec<I::Native>,
+ take_value_is_valid: &mut NullBufferBuilder,
+ values: &dyn Array,
+ physical: Option<usize>,
+) {
+ match physical.filter(|&idx| values.is_valid(idx)) {
+ Some(idx) => {
+ take_value_indices.push(I::Native::from_usize(idx).unwrap());
+ take_value_is_valid.append_non_null();
+ }
+ None => {
+ take_value_indices.push(I::Native::default());
+ take_value_is_valid.append_null();
+ }
+ }
Review Comment:
```suggestion
// Safe unwrap since physical indices came from a valid run array
let index = I::Native::from_usize(physical.unwrap_or_default()).unwrap();
take_value_indices.push(index);
let is_valid = physical.is_some_and(|idx| values.is_valid(idx));
take_value_is_valid.append(is_valid);
```
we could shorten like so
and then preferably inline this function, since its only used twice and the
signature is a bit gnarly (4 arguments, 2 of which are mutable references, just
for 5 lines)
--
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]
