refactor: dead code

Author: ilan-gold

src/zarr/codecs/sharding.py

@@ -234,7 +234,6 @@ def create_empty(cls, chunks_per_shard: tuple[int, ...]) -> _ShardIndex:
 class _ShardReader(ShardMapping):
     buf: Buffer
     index: _ShardIndex
-    order: SubchunkWriteOrder
 
     @classmethod
     async def from_bytes(
@@ -540,8 +539,10 @@ async def _decode_partial_single(
         else:
             return out
 
-    def _subchunk_order_iter(self, chunks_per_shard: tuple[int, ...]) -> Iterable[tuple[int, ...]]:
-        match self.subchunk_write_order:
+    def _subchunk_order_iter(
+        self, chunks_per_shard: tuple[int, ...], subchunk_write_order: SubchunkWriteOrder
+    ) -> Iterable[tuple[int, ...]]:
+        match subchunk_write_order:
             case "morton":
                 subchunk_iter = morton_order_iter(chunks_per_shard)
             case "lexicographic":
@@ -574,10 +575,6 @@ async def _encode_single(
             )
         )
         shard_builder = dict.fromkeys(np.array(list(np.ndindex(chunks_per_shard))))
-        assert (
-            shard_builder.keys()
-            == dict.fromkeys(self._subchunk_order_iter(chunks_per_shard)).keys()
-        )
 
         await self.codec_pipeline.write(
             [
@@ -618,7 +615,7 @@ async def _encode_partial_single(
         )
 
         if self._is_complete_shard_write(indexer, chunks_per_shard):
-            shard_dict = dict.fromkeys(np.ndindex(chunks_per_shard))
+            shard_dict = dict.fromkeys(self._subchunk_order_iter(chunks_per_shard, "lexicographic"))
         else:
             shard_reader = await self._load_full_shard_maybe(
                 byte_getter=byte_setter,
@@ -628,7 +625,7 @@ async def _encode_partial_single(
             shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard)
             # Use vectorized lookup for better performance
             shard_dict = shard_reader.to_dict_vectorized(
-                np.array(list(np.ndindex(chunks_per_shard)))
+                np.array(list(self._subchunk_order_iter(chunks_per_shard, "lexicographic")))
             )
 
         await self.codec_pipeline.write(
@@ -667,7 +664,7 @@ async def _encode_shard_dict(
 
         template = buffer_prototype.buffer.create_zero_length()
         chunk_start = 0
-        for chunk_coords in self._subchunk_order_iter(chunks_per_shard):
+        for chunk_coords in self._subchunk_order_iter(chunks_per_shard, self.subchunk_write_order):
             value = map.get(chunk_coords)
             if value is None:
                 continue

tests/test_codecs/test_sharding.py

@@ -619,7 +619,7 @@ async def test_encoded_subchunk_write_order(subchunk_write_order: SubchunkWriteO
 
     actual_order = await stored_data_and_get_order(codec, chunks_per_shard)
     if subchunk_write_order != "unordered":
-        expected_order = list(codec._subchunk_order_iter(chunks_per_shard))
+        expected_order = list(codec._subchunk_order_iter(chunks_per_shard, subchunk_write_order))
         assert actual_order == expected_order
     else:
         same_order_same_seed = list(
@@ -630,7 +630,7 @@ async def test_encoded_subchunk_write_order(subchunk_write_order: SubchunkWriteO
                 index_location=ShardingCodecIndexLocation.end,
                 subchunk_write_order=subchunk_write_order,
                 rng=np.random.default_rng(seed=seed),
-            )._subchunk_order_iter(chunks_per_shard)
+            )._subchunk_order_iter(chunks_per_shard, subchunk_write_order)
         )
         assert actual_order == same_order_same_seed