Add axis field to request data model

Author: sd109

benches/byte_order.rs

@@ -14,6 +14,7 @@ fn get_test_request_data() -> RequestData {
         offset: None,
         size: None,
         shape: None,
+        axis: None,
         order: None,
         selection: None,
         compression: None,

benches/operations.rs

@@ -19,6 +19,7 @@ fn get_test_request_data() -> RequestData {
         offset: None,
         size: None,
         shape: None,
+        axis: None,
         order: None,
         selection: None,
         compression: None,

src/models.rs

@@ -136,6 +136,8 @@ pub struct RequestData {
         custom = "validate_shape"
     )]
     pub shape: Option<Vec<usize>>,
+    /// Axis over which to perform the reduction operation
+    pub axis: Option<usize>,
     /// Order of the multi-dimensional array
     pub order: Option<Order>,
     /// Subset of the data to operate on
@@ -222,6 +224,7 @@ fn validate_request_data(request_data: &RequestData) -> Result<(), ValidationErr
             validate_raw_size(*size, request_data.dtype, &request_data.shape)?;
         }
     };
+    // Check selection is compatible with shape
     match (&request_data.shape, &request_data.selection) {
         (Some(shape), Some(selection)) => {
             validate_shape_selection(shape, selection)?;
@@ -233,6 +236,18 @@ fn validate_request_data(request_data: &RequestData) -> Result<(), ValidationErr
         }
         _ => (),
     };
+    // Check axis is compatible with shape
+    match (&request_data.shape, &request_data.axis) {
+        (Some(shape), Some(axis)) => {
+            if *axis > shape.len() - 1 {
+                return Err(ValidationError::new("Axis must be within shape"));
+            }
+        }
+        (None, Some(_)) => {
+            return Err(ValidationError::new("Axis requires shape to be specified"));
+        }
+        _ => (),
+    };
     if let Some(missing) = &request_data.missing {
         missing.validate(request_data.dtype)?;
     };
@@ -335,6 +350,9 @@ mod tests {
                 Token::U32(2),
                 Token::U32(5),
                 Token::SeqEnd,
+                Token::Str("axis"),
+                Token::Some,
+                Token::U32(1),
                 Token::Str("order"),
                 Token::Some,
                 Token::Enum { name: "Order" },
@@ -659,6 +677,23 @@ mod tests {
         request_data.validate().unwrap()
     }
 
+    #[test]
+    #[should_panic(expected = "Axis requires shape to be specified")]
+    fn test_axis_without_shape() {
+        let mut request_data = test_utils::get_test_request_data();
+        request_data.axis = Some(1);
+        request_data.validate().unwrap()
+    }
+
+    #[test]
+    #[should_panic(expected = "Axis must be within shape")]
+    fn test_axis_gt_shape() {
+        let mut request_data = test_utils::get_test_request_data();
+        request_data.axis = Some(2);
+        request_data.shape = Some(vec![2, 5]);
+        request_data.validate().unwrap()
+    }
+
     #[test]
     fn test_invalid_compression() {
         assert_de_tokens_error::<RequestData>(
@@ -731,7 +766,7 @@ mod tests {
             Token::Str("foo"),
             Token::StructEnd
             ],
-            "unknown field `foo`, expected one of `source`, `bucket`, `object`, `dtype`, `byte_order`, `offset`, `size`, `shape`, `order`, `selection`, `compression`, `filters`, `missing`"
+            "unknown field `foo`, expected one of `source`, `bucket`, `object`, `dtype`, `byte_order`, `offset`, `size`, `shape`, `axis`, `order`, `selection`, `compression`, `filters`, `missing`"
         )
     }
 
@@ -755,6 +790,7 @@ mod tests {
                         "offset": 4,
                         "size": 8,
                         "shape": [2, 5],
+                        "axis": 1,
                         "order": "C",
                         "selection": [[1, 2, 3], [4, 5, 6]],
                         "compression": {"id": "gzip"},
@@ -776,6 +812,7 @@ mod tests {
                         "offset": 4,
                         "size": 8,
                         "shape": [2, 5, 10],
+                        "axis": 2,
                         "order": "F",
                         "selection": [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                         "compression": {"id": "zlib"},
@@ -787,6 +824,7 @@ mod tests {
         expected.dtype = DType::Float64;
         expected.byte_order = Some(ByteOrder::Big);
         expected.shape = Some(vec![2, 5, 10]);
+        expected.axis = Some(2);
         expected.order = Some(Order::F);
         expected.selection = Some(vec![
             Slice::new(1, 2, 3),

src/test_utils.rs

@@ -13,6 +13,7 @@ pub(crate) fn get_test_request_data() -> RequestData {
         byte_order: None,
         offset: None,
         size: None,
+        axis: None,
         shape: None,
         order: None,
         selection: None,
@@ -32,6 +33,7 @@ pub(crate) fn get_test_request_data_optional() -> RequestData {
         byte_order: Some(ByteOrder::Little),
         offset: Some(4),
         size: Some(8),
+        axis: Some(1),
         shape: Some(vec![2, 5]),
         order: Some(Order::C),
         selection: Some(vec![Slice::new(1, 2, 3), Slice::new(4, 5, 6)]),