How to use repo.transaction context manager with icechunk.xarray.to_icechunk

[pjireland]

I'm trying to use the repo.transaction context manager together with icechunk.xarray.to_icechunk, as shown in the example below, where I'll be appending data along the time axis.

Is this possible (or even recommended)? Per the API docs, icechunk.xarray.to_icechunk expects the second argument to be of type icechunk.Session.

My motivation here is to append along the time axis and to use a context manager to handle the transactions. I'm open to other solutions if a different approach is recommended.

storage = icechunk.s3_storage(
    bucket="my-bucket",
    prefix="my-prefix",
    from_env=True,
)
if icechunk.Repository.exists(storage):
    repo = icechunk.Repository.open(storage)
else:
    repo = icechunk.Repository.create(storage)
with repo.transaction(
    "main",
    message="My commit message",
    rebase_with=icechunk.ConflictDetector(),
) as store:
    icechunk.xarray.to_icechunk(
        mydf.to_xarray(),
        store, # Expects icechunk.Session
        append_dim="time"
    )

[rabernat]

Welcome @pjireland! 👋

The context manager and xarray.to_icechunk are currently mutually exclusive. However, to_icechunk is only required with distributed writes. You have two options

Option 1: Use Context manager without to_xarray

with repo.transaction("main", message=message) as store:
    # ds needs to be in memory or using Dask threaded scheduler
    ds.to_zarr(store, zarr_format=3)

Option 2: Use to_xarray with Session API

session = repo.writable_session("main")
icechunk.xarray.to_icechunk(ds, session.store, zarr_format=3)
session.commit("message")

Keep in mind that the context manager is just syntactic sugar over the Session API. The implementation is extremely simple:

icechunk/icechunk-python/python/icechunk/repository.py

Lines 620 to 627 in 61c46a4

	session = self.writable_session(branch)
	yield session.store
	session.commit(
	message=message,
	metadata=metadata,
	rebase_with=rebase_with,
	rebase_tries=rebase_tries,
	)

[pjireland]

Thanks for the quick reply, @rabernat! That makes sense. I'll use option 2 for now. It looks like it works well in my test.