Add Variable::chunking

georust · Aug 9, 2024 · a4fd2c9 · a4fd2c9
1 parent e7c28ac
commit a4fd2c9
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/netcdf/src/variable.rs b/netcdf/src/variable.rs
@@ -206,6 +206,35 @@ impl<'g> Variable<'g> {
     pub fn access_collective(&self) -> error::Result<()> {
         self.access_mode(crate::par::AccessMode::Collective)
     }
+
+    /// Get the chunking for the variable.
+    /// Returns None for a contiguous variable.
+    ///
+    /// # Errors
+    ///
+    /// Not a `netCDF-4` file.
+    pub fn chunking(&self) -> error::Result<Option<Vec<usize>>> {
+        if self.dimensions.is_empty() {
+            return Ok(None);
+        }
+        let mut storage = std::mem::MaybeUninit::<std::ffi::c_int>::uninit();
+        let mut chunk_size = Vec::with_capacity(self.dimensions.len());
+        checked_with_lock(|| unsafe {
+            nc_inq_var_chunking(
+                self.ncid,
+                self.varid,
+                storage.as_mut_ptr(),
+                chunk_size.spare_capacity_mut().as_mut_ptr().cast(),
+            )
+        })?;
+        let storage = unsafe { storage.assume_init() };
+        if storage == NC_CHUNKED {
+            unsafe { chunk_size.set_len(self.dimensions.len()) };
+            Ok(Some(chunk_size))
+        } else {
+            Ok(None)
+        }
+    }
 }
 impl<'g> VariableMut<'g> {
     /// Sets compression on the variable. Must be set before filling in data.

diff --git a/netcdf/tests/group.rs b/netcdf/tests/group.rs
@@ -66,13 +66,16 @@ fn find_variable() {
     assert!(group.variable("vvvvv").is_none());
 
     for mut var in group.variables_mut() {
+        assert_eq!(var.chunking().unwrap(), None);
         if !var.dimensions().is_empty() {
             var.set_compression(3, false).unwrap();
         }
         if var.name() == "z" {
             var.set_chunking(&[1]).unwrap();
+            assert_eq!(var.chunking().unwrap(), Some(vec![1]));
         } else {
             var.set_chunking(&[]).unwrap();
+            assert_eq!(var.chunking().unwrap(), None);
         }
     }
 }

diff --git a/netcdf/tests/lib.rs b/netcdf/tests/lib.rs
@@ -843,17 +843,21 @@ fn use_compression_chunking() {
     file.add_dimension("x", 10).unwrap();
 
     let var = &mut file.add_variable::<i32>("compressed", &["x"]).unwrap();
+    assert_eq!(var.chunking().unwrap(), None);
     var.set_compression(5, false).unwrap();
     var.set_chunking(&[5]).unwrap();
+    assert_eq!(var.chunking().unwrap(), Some(vec![5]));
 
     let v = vec![0i32; 10];
     var.put_values(&v, ..).unwrap();
 
     let var = &mut file
         .add_variable::<i32>("compressed2", &["x", "x"])
         .unwrap();
+    assert_eq!(var.chunking().unwrap(), None);
     var.set_compression(9, true).unwrap();
     var.set_chunking(&[5, 5]).unwrap();
+    assert_eq!(var.chunking().unwrap(), Some(vec![5, 5]));
     var.put_values(&[1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10], (..10, ..1))
         .unwrap();
 
@@ -866,7 +870,9 @@ fn use_compression_chunking() {
     file.add_dimension("y", 0).unwrap();
     let var = &mut file.add_variable::<u8>("chunked4", &["y", "x"]).unwrap();
 
+    assert_eq!(var.chunking().unwrap(), Some(vec![1, 10])); // unlimited dim enables chunking
     var.set_chunking(&[100, 2]).unwrap();
+    assert_eq!(var.chunking().unwrap(), Some(vec![100, 2]));
 }
 
 #[test]