diff --git a/src/impl_methods.rs b/src/impl_methods.rs index 8046f4d1f..9d6a0e001 100644 --- a/src/impl_methods.rs +++ b/src/impl_methods.rs @@ -3310,16 +3310,18 @@ impl ArrayRef panic!("Partition index {} is out of bounds for axis {} of length {}", kth, axis.0, axis_len); } - // Check if the first lane is contiguous - let is_contiguous = result + // Check if the first lane can be borrowed as a slice. A lane with a + // negative stride (e.g. along an inverted axis) is contiguous but can't. + let is_slice = result .lanes_mut(axis) .into_iter() .next() // This unwrap shouldn't cause panics because the array isn't empty .unwrap() - .is_contiguous(); + .as_slice_mut() + .is_some(); - if is_contiguous { + if is_slice { result.lanes_mut(axis).into_iter().for_each(|mut lane| { lane.as_slice_mut().unwrap().select_nth_unstable(kth); }); @@ -3560,6 +3562,28 @@ mod tests } } + #[test] + fn test_partition_negative_stride() + { + // Lanes along an inverted axis are contiguous, but their stride is negative + let mut a = array![[7, 1, 5, 3], [2, 6, 0, 4]]; + a.invert_axis(Axis(1)); + let p = a.partition(1, Axis(1)); + assert_eq!(p[[0, 1]], 3); + assert_eq!(p[[1, 1]], 2); + for lane in p.rows() { + assert!(lane[0] <= lane[1]); + assert!(lane.slice(s![2..]).iter().all(|&x| x >= lane[1])); + } + + // Same for a reversed view + let a = arr1(&[3, 1, 4, 1, 5, 9, 2, 6]); + let p = a.slice(s![..;-1]).partition(3, Axis(0)); + assert_eq!(p[3], 3); + assert!(p.slice(s![..3]).iter().all(|&x| x <= p[3])); + assert!(p.slice(s![4..]).iter().all(|&x| x >= p[3])); + } + #[test] fn test_partition_empty() {