Use the last dim first to take advantage of c-ordered linearization

Author: jsignell

xarray/namedarray/parallelcompat.py

@@ -823,7 +823,7 @@ def preserve_chunks(
         ...     typesize=8,
         ...     previous_chunks=(128, 128, 1),
         ... )
-        (512, 256, 1)
+        (128, 1024, 1)
 
         >>> ChunkManagerEntrypoint.preserve_chunks(
         ...     chunks=("preserve", "preserve", 1),
@@ -861,6 +861,7 @@ def preserve_chunks(
                 if isinstance(previous_chunk, tuple):
                     # For uniform chunks just take the first item
                     if previous_chunk[1:-1] == previous_chunk[:-2]:
+                        new_chunks[i] = previous_chunk[0]
                         previous_chunk = previous_chunk[0]
                     # For non-uniform chunks, leave them alone
                     else:
@@ -880,14 +881,13 @@ def preserve_chunks(
             return chunks
 
         while True:
-            # Repeatedly look for the dim with the most chunks and multiply it by 2.
+            # Repeatedly look for the last dim with more than one chunk and multiply it by 2.
             # Stop when:
             # 1a. we are larger than the target chunk size OR
             # 1b. we are within 50% of the target chunk size OR
             # 2. the chunk covers the entire array
 
             num_chunks = np.array(shape) / max_chunks * auto_dims
-            idx = np.argmax(num_chunks)
             chunk_bytes = np.prod(max_chunks) * typesize
 
             if chunk_bytes > target or abs(chunk_bytes - target) / target < 0.5:
@@ -896,6 +896,8 @@ def preserve_chunks(
             if (num_chunks <= 1).all():
                 break
 
+            idx = int(np.nonzero(num_chunks > 1)[0][-1])
+
             new_chunks[idx] = min(new_chunks[idx] * 2, shape[idx])
             max_chunks[idx] = new_chunks[idx]
 

xarray/tests/test_backends.py

@@ -7439,7 +7439,7 @@ def test_chunking_consistency(chunks, tmp_path: Path) -> None:
 @pytest.mark.parametrize(
     "chunks,expected",
     [
-        ("preserve", (320, 320)),
+        ("preserve", (160, 500)),
         (-1, (500, 500)),
         ({}, (10, 10)),
         ({"x": "preserve"}, (500, 10)),