Skip to content

Commit dcd3ac9

Browse files
dcheriandcherian
committed
Refactor backend indexing tests
1 parent f58262a commit dcd3ac9

File tree

1 file changed

+138
-136
lines changed

1 file changed

+138
-136
lines changed

‎xarray/tests/test_backends.py‎

Lines changed: 138 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -351,7 +351,144 @@ def test_dtype_coercion_error(self) -> None:
351351
ds.to_netcdf(path, format=format)
352352

353353

354-
class DatasetIOBase:
354+
class BackendIndexingTestsMixin:
355+
def test_orthogonal_indexing(self) -> None:
356+
in_memory = create_test_data()
357+
with self.roundtrip(in_memory) as on_disk:
358+
indexers = {"dim1": [1, 2, 0], "dim2": [3, 2, 0, 3], "dim3": np.arange(5)}
359+
expected = in_memory.isel(indexers)
360+
actual = on_disk.isel(**indexers)
361+
# make sure the array is not yet loaded into memory
362+
assert not actual["var1"].variable._in_memory
363+
assert_identical(expected, actual)
364+
# do it twice, to make sure we're switched from orthogonal -> numpy
365+
# when we cached the values
366+
actual = on_disk.isel(**indexers)
367+
assert_identical(expected, actual)
368+
369+
def test_vectorized_indexing(self) -> None:
370+
in_memory = create_test_data()
371+
with self.roundtrip(in_memory) as on_disk:
372+
indexers = {
373+
"dim1": DataArray([0, 2, 0], dims="a"),
374+
"dim2": DataArray([0, 2, 3], dims="a"),
375+
}
376+
expected = in_memory.isel(indexers)
377+
actual = on_disk.isel(**indexers)
378+
# make sure the array is not yet loaded into memory
379+
assert not actual["var1"].variable._in_memory
380+
assert_identical(expected, actual.load())
381+
# do it twice, to make sure we're switched from
382+
# vectorized -> numpy when we cached the values
383+
actual = on_disk.isel(**indexers)
384+
assert_identical(expected, actual)
385+
386+
def multiple_indexing(indexers):
387+
# make sure a sequence of lazy indexings certainly works.
388+
with self.roundtrip(in_memory) as on_disk:
389+
actual = on_disk["var3"]
390+
expected = in_memory["var3"]
391+
for ind in indexers:
392+
actual = actual.isel(ind)
393+
expected = expected.isel(ind)
394+
# make sure the array is not yet loaded into memory
395+
assert not actual.variable._in_memory
396+
assert_identical(expected, actual.load())
397+
398+
# two-staged vectorized-indexing
399+
indexers2 = [
400+
{
401+
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
402+
"dim3": DataArray([[0, 4], [1, 3], [2, 2]], dims=["a", "b"]),
403+
},
404+
{"a": DataArray([0, 1], dims=["c"]), "b": DataArray([0, 1], dims=["c"])},
405+
]
406+
multiple_indexing(indexers2)
407+
408+
# vectorized-slice mixed
409+
indexers3 = [
410+
{
411+
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
412+
"dim3": slice(None, 10),
413+
}
414+
]
415+
multiple_indexing(indexers3)
416+
417+
# vectorized-integer mixed
418+
indexers4 = [
419+
{"dim3": 0},
420+
{"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"])},
421+
{"a": slice(None, None, 2)},
422+
]
423+
multiple_indexing(indexers4)
424+
425+
# vectorized-integer mixed
426+
indexers5 = [
427+
{"dim3": 0},
428+
{"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"])},
429+
{"a": 1, "b": 0},
430+
]
431+
multiple_indexing(indexers5)
432+
433+
def test_vectorized_indexing_negative_step(self) -> None:
434+
# use dask explicitly when present
435+
open_kwargs: dict[str, Any] | None
436+
if has_dask:
437+
open_kwargs = {"chunks": {}}
438+
else:
439+
open_kwargs = None
440+
in_memory = create_test_data()
441+
442+
def multiple_indexing(indexers):
443+
# make sure a sequence of lazy indexings certainly works.
444+
with self.roundtrip(in_memory, open_kwargs=open_kwargs) as on_disk:
445+
actual = on_disk["var3"]
446+
expected = in_memory["var3"]
447+
for ind in indexers:
448+
actual = actual.isel(ind)
449+
expected = expected.isel(ind)
450+
# make sure the array is not yet loaded into memory
451+
assert not actual.variable._in_memory
452+
assert_identical(expected, actual.load())
453+
454+
# with negative step slice.
455+
indexers = [
456+
{
457+
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
458+
"dim3": slice(-1, 1, -1),
459+
}
460+
]
461+
multiple_indexing(indexers)
462+
463+
# with negative step slice.
464+
indexers = [
465+
{
466+
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
467+
"dim3": slice(-1, 1, -2),
468+
}
469+
]
470+
multiple_indexing(indexers)
471+
472+
def test_outer_indexing_reversed(self) -> None:
473+
# regression test for GH6560
474+
ds = xr.Dataset(
475+
{"z": (("t", "p", "y", "x"), np.ones((1, 1, 31, 40)))},
476+
)
477+
478+
with self.roundtrip(ds) as on_disk:
479+
subset = on_disk.isel(t=[0], p=0).z[:, ::10, ::10][:, ::-1, :]
480+
assert subset.sizes == subset.load().sizes
481+
482+
def test_isel_dataarray(self) -> None:
483+
# Make sure isel works lazily. GH:issue:1688
484+
in_memory = create_test_data()
485+
with self.roundtrip(in_memory) as on_disk:
486+
expected = in_memory.isel(dim2=in_memory["dim2"] < 3)
487+
actual = on_disk.isel(dim2=on_disk["dim2"] < 3)
488+
assert_identical(expected, actual)
489+
490+
491+
class DatasetIOBase(BackendIndexingTestsMixin):
355492
engine: T_NetcdfEngine | None = None
356493
file_format: T_NetcdfTypes | None = None
357494

@@ -695,141 +832,6 @@ def test_roundtrip_boolean_dtype(self) -> None:
695832
assert_identical(original, actual2)
696833
assert actual2["x"].dtype == "bool"
697834

698-
def test_orthogonal_indexing(self) -> None:
699-
in_memory = create_test_data()
700-
with self.roundtrip(in_memory) as on_disk:
701-
indexers = {"dim1": [1, 2, 0], "dim2": [3, 2, 0, 3], "dim3": np.arange(5)}
702-
expected = in_memory.isel(indexers)
703-
actual = on_disk.isel(**indexers)
704-
# make sure the array is not yet loaded into memory
705-
assert not actual["var1"].variable._in_memory
706-
assert_identical(expected, actual)
707-
# do it twice, to make sure we're switched from orthogonal -> numpy
708-
# when we cached the values
709-
actual = on_disk.isel(**indexers)
710-
assert_identical(expected, actual)
711-
712-
def test_vectorized_indexing(self) -> None:
713-
in_memory = create_test_data()
714-
with self.roundtrip(in_memory) as on_disk:
715-
indexers = {
716-
"dim1": DataArray([0, 2, 0], dims="a"),
717-
"dim2": DataArray([0, 2, 3], dims="a"),
718-
}
719-
expected = in_memory.isel(indexers)
720-
actual = on_disk.isel(**indexers)
721-
# make sure the array is not yet loaded into memory
722-
assert not actual["var1"].variable._in_memory
723-
assert_identical(expected, actual.load())
724-
# do it twice, to make sure we're switched from
725-
# vectorized -> numpy when we cached the values
726-
actual = on_disk.isel(**indexers)
727-
assert_identical(expected, actual)
728-
729-
def multiple_indexing(indexers):
730-
# make sure a sequence of lazy indexings certainly works.
731-
with self.roundtrip(in_memory) as on_disk:
732-
actual = on_disk["var3"]
733-
expected = in_memory["var3"]
734-
for ind in indexers:
735-
actual = actual.isel(ind)
736-
expected = expected.isel(ind)
737-
# make sure the array is not yet loaded into memory
738-
assert not actual.variable._in_memory
739-
assert_identical(expected, actual.load())
740-
741-
# two-staged vectorized-indexing
742-
indexers2 = [
743-
{
744-
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
745-
"dim3": DataArray([[0, 4], [1, 3], [2, 2]], dims=["a", "b"]),
746-
},
747-
{"a": DataArray([0, 1], dims=["c"]), "b": DataArray([0, 1], dims=["c"])},
748-
]
749-
multiple_indexing(indexers2)
750-
751-
# vectorized-slice mixed
752-
indexers3 = [
753-
{
754-
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
755-
"dim3": slice(None, 10),
756-
}
757-
]
758-
multiple_indexing(indexers3)
759-
760-
# vectorized-integer mixed
761-
indexers4 = [
762-
{"dim3": 0},
763-
{"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"])},
764-
{"a": slice(None, None, 2)},
765-
]
766-
multiple_indexing(indexers4)
767-
768-
# vectorized-integer mixed
769-
indexers5 = [
770-
{"dim3": 0},
771-
{"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"])},
772-
{"a": 1, "b": 0},
773-
]
774-
multiple_indexing(indexers5)
775-
776-
def test_vectorized_indexing_negative_step(self) -> None:
777-
# use dask explicitly when present
778-
open_kwargs: dict[str, Any] | None
779-
if has_dask:
780-
open_kwargs = {"chunks": {}}
781-
else:
782-
open_kwargs = None
783-
in_memory = create_test_data()
784-
785-
def multiple_indexing(indexers):
786-
# make sure a sequence of lazy indexings certainly works.
787-
with self.roundtrip(in_memory, open_kwargs=open_kwargs) as on_disk:
788-
actual = on_disk["var3"]
789-
expected = in_memory["var3"]
790-
for ind in indexers:
791-
actual = actual.isel(ind)
792-
expected = expected.isel(ind)
793-
# make sure the array is not yet loaded into memory
794-
assert not actual.variable._in_memory
795-
assert_identical(expected, actual.load())
796-
797-
# with negative step slice.
798-
indexers = [
799-
{
800-
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
801-
"dim3": slice(-1, 1, -1),
802-
}
803-
]
804-
multiple_indexing(indexers)
805-
806-
# with negative step slice.
807-
indexers = [
808-
{
809-
"dim1": DataArray([[0, 7], [2, 6], [3, 5]], dims=["a", "b"]),
810-
"dim3": slice(-1, 1, -2),
811-
}
812-
]
813-
multiple_indexing(indexers)
814-
815-
def test_outer_indexing_reversed(self) -> None:
816-
# regression test for GH6560
817-
ds = xr.Dataset(
818-
{"z": (("t", "p", "y", "x"), np.ones((1, 1, 31, 40)))},
819-
)
820-
821-
with self.roundtrip(ds) as on_disk:
822-
subset = on_disk.isel(t=[0], p=0).z[:, ::10, ::10][:, ::-1, :]
823-
assert subset.sizes == subset.load().sizes
824-
825-
def test_isel_dataarray(self) -> None:
826-
# Make sure isel works lazily. GH:issue:1688
827-
in_memory = create_test_data()
828-
with self.roundtrip(in_memory) as on_disk:
829-
expected = in_memory.isel(dim2=in_memory["dim2"] < 3)
830-
actual = on_disk.isel(dim2=on_disk["dim2"] < 3)
831-
assert_identical(expected, actual)
832-
833835
def validate_array_type(self, ds):
834836
# Make sure that only NumpyIndexingAdapter stores a bare np.ndarray.
835837
def find_and_validate_array(obj):

0 commit comments

Comments
 (0)