Repository: earth-mover/icechunk (language: python)

Issue:
Don't allow to write chunks to invalid chunk coordinates



Discussion / hints:
@jhamman do you think implementing this could cause any kind of problem with zarr-python? For example, this would demand that an array (metadata) is resized *before* it's new chunks are set.
@paraseba Would like to try and give this one a shot! Let me know any context and background.
@emacollins Great!

So, the key to this one is in the `set_chunk_ref` method of the `Repository` impl. In there we get the Array and apply the requested change. After getting the array, but before applying the change we would need to check if the passed `coord: ChunkIndices` are valid for the array. If they are not, we should fail.

To know if the coordinates are valid, we need to look into the array metadata, which you can find in the response of `Repository.get_array`.

You'll probably need to add a new error case to `RepositoryError`.

Probably some of our tests fail after this change, and also, we'll want to add tests that verify the new behavior.

Let me know if that's enough to get you started or if you want more context. And ping me with any questions.

Thank you!



@paraseba Thanks for the detailed explanation! I think this is enough to get started. I'll familiarize myself about with each of these parts and see what I can do
Sounds great @emacollins ! Take you time and let me know if you have questions.
Another change we would need, but it could be done in a separate ticket, would be to make sure that if the user shrinks an array by changing its metadata, all the chunks that are outside of the array bounds disappear from the manifest.

Fix patch:
diff --git a/icechunk/src/format/snapshot.rs b/icechunk/src/format/snapshot.rs
index dc53fa21..22afc89d 100644
--- a/icechunk/src/format/snapshot.rs
+++ b/icechunk/src/format/snapshot.rs
@@ -14,9 +14,9 @@ use crate::metadata::{
 };
 
 use super::{
-    format_constants, manifest::ManifestRef, AttributesId, IcechunkFormatError,
-    IcechunkFormatVersion, IcechunkResult, ManifestId, NodeId, ObjectId, Path,
-    SnapshotId, TableOffset,
+    format_constants, manifest::ManifestRef, AttributesId, ChunkIndices,
+    IcechunkFormatError, IcechunkFormatVersion, IcechunkResult, ManifestId, NodeId,
+    ObjectId, Path, SnapshotId, TableOffset,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -49,6 +49,54 @@ pub struct ZarrArrayMetadata {
     pub dimension_names: Option<DimensionNames>,
 }
 
+impl ZarrArrayMetadata {
+    /// Returns an iterator over the maximum permitted chunk indices for the array.
+    ///
+    /// This function calculates the maximum chunk indices based on the shape of the array
+    /// and the chunk shape, using (shape - 1) / chunk_shape. Given integer division is truncating,
+    /// this will always result in proper indices at the boundaries.
+    ///
+    /// # Returns
+    ///
+    /// A ChunkIndices type containing the max chunk index for each dimension.
+    fn max_chunk_indices_permitted(&self) -> ChunkIndices {
+        debug_assert_eq!(self.shape.len(), self.chunk_shape.0.len());
+
+        ChunkIndices(
+            self.shape
+                .iter()
+                .zip(self.chunk_shape.0.iter())
+                .map(|(s, cs)| if *s == 0 { 0 } else { ((s - 1) / cs.get()) as u32 })
+                .collect(),
+        )
+    }
+
+    /// Validates the provided chunk coordinates for the array.
+    ///
+    /// This function checks if the provided chunk indices are valid for the array.
+    ///
+    /// # Arguments
+    ///
+    /// * `coord` - The chunk indices to validate.
+    ///
+    /// # Returns
+    ///
+    /// An bool indicating whether the chunk coordinates are valid.
+    ///
+    /// # Errors
+    ///
+    /// Returns false if the chunk coordinates are invalid.
+    pub fn valid_chunk_coord(&self, coord: &ChunkIndices) -> bool {
+        debug_assert_eq!(self.shape.len(), coord.0.len());
+
+        coord
+            .0
+            .iter()
+            .zip(self.max_chunk_indices_permitted().0)
+            .all(|(index, index_permitted)| *index <= index_permitted)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum NodeData {
     Array(ZarrArrayMetadata, Vec<ManifestRef>),
@@ -450,4 +498,48 @@ mod tests {
         );
         Ok(())
     }
+
+    #[test]
+    fn test_valid_chunk_coord() {
+        let zarr_meta1 = ZarrArrayMetadata {
+            shape: vec![10000, 10001, 9999],
+            data_type: DataType::Float32,
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(1000).unwrap(),
+                NonZeroU64::new(1000).unwrap(),
+                NonZeroU64::new(1000).unwrap(),
+            ]),
+            chunk_key_encoding: ChunkKeyEncoding::Slash,
+            fill_value: FillValue::Float32(0f32),
+
+            codecs: vec![Codec {
+                name: "mycodec".to_string(),
+                configuration: Some(HashMap::from_iter(iter::once((
+                    "foo".to_string(),
+                    serde_json::Value::from(42),
+                )))),
+            }],
+            storage_transformers: None,
+            dimension_names: None,
+        };
+
+        let zarr_meta2 = ZarrArrayMetadata {
+            shape: vec![0, 0, 0],
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(1000).unwrap(),
+                NonZeroU64::new(1000).unwrap(),
+                NonZeroU64::new(1000).unwrap(),
+            ]),
+            ..zarr_meta1.clone()
+        };
+
+        let coord1 = ChunkIndices(vec![9, 10, 9]);
+        let coord2 = ChunkIndices(vec![10, 11, 10]);
+        let coord3 = ChunkIndices(vec![0, 0, 0]);
+
+        assert!(zarr_meta1.valid_chunk_coord(&coord1));
+        assert!(!zarr_meta1.valid_chunk_coord(&coord2));
+
+        assert!(zarr_meta2.valid_chunk_coord(&coord3));
+    }
 }
diff --git a/icechunk/src/repository.rs b/icechunk/src/repository.rs
index 1dc7c075..d111558f 100644
--- a/icechunk/src/repository.rs
+++ b/icechunk/src/repository.rs
@@ -188,6 +188,10 @@ pub enum RepositoryError {
     ConflictingPathNotFound(NodeId),
     #[error("error in config deserialization `{0}`")]
     ConfigDeserializationError(#[from] serde_yml::Error),
+    #[error(
+        "invalid chunk index: coordinates {coords:?} are not valid for array at {path}"
+    )]
+    InvalidIndex { coords: ChunkIndices, path: Path },
 }
 
 pub type RepositoryResult<T> = Result<T, RepositoryError>;
@@ -510,9 +514,21 @@ impl Repository {
         coord: ChunkIndices,
         data: Option<ChunkPayload>,
     ) -> RepositoryResult<()> {
-        self.get_array(&path)
-            .await
-            .map(|node: NodeSnapshot| self.change_set.set_chunk_ref(node.id, coord, data))
+        let node_snapshot = self.get_array(&path).await?;
+
+        if let NodeData::Array(zarr_metadata, _) = node_snapshot.node_data {
+            if zarr_metadata.valid_chunk_coord(&coord) {
+                self.change_set.set_chunk_ref(node_snapshot.id, coord, data);
+                Ok(())
+            } else {
+                Err(RepositoryError::InvalidIndex { coords: coord, path: path.clone() })
+            }
+        } else {
+            Err(RepositoryError::NotAnArray {
+                node: node_snapshot,
+                message: "getting an array".to_string(),
+            })
+        }
     }
 
     pub async fn get_node(&self, path: &Path) -> RepositoryResult<NodeSnapshot> {
@@ -1807,7 +1823,11 @@ mod tests {
         let zarr_meta = ZarrArrayMetadata {
             shape: vec![1, 1, 2],
             data_type: DataType::Float16,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(1).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Float16(f32::NEG_INFINITY),
             codecs: vec![Codec { name: "mycodec".to_string(), configuration: None }],
@@ -2054,7 +2074,10 @@ mod tests {
         let zarr_meta = ZarrArrayMetadata {
             shape: vec![5, 5],
             data_type: DataType::Float16,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Float16(f32::NEG_INFINITY),
             codecs: vec![Codec { name: "mycodec".to_string(), configuration: None }],
@@ -2209,9 +2232,13 @@ mod tests {
         // add a new array and retrieve its node
         ds.add_group(Path::root()).await?;
         let zarr_meta = ZarrArrayMetadata {
-            shape: vec![1, 1, 2],
+            shape: vec![4, 2, 4],
             data_type: DataType::Int32,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(1).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Int32(0),
             codecs: vec![Codec { name: "mycodec".to_string(), configuration: None }],
@@ -2243,6 +2270,12 @@ mod tests {
             Some(ChunkPayload::Inline("hello".into())),
         )
         .await?;
+        ds.set_chunk_ref(
+            new_array_path.clone(),
+            ChunkIndices(vec![0, 1, 0]),
+            Some(ChunkPayload::Inline("hello".into())),
+        )
+        .await?;
         let snapshot_id = ds.flush("commit", SnapshotProperties::default()).await?;
         let ds = Repository::update(Arc::clone(&storage), snapshot_id).await?.build();
         let coords = ds
@@ -2256,7 +2289,8 @@ mod tests {
             vec![
                 ChunkIndices(vec![0, 0, 0]),
                 ChunkIndices(vec![0, 0, 1]),
-                ChunkIndices(vec![1, 0, 0])
+                ChunkIndices(vec![1, 0, 0]),
+                ChunkIndices(vec![0, 1, 0])
             ]
             .into_iter()
             .collect()
@@ -2301,7 +2335,11 @@ mod tests {
         let zarr_meta = ZarrArrayMetadata {
             shape: vec![1, 1, 2],
             data_type: DataType::Int32,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Int32(0),
             codecs: vec![Codec { name: "mycodec".to_string(), configuration: None }],
@@ -2393,6 +2431,72 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_setting_w_invalid_coords() -> Result<(), Box<dyn Error>> {
+        let in_mem_storage =
+            Arc::new(ObjectStorage::new_in_memory_store(Some("prefix".into())));
+        let storage: Arc<dyn Storage + Send + Sync> = in_mem_storage.clone();
+        let mut ds = Repository::init(Arc::clone(&storage), false).await?.build();
+
+        ds.add_group(Path::root()).await?;
+        let zarr_meta = ZarrArrayMetadata {
+            shape: vec![5, 5],
+            data_type: DataType::Float16,
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+            ]),
+            chunk_key_encoding: ChunkKeyEncoding::Slash,
+            fill_value: FillValue::Float16(f32::NEG_INFINITY),
+            codecs: vec![Codec { name: "mycodec".to_string(), configuration: None }],
+            storage_transformers: None,
+            dimension_names: None,
+        };
+
+        let apath: Path = "/array1".try_into()?;
+
+        ds.add_array(apath.clone(), zarr_meta.clone()).await?;
+
+        ds.commit("main", "first commit", None).await?;
+
+        // add 3 chunks
+        // First 2 chunks are valid, third will be invalid chunk indices
+
+        assert!(ds
+            .set_chunk_ref(
+                apath.clone(),
+                ChunkIndices(vec![0, 0]),
+                Some(ChunkPayload::Inline("hello".into())),
+            )
+            .await
+            .is_ok());
+        assert!(ds
+            .set_chunk_ref(
+                apath.clone(),
+                ChunkIndices(vec![2, 2]),
+                Some(ChunkPayload::Inline("hello".into())),
+            )
+            .await
+            .is_ok());
+
+        let bad_result = ds
+            .set_chunk_ref(
+                apath.clone(),
+                ChunkIndices(vec![3, 0]),
+                Some(ChunkPayload::Inline("hello".into())),
+            )
+            .await;
+
+        match bad_result {
+            Err(RepositoryError::InvalidIndex { coords, path }) => {
+                assert_eq!(coords, ChunkIndices(vec![3, 0]));
+                assert_eq!(path, apath);
+            }
+            _ => panic!("Expected InvalidIndex Error"),
+        }
+        Ok(())
+    }
+
     /// Construct two repos on the same storage, with a commit, a group and an array
     ///
     /// Group: /foo/bar
diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs
index b0fc66a5..3c02cad7 100644
--- a/icechunk/src/zarr.rs
+++ b/icechunk/src/zarr.rs
@@ -1940,8 +1940,12 @@ mod tests {
             store.delete("array/foo").await,
             Err(StoreError::InvalidKey { key }) if key == "array/foo",
         ));
-        // FIXME: deleting an invalid chunk should not be allowed.
-        store.delete("array/c/10/1/1").await.unwrap();
+
+        assert!(matches!(
+            store.delete("array/c/10/1/1").await,
+            Err(StoreError::RepositoryError(RepositoryError::InvalidIndex { coords, path }))
+                if path.to_string() == "/array" && coords == ChunkIndices([10, 1, 1].to_vec())
+        ));
 
         Ok(())
     }


Test patch:
diff --git a/icechunk/tests/test_distributed_writes.rs b/icechunk/tests/test_distributed_writes.rs
index 943f40ee..8ff16d1a 100644
--- a/icechunk/tests/test_distributed_writes.rs
+++ b/icechunk/tests/test_distributed_writes.rs
@@ -126,8 +126,8 @@ async fn test_distributed_writes() -> Result<(), Box<dyn std::error::Error + Sen
         shape: vec![SIZE as u64, SIZE as u64],
         data_type: DataType::Float64,
         chunk_shape: ChunkShape(vec![
-            NonZeroU64::new(2).unwrap(),
-            NonZeroU64::new(2).unwrap(),
+            NonZeroU64::new(1).unwrap(),
+            NonZeroU64::new(1).unwrap(),
         ]),
         chunk_key_encoding: ChunkKeyEncoding::Slash,
         fill_value: FillValue::Float64(f64::NAN),
diff --git a/icechunk/tests/test_virtual_refs.rs b/icechunk/tests/test_virtual_refs.rs
index b3980dae..7987d60f 100644
--- a/icechunk/tests/test_virtual_refs.rs
+++ b/icechunk/tests/test_virtual_refs.rs
@@ -143,7 +143,11 @@ mod tests {
         let zarr_meta = ZarrArrayMetadata {
             shape: vec![1, 1, 2],
             data_type: DataType::Int32,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(1).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Int32(0),
             codecs: vec![],
@@ -253,7 +257,11 @@ mod tests {
         let zarr_meta = ZarrArrayMetadata {
             shape: vec![1, 1, 2],
             data_type: DataType::Int32,
-            chunk_shape: ChunkShape(vec![NonZeroU64::new(2).unwrap()]),
+            chunk_shape: ChunkShape(vec![
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(2).unwrap(),
+                NonZeroU64::new(1).unwrap(),
+            ]),
             chunk_key_encoding: ChunkKeyEncoding::Slash,
             fill_value: FillValue::Int32(0),
             codecs: vec![],