tensogram-xarray 0.14.0__py3-none-any.whl

tensogram_xarray/scanner.py

@@ -0,0 +1,196 @@
+# (C) Copyright 2026- ECMWF and individual contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation nor
+# does it submit to any jurisdiction.
+
+"""File scanner: extract metadata from all messages/objects without decoding payloads.
+
+The scanner opens a ``.tgm`` file via ``tensogram.TensogramFile``, reads each
+message's metadata, and builds an index of per-object metadata dicts that
+downstream merge/split logic can consume.
+
+Per-object metadata is read from ``meta.base[i]`` (with ``_reserved_``
+filtered out) and supplemented by ``desc.params`` as fallback.
+Message-level metadata comes from ``meta.extra``.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+# The ``_reserved_`` key in base entries is populated by the encoder
+# with tensor info (ndim, shape, strides, dtype).  It must be excluded
+# from application-level metadata used for grouping and variable naming.
+RESERVED_KEY = "_reserved_"
+
+
+@dataclass
+class ObjectInfo:
+    """Metadata for a single data object within a message."""
+
+    msg_index: int
+    obj_index: int
+    ndim: int
+    shape: tuple[int, ...]
+    dtype: str
+    descriptor: Any  # tensogram.DataObjectDescriptor
+    per_object_meta: dict[str, Any]  # from meta.base[i] (filtered)
+    common_meta: dict[str, Any]  # from meta.extra
+
+    @property
+    def merged_meta(self) -> dict[str, Any]:
+        """Return common + per-object metadata merged (per-object wins)."""
+        merged: dict[str, Any] = {}
+        merged.update(self.common_meta)
+        merged.update(self.per_object_meta)
+        return merged
+
+
+@dataclass
+class FileIndex:
+    """Index over all messages and objects in a ``.tgm`` file."""
+
+    file_path: str
+    objects: list[ObjectInfo] = field(default_factory=list)
+
+    @property
+    def message_count(self) -> int:
+        if not self.objects:
+            return 0
+        return max(o.msg_index for o in self.objects) + 1
+
+
+def _desc_params(desc: Any) -> dict[str, Any]:
+    """Extract per-object metadata from a descriptor's params."""
+    params = getattr(desc, "params", None)
+    if params and isinstance(params, dict):
+        return dict(params)
+    return {}
+
+
+def _base_entry_from_meta(meta: Any, obj_index: int) -> dict[str, Any]:
+    """Extract per-object metadata from ``meta.base[obj_index]``.
+
+    The ``_reserved_`` key is filtered out since it contains
+    encoder-populated tensor info (ndim, shape, strides, dtype).
+    """
+    base = getattr(meta, "base", None)
+    if base and isinstance(base, list) and obj_index < len(base):
+        entry = base[obj_index]
+        if isinstance(entry, dict):
+            return {k: v for k, v in entry.items() if k != RESERVED_KEY}
+    return {}
+
+
+def _merge_per_object_meta(meta: Any, obj_index: int, desc: Any) -> dict[str, Any]:
+    """Build per-object metadata from base entry + descriptor params.
+
+    Base entry takes priority; descriptor params fill in any missing keys.
+    """
+    result = _base_entry_from_meta(meta, obj_index)
+    for k, v in _desc_params(desc).items():
+        if k not in result:
+            result[k] = v
+    return result
+
+
+def _extra_from_meta(meta: Any) -> dict[str, Any]:
+    """Extract message-level metadata from ``meta.extra``."""
+    result: dict[str, Any] = {}
+    extra = getattr(meta, "extra", None)
+    if extra and isinstance(extra, dict):
+        result.update(extra)
+    return result
+
+
+def scan_file(
+    file_path: str,
+    storage_options: dict[str, Any] | None = None,
+) -> FileIndex:
+    """Scan a ``.tgm`` file and return a :class:`FileIndex`.
+
+    Decodes each message to get descriptors, per-object metadata
+    (from ``meta.base`` and ``desc.params``), and extra metadata.
+
+    Parameters
+    ----------
+    file_path
+        Path or remote URL to the ``.tgm`` file.
+    storage_options
+        Key-value pairs forwarded to the object store backend for
+        remote URLs.  Ignored for local files.
+    """
+    import tensogram
+
+    is_remote = tensogram.is_remote_url(file_path)
+    resolved = file_path if is_remote else os.path.abspath(file_path)
+    index = FileIndex(file_path=resolved)
+
+    if is_remote:
+        f = tensogram.TensogramFile.open_remote(resolved, storage_options or {})
+    else:
+        f = tensogram.TensogramFile.open(resolved)
+
+    with f:
+        n_messages = len(f)
+        for msg_idx in range(n_messages):
+            if is_remote:
+                result = f.file_decode_descriptors(msg_idx)
+                meta = result["metadata"]
+                descriptors = result["descriptors"]
+            else:
+                raw = f.read_message(msg_idx)
+                meta = tensogram.decode_metadata(raw)
+                _, descriptors = tensogram.decode_descriptors(raw)
+
+            extra = _extra_from_meta(meta)
+
+            for obj_idx, desc in enumerate(descriptors):
+                per_obj = _merge_per_object_meta(meta, obj_idx, desc)
+                shape = tuple(desc.shape)
+                info = ObjectInfo(
+                    msg_index=msg_idx,
+                    obj_index=obj_idx,
+                    ndim=desc.ndim,
+                    shape=shape,
+                    dtype=desc.dtype,
+                    descriptor=desc,
+                    per_object_meta=per_obj,
+                    common_meta=dict(extra),
+                )
+                index.objects.append(info)
+
+    return index
+
+
+def scan_message(raw_msg: bytes) -> list[ObjectInfo]:
+    """Scan a single in-memory message and return :class:`ObjectInfo` list."""
+    import tensogram
+
+    meta = tensogram.decode_metadata(raw_msg)
+    extra = _extra_from_meta(meta)
+
+    _, descriptors = tensogram.decode_descriptors(raw_msg)
+
+    result: list[ObjectInfo] = []
+    for obj_idx, desc in enumerate(descriptors):
+        per_obj = _merge_per_object_meta(meta, obj_idx, desc)
+        shape = tuple(desc.shape)
+        info = ObjectInfo(
+            msg_index=0,
+            obj_index=obj_idx,
+            ndim=desc.ndim,
+            shape=shape,
+            dtype=desc.dtype,
+            descriptor=desc,
+            per_object_meta=per_obj,
+            common_meta=dict(extra),
+        )
+        result.append(info)
+
+    return result

tensogram_xarray/store.py

@@ -0,0 +1,383 @@
+# (C) Copyright 2026- ECMWF and individual contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation nor
+# does it submit to any jurisdiction.
+
+"""Data store: bridge between tensogram messages and xarray Variables.
+
+``TensogramDataStore`` reads a single tensogram message (identified by file
+path and message index) and produces :class:`xarray.Variable` objects with
+lazy-loaded data backed by :class:`TensogramBackendArray`.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+import xarray as xr
+from xarray.core import indexing
+
+from tensogram_xarray.array import TensogramBackendArray, _supports_range_decode
+from tensogram_xarray.coords import detect_coords
+from tensogram_xarray.mapping import resolve_dim_names, resolve_variable_name
+from tensogram_xarray.scanner import RESERVED_KEY
+
+logger = logging.getLogger(__name__)
+
+# Map tensogram dtype strings to numpy dtypes.
+try:
+    import ml_dtypes
+
+    _BFLOAT16_DTYPE = ml_dtypes.bfloat16
+except ImportError:
+    _BFLOAT16_DTYPE = np.dtype("uint16")  # fallback: raw 2-byte words
+
+_DTYPE_MAP: dict[str, np.dtype] = {
+    "float16": np.dtype("float16"),
+    "bfloat16": np.dtype(_BFLOAT16_DTYPE),
+    "float32": np.dtype("float32"),
+    "float64": np.dtype("float64"),
+    "complex64": np.dtype("complex64"),
+    "complex128": np.dtype("complex128"),
+    "int8": np.dtype("int8"),
+    "int16": np.dtype("int16"),
+    "int32": np.dtype("int32"),
+    "int64": np.dtype("int64"),
+    "uint8": np.dtype("uint8"),
+    "uint16": np.dtype("uint16"),
+    "uint32": np.dtype("uint32"),
+    "uint64": np.dtype("uint64"),
+    "bitmask": np.dtype("uint8"),  # packed bits as raw bytes
+}
+
+
+def _to_numpy_dtype(tgm_dtype: str) -> np.dtype:
+    """Convert a tensogram dtype string to a numpy dtype."""
+    key = tgm_dtype.lower()
+    if key in _DTYPE_MAP:
+        return _DTYPE_MAP[key]
+    # Fallback: try numpy directly.
+    try:
+        return np.dtype(tgm_dtype)
+    except TypeError as exc:
+        msg = f"unsupported tensogram dtype {tgm_dtype!r}: {exc}"
+        raise TypeError(msg) from exc
+
+
+class TensogramDataStore:
+    """Read-only data store for a single tensogram message.
+
+    Parameters
+    ----------
+    file_path
+        Path or remote URL to the ``.tgm`` file.
+    msg_index
+        Index of the message within the file.
+    dim_names
+        Optional user-specified dimension names for data variables.
+    variable_key
+        Optional dotted metadata path for variable naming.
+    verify_hash
+        Whether to verify object hashes on decode.
+    range_threshold
+        Maximum fraction of total array elements (0.0-1.0) for which
+        partial ``decode_range()`` is used.  Default ``0.5``.
+    storage_options
+        Key-value pairs forwarded to the object store backend when
+        ``file_path`` is a remote URL (S3, GCS, Azure, HTTP).  Used
+        for credentials, region, endpoint overrides, etc.  Ignored
+        for local files.
+    """
+
+    def __init__(
+        self,
+        file_path: str,
+        msg_index: int = 0,
+        dim_names: Sequence[str] | None = None,
+        variable_key: str | None = None,
+        verify_hash: bool = False,
+        range_threshold: float = 0.5,
+        storage_options: dict[str, Any] | None = None,
+    ):
+        import tensogram
+
+        self._is_remote = tensogram.is_remote_url(file_path)
+        self.file_path = file_path if self._is_remote else os.path.abspath(file_path)
+        self.msg_index = msg_index
+        self.dim_names = dim_names
+        self.variable_key = variable_key
+        self.verify_hash = verify_hash
+        self.range_threshold = range_threshold
+        self.storage_options = storage_options
+        self._lock = threading.Lock()
+        self._backend_arrays: list[TensogramBackendArray] = []
+
+        self._file = self._open_file()
+        self._meta, self._descriptors = self._read_metadata()
+
+    def _open_file(self) -> Any:
+        import tensogram
+
+        if self._is_remote:
+            return tensogram.TensogramFile.open_remote(self.file_path, self.storage_options or {})
+        return tensogram.TensogramFile.open(self.file_path)
+
+    def _read_metadata(self) -> tuple[Any, list]:
+        import tensogram
+
+        if self._is_remote:
+            result = self._file.file_decode_descriptors(self.msg_index)
+            return result["metadata"], result["descriptors"]
+
+        raw = self._file.read_message(self.msg_index)
+        meta = tensogram.decode_metadata(raw)
+        _, descriptors = tensogram.decode_descriptors(raw)
+        return meta, descriptors
+
+    def _get_common_meta(self) -> dict[str, Any]:
+        """Extract message-level metadata for Dataset attributes.
+
+        Reads from ``meta.extra`` (message-level annotations).
+        """
+        attrs: dict[str, Any] = {}
+        extra = getattr(self._meta, "extra", None)
+        if extra and isinstance(extra, dict):
+            attrs.update(extra)
+        attrs["tensogram_version"] = getattr(self._meta, "version", 2)
+        return attrs
+
+    def _get_per_object_meta(self, obj_index: int, desc: Any) -> dict[str, Any]:
+        """Extract per-object metadata.
+
+        Reads from ``meta.base[obj_index]`` (primary), filtering out the
+        ``_reserved_`` key (encoder-populated tensor info).  Then merges in
+        ``desc.params`` (fallback for extra keys in the descriptor dict).
+
+        If ``obj_index`` is out of range (fewer base entries than objects),
+        a warning is logged and the base entry is treated as empty.
+        """
+        meta: dict[str, Any] = {}
+        # Primary source: meta.base[i]
+        base = getattr(self._meta, "base", None)
+        if base is not None and isinstance(base, list):
+            if obj_index < len(base):
+                entry = base[obj_index]
+                if isinstance(entry, dict):
+                    for k, v in entry.items():
+                        if k != RESERVED_KEY:
+                            meta[k] = v
+            else:
+                logger.warning(
+                    "meta.base has %d entries but object index %d requested; "
+                    "per-object metadata will be empty for this object",
+                    len(base),
+                    obj_index,
+                )
+        # Fallback/supplement: desc.params (extra descriptor keys)
+        params = getattr(desc, "params", None)
+        if params and isinstance(params, dict):
+            for k, v in params.items():
+                if k not in meta:
+                    meta[k] = v
+        return meta
+
+    def build_dataset(
+        self,
+        drop_variables: set[str] | None = None,
+    ) -> xr.Dataset:
+        """Construct an :class:`xr.Dataset` from this message.
+
+        Coordinate objects are auto-detected by name matching.  Data objects
+        become lazy-loaded variables.  All metadata flows to attributes.
+        """
+        dataset_attrs = self._get_common_meta()
+
+        # Gather per-object metadata from payload + descriptor params.
+        obj_metas = [self._get_per_object_meta(i, d) for i, d in enumerate(self._descriptors)]
+
+        # Detect coordinates vs data variables.
+        coord_indices, var_indices, coord_dim_names = detect_coords(obj_metas)
+
+        # Build coordinate variables from detected coord objects.
+        coord_vars: dict[str, xr.Variable] = {}
+        for ci in coord_indices:
+            desc = self._descriptors[ci]
+            dim_name = coord_dim_names[ci]
+            np_dtype = _to_numpy_dtype(desc.dtype)
+            shape = tuple(desc.shape)
+
+            backend_array = TensogramBackendArray(
+                file_path=self.file_path,
+                msg_index=self.msg_index,
+                obj_index=ci,
+                shape=shape,
+                dtype=np_dtype,
+                supports_range=_supports_range_decode(desc),
+                verify_hash=self.verify_hash,
+                range_threshold=self.range_threshold,
+                lock=self._lock,
+                storage_options=self.storage_options,
+                shared_file=self._file,
+            )
+            self._backend_arrays.append(backend_array)
+            lazy_data = indexing.LazilyIndexedArray(backend_array)
+
+            coord_dims = (dim_name,)
+            coord_attrs = dict(obj_metas[ci])
+            coord_vars[dim_name] = xr.Variable(coord_dims, lazy_data, coord_attrs)
+
+        data_vars: dict[str, xr.Variable] = {}
+        for vi in var_indices:
+            desc = self._descriptors[vi]
+            var_name = resolve_variable_name(vi, obj_metas[vi], self.variable_key)
+            if drop_variables and var_name in drop_variables:
+                continue
+
+            np_dtype = _to_numpy_dtype(desc.dtype)
+            shape = tuple(desc.shape)
+
+            dims = self._resolve_dims_for_var(shape, coord_vars)
+
+            backend_array = TensogramBackendArray(
+                file_path=self.file_path,
+                msg_index=self.msg_index,
+                obj_index=vi,
+                shape=shape,
+                dtype=np_dtype,
+                supports_range=_supports_range_decode(desc),
+                verify_hash=self.verify_hash,
+                range_threshold=self.range_threshold,
+                lock=self._lock,
+                storage_options=self.storage_options,
+                shared_file=self._file,
+            )
+            self._backend_arrays.append(backend_array)
+            lazy_data = indexing.LazilyIndexedArray(backend_array)
+
+            var_attrs = dict(obj_metas[vi])
+            data_vars[var_name] = xr.Variable(dims, lazy_data, var_attrs)
+
+        # Assemble Dataset.
+        ds = xr.Dataset(data_vars, coords=coord_vars, attrs=dataset_attrs)
+        return ds
+
+    def _get_meta_dim_names(self, ndim: int) -> list[str] | dict[int, str]:
+        """Return dimension name hints from ``_extra_["dim_names"]``.
+
+        Any writer can embed hints at ``_extra_["dim_names"]`` to provide
+        meaningful dimension names without the reader passing ``dim_names``
+        explicitly. Two formats are accepted:
+
+        **List (preferred)** — axis-ordered names, one per dimension::
+
+            "_extra_": {"dim_names": ["values", "level"]}
+
+        This handles axes with identical sizes correctly because names
+        are assigned by position.
+
+        **Dict (legacy)** — size-to-name mapping::
+
+            "_extra_": {"dim_names": {"1000": "values", "50": "level"}}
+
+        A dict cannot disambiguate axes with the same size; only the
+        first matching axis receives the hint.
+
+        Returns an empty list/dict when the key is absent or malformed.
+        """
+        try:
+            raw = self._meta.extra["dim_names"]
+        except (AttributeError, KeyError, TypeError):
+            return {}
+
+        # List format: axis-ordered names.
+        if isinstance(raw, list):
+            try:
+                names = [str(n) for n in raw]
+                if len(names) == ndim:
+                    return names
+                # Length mismatch — ignore the hint rather than crash.
+                return {}
+            except (TypeError, ValueError):
+                return {}
+
+        # Dict format: size -> name (legacy).
+        if isinstance(raw, dict):
+            try:
+                return {int(k): str(v) for k, v in raw.items()}
+            except (TypeError, ValueError):
+                return {}
+
+        return {}
+
+    def _resolve_dims_for_var(
+        self,
+        shape: tuple[int, ...],
+        coord_vars: dict[str, xr.Variable],
+    ) -> tuple[str, ...]:
+        """Assign dimension names for a data variable.
+
+        Strategy:
+        1. If user provided ``dim_names``, use them directly.
+        2. Try to match each axis size against a known coordinate variable.
+        3. Use producer hints from ``_extra_["dim_names"]`` (list or dict).
+        4. Fall back to ``dim_0``, ``dim_1``, ...
+        """
+        ndim = len(shape)
+
+        # (1) Explicit user mapping.
+        if self.dim_names is not None:
+            return tuple(resolve_dim_names(ndim, self.dim_names))
+
+        # (2) Match by size against detected coordinates.
+        # Build a map: size -> list of coord dim names with that size.
+        size_to_coord: dict[int, list[str]] = {}
+        for cname, cvar in coord_vars.items():
+            csize = cvar.shape[0]
+            size_to_coord.setdefault(csize, []).append(cname)
+
+        # (3) Metadata hints written by the producer.
+        meta_hints = self._get_meta_dim_names(ndim)
+
+        # If the producer supplied an axis-ordered list, use it directly.
+        if isinstance(meta_hints, list):
+            return tuple(meta_hints)
+
+        # Otherwise meta_hints is a dict (size -> name); match by size.
+        dims: list[str] = []
+        used: set[str] = set()
+        for axis, axis_size in enumerate(shape):
+            matched = False
+            # Prefer a detected coordinate dimension.
+            if axis_size in size_to_coord:
+                for cname in size_to_coord[axis_size]:
+                    if cname not in used:
+                        dims.append(cname)
+                        used.add(cname)
+                        matched = True
+                        break
+            # Fall back to producer hint (dict).
+            if not matched and axis_size in meta_hints:
+                hint = meta_hints[axis_size]
+                if hint not in used:
+                    dims.append(hint)
+                    used.add(hint)
+                    matched = True
+            # Final fallback.
+            if not matched:
+                dims.append(f"dim_{axis}")
+
+        return tuple(dims)
+
+    def close(self) -> None:
+        for arr in self._backend_arrays:
+            arr._shared_file = None
+        self._backend_arrays.clear()
+        self._file = None

tensogram_xarray-0.14.0.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: tensogram-xarray
+Version: 0.14.0
+Summary: xarray backend engine for tensogram .tgm files
+Project-URL: Homepage, https://sites.ecmwf.int/docs/tensogram/main
+Project-URL: Repository, https://github.com/ecmwf/tensogram
+Project-URL: Documentation, https://sites.ecmwf.int/docs/tensogram/main
+Author-email: ECMWF <software@ecmwf.int>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
+Requires-Python: >=3.9
+Requires-Dist: numpy
+Requires-Dist: tensogram<0.15,>=0.14.0
+Requires-Dist: xarray>=2022.06
+Provides-Extra: dask
+Requires-Dist: dask[array]; extra == 'dask'
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'

tensogram_xarray-0.14.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+tensogram_xarray/__init__.py,sha256=UmIUhCzEZyCRC9zueuwHITobJcMXp34qS7QCjX9NAIw,895
+tensogram_xarray/array.py,sha256=MFb2m3JAfo2hQVnABq1zcdC9mMa55X6r0uGn0Te3eXY,14658
+tensogram_xarray/backend.py,sha256=JIrxmYXNNWHGn_kHtRUA_VngvtC9gzA2Z0LWdodzHhg,4970
+tensogram_xarray/coords.py,sha256=Vk7k7oF-GgS6EE5BwMsHanzWlXp4ersBuuuXVvyjayM,3398
+tensogram_xarray/mapping.py,sha256=i763IpvWxky1G9U27e9SWMPanCHLOlUl9DyHrfX2H54,3019
+tensogram_xarray/merge.py,sha256=zFr_AXwqtgz1R2Ag8RDws-kdfKkiuktTOh5XIPnJF4s,29206
+tensogram_xarray/scanner.py,sha256=uCeV-B_WBegNPP7qAuZPhUcgGUtQ8bEJfzMB5GOcKVQ,6439
+tensogram_xarray/store.py,sha256=-XfEnCyScpE0aibzoGQka1oZ7yay8PUnN3V1XqZBj1k,14083
+tensogram_xarray-0.14.0.dist-info/METADATA,sha256=Eb4lYe2yFFTif90BAIbSxBgGPumro-oxvU_Ef90adcw,936
+tensogram_xarray-0.14.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+tensogram_xarray-0.14.0.dist-info/entry_points.txt,sha256=l5kLmFeoDJheSrXui7OlHa3WeSkDu5y5lmUoC7kU0Rc,82
+tensogram_xarray-0.14.0.dist-info/RECORD,,

tensogram_xarray-0.14.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

tensogram_xarray-0.14.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[xarray.backends]
+tensogram = tensogram_xarray.backend:TensogramBackendEntrypoint