Jefffrey commented on code in PR #10909:
URL: https://github.com/apache/arrow-rs/pull/10909#discussion_r3888639951
##########
arrow-select/src/take.rs:
##########
@@ -1121,6 +1136,67 @@ 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(
+ prev_idx: Option<usize>,
+ cur_idx: Option<usize>,
+ values_cmp: &arrow_cmp::DynComparator,
+) -> bool {
+ match (prev_idx, cur_idx) {
+ (None, None) => false,
+ (Some(prev), Some(cur)) => 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 Vec<bool>,
+ physical: Option<usize>,
+) {
+ match physical {
+ Some(idx) => {
+ take_value_indices.push(I::Native::from_usize(idx).unwrap());
+ take_value_is_valid.push(true);
+ }
+ None => {
+ take_value_indices.push(I::Native::default());
+ take_value_is_valid.push(false);
Review Comment:
can we add a test for when the run values themselves have nulls too? i
suspect we might hit another edge case here since it seems we determine nulls
based only on the indices nulls, without considering the values nulls
##########
arrow-select/src/take.rs:
##########
@@ -347,7 +347,12 @@ fn take_impl<IndexType: ArrowPrimitiveType, const CHECKED:
bool>(
let values = values.as_any().downcast_ref::<UnionArray>().unwrap();
let type_ids =
<PrimitiveArray<Int8Type>>::try_new(take_native(values.type_ids(), indices),
None)?;
Review Comment:
would this type ids suffer from a similar issue? it seems itll place 0 for
null indices, but theres no guarantee we have a 0 child
##########
arrow-select/src/take.rs:
##########
@@ -1121,6 +1136,67 @@ 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(
+ prev_idx: Option<usize>,
+ cur_idx: Option<usize>,
+ values_cmp: &arrow_cmp::DynComparator,
+) -> bool {
+ match (prev_idx, cur_idx) {
+ (None, None) => false,
+ (Some(prev), Some(cur)) => prev != cur && values_cmp(cur,
prev).is_ne(),
Review Comment:
this might have a similar issue, though not a bug per se; if for example we
have a `None` index from indices, but we have a `Some(_)` index that points to
a null in the values, then its still technically a null run
(at worst this just generates an inefficient run array, so not necessarily a
correctness thing)
##########
arrow-select/src/take.rs:
##########
@@ -2781,6 +2857,38 @@ mod tests {
assert_eq!(take_out_values.values(), &[2, 1]);
}
+ #[test]
+ fn test_take_runs_null_indices() {
+ // [10, 10, 99, 1, 1, 1]; a null index must not become logical index 0.
+ let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+ builder.extend([10, 10, 99, 1, 1, 1].into_iter().map(Some));
+ let run_array = builder.finish();
+
+ let indices = Int32Array::from(vec![Some(0), None, Some(2)]);
+ let taken = take(&run_array, &indices, None).unwrap();
+ let logical: Vec<Option<i32>> = taken
+ .as_run::<Int32Type>()
+ .downcast::<Int32Array>()
+ .unwrap()
+ .into_iter()
+ .collect();
+ assert_eq!(logical, vec![Some(10), None, Some(99)]);
+
+ let primitive = Int32Array::from(vec![10, 10, 99, 1, 1, 1]);
+ let primitive_taken = take(&primitive, &indices, None).unwrap();
+ let primitive_logical: Vec<Option<i32>> =
+ primitive_taken.as_primitive::<Int32Type>().iter().collect();
+ assert_eq!(logical, primitive_logical);
Review Comment:
what is this case asserting? it seems to test a take on a primitive array,
not on a run array?
##########
arrow-select/src/take.rs:
##########
@@ -1094,22 +1101,30 @@ fn take_run<T: RunEndIndexType, I: ArrowPrimitiveType>(
for ix in 1..physical_indices.len() {
let prev_idx = physical_indices[ix - 1];
let cur_idx = physical_indices[ix];
- let is_new_run = cur_idx != prev_idx && values_cmp(cur_idx,
prev_idx).is_ne();
- if is_new_run {
- take_value_indices.push(I::Native::from_usize(prev_idx).unwrap());
+ if is_new_run_take(prev_idx, cur_idx, &values_cmp) {
+ push_run_take_index::<I>(&mut take_value_indices, &mut
take_value_is_valid, prev_idx);
new_run_ends.push(T::Native::from_usize(ix).unwrap());
}
}
- take_value_indices
- .push(I::Native::from_usize(physical_indices[physical_indices.len() -
1]).unwrap());
+ push_run_take_index::<I>(
+ &mut take_value_indices,
+ &mut take_value_is_valid,
+ physical_indices[physical_indices.len() - 1],
+ );
new_run_ends.push(T::Native::from_usize(physical_indices.len()).unwrap());
// SAFETY: run-ends are strictly increasing with last value == logical
length.
let run_ends = unsafe {
RunEndBuffer::new_unchecked(ScalarBuffer::from(new_run_ends), 0,
physical_indices.len())
};
- let take_value_indices =
PrimitiveArray::<I>::new(ScalarBuffer::from(take_value_indices), None);
+ let nulls = if take_value_is_valid.iter().all(|&v| v) {
Review Comment:
we can use a
[`NullBufferBuilder`](https://docs.rs/arrow/latest/arrow/array/struct.NullBufferBuilder.html)
for `take_value_is_valid` which has the benefit of doing this materialization
check for us, only more cheaply
--
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]
