Support byte shuffle filter

Both Zarr and HDF5 provide a byte shuffle filter which can be used to
improve compression ratios of certain data patterns.

There does not appear to be a Rust implementation of the byte shuffle
filter, so this change includes one.

It also adds a "filters" parameter to the API request data which accepts
a list of filter definitions. Each filter definition is a dict with an
"id" field and optional parameter fields specific to each filter. The
shuffle filter has one parameter, "element_size".

Author: markgoddard

README.md

@@ -77,7 +77,11 @@ with a JSON payload of the form:
 
     // Algorithm used to compress the data
     // - optional, defaults to no compression
-    "compression": "gzip|zlib"
+    "compression": {"id": "gzip|zlib"},
+
+    // List of algorithms used to filter the data
+    // - optional, defaults to no filters
+    "filters": [{"id": "shuffle", "element_size": 4}],
 }
 ```
 

scripts/client.py

@@ -37,7 +37,8 @@ def get_args() -> argparse.Namespace:
     parser.add_argument("--order", default="C") #, choices=["C", "F"]) allow invalid for testing
     parser.add_argument("--selection", type=str)
     parser.add_argument("--compression", type=str)
-    parser.add_argument("--show-response-headers", action=argparse.BooleanOptionalAction)
+    parser.add_argument("--shuffle", action=argparse.BooleanOptionalAction)
+    parser.add_argument("--verbose", action=argparse.BooleanOptionalAction)
     return parser.parse_args()
 
 
@@ -56,6 +57,14 @@ def build_request_data(args: argparse.Namespace) -> dict:
         request_data["shape"] = json.loads(args.shape)
     if args.selection:
         request_data["selection"] = json.loads(args.selection)
+    if args.compression:
+        request_data["compression"] = {"id": args.compression}
+    filters = []
+    if args.shuffle:
+        element_size = 4 if "32" in args.dtype else 8
+        filters.append({"id": "shuffle", "element_size": element_size})
+    if filters:
+        request_data["filters"] = filters
     return {k: v for k, v in request_data.items() if v is not None}
 
 
@@ -68,13 +77,13 @@ def request(url: str, username: str, password: str, request_data: dict):
     return response
 
 
-def display(response, show_headers=False):
+def display(response, verbose=False):
     #print(response.content)
     dtype = response.headers['x-activestorage-dtype']
     shape = json.loads(response.headers['x-activestorage-shape'])
     result = np.frombuffer(response.content, dtype=dtype)
     result = result.reshape(shape)
-    if show_headers:
+    if verbose:
         print("\nResponse headers:", response.headers)
         print("\nResult:", result)
     else:
@@ -92,10 +101,12 @@ def display_error(response):
 def main():
     args = get_args()
     request_data = build_request_data(args)
+    if args.verbose:
+        print("\nRequest data:", request_data)
     url = f'{args.server}/v1/{args.operation}/'
     response = request(url, args.username, args.password, request_data)
     if response.ok:
-        display(response, show_headers=args.show_response_headers)
+        display(response, verbose=args.verbose)
     else:
         display_error(response)
         sys.exit(1)

scripts/requirements.txt

@@ -1,3 +1,4 @@
+numcodecs
 numpy
 requests
 s3fs

scripts/upload_sample_data.py

@@ -1,5 +1,6 @@
 from enum import Enum
 import gzip
+import numcodecs
 import numpy as np
 import pathlib
 import s3fs
@@ -8,6 +9,7 @@
 NUM_ITEMS = 10
 OBJECT_PREFIX = "data"
 COMPRESSION_ALGS = [None, "gzip", "zlib"]
+FILTER_ALGS = [None, "shuffle"]
 
 #Use enum which also subclasses string type so that 
 # auto-generated OpenAPI schema can determine allowed dtypes
@@ -38,14 +40,18 @@ def n_bytes(self):
 # Create numpy arrays and upload to S3 as bytes
 for compression in COMPRESSION_ALGS:
     compression_suffix = f"-{compression}" if compression else ""
-    for d in AllowedDatatypes.__members__.keys():
-        obj_name = f'{OBJECT_PREFIX}-{d}{compression_suffix}.dat'
-        with s3_fs.open(bucket / obj_name, 'wb') as s3_file:
-            data = np.arange(NUM_ITEMS, dtype=d).tobytes()
-            if compression == "gzip":
-                data = gzip.compress(data)
-            elif compression == "zlib":
-                data = zlib.compress(data)
-            s3_file.write(data)
-
-print("Data upload successful. \nBucket contents:\n", s3_fs.ls(bucket))
+    for filter in FILTER_ALGS:
+        filter_suffix = f"-{filter}" if filter else ""
+        for d in AllowedDatatypes:
+            obj_name = f'{OBJECT_PREFIX}-{d}{compression_suffix}{filter_suffix}.dat'
+            with s3_fs.open(bucket / obj_name, 'wb') as s3_file:
+                data = np.arange(NUM_ITEMS, dtype=d).tobytes()
+                if filter == "shuffle":
+                    data = numcodecs.Shuffle(d.n_bytes()).encode(data)
+                if compression == "gzip":
+                    data = gzip.compress(data)
+                elif compression == "zlib":
+                    data = zlib.compress(data)
+                s3_file.write(data)
+
+print("Data upload successful. \nBucket contents:\n", "\n".join(s3_fs.ls(bucket)))

src/filter_pipeline.rs

@@ -2,6 +2,7 @@
 
 use crate::compression;
 use crate::error::ActiveStorageError;
+use crate::filters;
 use crate::models;
 
 use axum::body::Bytes;
@@ -13,24 +14,31 @@ use axum::body::Bytes;
 /// # Arguments
 ///
 /// * `request_data`: RequestData object for the request
-/// * `data`: Data to apply filter pipeline to.
+/// * `data`: Data [Bytes](axum::body::Bytes) to apply the pipeline to.
 pub fn filter_pipeline(
     request_data: &models::RequestData,
     data: &Bytes,
 ) -> Result<Bytes, ActiveStorageError> {
+    // Make a mutable shallow copy of the data.
+    let mut data = data.clone();
+    // First decompress.
     if let Some(compression) = request_data.compression {
-        compression::decompress(compression, data)
-    } else {
-        Ok(data.clone())
-    }
-    // TODO: Defilter
+        data = compression::decompress(compression, &data)?
+    };
+    // Then decode the filters in reverse order.
+    if let Some(filters) = &request_data.filters {
+        for filter in filters.iter().rev() {
+            data = filters::decode(filter, &data)?
+        }
+    };
+    Ok(data)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::test_utils;
-    use flate2::read::GzEncoder;
+    use flate2::read::{GzEncoder, ZlibEncoder};
     use flate2::Compression;
     use std::io::Read;
 
@@ -42,6 +50,14 @@ mod tests {
         result.into()
     }
 
+    fn compress_zlib(data: &[u8]) -> Bytes {
+        // Adapated from flate2 documentation.
+        let mut result = Vec::<u8>::new();
+        let mut deflater = ZlibEncoder::new(data, Compression::fast());
+        deflater.read_to_end(&mut result).unwrap();
+        result.into()
+    }
+
     #[test]
     fn test_filter_pipeline_noop() {
         let data = [1, 2, 3, 4];
@@ -60,4 +76,47 @@ mod tests {
         let result = filter_pipeline(&request_data, &bytes).unwrap();
         assert_eq!(data.as_ref(), result);
     }
+
+    #[test]
+    fn test_filter_pipeline_shuffle() {
+        let data = [1, 2, 3, 4, 5, 6, 7, 8];
+        let bytes = Bytes::copy_from_slice(&data);
+        let shuffled = filters::shuffle::test_utils::shuffle(&bytes, 4);
+        let mut request_data = test_utils::get_test_request_data();
+        request_data.filters = Some(vec![models::Filter::Shuffle { element_size: 4 }]);
+        let result = filter_pipeline(&request_data, &shuffled).unwrap();
+        assert_eq!(data.as_ref(), result);
+    }
+
+    #[test]
+    fn test_filter_pipeline_shuffle_zlib() {
+        let data: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
+        let bytes = Bytes::copy_from_slice(&data);
+        let shuffled = filters::shuffle::test_utils::shuffle(&bytes, 4);
+        let bytes = compress_zlib(shuffled.as_ref());
+        let mut request_data = test_utils::get_test_request_data();
+        request_data.compression = Some(models::Compression::Zlib);
+        request_data.filters = Some(vec![models::Filter::Shuffle { element_size: 4 }]);
+        let result = filter_pipeline(&request_data, &bytes).unwrap();
+        assert_eq!(data.as_ref(), result.as_ref());
+    }
+
+    #[test]
+    fn test_filter_pipeline_shuffle_x2_zlib() {
+        // Test multiple filters.
+        // Currently we only have shuffle, so run it twice with different element types.
+        let data: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
+        let bytes = Bytes::copy_from_slice(&data);
+        let shuffled = filters::shuffle::test_utils::shuffle(&bytes, 4);
+        let reshuffled = filters::shuffle::test_utils::shuffle(&shuffled, 2);
+        let bytes = compress_zlib(reshuffled.as_ref());
+        let mut request_data = test_utils::get_test_request_data();
+        request_data.compression = Some(models::Compression::Zlib);
+        request_data.filters = Some(vec![
+            models::Filter::Shuffle { element_size: 4 },
+            models::Filter::Shuffle { element_size: 2 },
+        ]);
+        let result = filter_pipeline(&request_data, &bytes).unwrap();
+        assert_eq!(data.as_ref(), result.as_ref());
+    }
 }

src/filters.rs

@@ -0,0 +1,36 @@
+//! Filter implementations.
+
+pub mod shuffle;
+
+use crate::error::ActiveStorageError;
+use crate::models;
+
+use axum::body::Bytes;
+
+/// Decodes some bytes using the specified filter and returns the result.
+///
+/// # Arguments
+///
+/// * `filter`: Filter algorithm
+/// * `data`: Filtered data [Bytes](axum::body::Bytes)
+pub fn decode(filter: &models::Filter, data: &Bytes) -> Result<Bytes, ActiveStorageError> {
+    match filter {
+        models::Filter::Shuffle { element_size } => Ok(shuffle::deshuffle(data, *element_size)),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::filters;
+
+    #[test]
+    fn test_decode_shuffle() {
+        let data = [1, 2, 3, 4, 5, 6, 7, 8];
+        let bytes = Bytes::copy_from_slice(&data);
+        let shuffled = filters::shuffle::test_utils::shuffle(&bytes, 4);
+        let filter = models::Filter::Shuffle { element_size: 4 };
+        let result = decode(&filter, &shuffled).unwrap();
+        assert_eq!(data.as_ref(), result);
+    }
+}