Revise array_write_read and async_array_write_read examples

LDeakin · Jan 2, 2024 · a2975b0 · a2975b0
1 parent 492f808
commit a2975b0
Show file tree

Hide file tree

Showing 9 changed files with 148 additions and 81 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  - Use macros to reduce common code patterns in `Array`
  - Separate `Array` methods into separate files for each storage trait
  - **Breaking**: Remove `_opt` and `par_` variants of `async_retrieve_array_subset` and `async_store_array_subset` (including `_elements` and `_ndarray` variants)
+ - Revise `array_write_read` and `async_array_write_read` examples
 
 ## [0.8.0] - 2023-12-26
 

diff --git a/examples/array_write_read.rs b/examples/array_write_read.rs
@@ -1,8 +1,9 @@
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
 fn array_write_read() -> Result<(), Box<dyn std::error::Error>> {
-    use rayon::prelude::{IntoParallelIterator, ParallelIterator};
     use std::sync::Arc;
     use zarrs::{
-        array::{chunk_grid::ChunkGridTraits, DataType, FillValue, ZARR_NAN_F32},
+        array::{DataType, FillValue, ZARR_NAN_F32},
         array_subset::ArraySubset,
         node::Node,
         storage::store,
@@ -49,73 +50,95 @@ fn array_write_read() -> Result<(), Box<dyn std::error::Error>> {
     // Write array metadata to store
     array.store_metadata()?;
 
-    // Write some chunks (in parallel)
-    (0..2).into_par_iter().try_for_each(|i| {
-        let chunk_grid: &Box<dyn ChunkGridTraits> = array.chunk_grid();
-        let chunk_indices: Vec<u64> = vec![i, 0];
-        if let Some(chunk_subset) = chunk_grid.subset(&chunk_indices, array.shape())? {
-            array.store_chunk_elements(
-                &chunk_indices,
-                vec![i as f32; chunk_subset.num_elements() as usize],
-            )
-            // let chunk_shape = chunk_grid.chunk_shape(&chunk_indices, &array.shape())?;
-            // let chunk_array = ndarray::ArrayD::<f32>::from_elem(chunk_shape.clone(), i as f32);
-            // array.store_chunk_ndarray(&chunk_indices, &chunk_array.view())
-        } else {
-            Err(zarrs::array::ArrayError::InvalidChunkGridIndicesError(
-                chunk_indices.to_vec(),
-            ))
-        }
-    })?;
-
     println!(
         "The array metadata is:\n{}\n",
         serde_json::to_string_pretty(&array.metadata()).unwrap()
     );
 
+    // Write some chunks
+    (0..2).into_par_iter().try_for_each(|i| {
+        let chunk_indices: Vec<u64> = vec![0, i];
+        let chunk_subset = array
+            .chunk_grid()
+            .subset(&chunk_indices, array.shape())?
+            .ok_or_else(|| {
+                zarrs::array::ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec())
+            })?;
+        array.store_chunk_elements(
+            &chunk_indices,
+            vec![i as f32 * 0.1; chunk_subset.num_elements() as usize],
+        )
+    })?;
+
+    let subset_all = ArraySubset::new_with_shape(array.shape().to_vec());
+    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
+    println!("store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n");
+
+    // Store multiple chunks
+    array.store_chunks_elements_opt::<f32>(
+        &ArraySubset::new_with_ranges(&[1..2, 0..2]),
+        vec![
+            //
+            1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
+            //
+            1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
+        ],
+        true,
+    )?;
+    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
+    println!("store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n");
+
     // Write a subset spanning multiple chunks, including updating chunks already written
     array.store_array_subset_elements::<f32>(
         &ArraySubset::new_with_ranges(&[3..6, 3..6]),
-        vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
+        vec![-3.3, -3.4, -3.5, -4.3, -4.4, -4.5, -5.3, -5.4, -5.5],
     )?;
+    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
+    println!("store_array_subset [3..6, 3..6]:\n{data_all:+4.1}\n");
 
-    // Store elements directly, in this case set the 7th column to 123.0
+    // Store array subset
     array.store_array_subset_elements::<f32>(
         &ArraySubset::new_with_ranges(&[0..8, 6..7]),
-        vec![123.0; 8],
+        vec![-0.6, -1.6, -2.6, -3.6, -4.6, -5.6, -6.6, -7.6],
     )?;
+    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
+    println!("store_array_subset [0..8, 6..7]:\n{data_all:+4.1}\n");
 
-    // Store elements directly in a chunk, in this case set the last row of the bottom right chunk
+    // Store chunk subset
     array.store_chunk_subset_elements::<f32>(
         // chunk indices
         &[1, 1],
         // subset within chunk
         &ArraySubset::new_with_ranges(&[3..4, 0..4]),
-        vec![-4.0; 4],
+        vec![-7.4, -7.5, -7.6, -7.7],
     )?;
+    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
+    println!("store_chunk_subset [3..4, 0..4] of chunk [1, 1]:\n{data_all:+4.1}\n");
 
     // Erase a chunk
-    array.erase_chunk(&[0, 1])?;
-
-    // Read the whole array
-    let subset_all = ArraySubset::new_with_shape(array.shape().to_vec());
+    array.erase_chunk(&[0, 0])?;
     let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
-    println!("The whole array is:\n{:?}\n", data_all);
+    println!("erase_chunk [0, 0]:\n{data_all:+4.1}\n");
 
-    // Read a chunk back from the store
-    let chunk_indices = vec![1, 0];
+    // Read a chunk
+    let chunk_indices = vec![0, 1];
     let data_chunk = array.retrieve_chunk_ndarray::<f32>(&chunk_indices)?;
-    println!("Chunk [1,0] is:\n{data_chunk:?}\n");
+    println!("retrieve_chunk [0, 1]:\n{data_chunk:+4.1}\n");
+
+    // Read chunks
+    let chunks = ArraySubset::new_with_ranges(&[0..2, 1..2]);
+    let data_chunks = array.retrieve_chunks_ndarray::<f32>(&chunks)?;
+    println!("retrieve_chunks [0..2, 1..2]:\n{data_chunks:+4.1}\n");
 
-    // Read the central 4x2 subset of the array
-    let subset_4x2 = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
-    let data_4x2 = array.retrieve_array_subset_ndarray::<f32>(&subset_4x2)?;
-    println!("The middle 4x2 subset is:\n{:?}\n", data_4x2);
+    // Retrieve an array subset
+    let subset = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
+    let data_subset = array.retrieve_array_subset_ndarray::<f32>(&subset)?;
+    println!("retrieve_array_subset [2..6, 3..5]:\n{data_subset:+4.1}\n");
 
     // Show the hierarchy
     let node = Node::new_with_store(&*store, "/").unwrap();
     let tree = node.hierarchy_tree();
-    println!("The zarr hierarchy tree is:\n{}", tree);
+    println!("hierarchy_tree:\n{}", tree);
 
     Ok(())
 }

diff --git a/examples/async_array_write_read.rs b/examples/async_array_write_read.rs
@@ -51,92 +51,129 @@ async fn async_array_write_read() -> Result<(), Box<dyn std::error::Error>> {
     // Write array metadata to store
     array.async_store_metadata().await?;
 
-    // Write some chunks (in parallel)
+    println!(
+        "The array metadata is:\n{}\n",
+        serde_json::to_string_pretty(&array.metadata()).unwrap()
+    );
+
+    // Write some chunks
     let subsets = (0..2)
         .map(|i| {
-            let chunk_indices: Vec<u64> = vec![i, 0];
-            if let Some(chunk_subset) = array.chunk_grid().subset(&chunk_indices, array.shape())? {
-                Ok((i, chunk_indices, chunk_subset))
-            } else {
-                Err(zarrs::array::ArrayError::InvalidChunkGridIndicesError(
-                    chunk_indices.to_vec(),
-                ))
-            }
+            let chunk_indices: Vec<u64> = vec![0, i];
+            array
+                .chunk_grid()
+                .subset(&chunk_indices, array.shape())?
+                .ok_or_else(|| {
+                    zarrs::array::ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec())
+                })
+                .map(|chunk_subset| (i, chunk_indices, chunk_subset))
         })
         .collect::<Result<Vec<_>, _>>()?;
     let mut futures = subsets
         .iter()
         .map(|(i, chunk_indices, chunk_subset)| {
             array.async_store_chunk_elements(
                 &chunk_indices,
-                vec![*i as f32; chunk_subset.num_elements() as usize],
+                vec![*i as f32 * 0.1; chunk_subset.num_elements() as usize],
             )
         })
         .collect::<FuturesUnordered<_>>();
     while let Some(item) = futures.next().await {
         item?;
     }
 
-    println!(
-        "The array metadata is:\n{}\n",
-        serde_json::to_string_pretty(&array.metadata()).unwrap()
-    );
+    let subset_all = ArraySubset::new_with_shape(array.shape().to_vec());
+    let data_all = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
+        .await?;
+    println!("async_store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n");
+
+    // Store multiple chunks
+    array
+        .async_store_chunks_elements::<f32>(
+            &ArraySubset::new_with_ranges(&[1..2, 0..2]),
+            vec![
+                //
+                1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
+                //
+                1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
+            ],
+        )
+        .await?;
+    let data_all = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
+        .await?;
+    println!("async_store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n");
 
     // Write a subset spanning multiple chunks, including updating chunks already written
     array
         .async_store_array_subset_elements::<f32>(
             &ArraySubset::new_with_ranges(&[3..6, 3..6]),
-            vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
+            vec![-3.3, -3.4, -3.5, -4.3, -4.4, -4.5, -5.3, -5.4, -5.5],
         )
         .await?;
+    let data_all = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
+        .await?;
+    println!("async_store_array_subset [3..6, 3..6]:\n{data_all:+4.1}\n");
 
-    // Store elements directly, in this case set the 7th column to 123.0
+    // Store array subset
     array
         .async_store_array_subset_elements::<f32>(
             &ArraySubset::new_with_ranges(&[0..8, 6..7]),
-            vec![123.0; 8],
+            vec![-0.6, -1.6, -2.6, -3.6, -4.6, -5.6, -6.6, -7.6],
         )
         .await?;
+    let data_all = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
+        .await?;
+    println!("async_store_array_subset [0..8, 6..7]:\n{data_all:+4.1}\n");
 
-    // Store elements directly in a chunk, in this case set the last row of the bottom right chunk
+    // Store chunk subset
     array
         .async_store_chunk_subset_elements::<f32>(
             // chunk indices
             &[1, 1],
             // subset within chunk
             &ArraySubset::new_with_ranges(&[3..4, 0..4]),
-            vec![-4.0; 4],
+            vec![-7.4, -7.5, -7.6, -7.7],
         )
         .await?;
+    let data_all = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
+        .await?;
+    println!("async_store_chunk_subset [3..4, 0..4] of chunk [1, 1]:\n{data_all:+4.1}\n");
 
     // Erase a chunk
-    array.async_erase_chunk(&[0, 1]).await?;
-
-    // Read the whole array
-    let subset_all = ArraySubset::new_with_shape(array.shape().to_vec());
+    array.async_erase_chunk(&[0, 0]).await?;
     let data_all = array
         .async_retrieve_array_subset_ndarray::<f32>(&subset_all)
         .await?;
-    println!("The whole array is:\n{:?}\n", data_all);
+    println!("async_erase_chunk [0, 0]:\n{data_all:+4.1}\n");
 
-    // Read a chunk back from the store
-    let chunk_indices = vec![1, 0];
+    // Read a chunk
+    let chunk_indices = vec![0, 1];
     let data_chunk = array
         .async_retrieve_chunk_ndarray::<f32>(&chunk_indices)
         .await?;
-    println!("Chunk [1,0] is:\n{data_chunk:?}\n");
+    println!("async_retrieve_chunk [0, 1]:\n{data_chunk:+4.1}\n");
+
+    // Read chunks
+    let chunks = ArraySubset::new_with_ranges(&[0..2, 1..2]);
+    let data_chunks = array.async_retrieve_chunks_ndarray::<f32>(&chunks).await?;
+    println!("async_retrieve_chunks [0..2, 1..2]:\n{data_chunks:+4.1}\n");
 
-    // Read the central 4x2 subset of the array
-    let subset_4x2 = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
-    let data_4x2 = array
-        .async_retrieve_array_subset_ndarray::<f32>(&subset_4x2)
+    // Retrieve an array subset
+    let subset = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
+    let data_subset = array
+        .async_retrieve_array_subset_ndarray::<f32>(&subset)
         .await?;
-    println!("The middle 4x2 subset is:\n{:?}\n", data_4x2);
+    println!("async_retrieve_array_subset [2..6, 3..5]:\n{data_subset:+4.1}\n");
 
     // Show the hierarchy
     let node = Node::async_new_with_store(&*store, "/").await.unwrap();
     let tree = node.hierarchy_tree();
-    println!("The zarr hierarchy tree is:\n{}", tree);
+    println!("hierarchy_tree:\n{}", tree);
 
     Ok(())
 }

diff --git a/src/storage/store/store_sync/http_store.rs b/src/storage/store/store_sync/http_store.rs
@@ -252,15 +252,14 @@ mod tests {
         let data_4x2 = array
             .retrieve_array_subset_elements::<f32>(&subset_4x2)
             .unwrap();
-        // assert_eq!(data_4x2, &[0.0, f32::NAN, 0.1, f32::NAN, 0.4, 0.5, 0.7, 0.8]);
-        assert_eq!(data_4x2[0], 0.0);
-        assert!(data_4x2[1].is_nan());
-        assert_eq!(data_4x2[2], 0.1);
-        assert!(data_4x2[3].is_nan());
-        assert_eq!(data_4x2[4], 0.4);
-        assert_eq!(data_4x2[5], 0.5);
-        assert_eq!(data_4x2[6], 0.7);
-        assert_eq!(data_4x2[7], 0.8);
+        assert!(data_4x2[0].is_nan());
+        assert_eq!(data_4x2[1], 0.1);
+        assert!(data_4x2[2].is_nan());
+        assert_eq!(data_4x2[3], -3.4);
+        assert_eq!(data_4x2[4], -4.3);
+        assert_eq!(data_4x2[5], -4.4);
+        assert_eq!(data_4x2[6], -5.3);
+        assert_eq!(data_4x2[7], -5.4);
 
         // let data = array.retrieve_array_subset_ndarray::<f32>(&ArraySubset::new_with_shape(array.shape().to_vec())).unwrap();
         // println!("{data:?}");

diff --git a/tests/data/array_write_read.zarr/group/array/c/0/0 b/tests/data/array_write_read.zarr/group/array/c/0/0
diff --git a/tests/data/array_write_read.zarr/group/array/c/0/1 b/tests/data/array_write_read.zarr/group/array/c/0/1
diff --git a/tests/data/array_write_read.zarr/group/array/c/1/0 b/tests/data/array_write_read.zarr/group/array/c/1/0
diff --git a/tests/data/array_write_read.zarr/group/array/c/1/1 b/tests/data/array_write_read.zarr/group/array/c/1/1
diff --git a/tests/data/array_write_read.zarr/group/array/zarr.json b/tests/data/array_write_read.zarr/group/array/zarr.json
@@ -30,6 +30,13 @@
       }
     }
   ],
+  "attributes": {
+    "_zarrs": {
+      "description": "This array was created with zarrs",
+      "repository": "https://github.com/LDeakin/zarrs",
+      "version": "0.8.1"
+    }
+  },
   "dimension_names": [
     "y",
     "x"