1"""
  2Icechunk Virtual Store Writer
  3=============================
  4
  5Provides :class:`IcechunkWriter`, which writes grib2io reference manifests
  6(Kerchunk v1 format) into an `Icechunk <https://icechunk.io/>`_ virtual
  7store, creating versioned Zarr v3 datasets backed by the original GRIB2 file
  8bytes.
  9
 10The writer translates each entry in the manifest's ``"refs"`` dict into
 11either a virtual chunk reference (for data chunks pointing at byte ranges
 12in GRIB2 files) or native Zarr metadata (for ``.zarray``, ``.zattrs``,
 13``.zgroup`` entries and inline base64-encoded coordinate arrays).
 14
 15Icechunk uses Zarr v3 format internally, so the writer converts Kerchunk v1
 16(Zarr v2) metadata and chunk keys to Zarr v3 equivalents:
 17
 18- ``.zgroup`` → ``zarr.json`` with ``node_type: "group"``
 19- ``.zarray`` + ``.zattrs`` → ``<var>/zarr.json`` with ``node_type: "array"``
 20- Chunk keys ``VAR/0.0.0`` → ``VAR/c/0/0/0``
 21
 22Example
 23-------
 24>>> from grib2io.kerchunk import ReferenceGenerator
 25>>> from grib2io.icechunk import IcechunkWriter
 26>>> gen = ReferenceGenerator("gfs.grib2")
 27>>> manifest = gen.generate()
 28>>> writer = IcechunkWriter("/tmp/gfs_icechunk")
 29>>> writer.write(manifest)
 30>>> snapshot_id = writer.commit("Initial ingest of GFS data")
 31"""
 32
 33from __future__ import annotations
 34
 35import base64
 36import json
 37import logging
 38from typing import Any, Dict, Optional, Set
 39
 40_logger = logging.getLogger(__name__)
 41
 42
 43# ---------------------------------------------------------------------------
 44# Lazy import guard
 45# ---------------------------------------------------------------------------
 46
 47
 48def _ensure_icechunk():
 49    """Raise ``ImportError`` if *icechunk* is not available."""
 50    try:
 51        import icechunk  # noqa: F401
 52    except ImportError:
 53        raise ImportError("icechunk is required for virtual store support. Install with: pip install grib2io[icechunk]")
 54
 55
 56# ---------------------------------------------------------------------------
 57# Helpers
 58# ---------------------------------------------------------------------------
 59
 60
 61def _is_metadata_key(key: str) -> bool:
 62    """Return ``True`` if *key* is a Zarr metadata key."""
 63    basename = key.rsplit("/", 1)[-1] if "/" in key else key
 64    return basename.startswith(".z")
 65
 66
 67def _is_inline_data(value) -> bool:
 68    """Return ``True`` if *value* is inline data (a JSON string or
 69    base64-encoded bytes), as opposed to a ``[uri, offset, length]``
 70    reference list."""
 71    return isinstance(value, str)
 72
 73
 74def _is_virtual_ref(value) -> bool:
 75    """Return ``True`` if *value* is a virtual chunk reference
 76    ``[uri, offset, length]``."""
 77    return isinstance(value, list) and len(value) == 3
 78
 79
 80def _decode_inline_value(value: str) -> bytes:
 81    """Decode an inline manifest value to raw bytes.
 82
 83    Inline values are either:
 84    - ``"base64:<data>"`` — base64-encoded binary data (coordinate arrays)
 85    - A plain JSON string (metadata like ``.zarray``, ``.zattrs``)
 86
 87    Parameters
 88    ----------
 89    value : str
 90        The inline value from the manifest.
 91
 92    Returns
 93    -------
 94    bytes
 95        The decoded bytes.
 96    """
 97    if value.startswith("base64:"):
 98        return base64.b64decode(value[7:])
 99    # Plain JSON string — encode to UTF-8 bytes for storage
100    return value.encode("utf-8")
101
102
103def _collect_virtual_chunk_prefixes(refs: Dict[str, Any]) -> Set[str]:
104    """Scan all virtual refs and collect the unique URI prefixes needed
105    for Icechunk virtual chunk containers.
106
107    For ``file://`` URIs, the prefix is the directory portion.
108    For ``s3://``, ``gcs://``, ``https://`` URIs, the prefix is the
109    bucket/host + path up to the last ``/``.
110
111    Parameters
112    ----------
113    refs : dict
114        The ``"refs"`` dict from a Kerchunk v1 manifest.
115
116    Returns
117    -------
118    set of str
119        Unique URI prefixes (each ending with ``/``).
120    """
121    prefixes: Set[str] = set()
122    for key, value in refs.items():
123        if _is_virtual_ref(value):
124            uri = value[0]
125            # Extract the directory prefix from the URI
126            last_slash = uri.rfind("/")
127            if last_slash > 0:
128                prefix = uri[: last_slash + 1]
129                prefixes.add(prefix)
130    return prefixes
131
132
133def _chunk_key_v2_to_v3(key: str) -> str:
134    """Convert a Kerchunk v1 (Zarr v2) chunk key to Zarr v3 format.
135
136    Zarr v2 chunk keys use dot-separated indices under the variable name:
137    ``"TMP/0.0.0.0"`` or ``"level/0"``.
138
139    Zarr v3 chunk keys use ``c/`` prefix with slash-separated indices:
140    ``"TMP/c/0/0/0/0"`` or ``"level/c/0"``.
141
142    Parameters
143    ----------
144    key : str
145        A Zarr v2 chunk key like ``"TMP/0.0.0"`` or ``"level/0"``.
146
147    Returns
148    -------
149    str
150        The Zarr v3 equivalent like ``"TMP/c/0/0/0"`` or ``"level/c/0"``.
151    """
152    if "/" not in key:
153        # Root-level key, no conversion needed
154        return key
155
156    parts = key.split("/")
157    var_name = parts[0]
158    index_str = parts[1]
159
160    # Split dot-separated indices and rejoin with slashes under c/
161    indices = index_str.split(".")
162    return var_name + "/c/" + "/".join(indices)
163
164
165def _make_zarr_v3_group_metadata(v2_zgroup_str: str) -> str:
166    """Convert a Zarr v2 ``.zgroup`` JSON string to Zarr v3 ``zarr.json``
167    group metadata.
168
169    Parameters
170    ----------
171    v2_zgroup_str : str
172        JSON string like ``'{"zarr_format": 2}'``.
173
174    Returns
175    -------
176    str
177        Zarr v3 group metadata JSON string.
178    """
179    return json.dumps(
180        {
181            "zarr_format": 3,
182            "node_type": "group",
183            "attributes": {},
184        }
185    )
186
187
188def _make_zarr_v3_array_metadata(
189    v2_zarray_str: str,
190    v2_zattrs_str: Optional[str] = None,
191) -> str:
192    """Convert Zarr v2 ``.zarray`` and ``.zattrs`` JSON strings to a
193    single Zarr v3 ``zarr.json`` array metadata string.
194
195    Parameters
196    ----------
197    v2_zarray_str : str
198        JSON string with Zarr v2 array metadata (shape, chunks, dtype, etc.).
199    v2_zattrs_str : str, optional
200        JSON string with Zarr v2 attributes (_ARRAY_DIMENSIONS, etc.).
201
202    Returns
203    -------
204    str
205        Zarr v3 array metadata JSON string.
206    """
207    v2_zarray = json.loads(v2_zarray_str)
208    v2_zattrs = json.loads(v2_zattrs_str) if v2_zattrs_str else {}
209
210    shape = v2_zarray.get("shape", [])
211    chunks = v2_zarray.get("chunks", shape)
212    dtype = v2_zarray.get("dtype", "<f4")
213    fill_value = v2_zarray.get("fill_value", None)
214
215    # Convert numpy dtype string to Zarr v3 data_type
216    dtype_map = {
217        "<f4": "float32",
218        "<f8": "float64",
219        "<i4": "int32",
220        "<i8": "int64",
221        "<i2": "int16",
222        "<u1": "uint8",
223        "<u2": "uint16",
224        "<u4": "uint32",
225        ">f4": "float32",
226        ">f8": "float64",
227        ">i4": "int32",
228        ">i8": "int64",
229    }
230    data_type = dtype_map.get(dtype, "float32")
231
232    # Determine endianness for the bytes codec
233    endian = "big" if dtype.startswith(">") else "little"
234
235    # Extract dimension names from attributes
236    dim_names = v2_zattrs.get("_ARRAY_DIMENSIONS", [])
237
238    # Build codecs — for Zarr v3 / Icechunk, we use a simple bytes codec.
239    # The grib2io compressor is a Zarr v2 concept used by the Kerchunk
240    # reference filesystem; Icechunk accesses data via virtual refs so
241    # the codec pipeline just needs to handle raw bytes.
242    codecs = [{"name": "bytes", "configuration": {"endian": endian}}]
243
244    # Preserve the grib2io compressor config in attributes so it can be
245    # recovered if needed, but don't put it in the codecs list where
246    # Zarr v3 would try to instantiate it.
247    compressor = v2_zarray.get("compressor", None)
248    if compressor and compressor.get("id") == "grib2io":
249        v2_zattrs["_grib2io_compressor"] = compressor
250
251    # Handle fill_value — "NaN" string is valid in Zarr v3
252    if fill_value == "NaN" or fill_value is None:
253        fill_value_v3 = "NaN"
254    else:
255        fill_value_v3 = fill_value
256
257    v3_meta = {
258        "zarr_format": 3,
259        "node_type": "array",
260        "shape": shape,
261        "data_type": data_type,
262        "chunk_grid": {
263            "name": "regular",
264            "configuration": {"chunk_shape": chunks},
265        },
266        "chunk_key_encoding": {
267            "name": "default",
268            "configuration": {"separator": "/"},
269        },
270        "fill_value": fill_value_v3,
271        "codecs": codecs,
272        "dimension_names": dim_names if dim_names else None,
273        "attributes": v2_zattrs,
274    }
275
276    return json.dumps(v3_meta)
277
278
279def _store_set_sync(store, key: str, data: bytes) -> None:
280    """Synchronously write bytes to an Icechunk store.
281
282    Icechunk v2's ``store.set()`` is an async coroutine. This helper
283    wraps it using the store's ``_sync`` method and converts raw bytes
284    to a Zarr ``Buffer`` object.
285
286    Parameters
287    ----------
288    store : IcechunkStore
289        The Icechunk store.
290    key : str
291        The Zarr key to write.
292    data : bytes
293        The raw bytes to write.
294    """
295    from zarr.core.buffer import default_buffer_prototype
296
297    buf = default_buffer_prototype().buffer.from_bytes(data)
298    store._sync(store.set(key, buf))
299
300
301def _store_get_sync(store, key: str):
302    """Synchronously read bytes from an Icechunk store.
303
304    Parameters
305    ----------
306    store : IcechunkStore
307        The Icechunk store.
308    key : str
309        The Zarr key to read.
310
311    Returns
312    -------
313    bytes or None
314        The raw bytes, or ``None`` if the key does not exist.
315    """
316    from zarr.core.buffer import default_buffer_prototype
317
318    buf = store._sync(store.get(key, default_buffer_prototype()))
319    if buf is not None:
320        return buf.to_bytes()
321    return None
322
323
324# ---------------------------------------------------------------------------
325# Public API
326# ---------------------------------------------------------------------------
327
328
329class IcechunkWriter:
330    """Write grib2io reference manifests into an Icechunk virtual store.
331
332    Parameters
333    ----------
334    store_path : str
335        Path or URI for the Icechunk store.  For local filesystem stores
336        this is a directory path.  For cloud stores, provide the
337        appropriate URI (e.g. ``s3://bucket/prefix``).
338    storage_config : optional
339        An Icechunk ``Storage`` object.  If ``None`` (the default), a
340        local filesystem storage is created at *store_path*.
341    """
342
343    def __init__(
344        self,
345        store_path: str,
346        storage_config: Optional[Any] = None,
347    ):
348        _ensure_icechunk()
349
350        self._store_path = store_path
351        self._storage_config = storage_config
352        self._repo = None
353        self._session = None
354
355    # ------------------------------------------------------------------
356    # Internal helpers
357    # ------------------------------------------------------------------
358
359    def _get_storage(self):
360        """Return an Icechunk ``Storage`` object."""
361        if self._storage_config is not None:
362            return self._storage_config
363
364        import icechunk
365
366        return icechunk.local_filesystem_storage(path=self._store_path)
367
368    def _create_or_open_repo(self, mode: str, virtual_prefixes: Set[str]):
369        """Create or open an Icechunk repository.
370
371        Parameters
372        ----------
373        mode : str
374            ``'w'`` to create a new repo (or overwrite), ``'a'`` to open
375            an existing repo for appending.
376        virtual_prefixes : set of str
377            URI prefixes that need to be registered as virtual chunk
378            containers.
379        """
380        import icechunk
381
382        storage = self._get_storage()
383
384        # Build repository config with virtual chunk containers
385        config = icechunk.config.RepositoryConfig.default()
386        for prefix in virtual_prefixes:
387            container_store = self._make_container_store(prefix)
388            config.set_virtual_chunk_container(icechunk.virtual.VirtualChunkContainer(prefix, container_store))
389
390        # Build authorize_virtual_chunk_access mapping
391        # Use None for credentials (will use environment or anonymous)
392        authorize = {prefix: None for prefix in virtual_prefixes}
393
394        if mode == "w":
395            self._repo = icechunk.Repository.open_or_create(
396                storage,
397                config=config,
398                authorize_virtual_chunk_access=authorize if authorize else None,
399            )
400        else:
401            # Append mode — open existing
402            self._repo = icechunk.Repository.open(
403                storage,
404                config=config,
405                authorize_virtual_chunk_access=authorize if authorize else None,
406            )
407
408        self._session = self._repo.writable_session("main")
409
410    def _make_container_store(self, prefix: str):
411        """Create an appropriate Icechunk storage backend for a virtual
412        chunk container based on the URI prefix scheme.
413
414        Parameters
415        ----------
416        prefix : str
417            URI prefix like ``file:///path/to/dir/`` or
418            ``s3://bucket/prefix/``.
419
420        Returns
421        -------
422        An Icechunk storage store object.
423        """
424        import icechunk
425
426        if prefix.startswith("file://"):
427            # Extract local path from file:// URI
428            local_path = prefix[7:]  # Remove "file://"
429            # Remove trailing slash for the store path
430            if local_path.endswith("/"):
431                local_path = local_path[:-1]
432            return icechunk.storage.local_filesystem_store(local_path)
433        elif prefix.startswith("s3://"):
434            return icechunk.storage.s3_store(region="us-east-1")
435        elif prefix.startswith("gcs://"):
436            return icechunk.storage.gcs_store(opts={})
437        elif prefix.startswith("http://") or prefix.startswith("https://"):
438            return icechunk.storage.http_store(opts={})
439        else:
440            # Default to local filesystem
441            return icechunk.storage.local_filesystem_store(prefix)
442
443    # ------------------------------------------------------------------
444    # Public methods
445    # ------------------------------------------------------------------
446
447    def write(
448        self,
449        manifest: dict,
450        mode: str = "w",
451        append_dim: Optional[str] = None,
452    ) -> None:
453        """Write a reference manifest into the Icechunk store.
454
455        Parameters
456        ----------
457        manifest : dict
458            Kerchunk v1 reference manifest dict with ``"version"`` and
459            ``"refs"`` keys.
460        mode : str
461            ``'w'`` for create/overwrite, ``'a'`` for append.
462        append_dim : str, optional
463            Dimension along which to append when ``mode='a'``.  Required
464            when appending multi-file data along a specific dimension
465            (e.g. ``"refDate"`` or ``"leadTime"``).
466        """
467        refs = manifest.get("refs", {})
468
469        # Collect virtual chunk URI prefixes for container registration
470        virtual_prefixes = _collect_virtual_chunk_prefixes(refs)
471
472        # Create or open the repository
473        self._create_or_open_repo(mode, virtual_prefixes)
474
475        store = self._session.store
476
477        if mode == "a" and append_dim is not None:
478            self._write_append(refs, store, append_dim)
479        else:
480            self._write_create(refs, store)
481
482    def _write_create(self, refs: Dict[str, Any], store) -> None:
483        """Write all refs to a fresh store using zarr's Python API.
484
485        Creates the root group and all arrays via ``zarr.open_group`` /
486        ``group.create_array`` so that Zarr v3 metadata is stored through
487        the proper path that icechunk and zarr recognise.  Coordinate
488        (inline) data is written via the zarr Array write API, and virtual
489        chunk refs for data variables are registered with
490        ``store.set_virtual_ref``.
491
492        Parameters
493        ----------
494        refs : dict
495            The ``"refs"`` dict from the manifest.
496        store : IcechunkStore
497            The Icechunk store to write to.
498        """
499        import numpy as np
500        import zarr
501
502        # Collect all metadata by variable
503        zarray_refs: Dict[str, str] = {}  # var_name -> .zarray JSON str
504        zattrs_refs: Dict[str, str] = {}  # var_name -> .zattrs JSON str
505        data_refs: Dict[str, Any] = {}  # chunk key -> [uri, offset, len]
506        inline_data_refs: Dict[str, str] = {}  # chunk key -> inline value
507
508        for key, value in refs.items():
509            if key == ".zgroup":
510                pass  # zarr.open_group handles the root group metadata
511            elif key.endswith("/.zarray"):
512                var_name = key.rsplit("/.zarray", 1)[0]
513                zarray_refs[var_name] = value
514            elif key.endswith("/.zattrs"):
515                var_name = key.rsplit("/.zattrs", 1)[0]
516                zattrs_refs[var_name] = value
517            elif _is_virtual_ref(value):
518                data_refs[key] = value
519            elif _is_inline_data(value) and not _is_metadata_key(key):
520                inline_data_refs[key] = value
521            elif _is_metadata_key(key):
522                pass  # skip other metadata keys
523            else:
524                _logger.warning(
525                    "Skipping unrecognized manifest entry: %s = %r",
526                    key,
527                    value,
528                )
529
530        # Create root group and all arrays via zarr's Python API so that
531        # metadata is stored in the location that zarr/icechunk actually
532        # reads, rather than as raw chunk bytes at the zarr.json key.
533        root = zarr.open_group(store, mode="w", zarr_format=3)
534
535        # Identify which variables have virtual (data) refs — those need
536        # Grib2SerializerCodec as their zarr v3 serializer.
537        data_var_names: set = {key.split("/")[0] for key in data_refs}
538
539        # Import the zarr v3 serializer codec (registered at import time)
540        from grib2io.codecs import Grib2SerializerCodec as _Grib2Ser
541
542        for var_name, zarray_str in zarray_refs.items():
543            v2_zarray = json.loads(zarray_str)
544            zattrs_str = zattrs_refs.get(var_name)
545            v2_zattrs = json.loads(zattrs_str) if zattrs_str else {}
546
547            shape = v2_zarray["shape"]
548            chunks = v2_zarray.get("chunks", shape)
549            dtype_str = v2_zarray.get("dtype", "<f4")
550            fill_value = v2_zarray.get("fill_value", None)
551            dim_names = v2_zattrs.get("_ARRAY_DIMENSIONS", [])
552
553            dtype = np.dtype(dtype_str)
554            if fill_value == "NaN" or fill_value is None:
555                fv: Any = float("nan") if np.issubdtype(dtype, np.floating) else 0
556            else:
557                fv = fill_value
558
559            if var_name in data_var_names:
560                # Data variable — use Grib2SerializerCodec so zarr v3 can
561                # decode the raw GRIB2 section 7 bytes from virtual chunks.
562                compressor = v2_zarray.get("compressor", {}) or {}
563                codec_cfg = {k: v for k, v in compressor.items() if k != "id"}
564                serializer = _Grib2Ser(**codec_cfg)
565                root.create_array(
566                    var_name,
567                    shape=shape,
568                    chunks=chunks,
569                    dtype=dtype,
570                    fill_value=fv,
571                    serializer=serializer,
572                    compressors=[],
573                    filters=[],
574                    dimension_names=dim_names if dim_names else None,
575                    attributes=v2_zattrs if v2_zattrs else None,
576                )
577            else:
578                # Coordinate array — plain bytes serializer (default)
579                root.create_array(
580                    var_name,
581                    shape=shape,
582                    chunks=chunks,
583                    dtype=dtype,
584                    fill_value=fv,
585                    dimension_names=dim_names if dim_names else None,
586                    attributes=v2_zattrs if v2_zattrs else None,
587                )
588
589        # Write coordinate (inline) data via zarr's Array write API so
590        # that encoding goes through the proper codec pipeline.
591        for key, value in inline_data_refs.items():
592            parts = key.split("/")
593            var_name = parts[0]
594            if var_name not in zarray_refs:
595                _logger.warning("Inline data for unknown variable %s, skipping", var_name)
596                continue
597
598            raw_bytes = _decode_inline_value(value)
599            v2_zarray = json.loads(zarray_refs[var_name])
600            dtype = np.dtype(v2_zarray["dtype"])
601            shape = v2_zarray["shape"]
602            coord_values = np.frombuffer(raw_bytes, dtype=dtype).reshape(shape)
603            root[var_name][...] = coord_values
604
605        # Write virtual chunk references for data arrays (unchanged path)
606        for key, value in data_refs.items():
607            uri, offset, length = value
608            v3_key = _chunk_key_v2_to_v3(key)
609            store.set_virtual_ref(
610                v3_key,
611                uri,
612                offset=int(offset),
613                length=int(length),
614                validate_container=True,
615            )
616
617    def _write_append(
618        self,
619        refs: Dict[str, Any],
620        store,
621        append_dim: str,
622    ) -> None:
623        """Append new data to an existing store along *append_dim*.
624
625        Opens the existing zarr group and uses zarr's Python API to resize
626        arrays and write new coordinate data, ensuring metadata is stored
627        correctly.
628
629        Parameters
630        ----------
631        refs : dict
632            The ``"refs"`` dict from the new manifest.
633        store : IcechunkStore
634            The Icechunk store to append to.
635        append_dim : str
636            Dimension name along which to append.
637        """
638        import numpy as np
639        import zarr
640
641        # Separate refs by type
642        zarray_refs: Dict[str, str] = {}
643        zattrs_refs: Dict[str, str] = {}
644        data_refs: Dict[str, Any] = {}
645        inline_refs: Dict[str, str] = {}
646
647        for key, value in refs.items():
648            if key in {".zgroup"}:
649                pass
650            elif key.endswith("/.zarray"):
651                var_name = key.rsplit("/.zarray", 1)[0]
652                zarray_refs[var_name] = value
653            elif key.endswith("/.zattrs"):
654                var_name = key.rsplit("/.zattrs", 1)[0]
655                zattrs_refs[var_name] = value
656            elif _is_virtual_ref(value):
657                data_refs[key] = value
658            elif _is_inline_data(value) and not _is_metadata_key(key):
659                inline_refs[key] = value
660
661        # Open the existing group via zarr API
662        root = zarr.open_group(store, mode="r+", zarr_format=3)
663
664        # Group data refs by variable
665        var_data_refs: Dict[str, Dict[str, list]] = {}
666        for key, value in data_refs.items():
667            var_name = key.split("/")[0]
668            var_data_refs.setdefault(var_name, {})[key] = value
669
670        # Extend each data variable and register new virtual refs
671        for var_name, chunk_refs in var_data_refs.items():
672            new_zarray_str = zarray_refs.get(var_name)
673            if new_zarray_str is None:
674                continue
675
676            new_zarray = json.loads(new_zarray_str)
677            new_shape_per_append = new_zarray["shape"]
678
679            # Determine append axis from dimension labels
680            new_zattrs_str = zattrs_refs.get(var_name)
681            dim_labels: list = []
682            if new_zattrs_str:
683                dim_labels = json.loads(new_zattrs_str).get("_ARRAY_DIMENSIONS", [])
684
685            append_axis = dim_labels.index(append_dim) if append_dim in dim_labels else None
686
687            if var_name in root and append_axis is not None:
688                zarr_arr = root[var_name]
689                existing_shape = list(zarr_arr.shape)
690                offset = existing_shape[append_axis]
691
692                # Resize array along append axis
693                new_size = existing_shape[append_axis] + new_shape_per_append[append_axis]
694                new_shape_full = list(existing_shape)
695                new_shape_full[append_axis] = new_size
696                zarr_arr.resize(new_shape_full)
697            else:
698                offset = 0
699
700            # Write virtual refs with adjusted chunk indices
701            for key, value in chunk_refs.items():
702                uri, chunk_offset, length = value
703                if append_axis is not None and offset > 0:
704                    adjusted_key = self._adjust_chunk_key(key, var_name, append_axis, offset)
705                else:
706                    adjusted_key = key
707                v3_key = _chunk_key_v2_to_v3(adjusted_key)
708                store.set_virtual_ref(
709                    v3_key,
710                    uri,
711                    offset=int(chunk_offset),
712                    length=int(length),
713                    validate_container=True,
714                )
715
716        # Handle coordinate arrays
717        for key, inline_value in inline_refs.items():
718            coord_name = key.split("/")[0] if "/" in key else None
719            if coord_name is None:
720                continue
721
722            raw_bytes = _decode_inline_value(inline_value)
723            new_zarray_str = zarray_refs.get(coord_name)
724            if new_zarray_str is None:
725                continue
726            dtype = np.dtype(json.loads(new_zarray_str)["dtype"])
727            new_values = np.frombuffer(raw_bytes, dtype=dtype)
728
729            if coord_name == append_dim and coord_name in root:
730                # Concatenate with existing coordinate values
731                existing_arr = root[coord_name][...]
732                combined = np.concatenate([existing_arr, new_values])
733                root[coord_name].resize(len(combined))
734                root[coord_name][...] = combined
735            elif coord_name not in root:
736                # Create new coordinate array via zarr API
737                v2_zattrs = json.loads(zattrs_refs[coord_name]) if coord_name in zattrs_refs else {}
738                dim_names = v2_zattrs.get("_ARRAY_DIMENSIONS", [])
739                shape = [len(new_values)]
740                root.create_array(
741                    coord_name,
742                    shape=shape,
743                    chunks=shape,
744                    dtype=dtype,
745                    fill_value=0,
746                    dimension_names=dim_names if dim_names else None,
747                    attributes=v2_zattrs if v2_zattrs else None,
748                )
749                root[coord_name][...] = new_values
750            # else: coordinate already exists and isn't the append dim — leave as-is
751
752    def _adjust_chunk_key(
753        self,
754        key: str,
755        var_name: str,
756        append_axis: int,
757        offset: int,
758    ) -> str:
759        """Adjust a chunk key's index along the append axis.
760
761        Parameters
762        ----------
763        key : str
764            Original chunk key like ``"TMP/0.0.0.0"``.
765        var_name : str
766            Variable name prefix.
767        append_axis : int
768            Axis index along which to offset.
769        offset : int
770            Number of existing chunks along the append axis.
771
772        Returns
773        -------
774        str
775            Adjusted chunk key (still in Zarr v2 format — caller
776            converts to v3).
777        """
778        # Extract the index portion after the variable name
779        prefix = var_name + "/"
780        if not key.startswith(prefix):
781            return key
782
783        index_str = key[len(prefix) :]
784        indices = index_str.split(".")
785        if append_axis < len(indices):
786            indices[append_axis] = str(int(indices[append_axis]) + offset)
787        return prefix + ".".join(indices)
788
789    def _extend_coord(
790        self,
791        store,
792        coord_name: str,
793        data_key: str,
794        new_value: str,
795        zarray_refs: Dict[str, str],
796        zattrs_refs: Dict[str, str],
797    ) -> None:
798        """Extend a coordinate array with new values for append mode.
799
800        Parameters
801        ----------
802        store : IcechunkStore
803            The Icechunk store.
804        coord_name : str
805            Coordinate name (e.g. ``"refDate"``).
806        data_key : str
807            The data chunk key in v2 format (e.g. ``"refDate/0"``).
808        new_value : str
809            The new inline value (base64-encoded).
810        zarray_refs : dict
811            Variable name -> .zarray JSON string from the new manifest.
812        zattrs_refs : dict
813            Variable name -> .zattrs JSON string from the new manifest.
814        """
815        import numpy as np
816
817        new_bytes = _decode_inline_value(new_value)
818        v3_data_key = _chunk_key_v2_to_v3(data_key)
819
820        # Try to read existing coordinate data
821        existing_bytes = None
822        try:
823            existing_bytes = _store_get_sync(store, v3_data_key)
824        except Exception:
825            pass
826
827        if existing_bytes is not None:
828            # Read existing zarr.json to get dtype
829            existing_v3_bytes = None
830            try:
831                existing_v3_bytes = _store_get_sync(store, f"{coord_name}/zarr.json")
832            except Exception:
833                pass
834
835            dtype = np.float64
836            if existing_v3_bytes is not None:
837                existing_v3_meta = json.loads(existing_v3_bytes.decode("utf-8"))
838                dt_str = existing_v3_meta.get("data_type", "float64")
839                dtype_map = {
840                    "float32": np.float32,
841                    "float64": np.float64,
842                    "int32": np.int32,
843                    "int64": np.int64,
844                }
845                dtype = dtype_map.get(dt_str, np.float64)
846
847            # Concatenate existing and new coordinate values
848            existing_arr = np.frombuffer(existing_bytes, dtype=dtype)
849            new_arr = np.frombuffer(new_bytes, dtype=dtype)
850            combined = np.concatenate([existing_arr, new_arr])
851
852            # Write combined data
853            _store_set_sync(store, v3_data_key, combined.tobytes())
854
855            # Update zarr.json with new shape
856            new_zarray_str = zarray_refs.get(coord_name)
857            if new_zarray_str:
858                coord_zarray = json.loads(new_zarray_str)
859            else:
860                # Reconstruct from existing v3 metadata
861                coord_zarray = {
862                    "shape": [len(combined)],
863                    "chunks": [len(combined)],
864                    "dtype": "<f8",
865                    "fill_value": "NaN",
866                    "order": "C",
867                    "zarr_format": 2,
868                    "compressor": None,
869                }
870
871            coord_zarray["shape"] = [len(combined)]
872            coord_zarray["chunks"] = [len(combined)]
873
874            zattrs_str = zattrs_refs.get(coord_name)
875            v3_array = _make_zarr_v3_array_metadata(json.dumps(coord_zarray), zattrs_str)
876            _store_set_sync(
877                store,
878                f"{coord_name}/zarr.json",
879                v3_array.encode("utf-8"),
880            )
881        else:
882            # No existing data — write as-is
883            _store_set_sync(store, v3_data_key, new_bytes)
884            zarray_str = zarray_refs.get(coord_name)
885            zattrs_str = zattrs_refs.get(coord_name)
886            if zarray_str:
887                v3_array = _make_zarr_v3_array_metadata(zarray_str, zattrs_str)
888                _store_set_sync(
889                    store,
890                    f"{coord_name}/zarr.json",
891                    v3_array.encode("utf-8"),
892                )
893
894    def commit(self, message: str = "") -> str:
895        """Commit the current transaction.
896
897        Parameters
898        ----------
899        message : str
900            Commit message describing the changes.
901
902        Returns
903        -------
904        str
905            The snapshot ID of the committed transaction.
906
907        Raises
908        ------
909        RuntimeError
910            If no session is active (i.e. :meth:`write` has not been
911            called).
912        """
913        if self._session is None:
914            raise RuntimeError("No active session. Call write() before commit().")
915        snapshot_id = self._session.commit(message)
916        return str(snapshot_id)