Skip to content

Commit 36a1bac

Browse files
authored
Fix specifying memory order in v2 arrays (#2951)
* Fix specifying memory order in v2 arrays * Re-work test_order * Fix getting array order in v3 * Fix order of arrays for v2 * Fix order with V3 arrays * Fix mypy * Remove errant print() * Fix order with v3 arrays * Fix v2 test * Add numpy order parametrization * Add changelog entry
1 parent 0c76778 commit 36a1bac

File tree

6 files changed

+68
-55
lines changed

6 files changed

+68
-55
lines changed

changes/2950.bufgix.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Specifying the memory order of Zarr format 2 arrays using the ``order`` keyword argument has been fixed.

src/zarr/api/asynchronous.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1040,15 +1040,13 @@ async def create(
10401040
)
10411041
warnings.warn(UserWarning(msg), stacklevel=1)
10421042
config_dict["write_empty_chunks"] = write_empty_chunks
1043-
if order is not None:
1044-
if config is not None:
1045-
msg = (
1046-
"Both order and config keyword arguments are set. "
1047-
"This is redundant. When both are set, order will be ignored and "
1048-
"config will be used."
1049-
)
1050-
warnings.warn(UserWarning(msg), stacklevel=1)
1051-
config_dict["order"] = order
1043+
if order is not None and config is not None:
1044+
msg = (
1045+
"Both order and config keyword arguments are set. "
1046+
"This is redundant. When both are set, order will be ignored and "
1047+
"config will be used."
1048+
)
1049+
warnings.warn(UserWarning(msg), stacklevel=1)
10521050

10531051
config_parsed = ArrayConfig.from_dict(config_dict)
10541052

@@ -1062,6 +1060,7 @@ async def create(
10621060
overwrite=overwrite,
10631061
filters=filters,
10641062
dimension_separator=dimension_separator,
1063+
order=order,
10651064
zarr_format=zarr_format,
10661065
chunk_shape=chunk_shape,
10671066
chunk_key_encoding=chunk_key_encoding,

src/zarr/core/array.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -609,6 +609,7 @@ async def _create(
609609

610610
if order is not None:
611611
_warn_order_kwarg()
612+
config_parsed = replace(config_parsed, order=order)
612613

613614
result = await cls._create_v3(
614615
store_path,
@@ -1044,7 +1045,10 @@ def order(self) -> MemoryOrder:
10441045
bool
10451046
Memory order of the array
10461047
"""
1047-
return self._config.order
1048+
if self.metadata.zarr_format == 2:
1049+
return self.metadata.order
1050+
else:
1051+
return self._config.order
10481052

10491053
@property
10501054
def attrs(self) -> dict[str, JSON]:
@@ -1276,14 +1280,14 @@ async def _get_selection(
12761280
out_buffer = prototype.nd_buffer.create(
12771281
shape=indexer.shape,
12781282
dtype=out_dtype,
1279-
order=self._config.order,
1283+
order=self.order,
12801284
fill_value=self.metadata.fill_value,
12811285
)
12821286
if product(indexer.shape) > 0:
12831287
# need to use the order from the metadata for v2
12841288
_config = self._config
12851289
if self.metadata.zarr_format == 2:
1286-
_config = replace(_config, order=self.metadata.order)
1290+
_config = replace(_config, order=self.order)
12871291

12881292
# reading chunks and decoding them
12891293
await self.codec_pipeline.read(
@@ -4256,6 +4260,11 @@ async def init_array(
42564260
chunks_out = chunk_shape_parsed
42574261
codecs_out = sub_codecs
42584262

4263+
if config is None:
4264+
config = {}
4265+
if order is not None and isinstance(config, dict):
4266+
config["order"] = config.get("order", order)
4267+
42594268
meta = AsyncArray._create_metadata_v3(
42604269
shape=shape_parsed,
42614270
dtype=dtype_parsed,

tests/test_api.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -326,13 +326,12 @@ def test_array_order(zarr_format: ZarrFormat) -> None:
326326
def test_array_order_warns(order: MemoryOrder | None, zarr_format: ZarrFormat) -> None:
327327
with pytest.warns(RuntimeWarning, match="The `order` keyword argument .*"):
328328
arr = zarr.ones(shape=(2, 2), order=order, zarr_format=zarr_format)
329-
expected = order or zarr.config.get("array.order")
330-
assert arr.order == expected
329+
assert arr.order == order
331330

332331
vals = np.asarray(arr)
333-
if expected == "C":
332+
if order == "C":
334333
assert vals.flags.c_contiguous
335-
elif expected == "F":
334+
elif order == "F":
336335
assert vals.flags.f_contiguous
337336
else:
338337
raise AssertionError

tests/test_array.py

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1262,39 +1262,43 @@ async def test_data_ignored_params(store: Store) -> None:
12621262
await create_array(store, data=data, shape=None, dtype=data.dtype, overwrite=True)
12631263

12641264
@staticmethod
1265-
@pytest.mark.parametrize("order_config", ["C", "F", None])
1265+
@pytest.mark.parametrize("order", ["C", "F", None])
1266+
@pytest.mark.parametrize("with_config", [True, False])
12661267
def test_order(
1267-
order_config: MemoryOrder | None,
1268+
order: MemoryOrder | None,
1269+
with_config: bool,
12681270
zarr_format: ZarrFormat,
12691271
store: MemoryStore,
12701272
) -> None:
12711273
"""
12721274
Test that the arrays generated by array indexing have a memory order defined by the config order
12731275
value, and that for zarr v2 arrays, the ``order`` field in the array metadata is set correctly.
12741276
"""
1275-
config: ArrayConfigLike = {}
1276-
if order_config is None:
1277+
config: ArrayConfigLike | None = {}
1278+
if order is None:
12771279
config = {}
12781280
expected = zarr.config.get("array.order")
12791281
else:
1280-
config = {"order": order_config}
1281-
expected = order_config
1282+
config = {"order": order}
1283+
expected = order
1284+
1285+
if not with_config:
1286+
# Test without passing config parameter
1287+
config = None
1288+
1289+
arr = zarr.create_array(
1290+
store=store,
1291+
shape=(2, 2),
1292+
zarr_format=zarr_format,
1293+
dtype="i4",
1294+
order=order,
1295+
config=config,
1296+
)
1297+
assert arr.order == expected
12821298
if zarr_format == 2:
1283-
arr = zarr.create_array(
1284-
store=store,
1285-
shape=(2, 2),
1286-
zarr_format=zarr_format,
1287-
dtype="i4",
1288-
order=expected,
1289-
config=config,
1290-
)
1291-
# guard for type checking
12921299
assert arr.metadata.zarr_format == 2
12931300
assert arr.metadata.order == expected
1294-
else:
1295-
arr = zarr.create_array(
1296-
store=store, shape=(2, 2), zarr_format=zarr_format, dtype="i4", config=config
1297-
)
1301+
12981302
vals = np.asarray(arr)
12991303
if expected == "C":
13001304
assert vals.flags.c_contiguous

tests/test_v2.py

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -195,12 +195,13 @@ def test_create_array_defaults(store: Store):
195195
)
196196

197197

198-
@pytest.mark.parametrize("array_order", ["C", "F"])
199-
@pytest.mark.parametrize("data_order", ["C", "F"])
200-
@pytest.mark.parametrize("memory_order", ["C", "F"])
201-
def test_v2_non_contiguous(
202-
array_order: Literal["C", "F"], data_order: Literal["C", "F"], memory_order: Literal["C", "F"]
203-
) -> None:
198+
@pytest.mark.parametrize("numpy_order", ["C", "F"])
199+
@pytest.mark.parametrize("zarr_order", ["C", "F"])
200+
def test_v2_non_contiguous(numpy_order: Literal["C", "F"], zarr_order: Literal["C", "F"]) -> None:
201+
"""
202+
Make sure zarr v2 arrays save data using the memory order given to the zarr array,
203+
not the memory order of the original numpy array.
204+
"""
204205
store = MemoryStore()
205206
arr = zarr.create_array(
206207
store,
@@ -212,26 +213,27 @@ def test_v2_non_contiguous(
212213
filters=None,
213214
compressors=None,
214215
overwrite=True,
215-
order=array_order,
216-
config={"order": memory_order},
216+
order=zarr_order,
217217
)
218218

219-
# Non-contiguous write
220-
a = np.arange(arr.shape[0] * arr.shape[1]).reshape(arr.shape, order=data_order)
219+
# Non-contiguous write, using numpy memory order
220+
a = np.arange(arr.shape[0] * arr.shape[1]).reshape(arr.shape, order=numpy_order)
221221
arr[6:9, 3:6] = a[6:9, 3:6] # The slice on the RHS is important
222222
np.testing.assert_array_equal(arr[6:9, 3:6], a[6:9, 3:6])
223223

224224
np.testing.assert_array_equal(
225225
a[6:9, 3:6],
226226
np.frombuffer(
227227
sync(store.get("2.1", default_buffer_prototype())).to_bytes(), dtype="float64"
228-
).reshape((3, 3), order=array_order),
228+
).reshape((3, 3), order=zarr_order),
229229
)
230-
if memory_order == "F":
230+
# After writing and reading from zarr array, order should be same as zarr order
231+
if zarr_order == "F":
231232
assert (arr[6:9, 3:6]).flags.f_contiguous
232233
else:
233234
assert (arr[6:9, 3:6]).flags.c_contiguous
234235

236+
# Contiguous write
235237
store = MemoryStore()
236238
arr = zarr.create_array(
237239
store,
@@ -243,18 +245,17 @@ def test_v2_non_contiguous(
243245
compressors=None,
244246
filters=None,
245247
overwrite=True,
246-
order=array_order,
247-
config={"order": memory_order},
248+
order=zarr_order,
248249
)
249250

250-
# Contiguous write
251-
a = np.arange(9).reshape((3, 3), order=data_order)
252-
if data_order == "F":
253-
assert a.flags.f_contiguous
254-
else:
255-
assert a.flags.c_contiguous
251+
a = np.arange(9).reshape((3, 3), order=numpy_order)
256252
arr[6:9, 3:6] = a
257253
np.testing.assert_array_equal(arr[6:9, 3:6], a)
254+
# After writing and reading from zarr array, order should be same as zarr order
255+
if zarr_order == "F":
256+
assert (arr[6:9, 3:6]).flags.f_contiguous
257+
else:
258+
assert (arr[6:9, 3:6]).flags.c_contiguous
258259

259260

260261
def test_default_compressor_deprecation_warning():

0 commit comments

Comments
 (0)