tests:Add single chunk write test for sharding

Author: mkitti

tests/benchmarks/test_indexing.py

@@ -221,67 +221,29 @@ def compute_morton_order() -> None:
     benchmark(compute_morton_order)
 
 
-# Benchmark for Morton order in write path
-# Morton order is used when encoding shards, so writes should show improvement
 @pytest.mark.parametrize("store", ["memory"], indirect=["store"])
 @pytest.mark.parametrize("shards", large_morton_shards, ids=str)
-def test_sharded_morton_write_full(
+def test_sharded_morton_write_single_chunk(
     store: Store,
     shards: tuple[int, ...],
     benchmark: BenchmarkFixture,
 ) -> None:
-    """Benchmark writing a full shard with Morton order.
+    """Benchmark writing a single chunk to a large shard.
 
-    Morton order is used in the write/encode path of the sharding codec.
-    This benchmark writes to a full shard with cache clearing to measure
-    the Morton optimization impact on writes.
-    """
-    import numpy as np
-
-    from zarr.core.indexing import _morton_order
-
-    shape = tuple(s * 2 for s in shards)
-    chunks = (1,) * 3
-
-    data = create_array(
-        store=store,
-        shape=shape,
-        dtype="uint8",
-        chunks=chunks,
-        shards=shards,
-        compressors=None,
-        filters=None,
-        fill_value=0,
-    )
-
-    # Pre-create write data
-    write_data = np.ones(shards, dtype="uint8")
-    indexer = (slice(shards[0]),) * 3
+    This is the clearest end-to-end demonstration of Morton order optimization.
+    Writing a single chunk to a shard with 32^3 = 32768 chunks requires
+    computing the full Morton order, but minimizes I/O overhead.
 
-    def write_with_cache_clear() -> None:
-        _morton_order.cache_clear()
-        data[indexer] = write_data
-
-    benchmark(write_with_cache_clear)
-
-
-@pytest.mark.parametrize("store", ["memory"], indirect=["store"])
-@pytest.mark.parametrize("shards", morton_iter_shapes, ids=str)
-def test_sharded_morton_write(
-    store: Store,
-    shards: tuple[int, ...],
-    benchmark: BenchmarkFixture,
-) -> None:
-    """Benchmark writing to sharded arrays with various shard sizes.
-
-    Tests Morton order impact on writes across different shard sizes.
+    Expected improvement: ~160ms (matching Morton computation speedup of ~178ms).
+    The Morton order cache is cleared before each iteration.
     """
     import numpy as np
 
     from zarr.core.indexing import _morton_order
 
-    shape = tuple(s * 2 for s in shards)
-    chunks = (1,) * len(shards)
+    # 1x1x1 chunks means chunks_per_shard equals shard shape
+    shape = tuple(s * 2 for s in shards)  # 2 shards per dimension
+    chunks = (1,) * 3  # 1x1x1 chunks: chunks_per_shard = shards
 
     data = create_array(
         store=store,
@@ -294,8 +256,9 @@ def test_sharded_morton_write(
         fill_value=0,
     )
 
-    write_data = np.ones(shards, dtype="uint8")
-    indexer = tuple(slice(s) for s in shards)
+    # Write data for a single chunk
+    write_data = np.ones((1, 1, 1), dtype="uint8")
+    indexer = (slice(1), slice(1), slice(1))
 
     def write_with_cache_clear() -> None:
         _morton_order.cache_clear()