Skip to content

Commit

Permalink
Support loadables without tests
Browse files Browse the repository at this point in the history
  • Loading branch information
mpiannucci committed Oct 21, 2024
1 parent 4d5c46a commit 5cb3d21
Showing 2 changed files with 13 additions and 8 deletions.
4 changes: 2 additions & 2 deletions virtualizarr/tests/test_writers/test_icechunk.py
Original file line number Diff line number Diff line change
@@ -20,15 +20,15 @@
from icechunk import IcechunkStore # type: ignore[import-not-found]


@pytest.fixture
@pytest.fixture(scope="function")
def icechunk_filestore(tmpdir) -> "IcechunkStore":
from icechunk import IcechunkStore, StorageConfig

storage = StorageConfig.filesystem(str(tmpdir))

# TODO if icechunk exposed a synchronous version of .open then we wouldn't need to use asyncio.run here
# TODO is this the correct mode to use?
store = asyncio.run(IcechunkStore.open(storage=storage, mode="a"))
store = IcechunkStore.create(storage=storage, mode="w")

# TODO instead yield store then store.close() ??
return store
17 changes: 11 additions & 6 deletions virtualizarr/writers/icechunk.py
Original file line number Diff line number Diff line change
@@ -64,14 +64,22 @@ def write_variables_to_icechunk_group(
store,
group,
):
for name, var in variables.items():
write_variable_to_icechunk(
print(variables)
virtual_variables = {name: var for name, var in variables.items() if isinstance(var.data, ManifestArray)}

for name, var in virtual_variables.items():
write_virtual_variable_to_icechunk(
store=store,
group=group,
name=name,
var=var,
)

loadable_variables = {name: var for name, var in variables.items() if name not in virtual_variables}

ds = Dataset(loadable_variables)
ds.to_zarr(store, zarr_format=3, consolidated=False, mode="r+")


def write_variable_to_icechunk(
store: "IcechunkStore",
@@ -80,7 +88,6 @@ def write_variable_to_icechunk(
var: Variable,
) -> None:
"""Write a single (possibly virtual) variable into an icechunk store"""

if isinstance(var.data, ManifestArray):
write_virtual_variable_to_icechunk(
store=store,
@@ -89,9 +96,7 @@ def write_variable_to_icechunk(
var=var,
)
else:
# TODO is writing loadable_variables just normal xarray ds.to_zarr?
# raise NotImplementedError()
print("skipping non-virtual variable", name)
raise ValueError("Cannot write non-virtual variables as virtual variables to Icechunk stores")


def write_virtual_variable_to_icechunk(

0 comments on commit 5cb3d21

Please sign in to comment.