icechunk/icechunk-python/tests/test_xarray.py at ea1ad1833fa35608e8e000263cea62a57c2639b0 · earth-mover/icechunk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
from collections.abc import Generator

import pytest

pytest.importorskip("xarray")

import contextlib
import string
import tempfile

import numpy as np
import pandas as pd

import xarray as xr
from icechunk import Repository, in_memory_storage, local_filesystem_storage
from icechunk.xarray import to_icechunk
from xarray.testing import assert_identical


def create_test_data(
    seed: int | None = None,
    add_attrs: bool = True,
    dim_sizes: tuple[int, int, int] = (8, 9, 10),
) -> xr.Dataset:
    rs = np.random.RandomState(seed)
    _vars = {
        "var1": ["dim1", "dim2"],
        "var2": ["dim1", "dim2"],
        "var3": ["dim3", "dim1"],
    }
    _dims = {"dim1": dim_sizes[0], "dim2": dim_sizes[1], "dim3": dim_sizes[2]}

    obj = xr.Dataset()
    obj["dim2"] = ("dim2", 0.5 * np.arange(_dims["dim2"]))
    if _dims["dim3"] > 26:
        raise RuntimeError(
            f"Not enough letters for filling this dimension size ({_dims['dim3']})"
        )
    obj["dim3"] = ("dim3", list(string.ascii_lowercase[0 : _dims["dim3"]]))
    obj["time"] = ("time", pd.date_range("2000-01-01", periods=20))
    for v, dims in sorted(_vars.items()):
        data = rs.normal(size=tuple(_dims[d] for d in dims))
        obj[v] = (dims, data)
        if add_attrs:
            obj[v].attrs = {"foo": "variable"}
    numbers_values = rs.randint(0, 3, _dims["dim3"], dtype="int64")
    obj.coords["numbers"] = ("dim3", numbers_values)
    obj.encoding = {"foo": "bar"}
    return obj


def create_test_datatree() -> xr.DataTree:
    return xr.DataTree.from_dict(
        {
            "/": xr.Dataset(
                data_vars={
                    "bar": ("x", ["hello", "world"]),
                },
                coords={
                    "x": (
                        "x",
                        [1, 2],
                    ),  # inherited dimension coordinate that can't be overridden
                    "w": (
                        "x",
                        [0.1, 0.2],
                    ),  # inherited non-dimension coordinate to override
                },
            ),
            "/a": xr.Dataset(
                data_vars={
                    "foo": ("x", ["alpha", "beta"]),
                },
                coords={
                    "w": ("x", [10, 20]),  # override inherited non-dimension coordinate
                    "z": ("z", ["alpha", "beta"]),  # non-inherited dimension coordinate
                },
            ),
            "/b": xr.Dataset(
                data_vars={
                    "foo": ("x", ["gamma", "delta"]),
                },
                coords={
                    "z": (
                        "z",
                        ["alpha", "beta", "gamma"],
                    ),  # override inherited non-dimension coordinate with different length (i.e. multi-resolution)
                },
            ),
        }
    )


@contextlib.contextmanager
def roundtrip(
    data: xr.Dataset, *, commit: bool = False
) -> Generator[xr.Dataset, None, None]:
    with tempfile.TemporaryDirectory() as tmpdir:
        repo = Repository.create(local_filesystem_storage(tmpdir))
        session = repo.writable_session("main")
        to_icechunk(data, session=session, mode="w")
        session.commit("write")
        with xr.open_zarr(session.store, consolidated=False) as ds:
            yield ds


def test_xarray_dataset_to_icechunk() -> None:
    ds = create_test_data()
    with roundtrip(ds) as actual:
        assert_identical(actual, ds)


@contextlib.contextmanager
def roundtrip_datatree(
    dt: xr.DataTree, *, commit: bool = False
) -> Generator[xr.DataTree, None, None]:
    with tempfile.TemporaryDirectory() as tmpdir:
        repo = Repository.create(local_filesystem_storage(tmpdir))
        session = repo.writable_session("main")
        to_icechunk(dt, session=session, mode="w")
        session.commit("write")
        with xr.open_datatree(session.store, consolidated=False, engine="zarr") as dt:
            yield dt


def test_xarray_datatree_to_icechunk() -> None:
    dt = create_test_datatree()
    with roundtrip_datatree(dt) as actual:
        assert_identical(actual, dt)


def test_repeated_to_icechunk_serial() -> None:
    ds = create_test_data()
    repo = Repository.create(in_memory_storage())
    session = repo.writable_session("main")
    to_icechunk(ds, session)
    to_icechunk(ds, session, mode="w")