jorgecarleitao commented on a change in pull request #236:
URL: https://github.com/apache/arrow-rs/pull/236#discussion_r623568928
##########
File path: arrow/src/compute/kernels/sort.rs
##########
@@ -1556,6 +1559,28 @@ mod tests {
Some(3),
vec![Some(1.0), Some(2.0), Some(3.0)],
);
+
+ // limit that includes some extra nulls
+ test_sort_primitive_arrays::<Float64Type>(
Review comment:
I would try to test the 4 branches:
* `(len < null.count(), null_first)`
* `(len < null.count(), null_last)`
* `(len > null.count(), null_first)`
* `(len > null.count(), null_last)`
just to make sure we cover all of them correctly.
##########
File path: arrow/src/compute/kernels/sort.rs
##########
@@ -487,24 +487,27 @@ where
len = limit.min(len);
}
if !descending {
- sort_by(&mut valids, len - nulls_len, |a, b| cmp(a.1, b.1));
+ sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
+ cmp(a.1, b.1)
+ });
} else {
- sort_by(&mut valids, len - nulls_len, |a, b| cmp(a.1, b.1).reverse());
+ sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
+ cmp(a.1, b.1).reverse()
+ });
// reverse to keep a stable ordering
nulls.reverse();
}
// collect results directly into a buffer instead of a vec to avoid
another aligned allocation
- let mut result = MutableBuffer::new(values.len() *
std::mem::size_of::<u32>());
+ let result_capacity = len * std::mem::size_of::<u32>();
+ let mut result = MutableBuffer::new(result_capacity);
// sets len to capacity so we can access the whole buffer as a typed slice
- result.resize(values.len() * std::mem::size_of::<u32>(), 0);
+ result.resize(result_capacity, 0);
let result_slice: &mut [u32] = result.typed_data_mut();
- debug_assert_eq!(result_slice.len(), nulls_len + valids_len);
Review comment:
There is still some invariant around here (may not exactly this one).
Should we still try to assert it, for readability purposes?
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
