stacking: Improve .select() with special case for 1D arrays

The 1D case is a simpler gather operation since we only select 1 element
per index. Special-case it.

The performance increase for this benchmark is that the benchmark
runtime changes by -92% with this change.

Author: bluss

benches/append.rs

@@ -22,3 +22,14 @@ fn select_axis1(bench: &mut Bencher) {
         a.select(Axis(1), &selectable)
     });
 }
+
+#[bench]
+fn select_1d(bench: &mut Bencher) {
+    let a = Array::<f32, _>::zeros(1024);
+    let mut selectable = (0..a.len()).step_by(17).collect::<Vec<_>>();
+    selectable.extend(selectable.clone().iter().rev());
+
+    bench.iter(|| {
+        a.select(Axis(0), &selectable)
+    });
+}

src/impl_methods.rs

@@ -877,16 +877,35 @@ where
         S: Data,
         D: RemoveAxis,
     {
-        let mut subs = vec![self.view(); indices.len()];
-        for (&i, sub) in zip(indices, &mut subs[..]) {
-            sub.collapse_axis(axis, i);
-        }
-        if subs.is_empty() {
-            let mut dim = self.raw_dim();
-            dim.set_axis(axis, 0);
-            unsafe { Array::from_shape_vec_unchecked(dim, vec![]) }
+        if self.ndim() == 1 {
+            // using .len_of(axis) means that we check if `axis` is in bounds too.
+            let axis_len = self.len_of(axis);
+            // bounds check the indices first
+            if let Some(max_index) = indices.iter().cloned().max() {
+                if max_index >= axis_len {
+                    panic!("ndarray: index {} is out of bounds in array of len {}",
+                           max_index, self.len_of(axis));
+                }
+            } // else: indices empty is ok
+            let view = self.view().into_dimensionality::<Ix1>().unwrap();
+            Array::from_iter(indices.iter().map(move |&index| {
+                // Safety: bounds checked indexes
+                unsafe {
+                    view.uget(index).clone()
+                }
+            })).into_dimensionality::<D>().unwrap()
         } else {
-            concatenate(axis, &subs).unwrap()
+            let mut subs = vec![self.view(); indices.len()];
+            for (&i, sub) in zip(indices, &mut subs[..]) {
+                sub.collapse_axis(axis, i);
+            }
+            if subs.is_empty() {
+                let mut dim = self.raw_dim();
+                dim.set_axis(axis, 0);
+                unsafe { Array::from_shape_vec_unchecked(dim, vec![]) }
+            } else {
+                concatenate(axis, &subs).unwrap()
+            }
         }
     }
 

src/lib.rs

@@ -13,6 +13,7 @@
     clippy::unreadable_literal,
     clippy::manual_map, // is not an error
     clippy::while_let_on_iterator, // is not an error
+    clippy::from_iter_instead_of_collect, // using from_iter is good style
 )]
 #![cfg_attr(not(feature = "std"), no_std)]
 

tests/array.rs

@@ -709,6 +709,19 @@ fn test_select() {
     assert_abs_diff_eq!(c, c_target);
 }
 
+#[test]
+fn test_select_1d() {
+    let x = arr1(&[0, 1, 2, 3, 4, 5, 6]);
+    let r1 = x.select(Axis(0), &[1, 3, 4, 2, 2, 5]);
+    assert_eq!(r1, arr1(&[1, 3, 4, 2, 2, 5]));
+    // select nothing
+    let r2 = x.select(Axis(0), &[]);
+    assert_eq!(r2, arr1(&[]));
+    // select nothing from empty
+    let r3 = r2.select(Axis(0), &[]);
+    assert_eq!(r3, arr1(&[]));
+}
+
 #[test]
 fn diag() {
     let d = arr2(&[[1., 2., 3.0f32]]).into_diag();