Merge pull request #41 from stackhpc/negative-selection

Support negative values in selection

Author: markgoddard

src/array.rs

@@ -59,29 +59,61 @@ fn build_array_from_shape<T>(
     ArrayView::<T, _>::from_shape(shape, data).map_err(ActiveStorageError::ShapeInvalid)
 }
 
-/// Returns an optional [ndarray] SliceInfo object corresponding to the selection.
+/// Returns an array index in numpy semantics to an index with ndarray semantics.
+///
+/// The resulting value will be clamped such that it is safe for indexing in ndarray.
+/// This allows us to accept selections with NumPy's less restrictive semantics.
+/// When the stride is negative (`reverse` is `true`), the result is offset by one to allow for
+/// Numpy's non-inclusive start and inclusive end in this scenario.
+///
+/// # Arguments
+///
+/// * `index`: Selection index
+/// * `length`: Length of corresponding axis
+/// * `reverse`: Whether the stride is negative
+fn to_ndarray_index(index: isize, length: usize, reverse: bool) -> isize {
+    let length_isize = length.try_into().expect("Length too large!");
+    let result = if reverse { index + 1 } else { index };
+    if index < 0 {
+        std::cmp::max(result + length_isize, 0)
+    } else {
+        std::cmp::min(result, length_isize)
+    }
+}
+
+/// Convert a [crate::models::Slice] object with indices in numpy semantics to an
+/// [ndarray::SliceInfoElem::Slice] with ndarray semantics.
+///
+/// See [ndarray docs](https://docs.rs/ndarray/0.15.6/ndarray/macro.s.html#negative-step) for
+/// information about ndarray's handling of negative strides.
+fn to_ndarray_slice(slice: &models::Slice, length: usize) -> ndarray::SliceInfoElem {
+    let reverse = slice.stride < 0;
+    let start = to_ndarray_index(slice.start, length, reverse);
+    let end = to_ndarray_index(slice.end, length, reverse);
+    let (start, end) = if reverse { (end, start) } else { (start, end) };
+    ndarray::SliceInfoElem::Slice {
+        start,
+        end: Some(end),
+        step: slice.stride,
+    }
+}
+
+/// Returns an [ndarray] SliceInfo object corresponding to the selection.
 pub fn build_slice_info<T>(
     selection: &Option<Vec<models::Slice>>,
     shape: &[usize],
 ) -> ndarray::SliceInfo<Vec<ndarray::SliceInfoElem>, ndarray::IxDyn, ndarray::IxDyn> {
     match selection {
         Some(selection) => {
-            let si: Vec<ndarray::SliceInfoElem> = selection
-                .iter()
-                .map(|slice| ndarray::SliceInfoElem::Slice {
-                    // FIXME: usize should be isize?
-                    start: slice.start as isize,
-                    end: Some(slice.end as isize),
-                    step: slice.stride as isize,
-                })
+            let si: Vec<ndarray::SliceInfoElem> = std::iter::zip(selection, shape)
+                .map(|(slice, length)| to_ndarray_slice(slice, *length))
                 .collect();
             ndarray::SliceInfo::try_from(si).expect("SliceInfo should not fail for IxDyn")
         }
         _ => {
             let si: Vec<ndarray::SliceInfoElem> = shape
                 .iter()
                 .map(|_| ndarray::SliceInfoElem::Slice {
-                    // FIXME: usize should be isize?
                     start: 0,
                     end: None,
                     step: 1,
@@ -309,7 +341,7 @@ mod tests {
     #[test]
     fn build_slice_info_1d_selection() {
         let selection = Some(vec![models::Slice::new(0, 1, 1)]);
-        let shape = [];
+        let shape = [1];
         let slice_info = build_slice_info::<u32>(&selection, &shape);
         assert_eq!(
             [ndarray::SliceInfoElem::Slice {
@@ -321,6 +353,51 @@ mod tests {
         );
     }
 
+    #[test]
+    fn build_slice_info_1d_selection_negative_stride() {
+        let selection = Some(vec![models::Slice::new(1, 0, -1)]);
+        let shape = [1];
+        let slice_info = build_slice_info::<u32>(&selection, &shape);
+        assert_eq!(
+            [ndarray::SliceInfoElem::Slice {
+                start: 1,
+                end: Some(1),
+                step: -1
+            }],
+            slice_info.as_ref()
+        );
+    }
+
+    #[test]
+    fn build_slice_info_1d_selection_negative_start() {
+        let selection = Some(vec![models::Slice::new(-1, 1, 1)]);
+        let shape = [1];
+        let slice_info = build_slice_info::<u32>(&selection, &shape);
+        assert_eq!(
+            [ndarray::SliceInfoElem::Slice {
+                start: 0,
+                end: Some(1),
+                step: 1
+            }],
+            slice_info.as_ref()
+        );
+    }
+
+    #[test]
+    fn build_slice_info_1d_selection_negative_end() {
+        let selection = Some(vec![models::Slice::new(0, -1, 1)]);
+        let shape = [1];
+        let slice_info = build_slice_info::<u32>(&selection, &shape);
+        assert_eq!(
+            [ndarray::SliceInfoElem::Slice {
+                start: 0,
+                end: Some(0),
+                step: 1
+            }],
+            slice_info.as_ref()
+        );
+    }
+
     #[test]
     fn build_slice_info_2d_no_selection() {
         let selection = None;
@@ -349,7 +426,7 @@ mod tests {
             models::Slice::new(0, 1, 1),
             models::Slice::new(0, 1, 1),
         ]);
-        let shape = [];
+        let shape = [1, 1];
         let slice_info = build_slice_info::<u32>(&selection, &shape);
         assert_eq!(
             [
@@ -405,4 +482,136 @@ mod tests {
         let array = build_array::<i64>(&request_data, &bytes).unwrap();
         assert_eq!(array![[0x04030201_i64], [0x08070605_i64]].into_dyn(), array);
     }
+
+    // Helper function for tests that slice an array using a selection.
+    fn test_selection(slice: models::Slice, expected: Array1<u32>) {
+        let data = [1, 2, 3, 4, 5, 6, 7, 8];
+        let request_data = models::RequestData {
+            source: Url::parse("http://example.com").unwrap(),
+            bucket: "bar".to_string(),
+            object: "baz".to_string(),
+            dtype: models::DType::Uint32,
+            offset: None,
+            size: None,
+            shape: None,
+            order: None,
+            selection: None,
+        };
+        let bytes = Bytes::copy_from_slice(&data);
+        let array = build_array::<u32>(&request_data, &bytes).unwrap();
+        let shape = vec![2];
+        let slice_info = build_slice_info::<u32>(&Some(vec![slice]), &shape);
+        let sliced = array.slice(slice_info);
+        assert_eq!(sliced, expected.into_dyn().view());
+    }
+
+    #[test]
+    fn build_array_with_selection_all() {
+        test_selection(
+            models::Slice::new(0, 2, 1),
+            array![0x04030201_u32, 0x08070605_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_negative_start() {
+        test_selection(
+            models::Slice::new(-2, 2, 1),
+            array![0x04030201_u32, 0x08070605_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_start_lt_negative_length() {
+        test_selection(
+            models::Slice::new(-3, 2, 1),
+            array![0x04030201_u32, 0x08070605_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_start_eq_length() {
+        test_selection(models::Slice::new(2, 2, 1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_start_gt_length() {
+        test_selection(models::Slice::new(3, 2, 1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_negative_end() {
+        test_selection(models::Slice::new(0, -1, 1), array![0x04030201_u32])
+    }
+
+    #[test]
+    fn build_array_with_selection_end_lt_negative_length() {
+        test_selection(models::Slice::new(0, -3, 1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_end_gt_length() {
+        test_selection(
+            models::Slice::new(0, 3, 1),
+            array![0x04030201_u32, 0x08070605_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_all_negative_stride() {
+        // Need to end at -3 to read first item.
+        // translates to [0, 2]
+        test_selection(
+            models::Slice::new(1, -3, -1),
+            array![0x08070605_u32, 0x04030201_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_negative_start_negative_stride() {
+        // translates to [0, 2]
+        test_selection(
+            models::Slice::new(-1, -3, -1),
+            array![0x08070605_u32, 0x04030201_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_start_lt_negative_length_negative_stride() {
+        // translates to [1, 0]
+        test_selection(models::Slice::new(-3, 0, -1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_start_eq_length_negative_stride() {
+        // translates to [2, 2]
+        test_selection(models::Slice::new(2, 1, -1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_start_gt_length_negative_stride() {
+        // translates to [2, 2]
+        test_selection(models::Slice::new(3, 1, -1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_negative_end_negative_stride() {
+        // translates to [2, 2]
+        test_selection(models::Slice::new(2, -1, -1), array![])
+    }
+
+    #[test]
+    fn build_array_with_selection_end_lt_negative_length_negative_stride() {
+        // translates to [0, 2]
+        test_selection(
+            models::Slice::new(1, -3, -1),
+            array![0x08070605_u32, 0x04030201_u32],
+        )
+    }
+
+    #[test]
+    fn build_array_with_selection_end_gt_length_negative_stride() {
+        // translates to [1, 2]
+        test_selection(models::Slice::new(3, 0, -1), array![0x08070605_u32])
+    }
 }