diff --git a/.github/workflows/z3fdb.yml b/.github/workflows/z3fdb.yml index 053f0a3f4..348ac6beb 100644 --- a/.github/workflows/z3fdb.yml +++ b/.github/workflows/z3fdb.yml @@ -24,8 +24,14 @@ on: jobs: prepare-deps: - runs-on: ubuntu-latest if: ${{ !github.event.pull_request.draft && (success() || failure()) && (!github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci') }} + runs-on: [self-hosted, Linux, platform-builder-docker-xl] + container: + image: eccr.ecmwf.int/platform-builder/platform-builder:ubuntu-24.04 + credentials: + username: ${{ secrets.ECMWF_DOCKER_REGISTRY_USERNAME }} + password: ${{ secrets.ECMWF_DOCKER_REGISTRY_ACCESS_TOKEN }} + options: '--user root' # see https://github.com/ecmwf/reusable-workflows/blob/main/.github/workflows/cd-system-package.yml steps: - name: Get ecbuild uses: actions/checkout@v5 @@ -108,11 +114,17 @@ jobs: retention-days: 1 build-wheels: needs: prepare-deps - runs-on: ubuntu-latest if: ${{ !github.event.pull_request.draft && (success() || failure()) && (!github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci') }} + runs-on: [self-hosted, Linux, platform-builder-docker-xl] + container: + image: eccr.ecmwf.int/platform-builder/platform-builder:ubuntu-24.04 + credentials: + username: ${{ secrets.ECMWF_DOCKER_REGISTRY_USERNAME }} + password: ${{ secrets.ECMWF_DOCKER_REGISTRY_ACCESS_TOKEN }} + options: '--user root' # see https://github.com/ecmwf/reusable-workflows/blob/main/.github/workflows/cd-system-package.yml strategy: matrix: - python-version: ['3.11', '3.12', '3.13'] # eccodes does not yet support 3.14 + python-version: ['3.11', '3.12', '3.13', '3.14'] # eccodes does not yet support 3.14 fail-fast: false # Continue running other versions if one fails steps: diff --git a/CMakeLists.txt b/CMakeLists.txt index 1a7805edd..e3d75a52c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -139,6 +139,15 @@ ecbuild_add_option( FEATURE PYTHON_ZARR_INTERFACE DEFAULT OFF REQUIRED_PACKAGES "NAME pybind11 VERSION 3.0.1" ) +### Enables the GribJump-backed extractor for the Zarr interface. +### NOTE: gribjump itself depends on fdb5, so this is only satisfiable in a bundle build where +### gribjump is a sibling project. It can never be satisfied by a standalone fdb build against +### an installed gribjump, which is why it defaults to OFF. +ecbuild_add_option( FEATURE ZARR_GRIBJUMP_EXTRACTOR + DESCRIPTION "Build the GribJump-backed extractor (bundle builds only)" + DEFAULT OFF + REQUIRED_PACKAGES "NAME gribjump" ) + # We need PyFDB in case z3fdb is enabled if(HAVE_PYTHON_ZARR_INTERFACE AND NOT HAVE_PYTHON_FDB_INTERFACE) message( WARN "FDB Zarr interface requires python FDB interface. I enable ENABLE_PYTHON_FDB_INTERFACE to build the PyFDB interface" ) diff --git a/cmake/z3fdb_setup.py.in b/cmake/z3fdb_setup.py.in index 059db00f9..0d8dacdd6 100644 --- a/cmake/z3fdb_setup.py.in +++ b/cmake/z3fdb_setup.py.in @@ -8,6 +8,16 @@ from wheel.bdist_wheel import bdist_wheel import sys +def _read_requirements(name): + """Requirement lines from a pip requirements file staged next to this setup.py. + + Comments and blank lines are dropped; everything else is passed through verbatim. + """ + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), name) + with open(path) as requirements_file: + return [line.strip() for line in requirements_file if line.strip() and not line.lstrip().startswith("#")] + + # NOTE this we need to correctly link with fdb5lib version on cd. For local builds feel free to ignore version_suffix = os.environ.get("VERSION_SUFFIX", "") if version_suffix: @@ -45,9 +55,17 @@ setup( "z3fdb._internal", "pychunked_data_view", "chunked_data_view_bindings", + # PEP 561 stub-only package. The hyphen is mandated by the PEP but is not a valid + # Python identifier, so setuptools cannot derive its directory from the name, + # package_dir below states it explicitly. + "chunked_data_view_bindings-stubs", ], + package_dir={ + "chunked_data_view_bindings-stubs": "chunked_data_view_bindings-stubs", + }, package_data={ - "chunked_data_view_bindings": ["*.so", "*.pyd", "*.pyi", "py.typed"], + "chunked_data_view_bindings": ["*.so", "*.pyd"], + "chunked_data_view_bindings-stubs": ["*.pyi", "py.typed"], }, license="Apache 2.0", license_files=["LICENSE"], @@ -65,7 +83,10 @@ setup( "Operating System :: OS Independent", "Topic :: Software Development :: Libraries", ], - install_requires=["numpy", "zarr~=3.1.6", "findlibs>=0.1.2"] + requires_extra, + # Runtime dependencies live in python/z3fdb-requirements.txt, staged next to this file by + # src/CMakeLists.txt. Build tooling stays out of the wheel metadata; it is provisioned into + # the uvx environment from python/z3fdb-build-requirements.txt instead. + install_requires=_read_requirements("z3fdb-requirements.txt") + requires_extra, python_requires=python_requires, has_ext_modules=lambda: True, **ext_kwargs[sys.platform], diff --git a/docs/z3fdb/api.rst b/docs/z3fdb/api.rst index 1708084b8..e2cc0ce27 100644 --- a/docs/z3fdb/api.rst +++ b/docs/z3fdb/api.rst @@ -21,6 +21,42 @@ z3fdb.Z3fdbError .. autoapiexception:: z3fdb.Z3fdbError +Extractor errors +^^^^^^^^^^^^^^^^ + +Raised by the extractor backends and re-exported from +:mod:`pychunked_data_view`, so they can be caught by type: + +``GribExtractorError`` + A GRIB field could not be retrieved or decoded. For example, FDB returned no field + for a sub-request, or a field's size does not match the rest of the view. + +``GribJumpExtractorError`` + GribJump extraction failed. For example, a field location carried no usable file + offset, or the request matched nothing. + +``MarsRequestFormattingError`` + A malformed MARS request string: a trailing comma, a missing comma between keys, or a + misspelled key. Raised from ``build()``; a subclass of ``RuntimeError``. + +``InternalError`` + Something inside ``pychunked_data_view`` is inconsistent. You should not see this. + +Note that other misconfiguration detected by the builder (unmapped axes, incompatible parts, an +invalid chunk size, parts disagreeing about the grid) surfaces as a plain ``RuntimeError``, +since it originates as an ``eckit::UserError``. + +Build capability +^^^^^^^^^^^^^^^^ + +.. py:data:: pychunked_data_view.has_gribjump_extractor + :type: bool + + Whether this build compiled the GribJump extractor. + :class:`~pychunked_data_view.ExtractorType.GribJump` can always be constructed, so this is + the way to find out whether it can actually be used. See + :ref:`z3fdb_gribjump_availability`. + Type aliases ------------ @@ -32,9 +68,9 @@ z3fdb.MarsSelection A MARS request expressed as a mapping. Keys are MARS keyword names (strings). Values may be: - * a single ``str``, ``int``, or ``float`` — e.g. ``"step": 0`` - * a list of ``str``, ``int``, or ``float`` — e.g. ``"param": [165, 166]`` - * a MARS range expression passed as a ``str`` — e.g. + * a single ``str``, ``int``, or ``float``, e.g. ``"step": 0`` + * a list of ``str``, ``int``, or ``float``, e.g. ``"param": [165, 166]`` + * a MARS range expression passed as a ``str``, e.g. ``"date": "2020-01-01/to/2020-01-04"`` Example:: @@ -58,9 +94,40 @@ Classes z3fdb.SimpleStoreBuilder ^^^^^^^^^^^^^^^^^^^^^^^^ +Creates a store whose root *is* the array. Equivalent to +:class:`~z3fdb.CustomStoreBuilder` restricted to ``path=None``, which is what it +delegates to. + .. autoapiclass:: z3fdb.SimpleStoreBuilder :members: +z3fdb.ChunkedDataView +^^^^^^^^^^^^^^^^^^^^^ + +The read-only array returned by ``build()`` on the lower-level +:class:`~pychunked_data_view.ChunkedDataViewBuilder`. Zarr normally drives it for you. + +.. note:: + + ``chunkShape()`` is deprecated in favour of ``chunk_shape()``; it still works but emits a + ``DeprecationWarning``. Every other accessor on the class is already snake_case. + +.. autoapiclass:: pychunked_data_view.ChunkedDataView + :members: + +z3fdb.CustomStoreBuilder +^^^^^^^^^^^^^^^^^^^^^^^^ + +Creates a store with an arbitrary group/array hierarchy: every method takes a +zarr-style *path* naming the array it applies to, and ``path=None`` addresses a +root array (mutually exclusive with any named path). + +.. seealso:: :ref:`tutorial_custom_store_mixed_extractors` for a worked example + building several arrays with different extractors in one store. + +.. autoapiclass:: z3fdb.CustomStoreBuilder + :members: + z3fdb.AxisDefinition ^^^^^^^^^^^^^^^^^^^^ @@ -70,21 +137,21 @@ to Zarr dimensions. .. autoapiclass:: pychunked_data_view.AxisDefinition :members: -Enums ------ +Chunking +-------- z3fdb.Chunking -^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^ .. autoapiclass:: pychunked_data_view.Chunking :members: -.. py:class:: pychunked_data_view.Chunking.FixedSizeChunk(chunkShape) +.. py:class:: pychunked_data_view.Chunking.FixedSizeChunk(chunk_shape) Specifies a custom chunk size along a single axis. This is a frozen dataclass nested inside :class:`~pychunked_data_view.Chunking`. - .. py:attribute:: chunkShape + .. py:attribute:: chunk_shape :type: int Number of consecutive axis values grouped into each chunk. @@ -97,14 +164,82 @@ z3fdb.Chunking .. code-block:: python # Chunk a 12-date axis into groups of 3 (gives 4 chunks) - AxisDefinition(["date"], Chunking.FixedSizeChunk(chunkShape=3)) + AxisDefinition(["date"], Chunking.FixedSizeChunk(chunk_shape=3)) - See :ref:`dimension_mapping:Chunking` for a full comparison of + See :ref:`z3fdb_chunking` for a full comparison of chunking modes and guidance on when to use each one. -z3fdb.ExtractorType -^^^^^^^^^^^^^^^^^^^ +Extractors +---------- -.. autoapiclass:: pychunked_data_view.ExtractorType - :members: +``ExtractorType`` is a namespace class, not an enum. Its nested classes +carry per-extractor configuration. Pass an *instance* to +:meth:`~z3fdb.SimpleStoreBuilder.add_part`. + +``add_part`` stores a *copy* of the configuration, so one instance can be reused across as many +parts and builders as you like, and the ``fdb_config`` a builder fills in for you is never +written back into your object. + +.. seealso:: :ref:`z3fdb_extractor_backends` for what the two backends do, their constraints, + and which builds provide GribJump. + +z3fdb.ExtractorType.Grib +^^^^^^^^^^^^^^^^^^^^^^^^ + +.. py:class:: pychunked_data_view.ExtractorType.Grib(*, fdb_config=None) + + Reads full GRIB fields from FDB and decodes them to ``float32`` via eccodes. + This is the default extractor for standard GRIB data. + + :param fdb_config: Path to an FDB configuration YAML file. + ``None`` (default) uses the path passed to :class:`~z3fdb.SimpleStoreBuilder`. + :type fdb_config: pathlib.Path or None + + **Example** + + .. code-block:: python + + builder.add_part(mars_request, axes, ExtractorType.Grib()) + + # With an explicit FDB config + builder.add_part(mars_request, axes, ExtractorType.Grib(fdb_config=Path("/etc/fdb/config.yaml"))) + +z3fdb.ExtractorType.GribJump +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. py:class:: pychunked_data_view.ExtractorType.GribJump(*, fdb_config=None, gribjump_config=None, field_chunking=None) + + Reads grid-point values from FDB using GribJump, a library that jumps + directly to the values inside the GRIB message without performing a full + decode. + + :param fdb_config: Path to an FDB configuration YAML file. + ``None`` (default) uses the path passed to :class:`~z3fdb.SimpleStoreBuilder`. + :type fdb_config: pathlib.Path or None + + :param gribjump_config: Path to a GribJump configuration YAML file. + ``None`` (default) reads the ``GRIBJUMP_CONFIG_FILE`` environment variable. + :type gribjump_config: pathlib.Path or None + + :param field_chunking: How to sub-divide the implicit (grid-point) dimension into + Zarr chunks. ``None`` (default) produces a single chunk covering the + full field. Pass :class:`pychunked_data_view.Chunking.FixedSizeChunk` + to split the implicit axis into equal-sized pieces; the size must divide + the grid exactly, as that dimension cannot be left ragged. + :type field_chunking: pychunked_data_view.Chunking.FixedSizeChunk or None + + **Example** + + .. code-block:: python + + # Full field: avoids eccodes decode + builder.add_part(mars_request, axes, ExtractorType.GribJump()) + + # Split the implicit grid-point axis into chunks of 1312 + builder.add_part(mars_request, axes, + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(1312))) + + .. seealso:: :ref:`z3fdb_extractor_backends` for how the two backends differ, their + constraints, and which builds provide GribJump; and + :ref:`tutorial_custom_store_mixed_extractors` for a worked example using both. diff --git a/docs/z3fdb/dimension_mapping.rst b/docs/z3fdb/dimension_mapping.rst index aff276dc5..53a738414 100644 --- a/docs/z3fdb/dimension_mapping.rst +++ b/docs/z3fdb/dimension_mapping.rst @@ -7,7 +7,7 @@ Dimension Mapping and Data Model A MARS request defines which data to retrieve from FDB. Each keyword with more than one value defines an axis and **must** be mapped to a Zarr dimension via :class:`~pychunked_data_view.AxisDefinition`. -Keywords with a single value **may** also be mapped — useful when MARS +Keywords with a single value **may** also be mapped. This is useful when MARS restricts a keyword to one value but you still want it as an explicit dimension in the resulting array. @@ -119,8 +119,8 @@ Coordinate System A ChunkedDataView exposes an ``(N+1)``-dimensional integer index space: -* **Axes 0 … N−1** — one per :class:`~pychunked_data_view.AxisDefinition`. -* **Axis N** — the implicit trailing dimension holding the decoded +* **Axes 0 to N-1**, one per :class:`~pychunked_data_view.AxisDefinition`. +* **Axis N**, the implicit trailing dimension holding the decoded grid-point float32 values for each field. The array **shape** is:: @@ -131,9 +131,17 @@ All axis indices are zero-based. Axis sizes are determined by the total number of distinct values held across all ``parts`` (see `Combining Multiple MARS Requests`_ below). +.. _z3fdb_chunking: + Chunking -------- +.. note:: + + A ``GribJump`` part can also chunk the implicit grid-point axis, via ``field_chunking``. That + size must divide the grid exactly, and only the default whole-axis setting can be combined + with a ``Grib`` part. See :ref:`z3fdb_extractor_backends`. + :class:`~pychunked_data_view.Chunking` determines how many values along a dimension are grouped into a single Zarr chunk: @@ -149,7 +157,7 @@ a dimension are grouped into a single Zarr chunk: * - :attr:`~pychunked_data_view.Chunking.WHOLE_AXIS` - The entire axis is stored in a single chunk - Full axis length - * - :class:`~pychunked_data_view.Chunking.FixedSizeChunk` ``(chunkShape=k)`` + * - :class:`~pychunked_data_view.Chunking.FixedSizeChunk` ``(chunk_shape=k)`` - Groups every ``k`` consecutive values along the axis into one chunk. ``k`` must divide the axis length evenly. - ``k`` (user-specified) @@ -167,18 +175,18 @@ values: # Chunk shape: (4, 1, N) [ - AxisDefinition(["date"], Chunking.FixedSizeChunk(chunkShape=2)), # chunk size = 2 + AxisDefinition(["date"], Chunking.FixedSizeChunk(chunk_shape=2)), # chunk size = 2 AxisDefinition(["param"], Chunking.SINGLE_VALUE), # chunk size = 1 ] # Array shape: (4, 3, N) # Chunk shape: (2, 1, N) ← two dates per chunk, four chunks total :class:`~pychunked_data_view.Chunking.FixedSizeChunk` is useful when -neither extreme fits — for instance, when you want to batch a temporal +neither extreme fits. For instance, when you want to batch a temporal axis into multi-day windows for efficient I/O while still keeping chunks small enough to fit in memory. -A chunk is addressed by a *chunk-index tuple* ``(c0, c1, …, cN-1)`` — +A chunk is addressed by a *chunk-index tuple* ``(c0, c1, ..., cN-1)``, one integer per MARS axis. The implicit values dimension is never chunked. Each chunk index ``ci`` maps to an axis range: @@ -188,11 +196,11 @@ chunked. Each chunk index ``ci`` maps to an axis range: * - Chunking - Axis range covered by chunk index ``ci`` * - ``SINGLE_VALUE`` - - ``[ci, ci]`` — exactly one slot + - ``[ci, ci]``, exactly one slot * - ``WHOLE_AXIS`` - - ``[0, size_i − 1]`` — the full axis (``ci`` is always 0) - * - ``FixedSizeChunk(chunkShape=k)`` - - ``[ci × k, (ci + 1) × k − 1]`` — a window of ``k`` consecutive values + - ``[0, size_i - 1]``, the full axis (``ci`` is always 0) + * - ``FixedSizeChunk(chunk_shape=k)`` + - ``[ci * k, (ci + 1) * k - 1]``, a window of ``k`` consecutive values The chunk's **bounding box** is the Cartesian product of these per-axis ranges. Its flat memory footprint is:: @@ -249,27 +257,33 @@ three axes reduces each chunk to a single field - Use when you always read the full axis in one go and want to reduce the number of FDB round-trips (e.g. a small ``step`` axis you always load entirely). - * - ``FixedSizeChunk(chunkShape=k)`` - - Use when you need a middle ground — for example, batching a + * - ``FixedSizeChunk(chunk_shape=k)`` + - Use when you need a middle ground. For example, batching a 365-day date axis into weekly (``k=7``) or monthly (``k=30``) windows. ``k`` must divide the axis length exactly. Fill Value ---------- -When a chunk is accessed and some of its fields are absent from FDB, -the missing slots are filled with a sentinel value. The default is -``float('nan')``. +The fill value is what Z3FDB writes in place of grid points a GRIB message flags as missing +through its bitmap. It is also written into the zarr array metadata as ``fill_value``. The +default is ``float('nan')``. See :ref:`z3fdb_missing_values` for what this means when +reading data. + +It is *not* a substitute for absent data: if a request matches fewer fields than the view +expects, accessing the chunk raises rather than filling the gap. A view describes data that +exists. - To override it, call :meth:`~z3fdb.SimpleStoreBuilder.fill_missing_value` (or use - :class:`~pychunked_data_view.ChunkedDataViewBuilder` directly): +To override it, call :meth:`~z3fdb.SimpleStoreBuilder.fill_missing_value` (or use +:class:`~pychunked_data_view.ChunkedDataViewBuilder` directly). The call configures an existing +array, so add a part first: .. code-block:: python from pychunked_data_view import ChunkedDataViewBuilder, AxisDefinition, Chunking, ExtractorType builder = ChunkedDataViewBuilder(fdb_config_file=None) - builder.add_part({...}, [...], ExtractorType.GRIB) + builder.add_part({...}, [...], ExtractorType.Grib()) builder.fill_missing_value(-999.0) # use -999.0 instead of NaN view = builder.build() print(view.fill_missing_value()) # -999.0 @@ -302,7 +316,7 @@ the same number of values across ``parts``. AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) # Part 2: pressure level parameters @@ -321,7 +335,7 @@ the same number of values across ``parts``. AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) # Extend on the param dimension (index 1) @@ -334,7 +348,7 @@ The datetime dimension (index 0) must have the same values in both 4 pressure-level combinations (2 params × 2 levels) = 6 entries total. Each ``part`` occupies a rectangular sub-region of the global index space -described by a closed bounding box — one ``[lower_i, upper_i]`` interval +described by a closed bounding box, one ``[lower_i, upper_i]`` interval per axis (both bounds **inclusive**). ``Parts`` tile the extension axis without overlap; their bounding boxes are identical on every other axis. @@ -342,8 +356,8 @@ without overlap; their bounding boxes are identical on every other axis. Global index space after extend_on_axis(1): - Axis 0 (date×time) varies along rows (0–3). - Axis 1 (param) varies along columns (0–5). + Axis 0 (date x time) varies along rows, 0 to 3. + Axis 1 (param) varies along columns, 0 to 5. 0 1 2 3 4 5 ┌───┬───┬───┬───┬───┬───┐ @@ -362,10 +376,10 @@ without overlap; their bounding boxes are identical on every other axis. .. seealso:: - :doc:`technical_insights/chunk_access` — how the library resolves a chunk + :doc:`technical_insights/chunk_access` for how the library resolves a chunk access into FDB sub-requests and writes each field into the correct buffer slot. - :doc:`technical_insights/buffer_layout` — the buffer position formula and + :doc:`technical_insights/buffer_layout` for the buffer position formula and flat-index calculation in full detail. diff --git a/docs/z3fdb/getting_started.rst b/docs/z3fdb/getting_started.rst index b737ada4a..3c4344fe4 100644 --- a/docs/z3fdb/getting_started.rst +++ b/docs/z3fdb/getting_started.rst @@ -11,7 +11,7 @@ understand how to index into it. Prerequisites ------------- -* Z3FDB installed — see :ref:`Z3FDB_Introduction` for build instructions. +* Z3FDB installed. See :ref:`Z3FDB_Introduction` for build instructions. * An FDB instance containing GRIB data accessible from your environment. * ``zarr`` installed. @@ -19,7 +19,7 @@ Your First Store ---------------- The example below creates a 3-dimensional Zarr array from two dates, four -time steps, and one surface parameter. Read it top to bottom — each step is +time steps, and one surface parameter. Read it top to bottom. Each step is explained immediately after the code block. .. code-block:: python @@ -46,11 +46,11 @@ explained immediately after the code block. AxisDefinition(["date"], Chunking.SINGLE_VALUE), # → Dim 0, size 2 AxisDefinition(["time"], Chunking.SINGLE_VALUE), # → Dim 1, size 4 ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) **What each piece does:** @@ -59,10 +59,10 @@ explained immediately after the code block. Call :meth:`~z3fdb.SimpleStoreBuilder.add_part` once per MARS request, then :meth:`~z3fdb.SimpleStoreBuilder.build` to produce a Zarr store. -``add_part(mars_request, axes, ExtractorType.GRIB)`` +``add_part(mars_request, axes, ExtractorType.Grib())`` Registers one MARS request. The ``axes`` list controls how MARS keywords - become Zarr dimensions. ``ExtractorType.GRIB`` tells Z3FDB the data is - encoded as GRIB. + become Zarr dimensions. ``ExtractorType.Grib()`` tells Z3FDB to decode the + GRIB fields in full using eccodes. ``AxisDefinition(keys, chunking)`` Maps one or more MARS keywords to **exactly one** Zarr dimension. @@ -82,9 +82,9 @@ After ``build()``, ``data`` is a 3-dimensional array: data.shape -> (2, 4, N) - Dim 0 — date: 2 entries (2020-01-01, 2020-01-02) - Dim 1 — time: 4 entries (0000, 0600, 1200, 1800) - Dim 2 — grid points: N float32 values decoded from the GRIB field [implicit] + Dim 0 (date): 2 entries (2020-01-01, 2020-01-02) + Dim 1 (time): 4 entries (0000, 0600, 1200, 1800) + Dim 2 (grid pts): N float32 values decoded from the GRIB field [implicit] The **implicit final dimension** always holds the decoded grid-point values for one field. Its size ``N`` is determined by the GRIB grid. @@ -103,10 +103,12 @@ is typically retrieved as a slice: # date index 1 (2020-01-02), time index 0 (0000) field = data[1, 0, :] -**No data is fetched from FDB until you index.** Building the store is cheap — -it validates your MARS request and pre-fetches layout metadata but does not +**No data is fetched from FDB until you index.** Building the store is cheap. +It validates your MARS request and pre-fetches layout metadata, but does not retrieve field values. +.. _z3fdb_missing_values: + Missing Data and Fill Values ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -125,7 +127,7 @@ detectable with :func:`numpy.isnan`: field = data[0, 2, :] missing = np.isnan(field) -To use a different sentinel — for example ``-999.0`` — call +To use a different sentinel, for example ``-999.0``, call :meth:`~z3fdb.SimpleStoreBuilder.fill_missing_value` on the builder before calling ``build()``: @@ -168,7 +170,7 @@ Pass both keys to one :class:`~pychunked_data_view.AxisDefinition`: [ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), # -> Dim 0, size 8 ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) This produces shape ``(8, N)``. The **rightmost key varies fastest** (C / NumPy @@ -203,7 +205,7 @@ dimension grows across the two ``parts``. builder = SimpleStoreBuilder() - # Part 1 — surface parameters (levtype=sfc): 2 params + # Part 1: surface parameters (levtype=sfc): 2 params builder.add_part( { "class": "ea", @@ -218,13 +220,13 @@ dimension grows across the two ``parts``. "param": [165, 166], # 2 surface params }, [ - AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), # Dim 0 — 8 entries - AxisDefinition(["param"], Chunking.SINGLE_VALUE), # Dim 1 — 2 entries + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), # Dim 0: 8 entries + AxisDefinition(["param"], Chunking.SINGLE_VALUE), # Dim 1: 2 entries ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) - # Part 2 — pressure-level parameters (levtype=pl): 2 params × 3 levels = 6 entries + # Part 2: pressure-level parameters (levtype=pl): 2 params × 3 levels = 6 entries builder.add_part( { "class": "ea", @@ -240,20 +242,20 @@ dimension grows across the two ``parts``. "levelist": [500, 850, 1000], # 3 levels }, [ - AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), # Dim 0 — must match Part 1 - AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), # Dim 1 — 6 entries + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), # Dim 0: must match Part 1 + AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), # Dim 1: 6 entries ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) - # Dim 1 (param) grows: Part 1 contributes indices 0–1, Part 2 contributes 2–7 + # Dim 1 (param) grows: Part 1 contributes indices 0 to 1, Part 2 contributes 2 to 7 builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) -The resulting array has shape ``(8, 8, N)`` — 8 datetime steps and 8 entries on the -param dimension (2 sfc + 6 pl). +The resulting array has shape ``(8, 8, N)``. That is 8 datetime steps, and 8 entries +on the param dimension (2 sfc + 6 pl). .. code-block:: python @@ -274,8 +276,8 @@ Common Pitfalls --------------- **MARS request ends with a comma** - ``"...,param=167,"`` — the trailing comma causes a parse error. - Omit the comma on the last key–value pair. + ``"...,param=167,"``. The trailing comma causes a parse error. + Omit the comma on the last key-value pair. **Multi-valued keyword not covered by any AxisDefinition** Every MARS keyword that has more than one value **must** appear in exactly one @@ -283,7 +285,7 @@ Common Pitfalls **Wrong array shape** If ``data.shape`` does not match what you expect, check the order of your - ``AxisDefinition`` list — position in the list is the dimension index. + ``AxisDefinition`` list. Position in the list is the dimension index. **Large chunk memory use** ``Chunking.WHOLE_AXIS`` on several axes can produce chunks of many gigabytes. diff --git a/docs/z3fdb/index.rst b/docs/z3fdb/index.rst index 009f0794d..45262b5aa 100644 --- a/docs/z3fdb/index.rst +++ b/docs/z3fdb/index.rst @@ -17,7 +17,14 @@ Z3FDB getting_started dimension_mapping architecture - api + +.. toctree:: + :maxdepth: 1 + :caption: Tutorials: + :hidden: + + tutorials/custom_store_mixed_extractors + tutorials/ensemble_timeseries .. toctree:: :maxdepth: 1 @@ -28,6 +35,14 @@ Z3FDB technical_insights/chunk_access technical_insights/buffer_layout technical_insights/extractor + technical_insights/extractor_backends + +.. toctree:: + :maxdepth: 2 + :caption: Reference: + :hidden: + + api @@ -50,13 +65,13 @@ data is decoded to float32 in memory. When to use Z3FDB ----------------- -Z3FDB is a good fit when a 2–5× slowdown compared to on-disk Zarr is +Z3FDB is a good fit when a 2-5x slowdown compared to on-disk Zarr is acceptable. That trade-off makes sense in two situations: **Prototyping with an existing FDB** You have GRIB data in FDB and want a Zarr interface without first writing everything to disk. Redefining the store is a matter of changing a few lines - of Python — no data copy required. Z3FDB stores are cheap to create because + of Python, with no data copy required. Z3FDB stores are cheap to create because data is only fetched when a chunk is accessed. **Very large datasets** @@ -72,7 +87,11 @@ including how to run the test suite. Next steps ---------- -New to Z3FDB? Start with the step-by-step tutorial: :doc:`getting_started`. +New to Z3FDB? Start with the step-by-step guide: :doc:`getting_started`. + +Ready to explore a realistic example? The +:doc:`tutorials/ensemble_timeseries` tutorial shows how to access ensemble +forecast data with GribJump and plot a temperature time series. For a full reference on how MARS keywords map to Zarr dimensions, chunking strategies, fill values, and multi-part views, see :doc:`dimension_mapping`. diff --git a/docs/z3fdb/installation.rst b/docs/z3fdb/installation.rst index 725a3a817..dd53e5ceb 100644 --- a/docs/z3fdb/installation.rst +++ b/docs/z3fdb/installation.rst @@ -20,6 +20,43 @@ Verify the installation by importing the package in a Python REPL: print(z3fdb) +.. _z3fdb_gribjump_availability: + +Optional: the GribJump Extractor +-------------------------------- + +:class:`~pychunked_data_view.ExtractorType.Grib` works in every installation. +:class:`~pychunked_data_view.ExtractorType.GribJump` does not: it is compiled only when fdb is +configured with + +.. code-block:: bash + + -DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON + +which is **off by default**, so a wheel from PyPI will not normally have it. The feature needs a +bundle build in which gribjump is present as a sibling project. Gribjump itself depends on +fdb5, so it cannot simply be resolved as an ordinary installed dependency. + +The class is importable and constructible either way, so code does not have to branch on how the +wheel was built. Only building a view from it fails: + +.. code-block:: text + + RuntimeError: GribJumpExtractorDefinition: this build has no GribJump support. Rebuild fdb + with -DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON (requires a bundle build providing gribjump). + +To check before you get there: + +.. code-block:: python + + from pychunked_data_view import has_gribjump_extractor + + print(has_gribjump_extractor) # False on a default build + +The two tutorials both use GribJump; everything else in this documentation works without it. + +.. seealso:: :ref:`z3fdb_extractor_backends` for what the two backends do and how to choose. + Running the Tests ----------------- diff --git a/docs/z3fdb/technical_insights/buffer_layout.rst b/docs/z3fdb/technical_insights/buffer_layout.rst index 5ebd03695..83f075dc6 100644 --- a/docs/z3fdb/technical_insights/buffer_layout.rst +++ b/docs/z3fdb/technical_insights/buffer_layout.rst @@ -36,8 +36,11 @@ grid-point values ``num_values``: * - ``FixedSizeChunk(k)`` - ``k`` -The grid-point dimension (``num_values``) is always the trailing dimension and -is never chunked. +The grid-point dimension (``num_values``) is always the trailing dimension. It is a single +chunk covering the whole field unless the part is served by a ``GribJump`` extractor configured +with ``field_chunking``, in which case it is split into equally sized chunks and the formulas +below use that per-chunk size rather than the full field size. See +:ref:`z3fdb_extractor_backends`. Buffer Position Formula ----------------------- @@ -47,24 +50,24 @@ position along each axis ``i`` is: .. code-block:: text - local_i = axis.index(key_i) − partAxisOffset[i] + local_i = axis.index(key_i) - partAxisOffset[i] bufPos_i = local_i + bufferOffset[i] Where: ``axis.index(key_i)`` Zero-based position of the MARS key value returned by FDB within the - ``Part``'s local axis (range: ``0`` to ``axisSize_i − 1``). + ``Part``'s local axis (range: ``0`` to ``axisSize_i - 1``). ``partAxisOffset[i]`` Start of the intersection within the ``Part``'s own local axis:: - partAxisOffset[i] = intersection.lower[i] − partBoundingBox.lower[i] + partAxisOffset[i] = intersection.lower[i] - partBoundingBox.lower[i] ``bufferOffset[i]`` Start of the intersection within the chunk buffer:: - bufferOffset[i] = intersection.lower[i] − chunkBoundingBox.lower[i] + bufferOffset[i] = intersection.lower[i] - chunkBoundingBox.lower[i] ``local_i`` is the zero-based position of the field *within the intersection* along axis ``i``. Adding ``bufferOffset[i]`` shifts it to the correct slot @@ -74,7 +77,7 @@ Flat Buffer Index ----------------- The per-axis positions are combined into a single flat index using C-order -(row-major) arithmetic — the rightmost axis varies fastest: +(row-major) arithmetic. The rightmost axis varies fastest: .. code-block:: text @@ -86,7 +89,7 @@ The per-axis positions are combined into a single flat index using C-order flatIndex = sum(bufPos_i × stride_i for i in 0..N-1) Each position in the flat index corresponds to ``num_values`` consecutive -``float32`` values — the decoded grid-point values for that field. +``float32`` values, the decoded grid-point values for that field. .. note:: diff --git a/docs/z3fdb/technical_insights/chunk_access.rst b/docs/z3fdb/technical_insights/chunk_access.rst index 2c1496883..c5416d9a2 100644 --- a/docs/z3fdb/technical_insights/chunk_access.rst +++ b/docs/z3fdb/technical_insights/chunk_access.rst @@ -10,15 +10,15 @@ When a Zarr chunk is read from a Z3FDB store, the library executes three steps for every ``Part`` (one per :meth:`~z3fdb.SimpleStoreBuilder.add_part` call): -1. **Intersection** — compute the overlap between the requested chunk's +1. **Intersection.** Compute the overlap between the requested chunk's bounding box and the ``Part``'s bounding box. Parts with no overlap are skipped immediately. -2. **FDB retrieve** — issue a sub-request to FDB for exactly the fields +2. **FDB retrieve.** Issue a sub-request to FDB for exactly the fields inside the intersection. Fields outside it are never fetched. -3. **Buffer fill** — decode each returned GRIB field to ``float32`` and write +3. **Buffer fill.** Decode each returned GRIB field to ``float32`` and write it into the correct position in the flat chunk buffer (row-major / C order). -The examples below use the two-part view from :doc:`/z3fdb/dimension_mapping`: +The examples below use the two-part view from :doc:`../dimension_mapping`: **Part A** covers surface parameters (``sfc``, 2 params, axis 1 = [0, 1]) and **Part B** covers pressure-level parameters (``pl``, 4 params, axis 1 = [2, 5]). Both parts share 4 date × time values on axis 0. @@ -51,10 +51,10 @@ region of the buffer. Buffer extent: [4, 6] - Part A — partAxisOffset = [0, 0], bufferOffset = [0, 0] + Part A: partAxisOffset = [0, 0], bufferOffset = [0, 0] The intersection starts at A's local origin and at the buffer corner. - Part B — partAxisOffset = [0, 0], bufferOffset = [0, 2] + Part B: partAxisOffset = [0, 0], bufferOffset = [0, 2] The intersection starts at B's local origin but at buffer column 2, because B begins at global index 2 on axis 1. @@ -80,10 +80,10 @@ With ``SINGLE_VALUE`` every chunk holds exactly one field. Accessing chunk 3 │ │ │ │ │ │ │ └───┴───┴───┴───┴───┴───┘ - Intersection with Part A: empty — skipped. + Intersection with Part A: empty, so it is skipped. Intersection with Part B: axis0 = [1, 1], axis1 = [2, 2] - Part B — partAxisOffset = [1, 0], bufferOffset = [0, 0], bufferExtent = [1, 1] + Part B: partAxisOffset = [1, 0], bufferOffset = [0, 0], bufferExtent = [1, 1] axis 0: partAxisOffset = 1 because the intersection starts at date×time index 1 within Part B's local axis. @@ -95,12 +95,12 @@ With ``SINGLE_VALUE`` every chunk holds exactly one field. Accessing chunk FDB returns one field. Within Part B, axis.index(key) = [1, 0]: - axis 0: local = 1 − 1 = 0, bufPos = 0 + 0 = 0 - axis 1: local = 0 − 0 = 0, bufPos = 0 + 0 = 0 + axis 0: local = 1 - 1 = 0, bufPos = 0 + 0 = 0 + axis 1: local = 0 - 0 = 0, bufPos = 0 + 0 = 0 → written to buffer slot (0, 0) -``FixedSizeChunking`` — cross-part chunk example -------------------------------------------------- +``FixedSizeChunking``, cross-part chunk example +----------------------------------------------- With ``FixedSizeChunk(2)`` on axis 0 and ``FixedSizeChunk(3)`` on axis 1, the chunk grid is 2 × 2. Chunk ``(0, 0)`` covers two date×time steps and the first @@ -129,13 +129,13 @@ three param slots, which straddles the boundary between Part A and Part B. Buffer extent: [2, 3] - Part A — partAxisOffset = [0, 0], bufferOffset = [0, 0] + Part A: partAxisOffset = [0, 0], bufferOffset = [0, 0] Intersection starts at A's local origin and at the buffer corner. - Part B — partAxisOffset = [0, 0], bufferOffset = [0, 2] + Part B: partAxisOffset = [0, 0], bufferOffset = [0, 2] B's local axis1 starts at global index 2, so global [2, 2] maps to local [0, 0]. The intersection lands at buffer column 2 - because 2 − 0 (chunk lower bound) = 2. + because 2 - 0 (chunk lower bound) = 2. Buffer layout (2 rows × 3 columns): @@ -146,26 +146,32 @@ three param slots, which straddles the boundary between Part A and Part B. 1 │ A │ A │ B │ └───┴───┴───┘ -FDB issues two sub-requests — one for Part A, one for Part B. Each field is +FDB issues two sub-requests, one for Part A and one for Part B. Each field is placed using the buffer-position formula. For a field returned by Part A with ``axis.index(key) = [1, 1]`` (second date×time, second sfc param): .. code-block:: text - axis 0: local = 1 − 0 = 1, bufPos = 1 + 0 = 1 - axis 1: local = 1 − 0 = 1, bufPos = 1 + 0 = 1 + axis 0: local = 1 - 0 = 1, bufPos = 1 + 0 = 1 + axis 1: local = 1 - 0 = 1, bufPos = 1 + 0 = 1 → written to buffer slot (1, 1) For a field returned by Part B with ``axis.index(key) = [0, 0]`` (first -date×time, first pl param — which is global param index 2): +date x time, first pl param, which is global param index 2): .. code-block:: text - axis 0: local = 0 − 0 = 0, bufPos = 0 + 0 = 0 - axis 1: local = 0 − 0 = 0, bufPos = 0 + 2 = 2 + axis 0: local = 0 - 0 = 0, bufPos = 0 + 0 = 0 + axis 1: local = 0 - 0 = 0, bufPos = 0 + 2 = 2 → written to buffer slot (0, 2) .. seealso:: :doc:`buffer_layout` for the general buffer-position formula and how the flat buffer index is computed from the per-axis positions. + +.. seealso:: + + :ref:`z3fdb_extractor_backends` for what the two extractor backends do inside step three, + and for the concurrency guarantee: each extractor serialises its own ``extractInto`` calls, + so the unit of parallelism is the part rather than the chunk. diff --git a/docs/z3fdb/technical_insights/dev_setup.rst b/docs/z3fdb/technical_insights/dev_setup.rst index b69f02636..8b0196fc9 100644 --- a/docs/z3fdb/technical_insights/dev_setup.rst +++ b/docs/z3fdb/technical_insights/dev_setup.rst @@ -45,7 +45,7 @@ A minimal bundle for Z3FDB development looks like this: Place this file in a ``bundle`` directory alongside the cloned source trees. cmake needs a virtual environment with the ``build`` package present to invoke the Python wheel builder at the end of the build. ``pybind11-stubgen`` is also -required — the build generates ``.pyi`` stub files for the C++ extension during +required. The build generates ``.pyi`` stub files for the C++ extension during the cmake build. The following block can be pasted directly into a shell: .. code-block:: bash @@ -188,14 +188,14 @@ FAQ export FDB5_HOME=/path/to/build -**pytest crashes with** ``abort`` **— wrong binaries picked up** +**pytest crashes with** ``abort``**: wrong binaries picked up** .. code-block:: text [1] 39253 abort pytest z3fdb A hard crash (SIGABRT or similar) usually means that a library loaded at - runtime does not match the one it was compiled against — for example, a + runtime does not match the one it was compiled against. For example, a system-installed ``libfdb5`` or ``libeccodes`` is found instead of the one from the build tree. diff --git a/docs/z3fdb/technical_insights/extractor.rst b/docs/z3fdb/technical_insights/extractor.rst index f595cebc1..6383d74cf 100644 --- a/docs/z3fdb/technical_insights/extractor.rst +++ b/docs/z3fdb/technical_insights/extractor.rst @@ -14,13 +14,8 @@ buffer. Ownership Model --------------- -Extractors are held via ``std::shared_ptr``. This enables two -important properties: - -**Sharing across parts** - Multiple ``ViewPart`` objects can share a single ``Extractor`` instance — - for example, two parts reading from the same FDB store can reuse the same - open handle without duplicating it. +Extractors are held via ``std::unique_ptr``. Each ``ViewPart`` +owns its extractor exclusively. There is no sharing between parts. **Tied lifetime** Each extractor's lifetime is bound to the ``ChunkedDataView`` that owns it. @@ -32,31 +27,64 @@ Why Extractors Are Non-Copyable Concrete extractor implementations are **stateful and non-copyable**. An FDB-backed extractor owns an open FDB connection handle. Copying such a handle -would duplicate a live network or file-system connection, which is unsafe — -both copies would race on the same underlying state. +would duplicate a live network or file-system connection, which is unsafe. +Both copies would race on the same underlying state. + +``std::unique_ptr`` makes this ownership explicit: each extractor belongs to +exactly one ``ViewPart``, and is destroyed exactly once when the view is +destroyed. + +ExtractorDefinition Factory +---------------------------- + +Extractors are not constructed directly in ``addPart``. Instead, each part +records an ``ExtractorDefinition``, a lightweight configuration object that +implements a single factory method: + +.. code-block:: cpp + + virtual std::unique_ptr buildExtractor( + const metkit::mars::MarsRequest& request) const = 0; -``std::shared_ptr`` provides shared ownership without copying: all parts that -reference the same extractor share one instance, and the instance is destroyed -exactly once when the last shared reference is dropped. +``ChunkedDataViewBuilder::build()`` calls ``buildExtractor(request)`` once per +part after the MARS request string has been parsed. This defers FDB and +GribJump initialisation to ``build()`` time, so any configuration errors are +raised there rather than in ``addPart``. + +``ChunkedDataViewBuilder`` itself is non-copyable (its copy constructor and +copy-assignment operator are explicitly deleted) because it stores +``std::unique_ptr`` objects that cannot be duplicated. Extractor Interface -------------------- -All extractors implement two methods called by the core during a chunk access: +All extractors implement one method called by the core during a chunk access, +and expose a ``DataLayout`` computed eagerly in the constructor: + +``DataLayout layout_`` + Set during construction by issuing a sample FDB retrieve for the part's + MARS request. Records the field's grid-point count, bytes-per-value, and + chunk shape so that ``ChunkedDataViewBuilder::build()`` can validate axis + compatibility before committing to the view. + +``extractInto(part, chunkBB, intersectionBB, ptr, len)`` + Called during each chunk access to retrieve the fields matching the part's + MARS request from FDB, decode them (or partially decode), and write the + float32 values into the provided buffer at the correct offset. + +.. seealso:: -``layout(request)`` - Called once per part during :meth:`~z3fdb.SimpleStoreBuilder.build` to - probe the field layout — grid size and axis ordering — without reading - actual data values. + :doc:`chunk_access` for how the extractor's ``extractInto`` method is + invoked as part of the three-step chunk-access pipeline. -``extract(request, buffer, offset)`` - Called during each chunk access to retrieve the fields matching *request* - from FDB, decode them, and write them into *buffer* at *offset*. +The Two Backends +---------------- -The only extractor currently shipped is ``GribExtractor``, which reads GRIB -messages from FDB and decodes them to ``float32`` via eccodes. +``GribExtractor`` and ``GribJumpExtractor`` are documented together in +:doc:`extractor_backends`: what each one does, how to choose between them, their configuration, +the constraints they impose on a view, and which builds provide them. .. seealso:: - :doc:`chunk_access` for how the extractor's ``extract`` method is invoked - as part of the three-step chunk-access pipeline. + :ref:`z3fdb_extractor_backends` for the backends themselves, and + :ref:`tutorial_custom_store_mixed_extractors` for a store that uses both. diff --git a/docs/z3fdb/technical_insights/extractor_backends.rst b/docs/z3fdb/technical_insights/extractor_backends.rst new file mode 100644 index 000000000..cd3f83536 --- /dev/null +++ b/docs/z3fdb/technical_insights/extractor_backends.rst @@ -0,0 +1,173 @@ +.. SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +.. SPDX-License-Identifier: Apache-2.0 + +.. _z3fdb_extractor_backends: + +GRIB and GribJump Extractors +============================ + +Z3FDB ships two extractor backends. Both implement the same interface and both read the same +data from the same FDB; the only difference is how the values leave a GRIB message. This page +covers what each one does, how to choose, and the constraints each imposes on a view. + +For the machinery around them, see :doc:`extractor`. That covers ownership, the factory, and +when the layout is established. + +.. contents:: On this page + :local: + :depth: 1 + +Choosing between them +--------------------- + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - + - ``ExtractorType.Grib`` + - ``ExtractorType.GribJump`` + * - How values are read + - Full eccodes decode of the message + - Jumps to the values inside the message + * - Grid-point axis + - Always one chunk: the whole field + - Optionally split, via ``field_chunking`` + * - Best for + - Reading whole fields + - Reading a sub-range of each field, repeatedly + * - Availability + - Always + - Opt-in build feature (see `Availability`_) + +The rule of thumb: if you read whole fields, ``Grib`` is the simpler choice and there is nothing +to gain from ``GribJump``. If you slice into the grid-point axis, ``GribJump`` avoids decoding the values you are +going to discard. A time series at one location, or a region out of a global field, are the +typical cases, and ``field_chunking`` is what lets zarr fetch only the piece you asked for. + +:ref:`tutorial_custom_store_mixed_extractors` builds one store containing both, which is a +reasonable pattern when different consumers of the same data have different access patterns. + +GribExtractor +------------- + +Reads GRIB messages from FDB and decodes them to ``float32`` via eccodes. The entire field is +decoded on every chunk access, so the implicit grid-point dimension is always a single chunk +covering the whole field. + +Where a message carries a bitmap, the masked grid points are replaced with the view's fill +value (see :meth:`~z3fdb.SimpleStoreBuilder.fill_missing_value`). + +GribJumpExtractor +----------------- + +Uses the GribJump library to read grid-point values without performing a full GRIB decode. + +Field enumeration still goes through FDB: the extractor inspects FDB to learn which fields +match the part's sub-request and in what order, then asks GribJump for the values. The order +matters, because it is what maps each field onto its slot in the chunk buffer. + +**Field chunking.** ``field_chunking`` subdivides the implicit grid-point dimension into equally +sized zarr chunks:: + + # 5248 grid points, split into 4 chunks of 1312 + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=1312)) + +Each ``extractInto`` call derives the range it needs from the chunk it was handed; nothing is +cached between calls, and there is no separate sub-range selection. + +**Missing values.** GribJump returns a bitmask alongside the values, in which a *set* bit means +the point is valid and a clear bit means it is missing. Masked points are written as the view's +fill value. + +Configuration +------------- + +``fdb_config`` (both backends) + Path to an FDB configuration YAML. When left as ``None``, the extractor inherits the path + given to the store builder; if that is also unset, FDB resolves its own configuration from + ``FDB5_CONFIG`` / ``FDB_HOME``. + +``gribjump_config`` (GribJump only) + Path to a GribJump configuration YAML. When set, the binding exports it as + ``GRIBJUMP_CONFIG_FILE`` before constructing the GribJump object; when unset, that + environment variable is used as-is. + +``field_chunking`` (GribJump only) + Chunking of the implicit grid-point dimension. Defaults to a single chunk covering the whole + field. That is the only setting which can be mixed with a ``Grib`` part, see below. + +An extractor configuration object is *copied* when it is handed to ``add_part``, so one object +can be passed to as many parts and as many builders as you like, and the ``fdb_config`` a +builder fills in never leaks back into your object. + +Constraints +----------- + +Three rules are enforced when the view is built. All three surface as ``RuntimeError`` in +Python, carrying the message quoted below. + +**A field chunk size must divide the grid exactly.** The grid-point dimension is the one +dimension a zarr array cannot leave ragged, so 5248 points can be split into chunks of 1312 or +2624, but not 1000: + +.. code-block:: text + + GribJumpExtractor: field chunk size 1000 does not evenly divide the window size 5248. + +A size of zero is rejected earlier still, by ``Chunking.FixedSizeChunk`` itself. + +**Every part of a view must cover the same grid.** The grid-point dimension is never the +extension axis. Like every other non-extension axis, all parts have to agree on it: + +.. code-block:: text + + ChunkedDataViewBuilder::build: part 1 has 2048 grid points but part 0 has 5248. The + grid-point dimension is never the extension axis, so every part must cover the same grid; + a view cannot have a ragged last dimension. + +**A GribJump part mixed with a Grib part must use whole-axis field chunking.** ``Grib`` always +returns the whole field, so its chunk on the grid-point axis is the full grid; a ``GribJump`` +part using ``FixedSizeChunk`` would write a smaller block into a buffer laid out for a larger +one: + +.. code-block:: text + + ChunkedDataViewBuilder::build: part 1 splits the grid-point dimension into chunks of 1312 + values but part 0 uses 5248. All parts must agree on the field chunking, so a GribJump part + mixed with a Grib part has to use the default WholeAxisChunking. + +Availability +------------ + +``GribExtractor`` is always present. ``GribJumpExtractor`` is built only when fdb is configured +with ``-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON``, which is **off by default**. See +:ref:`z3fdb_installation`. + +``ExtractorType.GribJump`` is importable and constructible in every build, so your code does not +have to branch on how the wheel was compiled. Only building a view from it fails: + +.. code-block:: python + + from pychunked_data_view import has_gribjump_extractor + + if not has_gribjump_extractor: + ... # fall back to ExtractorType.Grib() + +Concurrency +----------- + +Each extractor serialises its own ``extractInto`` calls with a mutex: it drives a shared FDB +handle (and, for GribJump, a shared GribJump object) while the Python layer has released the +GIL, so concurrent chunk reads would otherwise race. + +The unit of parallelism is therefore the *part*, not the chunk. Two threads reading chunks that +are served by the same part will take turns; two threads reading chunks served by different +parts proceed at the same time. Worth knowing when sizing a dask cluster against a +single-part store. + +.. seealso:: + + :doc:`extractor` for the ownership model and the ``ExtractorDefinition`` factory, + :doc:`chunk_access` for where ``extractInto`` sits in the chunk-access pipeline, and + :doc:`../api` for the full configuration reference. diff --git a/docs/z3fdb/tutorials/custom_store_mixed_extractors.rst b/docs/z3fdb/tutorials/custom_store_mixed_extractors.rst new file mode 100644 index 000000000..aade4d889 --- /dev/null +++ b/docs/z3fdb/tutorials/custom_store_mixed_extractors.rst @@ -0,0 +1,200 @@ +.. SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +.. SPDX-License-Identifier: Apache-2.0 + +.. _tutorial_custom_store_mixed_extractors: + +Mixed-Extractor Custom Store +============================= + +This tutorial shows how to use :class:`~z3fdb.CustomStoreBuilder` to create +a single Zarr store that contains multiple named arrays, each backed by a +different extraction strategy. You will build a store with two arrays holding +the same 10 m u-wind field. One uses the standard ``Grib`` extractor for +efficient full-field access, the other uses ``GribJump``. You then plot the field +as a global map. + +.. note:: + + This tutorial uses ``ExtractorType.GribJump``, which is built only when fdb is configured + with ``-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON``, off by default. On a build without it, + ``build()`` raises and names the flag. Check with + ``pychunked_data_view.has_gribjump_extractor``, and see :ref:`z3fdb_gribjump_availability`. + +.. contents:: On this page + :local: + :depth: 1 + +When to use CustomStoreBuilder +------------------------------- + +:class:`~z3fdb.SimpleStoreBuilder` always places a single array at the store +root. :class:`~z3fdb.CustomStoreBuilder` lifts that restriction: you can +register any number of arrays at named paths inside one store and open them +as a Zarr group. + +This is useful when: + +* You want to expose multiple parameters or level types through one store + object. +* Different arrays in the store benefit from different extractors. For + example, surface fields used for map visualisation work well with a + ``Grib`` extractor (one FDB retrieve returns the full field), while the + same field accessed at individual grid points is better served by + ``GribJump``. + +The MARS Request +----------------- + +The example below retrieves the 10 m u-wind (param ``165.128``) from a +single ensemble member at analysis time: + +.. code-block:: python + + _REQUEST = { + "class": "od", + "date": "20260818", + "expver": "0001", + "levtype": "sfc", + "domain": "g", + "stream": "enfo", + "type": "pf", + "number": [1], # one ensemble member + "param": ["165.128"], # 10 m u-wind + "step": [0], # analysis step + "time": ["00:00:00"], + } + +With a single value on every axis the resulting array has shape +``(1, 1, 1, 1, N)`` where ``N`` is the number of grid points in the GRIB +field. + +Building the Custom Store +-------------------------- + +.. code-block:: python + + import zarr + import numpy as np + import matplotlib.pyplot as plt + + zarr.config.set({"async.concurrency": 1, "threading.max_workers": 1}) + + from z3fdb import AxisDefinition, Chunking, ExtractorType, CustomStoreBuilder + + _AXES = [ + AxisDefinition(keys=["date", "time"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["number"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["step"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["param"], chunking=Chunking.SINGLE_VALUE), + ] + + builder = CustomStoreBuilder() + + # Array 1 - full-field Grib extraction: one FDB retrieve gives the complete field. + # Best choice when you need all grid points (e.g. for maps). + builder.add_part("grib/u10", _REQUEST, axes=_AXES, extractor=ExtractorType.Grib()) + + # Array 2 - GribJump extraction with one grid point per chunk. + # Best choice for sparse access (e.g. time series at a station). + builder.add_part("gribjump/u10", _REQUEST, axes=_AXES, + extractor=ExtractorType.GribJump( + field_chunking=Chunking.FixedSizeChunk(chunk_shape=1) + )) + + store = builder.build() + +``add_part("grib/u10", ...)`` + Registers an array at the path ``grib/u10`` inside the store. The path + follows zarr conventions: ``/``-separated segments, leading slash + optional. Here ``grib`` becomes an intermediate group and ``u10`` is the + array name. + +``ExtractorType.Grib()`` + Standard extraction: one FDB retrieve per chunk, returning a full GRIB + field decoded to float32. With ``SINGLE_VALUE`` chunking on every explicit + axis, each chunk holds one complete field. That is ideal for reading all grid + points at once. + +``ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=1))`` + GribJump jumps to the bytes inside each GRIB message that correspond to + the requested grid points, without decoding the rest. ``FixedSizeChunk(1)`` + makes every grid point its own chunk: reading ``arr[..., k]`` extracts + exactly the value at grid point ``k``. This is expensive for full-field + access but efficient for sparse access across many time steps. + +Opening the Store +----------------- + +Because the store contains multiple named arrays, open it as a **group**, not +an array: + +.. code-block:: python + + grp = zarr.open_group(store, mode="r") + + # Access each array by its registered path + grib_arr = grp["grib/u10"] + gribjump_arr = grp["gribjump/u10"] + + print("shape :", grib_arr.shape) # (1, 1, 1, 1, N) + print("chunks:", grib_arr.chunks) # (1, 1, 1, 1, N) - whole field per chunk + +.. code-block:: text + + grp["grib/u10"][dt, member, step, param, grid_point] + ^ ^ ^ ^ ^ + | | | | grid point [implicit, size N] + | | | param [size 1] + | | step [size 1] + | member [size 1] + date x time [size 1] + +**No data is fetched from FDB until you index.** Building and opening the +store only validates the MARS request and determines the field layout. + +Plotting the 10 m u-wind Map +------------------------------ + +Read the full field from the ``Grib``-backed array and plot it as a global +map. Z3FDB returns a flat 1-D float32 array of ``N`` grid-point values; +reshaping it to 2-D requires the latitude and longitude coordinates for your +grid. These can be obtained from the GRIB message metadata via ``eccodes`` +or ``cfgrib``. + +.. code-block:: python + + # Retrieve the full field - shape (1, 1, 1, 1, N), index to (N,) + u10_flat = grib_arr[0, 0, 0, 0, :] + + # Obtain lat/lon coordinates for your grid. + # The example below assumes a regular 0.25 deg global grid (1440 x 721). + # Replace nlat, nlon, lats, lons with values matching your actual grid. + nlat, nlon = 721, 1440 + lats = np.linspace(90, -90, nlat) + lons = np.linspace(0, 360, nlon, endpoint=False) + u10_2d = u10_flat.reshape(nlat, nlon) + + fig, ax = plt.subplots(figsize=(12, 5)) + img = ax.pcolormesh(lons, lats, u10_2d, cmap="RdBu_r", shading="auto") + fig.colorbar(img, ax=ax, label="10 m u-wind (m/s)") + ax.set_xlabel("Longitude (deg)") + ax.set_ylabel("Latitude (deg)") + ax.set_title(f"10 m u-wind, {_REQUEST['date']} {_REQUEST['time'][0]}, step 0, member 1") + plt.tight_layout() + plt.savefig("u10_map.png", dpi=150) + plt.show() + +.. note:: + + For non-regular grids (e.g. the ECMWF O1280 reduced Gaussian grid) you + cannot reshape to a simple 2-D array. In that case retrieve the + ``(lat, lon)`` coordinates for each grid point from an ``eccodes``-opened + sample field and pass them to :func:`matplotlib.pyplot.tricontourf`. + +Next Steps +---------- + +* :doc:`ensemble_timeseries` for using ``GribJump`` with ``FixedSizeChunk(1)`` + across a full ensemble and plot time series. +* :doc:`../dimension_mapping` for the complete reference on axis mapping, + chunking strategies, and multi-part views. diff --git a/docs/z3fdb/tutorials/ensemble_timeseries.rst b/docs/z3fdb/tutorials/ensemble_timeseries.rst new file mode 100644 index 000000000..e26a37c0c --- /dev/null +++ b/docs/z3fdb/tutorials/ensemble_timeseries.rst @@ -0,0 +1,202 @@ +.. SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +.. SPDX-License-Identifier: Apache-2.0 + +.. _tutorial_ensemble_timeseries: + +Ensemble Forecast Time Series with GribJump +=========================================== + +This tutorial walks you through accessing ensemble forecast data from FDB +as a Zarr array. You will build a Z3FDB store backed by 50 ensemble members, +plot the 2 m temperature time series for every member, and overlay the +ensemble mean. + +.. note:: + + This tutorial uses ``ExtractorType.GribJump``, which is built only when fdb is configured + with ``-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON``, off by default. On a build without it, + ``build()`` raises and names the flag. Check with + ``pychunked_data_view.has_gribjump_extractor``, and see :ref:`z3fdb_gribjump_availability`. + +.. contents:: On this page + :local: + :depth: 1 + +When to use GribJump +-------------------- + +The standard :class:`~z3fdb.ExtractorType.Grib` extractor decodes an entire +GRIB field, typically several million grid-point values, just to return the +portion you requested. For time series work where you need values at **one +or a few grid points**, this is wasteful. + +:class:`~z3fdb.ExtractorType.GribJump` solves this by jumping directly to the +bytes inside each GRIB message that correspond to the requested grid points, +without decoding the rest. The trade-off: each grid-point chunk triggers an +individual GribJump lookup, so this approach shines for sparse access (a few +points across many time steps) but is slower than ``Grib`` for dense access +(full fields). + +The MARS Request +---------------- + +The example retrieves 2 m temperature and three wind/cloud parameters from +an ECMWF ensemble forecast: + +.. code-block:: python + + _STEPS = list(range(0, 91)) + list(range(93, 145, 3)) + # 0-90 h (hourly), 93-144 h (3-hourly) -> 109 steps total + + _REQUEST = { + "class": "od", + "date": "20260818", + "expver": "0001", + "levtype": "sfc", + "domain": "g", + "stream": "enfo", # ensemble forecast stream + "type": "pf", # perturbed forecast members + "number": list(range(1, 51)), # 50 members + "param": ["167.128", "165.128", "166.128", "164.128"], + # 2m T, 10m u, 10m v, total cloud cover + "step": _STEPS, + "time": ["00:00:00", "06:00:00", "12:00:00", "18:00:00"], + } + +Key points: + +``stream=enfo``, ``type=pf`` + Selects the ensemble stream; ``pf`` (perturbed forecast) gives you the 50 + individual members identified by ``number``. + +``number`` + Each integer from 1 to 50 identifies one ensemble member. + +``step`` + The forecast lead time in hours. Here we combine hourly output for the + first 90 hours with 3-hourly output from hour 93 to 144, giving 109 steps. + +``time`` + The four daily analysis times that serve as initialisation times. + +Building the Store +------------------ + +.. code-block:: python + + import zarr + import numpy as np + import matplotlib.pyplot as plt + + zarr.config.set({"async.concurrency": 1, "threading.max_workers": 1}) + + from z3fdb import AxisDefinition, Chunking, ExtractorType, SimpleStoreBuilder + + builder = SimpleStoreBuilder() + builder.add_part( + _REQUEST, + axes=[ + AxisDefinition(keys=["date", "time"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["number"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["step"], chunking=Chunking.SINGLE_VALUE), + AxisDefinition(keys=["param"], chunking=Chunking.SINGLE_VALUE), + ], + extractor=ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=1)), + ) + store = builder.build() + arr = zarr.open_array(store, mode="r") + +``AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE)`` + Combines date and time into a single dimension. With one date and four + init times this gives 4 entries, ordered as time cycles within each date. + +``ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=1))`` + Uses GribJump as the extraction backend. ``FixedSizeChunk(chunk_shape=1)`` + makes each grid point its own chunk: accessing ``arr[..., k]`` retrieves + exactly the value at grid point ``k`` without decoding the full field. + +Array Shape +----------- + +.. code-block:: python + + print("shape :", arr.shape) # (4, 50, 109, 4, N) + print("chunks:", arr.chunks) # (1, 1, 1, 1, 1) + +.. code-block:: text + + arr[dt, member, step, param, grid_point] + ^ ^ ^ ^ ^ + | | | | grid point index [implicit, size N] + | | | param index [size 4] + | | step index [size 109] + | member index [size 50] + date x init time [size 4] + +The **implicit final dimension** always holds the decoded grid-point values. +Its size ``N`` is determined by the GRIB grid (for a global O1280 grid, +``N ~ 6 600 000``). + +**No data is fetched from FDB until you index.** Building the store is +cheap. It probes one representative field to determine the layout, but does +not retrieve the full dataset. + +Plotting the Temperature Time Series +------------------------------------- + +The block below retrieves the 2 m temperature time series for all 50 ensemble +members at a single grid point, computes the ensemble mean, and produces a +plot. + +.. code-block:: python + + T2M = 0 # "167.128" is first in the param list + INIT_TIME = 0 # 00 UTC initialisation + GRID_POINT = 1_000_000 # replace with any valid grid-point index + + # Fetch all members at once - shape (50, 109) + all_members = arr[INIT_TIME, :, :, T2M, GRID_POINT] + ensemble_mean = all_members.mean(axis=0) # shape (109,) + + print(f"Ensemble-mean 2m temperature: {ensemble_mean.mean():.2f} K") + + fig, ax = plt.subplots(figsize=(10, 4)) + + # Individual members - thin, semi-transparent + for member_ts in all_members: + ax.plot(_STEPS, member_ts, color="steelblue", alpha=0.2, linewidth=0.7) + + # Ensemble mean - bold + ax.plot( + _STEPS, ensemble_mean, + color="darkred", linewidth=2, + label=f"Ensemble mean ({ensemble_mean.mean():.1f} K)", + ) + + ax.set_xlabel("Forecast step (hours)") + ax.set_ylabel("2m Temperature (K)") + ax.set_title("2m Temperature, 50 ensemble members and mean\n" + f"Grid point {GRID_POINT}, init {_REQUEST['date']} {_REQUEST['time'][INIT_TIME]}") + ax.legend() + plt.tight_layout() + plt.savefig("t2m_timeseries.png", dpi=150) + plt.show() + +The statement ``arr[INIT_TIME, :, :, T2M, GRID_POINT]`` is a single zarr +read that triggers 50 x 109 = 5 450 GribJump extractions, one per +(member, step) combination. It returns a ``(50, 109)`` NumPy array. + +.. note:: + + Accessing many individual grid-point chunks in sequence can be slow for + large ensemble x step combinations. If you need values at many grid points, + consider increasing ``chunk_shape`` in ``FixedSizeChunk`` to batch multiple + grid points into one GribJump call. + +Next Steps +---------- + +* :doc:`../dimension_mapping` for the full reference on axis mapping, chunking + strategies, fill values, and multi-part views. +* :doc:`../getting_started` for an introduction to ``SimpleStoreBuilder`` + covering surface and pressure-level data in a single array. diff --git a/python/fdb5lib/buildconfig b/python/fdb5lib/buildconfig index 6b37c2a8d..07c530f80 100644 --- a/python/fdb5lib/buildconfig +++ b/python/fdb5lib/buildconfig @@ -14,7 +14,7 @@ NAME="fdb" # NOTE zarr interface is dependent on python 3.11+ -- but we dont activate the venv at the time this is sourced, that's why we rely on PYVERSION. Remove this whole part around October 2026 when 3.10 goes EoL # if [ "True" = "$(python -c 'import sys; print(sys.version_info[0:2] >= (3, 11))')" ] ; then if [ "3.10" != "$PYVERSION" ] ; then - ZARR_IFACE="-DENABLE_PYTHON_ZARR_INTERFACE=ON" + ZARR_IFACE="-DENABLE_PYTHON_ZARR_INTERFACE=ON -DENABLE_ZARR_GRIBJUMP_EXTRACTOR=OFF" else >&2 echo "Not enabling zarr interface because python version is $PYVERSION" ZARR_IFACE="" diff --git a/python/z3fdb-build-requirements.txt b/python/z3fdb-build-requirements.txt new file mode 100644 index 000000000..72b54e265 --- /dev/null +++ b/python/z3fdb-build-requirements.txt @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +# +# Build-time dependencies for the z3fdb wheel. NOT runtime dependencies: those live in +# z3fdb-requirements.txt, which cmake/z3fdb_setup.py.in reads into install_requires. +# +# Consumed by the cmake targets in src/CMakeLists.txt via +# uvx --with-requirements python/z3fdb-build-requirements.txt ... +# so that a CI runner shipping nothing but a Python interpreter can still build the wheel. + +# Wheel construction. +# +# NOTE: while `build` runs with isolation, it provisions setuptools and wheel itself in the +# environment it creates per build, so those two are not strictly required here. They are listed +# because they are genuine build dependencies of this project (z3fdb_setup.py.in imports +# wheel.bdist_wheel directly), and so the set stays complete if `--no-isolation` is ever used. +build +setuptools +wheel + +# Stub generation for the pybind11 extension module. +pybind11-stubgen + +# pybind11-stubgen imports chunked_data_view_bindings to introspect it, so whatever that +# import needs has to be available too. These are also runtime dependencies of the wheel; +# they appear here because stub generation cannot import the module without them. +# chunked_data_view_bindings/__init__.py -> findlibs.load("fdb5") +# the extension registers py::array_t types -> numpy +findlibs>=0.1.2 +numpy diff --git a/python/z3fdb-requirements.txt b/python/z3fdb-requirements.txt new file mode 100644 index 000000000..60eef59ed --- /dev/null +++ b/python/z3fdb-requirements.txt @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +# +# Runtime dependencies of the z3fdb wheel. Read into install_requires by +# cmake/z3fdb_setup.py.in, which stages a copy next to the generated setup.py. +# +# Build-time tooling belongs in z3fdb-build-requirements.txt, not here. Anything listed here is +# installed by every `pip install z3fdb`. +# +# This is the complete set of third-party imports across z3fdb, pychunked_data_view and +# chunked_data_view_bindings. The fdb5lib pin is added separately by setup.py, because it only +# applies when VERSION_SUFFIX is set on the CD builds. + +# Zarr store implementation. z3fdb._internal.zarr subclasses zarr.abc.store.Store and uses +# zarr.core.buffer, so it is bound to the v3 API. +zarr~=3.1.6 + +# Chunk buffers are returned as float32 arrays. +numpy + +# chunked_data_view_bindings/__init__.py calls findlibs.load("fdb5") to locate libfdb5 +# before importing the extension module. +findlibs>=0.1.2 diff --git a/python/z3fdb-testing-requirements.txt b/python/z3fdb-testing-requirements.txt new file mode 100644 index 000000000..144f0eeaa --- /dev/null +++ b/python/z3fdb-testing-requirements.txt @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +# +# Dependencies for running the z3fdb and pychunked_data_view test suites. +# +# The runtime dependencies in z3fdb-requirements.txt are needed as well; install both. Build +# tooling lives in z3fdb-build-requirements.txt and is not needed to run the tests against an +# already-built package. + +# Test runner. pytest-asyncio drives the async parts of the zarr store API, exercised by +# tests/z3fdb/interface/zarr_interface_conformity. +pytest +pytest-asyncio + +# Test data. The fixtures in tests/conftest.py build GRIB messages with eccodes and write the +# FDB config as YAML. +eccodes +PyYAML +GitPython + +# tests/z3fdb/interface/zarr_interface_conformity/_mocks.py uses @override, which is only in +# typing from Python 3.12. +typing_extensions + +# Optional consumers, pulled in through pytest.importorskip. Without them +# tests/z3fdb/interface/user_tests/test_dask_access.py and test_xarray_access.py skip silently +# rather than fail, so a suite run without these covers less than it appears to. +dask +xarray diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index bae285628..ca3ab56c3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -8,19 +8,91 @@ add_subdirectory(chunked_data_view) if (HAVE_PYTHON_ZARR_INTERFACE) # We create the complete python package layout at this location. - # This allows us to run python wheel creation at this path and + # This allows us to run python wheel creation at this path and # to put this path on the PYTHONPATH to allow direct use of # z3fdb, e.g. for testing or local exploration. set(Z3FDB_STAGING "${CMAKE_BINARY_DIR}/z3fdb-python-package-staging") - file(MAKE_DIRECTORY "${Z3FDB_STAGING}") - file(CREATE_LINK - "${CMAKE_CURRENT_SOURCE_DIR}/z3fdb" - "${Z3FDB_STAGING}/z3fdb" SYMBOLIC + + # Python build tooling is run through uvx rather than the build machine's interpreter, so a + # CI runner that ships only an interpreter can still build the wheel. uvx resolves each tool + # into an ephemeral environment; build-requirements.txt lists what has to be in it. + set(Z3FDB_BUILD_REQUIREMENTS "${CMAKE_CURRENT_SOURCE_DIR}/../python/z3fdb-build-requirements.txt") + if(NOT EXISTS "${Z3FDB_BUILD_REQUIREMENTS}") + message(FATAL_ERROR "Missing ${Z3FDB_BUILD_REQUIREMENTS}, required to build the z3fdb wheel.") + endif() + + # Runtime dependencies of the wheel. setup.py reads this into install_requires, so it is + # staged next to the generated setup.py below. Kept separate from the build requirements so + # that build tooling never reaches the wheel metadata. + set(Z3FDB_RUNTIME_REQUIREMENTS "${CMAKE_CURRENT_SOURCE_DIR}/../python/z3fdb-requirements.txt") + if(NOT EXISTS "${Z3FDB_RUNTIME_REQUIREMENTS}") + message(FATAL_ERROR "Missing ${Z3FDB_RUNTIME_REQUIREMENTS}, required to build the z3fdb wheel.") + endif() + + find_program(UVX_EXECUTABLE NAMES uvx + DOC "uv tool runner (uvx), used to run the Python build tooling") + if(NOT UVX_EXECUTABLE) + message(FATAL_ERROR + "uvx not found, but ENABLE_PYTHON_ZARR_INTERFACE=ON needs it to build the wheel. " + "Install uv (https://docs.astral.sh/uv/) or point -DUVX_EXECUTABLE=/path/to/uvx at it.") + endif() + + # Create all staging subdirectories at configure time so build-time copy + # commands do not need to worry about missing parent directories. + file(MAKE_DIRECTORY + "${Z3FDB_STAGING}" + "${Z3FDB_STAGING}/z3fdb" + "${Z3FDB_STAGING}/z3fdb/_internal" + "${Z3FDB_STAGING}/pychunked_data_view" + "${Z3FDB_STAGING}/chunked_data_view_bindings" + "${Z3FDB_STAGING}/chunked_data_view_bindings-stubs" ) - file(CREATE_LINK - "${CMAKE_CURRENT_SOURCE_DIR}/pychunked_data_view" - "${Z3FDB_STAGING}/pychunked_data_view" SYMBOLIC + + # ---- Per-file copy commands (tracked: re-copy when source changes) ---- + # Each add_custom_command pairs one source file with its staging destination. + # cmake re-runs the copy whenever the source is newer than the destination. + set(_z3fdb_py_sources + z3fdb/__init__.py + z3fdb/simple_store_builder.py + z3fdb/custom_store_builder.py + z3fdb/z3fdb_error.py + z3fdb/_internal/__init__.py + z3fdb/_internal/zarr.py + pychunked_data_view/__init__.py + pychunked_data_view/chunked_data_view.py + pychunked_data_view/exceptions.py + chunked_data_view_bindings/__init__.py ) + + # A module that exists in src/ but is missing from the list above is silently absent from + # the staging tree and the wheel: nothing fails, the import just goes missing at runtime. + # Fail at configure time instead. + file(GLOB_RECURSE _z3fdb_py_found + RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" + "z3fdb/*.py" + "pychunked_data_view/*.py" + "chunked_data_view_bindings/*.py") + foreach(_found IN LISTS _z3fdb_py_found) + if(NOT "${_found}" IN_LIST _z3fdb_py_sources) + message(FATAL_ERROR + "${_found} is not listed in _z3fdb_py_sources (src/CMakeLists.txt), so it " + "would never reach the z3fdb staging tree or the wheel. Add it there.") + endif() + endforeach() + + set(_z3fdb_staged_py_files) + foreach(_rel IN LISTS _z3fdb_py_sources) + set(_src "${CMAKE_CURRENT_SOURCE_DIR}/${_rel}") + set(_dst "${Z3FDB_STAGING}/${_rel}") + add_custom_command( + OUTPUT "${_dst}" + COMMAND ${CMAKE_COMMAND} -E copy "${_src}" "${_dst}" + DEPENDS "${_src}" + COMMENT "Copying ${_rel} to staging..." + ) + list(APPEND _z3fdb_staged_py_files "${_dst}") + endforeach() + # Copy README.md and LICENSE at build time so changes are picked up # without needing to re-run cmake configure add_custom_command( @@ -39,6 +111,7 @@ if (HAVE_PYTHON_ZARR_INTERFACE) DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/../LICENSE" COMMENT "Copying LICENSE to staging..." ) + configure_file( ${CMAKE_CURRENT_SOURCE_DIR}/../cmake/z3fdb_setup.py.in ${Z3FDB_STAGING}/setup.py @@ -49,64 +122,75 @@ if (HAVE_PYTHON_ZARR_INTERFACE) ${Z3FDB_STAGING}/setup.cfg @ONLY ) - add_subdirectory(chunked_data_view_bindings) - set(_z3fdb_package_files - "z3fdb/simple_store_builder.py" - "z3fdb/z3fdb_error.py" - "z3fdb/__init__.py" - "z3fdb/_internal/zarr.py" - "pychunked_data_view/__init__.py" - "pychunked_data_view/chunked_data_view.py" - "pychunked_data_view/exceptions.py" - "chunked_data_view_bindings/__init__.py" - ) - list(APPEND _z3fdb_package_files - "${Z3FDB_STAGING}/README.md" - "${Z3FDB_STAGING}/LICENSE" + + # setup.py reads install_requires from a file next to itself, so stage a copy. + add_custom_command( + OUTPUT ${Z3FDB_STAGING}/z3fdb-requirements.txt + COMMAND ${CMAKE_COMMAND} -E copy + ${Z3FDB_RUNTIME_REQUIREMENTS} + ${Z3FDB_STAGING}/z3fdb-requirements.txt + DEPENDS ${Z3FDB_RUNTIME_REQUIREMENTS} + COMMENT "Copying z3fdb-requirements.txt to staging..." ) + add_subdirectory(chunked_data_view_bindings) + + # ---- Stub generation ---- # chunked_data_view_bindings/__init__.py calls findlibs.load("fdb5") before - # the extension module is imported, and findlibs to locate the library. - # FDB5_HOME points findlibs directly at the fdb5 build output directory so - # it can find fdb5lib there. + # the extension module is imported. FDB5_HOME points findlibs directly at the + # fdb5 build output directory so it can find fdb5lib there. # - # Stubgen writes into a temporary directory; only chunked_data_view_bindings.pyi - # is copied next to the *.so. The temp dir (including any __init__.pyi) is - # removed afterwards. - # TODO(TKR): Enable this once the isolated build environment for python wheels is in place - # add_custom_command( - # OUTPUT ${CMAKE_BINARY_DIR}/stubs.stamp - # COMMAND ${CMAKE_COMMAND} -E env - # "PYTHONPATH=${Z3FDB_STAGING}" - # "FDB5_HOME=${CMAKE_BINARY_DIR}" - # ${Python_EXECUTABLE} -m pybind11_stubgen - # chunked_data_view_bindings.chunked_data_view_bindings - # -o ${CMAKE_BINARY_DIR}/stubs_tmp - # COMMAND ${CMAKE_COMMAND} -E copy - # ${CMAKE_BINARY_DIR}/stubs_tmp/chunked_data_view_bindings/chunked_data_view_bindings.pyi - # ${Z3FDB_STAGING}/chunked_data_view_bindings/chunked_data_view_bindings.pyi - # COMMAND ${CMAKE_COMMAND} -E rm -rf ${CMAKE_BINARY_DIR}/stubs_tmp - # COMMAND ${CMAKE_COMMAND} -E touch - # ${Z3FDB_STAGING}/chunked_data_view_bindings/py.typed - # COMMAND ${CMAKE_COMMAND} -E touch stubs.stamp - # WORKING_DIRECTORY ${CMAKE_BINARY_DIR} - # DEPENDS - # "chunked_data_view_bindings/__init__.py" - # chunked_data_view_bindings - # fdb5 - # COMMENT "Generating stubs and py.typed for chunked_data_view_bindings..." - # ) - # list(APPEND _z3fdb_package_files - # ${CMAKE_BINARY_DIR}/stubs.stamp - # ) + # Stubs are written to a PEP 561 stub-only package directory + # (chunked_data_view_bindings-stubs/) in the staging tree. This keeps the + # runtime package (chunked_data_view_bindings/) free of .pyi files and lets + # type checkers discover stubs via the standard stub-package lookup. + add_custom_command( + OUTPUT ${CMAKE_BINARY_DIR}/z3fdb.stubs.stamp + COMMAND ${CMAKE_COMMAND} -E env + "PYTHONPATH=${Z3FDB_STAGING}" + "FDB5_HOME=${CMAKE_BINARY_DIR}" + ${UVX_EXECUTABLE} + --python ${Python_VERSION} + --from pybind11-stubgen + --with-requirements ${Z3FDB_BUILD_REQUIREMENTS} + --isolated + pybind11-stubgen + chunked_data_view_bindings + -o ${CMAKE_BINARY_DIR}/stubs_tmp + COMMAND ${CMAKE_COMMAND} -E copy + ${CMAKE_BINARY_DIR}/stubs_tmp/chunked_data_view_bindings/chunked_data_view_bindings.pyi + ${Z3FDB_STAGING}/chunked_data_view_bindings-stubs/chunked_data_view_bindings.pyi + COMMAND ${CMAKE_COMMAND} -E copy + ${CMAKE_BINARY_DIR}/stubs_tmp/chunked_data_view_bindings/__init__.pyi + ${Z3FDB_STAGING}/chunked_data_view_bindings-stubs/__init__.pyi + COMMAND ${CMAKE_COMMAND} -E rm -rf ${CMAKE_BINARY_DIR}/stubs_tmp + COMMAND ${CMAKE_COMMAND} -E touch + ${Z3FDB_STAGING}/chunked_data_view_bindings-stubs/py.typed + COMMAND ${CMAKE_COMMAND} -E touch z3fdb.stubs.stamp + WORKING_DIRECTORY ${CMAKE_BINARY_DIR} + DEPENDS + ${Z3FDB_STAGING}/chunked_data_view_bindings/__init__.py + chunked_data_view_bindings + fdb5 + COMMENT "Generating stubs for chunked_data_view_bindings (-> chunked_data_view_bindings-stubs/)..." + ) + # ---- Wheel ---- add_custom_command( OUTPUT ${CMAKE_BINARY_DIR}/wheel.stamp - COMMAND ${Python_EXECUTABLE} -m build --wheel ${Z3FDB_STAGING} -o . + COMMAND ${UVX_EXECUTABLE} + --python ${Python_VERSION} + --from build + --isolated + --with-requirements ${Z3FDB_BUILD_REQUIREMENTS} + pyproject-build --wheel ${Z3FDB_STAGING} -o . COMMAND ${CMAKE_COMMAND} -E touch wheel.stamp WORKING_DIRECTORY ${CMAKE_BINARY_DIR} DEPENDS - ${_z3fdb_package_files} - chunked_data_view_bindings + ${_z3fdb_staged_py_files} + ${Z3FDB_STAGING}/README.md + ${Z3FDB_STAGING}/LICENSE + ${Z3FDB_STAGING}/z3fdb-requirements.txt + ${CMAKE_BINARY_DIR}/z3fdb.stubs.stamp COMMENT "Building Python wheel for z3fdb..." ) add_custom_target(z3fdb-wheel ALL DEPENDS ${CMAKE_BINARY_DIR}/wheel.stamp) diff --git a/src/chunked_data_view/Axis.h b/src/chunked_data_view/Axis.h index 26a78e8b3..0215ed219 100644 --- a/src/chunked_data_view/Axis.h +++ b/src/chunked_data_view/Axis.h @@ -39,14 +39,17 @@ class Parameter { /// Describes how one axis is divided into Zarr chunks. /// /// Stores one extent entry per chunk; for uniform chunking all entries are equal. -/// The extensible flag marks axes using WholeAxisChunking: their single chunk grows when -/// additional parts are stitched onto the view along this axis. +/// The singleGrowingChunk flag marks axes using WholeAxisChunking: their single chunk +/// grows to encompass all parts' extents when multiple parts are stitched together. +/// SingleValueChunking and FixedSizeChunking use fixed chunk sizes and produce more +/// chunks (not larger ones) as parts are added — singleGrowingChunk is false for these. class AxisChunks { public: - explicit AxisChunks(const std::vector>>& chunks, bool extensible) : - extensible_(extensible) { + explicit AxisChunks(const std::vector>>& chunks, + bool singleGrowingChunk) : + singleGrowingChunk_(singleGrowingChunk) { for (const auto& element : chunks) { if (std::holds_alternative(element)) { extensions_.emplace_back(std::get(element)); @@ -61,16 +64,18 @@ class AxisChunks { } /// Convenience constructor: @p amount chunks each of size @p chunk_extension. - AxisChunks(size_t chunk_extension, size_t amount, bool extensible) : + AxisChunks(size_t chunk_extension, size_t amount, bool singleGrowingChunk) : AxisChunks(std::vector>>{std::tuple{ chunk_extension, amount}}, - extensible) {}; + singleGrowingChunk) {}; /// Number of chunks along this axis. size_t size() const { return extensions_.size(); } - /// True for WholeAxisChunking axes, whose combined extent is the sum of all parts' extents. - bool isExtensible() const { return extensible_; } + /// True for WholeAxisChunking axes only: the single chunk's extent is the sum of all + /// parts' extents. False for SingleValueChunking and FixedSizeChunking, where the chunk + /// size stays fixed and more chunks accumulate as parts are added. + bool isSingleGrowingChunk() const { return singleGrowingChunk_; } /// Per-chunk extents; each entry is the number of axis elements in that chunk. const std::vector& extensions() const { return extensions_; } @@ -81,7 +86,7 @@ class AxisChunks { private: std::vector extensions_{}; - bool extensible_; + bool singleGrowingChunk_; }; diff --git a/src/chunked_data_view/CMakeLists.txt b/src/chunked_data_view/CMakeLists.txt index b848822d1..eb5ee962c 100644 --- a/src/chunked_data_view/CMakeLists.txt +++ b/src/chunked_data_view/CMakeLists.txt @@ -15,8 +15,12 @@ add_library(chunked_data_view STATIC ChunkedDataViewImpl.cc ChunkedDataViewImpl.h Fdb.cc - GribExtractor.cc - GribExtractor.h + extractors/grib/GribExtractor.cc + extractors/grib/GribExtractor.h + extractors/grib/GribExtractorDefinition.cc + extractors/grib/GribExtractorDefinition.h + extractors/gribjump/GribJumpExtractorDefinition.cc + extractors/gribjump/GribJumpExtractorDefinition.h LibChunkedDataView.cc ListIterator.cc RequestManipulation.cc @@ -28,6 +32,7 @@ add_library(chunked_data_view STATIC mapping/IndexMapper.cc mapping/IndexMapper.h exception/GribExtractorException.cc + exception/GribJumpExtractorException.cc exception/UnknownExtractorException.cc exception/BoundingBoxException.cc exception/RequestManipulationException.cc @@ -42,6 +47,7 @@ add_library(chunked_data_view STATIC include/chunked_data_view/ListIterator.h include/chunked_data_view/Types.h include/chunked_data_view/exception/GribExtractorException.h + include/chunked_data_view/exception/GribJumpExtractorException.h include/chunked_data_view/exception/UnknownExtractorException.h include/chunked_data_view/exception/BoundingBoxException.h include/chunked_data_view/exception/RequestManipulationException.h @@ -55,6 +61,19 @@ target_include_directories(chunked_data_view ${CMAKE_CURRENT_BINARY_DIR} ) +# GribJumpExtractor is the only part of this library that needs gribjump headers. The +# definition (GribJumpExtractorDefinition) and its exception are gribjump-free and stay +# unconditional, so the type remains complete — and bindable from Python — in every build. +if(HAVE_ZARR_GRIBJUMP_EXTRACTOR) + target_sources(chunked_data_view PRIVATE + extractors/gribjump/GribJumpExtractor.cc + extractors/gribjump/GribJumpExtractor.h + ) + target_link_libraries(chunked_data_view PRIVATE gribjump) + # PUBLIC: GribJumpExtractorDefinition.cc and chunked_data_view_bindings both switch on this. + target_compile_definitions(chunked_data_view PUBLIC HAVE_ZARR_GRIBJUMP_EXTRACTOR) +endif() + target_link_libraries(chunked_data_view PRIVATE eckit diff --git a/src/chunked_data_view/ChunkedDataViewBuilder.cc b/src/chunked_data_view/ChunkedDataViewBuilder.cc index b5e33603d..2da2c26d7 100644 --- a/src/chunked_data_view/ChunkedDataViewBuilder.cc +++ b/src/chunked_data_view/ChunkedDataViewBuilder.cc @@ -5,6 +5,7 @@ #include "ChunkedDataViewImpl.h" #include "chunked_data_view/AxisDefinition.h" #include "chunked_data_view/ChunkedDataView.h" +#include "chunked_data_view/DataLayout.h" #include "chunked_data_view/Extractor.h" #include "chunked_data_view/ViewPart.h" #include "chunked_data_view/mapping/AxisMapper.h" @@ -28,8 +29,10 @@ ChunkedDataViewBuilder::ChunkedDataViewBuilder(const std::optional axes, - std::shared_ptr extractor) { - parts_.emplace_back(std::move(marsRequestKeyValues), std::move(axes), std::move(extractor)); + const ExtractorDefinition& definition) { + auto copiedDefinition = definition.copy(); + copiedDefinition->setDefaultIfUnset(configPath_); // Set the default + parts_.emplace_back(std::move(marsRequestKeyValues), std::move(axes), std::move(copiedDefinition)); return *this; } @@ -44,7 +47,7 @@ ChunkedDataViewBuilder& ChunkedDataViewBuilder::fillMissingValue(float fillValue } bool ChunkedDataViewBuilder::chunkingConsistencyCheck( - const std::vector>>& viewParts) { + const std::vector>>& viewParts) { if (viewParts.size() <= 1) { return true; @@ -54,9 +57,9 @@ bool ChunkedDataViewBuilder::chunkingConsistencyCheck( const size_t numAxes = refChunks.size(); for (size_t axisIdx = 0; axisIdx < numAxes; ++axisIdx) { - // WholeAxisChunking axes are extensible: their single chunk grows when parts are - // stitched together, so differing extents per part are expected and correct. - if (refChunks[axisIdx].isExtensible()) { + // WholeAxisChunking axes have one chunk whose size grows across parts; + // differing extents per part are expected and correct, so skip the check. + if (refChunks[axisIdx].isSingleGrowingChunk()) { continue; } @@ -70,8 +73,36 @@ bool ChunkedDataViewBuilder::chunkingConsistencyCheck( return true; } +void ChunkedDataViewBuilder::validateLayouts( + const std::vector>>& viewParts) { + + const DataLayout& reference = viewParts[0].second->layout(); + + for (size_t index = 1; index < viewParts.size(); ++index) { + const DataLayout& current = viewParts[index].second->layout(); + + if (current.countValues != reference.countValues) { + std::ostringstream ss; + ss << "ChunkedDataViewBuilder::build: part " << index << " has " << current.countValues + << " grid points but part 0 has " << reference.countValues + << ". The grid-point dimension is never the extension axis, so every part must cover the same " + "grid; a view cannot have a ragged last dimension."; + throw eckit::UserError(ss.str()); + } + + if (current.countChunkValues != reference.countChunkValues) { + std::ostringstream ss; + ss << "ChunkedDataViewBuilder::build: part " << index << " splits the grid-point dimension into chunks of " + << current.countChunkValues << " values but part 0 uses " << reference.countChunkValues + << ". All parts must agree on the field chunking, so a GribJump part mixed with a Grib part has to " + "use the default WholeAxisChunking."; + throw eckit::UserError(ss.str()); + } + } +} + bool ChunkedDataViewBuilder::doPartsAlign( - const std::vector>>& viewParts) { + const std::vector>>& viewParts) { const ViewPart& first = std::get<0>(viewParts[0]); bool extensible = true; for (const auto& [viewPart, _] : viewParts) { @@ -96,15 +127,13 @@ std::unique_ptr ChunkedDataViewBuilder::build() { } } - std::vector>> viewParts{}; + std::vector>> viewParts{}; viewParts.reserve(parts_.size()); // Offset is one-dimensional along the extension axis std::vector part_offsets = {0}; - for (auto& [req, defs, ext] : parts_) { - ext->setFillValue(fillValue_); - + for (auto& [req, defs, extDef] : parts_) { const auto requests = fdb5::FDBToolRequest::requestsFromString(req); if (requests.size() > 1) { @@ -116,17 +145,18 @@ std::unique_ptr ChunkedDataViewBuilder::build() { const auto request = requests.at(0).request(); try { - const auto layout = ext->layout(request); const auto axes = AxisMapper::mapRequestToAxis(request, defs); // Create offset vector std::vector offsetInChunkedDataView(axes.size(), 0); offsetInChunkedDataView[extensionAxisIndex_.value_or(0)] = part_offsets[part_offsets.size() - 1]; - ViewPart vp(std::move(request), layout, axes, offsetInChunkedDataView); + ViewPart vp(std::move(request), axes, offsetInChunkedDataView); part_offsets.push_back(part_offsets.back() + vp.extension()[extensionAxisIndex_.value_or(0)]); - viewParts.emplace_back(std::move(vp), ext); + auto ext = extDef->buildExtractor(request); + ext->setFillValue(fillValue_); + viewParts.emplace_back(std::move(vp), std::move(ext)); } catch (const std::exception& e) { std::ostringstream ss; @@ -136,6 +166,8 @@ std::unique_ptr ChunkedDataViewBuilder::build() { } } + validateLayouts(viewParts); + if (!doPartsAlign(viewParts)) { throw eckit::UserError("Shape of all parts must be identical except for the extension axis index."); } @@ -147,7 +179,7 @@ std::unique_ptr ChunkedDataViewBuilder::build() { "boundaries coincide with Zarr chunk boundaries."); } - return std::make_unique(viewParts, fillValue_, extensionAxisIndex_.value_or(0)); + return std::make_unique(std::move(viewParts), fillValue_, extensionAxisIndex_.value_or(0)); } }; // namespace chunked_data_view diff --git a/src/chunked_data_view/ChunkedDataViewImpl.cc b/src/chunked_data_view/ChunkedDataViewImpl.cc index 3b7848353..818ab8fbd 100644 --- a/src/chunked_data_view/ChunkedDataViewImpl.cc +++ b/src/chunked_data_view/ChunkedDataViewImpl.cc @@ -9,22 +9,26 @@ #include "ChunkedDataViewImpl.h" +#include "chunked_data_view/Extractor.h" +#include "chunked_data_view/Types.h" #include "chunked_data_view/ViewPart.h" #include "eckit/exception/Exceptions.h" namespace chunked_data_view { -namespace {} // namespace -bool checkForEqualChunking(const std::vector>>& parts) { +// Duplicates ChunkedDataViewBuilder::chunkingConsistencyCheck; kept as a belt-and-braces +// check because the invariant is what chunkShape() relies on. Static: it is not API. +static bool checkForEqualChunking(const std::vector>>& parts) { const auto reference_chunks = parts[0].first.chunks(); for (const auto& [part, _] : parts) { for (size_t i = 0; i < part.axes().size(); ++i) { if (part.chunks()[i].representativeExtent() != reference_chunks[i].representativeExtent()) { - // If the axis is extensible along this axis, skip as we are fetching - if (!reference_chunks[i].isExtensible()) { + // WholeAxisChunking axes have one growing chunk per part; differing extents + // are expected and correct, so skip the consistency check for those. + if (!reference_chunks[i].isSingleGrowingChunk()) { return false; } } @@ -35,7 +39,7 @@ bool checkForEqualChunking(const std::vector ChunkedDataViewImpl::chunkShape( - const std::vector>>& parts) { + const std::vector>>& parts) { const ViewPart& reference_part = (parts[0].first); std::vector reference_extensions; @@ -43,10 +47,11 @@ std::vector ChunkedDataViewImpl::chunkShape( for (size_t i = 0; i < reference_part.axes().size(); ++i) { reference_extensions.push_back(reference_part.chunks()[i].representativeExtent()); } - reference_extensions.push_back(reference_part.layout().countValues); // Add the size of the fields + reference_extensions.push_back(parts[0].second->layout().countChunkValues); // Add the size of the (sub)fields - // Check for merging in case of extension axis - if (reference_part.isExtensible(extensionAxisIndex_)) { + // WholeAxisChunking: merge all parts' extents into one growing chunk. + // SingleValueChunking / FixedSizeChunking: chunk size is fixed; the number of chunks grows. + if (reference_part.isSingleGrowingChunk(extensionAxisIndex_)) { reference_extensions[extensionAxisIndex_] = 0; for (const auto& [part, extractor] : parts) { reference_extensions[extensionAxisIndex_] += part.extension()[extensionAxisIndex_]; @@ -56,7 +61,7 @@ std::vector ChunkedDataViewImpl::chunkShape( return reference_extensions; } -ChunkedDataViewImpl::ChunkedDataViewImpl(std::vector>>& parts, +ChunkedDataViewImpl::ChunkedDataViewImpl(std::vector>> parts, float fillValue, size_t extensionAxisIndex) : parts_(std::move(parts)), extensionAxisIndex_(extensionAxisIndex), fillValue_(fillValue) { @@ -88,7 +93,7 @@ ChunkedDataViewImpl::ChunkedDataViewImpl(std::vectorlayout().countValues); if (!checkForEqualChunking(parts_)) { throw eckit::UserError("ChunkedDataViewImpl::constructor: view parts need to have same chunking extensions."); @@ -97,14 +102,11 @@ ChunkedDataViewImpl::ChunkedDataViewImpl(std::vector(chunkShape_.size(), 0); - // The last dimension is implicitly created for the number of values in a field, i.e. there is no representation in - // the axes. And the dimension of fields is never chunked I.e. fields are always returned whole. - for (size_t index = 0; index < chunkShape_.size() - 1; ++index) { + for (size_t index = 0; index < chunkShape_.size(); ++index) { // Integer ceil chunks_[index] = chunkedDataViewShape_[index] / chunkShape_[index] + ((chunkedDataViewShape_[index] % chunkShape_[index]) != 0); } - chunks_.back() = 1; // Make the implicit dimension always single chunked } @@ -126,19 +128,19 @@ void ChunkedDataViewImpl::at(const std::vector& chunkIndex, float* ptr, } } - std::vector chunkLower(chunkShape_.size() - 1, 0); - std::vector chunkUpper(chunkShape_.size() - 1, 0); + std::vector chunkLower(chunkShape_.size(), 0); + std::vector chunkUpper(chunkShape_.size(), 0); - for (size_t i = 0; i < chunkShape_.size() - 1; ++i) { + for (size_t i = 0; i < chunkShape_.size(); ++i) { chunkLower[i] = chunkShape_[i] * chunkIndex[i]; chunkUpper[i] = chunkLower[i] + chunkShape_[i] - 1; } - ChunkedDataViewPartBoundingBox chunkBoundingBox{chunkLower, chunkUpper}; + ChunkBoundingBox chunkBoundingBox{chunkLower, chunkUpper}; for (const auto& [part, extractor] : parts_) { const std::optional intersectionBoundingBox = - part.boundingBox().intersect(chunkBoundingBox); + part.boundingBox().intersect(chunkBoundingBox.dropLastDimension()); // Skip the part if it doesn't contribute to the buffer if (!intersectionBoundingBox.has_value()) { diff --git a/src/chunked_data_view/ChunkedDataViewImpl.h b/src/chunked_data_view/ChunkedDataViewImpl.h index 6bf977c10..6bc1f1084 100644 --- a/src/chunked_data_view/ChunkedDataViewImpl.h +++ b/src/chunked_data_view/ChunkedDataViewImpl.h @@ -19,7 +19,7 @@ namespace chunked_data_view { class ChunkedDataViewImpl : public ChunkedDataView { public: - ChunkedDataViewImpl(std::vector>>& partialViews, float fillValue, + ChunkedDataViewImpl(std::vector>> partialViews, float fillValue, size_t extensionAxisIndex); /// Fills @p ptr with the float values of the chunk at @p chunkIndex. @@ -57,14 +57,16 @@ class ChunkedDataViewImpl : public ChunkedDataView { std::vector chunkShape_{}; std::vector chunkedDataViewShape_{}; std::vector chunks_{}; - std::vector>> parts_{}; + std::vector>> parts_{}; size_t extensionAxisIndex_{}; float fillValue_; private: // methods - /// Computes chunkShape_ from the parts, summing extensible-axis extents across all parts. - std::vector chunkShape(const std::vector>>& parts); + /// Computes chunkShape_ from the parts. For WholeAxisChunking on the extension axis, + /// all parts' extents are summed into one growing chunk. For SingleValueChunking and + /// FixedSizeChunking the representative chunk size is taken from the first part. + std::vector chunkShape(const std::vector>>& parts); }; } // namespace chunked_data_view diff --git a/src/chunked_data_view/ListIterator.cc b/src/chunked_data_view/ListIterator.cc index 403441e68..6c463f564 100644 --- a/src/chunked_data_view/ListIterator.cc +++ b/src/chunked_data_view/ListIterator.cc @@ -2,25 +2,18 @@ // SPDX-License-Identifier: Apache-2.0 #include "chunked_data_view/ListIterator.h" -#include "eckit/io/DataHandle.h" #include "fdb5/api/helpers/ListElement.h" #include "fdb5/api/helpers/ListIterator.h" -#include "fdb5/database/Key.h" -#include #include -#include -#include namespace chunked_data_view { -std::optional>> ListIteratorWrapperImpl::next() { +std::optional ListIteratorWrapperImpl::next() { fdb5::ListElement elem; - auto has_next = listIterator_.next(elem); - - if (has_next) { - return std::make_tuple(elem.combinedKey(), std::unique_ptr(elem.location().dataHandle())); + if (listIterator_.next(elem)) { + return ListElement{elem.combinedKey(), elem.sharedLocation()}; } return std::nullopt; diff --git a/src/chunked_data_view/ViewPart.cc b/src/chunked_data_view/ViewPart.cc index 2da8c0cf2..b2fa5c889 100644 --- a/src/chunked_data_view/ViewPart.cc +++ b/src/chunked_data_view/ViewPart.cc @@ -99,6 +99,7 @@ std::optional BoundingBox::intersect(const BoundingBox& other) cons // u1 < l2 or l1 > u2 (separating axis) // If there is an intersection it's [max(l1, l2), min(u1, u2)] // For every axis + assert(dimensions() == other.dimensions()); std::vector lower; std::vector upper; @@ -123,13 +124,13 @@ std::optional BoundingBox::intersect(const BoundingBox& other) cons return std::make_optional<>(BoundingBox(lower, upper)); } -ViewPart::ViewPart(const metkit::mars::MarsRequest& request, const DataLayout& data_layout, - const std::vector>& axes, const std::vector& offset) : - request_(request), layout_(data_layout), offset_(offset) { +ViewPart::ViewPart(const metkit::mars::MarsRequest& request, const std::vector>& axes, + const std::vector& offset) : + request_(request), offset_(offset) { - extension_.reserve(axes_.size()); - chunks_.reserve(axes_.size()); - axes_.reserve(axes_.size()); + extension_.reserve(axes.size()); + chunks_.reserve(axes.size()); + axes_.reserve(axes.size()); for (const auto& [axis, axis_chunks] : axes) { axes_.push_back(axis); diff --git a/src/chunked_data_view/ViewPart.h b/src/chunked_data_view/ViewPart.h index 275c65021..cbd7c9667 100644 --- a/src/chunked_data_view/ViewPart.h +++ b/src/chunked_data_view/ViewPart.h @@ -108,13 +108,12 @@ class ViewPart { /// Constructs a ViewPart. /// @param request The MARS request that describes the data covered by this part. - /// @param data_layout Number of values and bytes-per-value for each field. /// @param axes Ordered list of (Axis, AxisChunks) pairs, one per non-values dimension. /// Each keyword with more than one value must be covered by exactly one axis. /// @param offset Position of the lower corner of this part in the global view index space, /// one entry per axis (excluding the implicit values dimension). - ViewPart(const metkit::mars::MarsRequest& request, const DataLayout& data_layout, - const std::vector>& axes, const std::vector& offset); + ViewPart(const metkit::mars::MarsRequest& request, const std::vector>& axes, + const std::vector& offset); ~ViewPart() = default; @@ -131,18 +130,19 @@ class ViewPart { /// Chunking descriptors for each axis (excluding the implicit values dimension). std::vector chunks() const { return chunks_; } - /// Returns true if the axis at @p axisIndex is marked as extensible, - /// i.e. additional parts may be stitched onto this part along that axis. - bool isExtensible(const size_t axisIndex) const { return chunks_[axisIndex].isExtensible(); } - - /// Field layout (countValues and bytesPerValue) shared by all fields in this part. - const DataLayout& layout() const { return layout_; } + /// Returns true if the axis at @p axisIndex uses WholeAxisChunking, meaning its single + /// chunk grows to cover the combined extent of all stitched parts. False for + /// SingleValueChunking and FixedSizeChunking, which accumulate more fixed-size chunks. + bool isSingleGrowingChunk(const size_t axisIndex) const { return chunks_[axisIndex].isSingleGrowingChunk(); } /// Ordered axes that define the non-values dimensions of this part. const std::vector& axes() const { return axes_; } - /// Number of entries (fields or values) along each dimension, including the implicit values dimension as the last - /// entry. + /// Number of entries along each MARS-derived dimension, one per axis. + /// + /// Does *not* include the implicit values dimension: ChunkedDataViewImpl appends that from + /// the extractor's DataLayout after using this to size the view. Consequently the extension + /// axis index is always a valid index into this vector. std::vector extension() const { return extension_; } /// Position of the lower corner of this part in the global view index space. @@ -154,7 +154,6 @@ class ViewPart { /// Offset of this part along a single axis in the global view index space. size_t offsetOnAxis(size_t axisIndex) const { return offset_[axisIndex]; } - bool isAxisChunked(size_t index) const { return true; }; /// Returns true if this part and @p other can be stitched together along @p extension_axis, /// i.e. their extents match on every axis except the extension axis. @@ -169,7 +168,6 @@ class ViewPart { metkit::mars::MarsRequest request_{}; std::vector axes_{}; std::vector chunks_; - DataLayout layout_{}; std::vector extension_{}; // extension in each dimension, counting entries std::vector offset_{}; // offset in chunked data view diff --git a/src/chunked_data_view/exception/GribJumpExtractorException.cc b/src/chunked_data_view/exception/GribJumpExtractorException.cc new file mode 100644 index 000000000..75a4c9242 --- /dev/null +++ b/src/chunked_data_view/exception/GribJumpExtractorException.cc @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#include "chunked_data_view/exception/GribJumpExtractorException.h" + +namespace chunked_data_view { +GribJumpExtractorException::GribJumpExtractorException(const std::string& w) : Exception(w) {} + +GribJumpExtractorException::GribJumpExtractorException(const std::string& w, const eckit::CodeLocation& l) : + Exception(w, l) {} + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/GribExtractor.cc b/src/chunked_data_view/extractors/grib/GribExtractor.cc similarity index 74% rename from src/chunked_data_view/GribExtractor.cc rename to src/chunked_data_view/extractors/grib/GribExtractor.cc index 7f15087e0..a6a7581cf 100644 --- a/src/chunked_data_view/GribExtractor.cc +++ b/src/chunked_data_view/extractors/grib/GribExtractor.cc @@ -1,12 +1,13 @@ -// SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) // SPDX-License-Identifier: Apache-2.0 -#include "GribExtractor.h" +#include "chunked_data_view/extractors/grib/GribExtractor.h" #include "chunked_data_view/DataLayout.h" #include "chunked_data_view/Extractor.h" #include "chunked_data_view/Fdb.h" #include "chunked_data_view/ListIterator.h" #include "chunked_data_view/RequestManipulation.h" +#include "chunked_data_view/Types.h" #include "chunked_data_view/ViewPart.h" #include "chunked_data_view/exception/GribExtractorException.h" #include "chunked_data_view/mapping/IndexMapper.h" @@ -14,24 +15,25 @@ #include "eckit/exception/Exceptions.h" #include "eckit/message/Reader.h" #include "fdb5/database/Key.h" +#include "metkit/mars/MarsRequest.h" #include #include #include +#include #include #include namespace chunked_data_view { -GribExtractor::GribExtractor(const std::shared_ptr fdb) : fdb_(fdb) {} - -DataLayout GribExtractor::layout(const metkit::mars::MarsRequest& mars_request) const { +GribExtractor::GribExtractor(std::unique_ptr fdb, const metkit::mars::MarsRequest& marsRequest) : + fdb_(std::move(fdb)) { // Use a minimal sample request: all requested params but only the first value of every // other key. This lets us verify every param the user asked for without retrieving the // full data volume. const metkit::mars::MarsRequest sampleRequest = - mars_request.has("param") ? RequestManipulation::allParamRequest(mars_request) : mars_request; + marsRequest.has("param") ? RequestManipulation::allParamRequest(marsRequest) : marsRequest; const auto& handle = fdb_->retrieve(sampleRequest); eckit::message::Reader reader(*handle); @@ -50,12 +52,13 @@ DataLayout GribExtractor::layout(const metkit::mars::MarsRequest& mars_request) // the request exactly. Checking all messages (not just the first) catches cases where // the mismatch only appears for certain parameters. // - // Note: mars_request was produced by FDBToolRequest::requestsFromString() which runs - // metkit's TypeParam expansion pass. Short param names (e.g. "v", "vo") are resolved - // to numeric paramId strings (e.g. "132", "138") before this point, so the comparison - // against std::to_string(msg.getLong("paramId")) is always numeric-vs-numeric. - if (mars_request.has("param")) { - const auto& requestedParams = mars_request.values("param"); + // Note: marsRequest has already been passed through FDBToolRequest::requestsFromString() + // by ChunkedDataViewBuilder::build(), which runs metkit's TypeParam expansion pass. + // Short param names (e.g. "v", "vo") are resolved to numeric paramId strings (e.g. "132", + // "138") before this point, so the comparison against std::to_string(msg.getLong("paramId")) + // is always numeric-vs-numeric. + if (marsRequest.has("param")) { + const auto& requestedParams = marsRequest.values("param"); do { const std::string returnedParam = std::to_string(msg.getLong("paramId")); if (std::find(requestedParams.begin(), requestedParams.end(), returnedParam) == requestedParams.end()) { @@ -74,8 +77,7 @@ DataLayout GribExtractor::layout(const metkit::mars::MarsRequest& mars_request) } while ((msg = reader.next())); } - - return {countValues, 4}; + layout_ = {countValues, 4, countValues}; } @@ -117,30 +119,30 @@ size_t GribExtractor::writeInto(std::unique_ptr list_iter } iterator_empty = false; - const auto& key = std::get<0>(*res); - auto& data_handle = std::get<1>(*res); + const auto& key = res->key; + auto data_handle = res->dataHandle(); const size_t msgIndex = index_mapping::computeBufferIndex(ctx.axes, key, ctx.partAxisOffset, ctx.bufferOffset, ctx.bufferExtent); eckit::message::Reader reader(*data_handle); eckit::message::Message msg{}; - auto copyInto = ptr + msgIndex * ctx.layout.countValues; - const auto end = copyInto + ctx.layout.countValues; + auto copyInto = ptr + msgIndex * ctx.layout.countChunkValues; + const auto end = copyInto + ctx.layout.countChunkValues; ASSERT(end - ptr <= len); while ((msg = reader.next())) { - if (const auto size = msg.getSize("values"); size != ctx.layout.countValues) { + if (const auto size = msg.getSize("values"); size != ctx.layout.countChunkValues) { std::ostringstream ss; - ss << "GribExractor: Unexpected field size found in GRIB message for key: " << key - << " expected: " << ctx.layout.countValues << " found: " << size + ss << "GribExtractor: Unexpected field size found in GRIB message for key: " << key + << " expected: " << ctx.layout.countChunkValues << " found: " << size << ". All fields in your view need to be of equal size."; throw eckit::Exception(ss.str()); } - msg.getFloatArray("values", copyInto, ctx.layout.countValues); + msg.getFloatArray("values", copyInto, ctx.layout.countChunkValues); if (msg.getLong("bitmapPresent") != 0) { const auto gribMissing = static_cast(msg.getDouble("missingValue")); - std::replace(copyInto, copyInto + ctx.layout.countValues, gribMissing, fillValue_); + std::replace(copyInto, copyInto + ctx.layout.countChunkValues, gribMissing, fillValue_); } messagesWritten++; } @@ -154,21 +156,28 @@ size_t GribExtractor::writeInto(std::unique_ptr list_iter return messagesWritten; } -size_t GribExtractor::extractInto(const ViewPart& part, const ChunkedDataViewPartBoundingBox& chunkBoundingBox, +size_t GribExtractor::extractInto(const ViewPart& part, const ChunkBoundingBox& chunkBoundingBox, const ChunkedDataViewPartBoundingBox& intersectionBoundingBox, float* ptr, size_t len) const { - ASSERT(chunkBoundingBox.contains(intersectionBoundingBox)); + + const auto& chunkPartBoundingBox = chunkBoundingBox.dropLastDimension(); + + ASSERT(chunkPartBoundingBox.contains(intersectionBoundingBox)); ASSERT(part.boundingBox().contains(intersectionBoundingBox)); const PartBoundingBox& partRelativeBoundingBox = intersectionBoundingBox.subtract(part.boundingBox().lower()); const auto& request = part.at(partRelativeBoundingBox); + + // fdb_ is shared mutable state; see mutex_ in the header. + const std::lock_guard lock(mutex_); + auto listIterator = fdb_->inspect(request); - const BufferBoundingBox& bufferRelativBoundingBox = intersectionBoundingBox.subtract(chunkBoundingBox.lower()); + const BufferBoundingBox& bufferRelativBoundingBox = intersectionBoundingBox.subtract(chunkPartBoundingBox.lower()); - const WriteContext ctx{part.axes(), part.layout(), partRelativeBoundingBox.lower(), - bufferRelativBoundingBox.lower(), chunkBoundingBox.extent()}; + const WriteContext ctx{part.axes(), layout_, partRelativeBoundingBox.lower(), bufferRelativBoundingBox.lower(), + chunkPartBoundingBox.extent()}; try { size_t written = writeInto(std::move(listIterator), ctx, ptr, len); @@ -176,6 +185,7 @@ size_t GribExtractor::extractInto(const ViewPart& part, const ChunkedDataViewPar } catch (GribExtractorException& exception) { std::ostringstream ss; + ss << "GribExtractor::extractInto: "; ss << exception.what(); ss << "Request was: " << part.at(partRelativeBoundingBox) << std::endl; throw GribExtractorException(ss.str()); diff --git a/src/chunked_data_view/GribExtractor.h b/src/chunked_data_view/extractors/grib/GribExtractor.h similarity index 71% rename from src/chunked_data_view/GribExtractor.h rename to src/chunked_data_view/extractors/grib/GribExtractor.h index dfbc092cc..cf17756c1 100644 --- a/src/chunked_data_view/GribExtractor.h +++ b/src/chunked_data_view/extractors/grib/GribExtractor.h @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) // SPDX-License-Identifier: Apache-2.0 #pragma once @@ -7,11 +7,15 @@ #include "chunked_data_view/Extractor.h" #include "chunked_data_view/Fdb.h" #include "chunked_data_view/ListIterator.h" +#include "chunked_data_view/Types.h" #include "chunked_data_view/ViewPart.h" +#include "metkit/mars/MarsRequest.h" + #include #include #include +#include #include namespace chunked_data_view { @@ -23,17 +27,16 @@ namespace chunked_data_view { class GribExtractor final : public Extractor { public: - explicit GribExtractor(const std::shared_ptr fdb); + /// Constructs a GribExtractor and eagerly determines the DataLayout by retrieving + /// a representative field from FDB for the given MARS request. + explicit GribExtractor(std::unique_ptr fdb, const metkit::mars::MarsRequest& marsRequest); /// Sets the fill value written in place of bitmap-masked (missing) grid points. void setFillValue(float v) override { fillValue_ = v; } - /// Retrieves one representative field to determine countValues and bytesPerValue. - DataLayout layout(const metkit::mars::MarsRequest& mars_request) const override; - /// Copies all fields in @p intersectionBoundingBox into @p ptr. /// @p chunkBoundingBox defines the origin for computing buffer offsets. - size_t extractInto(const ViewPart& part, const ChunkedDataViewPartBoundingBox& chunkBoundingBox, + size_t extractInto(const ViewPart& part, const ChunkBoundingBox& chunkBoundingBox, const ChunkedDataViewPartBoundingBox& intersectionBoundingBox, float* ptr, size_t len) const override; @@ -51,9 +54,15 @@ class GribExtractor final : public Extractor { private: // members - std::shared_ptr fdb_; + std::unique_ptr fdb_; float fillValue_ = std::numeric_limits::quiet_NaN(); + /// Serialises extractInto(). fdb_ is shared mutable backend state, but extractInto() is + /// const and the pybind11 layer releases the GIL around it, so a threaded zarr consumer + /// (e.g. dask) can enter it concurrently on one view. Reads serialise within a part; + /// separate parts own separate extractors and still proceed in parallel. + mutable std::mutex mutex_; + private: // methods /// Iterates over @p list_iterator, maps each field's key to a buffer slot via diff --git a/src/chunked_data_view/extractors/grib/GribExtractorDefinition.cc b/src/chunked_data_view/extractors/grib/GribExtractorDefinition.cc new file mode 100644 index 000000000..14d80cd9b --- /dev/null +++ b/src/chunked_data_view/extractors/grib/GribExtractorDefinition.cc @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#include "chunked_data_view/extractors/grib/GribExtractorDefinition.h" + +#include "chunked_data_view/Fdb.h" +#include "chunked_data_view/extractors/grib/GribExtractor.h" + +#include + +namespace chunked_data_view { + +void GribExtractorDefinition::setDefaultIfUnset(const std::optional& fdbConfigPath) { + if (!fdbConfig.has_value()) { + fdbConfig = fdbConfigPath; + } +} + +std::unique_ptr GribExtractorDefinition::copy() const { + return std::make_unique(*this); +} + +std::unique_ptr GribExtractorDefinition::buildExtractor(const metkit::mars::MarsRequest& request) const { + auto fdb = makeFdb(fdbConfig); + auto ext = std::make_unique(std::move(fdb), request); + return ext; +} + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/extractors/grib/GribExtractorDefinition.h b/src/chunked_data_view/extractors/grib/GribExtractorDefinition.h new file mode 100644 index 000000000..fc33a85ab --- /dev/null +++ b/src/chunked_data_view/extractors/grib/GribExtractorDefinition.h @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "chunked_data_view/Extractor.h" + +#include +#include +#include + +namespace chunked_data_view { + +/// ExtractorDefinition for full-field GRIB extraction via FDB. +/// +/// buildExtractor() creates a FdbInterface and GribExtractor for the given MARS request, +/// applying the configured fill value. +class GribExtractorDefinition : public ExtractorDefinition { +public: + + GribExtractorDefinition() = default; + GribExtractorDefinition(const GribExtractorDefinition&) = default; + GribExtractorDefinition& operator=(const GribExtractorDefinition&) = default; + GribExtractorDefinition(GribExtractorDefinition&&) = default; + GribExtractorDefinition& operator=(GribExtractorDefinition&&) = default; + + /// FDB config path. std::nullopt uses the environment default (FDB5_CONFIG / FDB_HOME). + std::optional fdbConfig; + + void setDefaultIfUnset(const std::optional& fdbConfigPath) override; + + std::unique_ptr copy() const override; + + std::unique_ptr buildExtractor(const metkit::mars::MarsRequest& request) const override; +}; + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.cc b/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.cc new file mode 100644 index 000000000..75abc2c20 --- /dev/null +++ b/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.cc @@ -0,0 +1,225 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#include "chunked_data_view/extractors/gribjump/GribJumpExtractor.h" + +#include "chunked_data_view/AxisDefinition.h" +#include "chunked_data_view/DataLayout.h" +#include "chunked_data_view/Fdb.h" +#include "chunked_data_view/ListIterator.h" +#include "chunked_data_view/RequestManipulation.h" +#include "chunked_data_view/Types.h" +#include "chunked_data_view/ViewPart.h" +#include "chunked_data_view/exception/GribJumpExtractorException.h" +#include "chunked_data_view/mapping/IndexMapper.h" + +#include "eckit/exception/Exceptions.h" +#include "eckit/message/Reader.h" +#include "fdb5/database/Key.h" +#include "metkit/mars/MarsRequest.h" + +#include "gribjump/ExtractionData.h" +#include "gribjump/GribJump.h" +#include "gribjump/api/ExtractionIterator.h" + +#include +#include +#include +#include +#include +#include + +namespace chunked_data_view { + +GribJumpExtractor::GribJumpExtractor(std::unique_ptr fdb, std::unique_ptr gj, + const metkit::mars::MarsRequest& marsRequest, + AxisDefinition::ChunkingType fieldChunking) : + fdb_(std::move(fdb)), gj_(std::move(gj)) { + + // Use a minimal sample request (mirrors GribExtractor constructor). + const metkit::mars::MarsRequest sampleRequest = + marsRequest.has("param") ? RequestManipulation::allParamRequest(marsRequest) : marsRequest; + + const auto& handle = fdb_->retrieve(sampleRequest); + + eckit::message::Reader reader(*handle); + eckit::message::Message msg = reader.next(); + + if (!msg) { + throw eckit::Exception("GribJumpExtractor::layout: Couldn't read GRIB message."); + } + + const size_t countValues = msg.getSize("values"); + + // Same paramId sanity check as GribExtractor constructor. + if (marsRequest.has("param")) { + const auto& requestedParams = marsRequest.values("param"); + do { + const std::string returnedParam = std::to_string(msg.getLong("paramId")); + if (std::find(requestedParams.begin(), requestedParams.end(), returnedParam) == requestedParams.end()) { + std::ostringstream buf; + buf << "GribJumpExtractor::layout: FDB returned paramId=" << returnedParam + << " which is not among the requested params ["; + for (size_t i = 0; i < requestedParams.size(); ++i) { + if (i > 0) { + buf << ", "; + } + buf << requestedParams[i]; + } + buf << "]. On-the-fly field derivation (e.g. u/v from vo/d) is not supported."; + throw GribJumpExtractorException(buf.str()); + } + } while ((msg = reader.next())); + } + + // The whole field is the window: there is no sub-range selection. + const size_t windowSize = countValues; + + // Resolve the per-chunk size from the fieldChunking variant. + struct ChunkSizeVisitor { + size_t windowSize; + size_t operator()(const AxisDefinition::WholeAxisChunking&) const { return windowSize; } + size_t operator()(const AxisDefinition::SingleValueChunking&) const { return 1; } + size_t operator()(const AxisDefinition::FixedSizeChunking& c) const { return c.chunkSize; } + }; + const size_t fieldChunkSize = std::visit(ChunkSizeVisitor{windowSize}, fieldChunking); + + if (fieldChunkSize == 0) { + throw eckit::UserError("GribJumpExtractor: field chunk size must be greater than zero."); + } + + if (windowSize % fieldChunkSize != 0) { + std::ostringstream ss; + ss << "GribJumpExtractor: field chunk size " << fieldChunkSize << " does not evenly divide the window size " + << windowSize << "."; + throw eckit::UserError(ss.str()); + } + + // countValues = total window size (implicit dimension extent in the Zarr array). + // countChunkValues = per-chunk size (what extractInto() writes per field per call). + layout_ = {windowSize, 4, fieldChunkSize}; +} + +size_t GribJumpExtractor::writeInto(const std::vector& gj_keys, gribjump::ExtractionIterator& gj_it, + const WriteContext& ctx, float* ptr, size_t len) const { + size_t messagesWritten = 0; + + for (size_t i = 0; gj_it.hasNext(); ++i) { + auto result = gj_it.next(); // unique_ptr + if (!result) { + break; + } + + // gribjump::LocalGribJump::collect_results() builds its vector by walking the requests + // in order, so result i belongs to gj_requests[i] and hence gj_keys[i]. Checked, not + // assumed: a change upstream would otherwise scatter fields into the wrong slots. + ASSERT(i < gj_keys.size()); + + const fdb5::Key& key = gj_keys.at(i); + const size_t msgIndex = + index_mapping::computeBufferIndex(ctx.axes, key, ctx.partAxisOffset, ctx.bufferOffset, ctx.bufferExtent); + + float* dst = ptr + msgIndex * ctx.layout.countChunkValues; + const float* end = dst + ctx.layout.countChunkValues; + ASSERT(end - ptr <= static_cast(len)); + + ASSERT(result->values().size() == 1); + ASSERT(result->mask().size() == 1); + + const auto& vals = result->values()[0]; // vector + const auto& mask = result->mask()[0]; // vector> + + ASSERT(vals.size() == ctx.layout.countChunkValues); + + for (size_t j = 0; j < vals.size(); ++j) { + const size_t word = j / 64; + const size_t bit = j % 64; + // GribJump bitmap convention: bit set (1) = valid, bit clear (0) = missing. + dst[j] = mask[word][bit] ? static_cast(vals[j]) : fillValue_; + } + + ++messagesWritten; + } + + return messagesWritten; +} + +size_t GribJumpExtractor::extractInto(const ViewPart& part, const ChunkBoundingBox& chunkBoundingBox, + const ChunkedDataViewPartBoundingBox& intersectionBB, float* ptr, + size_t len) const { + const ChunkedDataViewPartBoundingBox chunkPartBoundingBox = chunkBoundingBox.dropLastDimension(); + + ASSERT(chunkPartBoundingBox.contains(intersectionBB)); + ASSERT(part.boundingBox().contains(intersectionBB)); + + const PartBoundingBox& partRelBB = intersectionBB.subtract(part.boundingBox().lower()); + const BufferBoundingBox& bufRelBB = intersectionBB.subtract(chunkPartBoundingBox.lower()); + + const metkit::mars::MarsRequest request = part.at(partRelBB); + + // Derive the extraction range from the implicit dimension of chunkBoundingBox. It is the + // same for every field in this chunk, so it is built once rather than per field. + // chunkBoundingBox.lower().back() = chunkIndex.back() * layout_.countChunkValues + // chunkBoundingBox.upper().back() = lower.back() + countChunkValues - 1 (inclusive) + // gribjump::Range is half-open, hence the +1 on the upper bound. + const gribjump::Range chunkRange{chunkBoundingBox.lower().back(), chunkBoundingBox.upper().back() + 1}; + + // fdb_ and gj_ are shared mutable state; see mutex_ in the header. + const std::lock_guard lock(mutex_); + + auto listIt = fdb_->inspect(request); + + std::vector gj_keys; + std::vector gj_requests; + + while (const auto res = listIt->next()) { + + gj_keys.push_back(res->key); + + const auto& location_uri = res->location->fullUri(); + + if (location_uri.fragment().empty()) { + std::ostringstream ss; + ss << "GribJumpExtractor: Empty fragment for location uri in request " << request + << ". Can't forward the file offset."; + throw GribJumpExtractorException(ss.str()); + } + + // File offsets can exceed 2 GB, therefore use stoll. + size_t fieldOffset = 0; + try { + fieldOffset = static_cast(std::stoll(location_uri.fragment())); + } + catch (const std::exception& e) { + std::ostringstream ss; + ss << "GribJumpExtractor: Could not parse the file offset '" << location_uri.fragment() + << "' from the location uri in request " << request << ": " << e.what(); + throw GribJumpExtractorException(ss.str()); + } + + gj_requests.emplace_back(location_uri.path(), location_uri.scheme(), fieldOffset, location_uri.host(), + location_uri.port() > 0 ? location_uri.port() : 0, + std::vector{chunkRange}); + } + + if (gj_keys.empty()) { + std::ostringstream ss; + ss << "GribJumpExtractor: Empty iterator for request " << request << ". Is the request correctly specified?"; + throw GribJumpExtractorException(ss.str()); + } + + try { + gribjump::ExtractionIterator gj_it = gj_->extract(gj_requests); + + const WriteContext ctx{part.axes(), layout_, partRelBB.lower(), bufRelBB.lower(), + chunkPartBoundingBox.extent()}; + + return writeInto(gj_keys, gj_it, ctx, ptr, len); + } + catch (eckit::SeriousBug& exception) { + std::ostringstream buf; + buf << "GribJumpExtractor::extractInto: " << exception.what() << ". Request was: " << request; + throw GribJumpExtractorException(buf.str()); + } +} + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.h b/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.h new file mode 100644 index 000000000..f48e109c7 --- /dev/null +++ b/src/chunked_data_view/extractors/gribjump/GribJumpExtractor.h @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "chunked_data_view/Axis.h" +#include "chunked_data_view/DataLayout.h" +#include "chunked_data_view/Extractor.h" +#include "chunked_data_view/Fdb.h" +#include "chunked_data_view/ViewPart.h" + +#include "gribjump/Types.h" // gribjump::Range +#include "metkit/mars/MarsRequest.h" + +#include +#include +#include +#include +#include + +namespace gribjump { +class GribJump; +class ExtractionIterator; +} // namespace gribjump + +namespace fdb5 { +class Key; +} // namespace fdb5 + +namespace chunked_data_view { + +// --------------------------------------------------------------------------- +// GribJumpExtractor +// --------------------------------------------------------------------------- + +/// Concrete Extractor that uses GribJump as the data-retrieval backend. +/// +/// Field enumeration (key ordering for computeBufferIndex) still uses +/// FdbInterface::inspect(). Actual value extraction uses +/// gribjump::GribJump::extract(), avoiding full GRIB decode. +/// +/// An optional fieldChunking controls how the implicit (grid-point) dimension +/// is sub-divided into Zarr chunks. +class GribJumpExtractor final : public Extractor { +public: + + /// Constructs a GribJumpExtractor and eagerly determines the DataLayout by retrieving + /// a representative field from FDB for the given MARS request. + explicit GribJumpExtractor(std::unique_ptr fdb, std::unique_ptr gj, + const metkit::mars::MarsRequest& marsRequest, + AxisDefinition::ChunkingType fieldChunking = AxisDefinition::WholeAxisChunking{}); + + void setFillValue(float v) override { fillValue_ = v; } + + /// Copies the GribJump-extracted values for all fields in @p intersectionBB + /// into the output buffer. Derives the per-chunk grid-point range from the + /// last dimension of @p chunkBB. + size_t extractInto(const ViewPart& part, const ChunkBoundingBox& chunkBB, + const ChunkedDataViewPartBoundingBox& intersectionBB, float* ptr, size_t len) const override; + +private: // types + + /// Bundles all index-mapping and field metadata needed by writeInto(). + /// All members are references; the struct must not outlive extractInto(). + struct WriteContext { + const std::vector& axes; + const DataLayout& layout; + const std::vector& partAxisOffset; ///< Intersection start in part-local axis space. + const std::vector& bufferOffset; ///< Intersection start in chunk-buffer space. + const std::vector& bufferExtent; ///< Per-axis size of the chunk buffer. + }; + +private: // members + + std::unique_ptr fdb_; + std::unique_ptr gj_; + float fillValue_ = std::numeric_limits::quiet_NaN(); + + /// Serialises extractInto(). fdb_ and gj_ are shared mutable backend state, but + /// extractInto() is const and the pybind11 layer releases the GIL around it, so a threaded + /// zarr consumer (e.g. dask) can enter it concurrently on one view. Reads serialise within + /// a part; separate parts own separate extractors and still proceed in parallel. + mutable std::mutex mutex_; + + // NOTE: the grid-point window and the per-chunk size are fully described by layout_ + // (countValues / countChunkValues) from the base class, so they are deliberately not + // duplicated here. extractInto() derives the range it needs from its chunk bounding box. + +private: // methods + + /// Iterate @p gj_it, map each result to the correct buffer slot, and write + /// the double values (with GribJump bitmap) as float32. + size_t writeInto(const std::vector& gj_keys, gribjump::ExtractionIterator& gj_it, + const WriteContext& ctx, float* ptr, size_t len) const; +}; + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.cc b/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.cc new file mode 100644 index 000000000..4264bfd82 --- /dev/null +++ b/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.cc @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#include "chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.h" + +#include "chunked_data_view/Fdb.h" +#ifdef HAVE_ZARR_GRIBJUMP_EXTRACTOR +#include "chunked_data_view/extractors/gribjump/GribJumpExtractor.h" +#include "gribjump/GribJump.h" +#endif + +#include "eckit/exception/Exceptions.h" + +#include +#include + +namespace chunked_data_view { + +void GribJumpExtractorDefinition::setDefaultIfUnset(const std::optional& fdbConfigPath) { + if (!fdbConfig.has_value()) { + fdbConfig = fdbConfigPath; + } +} + +std::unique_ptr GribJumpExtractorDefinition::copy() const { + return std::make_unique(*this); +} + +std::unique_ptr GribJumpExtractorDefinition::buildExtractor(const metkit::mars::MarsRequest& request) const { +#ifndef HAVE_ZARR_GRIBJUMP_EXTRACTOR + // This configuration type stays available in every build, so that user code does not depend + // on how fdb was compiled. Only building the extractor fails, and it says why. + throw eckit::UserError( + "GribJumpExtractorDefinition: this build has no GribJump support. Rebuild fdb with " + "-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON (requires a bundle build providing gribjump)."); +#else + if (gribjumpConfig) { + ::setenv("GRIBJUMP_CONFIG_FILE", gribjumpConfig->c_str(), /*overwrite=*/1); + } + auto fdb = makeFdb(fdbConfig); + auto gj = std::make_unique(); + auto ext = std::make_unique(std::move(fdb), std::move(gj), request, fieldChunking); + return ext; +#endif +} + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.h b/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.h new file mode 100644 index 000000000..c7a32e2b1 --- /dev/null +++ b/src/chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.h @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "chunked_data_view/AxisDefinition.h" +#include "chunked_data_view/Extractor.h" + +#include +#include +#include + +namespace chunked_data_view { + +/// ExtractorDefinition for partial-field extraction via GribJump. +/// +/// buildExtractor() sets GRIBJUMP_CONFIG_FILE (if gribjumpConfig is set), creates +/// FdbInterface + GribJump, and constructs a GribJumpExtractor for the given MARS request. +class GribJumpExtractorDefinition : public ExtractorDefinition { +public: + + GribJumpExtractorDefinition() = default; + GribJumpExtractorDefinition(const GribJumpExtractorDefinition&) = default; + GribJumpExtractorDefinition& operator=(const GribJumpExtractorDefinition&) = default; + GribJumpExtractorDefinition(GribJumpExtractorDefinition&&) = default; + GribJumpExtractorDefinition& operator=(GribJumpExtractorDefinition&&) = default; + + /// FDB config path. std::nullopt uses the environment default (FDB5_CONFIG / FDB_HOME). + std::optional fdbConfig; + /// GribJump config path. std::nullopt reads the GRIBJUMP_CONFIG_FILE env var. + std::optional gribjumpConfig; + /// Chunking strategy for the implicit (grid-point) dimension. + /// Defaults to WholeAxisChunking (single chunk covering the full field/window). + AxisDefinition::ChunkingType fieldChunking{AxisDefinition::WholeAxisChunking{}}; + + void setDefaultIfUnset(const std::optional& fdbConfigPath) override; + + std::unique_ptr copy() const override; + + std::unique_ptr buildExtractor(const metkit::mars::MarsRequest& request) const override; +}; + +} // namespace chunked_data_view diff --git a/src/chunked_data_view/include/chunked_data_view/AxisDefinition.h b/src/chunked_data_view/include/chunked_data_view/AxisDefinition.h index f5e6edd66..ea6d4f0ce 100644 --- a/src/chunked_data_view/include/chunked_data_view/AxisDefinition.h +++ b/src/chunked_data_view/include/chunked_data_view/AxisDefinition.h @@ -2,9 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#include #include #include #include +#include "eckit/exception/Exceptions.h" namespace chunked_data_view { @@ -29,6 +31,7 @@ struct AxisDefinition { /// @p chunkSize must evenly divide the combined axis size, or evenly divide the /// fastest-varying constituent key's value count. struct FixedSizeChunking { + FixedSizeChunking(size_t chunkSize) : chunkSize(chunkSize) { ASSERT(chunkSize > 0); } size_t chunkSize; }; @@ -41,6 +44,9 @@ struct AxisDefinition { /// Chunking strategy applied to this axis. ChunkingType chunking{}; + + /// Optional label used as the zarr dimension name. + std::optional name = std::nullopt; }; } // namespace chunked_data_view diff --git a/src/chunked_data_view/include/chunked_data_view/ChunkedDataView.h b/src/chunked_data_view/include/chunked_data_view/ChunkedDataView.h index dc73857a5..318544536 100644 --- a/src/chunked_data_view/include/chunked_data_view/ChunkedDataView.h +++ b/src/chunked_data_view/include/chunked_data_view/ChunkedDataView.h @@ -14,7 +14,10 @@ namespace chunked_data_view { /// corresponding field values from FDB. /// /// The last dimension is always the implicit field-values dimension (one entry per grid -/// point in a GRIB message); it is never chunked and is always returned as a contiguous block. +/// point in a GRIB message). It forms a single chunk by default, but a GribJump-backed part +/// may subdivide it (see ExtractorType::GribJump's field chunking), in which case its chunk +/// index varies like any other dimension. Its chunk size must divide the grid exactly: that +/// dimension cannot be left ragged. class ChunkedDataView { public: diff --git a/src/chunked_data_view/include/chunked_data_view/ChunkedDataViewBuilder.h b/src/chunked_data_view/include/chunked_data_view/ChunkedDataViewBuilder.h index ab9b592bc..08f119301 100644 --- a/src/chunked_data_view/include/chunked_data_view/ChunkedDataViewBuilder.h +++ b/src/chunked_data_view/include/chunked_data_view/ChunkedDataViewBuilder.h @@ -34,9 +34,18 @@ class ChunkedDataViewBuilder { public: /// @param configPath Optional path to an FDB config file. Passed through to the FDB - /// instance created when building with ExtractorType::GRIB. + /// instance created when building with ExtractorType::Grib. explicit ChunkedDataViewBuilder(const std::optional& configPath = std::nullopt); + + // Deleted to make the ChunkedDataViewBuilder non-copyable. This is also needed by pybind11 layer. + ChunkedDataViewBuilder(const ChunkedDataViewBuilder&) = delete; + ChunkedDataViewBuilder& operator=(const ChunkedDataViewBuilder&) = delete; + + // declaring the deleted copy suppresses implicit moves, so restore them: + ChunkedDataViewBuilder(ChunkedDataViewBuilder&&) = default; + ChunkedDataViewBuilder& operator=(ChunkedDataViewBuilder&&) = default; + /// Registers one data region (part) of the view. /// /// Each keyword in @p marsRequestKeyValues that has more than one value must appear in @@ -46,13 +55,15 @@ class ChunkedDataViewBuilder { /// Multiple parts can cover different variable types (e.g. surface and pressure-level /// fields) and are stitched together along the extension axis specified by extendOnAxis(). /// - /// @p extractor is taken as a shared_ptr because the assembled ChunkedDataViewImpl pairs - /// each ViewPart with its Extractor and multiple ViewParts may legally share one Extractor - /// instance (e.g. when two parts draw from the same FDB store). Shared ownership avoids - /// copying the (potentially stateful, non-copyable) extractor while guaranteeing its - /// lifetime extends at least as long as the built ChunkedDataView. + /// @p definition is an ExtractorDefinition whose buildExtractor() is called once per part + /// inside build() after the MARS request string has been parsed. + /// + /// The builder stores a *copy* of @p definition, so the caller keeps ownership of its + /// object and may register the same configuration on several parts, or on several + /// builders: the per-part defaults the builder applies (see setDefaultIfUnset) are written + /// to the copy only. ChunkedDataViewBuilder& addPart(std::string marsRequestKeyValues, std::vector axes, - std::shared_ptr extractor); + const ExtractorDefinition& definition); /// Sets the axis index along which multiple parts are concatenated. /// @@ -70,25 +81,43 @@ class ChunkedDataViewBuilder { /// @throws eckit::UserError on misconfiguration (missing parts, axis mismatch, etc.). std::unique_ptr build(); - /// Returns the FDB config path supplied at construction, if any. - std::optional getFdbConfigPath() const { return configPath_; } - private: std::optional configPath_{}; - std::vector, std::shared_ptr>> parts_{}; + std::vector, std::unique_ptr>> parts_{}; std::optional extensionAxisIndex_ = std::nullopt; float fillValue_ = std::numeric_limits::quiet_NaN(); - bool doPartsAlign(const std::vector>>& viewParts); + bool doPartsAlign(const std::vector>>& viewParts); - /// Returns true if all parts use the same chunk size on every non-extensible axis. + /// Returns true if all parts use the same chunk size on every axis that is not + /// WholeAxisChunking. /// /// For FixedSizeChunking axes the chunk size must be identical across all parts so /// that part boundaries coincide with Zarr chunk boundaries. WholeAxisChunking axes - /// (extensible=true) grow when parts are stitched together and are therefore exempt - /// from this check. SingleValueChunking always produces chunk size 1 and is trivially - /// consistent. - static bool chunkingConsistencyCheck(const std::vector>>& viewParts); + /// (isSingleGrowingChunk=true) have one chunk per part whose size grows with the + /// combined extent, so differing sizes are expected and exempt from this check. + /// SingleValueChunking always produces chunk size 1 and is trivially consistent. + /// + /// Note: extension (adding parts) is supported for all three chunking types. + static bool chunkingConsistencyCheck(const std::vector>>& viewParts); + + /// Throws unless every part agrees about the implicit (grid-point) dimension. + /// + /// That dimension is never the extension axis, so — like every other non-extension axis — + /// all parts must match on it. Unlike the others it is not derived from the AxisDefinitions + /// but from each extractor's DataLayout, and ChunkedDataViewImpl takes both of its values + /// from the first part alone: + /// - countValues the array's last extent. Differing grids have no representation in + /// one zarr array (the last dimension would have to be ragged). + /// - countChunkValues the chunk's last extent. Differing field chunking means a part + /// writes a differently sized block than the buffer is laid out for. + /// + /// Both are unchecked anywhere else, and both fail *silently* at read time: each extractor + /// sizes its writes from its own layout, so a part with a larger field overruns its slots + /// and still reports the expected message count. + /// + /// @throws eckit::UserError naming the offending part and both values. + static void validateLayouts(const std::vector>>& viewParts); }; } // namespace chunked_data_view diff --git a/src/chunked_data_view/include/chunked_data_view/DataLayout.h b/src/chunked_data_view/include/chunked_data_view/DataLayout.h index 0a1d97411..c73611dc6 100644 --- a/src/chunked_data_view/include/chunked_data_view/DataLayout.h +++ b/src/chunked_data_view/include/chunked_data_view/DataLayout.h @@ -9,8 +9,9 @@ namespace chunked_data_view { /// Describes the binary layout of a single GRIB field's value array. /// All fields within one ViewPart are required to share the same layout. struct DataLayout { - size_t countValues{}; ///< Number of floating-point values in the field (e.g. grid points). - size_t bytesPerValue{}; ///< Storage size of each value in bytes (typically 4 for float32). + size_t countValues{}; ///< Number of floating-point values in the field (e.g. grid points). + size_t bytesPerValue{}; ///< Storage size of each value in bytes (typically 4 for float32). + size_t countChunkValues{}; ///< Number of floating-point values in the chunk (e.g. grid points). }; } // namespace chunked_data_view diff --git a/src/chunked_data_view/include/chunked_data_view/Extractor.h b/src/chunked_data_view/include/chunked_data_view/Extractor.h index 2f4a02ee8..528bf51a4 100644 --- a/src/chunked_data_view/include/chunked_data_view/Extractor.h +++ b/src/chunked_data_view/include/chunked_data_view/Extractor.h @@ -2,12 +2,22 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#include "chunked_data_view/AxisDefinition.h" #include "chunked_data_view/DataLayout.h" +#include "chunked_data_view/Types.h" #include "chunked_data_view/ViewPart.h" + #include "metkit/mars/MarsRequest.h" +#include "chunked_data_view/exception/GribExtractorException.h" +#include "chunked_data_view/exception/GribJumpExtractorException.h" + #include +#include +#include +#include +#include namespace eckit { class DataHandle; @@ -19,17 +29,16 @@ namespace chunked_data_view { /// /// Concrete implementations are typically stateful (e.g. they hold an open FDB handle or /// an HTTP client) and non-copyable by design. For this reason the builder and the assembled -/// ChunkedDataViewImpl always hold extractors via std::shared_ptr: shared ownership lets -/// multiple ViewParts reference the same backing store without copying, and it guarantees -/// the extractor's lifetime is tied to the view rather than to any single part. +/// ChunkedDataViewImpl always hold extractors via std::unique_ptr: each ViewPart owns its +/// Extractor exclusively, and the extractor's lifetime is tied to the view. class Extractor { public: virtual ~Extractor() = default; - /// Retrieves one field for @p req and returns its layout (number of values and bytes per value). - /// All fields in a part are expected to share the same layout; this method establishes it. - virtual DataLayout layout(const metkit::mars::MarsRequest& req) const = 0; + /// Returns the DataLayout established by the concrete extractor at construction time. + /// Describes the size of the implicit (grid-point) dimension and the per-chunk size. + DataLayout layout() const { return layout_; } /// Sets the fill value used to replace GRIB bitmap missing-value sentinels. /// Default no-op; override in concrete extractors that read real field data. @@ -50,12 +59,62 @@ class Extractor { /// @param ptr output buffer to write field values into /// @param len capacity of the output buffer in number of floats /// @return number of GRIB messages written into the buffer - virtual size_t extractInto(const ViewPart& part, const ChunkedDataViewPartBoundingBox& chunkBoundingBox, + virtual size_t extractInto(const ViewPart& part, const ChunkBoundingBox& chunkBoundingBox, const ChunkedDataViewPartBoundingBox& intersectionBoundingBox, float* ptr, size_t len) const = 0; + +protected: + + /// Populated by each concrete extractor's constructor; returned by layout(). + DataLayout layout_{}; }; -enum class ExtractorType { - GRIB + +/// Abstract factory that produces a concrete Extractor for a given (already-parsed) +/// MARS request. ChunkedDataViewBuilder::build() calls buildExtractor() once per +/// registered part after it has parsed and validated the request string. +/// +/// Concrete definitions (GribExtractorDefinition, GribJumpExtractorDefinition) carry +/// their own backend-specific configuration (FDB path, fill value, etc.) as public +/// data members. They are also the user-facing configuration objects: the Python layer +/// binds them as ExtractorType.Grib / ExtractorType.GribJump. +class ExtractorDefinition { +public: + + virtual ~ExtractorDefinition() = default; + + /// Adopts @p fdbConfigPath as this definition's FDB config unless one was set explicitly. + /// + /// Called by ChunkedDataViewBuilder::addPart() so that a definition which names no FDB + /// config inherits the builder's, without the builder having to know which backend it + /// is holding. An empty path on both sides leaves FDB to resolve its own configuration + /// from the environment (FDB5_CONFIG / FDB_HOME). + virtual void setDefaultIfUnset(const std::optional& fdbConfigPath) = 0; + + /// Returns an independent copy of this definition. + /// + /// ChunkedDataViewBuilder::addPart() stores a copy rather than the caller's object, so one + /// configuration can be registered on several parts (and several builders) without the + /// builder's own defaults leaking back into it. + virtual std::unique_ptr copy() const = 0; + + /// Construct the concrete Extractor for the given MARS request. + /// Called exactly once per part by ChunkedDataViewBuilder::build(). + virtual std::unique_ptr buildExtractor(const metkit::mars::MarsRequest& request) const = 0; + +protected: + + /// Definitions are plain, copyable configuration objects. ChunkedDataViewBuilder::addPart() + /// takes ownership of the definition it is handed, so a caller that wants to reuse one + /// configuration across several parts hands over a copy of its concrete definition. + /// + /// These members are protected rather than public so that a copy can only be made through + /// a concrete definition; copying via an ExtractorDefinition& would slice. + ExtractorDefinition() = default; + ExtractorDefinition(const ExtractorDefinition&) = default; + ExtractorDefinition& operator=(const ExtractorDefinition&) = default; + ExtractorDefinition(ExtractorDefinition&&) = default; + ExtractorDefinition& operator=(ExtractorDefinition&&) = default; }; + } // namespace chunked_data_view diff --git a/src/chunked_data_view/include/chunked_data_view/ListIterator.h b/src/chunked_data_view/include/chunked_data_view/ListIterator.h index e9a8a2c1a..7fa1fda14 100644 --- a/src/chunked_data_view/include/chunked_data_view/ListIterator.h +++ b/src/chunked_data_view/include/chunked_data_view/ListIterator.h @@ -9,14 +9,24 @@ #include #include -#include -#include namespace chunked_data_view { +/// A field entry returned by a list/inspect call: the MARS key identifying the field +/// and a pointer to its storage location (from which a DataHandle can be opened). +struct ListElement { + fdb5::Key key; + std::shared_ptr location; + + /// Opens a new DataHandle for this field. + std::unique_ptr dataHandle() const { + return std::unique_ptr(location->dataHandle()); + } +}; + /// Abstract iterator over FDB fields matching a MARS request. /// -/// Each call to next() yields the MARS key and a data handle for one matching field, +/// Each call to next() yields a ListElement for one matching field, /// or std::nullopt when the sequence is exhausted. class ListIteratorInterface { @@ -24,8 +34,8 @@ class ListIteratorInterface { virtual ~ListIteratorInterface() = default; - /// Returns the next (key, data-handle) pair, or std::nullopt if there are no more fields. - virtual std::optional>> next() = 0; + /// Returns the next field entry, or std::nullopt if there are no more fields. + virtual std::optional next() = 0; }; @@ -37,7 +47,7 @@ class ListIteratorWrapperImpl : public ListIteratorInterface { public: explicit ListIteratorWrapperImpl(fdb5::ListIterator listIterator) : listIterator_(std::move(listIterator)) {}; - std::optional>> next() override; + std::optional next() override; }; /// Wraps @p listIterator in a ListIteratorInterface-compatible heap object. diff --git a/src/chunked_data_view/include/chunked_data_view/Types.h b/src/chunked_data_view/include/chunked_data_view/Types.h index aad322401..fd5465168 100644 --- a/src/chunked_data_view/include/chunked_data_view/Types.h +++ b/src/chunked_data_view/include/chunked_data_view/Types.h @@ -10,6 +10,9 @@ class BoundingBox; /// Bounding box of a Zarr chunk in the global ChunkedDataView index space. using ChunkedDataViewPartBoundingBox = chunked_data_view::BoundingBox; +/// Bounding box of a Zarr chunk in the global ChunkedDataView index space (incl. implicit dimension). +using ChunkBoundingBox = chunked_data_view::BoundingBox; + /// Bounding box expressed in the coordinate space of a single ViewPart /// (i.e. with the part's offset subtracted so that the part's own lower corner is the origin). using PartBoundingBox = chunked_data_view::BoundingBox; diff --git a/src/chunked_data_view/include/chunked_data_view/exception/GribJumpExtractorException.h b/src/chunked_data_view/include/chunked_data_view/exception/GribJumpExtractorException.h new file mode 100644 index 000000000..84d296823 --- /dev/null +++ b/src/chunked_data_view/include/chunked_data_view/exception/GribJumpExtractorException.h @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 +#pragma once +#include "eckit/exception/Exceptions.h" + +namespace chunked_data_view { + +/// Thrown by GribJumpExtractor when field retrieval fails (e.g. the FDB iterator returns no +/// results for a request, or a GribJump extraction result does not match the expected layout). +class GribJumpExtractorException : public eckit::Exception { + +public: + + GribJumpExtractorException(const std::string&); + GribJumpExtractorException(const std::string&, const eckit::CodeLocation&); +}; +} // namespace chunked_data_view diff --git a/src/chunked_data_view_bindings/CMakeLists.txt b/src/chunked_data_view_bindings/CMakeLists.txt index 9a772c54c..8f934d710 100644 --- a/src/chunked_data_view_bindings/CMakeLists.txt +++ b/src/chunked_data_view_bindings/CMakeLists.txt @@ -1,8 +1,14 @@ pybind11_add_module(chunked_data_view_bindings MODULE bindings.cc) +# bindings.cc itself needs no gribjump header; the link is only to resolve the symbols pulled +# in from the chunked_data_view static library when the extractor is compiled. +if(HAVE_ZARR_GRIBJUMP_EXTRACTOR) + target_link_libraries(chunked_data_view_bindings PRIVATE gribjump) +endif() + target_link_libraries(chunked_data_view_bindings - PRIVATE pybind11::module - pybind11::lto + PRIVATE pybind11::module + pybind11::lto chunked_data_view eckit metkit @@ -21,10 +27,4 @@ set_target_properties(chunked_data_view_bindings LIBRARY_OUTPUT_DIRECTORY ${Z3FDB_STAGING}/chunked_data_view_bindings ) -add_custom_command(TARGET chunked_data_view_bindings POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - ${CMAKE_CURRENT_SOURCE_DIR}/__init__.py - $ - COMMENT "Copying __init__.py for chunked_data_view_bindings" -) diff --git a/src/chunked_data_view_bindings/__init__.py b/src/chunked_data_view_bindings/__init__.py index 45216d72f..ef1dd12da 100644 --- a/src/chunked_data_view_bindings/__init__.py +++ b/src/chunked_data_view_bindings/__init__.py @@ -7,11 +7,15 @@ findlibs.load("fdb5") -from chunked_data_view_bindings.chunked_data_view_bindings import ( +from .chunked_data_view_bindings import ( init_bindings, AxisDefinition, ChunkedDataView, ChunkedDataViewBuilder, + ExtractorDefinition, + GribExtractorError, + GribJumpExtractorError, + has_gribjump_extractor, ExtractorType, ) @@ -23,5 +27,9 @@ "AxisDefinition", "ChunkedDataView", "ChunkedDataViewBuilder", + "ExtractorDefinition", + "GribExtractorError", + "GribJumpExtractorError", + "has_gribjump_extractor", "ExtractorType", ] diff --git a/src/chunked_data_view_bindings/bindings.cc b/src/chunked_data_view_bindings/bindings.cc index b33b73392..3841dbf21 100644 --- a/src/chunked_data_view_bindings/bindings.cc +++ b/src/chunked_data_view_bindings/bindings.cc @@ -1,8 +1,9 @@ // SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) // SPDX-License-Identifier: Apache-2.0 -#include -#include +#include +#include +#include #include #include @@ -15,26 +16,35 @@ #include "chunked_data_view/ChunkedDataView.h" #include "chunked_data_view/ChunkedDataViewBuilder.h" #include "chunked_data_view/Extractor.h" -#include "chunked_data_view/Fdb.h" -#include "chunked_data_view/GribExtractor.h" #include "chunked_data_view/LibChunkedDataView.h" -#include "chunked_data_view/exception/UnknownExtractorException.h" +#include "chunked_data_view/extractors/grib/GribExtractorDefinition.h" +#include "chunked_data_view/extractors/gribjump/GribJumpExtractorDefinition.h" namespace py = pybind11; namespace cdv = chunked_data_view; +namespace docs { + +inline constexpr auto chunked_data_view_module_doc = R"doc( +Low-level pybind11 bindings for the chunked_data_view C++ library. +This module exposes the core types needed to build a Zarr-compatible N-dimensional +view over FDB data. It is not part of the public API; users should import from +``pychunked_data_view`` or ``z3fdb`` instead. +Typical call sequence:: + init_bindings() + builder = ChunkedDataViewBuilder(fdb_config_path) + builder.add_part(mars_request_string, [AxisDefinition(...)], ExtractorType.Grib()) + view = builder.build() + chunk = view.at([0, 1, 0]) # numpy array of float32 +)doc"; +}; + +/// Empty tag type used as the pybind11 handle for the ExtractorType namespace object. +/// cdv::ExtractorType itself has no data members +struct ExtractorTypeNamespace {}; + PYBIND11_MODULE(chunked_data_view_bindings, m) { - m.doc() = - "Low-level pybind11 bindings for the chunked_data_view C++ library.\n\n" - "This module exposes the core types needed to build a Zarr-compatible N-dimensional\n" - "view over FDB data. It is not part of the public API; users should import from\n" - "``pychunked_data_view`` or ``z3fdb`` instead.\n\n" - "Typical call sequence::\n\n" - " init_bindings()\n" - " builder = ChunkedDataViewBuilder(fdb_config_path)\n" - " builder.add_part(mars_request_string, [AxisDefinition(...)], ExtractorType.GRIB)\n" - " view = builder.build()\n" - " chunk = view.at([0, 1, 0]) # numpy array of float32\n"; + m.doc() = docs::chunked_data_view_module_doc; m.def( "init_bindings", []() { cdv::init_eckit_main(); }, @@ -42,6 +52,19 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { "Must be called exactly once before any other function or class in this module is\n" "used. Subsequent calls are safe but have no effect."); + // Build capability. ExtractorType.GribJump is always present so that user code does not + // depend on build flags; this reports whether it can actually be used. +#ifdef HAVE_ZARR_GRIBJUMP_EXTRACTOR + m.attr("has_gribjump_extractor") = true; +#else + m.attr("has_gribjump_extractor") = false; +#endif + + // Exception registration. Names are re-exported from chunked_data_view_bindings/__init__.py + // and pychunked_data_view, so user code can catch them by type. + py::register_local_exception(m, "GribJumpExtractorError"); + py::register_local_exception(m, "GribExtractorError"); + // Axis Definition and subclasses auto axis_definition = py::class_( m, "AxisDefinition", @@ -75,40 +98,29 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { py::class_( axis_definition, "FixedSizeChunking", "Chunking strategy where the axis is divided into chunks of a fixed size.\n\n" - "The chunk size must be a valid trailing-product divisor of the combined axis\n" - "size; see ``AxisMapper::chunkSizeCheck`` for the exact rule. Use this when\n" - "a tuned chunk size is needed to balance read amplification against request\n" - "count (e.g. grouping every 3 time steps into one chunk).\n\n" + "The chunk size must be a valid trailing-product divisor of the combined axis count\n" + "(e.g. grouping every 3 time steps into one chunk). Use this when a tuned chunk size\n" + "is needed to balance read amplification against the number of requests.\n\n" "Args:\n" - " chunk_size: Number of axis elements per chunk.") - .def(py::init<>([](size_t& chunkExtension) { - return cdv::AxisDefinition::FixedSizeChunking{.chunkSize = chunkExtension}; - }), + " chunk_size: Number of axis elements per chunk. Must be greater than zero.") + .def(py::init<>([](size_t& chunkExtension) { return cdv::AxisDefinition::FixedSizeChunking(chunkExtension); }), py::arg("chunk_size")) - .def( + .def_property_readonly( "chunk_shape", - [](const cdv::AxisDefinition::FixedSizeChunking* fixedSizeChunking) { - return fixedSizeChunking->chunkSize; - }, + [](const cdv::AxisDefinition::FixedSizeChunking& fixedSizeChunking) { return fixedSizeChunking.chunkSize; }, "int: Number of axis elements per chunk."); axis_definition - .def(py::init([](std::vector keys, cdv::AxisDefinition::ChunkingType chunking) { - return cdv::AxisDefinition{std::move(keys), chunking}; + .def(py::init([](std::vector keys, cdv::AxisDefinition::ChunkingType chunking, + std::optional name) { + return cdv::AxisDefinition{std::move(keys), chunking, std::move(name)}; }), - py::kw_only(), py::arg("keys"), py::arg("chunking")) - .def(py::init([](std::vector keys, cdv::AxisDefinition::FixedSizeChunking chunking) { - return cdv::AxisDefinition{std::move(keys), chunking}; - }), - py::kw_only(), py::arg("keys"), py::arg("chunking")) - .def(py::init([](std::vector keys, cdv::AxisDefinition::WholeAxisChunking chunking) { - return cdv::AxisDefinition{std::move(keys), chunking}; - }), - py::kw_only(), py::arg("keys"), py::arg("chunking")) + py::kw_only(), py::arg("keys"), py::arg("chunking"), py::arg("name") = py::none()) .def_readwrite("keys", &cdv::AxisDefinition::keys, "list[str]: Ordered MARS keyword names that form this axis.") .def_readwrite("chunking", &cdv::AxisDefinition::chunking, "WholeAxisChunking | SingleValueChunking | FixedSizeChunking: " - "Chunking strategy applied to this axis."); + "Chunking strategy applied to this axis.") + .def_readwrite("name", &cdv::AxisDefinition::name, "str | None: Optional zarr dimension name."); // ChunkedDataView py::class_(m, "ChunkedDataView", @@ -118,7 +130,9 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { "triggers one or more FDB retrievals and assembles the field values into\n" "a contiguous ``float32`` numpy array.\n\n" "The last dimension is the implicit grid-point dimension (number of values\n" - "per GRIB message). It is never chunked and is always returned in full.\n\n" + "per GRIB message). It is a single chunk by default; a GribJump-backed part\n" + "may subdivide it via field_chunking, in which case its chunk index varies\n" + "like any other dimension.\n\n" "Instances are created exclusively by :meth:`ChunkedDataViewBuilder.build`.") .def( "at", @@ -126,15 +140,18 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { const auto len = view->countChunkValues(); py::array_t arr(len); float* p = arr.mutable_data(); - view->at(index, p, len); + { + py::gil_scoped_release release; + view->at(index, p, len); + } return arr; }, py::arg("index"), "Return the data for the chunk at *index* as a 1-D ``float32`` numpy array.\n\n" "Args:\n" - " index (list[int]): Chunk-grid coordinates, one entry per dimension\n" - " (including the implicit grid-point dimension whose\n" - " index must always be 0).\n\n" + " index (list[int]): Chunk-grid coordinates, one entry per dimension,\n" + " including the implicit grid-point dimension (whose index\n" + " is 0 unless the part uses field_chunking).\n\n" "Returns:\n" " numpy.ndarray: 1-D float32 array of length ``chunk_shape()[-1]``\n" " containing the field values for this chunk.\n\n" @@ -142,40 +159,114 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { " RuntimeError: If the FDB retrieval fails or the index is out of bounds.") .def( "chunk_shape", [](const cdv::ChunkedDataView* view) { return view->chunkShape(); }, + py::call_guard(), "tuple[int, ...]: Number of elements per chunk in each dimension (Zarr chunk shape).") .def( "chunks", [](const cdv::ChunkedDataView* view) { return view->chunks(); }, + py::call_guard(), "tuple[int, ...]: Number of chunks along each dimension of the chunk grid.") .def( "shape", [](const cdv::ChunkedDataView* view) { return view->shape(); }, + py::call_guard(), "tuple[int, ...]: Total number of elements along each dimension of the full array.") .def( "fill_missing_value", [](const cdv::ChunkedDataView* view) { return view->fillMissingValue(); }, + py::call_guard(), "float: Value written for array positions not covered by any data part\n" " (default: ``float('nan')``)."); - // ExtractorType - py::enum_(m, "ExtractorType", - "Selects the data format expected in FDB.\n\n" - "Passed to :meth:`ChunkedDataViewBuilder.add_part` to control which\n" - ":class:`Extractor` implementation is instantiated for a data part.") - .value("GRIB", cdv::ExtractorType::GRIB, - "Data is stored as GRIB messages. The extractor reads each message,\n" - "validates the returned paramIds against the request, and copies the\n" - "``values`` array into the chunk buffer."); - - // ChunkedDataViewBuilder + // Extractor definitions. These *are* the user-facing configuration objects: the abstract + // base is registered so that pybind11 knows the hierarchy, and each concrete definition is + // exposed under the ExtractorType namespace object. + py::class_( + m, "ExtractorDefinition", + "Base class of every extractor configuration.\n\n" + "Not instantiable — use :class:`ExtractorType.Grib` or :class:`ExtractorType.GribJump`."); + + auto extractor_type_ns = py::class_( + m, "ExtractorType", + "Namespace for extractor configuration types.\n\n" + "* :class:`ExtractorType.Grib` — standard full-field GRIB extraction.\n" + "* :class:`ExtractorType.GribJump` — partial-field extraction via GribJump."); + + py::class_( + extractor_type_ns, "Grib", + "Configuration for full-field GRIB extraction.\n\n" + "Args:\n" + " fdb_config (pathlib.Path | None): Path to the FDB configuration YAML.\n" + " ``None`` (default) uses the builder's FDB config.") + .def(py::init([](std::optional fdbConfig) { + cdv::GribExtractorDefinition definition{}; + definition.fdbConfig = std::move(fdbConfig); + return definition; + }), + py::kw_only(), py::arg("fdb_config") = py::none()) + .def_readwrite("fdb_config", &cdv::GribExtractorDefinition::fdbConfig, + "pathlib.Path | None: Path to the FDB configuration YAML."); + + py::class_( + extractor_type_ns, "GribJump", + "Configuration for partial-field extraction via GribJump.\n\n" + "GribJump avoids a full GRIB decode by jumping directly to the grid-point\n" + "values inside each message. The implicit (grid-point) dimension can be\n" + "split into equal-sized Zarr sub-chunks via *field_chunking*::\n\n" + " ExtractorType.GribJump(\n" + " field_chunking=FixedSizeChunking(1312), # splits 5248 values into 4 chunks\n" + " )\n\n" + "Args:\n" + " fdb_config (pathlib.Path | None): Path to the FDB configuration YAML.\n" + " ``None`` (default) uses the builder's FDB config.\n" + " gribjump_config (pathlib.Path | None): Path to the GribJump configuration YAML.\n" + " ``None`` (default) reads the ``GRIBJUMP_CONFIG_FILE`` environment variable.\n" + " field_chunking (WholeAxisChunking | SingleValueChunking | FixedSizeChunking | None):\n" + " How to sub-divide the implicit (grid-point) dimension into Zarr chunks.\n" + " ``None`` or :class:`WholeAxisChunking` (default) produces a single chunk\n" + " covering the full field. :class:`FixedSizeChunking` splits it into\n" + " equal-sized pieces.") + .def(py::init([](std::optional fdbConfig, + std::optional gribjumpConfig, + std::optional fieldChunking) { + cdv::GribJumpExtractorDefinition definition{}; + definition.fdbConfig = std::move(fdbConfig); + definition.gribjumpConfig = std::move(gribjumpConfig); + // The variant caster rejects anything that is not a chunking type, so no + // per-type checking is needed here. + if (fieldChunking.has_value()) { + definition.fieldChunking = std::move(*fieldChunking); + } + return definition; + }), + py::kw_only(), py::arg("fdb_config") = py::none(), py::arg("gribjump_config") = py::none(), + py::arg("field_chunking") = py::none()) + .def_readwrite("fdb_config", &cdv::GribJumpExtractorDefinition::fdbConfig, + "pathlib.Path | None: Path to the FDB configuration YAML.") + .def_readwrite("gribjump_config", &cdv::GribJumpExtractorDefinition::gribjumpConfig, + "pathlib.Path | None: Path to the GribJump configuration YAML.") + .def_readwrite("field_chunking", &cdv::GribJumpExtractorDefinition::fieldChunking, + "WholeAxisChunking | SingleValueChunking | FixedSizeChunking: " + "Chunking of the implicit grid-point dimension."); + py::class_( m, "ChunkedDataViewBuilder", "Fluent builder that constructs a :class:`ChunkedDataView` from one or more\n" "MARS data parts.\n\n" - "Usage::\n\n" + "**Standard GRIB extraction** (full field, one chunk per field)::\n\n" " builder = ChunkedDataViewBuilder(fdb_config_path)\n" " builder.add_part(\n" " 'class=ea,type=an,...,param=167/131',\n" " [AxisDefinition(keys=['date','time'], chunking=SingleValueChunking()),\n" " AxisDefinition(keys=['param'], chunking=SingleValueChunking())],\n" - " ExtractorType.GRIB,\n" + " ExtractorType.Grib(),\n" + " )\n" + " view = builder.build()\n\n" + "**GribJump extraction with implicit-axis chunking** — splits the grid-point\n" + "dimension into fixed-size Zarr chunks::\n\n" + " builder = ChunkedDataViewBuilder(fdb_config_path)\n" + " builder.add_part(\n" + " 'class=ea,type=an,...,param=167/131',\n" + " [AxisDefinition(keys=['date','time'], chunking=FixedSizeChunking(8)),\n" + " AxisDefinition(keys=['param'], chunking=SingleValueChunking())],\n" + " ExtractorType.GribJump(field_chunking=FixedSizeChunking(1312)),\n" " )\n" " view = builder.build()\n\n" "When more than one part is added (e.g. surface and pressure-level fields),\n" @@ -189,41 +280,35 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { " fdb_config_path (pathlib.Path | None): Path to an FDB configuration file.\n" " ``None`` (default) lets FDB resolve its configuration from the\n" " environment (``FDB5_CONFIG`` / ``FDB_HOME``).") - .def( - "add_part", - [](cdv::ChunkedDataViewBuilder& builder, std::string marsRequestKeyValues, - std::vector axes, const cdv::ExtractorType extractorType) { - switch (extractorType) { - case chunked_data_view::ExtractorType::GRIB: { - auto fdb = cdv::makeFdb(builder.getFdbConfigPath()); - auto extractor = std::make_shared( - chunked_data_view::GribExtractor(std::move(fdb))); - builder.addPart(std::move(marsRequestKeyValues), std::move(axes), std::move(extractor)); - break; - } - default: - std::stringstream buf; - buf << "ChunkedDataViewBuidler::add_part: Unknown Extractor of type " << std::endl; - throw cdv::UnknownExtractorException(buf.str()); - } - }, - py::arg("mars_request"), py::arg("axes"), py::arg("extractor_type"), - "Register one data region (part) of the view.\n\n" - "Each MARS keyword in *mars_request* that carries more than one value must\n" - "appear in exactly one :class:`AxisDefinition` in *axes*. Multiple keywords\n" - "may share one axis and are combined as a Cartesian product.\n\n" - "Args:\n" - " mars_request (str): Comma-separated ``key=value[/value...]``\n" - " MARS request string, e.g.\n" - " ``'class=ea,param=167/131,date=20200101/20200102'``.\n" - " axes (list[AxisDefinition]): Axis definitions covering every\n" - " multi-valued key in the request.\n" - " extractor_type (ExtractorType): Format of the data in FDB.\n\n" - "Raises:\n" - " RuntimeError: If FDB cannot retrieve a sample field for the request, or if\n" - " a returned paramId does not match the requested params (which\n" - " indicates unsupported on-the-fly field derivation).") + .def("add_part", &cdv::ChunkedDataViewBuilder::addPart, py::arg("mars_request"), py::arg("axes"), + py::arg("extractor_type"), py::return_value_policy::reference_internal, + "Register one part of the view.\n\n" + "Each MARS keyword in *mars_request* that carries more than one value must\n" + "appear in exactly one :class:`AxisDefinition` in *axes*. Multiple keywords\n" + "may share one axis and are combined as a Cartesian product.\n\n" + "Args:\n" + " mars_request (str): Comma-separated ``key=value[/value...]``\n" + " MARS request string, e.g.\n" + " ``'class=ea,param=167/131,date=20200101/20200102'``.\n" + " axes (list[AxisDefinition]): Axis definitions covering every\n" + " multi-valued key in the request.\n" + " extractor_type (ExtractorType.Grib | ExtractorType.GribJump):\n" + " Extractor configuration object.\n" + " ``ExtractorType.Grib`` reads the full GRIB field.\n" + " ``ExtractorType.GribJump`` supports splitting\n" + " the implicit grid-point axis into Zarr sub-chunks\n" + " (``field_chunking``).\n\n" + " The builder stores a copy, so one extractor can\n" + " be reused across parts and builders; defaults the\n" + " builder applies (e.g. its fdb_config) are not\n" + " written back into your object.\n\n" + "Raises:\n" + " RuntimeError: If FDB cannot retrieve a sample field for the request, or if\n" + " a returned paramId does not match the requested params.\n" + " TypeError: If *extractor* is not an ``ExtractorType.Grib`` or\n" + " ``ExtractorType.GribJump`` instance.") .def("extend_on_axis", &cdv::ChunkedDataViewBuilder::extendOnAxis, py::arg("axis"), + py::return_value_policy::reference_internal, "Declare the axis index along which multiple parts are concatenated.\n\n" "Required when more than one part is added; ignored for single-part views.\n" "All parts must have identical extents on every axis except this one.\n\n" @@ -232,6 +317,7 @@ PYBIND11_MODULE(chunked_data_view_bindings, m) { "Raises:\n" " RuntimeError: If *axis* is out of range for the first part's axis list.") .def("fill_missing_value", &cdv::ChunkedDataViewBuilder::fillMissingValue, py::arg("value"), + py::return_value_policy::reference_internal, "Set the fill value for array positions not covered by any data part.\n\n" "Args:\n" " value (float): Fill value (default: ``float('nan')``).") diff --git a/src/pychunked_data_view/__init__.py b/src/pychunked_data_view/__init__.py index d4d53f384..e6b8893c7 100644 --- a/src/pychunked_data_view/__init__.py +++ b/src/pychunked_data_view/__init__.py @@ -1,6 +1,11 @@ # SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) # SPDX-License-Identifier: Apache-2.0 +from chunked_data_view_bindings import ( # noqa: E402 + GribExtractorError, + GribJumpExtractorError, + has_gribjump_extractor, +) from pychunked_data_view.chunked_data_view import ( # noqa: E402 AxisDefinition, ChunkedDataView, @@ -9,6 +14,7 @@ ExtractorType, MarsSelection, ) +from pychunked_data_view.exceptions import InternalError, MarsRequestFormattingError # noqa: E402 __all__ = [ "AxisDefinition", @@ -16,5 +22,10 @@ "ChunkedDataViewBuilder", "Chunking", "ExtractorType", + "GribExtractorError", + "GribJumpExtractorError", + "InternalError", + "MarsRequestFormattingError", "MarsSelection", + "has_gribjump_extractor", ] diff --git a/src/pychunked_data_view/chunked_data_view.py b/src/pychunked_data_view/chunked_data_view.py index ab0b756df..ccc550bb1 100644 --- a/src/pychunked_data_view/chunked_data_view.py +++ b/src/pychunked_data_view/chunked_data_view.py @@ -1,21 +1,35 @@ # SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Collection, Mapping -from dataclasses import dataclass import enum import pathlib +import warnings +from collections.abc import Collection, Mapping from typing import TypeAlias -import chunked_data_view_bindings.chunked_data_view_bindings as pdv +import numpy -from pychunked_data_view.exceptions import MarsRequestFormattingError, InternalError +import chunked_data_view_bindings as pdv +from chunked_data_view_bindings import ( # noqa: E402 + GribExtractorError as GribExtractorError, + GribJumpExtractorError as GribJumpExtractorError, + has_gribjump_extractor as has_gribjump_extractor, +) +from pychunked_data_view.exceptions import InternalError, MarsRequestFormattingError MarsSelection: TypeAlias = Mapping[str, str | int | float | Collection[str | int | float]] -# Mapping functionality for MarsSelection def _mars_selection_to_string(request: MarsSelection) -> str: + """Serialise a :class:`MarsSelection` dict to a ``key=value,...`` MARS string. + + Args: + request (MarsSelection): MARS key-value mapping to serialise. + + Returns: + str: Comma-separated ``key=value`` pairs, with multi-valued entries + joined by ``/``. + """ parts = [] for key, value in request.items(): if isinstance(value, (str, int, float)): @@ -32,19 +46,22 @@ class Chunking(enum.Enum): Attributes: WHOLE_AXIS: The entire axis is a single chunk; accessing any value loads all values on that axis. SINGLE_VALUE: Each value along the axis is its own chunk. - FixedSizeChunk: Groups every ``chunkShape`` consecutive values into one chunk. + FixedSizeChunk: Groups every ``chunk_shape`` consecutive values into one chunk. """ WHOLE_AXIS = enum.auto() SINGLE_VALUE = enum.auto() @enum.nonmember - @dataclass(frozen=True) class FixedSizeChunk: - chunkShape: int + def __init__(self, chunk_shape: int) -> None: + assert chunk_shape > 0, "The supplied chunk shape needs to be positive" + self.chunk_shape = chunk_shape class AxisDefinition: + """Maps one or more MARS keys to a single zarr array axis with a given chunking strategy.""" + @staticmethod def _translate_chunking( chunking: Chunking | Chunking.FixedSizeChunk, @@ -53,8 +70,20 @@ def _translate_chunking( | pdv.AxisDefinition.SingleValueChunking | pdv.AxisDefinition.FixedSizeChunking ): + """Convert a Python :class:`Chunking` value to the corresponding C++ binding type. + + Args: + chunking (~pychunked_data_view.Chunking | ~pychunked_data_view.Chunking.FixedSizeChunk): + Chunking strategy to translate. + + Returns: + The matching ``pdv.AxisDefinition`` chunking object. + + Raises: + TypeError: If *chunking* is not a recognised :class:`Chunking` value. + """ if isinstance(chunking, Chunking.FixedSizeChunk): - return pdv.AxisDefinition.FixedSizeChunking(chunking.chunkShape) + return pdv.AxisDefinition.FixedSizeChunking(chunking.chunk_shape) elif chunking is Chunking.WHOLE_AXIS: return pdv.AxisDefinition.WholeAxisChunking() elif chunking is Chunking.SINGLE_VALUE: @@ -65,19 +94,34 @@ def _translate_chunking( f"Chunking.FixedSizeChunk, got {type(chunking).__qualname__!r}" ) - def __init__(self, keys: list[str], chunking: Chunking | Chunking.FixedSizeChunk): - """Defines which axis from a MARS Request form an axis in the Zarr array. - - Also defines if the data is to be chunked. + def __init__( + self, + keys: list[str], + chunking: Chunking | Chunking.FixedSizeChunk, + name: str | None = None, + ): + """Defines which MARS keys form an axis in the zarr array, and how it is chunked. Args: - keys(list of str): mars keys that for this axis. - chunking ( Chunking): Define how this axis shall be chunked + keys (list[str]): MARS keys that form this axis. + chunking (~pychunked_data_view.Chunking | ~pychunked_data_view.Chunking.FixedSizeChunk): + How this axis shall be chunked. + name (str | None): Zarr dimension name. Defaults to the keys joined by ``"_"``. """ - self._obj = pdv.AxisDefinition(keys=keys, chunking=self._translate_chunking(chunking)) + self._obj = pdv.AxisDefinition(keys=keys, chunking=self._translate_chunking(chunking), name=name) + + @property + def name(self) -> str | None: + """The zarr dimension name for this axis, or None to derive it from the keys.""" + return self._obj.name + + @name.setter + def name(self, name: str | None) -> None: + self._obj.name = name @property def keys(self) -> list[str]: + """The MARS keys that form this axis.""" return self._obj.keys @keys.setter @@ -85,14 +129,20 @@ def keys(self, keys: list[str]) -> None: self._obj.keys = keys @property - def chunking(self): + def chunking(self) -> Chunking | Chunking.FixedSizeChunk: + """The chunking strategy for this axis. + + Raises: + ~pychunked_data_view.exceptions.InternalError: If the underlying C++ chunking + type is unrecognised. + """ chunking = self._obj.chunking if isinstance(chunking, pdv.AxisDefinition.WholeAxisChunking): return Chunking.WHOLE_AXIS elif isinstance(chunking, pdv.AxisDefinition.SingleValueChunking): return Chunking.SINGLE_VALUE elif isinstance(chunking, pdv.AxisDefinition.FixedSizeChunking): - return Chunking.FixedSizeChunk(chunking.chunk_shape()) + return Chunking.FixedSizeChunk(chunking.chunk_shape) else: raise InternalError() @@ -102,58 +152,236 @@ def chunking(self, chunking: Chunking | Chunking.FixedSizeChunk) -> None: class ChunkedDataView: + """Python wrapper around the C++ ``ChunkedDataView``. + + Provides shape and chunk-count metadata, and per-chunk data access via + :meth:`at`. Instances are returned by :meth:`ChunkedDataViewBuilder.build`. + """ + def __init__(self, obj: pdv.ChunkedDataView): self._obj = obj - def at(self, index: list[int] | tuple[int, ...]): + def at(self, index: list[int] | tuple[int, ...]) -> "numpy.ndarray": + """Return the values of the chunk at *index*. + + Args: + index (list[int] | tuple[int, ...]): Per-dimension chunk coordinates, including + the implicit grid-point dimension. + + Returns: + numpy.ndarray: 1-D ``float32`` array of ``chunk_shape()`` values, C-order. + + Raises: + RuntimeError: If *index* is out of bounds or the FDB retrieval fails. + """ return self._obj.at(index) - def chunkShape(self): + def chunk_shape(self) -> tuple[int, ...]: + """Return the per-dimension element count of one chunk. + + Returns: + tuple[int, ...]: Number of elements along each dimension within a single chunk. + """ return self._obj.chunk_shape() - def chunks(self): + def chunkShape(self) -> tuple[int, ...]: + """Deprecated alias of :meth:`chunk_shape`. + + Kept so existing callers keep working; every other method on this class is + snake_case. + """ + warnings.warn( + "ChunkedDataView.chunkShape() is deprecated, use chunk_shape() instead", + DeprecationWarning, + stacklevel=2, + ) + return self.chunk_shape() + + def chunks(self) -> tuple[int, ...]: + """Return the per-dimension number of chunks. + + Returns: + tuple[int, ...]: Number of chunks along each dimension. + """ return self._obj.chunks() - def shape(self): + def shape(self) -> tuple[int, ...]: + """Return the total array shape in elements (not chunks). + + Returns: + tuple[int, ...]: Total number of elements along each dimension. + """ return self._obj.shape() - def fill_missing_value(self): + def fill_missing_value(self) -> float: + """Return the fill value used for bitmap-masked grid points. + + Returns: + float: Value written into positions flagged as missing by the GRIB bitmap. + """ return self._obj.fill_missing_value() -class ExtractorType(enum.Enum): - """Suported data extractors. +class ExtractorType: + """Namespace for extractor configuration types. + + * :class:`ExtractorType.Grib` - standard full-field GRIB extraction. + * :class:`ExtractorType.GribJump` - partial-field extraction via GribJump. + + Each class wraps the matching C++ ``ExtractorDefinition``, which is what + :meth:`ChunkedDataViewBuilder.add_part` takes. - Defines what storage format the caller expects to be stored in FDB. + One instance may be reused across as many parts and builders as you like: + ``add_part`` stores a copy, so defaults the builder applies (e.g. its + ``fdb_config``) are never written back into your object. """ - GRIB = pdv.ExtractorType.GRIB - """Extract data from GRIB""" + class Grib: + """Reads full GRIB fields from FDB. + + Args: + fdb_config (pathlib.Path | None): Path to the FDB configuration YAML. + ``None`` (default) uses the builder's FDB config. + """ + + def __init__(self, fdb_config: pathlib.Path | None = None): + self._obj = pdv.ExtractorType.Grib(fdb_config=fdb_config) + + class GribJump: + """Reads grid-point values from FDB via GribJump. + + GribJump avoids a full GRIB decode by jumping directly to the + grid-point values inside each message. + + Args: + fdb_config (pathlib.Path | None): Path to the FDB configuration YAML. + ``None`` (default) uses the builder's FDB config. + gribjump_config (pathlib.Path | None): Path to the GribJump configuration YAML. + ``None`` (default) reads the ``GRIBJUMP_CONFIG_FILE`` environment variable. + field_chunking (~pychunked_data_view.Chunking | ~pychunked_data_view.Chunking.FixedSizeChunk | None): + How to sub-divide the implicit (grid-point) dimension into Zarr chunks. + ``None`` (default) produces a single chunk covering the full field. The size + must divide the grid exactly -- that dimension cannot be left ragged. + """ + + def __init__( + self, + fdb_config: pathlib.Path | None = None, + gribjump_config: pathlib.Path | None = None, + field_chunking: "Chunking | Chunking.FixedSizeChunk | None" = None, + ): + self._obj = pdv.ExtractorType.GribJump( + fdb_config=fdb_config, + gribjump_config=gribjump_config, + field_chunking=( + AxisDefinition._translate_chunking(field_chunking) if field_chunking is not None else None + ), + ) class ChunkedDataViewBuilder: + """Collects MARS request parts and builds a :class:`ChunkedDataView`. + + Wraps the C++ ``ChunkedDataViewBuilder``. Call :meth:`add_part` one or + more times, then :meth:`build` to obtain the view. + + Args: + fdb_config_file (pathlib.Path | None): Path to the FDB configuration YAML. + ``None`` lets FDB resolve its configuration from the environment. + """ + def __init__(self, fdb_config_file: pathlib.Path | None): self._obj = pdv.ChunkedDataViewBuilder(fdb_config_file) + self._dim_names: list[str] | None = None def add_part( self, mars_request: MarsSelection, axes: list[AxisDefinition], - extractor_type: ExtractorType, - ): + extractor: ExtractorType.Grib | ExtractorType.GribJump, + ) -> "ChunkedDataViewBuilder": + """Validate *axes* against *mars_request*, record dimension names, and register the part. + + Args: + mars_request (MarsSelection): MARS key-value mapping describing the data to retrieve. + axes (list[AxisDefinition]): Axis definitions; each must reference keys present in + *mars_request*. + extractor (ExtractorType.Grib | ExtractorType.GribJump): Extraction backend to use. + + Returns: + ChunkedDataViewBuilder: ``self``, for method chaining. + + Raises: + ValueError: If any axis key is not present in *mars_request*. + + Note: + Only the axis-key check happens here. Everything that needs FDB -- field sizes, + axis mapping, whether the parts fit together -- is validated by :meth:`build`, + which raises ``RuntimeError`` on any of it. + """ + for ax in axes: + missing = [k for k in ax.keys if k not in mars_request] + if missing: + raise ValueError( + f"ChunkedDataViewBuilder::add_part: Axis key(s) {missing!r} not found in the MARS request. " + f"Available keys: {sorted(mars_request.keys())!r}. " + f"Check for typos in the AxisDefinition." + ) + if self._dim_names is None: + self._dim_names = [ax.name if ax.name is not None else "_".join(ax.keys) for ax in axes] + self._dim_names.append("values") # implicit grid-point axis self._obj.add_part( _mars_selection_to_string(mars_request), [ax._obj for ax in axes], - extractor_type.value, + extractor._obj, ) + return self - def extend_on_axis(self, axis: int): + def dim_names(self) -> list[str]: + """Return the zarr dimension names derived from the first registered part. + + Returns: + list[str]: One name per axis (MARS keys joined by ``_``), plus ``"values"`` + for the implicit grid-point axis. Empty if no part has been added yet. + """ + return self._dim_names or [] + + def extend_on_axis(self, axis: int) -> "ChunkedDataViewBuilder": + """Set *axis* as the extension axis when multiple parts are added. + + Args: + axis (int): Zero-based index of the axis along which parts are concatenated. + + Returns: + ChunkedDataViewBuilder: ``self``, for method chaining. + """ self._obj.extend_on_axis(axis) + return self + + def fill_missing_value(self, value: float) -> "ChunkedDataViewBuilder": + """Set the fill value for bitmap-masked grid points. + + Args: + value (float): Value written into positions flagged as missing by the GRIB bitmap. + Also used as the zarr array ``fill_value``. - def fill_missing_value(self, value: float): + Returns: + ChunkedDataViewBuilder: ``self``, for method chaining. + """ self._obj.fill_missing_value(value) + return self + + def build(self) -> ChunkedDataView: + """Build and return the :class:`ChunkedDataView`. - def build(self): + Returns: + ChunkedDataView: The assembled view, ready for chunk-level data access. + + Raises: + ~pychunked_data_view.exceptions.MarsRequestFormattingError: If the MARS request + string is malformed + (trailing comma, missing comma between keys, or misspelled key). + """ try: return ChunkedDataView(self._obj.build()) except RuntimeError as re: diff --git a/src/z3fdb/__init__.py b/src/z3fdb/__init__.py index bd8ec99ac..23a77a1f5 100644 --- a/src/z3fdb/__init__.py +++ b/src/z3fdb/__init__.py @@ -47,7 +47,7 @@ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE) ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( { @@ -67,7 +67,7 @@ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE) ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() @@ -92,15 +92,14 @@ ExtractorType, MarsSelection, ) -from z3fdb.simple_store_builder import ( - SimpleStoreBuilder, -) - +from z3fdb.simple_store_builder import SimpleStoreBuilder +from z3fdb.custom_store_builder import CustomStoreBuilder from z3fdb.z3fdb_error import Z3fdbError __all__ = [ "AxisDefinition", "Chunking", + "CustomStoreBuilder", "ExtractorType", "MarsSelection", "SimpleStoreBuilder", diff --git a/src/z3fdb/_internal/zarr.py b/src/z3fdb/_internal/zarr.py index 4d78efebd..3fe38f6d0 100644 --- a/src/z3fdb/_internal/zarr.py +++ b/src/z3fdb/_internal/zarr.py @@ -17,9 +17,9 @@ Buffer = Union[bytes, bytearray, memoryview] -from typing import AsyncIterator, Iterable, Literal +from typing import AsyncIterator, Iterable, Iterator, Literal -import numpy as np +import math import itertools from zarr.abc import store @@ -29,7 +29,6 @@ from zarr.core.buffer.cpu import Buffer as CpuBuffer from zarr.core.common import BytesLike -from functools import cache from typing import Self from pychunked_data_view import ( @@ -42,15 +41,19 @@ def to_cpu_buffer(d: dict) -> CpuBuffer: + """Serialise *d* to JSON and wrap it in a zarr CPU buffer.""" return CpuBuffer.from_bytes(json.dumps(d).encode("utf-8")) def from_cpu_buffer(buf: CpuBuffer) -> dict: + """Deserialise a zarr CPU buffer to a dict.""" return json.loads(buf.to_bytes().decode("utf-8")) @dataclass(frozen=True) class DotZarrAttributes: + """Zarr v3 node attributes attached to every group and array node.""" + _: KW_ONLY copyright: str = "ecmwf" zarr_format: int = 3 @@ -78,13 +81,10 @@ def __init__(self, chunks) -> None: @dataclass class DotZarrArrayJson: - """ - Generates the .zarr metadata for an array. + """Zarr v3 array metadata written to ``zarr.json`` at each array node. - If additional fields are introduced, read the documentation about must_understand + If you add fields, check the ``must_understand`` rules in the spec: https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#id13 - - Most of what happens here has been reverse-engineered from the zarr-python code. """ _: KW_ONLY @@ -107,13 +107,10 @@ class DotZarrArrayJson: @dataclass() class DotZarrGroupJson: - """ - Generates the .zarr metadata for a group. + """Zarr v3 group metadata written to ``zarr.json`` at each group node. - If additional fields are introduced, read the documentation about must_understand + If you add fields, check the ``must_understand`` rules in the spec: https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#id13 - - Most of what happens here has been reverse-engineered from the zarr-python code. """ _: KW_ONLY @@ -125,23 +122,27 @@ class DotZarrGroupJson: class FdbSource: - """ - Uses FDB as a backend. - Data is retrieved from FDB and assembled on each access. + """Bridge between the C++ ``ChunkedDataView`` and zarr's array protocol. + + Caches shape and chunk-count metadata from the view at construction; + serves raw chunk buffers on demand via :meth:`__getitem__`. """ def __init__( self, chunked_data_view: ChunkedDataView, + dim_names: list[str] | None = None, ) -> None: self._chunked_data_view = chunked_data_view + self._dim_names = dim_names self._shape = self._chunked_data_view.shape() - self._chunks = self._chunked_data_view.chunkShape() + self._chunks = self._chunked_data_view.chunk_shape() self._chunks_per_dimension = self._chunked_data_view.chunks() self._fill_value = self._chunked_data_view.fill_missing_value() def create_dot_zarr_json(self) -> CpuBuffer: + """Return the ``zarr.json`` metadata buffer for this array.""" return to_cpu_buffer( asdict( DotZarrArrayJson( @@ -149,21 +150,26 @@ def create_dot_zarr_json(self) -> CpuBuffer: chunk_grid=ChunkGridMetadata(chunks=self._chunks), data_type="float32", fill_value=self._fill_value, + dimension_names=self._dim_names, ) ) ) - def __contains__(self, key: tuple[int, ...]) -> bool: - if len(key) != len(self._shape): + def contains_chunk(self, coords: tuple[int, ...]) -> bool: + """Return True if *coords* is a valid chunk index for this array.""" + if len(coords) != len(self._chunks_per_dimension): return False - if any(k < 0 or k >= limit for k, limit in zip(key, self._chunks_per_dimension)): - return False - return True + return all(0 <= c < limit for c, limit in zip(coords, self._chunks_per_dimension)) + + def __contains__(self, key: tuple[int, ...]) -> bool: + return self.contains_chunk(key) def chunks(self) -> tuple[int, ...]: + """Return the per-dimension chunk counts.""" return self._chunks_per_dimension def __getitem__(self, key: tuple[int, ...]) -> CpuBuffer: + """Fetch and return chunk *key* as a CPU buffer.""" if len(key) != len(self._shape): raise KeyError if any(k < 0 or k >= limit for k, limit in zip(key, self._chunks_per_dimension)): @@ -172,12 +178,15 @@ def __getitem__(self, key: tuple[int, ...]) -> CpuBuffer: class FdbZarrArray: + """Zarr v3 array node backed by an :class:`FdbSource`.""" + def __init__(self, *, name: str = "", datasource: FdbSource): self._name = name self._datasource = datasource self._metadata = self._datasource.create_dot_zarr_json() def __getitem__(self, key: str) -> AbstractBuffer | None: + """Route ``zarr.json`` to metadata and ``c//...`` keys to chunk data.""" if key == "zarr.json": return self._metadata if key.startswith("c/"): @@ -189,25 +198,14 @@ def __getitem__(self, key: str) -> AbstractBuffer | None: def name(self) -> str: return self._name - @cache def paths(self) -> list[str]: - """ - Zarr paths associated to this array, this includes .zarray, .zattrs and all chunks. - - Returns - ------- - list[str] - A list of paths belonging to this group - """ - files = ["zarr.json"] - if len(chunks_per_axis := self._datasource.chunks()) > 0: - tuples = itertools.product(*[np.arange(0, x) for x in chunks_per_axis]) - chunk_names = ["/".join([str(i) for i in ["c", *t]]) for t in tuples] - files += chunk_names - return files + """Return ``['zarr.json']``; chunk keys are not enumerated here.""" + return ["zarr.json"] class FdbZarrGroup: + """Zarr v3 group node; delegates key lookups to its named children.""" + def __init__( self, *, @@ -243,27 +241,62 @@ def children(self) -> list["FdbZarrArray | FdbZarrGroup"]: return list(self._children.values()) def paths(self) -> list[str]: - """ - Zarr paths associated to this group, excluding child groups or arrays. - - Returns - ------- - list[str] - A list of paths belonging to this group - """ + """Return ``['zarr.json']`` for this group (children contribute their own paths).""" return ["zarr.json"] class FdbZarrStore(store.Store): - """Provide access to FDB.""" + """Read-only zarr v3 store that virtualises an arbitrary group/array hierarchy. + + ``_known_paths`` holds metadata keys only (one ``zarr.json`` per node). + Chunk keys are validated O(1) via :meth:`_chunk_key_exists` and generated + lazily by :meth:`_iter_chunk_paths` - they are never stored. + """ def __init__(self, child: FdbZarrGroup | FdbZarrArray): super().__init__(read_only=True) self._child = child + # Metadata paths only - chunk keys are generated on demand. self._known_paths = self._build_paths(self._child) + # Flat map: absolute array path -> FdbZarrArray, for chunk operations. + self._arrays: dict[str, FdbZarrArray] = self._collect_arrays(self._child) self._root_zarr_json = self._build_root_zarr_json() + def _collect_arrays( + self, + item: "FdbZarrGroup | FdbZarrArray", + parent_path: str = "", + ) -> "dict[str, FdbZarrArray]": + """Return a flat dict mapping each array's absolute path to its node.""" + path = f"{parent_path}/{item.name}" if parent_path else item.name + if isinstance(item, FdbZarrArray): + return {path: item} + result: dict[str, FdbZarrArray] = {} + for c in item.children: + result.update(self._collect_arrays(c, path)) + return result + + def _chunk_key_exists(self, key: str) -> bool: + """Return True if *key* is a valid chunk key for any array, without enumeration.""" + for arr_path, arr in self._arrays.items(): + chunk_prefix = (arr_path + "/c/") if arr_path else "c/" + if key.startswith(chunk_prefix): + try: + coords = tuple(int(c) for c in key[len(chunk_prefix) :].split("/")) + except ValueError: + return False + return arr._datasource.contains_chunk(coords) + return False + + def _iter_chunk_paths(self) -> Iterator[str]: + """Yield every chunk key for every array, built at call time.""" + for arr_path, arr in self._arrays.items(): + prefix = (arr_path + "/c/") if arr_path else "c/" + for coords in itertools.product(*[range(n) for n in arr._datasource.chunks()]): + yield prefix + "/".join(str(c) for c in coords) + def _build_paths(self, item, parent_path=None) -> list[str]: + """Collect metadata paths for *item* and all its descendants.""" path = f"{parent_path}/{item.name}" if parent_path else item.name files = [f"{path}/{f}" if path != "" else f for f in item.paths()] @@ -274,6 +307,7 @@ def _build_paths(self, item, parent_path=None) -> list[str]: return files def _build_root_zarr_json(self) -> CpuBuffer: + """Build the root ``zarr.json`` with inline ``consolidated_metadata`` for the full hierarchy.""" # Only root groups can carry consolidated metadata in zarr v3. # Root arrays are terminal nodes; no consolidated_metadata concept applies. if not isinstance(self._child, FdbZarrGroup): @@ -309,7 +343,7 @@ def _descendants(group_abs_path: str) -> dict[str, dict]: # Build the flat metadata dict. Every group gets a consolidated_metadata # whose metadata contains all its descendants with relative paths and plain - # (no consolidated_metadata) values — groups within are raw group JSON only. + # (no consolidated_metadata) values - groups within are raw group JSON only. flat: dict[str, dict] = {} for abs_path, meta in raw.items(): if not abs_path: @@ -340,10 +374,12 @@ async def __getitem__(self, key) -> AbstractBuffer | None: return self._child[key] def __iter__(self): - yield from iter(self._known_paths) + yield from self._known_paths + yield from self._iter_chunk_paths() def __len__(self): - return len(self._known_paths) + chunk_total = sum(math.prod(arr._datasource.chunks()) for arr in self._arrays.values()) + return len(self._known_paths) + chunk_total def __setitem__(self, _k, _v): raise Z3fdbError("Views into FDB are not writable") @@ -352,7 +388,7 @@ def __delitem__(self, _k): raise Z3fdbError("Views into FDB are not writable") def __contains__(self, key) -> bool: - return key in self._known_paths + return key in self._known_paths or self._chunk_key_exists(key) def __eq__(self, value: object) -> bool: if not isinstance(value, FdbZarrStore): @@ -414,13 +450,18 @@ def supports_listing(self) -> bool: return True async def list(self) -> AsyncIterator[str]: - for i in self._known_paths: - yield i + for path in self._known_paths: + yield path + for path in self._iter_chunk_paths(): + yield path async def list_prefix(self, prefix: str) -> AsyncIterator[str]: for path in self._known_paths: if path.startswith(prefix): yield path + for path in self._iter_chunk_paths(): + if path.startswith(prefix): + yield path async def list_dir(self, prefix: str) -> AsyncIterator[str]: # Normalize so the scan prefix ends with "/" @@ -432,3 +473,9 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: if child and child not in seen: seen.add(child) yield child + for path in self._iter_chunk_paths(): + if path.startswith(scan_prefix): + child = path[len(scan_prefix) :].split("/")[0] + if child and child not in seen: + seen.add(child) + yield child diff --git a/src/z3fdb/custom_store_builder.py b/src/z3fdb/custom_store_builder.py new file mode 100644 index 000000000..30af077f1 --- /dev/null +++ b/src/z3fdb/custom_store_builder.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from pathlib import Path + +from z3fdb._internal.zarr import FdbZarrArray, FdbZarrGroup, FdbZarrStore, FdbSource +from z3fdb.z3fdb_error import Z3fdbError +from pychunked_data_view import ( + ChunkedDataViewBuilder, + AxisDefinition, + ExtractorType, + MarsSelection, +) + + +@dataclass +class _VArray: + name: str + parent: "_VGroup" + builder: ChunkedDataViewBuilder + + +@dataclass +class _VGroup: + parents: "list[_VGroup] | None" + name: str + children: "list[_VGroup | _VArray]" = field(default_factory=list) + + @staticmethod + def _join_path(group: "_VGroup") -> str: + return "/".join(p.name for p in group.parents) if group.parents else "" + + def __eq__(self, value: object) -> bool: + if not isinstance(value, _VGroup): + return False + return self._join_path(self) == self._join_path(value) and self.name == value.name + + def descent(self, name: str) -> "_VGroup": + """Return the single child group called *name*. + + Raises: + ~z3fdb.Z3fdbError: If there is not exactly one. Names are unique by construction, so + this is an internal invariant, a bare ``assert`` would vanish under + ``python -O``. + """ + matches = [c for c in self.children if isinstance(c, _VGroup) and c.name == name] + if len(matches) != 1: + raise Z3fdbError( + f"CustomStoreBuilder: expected exactly one group named {name!r} under " + f"{self.name!r}, found {len(matches)}." + ) + return matches[0] + + +class CustomStoreBuilder: + """Builds a zarr store backed by FDB with an arbitrary group/array hierarchy. + + Use :meth:`add_part` to register one or more MARS request parts (each + producing a virtual zarr array) at arbitrary nested paths, then call + :meth:`build` to obtain a read-only :class:`FdbZarrStore` that zarr can + open directly. + + Args: + fdb_config_file: Optional path to an FDB config file. ``None`` (default) + lets FDB resolve its configuration from the environment. + """ + + _ROOT_KEY = "" # structure-dict key reserved for the root array (path=None) + + def __init__(self, fdb_config_file: Path | None = None): + self._config = fdb_config_file + self._structure: dict[str, _VArray] = {} + self._root = _VGroup(parents=None, name="/", children=[]) + + def _merge_vgroup(self, vgroup: _VGroup) -> None: + """Insert *vgroup* into the virtual group tree rooted at self._root.""" + if vgroup.parents is None: + raise RuntimeError("Cannot have two root groups.") + + current_group = self._root + for parent in vgroup.parents[1:]: + # _build_structure creates parents top-down, so every ancestor already exists. + if parent.name not in [c.name for c in current_group.children]: + raise Z3fdbError( + f"CustomStoreBuilder: parent group {parent.name!r} of {vgroup.name!r} does not exist yet." + ) + current_group = current_group.descent(parent.name) + + if vgroup not in current_group.children: + current_group.children.append(vgroup) + + @staticmethod + def _parse_path(path: str) -> list[str]: + """Convert a zarr-style path string to a list of name segments. + + Leading and trailing slashes are stripped; multiple consecutive slashes + are collapsed. An empty result (e.g. ``""`` or ``"/"``) raises + :exc:`ValueError`. + + Examples:: + + "sfc/wind" -> ["sfc", "wind"] + "/sfc/wind" -> ["sfc", "wind"] # leading slash accepted + "t2m" -> ["t2m"] # top-level array + "" -> ValueError + "/" -> ValueError + """ + parts = [p for p in path.split("/") if p] + if not parts: + raise ValueError( + f"CustomStoreBuilder: path must not be empty. " + f"Got {path!r}. Use a zarr-style path like 'group/array', " + f"'/group/array', or 'array' for a top-level array." + ) + return parts + + def _build_structure(self, path: list[str] | None) -> ChunkedDataViewBuilder: + """Return the ChunkedDataViewBuilder for *path*, creating it if needed. + + Pass ``None`` to obtain the builder for the root array. + """ + if path is None: + # Root array - the store root is itself an array, not a group. + if self._root.children: + raise ValueError( + "CustomStoreBuilder: cannot register a root array (path=None) when " + "named paths are already registered. Use a named path instead." + ) + if self._ROOT_KEY not in self._structure: + builder = ChunkedDataViewBuilder(self._config) + self._structure[self._ROOT_KEY] = _VArray(name="", parent=self._root, builder=builder) + return self._structure[self._ROOT_KEY].builder + + # Named path - cannot mix with a root array. + if self._ROOT_KEY in self._structure: + raise ValueError( + "CustomStoreBuilder: cannot register a named path when a root array " + "(path=None) is already registered. Use path=None to add more parts " + "to the root array." + ) + + key = "/".join(path) + if key in self._structure: + return self._structure[key].builder + + group_names = path[:-1] + array_name = path[-1] + + parents: list[_VGroup] = [self._root] + for group_name in group_names: + # Collision: a _VArray already occupies this name - cannot reuse as a group. + if any(isinstance(c, _VArray) and c.name == group_name for c in parents[-1].children): + raise ValueError( + f"CustomStoreBuilder: '{group_name}' is already registered as an array " + f"and cannot also be used as a group." + ) + vgroup = _VGroup(parents=list(parents), name=group_name) + self._merge_vgroup(vgroup) + actual = next(c for c in parents[-1].children if isinstance(c, _VGroup) and c.name == group_name) + parents.append(actual) + + parent_group = parents[-1] + # Collision: a _VGroup already occupies this name - cannot reuse as an array. + if any(isinstance(c, _VGroup) and c.name == array_name for c in parent_group.children): + raise ValueError( + f"CustomStoreBuilder: '{array_name}' is already registered as a group " + f"and cannot also be used as an array." + ) + builder = ChunkedDataViewBuilder(self._config) + varray = _VArray(name=array_name, parent=parent_group, builder=builder) + parent_group.children.append(varray) + self._structure[key] = varray + return builder + + def add_part( + self, + path: str | None, + mars_request: MarsSelection, + axes: list[AxisDefinition], + extractor: ExtractorType.Grib | ExtractorType.GribJump, + ) -> None: + """Register a MARS request as a part of a virtual zarr array at *path*. + + Calling this method multiple times with the same *path* adds further + parts to the same array (equivalent to + :meth:`ChunkedDataViewBuilder.add_part` called repeatedly). + + Args: + path: Zarr-style path of the array in the hierarchy, + e.g. ``"group_a/sub_group/my_array"`` or ``"t2m"`` for a + top-level (no-group) array. A leading ``/`` is accepted and + ignored. Pass ``None`` to place the array at the store root + (accessible via ``zarr.open_array(store)``); this is mutually + exclusive with any named path. + mars_request: MARS request as a dict mapping keys to values. + axes: Axis definitions describing how the request dimensions map + to zarr array dimensions. + extractor: Extractor configuration (``ExtractorType.Grib`` or + ``ExtractorType.GribJump``). + """ + parts = None if path is None else self._parse_path(path) + builder = self._build_structure(parts) + builder.add_part(mars_request, axes, extractor) + + def _existing_array(self, path: list[str] | None) -> ChunkedDataViewBuilder: + """Return the builder for an array already registered at *path*. + + Unlike :meth:`_build_structure` this never creates one. :meth:`extend_on_axis` and + :meth:`fill_missing_value` *configure* an existing array, so an unknown path is a + mistake -- usually a typo -- rather than a request for a new empty array. Creating one + silently would only surface much later, as "must add at least one part" from + :meth:`build`. + + Args: + path: Path segments, or ``None`` for the root array. + + Returns: + ChunkedDataViewBuilder: The builder registered at *path*. + + Raises: + ValueError: If no array is registered at *path*. + """ + key = self._ROOT_KEY if path is None else "/".join(path) + if key not in self._structure: + known = sorted(k or "" for k in self._structure) or ["none"] + where = "the root array (path=None)" if path is None else repr("/".join(path)) + raise ValueError( + f"CustomStoreBuilder: no array registered at {where}. Call add_part first. " + f"Registered arrays: {', '.join(known)}." + ) + return self._structure[key].builder + + def extend_on_axis(self, path: str | None, axis: int) -> None: + """Declare the extension axis of the array at *path*. + + The array must already exist: call :meth:`add_part` for *path* first. + + Args: + path: Zarr-style path (same format as :meth:`add_part`). + ``None`` refers to the root array. + axis: Zero-based index of the axis to extend. + + Raises: + ValueError: If no array is registered at *path*. + """ + parts = None if path is None else self._parse_path(path) + self._existing_array(parts).extend_on_axis(axis) + + def fill_missing_value(self, path: str | None, value: float) -> None: + """Set the fill value for the array at *path*. + + The array must already exist: call :meth:`add_part` for *path* first. + + Args: + path: Zarr-style path (same format as :meth:`add_part`). ``None`` refers to the + root array. + value: Value written into positions flagged as missing by the GRIB bitmap. Also + becomes the zarr array's ``fill_value``. Defaults to NaN when not set. + + Raises: + ValueError: If no array is registered at *path*. + """ + parts = None if path is None else self._parse_path(path) + self._existing_array(parts).fill_missing_value(value) + + def build(self) -> FdbZarrStore: + """Assemble all registered views into a read-only :class:`FdbZarrStore`. + + Returns: + A zarr-compatible store that can be opened with + ``zarr.open(store)`` (group hierarchy) or + ``zarr.open_array(store)`` (when built from a single root array + registered via ``path=None``). + """ + # Root-array shortcut: the store root is itself an array. + if self._ROOT_KEY in self._structure: + varray = self._structure[self._ROOT_KEY] + view = varray.builder.build() + return FdbZarrStore( + FdbZarrArray( + name="", + datasource=FdbSource(view, dim_names=varray.builder.dim_names()), + ) + ) + + def _to_fdb_node(node: _VGroup | _VArray) -> FdbZarrGroup | FdbZarrArray: + if isinstance(node, _VArray): + view = node.builder.build() + return FdbZarrArray( + name=node.name, + datasource=FdbSource(view, dim_names=node.builder.dim_names()), + ) + children = [_to_fdb_node(c) for c in node.children] + return FdbZarrGroup(name=node.name, children=children) + + root_children = [_to_fdb_node(c) for c in self._root.children] + return FdbZarrStore(FdbZarrGroup(name="", children=root_children)) diff --git a/src/z3fdb/simple_store_builder.py b/src/z3fdb/simple_store_builder.py index 2a84479d0..f032364b8 100644 --- a/src/z3fdb/simple_store_builder.py +++ b/src/z3fdb/simple_store_builder.py @@ -1,12 +1,11 @@ # SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) # SPDX-License-Identifier: Apache-2.0 -from zarr.abc.store import Store from pathlib import Path -from z3fdb._internal.zarr import FdbZarrStore, FdbZarrArray, FdbSource +from z3fdb._internal.zarr import FdbZarrStore +from z3fdb.custom_store_builder import CustomStoreBuilder from pychunked_data_view import ( - ChunkedDataViewBuilder, AxisDefinition, ExtractorType, MarsSelection, @@ -19,19 +18,22 @@ class SimpleStoreBuilder: This builder will create a Zarr store with a Zarr Array at its root ("/") containing the data from your MARS request(s). + It is exactly :class:`~z3fdb.CustomStoreBuilder` restricted to the root array, and + delegates to it -- so the two cannot drift apart. + Args: fdb_config_file: Optional path to FDB config file. If not set normal FDB config file resolution is applied. """ def __init__(self, fdb_config_file: Path | None = None): - self._builder = ChunkedDataViewBuilder(fdb_config_file) + self._builder = CustomStoreBuilder(fdb_config_file) def add_part( self, mars_request: MarsSelection, axes: list[AxisDefinition], - extractor_type: ExtractorType, + extractor: ExtractorType.Grib | ExtractorType.GribJump, ) -> None: """Add a MARS request to the view. @@ -61,39 +63,50 @@ def add_part( axes(:obj:`list` of :obj:`AxisDefinition`): List of AxisDefinitions that describe how axis in the MARS request are mapped to axis in the Zarr array. - extractor_type: Defines how to extract data from FDB. Currently - only ExtractorType.GRIB is supported. + extractor: Extractor configuration object. Use + ``ExtractorType.Grib()`` for full-field GRIB extraction or + ``ExtractorType.GribJump(...)`` for partial-field extraction. """ - self._builder.add_part(mars_request, axes, extractor_type) + self._builder.add_part(None, mars_request, axes, extractor) def fill_missing_value(self, value: float) -> None: - """Set the fill value used for missing / bitmap-masked grid points. + """Set the fill value used for bitmap-masked grid points. + + Call :meth:`add_part` first: this configures the array, so there has to be one. Args: value(float): Fill value written into array positions that carry a GRIB bitmap missing flag. Also used as the zarr array fill_value. + + Raises: + ValueError: If no part has been added yet. """ - self._builder.fill_missing_value(value) + self._builder.fill_missing_value(None, value) def extend_on_axis(self, axis: int) -> None: """Defines the extension axis when multiple parts are added. + Call :meth:`add_part` first: this configures the array, so there has to be one. + Args: - axis(int): Index of the axis that is extendet when multiple parts + axis(int): Index of the axis that is extended when multiple parts have been added. + Raises: + ValueError: If no part has been added yet. """ - self._builder.extend_on_axis(axis) + self._builder.extend_on_axis(None, axis) def build(self) -> FdbZarrStore: - """Build the store from the inputs. + """Build the store from the registered parts. - Raises: - Z3fdbError if store cannot be created. + Returns: + :class:`~z3fdb._internal.zarr.FdbZarrStore` ready to pass to + ``zarr.open_array()``. + Raises: + RuntimeError: If the view is misconfigured -- no parts, a missing extension axis, + incompatible part shapes, or parts that disagree about the grid. These + originate as ``eckit::UserError`` in the C++ layer. """ - return FdbZarrStore( - FdbZarrArray( - datasource=FdbSource(self._builder.build()), - ) - ) + return self._builder.build() diff --git a/tests/chunked_data_view/CMakeLists.txt b/tests/chunked_data_view/CMakeLists.txt index dae7b33bf..00e2f7348 100644 --- a/tests/chunked_data_view/CMakeLists.txt +++ b/tests/chunked_data_view/CMakeLists.txt @@ -4,6 +4,7 @@ list(APPEND tests index_mapper request_manipulation_bounding_box bounding_box + list_element view_individual_chunking ) diff --git a/tests/chunked_data_view/test_list_element.cc b/tests/chunked_data_view/test_list_element.cc new file mode 100644 index 000000000..257890602 --- /dev/null +++ b/tests/chunked_data_view/test_list_element.cc @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +// SPDX-License-Identifier: Apache-2.0 + +#include "chunked_data_view/ListIterator.h" + +#include "eckit/testing/Test.h" + +#include "test_mock_helpers.h" + +#include +#include +#include + +// Regression guard for a leak: fdb5::FieldLocation::dataHandle() returns a raw *owning* +// pointer, and ListElement::dataHandle() used to hand it straight through, so every caller +// leaked one DataHandle (and its file descriptor) per field per chunk read. +CASE("ListElement | dataHandle returns an owning unique_ptr") { + + using Returned = decltype(std::declval().dataHandle()); + + static_assert(std::is_same_v>, + "ListElement::dataHandle() must return an owning unique_ptr, never a raw " + "eckit::DataHandle* — a raw pointer here leaks on every chunk read."); + + const std::vector values = {1, 2, 3, 4}; + const chunked_data_view::ListElement element{fdb5::Key(), std::make_shared(values)}; + + auto first = element.dataHandle(); + EXPECT(first != nullptr); + + // Each call must open its own handle; the caller owns and releases each one. + auto second = element.dataHandle(); + EXPECT(second != nullptr); + EXPECT(first.get() != second.get()); +}; + +CASE("ListElement | the handle from dataHandle is readable and self-owned") { + + const std::vector values = {7, 8, 9}; + const chunked_data_view::ListElement element{fdb5::Key(), std::make_shared(values)}; + + size_t countValues = 0; + size_t bytesPerValue = 0; + { + auto handle = element.dataHandle(); + handle->openForRead(); + EXPECT_EQUAL(handle->read(&countValues, sizeof(countValues)), sizeof(countValues)); + EXPECT_EQUAL(handle->read(&bytesPerValue, sizeof(bytesPerValue)), sizeof(bytesPerValue)); + handle->close(); + } // handle destroyed here — under -fsanitize=leak this scope must not report a leak + + EXPECT_EQUAL(countValues, values.size()); + EXPECT_EQUAL(bytesPerValue, 8); +}; + +CASE("ListElement | MockListIterator yields a usable location") { + + // The mock used to return a null location, which is why the leak above went unnoticed. + // NOTE: MockListIterator pre-increments before returning, so createMockFDB(n) yields + // n - 1 elements (see the comment in test_view.cc). + auto fdb = createMockFDB(/* fieldAmount = */ 2); + auto iterator = fdb->inspect(metkit::mars::MarsRequest{}); + + size_t seen = 0; + while (const auto element = iterator->next()) { + EXPECT(element->location != nullptr); + EXPECT(element->dataHandle() != nullptr); + ++seen; + } + EXPECT_EQUAL(seen, 1); +}; + +int main(int argc, char** argv) { + return ::eckit::testing::run_tests(argc, argv); +} diff --git a/tests/chunked_data_view/test_mock_helpers.h b/tests/chunked_data_view/test_mock_helpers.h index a2685662b..349aa2d28 100644 --- a/tests/chunked_data_view/test_mock_helpers.h +++ b/tests/chunked_data_view/test_mock_helpers.h @@ -12,11 +12,13 @@ #include #include #include +#include #include #include #include #include #include "chunked_data_view/ListIterator.h" +#include "chunked_data_view/Types.h" #include "chunked_data_view/ViewPart.h" #include "eckit/io/DataHandle.h" #include "fdb5/database/FieldLocation.h" @@ -40,6 +42,22 @@ inline std::unique_ptr makeHandle(const std::vector& return handle; }; +/// Minimal concrete fdb5::FieldLocation serving an in-memory field. +struct MockFieldLocation final : public fdb5::FieldLocation { + + explicit MockFieldLocation(std::vector values) : values_(std::move(values)) {} + + eckit::DataHandle* dataHandle() const override { return makeHandle(values_).release(); } + + std::shared_ptr make_shared() const override { + return std::make_shared(values_); + } + + void visit(fdb5::FieldLocationVisitor&) const override {} + + std::vector values_; +}; + struct MockListIterator final : public chunked_data_view::ListIteratorInterface { using vec2 = std::vector>>; @@ -48,14 +66,15 @@ struct MockListIterator final : public chunked_data_view::ListIteratorInterface MockListIterator(vec2 data) : data_(std::move(data)), iter_(std::begin(data_)) {}; - std::optional>> next() { + std::optional next() override { iter_++; if (std::end(data_) == iter_) { return std::nullopt; } - return std::make_tuple(std::get<0>(*iter_), makeHandle(std::get<1>(*iter_))); + return chunked_data_view::ListElement{std::get<0>(*iter_), + std::make_shared(std::get<1>(*iter_))}; }; }; @@ -80,8 +99,13 @@ struct MockFdb final : public cdv::FdbInterface { InsFunc insFn{}; }; -inline std::shared_ptr createMockFDB(size_t fieldAmount = 1) { - const std::vector values = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; +/// @param fieldAmount how many entries the inspect() iterator holds. NOTE: MockListIterator +/// pre-increments before returning, so it yields fieldAmount - 1 elements. +/// @param countValues values per field, i.e. the extent of the implicit (grid-point) axis. +/// Vary it to build parts whose grids disagree. +inline std::shared_ptr createMockFDB(size_t fieldAmount = 1, size_t countValues = 10) { + std::vector values(countValues); + std::iota(values.begin(), values.end(), 1.0); return std::make_shared( [values](auto& _) { return makeHandle(values); }, [fieldAmount, values](auto& _) -> std::unique_ptr { @@ -94,20 +118,19 @@ struct FakeExtractor : public cdv::Extractor { std::shared_ptr mock_; - explicit FakeExtractor(std::shared_ptr mock_fdb) : mock_(mock_fdb) {} - - cdv::DataLayout layout(const metkit::mars::MarsRequest& mars_request) const override { - const auto handle = mock_->retrieve(mars_request); - cdv::DataLayout layout{}; + explicit FakeExtractor(std::shared_ptr mock_fdb) : mock_(mock_fdb) { + // The mock's retrieve() ignores the request, so a default-constructed one suffices. + const auto handle = mock_->retrieve(metkit::mars::MarsRequest{}); handle->openForRead(); - EXPECT_EQUAL(handle->read(&layout.countValues, sizeof(layout.countValues)), sizeof(layout.countValues)); - EXPECT_EQUAL(handle->read(&layout.bytesPerValue, sizeof(layout.bytesPerValue)), sizeof(layout.bytesPerValue)); + EXPECT_EQUAL(handle->read(&layout_.countValues, sizeof(layout_.countValues)), sizeof(layout_.countValues)); + EXPECT_EQUAL(handle->read(&layout_.bytesPerValue, sizeof(layout_.bytesPerValue)), + sizeof(layout_.bytesPerValue)); handle->close(); - return layout; + layout_.countChunkValues = layout_.countValues; // full field = single implicit chunk } size_t extractInto(const chunked_data_view::ViewPart& part, - const chunked_data_view::ChunkedDataViewPartBoundingBox& chunkBoundingBox, + const chunked_data_view::ChunkBoundingBox& chunkBoundingBox, const chunked_data_view::ChunkedDataViewPartBoundingBox& intersectionBoundingBox, float* ptr, size_t len) const override { @@ -126,3 +149,25 @@ struct FakeExtractor : public cdv::Extractor { return written; }; }; + +/// ExtractorDefinition backed by a shared mock FDB. +/// buildExtractor() creates a fresh FakeExtractor each call (satisfying unique_ptr +/// ownership), while all copies share the same mock FDB instance so multi-part tests +/// can share a single createMockFDB() result across several addPart() calls. +struct FakeExtractorDefinition : public cdv::ExtractorDefinition { + std::shared_ptr mock_fdb_; + + explicit FakeExtractorDefinition(std::shared_ptr mock_fdb) : + mock_fdb_(std::move(mock_fdb)) {} + + /// No-op: the mock FDB ignores configuration entirely. + void setDefaultIfUnset(const std::optional& fdbConfigPath) override {} + + std::unique_ptr copy() const override { + return std::make_unique(*this); + } + + std::unique_ptr buildExtractor(const metkit::mars::MarsRequest&) const override { + return std::make_unique(mock_fdb_); + } +}; diff --git a/tests/chunked_data_view/test_view.cc b/tests/chunked_data_view/test_view.cc index 29842c466..02538a399 100644 --- a/tests/chunked_data_view/test_view.cc +++ b/tests/chunked_data_view/test_view.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -27,7 +28,7 @@ CASE("ChunkedDataView | View from 1 request | Can compute shape") { "time=0/6/12/18"}; // If - auto mock_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); // Then const auto view = cdv::ChunkedDataViewBuilder() @@ -35,7 +36,7 @@ CASE("ChunkedDataView | View from 1 request | Can compute shape") { {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // Expect to get: 4 dates, 4 times, 2 fields, 10 values per field (implicit axis) @@ -56,19 +57,19 @@ CASE("ChunkedDataView | View from 2 requests | Can compute shape") { "time=0/6/12/18"}; - auto fake_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - fake_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - fake_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(2) .build(); // Expect to get: 4 dates, 4 times, 2*2 fields (2 per request), 10 values per field (implicit axis) @@ -86,19 +87,19 @@ CASE("ChunkedDataView | View from 2 requests | WholeAxisChunking on extension ax "param=v/u," "time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(3)); + auto mock_fdb = createMockFDB(3); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(2) .build(); @@ -123,17 +124,17 @@ CASE("ChunkedDataView | View from 2 requests | Can compute shape, combined axis" "param=v/u," "time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(1) .build(); // Expect to get: 16 date/times (4 dates * 4 times), 2*2 fields (2 per request), 10 values per field (implicit @@ -141,21 +142,21 @@ CASE("ChunkedDataView | View from 2 requests | Can compute shape, combined axis" } CASE("ChunkedDataViewBuilder | build | Calling build twice throws") { - // build() used to move the shared_ptr out of the builder's internal parts_ list, making a second call - // UB. The builder now retains ownership so build() can be called multiple times safely. + // The ExtractorDefinition is retained in parts_ (not consumed on build()), so build() can + // be called multiple times. Each call invokes buildExtractor() to produce a fresh Extractor. const std::string keys{ "type=an,domain=g,expver=0001,stream=oper," "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); cdv::ChunkedDataViewBuilder builder; builder.addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor); + FakeExtractorDefinition{mock_fdb}); EXPECT_NO_THROW(builder.build()); EXPECT_NO_THROW(builder.build()); @@ -172,14 +173,14 @@ CASE("ChunkedDataView - Can build") { "param=v/u," "time=0/6/12/18"}; - auto fake_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); EXPECT_NO_THROW(cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - fake_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); } @@ -193,14 +194,12 @@ CASE("ChunkedDataView | build | No data in FDB throws user-facing error") { [](auto& _) -> std::unique_ptr { throw eckit::Exception("FDB: no data"); }, [](auto& _) -> std::unique_ptr { return nullptr; }); - auto fake_extractor = std::make_shared(mock_fdb); - EXPECT_THROWS(cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - fake_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); } @@ -210,14 +209,14 @@ CASE("ChunkedDataView | at | Wrong index dimension throws") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(3)); + auto mock_fdb = createMockFDB(3); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // View has 4 dimensions (date, time, param, values) -> chunks has 4 entries @@ -236,14 +235,14 @@ CASE("ChunkedDataView | at | Out-of-bounds chunk index throws") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(3)); + auto mock_fdb = createMockFDB(3); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // chunks = {4, 4, 1, 1} (4 dates, 4 times, 1 param-chunk, 1 value-chunk) @@ -271,14 +270,14 @@ CASE("ChunkedDataView | at | Partial read throws") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(2)); + auto mock_fdb = createMockFDB(2); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); std::vector buf(view->countChunkValues()); @@ -291,19 +290,19 @@ CASE("ChunkedDataView | at | Partial read in multi-part extension throws") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(2)); + auto mock_fdb = createMockFDB(2); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(2) .build(); @@ -323,19 +322,19 @@ CASE("ChunkedDataView | at | Partial read error path uses part-local coordinates // createMockFDB(1) returns 0 messages (MockListIterator pre-increments before returning). // SingleValueChunking expects 1 message per chunk -> mismatch triggers the error path. - auto mock_extractor = std::make_shared(createMockFDB(1)); + auto mock_fdb = createMockFDB(1); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(2) .build(); @@ -358,24 +357,24 @@ CASE("ChunkedDataView | View from 3 requests | Can compute shape and access") { "param=v/u," "time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(3)); + auto mock_fdb = createMockFDB(3); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(2) .build(); @@ -388,6 +387,70 @@ CASE("ChunkedDataView | View from 3 requests | Can compute shape and access") { EXPECT_NO_THROW(view->at({0, 0, 0, 0}, buf.data(), buf.size())); } +CASE("ChunkedDataView | Parts whose grids differ are rejected at build time") { + const std::string keys{ + "type=an," + "domain=g," + "expver=0001," + "stream=oper," + "date=2020-01-01/to/2020-01-04," + "levtype=sfc," + "param=v/u," + "time=0/6/12/18"}; + + // The implicit (grid-point) axis is never the extension axis, so like every other + // non-extension axis all parts must agree on it — and a zarr array cannot have a ragged + // last dimension. Left unchecked this fails silently at read time: each extractor sizes + // its writes from its own layout, so the part with the larger field overruns its slots + // while still reporting the expected message count. + auto mock_10 = createMockFDB(/* fieldAmount = */ 3, /* countValues = */ 10); + auto mock_20 = createMockFDB(/* fieldAmount = */ 3, /* countValues = */ 20); + + const auto axes = [] { + return std::vector{ + cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, + cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, + cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}; + }; + + EXPECT_THROWS_AS(cdv::ChunkedDataViewBuilder() + .addPart(keys, axes(), FakeExtractorDefinition{mock_10}) + .addPart(keys, axes(), FakeExtractorDefinition{mock_20}) + .extendOnAxis(2) + .build(), + eckit::UserError); +} + +CASE("ChunkedDataView | Parts sharing a grid still build") { + const std::string keys{ + "type=an," + "domain=g," + "expver=0001," + "stream=oper," + "date=2020-01-01/to/2020-01-04," + "levtype=sfc," + "param=v/u," + "time=0/6/12/18"}; + + // Guards the check above against over-rejecting: same grid, two parts, must succeed. + auto mock_fdb = createMockFDB(/* fieldAmount = */ 3, /* countValues = */ 20); + + const auto axes = [] { + return std::vector{ + cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::SingleValueChunking{}}, + cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, + cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::SingleValueChunking{}}}; + }; + + const auto view = cdv::ChunkedDataViewBuilder() + .addPart(keys, axes(), FakeExtractorDefinition{mock_fdb}) + .addPart(keys, axes(), FakeExtractorDefinition{mock_fdb}) + .extendOnAxis(2) + .build(); + + EXPECT_EQUAL(view->shape(), (std::vector{4, 4, 4, 20})); +} + int main(int argc, char** argv) { return ::eckit::testing::run_tests(argc, argv); } diff --git a/tests/chunked_data_view/test_view_individual_chunking.cc b/tests/chunked_data_view/test_view_individual_chunking.cc index 450b809aa..d9d1a3ff4 100644 --- a/tests/chunked_data_view/test_view_individual_chunking.cc +++ b/tests/chunked_data_view/test_view_individual_chunking.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -20,14 +21,14 @@ CASE("ChunkedDataView | FixedSizeChunking | Can compute shape") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::FixedSizeChunking{2}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // 4 dates (chunkSize=2 -> 2 chunks), 4 single times, 2 params whole, 10 values @@ -42,16 +43,12 @@ CASE("ChunkedDataView | FixedSizeChunking | Invalid chunk size throws") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB()); + auto mock_fdb = createMockFDB(); - // chunkSize=0 -- explicitly forbidden - EXPECT_THROWS(cdv::ChunkedDataViewBuilder() - .addPart(keys, - {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::FixedSizeChunking{0}}, - cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, - cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) - .build()); + // chunkSize=0 -- rejected by FixedSizeChunking's own constructor, so it never reaches + // addPart() or build(). Asserted directly: wrapping it in a builder chain would still + // pass, but only because the throw happens while evaluating the argument. + EXPECT_THROWS_AS(cdv::AxisDefinition::FixedSizeChunking{0}, eckit::AssertionFailed); // chunkSize=1 -- same as SingleValueChunking, redundant, but ok EXPECT_NO_THROW(cdv::ChunkedDataViewBuilder() @@ -59,7 +56,7 @@ CASE("ChunkedDataView | FixedSizeChunking | Invalid chunk size throws") { {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::FixedSizeChunking{1}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); // chunkSize==axis size (4) -- same as WholeAxisChunking, redundant, but ok @@ -68,7 +65,7 @@ CASE("ChunkedDataView | FixedSizeChunking | Invalid chunk size throws") { {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::FixedSizeChunking{4}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); } @@ -79,14 +76,14 @@ CASE("ChunkedDataView | FixedSizeChunking | at() accesses each chunk") { "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=v/u,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(5)); + auto mock_fdb = createMockFDB(5); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date"}, cdv::AxisDefinition::FixedSizeChunking{2}}, cdv::AxisDefinition{{"time"}, cdv::AxisDefinition::SingleValueChunking{}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::WholeAxisChunking{}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); std::vector buf(view->countChunkValues()); @@ -113,13 +110,13 @@ CASE("ChunkedDataView | FixedSizeChunking | Combined axis with differently-sized "date=2020-01-01/to/2020-01-04,levtype=sfc," "param=10/20/30/40,time=0/6/12/18"}; - auto mock_extractor = std::make_shared(createMockFDB(9)); + auto mock_fdb = createMockFDB(9); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{4}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::FixedSizeChunking{2}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // Combined datextime axis: 16 values, chunkSize=4 -> 4 chunks. @@ -154,14 +151,14 @@ CASE("ChunkedDataView | FixedSizeChunking | Combined axis with differently-sized "date=2020-01-01/to/2020-01-04,levtype=ml,levelist=100/200/300/400/500/600/700/800/900/1000," "param=10/20/30/40/50/60,time=0/6/12"}; - auto mock_extractor = std::make_shared(createMockFDB(46)); + auto mock_fdb = createMockFDB(46); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{3}}, cdv::AxisDefinition{{"levelist"}, cdv::AxisDefinition::FixedSizeChunking{5}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::FixedSizeChunking{3}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build(); // Combined datextime: 12, chunkSize=3 -> 4 chunks @@ -197,7 +194,7 @@ CASE("ChunkedDataView | FixedSizeChunking | Combined axis with differently-sized {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{3}}, cdv::AxisDefinition{{"levelist"}, cdv::AxisDefinition::FixedSizeChunking{3}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::FixedSizeChunking{3}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); // param has 6 values -> FixedSizeChunking{4} does not divide evenly @@ -206,7 +203,7 @@ CASE("ChunkedDataView | FixedSizeChunking | Combined axis with differently-sized {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{3}}, cdv::AxisDefinition{{"levelist"}, cdv::AxisDefinition::FixedSizeChunking{5}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::FixedSizeChunking{4}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .build()); } @@ -227,18 +224,18 @@ CASE("ChunkedDataView | FixedSizeChunking | Multi part with combined axis with d "date=2020-01-01/to/2020-01-04,levtype=ml,levelist=100/200/300/400/500/600/700/800/900/1000," "param=10/20/30/40/50/60,time=0/6/12"}; - auto mock_extractor = std::make_shared(createMockFDB(10)); + auto mock_fdb = createMockFDB(10); const auto view = cdv::ChunkedDataViewBuilder() .addPart(keys_sfc, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{3}}, cdv::AxisDefinition{{"param"}, cdv::AxisDefinition::FixedSizeChunking{3}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .addPart(keys_ml, {cdv::AxisDefinition{{"date", "time"}, cdv::AxisDefinition::FixedSizeChunking{3}}, {cdv::AxisDefinition{{"levelist", "param"}, cdv::AxisDefinition::FixedSizeChunking{3}}}}, - mock_extractor) + FakeExtractorDefinition{mock_fdb}) .extendOnAxis(1) .build(); diff --git a/tests/conftest.py b/tests/conftest.py index b7f68ebdf..4081d567d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,7 @@ import shutil import eccodes as ec +import numpy import pytest import yaml from numpy import repeat @@ -28,6 +29,16 @@ def pytest_configure(config): config.addinivalue_line( "markers", "online: tests that require network access to ECMWF open data" ) + config.addinivalue_line( + "markers", + "zfdb_user_tests: end-to-end user-facing tests showing how to consume a z3fdb store " + "(e.g. via dask, xarray) without FDB internals knowledge", + ) + config.addinivalue_line( + "markers", + "gribjump: tests requiring the GribJump extractor (cmake feature " + "ZARR_GRIBJUMP_EXTRACTOR); skipped when the build does not provide it", + ) def create_fdb(root: pathlib.Path, schema_src: pathlib.Path) -> pathlib.Path: @@ -461,3 +472,205 @@ def build_pattern_grib_messages(data_path, session_tmp) -> pathlib.Path: ec.codes_release(gid) return messages + + +# --------------------------------------------------------------------------- +# Ramp-pattern data: distinguishes fields *and* grid points +# --------------------------------------------------------------------------- +# +# build_grib_messages sets the values once before its loop, so every field carries the +# identical ramp -- a test cannot tell whether a field landed in the right buffer slot. +# build_pattern_grib_messages gives each field a distinct *constant*, so it cannot tell +# whether a sub-range of the implicit (grid-point) axis was read correctly. +# +# This fixture varies along both axes at once: +# +# value(field, i) = field * _RAMP_FIELD_STRIDE + i +# +# The per-field term detects field -> slot mis-mapping; the +i ramp detects a wrong +# grid-point range. field is the index into product(dates, times, params), which for the +# z3fdb axes [date, time] x [param] is (dt_idx * len(params)) + param_idx. + +_RAMP_FIELD_STRIDE = 10_000 # > grid size, so the two components never alias + + +@pytest.fixture(scope="session") +def build_ramp_pattern_grib_messages(data_path, session_tmp) -> pathlib.Path: + """GRIB messages whose values identify both the field and the grid point. + + Uses the same date/time/param span as :func:`build_grib_messages`, so views built on it + keep the same shape and only the values differ. + + bitsPerValue is raised to 32: with simple packing the quantum is + ``(max - min) / 2**bitsPerValue``, and the default from the template is too coarse to + round-trip values approaching a million. Tests should still compare with ``atol=0.5`` + rather than exactly -- that absorbs any packing error while staying far below the + smallest difference that means anything here (1, one grid point). + """ + tmp = session_tmp / "build_ramp_pattern_grib_messages" + tmp.mkdir() + template_grib = data_path / "template.grib" + assert template_grib.is_file() + with open(template_grib, "rb") as template_grib_fd: + gid = ec.codes_grib_new_from_file(template_grib_fd) + + count_values = int(ec.codes_get(gid, "numberOfValues")) + assert int(ec.codes_get(gid, "numberOfDataPoints")) == count_values + assert int(ec.codes_get(gid, "numberOfMissing")) == 0 + assert count_values < _RAMP_FIELD_STRIDE, "field stride must exceed the grid size" + + ec.codes_set_string(gid, "type", "an") + ec.codes_set_string(gid, "class", "ea") + ec.codes_set_string(gid, "expver", "0001") + ec.codes_set_string(gid, "stream", "oper") + ec.codes_set_string(gid, "levtype", "sfc") + ec.codes_set(gid, "bitsPerValue", 32) + + dates = [20200101, 20200102, 20200103, 20200104] + times = [0, 300, 600, 900, 1200, 1500, 1800, 2100] + parameters = [167, 131, 132] + + ramp = numpy.arange(count_values) + + messages = tmp / "test_data_ramp.grib" + with open(messages, "wb") as out: + for field, (date, time, parameter) in enumerate( + itertools.product(dates, times, parameters) + ): + ec.codes_set(gid, "date", date) + ec.codes_set(gid, "time", time) + ec.codes_set(gid, "paramId", parameter) + ec.codes_set_values(gid, field * _RAMP_FIELD_STRIDE + ramp) + ec.codes_write(gid, out) + + ec.codes_release(gid) + return messages + + +@pytest.fixture(scope="session") +def read_only_fdb_ramp_setup(data_path, session_tmp, build_ramp_pattern_grib_messages) -> pathlib.Path: + """Read-only FDB holding the ramp-pattern data. See build_ramp_pattern_grib_messages.""" + fdb_root = session_tmp / "ramp-pattern-fdb" + fdb_root.mkdir() + cfg = create_fdb(fdb_root, data_path / "schema") + populate_fdb(cfg, [build_ramp_pattern_grib_messages]) + return cfg + + +@pytest.fixture(scope="session") +def ramp_expected(): + """Callable giving the expected values of one (datetime, param) slice of the ramp data. + + A fixture rather than a plain function so that tests in any subdirectory can reach it + without importing across directories. + + expected = ramp_expected(dt_idx, param_idx, count_params, count_values) + """ + + def _expected(dt_idx: int, param_idx: int, count_params: int, count_values: int): + field = dt_idx * count_params + param_idx + return field * _RAMP_FIELD_STRIDE + numpy.arange(count_values, dtype="float64") + + return _expected + + +# --------------------------------------------------------------------------- +# Bitmap data: exercises the missing-value path +# --------------------------------------------------------------------------- +# +# Every other fixture asserts numberOfMissing == 0, so nothing covers the code that turns +# GRIB bitmap sentinels into the configured fill value -- neither GribExtractor's +# std::replace nor GribJumpExtractor's mask[j/64][j%64] bit indexing, which is the most +# intricate line in the extractor. + +_BITMAP_SENTINEL = -9999.0 + +# Chosen to land in several distinct 64-bit mask words, including both sides of the +# 1312-value field-chunk boundary used by the GribJump tests, and the very first and last +# grid point. A word/bit indexing error shows up as a shifted mask. +_BITMAP_MISSING_INDICES = (0, 1, 63, 64, 65, 127, 1311, 1312, 1313, 2624, 5246, 5247) + + +@pytest.fixture(scope="session") +def bitmap_missing_indices() -> tuple: + """Grid-point indices flagged missing by the bitmap fixture. + + A fixture rather than a plain constant so that tests in any subdirectory can reach it + without importing across directories. + """ + return _BITMAP_MISSING_INDICES + + +@pytest.fixture(scope="session") +def build_bitmap_grib_messages(data_path, session_tmp) -> pathlib.Path: + """GRIB messages carrying a bitmap, over the ramp pattern. + + Present points follow the same ``field * stride + i`` ramp as + :func:`build_ramp_pattern_grib_messages`, so a test can check that masking did not + shift the surviving values. The points in :data:`_BITMAP_MISSING_INDICES` are flagged + missing and must read back as the view's fill value, not as the sentinel. + """ + tmp = session_tmp / "build_bitmap_grib_messages" + tmp.mkdir() + template_grib = data_path / "template.grib" + assert template_grib.is_file() + with open(template_grib, "rb") as template_grib_fd: + gid = ec.codes_grib_new_from_file(template_grib_fd) + + count_values = int(ec.codes_get(gid, "numberOfValues")) + assert max(_BITMAP_MISSING_INDICES) < count_values + + ec.codes_set_string(gid, "type", "an") + ec.codes_set_string(gid, "class", "ea") + ec.codes_set_string(gid, "expver", "0001") + ec.codes_set_string(gid, "stream", "oper") + ec.codes_set_string(gid, "levtype", "sfc") + ec.codes_set(gid, "bitsPerValue", 32) + + # Order matters: the sentinel and the bitmap flag must be set before the values, so + # that eccodes builds the bitmap from them. + ec.codes_set(gid, "missingValue", _BITMAP_SENTINEL) + ec.codes_set(gid, "bitmapPresent", 1) + + dates = [20200101, 20200102, 20200103, 20200104] + times = [0, 300, 600, 900, 1200, 1500, 1800, 2100] + parameters = [167, 131, 132] + + ramp = numpy.arange(count_values) + + messages = tmp / "test_data_bitmap.grib" + with open(messages, "wb") as out: + for field, (date, time, parameter) in enumerate( + itertools.product(dates, times, parameters) + ): + ec.codes_set(gid, "date", date) + ec.codes_set(gid, "time", time) + ec.codes_set(gid, "paramId", parameter) + + values = (field * _RAMP_FIELD_STRIDE + ramp).astype("float64") + values[list(_BITMAP_MISSING_INDICES)] = _BITMAP_SENTINEL + ec.codes_set_values(gid, values) + + # Without this the whole bitmap test suite passes vacuously: if eccodes did not + # actually build a bitmap, every point reads back as present and the assertions + # about masked points never get exercised. + assert int(ec.codes_get(gid, "bitmapPresent")) != 0, "no bitmap in the message" + assert int(ec.codes_get(gid, "numberOfMissing")) == len(_BITMAP_MISSING_INDICES), ( + f"expected {len(_BITMAP_MISSING_INDICES)} missing values, " + f"got {int(ec.codes_get(gid, 'numberOfMissing'))}" + ) + + ec.codes_write(gid, out) + + ec.codes_release(gid) + return messages + + +@pytest.fixture(scope="session") +def read_only_fdb_bitmap_setup(data_path, session_tmp, build_bitmap_grib_messages) -> pathlib.Path: + """Read-only FDB holding bitmapped data. See build_bitmap_grib_messages.""" + fdb_root = session_tmp / "bitmap-fdb" + fdb_root.mkdir() + cfg = create_fdb(fdb_root, data_path / "schema") + populate_fdb(cfg, [build_bitmap_grib_messages]) + return cfg diff --git a/tests/pychunked_data_view/CMakeLists.txt b/tests/pychunked_data_view/CMakeLists.txt index 3bbc99807..13a96f103 100644 --- a/tests/pychunked_data_view/CMakeLists.txt +++ b/tests/pychunked_data_view/CMakeLists.txt @@ -25,7 +25,7 @@ endfunction() set(test_files - test_chunked_data_view.py + fdb/test_chunked_data_view_fdb.py test_chunked_data_view_errors.py test_mars_selection.py ) diff --git a/tests/pychunked_data_view/test_chunked_data_view.py b/tests/pychunked_data_view/fdb/test_chunked_data_view_fdb.py similarity index 96% rename from tests/pychunked_data_view/test_chunked_data_view.py rename to tests/pychunked_data_view/fdb/test_chunked_data_view_fdb.py index 545713557..5f07805b6 100644 --- a/tests/pychunked_data_view/test_chunked_data_view.py +++ b/tests/pychunked_data_view/fdb/test_chunked_data_view_fdb.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: 2025 European Centre for Medium-Range Weather Forecasts (ECMWF) # SPDX-License-Identifier: Apache-2.0 -import numpy as np import itertools +import numpy as np from pychunked_data_view import ( AxisDefinition, @@ -40,7 +40,7 @@ def test_axis_definition_individual_chunk(): assert obj.chunking == Chunking.FixedSizeChunk(2) -def test_builder(read_only_fdb_setup): +def test_builder_fdb(read_only_fdb_setup): builder = ChunkedDataViewBuilder(read_only_fdb_setup) builder.add_part( { @@ -59,7 +59,7 @@ def test_builder(read_only_fdb_setup): AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.fill_missing_value(-20.0) view = builder.build() diff --git a/tests/pychunked_data_view/test_chunked_data_view_errors.py b/tests/pychunked_data_view/test_chunked_data_view_errors.py index f02731e1a..9ac84b5b9 100644 --- a/tests/pychunked_data_view/test_chunked_data_view_errors.py +++ b/tests/pychunked_data_view/test_chunked_data_view_errors.py @@ -34,7 +34,7 @@ def test_malformed_mars_string_raises(read_only_fdb_setup, malformed_request, ex builder = ChunkedDataViewBuilder(read_only_fdb_setup) # Inject the malformed MARS string at the bindings level — addPart() stores it # without parsing, so no error fires here; it fires in build(). - builder._obj.add_part(malformed_request, [], ExtractorType.GRIB.value) + builder._obj.add_part(malformed_request, [], ExtractorType.Grib()._obj) with pytest.raises(MarsRequestFormattingError, match=expected_hint): builder.build() @@ -62,7 +62,7 @@ def test_misspelled_mars_key_raises(read_only_fdb_setup) -> None: "klasse": "ea", # 'klasse' is not a valid MARS key }, [AxisDefinition(["param"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(MarsRequestFormattingError, match="Did you misspell a MARS key"): builder.build() @@ -75,3 +75,16 @@ def test_invalid_chunking_type_raises() -> None: """ with pytest.raises(TypeError, match="chunking must be Chunking"): AxisDefinition(["param"], object()) + + +@pytest.mark.parametrize("chunk_shape", [0], ids=["zero"]) +def test_invalid_fixed_chunk_size_rejected_on_construction(chunk_shape) -> None: + """A non-positive chunk size is refused by FixedSizeChunking itself. + + AxisDefinition::FixedSizeChunking asserts chunkSize > 0 in its constructor, so this fails + as soon as the chunking object is built — before any AxisDefinition, builder, request or + FDB is involved. Placed here rather than with the extractor tests because it is a property + of the chunking type, independent of which extractor consumes it. + """ + with pytest.raises(Exception, match="The supplied chunk shape needs to be positive"): + Chunking.FixedSizeChunk(chunk_shape=chunk_shape) diff --git a/tests/z3fdb/CMakeLists.txt b/tests/z3fdb/CMakeLists.txt index 01671f66a..099bb55a3 100644 --- a/tests/z3fdb/CMakeLists.txt +++ b/tests/z3fdb/CMakeLists.txt @@ -43,16 +43,31 @@ function(add_py_online_test test) endfunction() -set(unit_tests - unit/test_store_v3.py - unit/test_store_v3_errors.py - unit/test_u_v_vo_d_retrieval.py +set(store_tests + interface/store/test_store_v3_simple_builder.py + interface/store/test_store_v3_errors_simple_builder.py + interface/store/test_u_v_vo_d_retrieval_simple_builder.py + interface/store/test_custom_store_builder.py +) + +# GribJump-backed integration tests. Registered unconditionally on purpose: the conftest in +# integration/gribjump/ marks them and tests/z3fdb/conftest.py skips them when the build has no +# GribJump extractor, so ctest reports a skip instead of the test silently not existing. +set(integration_tests + integration/gribjump/test_gribjump_extractor.py + integration/gribjump/test_mixed_extractors.py ) set(zarr_interface_conformity_tests - zarr_interface_conformity/test_z3fdb_store_interface.py - zarr_interface_conformity/test_z3fdb_store_path_correctness.py - zarr_interface_conformity/test_z3fdb_consolidated_metadata.py + interface/zarr_interface_conformity/test_z3fdb_store_interface.py + interface/zarr_interface_conformity/test_z3fdb_store_path_correctness.py + interface/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py +) + +set(user_tests + interface/user_tests/test_dask_access.py + interface/user_tests/test_xarray_access.py + interface/user_tests/test_metadata_mapping.py ) set(permutation_tests @@ -60,14 +75,14 @@ set(permutation_tests permutation_tests/test_axis_definition_ordering.py permutation_tests/test_axis_definition_permutations.py permutation_tests/test_store_fixed_size_chunking.py + permutation_tests/test_builder_extension.py permutation_tests/test_store_multiple_parts.py ) -foreach(test_file ${unit_tests} ${zarr_interface_conformity_tests} ${permutation_tests}) +foreach(test_file ${store_tests} ${integration_tests} ${zarr_interface_conformity_tests} ${user_tests} ${permutation_tests}) add_py_test(${test_file}) endforeach() # Online tests: require network access to ECMWF open data. # Run with: ctest -L z3fdb_online -add_py_online_test(test_store_v3_online.py) add_py_online_test(permutation_tests/test_store_missing_values.py) diff --git a/tests/z3fdb/conftest.py b/tests/z3fdb/conftest.py new file mode 100644 index 000000000..c6c1487ab --- /dev/null +++ b/tests/z3fdb/conftest.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +"""Shared gating for z3fdb tests that depend on optional build features.""" + +import pytest + + +def _has_gribjump_extractor() -> bool: + """Return True if this build compiled the GribJump extractor. + + ``ExtractorType.GribJump`` is bound in every build so that user code does not depend on + cmake flags, so its mere presence proves nothing — the extension module reports the build + capability separately. + """ + try: + import chunked_data_view_bindings as pdv + except ImportError: + return False + return bool(getattr(pdv, "has_gribjump_extractor", False)) + + +HAS_GRIBJUMP_EXTRACTOR = _has_gribjump_extractor() + +_GRIBJUMP_SKIP_REASON = ( + "build has no GribJump extractor; configure with -DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON " + "(requires a bundle build providing gribjump)" +) + + +def pytest_runtest_setup(item): + """Skip tests marked ``gribjump`` when the extractor was not compiled. + + Applies to every test under ``tests/z3fdb``, wherever it lives — mark an individual test + with ``@pytest.mark.gribjump``, or a whole folder via a conftest that adds the marker (see + ``integration/gribjump/conftest.py``). + """ + if not HAS_GRIBJUMP_EXTRACTOR and item.get_closest_marker("gribjump") is not None: + pytest.skip(_GRIBJUMP_SKIP_REASON) diff --git a/tests/z3fdb/integration/gribjump/test_gribjump_extractor.py b/tests/z3fdb/integration/gribjump/test_gribjump_extractor.py new file mode 100644 index 000000000..767fe9179 --- /dev/null +++ b/tests/z3fdb/integration/gribjump/test_gribjump_extractor.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +"""Integration tests for the GribJump extractor via SimpleStoreBuilder.""" + +import logging + +import numpy as np +import pytest +import zarr + +from chunked_data_view_bindings import has_gribjump_extractor +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +log = logging.getLogger(__name__) + +pytestmark = [ + pytest.mark.offline, + pytest.mark.gribjump, + pytest.mark.skipif( + not has_gribjump_extractor, + reason=( + "build has no GribJump extractor; configure with " + "-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON (requires a bundle build providing gribjump)" + ), + ), +] + +_TOTAL_VALUES = 5248 + +_COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": "2020-01-01/to/2020-01-04", + "levtype": "sfc", + "step": 0, + "param": [167, 131, 132], + "time": "0/to/21/by/3", +} + + +@pytest.mark.parametrize( + ("chunk_size", "n_implicit_chunks"), + [ + (1312, 4), # 5248 / 1312 = 4 (exact) + (2624, 2), # 5248 / 2624 = 2 (exact) + ], + ids=["chunk-1312", "chunk-2624"], +) +def test_gribjump_field_chunking_shape(read_only_fdb_ramp_setup, chunk_size, n_implicit_chunks) -> None: + """FixedSizeChunking on the implicit axis produces the expected zarr chunk shape. + + The implicit (grid-point) dimension must be split into chunks of *chunk_size*, + leaving the overall shape unchanged and reflecting the chunk size in the zarr + metadata. + """ + builder = SimpleStoreBuilder(read_only_fdb_ramp_setup) + builder.add_part( + _COMMON, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=8)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=chunk_size)), + ) + store = builder.build() + data = zarr.open_array(store) + + assert data.shape == (32, 3, _TOTAL_VALUES), f"unexpected shape {data.shape}" + assert data.chunks == (8, 1, chunk_size), f"unexpected chunk shape {data.chunks}" + assert data.nchunks == 4 * 3 * n_implicit_chunks, f"unexpected chunk count {data.nchunks}" + + +@pytest.mark.parametrize( + "chunk_size", + [1312, 2624], + ids=["chunk-1312", "chunk-2624"], +) +def test_gribjump_field_chunking_values(read_only_fdb_ramp_setup, ramp_expected, chunk_size) -> None: + """Every field lands in the right slot, with the right grid-point range. + + Uses the ramp fixture, where ``value(field, i) = field * 10000 + i``. That matters: with + the older fixture every field carried an identical ramp, so this test passed even if the + key -> buffer-slot mapping was completely scrambled. The per-field term now pins the slot + and the ``+ i`` term pins the sub-range, so both failure modes are visible. + + Compared with atol=0.5 to absorb GRIB packing, which is far below the smallest difference + that means anything here: 1 for a neighbouring grid point, 10000 for a wrong field. + """ + builder = SimpleStoreBuilder(read_only_fdb_ramp_setup) + builder.add_part( + _COMMON, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=8)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=chunk_size)), + ) + store = builder.build() + data = zarr.open_array(store) + + count_params = data.shape[1] + n_chunks = _TOTAL_VALUES // chunk_size + + # Full assembled slice per (datetime, param): catches a field in the wrong slot. + for dt_idx in range(data.shape[0]): + for param_idx in range(count_params): + np.testing.assert_allclose( + data[dt_idx, param_idx], + ramp_expected(dt_idx, param_idx, count_params, _TOTAL_VALUES), + atol=0.5, + err_msg=f"full-field mismatch at datetime={dt_idx}, param={param_idx}", + ) + + # Each implicit chunk read on its own: catches a wrong grid-point range. + expected_last = ramp_expected(data.shape[0] - 1, count_params - 1, count_params, _TOTAL_VALUES) + for chunk_k in range(n_chunks): + start = chunk_k * chunk_size + end = start + chunk_size + logging.debug(f"Chunk {chunk_k}: {data[-1, -1, start:end]}") + np.testing.assert_allclose( + data[-1, -1, start:end], + expected_last[start:end], + atol=0.5, + err_msg=f"sub-range mismatch for implicit chunk {chunk_k} [{start}:{end}]", + ) + + +@pytest.mark.parametrize( + "chunk_size", + [ + 1000, # 5248 % 1000 != 0 + 5247, # one short of the whole field + ], + ids=["not-a-divisor", "off-by-one"], +) +def test_gribjump_field_chunk_size_must_divide_the_grid(read_only_fdb_ramp_setup, chunk_size) -> None: + """A field chunk size that cannot tile the grid must fail at build() and say why. + + The implicit axis is the one dimension a view cannot leave ragged, so the size has to + divide the grid exactly. + """ + builder = SimpleStoreBuilder(read_only_fdb_ramp_setup) + builder.add_part( + _COMMON, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=8)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=chunk_size)), + ) + with pytest.raises(RuntimeError, match="does not evenly divide"): + builder.build() + + + +@pytest.mark.parametrize("chunk_size", [None, 1312], ids=["whole-axis", "chunk-1312"]) +def test_gribjump_bitmap_missing_values( + read_only_fdb_bitmap_setup, ramp_expected, bitmap_missing_indices, chunk_size +) -> None: + """Bitmap-masked grid points read back as the fill value, and the rest are not shifted. + + This is the only coverage of GribJumpExtractor's ``mask[j / 64][j % 64]`` indexing. The + masked indices deliberately straddle 64-bit word boundaries and the 1312 field-chunk + boundary, so a word/bit mix-up shows up as a shifted mask rather than passing by luck. + Parametrised over the chunked case too, since that is where the range offset and the mask + offset could disagree. + """ + fill = -1234.0 + field_chunking = None if chunk_size is None else Chunking.FixedSizeChunk(chunk_shape=chunk_size) + + builder = SimpleStoreBuilder(read_only_fdb_bitmap_setup) + builder.add_part( + _COMMON, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=8)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.GribJump(field_chunking=field_chunking), + ) + builder.fill_missing_value(fill) + data = zarr.open_array(builder.build()) + + missing = np.array(bitmap_missing_indices) + present = np.setdiff1d(np.arange(_TOTAL_VALUES), missing) + + for dt_idx, param_idx in ((0, 0), (31, 2)): + slice_ = data[dt_idx, param_idx] + expected = ramp_expected(dt_idx, param_idx, data.shape[1], _TOTAL_VALUES) + + np.testing.assert_array_equal( + slice_[missing], + np.full(missing.size, fill, dtype=np.float32), + err_msg=f"masked points not filled at datetime={dt_idx}, param={param_idx}", + ) + np.testing.assert_allclose( + slice_[present], + expected[present], + atol=0.5, + err_msg=f"unmasked points shifted at datetime={dt_idx}, param={param_idx}", + ) diff --git a/tests/z3fdb/integration/gribjump/test_mixed_extractors.py b/tests/z3fdb/integration/gribjump/test_mixed_extractors.py new file mode 100644 index 000000000..6600a1421 --- /dev/null +++ b/tests/z3fdb/integration/gribjump/test_mixed_extractors.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +"""Integration test: SimpleStoreBuilder with one GRIB part and one GribJump part. + +Data layout (from ``build_grib_messages`` in conftest.py): + - 4 dates (2020-01-01 .. 2020-01-04) x 8 times (0,3,..,21 h) = 32 datetime entries + - 3 params: 167, 131, 132 (all fields carry values list(range(0, grid_points))) + +The test splits params across two extractors: + - Part 1 (Grib): params [167, 131] + - Part 2 (GribJump): param [132] + +The datetime axis uses FixedSizeChunking(8) -> 4 chunks of 8. +The param axis uses SingleValueChunking and is the extension axis. + +Resulting view shape: (32, 3, grid_points) +Chunk shape: (8, 1, grid_points) +""" + +import logging + +import numpy as np +import pytest +import zarr + +from chunked_data_view_bindings import has_gribjump_extractor +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +log = logging.getLogger(__name__) + +pytestmark = [ + pytest.mark.offline, + pytest.mark.gribjump, + pytest.mark.skipif( + not has_gribjump_extractor, + reason=( + "build has no GribJump extractor; configure with " + "-DENABLE_ZARR_GRIBJUMP_EXTRACTOR=ON (requires a bundle build providing gribjump)" + ), + ), +] + +_COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": "2020-01-01/to/2020-01-04", + "levtype": "sfc", + "step": 0, + "time": "0/to/21/by/3", +} + +# 4 dates x 8 times = 32 entries -> 4 chunks of 8 +_DATETIME_AXIS = AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=8)) +_PARAM_AXIS = AxisDefinition(["param"], Chunking.SINGLE_VALUE) + + +def test_grib_and_gribjump_parts(read_only_fdb_ramp_setup, ramp_expected) -> None: + """A GRIB part and a GribJump part agree on both values and field ordering. + + Both extractors read the same FDB; GRIB does a full decode, GribJump jumps to the values. + The ramp fixture gives every field a distinct base, so this checks the two extractors put + the *same field* in the *same slot* — with the older fixture every field carried an + identical ramp, so a scrambled key -> slot mapping would still have passed. + + Compared with atol=0.5 to absorb GRIB packing, well under the 1 / 10000 differences that + a wrong grid point or a wrong field would produce. + """ + builder = SimpleStoreBuilder(read_only_fdb_ramp_setup) + + # Part 1 - standard GRIB decode (params 167 and 131) + builder.add_part( + {**_COMMON, "param": [167, 131]}, + [_DATETIME_AXIS, _PARAM_AXIS], + ExtractorType.Grib(), + ) + + # Part 2 - GribJump partial decode (param 132) + builder.add_part( + {**_COMMON, "param": [132]}, + [_DATETIME_AXIS, _PARAM_AXIS], + ExtractorType.GribJump(), + ) + + builder.extend_on_axis(1) + store = builder.build() + + data = zarr.open_array(store) + log.debug("shape=%s chunks=%s", data.shape, data.chunks) + + assert data.shape[:2] == (32, 3), f"expected (32, 3, ...), got {data.shape}" + assert data.chunks[:2] == (8, 1), f"expected chunks (8, 1, ...), got {data.chunks}" + + # param axis order is [167, 131] from the GRIB part then [132] from the GribJump part, + # which is the order the ramp fixture wrote them in, so the expectation is the plain + # field index. Slots 0 and 1 are served by GribExtractor, slot 2 by GribJumpExtractor. + grid_points = data.shape[-1] + count_params = data.shape[1] + + for dt_idx in range(data.shape[0]): + for param_idx in range(count_params): + np.testing.assert_allclose( + data[dt_idx, param_idx], + ramp_expected(dt_idx, param_idx, count_params, grid_points), + atol=0.5, + err_msg=f"mismatch at datetime={dt_idx}, param={param_idx}", + ) + + +@pytest.mark.parametrize("gribjump_first", [False, True], ids=["grib-first", "gribjump-first"]) +def test_field_chunking_mixed_with_grib_is_rejected(read_only_fdb_ramp_setup, gribjump_first) -> None: + """A GribJump part using field chunking cannot be mixed with a GRIB part. + + ``GribExtractor`` always returns the whole field, so its chunk on the implicit axis is the + full grid; a ``GribJumpExtractor`` with ``FixedSizeChunk`` writes a smaller block. The view + takes that extent from the *first* part only, so the mismatch must be rejected at build() + time — otherwise the part whose block is larger overruns its slots at read time while still + reporting the expected message count. + """ + grib_part = ({**_COMMON, "param": [167, 131]}, ExtractorType.Grib()) + gribjump_part = ( + {**_COMMON, "param": [132]}, + ExtractorType.GribJump(field_chunking=Chunking.FixedSizeChunk(chunk_shape=1312)), + ) + + first, second = (gribjump_part, grib_part) if gribjump_first else (grib_part, gribjump_part) + + builder = SimpleStoreBuilder(read_only_fdb_ramp_setup) + builder.add_part(first[0], [_DATETIME_AXIS, _PARAM_AXIS], first[1]) + builder.add_part(second[0], [_DATETIME_AXIS, _PARAM_AXIS], second[1]) + builder.extend_on_axis(1) + + with pytest.raises(RuntimeError, match="WholeAxisChunking") as exc: + builder.build() + + logging.debug(f"{exc.value.args[0]}") diff --git a/tests/z3fdb/interface/store/test_custom_store_builder.py b/tests/z3fdb/interface/store/test_custom_store_builder.py new file mode 100644 index 000000000..d2570e447 --- /dev/null +++ b/tests/z3fdb/interface/store/test_custom_store_builder.py @@ -0,0 +1,456 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 +"""Tests for CustomStoreBuilder: hierarchy construction, repeated parts, extension, and +cross-type name-collision detection. + +Fixture data (from conftest build_pattern_grib_messages): + dates = [20200101, 20200102, 20200103] (3 dates, index d) + times = [0, 600, 1200, 1800] (4 times, index t) + params_sfc = [165, 166, 167] (3 params, index p) + params_pl = [131, 132, 133] (3 params, index p) + levels = [50, 100, 150] (3 levels, index l) + +Field-value formulas (all indices 0-based): + sfc_value(d, t, p) = d*12 + t*3 + p (sequential over product(dates,times,params_sfc)) + pl_value(d, t, p, l) = 36 + d*36 + t*9 + p*3 + l +""" + +import numpy as np +import pytest +import zarr + +from z3fdb import AxisDefinition, Chunking, CustomStoreBuilder, ExtractorType + +pytestmark = [pytest.mark.offline, pytest.mark.zfdb_user_tests] + +# --------------------------------------------------------------------------- +# Shared request helpers +# --------------------------------------------------------------------------- + +_COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, +} + +_SFC_REQUEST = { + **_COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], +} + +_SFC_AXES = [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), +] + +_PL_REQUEST = { + **_COMMON, + "levtype": "pl", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [131, 132, 133], + "levelist": [50, 100, 150], +} + +_PL_AXES = [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), +] + + +# Structural / data tests (require read_only_fdb_pattern_setup) +def test_single_group_single_array(read_only_fdb_pattern_setup) -> None: + """One group containing one array — smoke test for the minimal hierarchy. + + Verifies: + - zarr can navigate into the group and open the array. + - Array shape reflects 12 date×time entries, 3 params, plus the implicit axis. + - Spot-check values match the known sfc_value formula. + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part("sfc/wind", _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + store = builder.build() + + root = zarr.open_group(store, mode="r") + assert "sfc" in root, "top-level group 'sfc' missing from store" + arr = root["sfc"]["wind"] + + assert arr.shape[:2] == (12, 3), f"unexpected shape {arr.shape}" + + # sfc_value(d=0, t=0, p=0) = 0 + assert np.all(arr[0, 0] == 0), f"unexpected value at [0,0]: {arr[0, 0, 0]}" + # sfc_value(d=2, t=3, p=2) = 24+9+2 = 35 + assert np.all(arr[11, 2] == 35), f"unexpected value at [11,2]: {arr[11, 2, 0]}" + + +def test_sibling_arrays_same_group(read_only_fdb_pattern_setup) -> None: + """Two sibling arrays under one parent group — both must be accessible and correct. + + 'sfc' carries 3 params (shape (12, 3, N)). + 'pl' carries 3 params × 3 levels combined (shape (12, 9, N)). + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part("weather/sfc", _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + builder.add_part("weather/pl", _PL_REQUEST, _PL_AXES, ExtractorType.Grib()) + store = builder.build() + + root = zarr.open_group(store, mode="r") + weather = root["weather"] + + assert "sfc" in weather, "'sfc' array missing from 'weather' group" + assert "pl" in weather, "'pl' array missing from 'weather' group" + + sfc = weather["sfc"] + pl = weather["pl"] + + assert sfc.shape[:2] == (12, 3), f"sfc shape {sfc.shape}" + assert pl.shape[:2] == (12, 9), f"pl shape {pl.shape}" + + # sfc spot-check: sfc_value(d=0, t=0, p=0) = 0 + assert np.all(sfc[0, 0] == 0) + # sfc spot-check: sfc_value(d=2, t=3, p=2) = 35 + assert np.all(sfc[11, 2] == 35) + + # pl spot-check: pl_value(d=0, t=0, p=0, l=0) = 36 + # combined pl_idx = p*3+l = 0 + assert np.all(pl[0, 0] == 36) + # pl spot-check: pl_value(d=2, t=3, p=2, l=2) = 36+72+27+6+2 = 143 + # combined dt_idx=11, pl_idx=8 + assert np.all(pl[11, 8] == 143) + + +def test_deep_nesting(read_only_fdb_pattern_setup) -> None: + """Arrays at depth-3 paths sharing a common grandparent group. + + Path A: ["analysis", "sfc", "wind"] + Path B: ["analysis", "pl", "wind"] + + Both arrays must be accessible; group "analysis" must contain exactly + the two sub-groups "sfc" and "pl". + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part("analysis/sfc/wind", _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + builder.add_part("analysis/pl/wind", _PL_REQUEST, _PL_AXES, ExtractorType.Grib()) + store = builder.build() + + root = zarr.open_group(store, mode="r") + assert "analysis" in root + + # The docstring's claim, actually asserted: exactly these two sub-groups, and no stray + # array left at the group level by the path-splitting logic. + analysis = root["analysis"] + assert sorted(analysis.group_keys()) == ["pl", "sfc"] + assert sorted(analysis.array_keys()) == [] + + sfc_wind = analysis["sfc"]["wind"] + pl_wind = analysis["pl"]["wind"] + + assert sfc_wind.shape[:2] == (12, 3) + assert pl_wind.shape[:2] == (12, 9) + + assert np.all(sfc_wind[0, 0] == 0) + assert np.all(pl_wind[0, 0] == 36) + + +def test_repeated_add_part_multi_part(read_only_fdb_pattern_setup) -> None: + """Calling add_part twice with the same path registers two parts on one builder. + + Part 1: param=165 → param-axis index 0 → values = d*12 + t*3 + 0 + Part 2: param=166 → param-axis index 1 → values = d*12 + t*3 + 1 + + extend_on_axis(path, 1) declares that the param axis is the extension axis. + The assembled array must have shape (12, 2, N). + """ + path = "data/wind" + axes_one_param = [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + path, + { + **_COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165], + }, + axes_one_param, + ExtractorType.Grib(), + ) + builder.extend_on_axis(path, 1) + builder.add_part( + path, + { + **_COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [166], + }, + axes_one_param, + ExtractorType.Grib(), + ) + store = builder.build() + + root = zarr.open_group(store, mode="r") + arr = root["data"]["wind"] + + assert arr.shape[:2] == (12, 2), f"expected (12, 2, N), got {arr.shape}" + + # param=165 is axis-1 index 0; sfc_value(d=0,t=0,p=0) = 0 + assert np.all(arr[0, 0] == 0) + # param=166 is axis-1 index 1; sfc_value(d=0,t=0,p=1) = 1 + assert np.all(arr[0, 1] == 1) + # sfc_value(d=2,t=3,p=0) = 33 and sfc_value(d=2,t=3,p=1) = 34 + assert np.all(arr[11, 0] == 33) + assert np.all(arr[11, 1] == 34) + + +def test_extend_on_axis_shape(read_only_fdb_pattern_setup) -> None: + """extend_on_axis via CustomStoreBuilder correctly stitches two SFC + PL parts. + + SFC part: 3 params → axis-1 indices 0..2 + PL part: 9 (param×level) → axis-1 indices 3..11 + + Both parts share the same date×time axis (12 entries). + """ + path = "combined/all" + dt_ax = AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE) + + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + path, + _SFC_REQUEST, + [dt_ax, AxisDefinition(["param"], Chunking.SINGLE_VALUE)], + ExtractorType.Grib(), + ) + builder.extend_on_axis(path, 1) + builder.add_part( + path, + _PL_REQUEST, + [dt_ax, AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE)], + ExtractorType.Grib(), + ) + store = builder.build() + + arr = zarr.open_group(store, mode="r")["combined"]["all"] + + assert arr.shape[:2] == (12, 12), f"expected (12, 12, N), got {arr.shape}" + + # SFC values occupy axis-1 indices 0..2 + assert np.all(arr[0, 0] == 0) # sfc_value(0,0,0) + assert np.all(arr[11, 2] == 35) # sfc_value(2,3,2) + + # PL values occupy axis-1 indices 3..11 + assert np.all(arr[0, 3] == 36) # pl_value(0,0,0,0) + assert np.all(arr[11, 11] == 143) # pl_value(2,3,2,2) + + +# Collision detection tests (no FDB access needed — errors fire before build()) +_DUMMY_REQUEST = {"type": "an", "class": "ea"} +_DUMMY_AXES: list[AxisDefinition] = [] +_DUMMY_EXTRACTOR = ExtractorType.Grib() + + +def test_array_then_group_collision_raises() -> None: + """Registering an array at 'foo' and then requesting 'foo/bar' must raise ValueError. + + The first call registers 'foo' as a leaf array under the root group. + The second call attempts to treat 'foo' as a parent group — a cross-type collision. + """ + builder = CustomStoreBuilder() + builder.add_part("foo", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + with pytest.raises(ValueError, match="foo"): + builder.add_part("foo/bar", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_group_then_array_collision_raises() -> None: + """Registering 'foo/bar' and then requesting 'foo' as an array must raise ValueError. + + The first call creates an intermediate group named 'foo'. + The second call attempts to register an array also named 'foo' — a cross-type collision. + """ + builder = CustomStoreBuilder() + builder.add_part("foo/bar", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + with pytest.raises(ValueError, match="foo"): + builder.add_part("foo", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_deeply_nested_collision_raises() -> None: + """Cross-type collision detected at a non-root level. + + Registers 'a/b/arr', making 'a' and 'b' intermediate groups. + Then tries to register 'a/b' as an array — 'b' is already a group. + """ + builder = CustomStoreBuilder() + builder.add_part("a/b/arr", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + with pytest.raises(ValueError, match="b"): + builder.add_part("a/b", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_empty_path_raises() -> None: + """An empty path string must raise ValueError before any tree mutation.""" + builder = CustomStoreBuilder() + with pytest.raises(ValueError, match="empty"): + builder.add_part("", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_root_array(read_only_fdb_pattern_setup) -> None: + """path=None places the array at the store root — no enclosing group. + + The store must be openable as a zarr array directly via + ``zarr.open_array(store)`` and its shape and values must match the + SFC dataset. + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part(None, _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + store = builder.build() + + arr = zarr.open_array(store, mode="r") + assert arr.shape[:2] == (12, 3), f"unexpected root-array shape {arr.shape}" + # sfc_value(d=0, t=0, p=0) = 0 + assert np.all(arr[0, 0] == 0) + # sfc_value(d=2, t=3, p=2) = 35 + assert np.all(arr[11, 2] == 35) + + +def test_root_array_multi_part(read_only_fdb_pattern_setup) -> None: + """Two add_part(None, ...) calls with extend_on_axis(None, ...) between them. + + Part 1: param=165 → axis-1 index 0 + Part 2: param=166 → axis-1 index 1 + + The assembled root array must have shape (12, 2, N). + """ + axes_one_param = [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + None, + {**_COMMON, "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], "param": [165]}, + axes_one_param, + ExtractorType.Grib(), + ) + builder.extend_on_axis(None, 1) + builder.add_part( + None, + {**_COMMON, "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], "param": [166]}, + axes_one_param, + ExtractorType.Grib(), + ) + store = builder.build() + + arr = zarr.open_array(store, mode="r") + assert arr.shape[:2] == (12, 2), f"expected (12, 2, N), got {arr.shape}" + assert np.all(arr[0, 0] == 0) # sfc_value(0, 0, 0) + assert np.all(arr[0, 1] == 1) # sfc_value(0, 0, 1) + + +def test_root_array_then_named_raises() -> None: + """Registering a root array (path=None) then a named path must raise ValueError.""" + builder = CustomStoreBuilder() + builder.add_part(None, _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + with pytest.raises(ValueError, match="root array"): + builder.add_part("other", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_named_then_root_raises() -> None: + """Registering a named path then a root array (path=None) must raise ValueError.""" + builder = CustomStoreBuilder() + builder.add_part("named", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + with pytest.raises(ValueError, match="root array"): + builder.add_part(None, _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + +def test_top_level_arrays(read_only_fdb_pattern_setup) -> None: + """Two arrays registered directly at the store root (no parent group). + + Single-segment paths ``"sfc"`` and ``"pl"`` must produce zarr arrays + accessible directly on the root group — no intervening group level. + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part("sfc", _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + builder.add_part("pl", _PL_REQUEST, _PL_AXES, ExtractorType.Grib()) + store = builder.build() + + root = zarr.open_group(store, mode="r") + assert "sfc" in root, "top-level array 'sfc' missing from store root" + assert "pl" in root, "top-level array 'pl' missing from store root" + + sfc = root["sfc"] + pl = root["pl"] + + assert sfc.shape[:2] == (12, 3), f"sfc shape {sfc.shape}" + assert pl.shape[:2] == (12, 9), f"pl shape {pl.shape}" + + # sfc_value(d=0, t=0, p=0) = 0 + assert np.all(sfc[0, 0] == 0) + # pl_value(d=0, t=0, p=0, l=0) = 36 + assert np.all(pl[0, 0] == 36) + + +# extend_on_axis / fill_missing_value configure an existing array, so an unknown path is a +# mistake. Before this was enforced they created an empty array instead, and the mistake only +# surfaced at build() as "must add at least one part". +def test_extend_on_axis_unknown_path_raises() -> None: + """Addressing a path that was never given an add_part must raise, not create it.""" + builder = CustomStoreBuilder() + with pytest.raises(ValueError, match="no array registered"): + builder.extend_on_axis("never/registered", 0) + + +def test_fill_missing_value_unknown_path_raises() -> None: + """Same rule for fill_missing_value.""" + builder = CustomStoreBuilder() + with pytest.raises(ValueError, match="no array registered"): + builder.fill_missing_value("never/registered", -999.0) + + +def test_configuring_a_mistyped_path_raises_and_names_the_real_one() -> None: + """The realistic case: a typo in an otherwise valid path. + + The message must name the registered arrays, otherwise the reader has no way to spot the + transposition without re-reading their own code. + """ + builder = CustomStoreBuilder() + builder.add_part("sfc/wind", _DUMMY_REQUEST, _DUMMY_AXES, _DUMMY_EXTRACTOR) + + with pytest.raises(ValueError, match="sfc/wind") as excinfo: + builder.extend_on_axis("sfc/wnid", 1) + assert "sfc/wnid" in str(excinfo.value), "the rejected path should be quoted too" + + +def test_configuring_root_before_add_part_raises() -> None: + """path=None is not exempt: the root array has to be registered first.""" + builder = CustomStoreBuilder() + with pytest.raises(ValueError, match="root array"): + builder.fill_missing_value(None, -999.0) + + +def test_fill_missing_value_reaches_the_zarr_metadata(read_only_fdb_pattern_setup) -> None: + """The fill value is threaded through to the array metadata, not merely accepted. + + Nothing else exercises CustomStoreBuilder.fill_missing_value, so without this the plumbing + through FdbSource to DotZarrArrayJson is untested for the custom builder. + """ + builder = CustomStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part("sfc/wind", _SFC_REQUEST, _SFC_AXES, ExtractorType.Grib()) + builder.fill_missing_value("sfc/wind", -999.0) + + arr = zarr.open_group(builder.build(), mode="r")["sfc"]["wind"] + assert arr.fill_value == -999.0, f"unexpected fill_value {arr.fill_value}" diff --git a/tests/z3fdb/unit/test_store_v3_errors.py b/tests/z3fdb/interface/store/test_store_v3_errors_simple_builder.py similarity index 51% rename from tests/z3fdb/unit/test_store_v3_errors.py rename to tests/z3fdb/interface/store/test_store_v3_errors_simple_builder.py index 65f5511e8..2de92043a 100644 --- a/tests/z3fdb/unit/test_store_v3_errors.py +++ b/tests/z3fdb/interface/store/test_store_v3_errors_simple_builder.py @@ -6,12 +6,9 @@ import pytest from pychunked_data_view.exceptions import MarsRequestFormattingError -from z3fdb import ( - AxisDefinition, - Chunking, - ExtractorType, - SimpleStoreBuilder, -) +from z3fdb import AxisDefinition, Chunking, ExtractorType, SimpleStoreBuilder + +from chunked_data_view_bindings import has_gribjump_extractor logging.basicConfig(level=logging.DEBUG) @@ -46,7 +43,7 @@ def test_extend_on_invalid_axis_raises( AxisDefinition(["time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(invalid_axis) @@ -80,8 +77,70 @@ def test_wrong_key( AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["date"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(MarsRequestFormattingError): builder.build() + + +@pytest.mark.parametrize( + ("bad_keys", "match"), + [ + (["steps"], "steps"), # single typo + (["date", "times"], "times"), # one correct, one typo + ], +) +def test_axis_key_absent_from_request_raises(bad_keys, match) -> None: + """AxisDefinition keys that are not present in the MARS request must raise + ValueError at add_part time — before any FDB connection is attempted.""" + builder = SimpleStoreBuilder() + with pytest.raises(ValueError, match=match): + builder.add_part( + { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "levtype": "sfc", + "date": "2020-01-01", + "time": 0, + "step": 0, + "param": 167, + }, + [AxisDefinition(bad_keys, Chunking.SINGLE_VALUE)], + ExtractorType.Grib(), + ) + + +@pytest.mark.skipif( + has_gribjump_extractor, + reason="build compiled the GribJump extractor", +) +def test_gribjump_without_build_support_raises(read_only_fdb_setup) -> None: + """Using GribJump in a build without it must fail at build() and name the cmake flag. + + ``ExtractorType.GribJump`` is bound in every build so that user code does not depend on how + fdb was compiled — which means the failure has to be actionable. This is the inverse of the + ``gribjump`` marker: it runs exactly where integration/gribjump/ is skipped. + """ + builder = SimpleStoreBuilder(read_only_fdb_setup) + builder.add_part( + { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "levtype": "sfc", + "date": "2020-01-01", + "time": 0, + "step": 0, + "param": 167, + }, + [AxisDefinition(["param"], Chunking.SINGLE_VALUE)], + ExtractorType.GribJump(), + ) + with pytest.raises(RuntimeError, match="ENABLE_ZARR_GRIBJUMP_EXTRACTOR"): + builder.build() diff --git a/tests/z3fdb/unit/test_store_v3.py b/tests/z3fdb/interface/store/test_store_v3_simple_builder.py similarity index 92% rename from tests/z3fdb/unit/test_store_v3.py rename to tests/z3fdb/interface/store/test_store_v3_simple_builder.py index c6597666f..daac720d7 100644 --- a/tests/z3fdb/unit/test_store_v3.py +++ b/tests/z3fdb/interface/store/test_store_v3_simple_builder.py @@ -46,9 +46,9 @@ def test_zarr_use_spec_v2(read_only_fdb_setup) -> None: def test_access(read_only_fdb_setup) -> None: builder = SimpleStoreBuilder(read_only_fdb_setup) - builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.GRIB) + builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.Grib()) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data log.debug("shape=%s", data.shape) log.debug("data[:, :]=%s", data[:, :]) @@ -77,10 +77,10 @@ def test_axis_check_out_of_bounds(read_only_fdb_setup_for_sfc_pl_example) -> Non AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -102,7 +102,7 @@ def test_store_list_chunks_complete(read_only_fdb_setup) -> None: - All fields carry identical values: range(0, grid_points) as float32 """ builder = SimpleStoreBuilder(read_only_fdb_setup) - builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.GRIB) + builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.Grib()) store = builder.build() # Expected key counts derived from the fixture's data generation @@ -124,7 +124,7 @@ def test_store_list_chunks_complete(read_only_fdb_setup) -> None: assert len(chunk_keys) == EXPECTED_CHUNKS, f"expected {EXPECTED_CHUNKS} chunk keys, got {len(chunk_keys)}" # Determine grid size from the zarr metadata so we don't hard-code it - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) grid_points = data.shape[-1] expected_bytes = grid_points * 4 # float32 = 4 bytes per value log.debug( diff --git a/tests/z3fdb/unit/test_u_v_vo_d_retrieval.py b/tests/z3fdb/interface/store/test_u_v_vo_d_retrieval_simple_builder.py similarity index 85% rename from tests/z3fdb/unit/test_u_v_vo_d_retrieval.py rename to tests/z3fdb/interface/store/test_u_v_vo_d_retrieval_simple_builder.py index 8f6e47dc8..d41e4a395 100644 --- a/tests/z3fdb/unit/test_u_v_vo_d_retrieval.py +++ b/tests/z3fdb/interface/store/test_u_v_vo_d_retrieval_simple_builder.py @@ -50,7 +50,7 @@ def test_retrieve_u_and_v(read_only_fdb_setup_for_div_vo_example) -> None: AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(RuntimeError, match=".*FDB returned paramId=.*"): @@ -70,7 +70,7 @@ def test_retrieve_u_and_v_with_short_names(read_only_fdb_setup_for_div_vo_exampl AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(RuntimeError, match=".*FDB returned paramId=.*"): @@ -85,7 +85,7 @@ def test_retrieve_vo_and_v(read_only_fdb_setup_for_div_vo_example) -> None: AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(RuntimeError, match=".*FDB returned paramId=.*"): @@ -105,15 +105,15 @@ def test_retrieve_vo_and_d_with_short_names(read_only_fdb_setup_for_div_vo_examp AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) - data = zarr.open_array(builder.build(), mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(builder.build()) assert data.shape[:2] == (16, 4) - assert all(data[0, 0] == 1) # vo, level=50: 0+0+0+1 = 1 - assert all(data[0, 1] == 3) # vo, level=100: 0+0+2+1 = 3 - assert all(data[0, 2] == 0) # d, level=50: 0+0+0+0 = 0 - assert all(data[0, 3] == 2) # d, level=100: 0+0+2+0 = 2 + assert all(data[0, 0] == 1) # vo, level=50: 0+0+0+1 = 1 + assert all(data[0, 1] == 3) # vo, level=100: 0+0+2+1 = 3 + assert all(data[0, 2] == 0) # d, level=50: 0+0+0+0 = 0 + assert all(data[0, 3] == 2) # d, level=100: 0+0+2+0 = 2 assert all(data[8, 0] == 33) # vo, level=50: 32+0+0+1 = 33 assert all(data[8, 2] == 32) # d, level=50: 32+0+0+0 = 32 @@ -137,9 +137,9 @@ def test_retrieve_vo_and_d(read_only_fdb_setup_for_div_vo_example) -> None: AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) - data = zarr.open_array(builder.build(), mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(builder.build()) assert data.shape[:2] == (16, 4) @@ -148,12 +148,12 @@ def test_retrieve_vo_and_d(read_only_fdb_setup_for_div_vo_example) -> None: # level_idx: 0=50, 1=100 # date=2020-01-01 (idx 0), time=0 (idx 0) - assert all(data[0, 0] == 1) # vo, level=50: 0+0+0+1 = 1 - assert all(data[0, 1] == 3) # vo, level=100: 0+0+2+1 = 3 - assert all(data[0, 2] == 0) # d, level=50: 0+0+0+0 = 0 - assert all(data[0, 3] == 2) # d, level=100: 0+0+2+0 = 2 + assert all(data[0, 0] == 1) # vo, level=50: 0+0+0+1 = 1 + assert all(data[0, 1] == 3) # vo, level=100: 0+0+2+1 = 3 + assert all(data[0, 2] == 0) # d, level=50: 0+0+0+0 = 0 + assert all(data[0, 3] == 2) # d, level=100: 0+0+2+0 = 2 # date=2020-01-01, time=300 (idx 1) - assert all(data[1, 0] == 5) # vo, level=50: 0+4+0+1 = 5 + assert all(data[1, 0] == 5) # vo, level=50: 0+4+0+1 = 5 # date=2020-01-02 (idx 1), time=0 assert all(data[8, 0] == 33) # vo, level=50: 32+0+0+1 = 33 assert all(data[8, 2] == 32) # d, level=50: 32+0+0+0 = 32 diff --git a/tests/z3fdb/interface/user_tests/test_dask_access.py b/tests/z3fdb/interface/user_tests/test_dask_access.py new file mode 100644 index 000000000..3e35e229f --- /dev/null +++ b/tests/z3fdb/interface/user_tests/test_dask_access.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +"""User-facing example: accessing a z3fdb store via dask. + +Dask treats a Zarr array as a chunked graph — each Zarr chunk becomes one +dask task. Because z3fdb exposes a fully conformant Zarr v3 store, dask.array +can open it directly and schedule FDB retrievals lazily. + +Fixture data (from conftest.py build_pattern_grib_messages): + dates = [20200101, 20200102, 20200103] (3 values, 0-based index d) + times = [0, 600, 1200, 1800] (4 values, 0-based index t) + params_sfc = [165, 166, 167] (3 values, 0-based index p) + params_pl = [131, 132, 133] (3 values, 0-based index p) + levels = [50, 100, 150] (3 values, 0-based index l) + +Field-value formulas (all indices 0-based): + sfc_value(d, t, p) = d*12 + t*3 + p + pl_value(d, t, p, l) = 36 + d*36 + t*9 + p*3 + l +""" + +import numpy as np +import pytest +import zarr + +da = pytest.importorskip("dask.array", reason="dask is not installed") + +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +pytestmark = [pytest.mark.offline, pytest.mark.zfdb_user_tests] + +COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, +} + + +def test_dask_open_single_part(read_only_fdb_pattern_setup) -> None: + """A single-part SFC view opened as a dask array. + + Verifies that: + - dask.array.from_zarr accepts an FdbZarrStore without materialising data. + - The dask graph chunk shape matches the z3fdb chunk configuration. + - Computing individual chunks triggers FDB retrieval and returns correct values. + + Array layout: + data[d, t, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE), + AxisDefinition(["time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + # Open via zarr first, then wrap with dask — the standard pattern + zarr_array = zarr.open_array(store) + dask_array = da.from_zarr(zarr_array) + + # Lazy: no FDB retrieval yet + assert isinstance(dask_array, da.Array) + assert dask_array.shape[:3] == (3, 4, 3) + assert dask_array.chunks[:3] == ((1, 1, 1), (1, 1, 1, 1), (1, 1, 1)) + + # Compute a single scalar (triggers one FDB retrieve for that chunk) + val = dask_array[0, 0, 0].compute() + assert np.all(val == 0) # sfc_value(0, 0, 0) + + val = dask_array[2, 3, 2].compute() + assert np.all(val == 35) # sfc_value(2, 3, 2) = 24+9+2 + + +def test_dask_open_fixed_size_chunking(read_only_fdb_pattern_setup) -> None: + """FixedSizeChunking on the combined date+time axis opened as a dask array. + + FixedSizeChunk(4) on the ["date","time"] axis groups 4 consecutive + date×time combinations into one Zarr chunk → one dask task. + + Array layout (combined axis 0 index i = d*4 + t): + data[i, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=4)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + zarr_array = zarr.open_array(store) + dask_array = da.from_zarr(zarr_array) + + # 12 date×time values in 3 chunks of 4; 3 param values each in 1 chunk + assert dask_array.shape[:2] == (12, 3) + assert dask_array.chunks[:2] == ((4, 4, 4), (1, 1, 1)) + + # Spot-check: first element of first chunk + assert np.all(dask_array[0, 0].compute() == 0) # sfc_value(0, 0, 0) + # Spot-check: last element of last chunk (chunk boundary) + assert np.all(dask_array[11, 2].compute() == 35) # sfc_value(2, 3, 2) + # Chunk boundary: index 4 is the first element of the second chunk + assert np.all(dask_array[4, 0].compute() == 12) # sfc_value(1, 0, 0) + + +def test_dask_multi_part_sfc_pl(read_only_fdb_pattern_setup) -> None: + """Multi-part SFC + PL view with two-axis dask graph. + + Both parts use SingleValueChunking. Part 1 (SFC) occupies param indices + 0–2; Part 2 (PL) occupies param indices 3–11 on the extension axis. + + Array layout: + Axis 0: combined date×time (12 values, SINGLE_VALUE → 12 chunks of 1) + Axis 1: SFC param / PL param×levelist (3 + 9 = 12 values, SINGLE_VALUE) + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "pl", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [131, 132, 133], + "levelist": [50, 100, 150], + }, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + builder.extend_on_axis(1) + store = builder.build() + + zarr_array = zarr.open_array(store) + dask_array = da.from_zarr(zarr_array) + + assert dask_array.shape[:2] == (12, 12) + + # SFC spot-check: (d=0,t=0) → combined index 0; param=165 → axis-1 index 0 + # sfc_value(0, 0, 0) = 0 + assert np.all(dask_array[0, 0].compute() == 0) + + # PL spot-check: (d=0,t=0) → combined index 0; param=131,levelist=50 → axis-1 index 3 + # pl_value(0, 0, 0, 0) = 36 + assert np.all(dask_array[0, 3].compute() == 36) + + # PL corner: (d=2,t=3) → combined index 11; param=133,levelist=150 → axis-1 index 11 + # pl_value(2, 3, 2, 2) = 36 + 72 + 27 + 6 + 2 = 143 + assert np.all(dask_array[11, 11].compute() == 143) + + +def test_dask_reduction(read_only_fdb_pattern_setup) -> None: + """Dask reduction across the full SFC array. + + Computes the mean over the time axis using dask's lazy graph. The result + is compared against the numpy reference computed from the known value formula. + + sfc_value(d, t, p) = d*12 + t*3 + p + mean over t in [0,1,2,3]: (d*12 + 0 + d*12 + 3 + d*12 + 6 + d*12 + 9 + 4*p) / 4 + = d*12 + 4.5 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE), + AxisDefinition(["time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + zarr_array = zarr.open_array(store) + dask_array = da.from_zarr(zarr_array) + + # Mean over time axis (axis 1), result shape: (3 dates, 3 params, N grid points) + time_mean = dask_array.mean(axis=1).compute() + + # Expected: mean_value(d, p) = d*12 + 4.5 + p + for d in range(3): + for p in range(3): + expected = d * 12 + 4.5 + p + assert np.allclose(time_mean[d, p], expected), ( + f"time mean mismatch at d={d}, p={p}: got {time_mean[d, p, 0]:.2f}, expected {expected:.2f}" + ) diff --git a/tests/z3fdb/interface/user_tests/test_metadata_mapping.py b/tests/z3fdb/interface/user_tests/test_metadata_mapping.py new file mode 100644 index 000000000..94ac41399 --- /dev/null +++ b/tests/z3fdb/interface/user_tests/test_metadata_mapping.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +"""User-facing example: zarr dimension_names from AxisDefinition.name. + +When an ``AxisDefinition`` is given a ``name``, that label appears in the +zarr array metadata as a ``dimension_names`` entry. When no name is given, +one is auto-derived by joining the axis keys with ``"_"``. The implicit +grid-point axis is always named ``"values"``. + +Fixture data (from conftest.py build_pattern_grib_messages): + dates = [20200101, 20200102, 20200103] (3 values) + times = [0, 600, 1200, 1800] (4 values) + params = [165, 166, 167] (3 values, surface) +""" + +import logging +import pytest +import zarr + +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +pytestmark = [pytest.mark.offline, pytest.mark.zfdb_user_tests] + +COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "levtype": "sfc", + "step": 0, + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], +} + + +def test_explicit_axis_names_appear_in_zarr_metadata(read_only_fdb_pattern_setup) -> None: + """Explicit names on AxisDefinition are written into zarr dimension_names. + + With ``name="datetime"`` on the date+time axis and ``name="param"`` on + the param axis, the resulting zarr array metadata must contain + ``["datetime", "param", "values"]``. + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + COMMON, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE, name="datetime"), + AxisDefinition(["param"], Chunking.SINGLE_VALUE, name="param"), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + arr = zarr.open_array(store) + logging.debug(arr.metadata) + assert arr.metadata.dimension_names == ("datetime", "param", "values") + + +def test_auto_derived_axis_names_from_keys(read_only_fdb_pattern_setup) -> None: + """When no name is given, axis names are auto-derived from the MARS keys. + + A compound axis with ``keys=["date", "time"]`` becomes ``"date_time"``; + a single-key axis with ``keys=["param"]`` becomes ``"param"``. + The implicit grid-point axis is always ``"values"``. + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + COMMON, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + arr = zarr.open_array(store) + assert arr.metadata.dimension_names == ("date_time", "param", "values") diff --git a/tests/z3fdb/interface/user_tests/test_xarray_access.py b/tests/z3fdb/interface/user_tests/test_xarray_access.py new file mode 100644 index 000000000..b85dccbf3 --- /dev/null +++ b/tests/z3fdb/interface/user_tests/test_xarray_access.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +"""User-facing example: accessing a z3fdb store via xarray. + +xarray wraps a zarr array as a labeled DataArray. Because the zarr metadata +now carries ``dimension_names``, xarray automatically assigns those names as +dimension labels — enabling ``.isel()``, ``.mean(dim=...)``, and other +label-based operations without any manual wiring. + +Fixture data (from conftest.py build_pattern_grib_messages): + dates = [20200101, 20200102, 20200103] (3 values, 0-based index d) + times = [0, 600, 1200, 1800] (4 values, 0-based index t) + params_sfc = [165, 166, 167] (3 values, 0-based index p) + params_pl = [131, 132, 133] (3 values, 0-based index p) + levels = [50, 100, 150] (3 values, 0-based index l) + +Field-value formulas (all indices 0-based): + sfc_value(d, t, p) = d*12 + t*3 + p + pl_value(d, t, p, l) = 36 + d*36 + t*9 + p*3 + l +""" + +import numpy as np +import pytest +import zarr + +xr = pytest.importorskip("xarray", reason="xarray is not installed") + +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +pytestmark = [pytest.mark.offline, pytest.mark.zfdb_user_tests] + +COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, +} + + +def _open_as_dataarray(store) -> "xr.DataArray": + """Open an FdbZarrStore as a labeled xarray DataArray. + + Opens the underlying zarr array, reads ``dimension_names`` from its + metadata, and wraps it as a DataArray with those dimension labels. + """ + arr = zarr.open_array(store) + dims = list(arr.metadata.dimension_names) + return xr.DataArray(arr, dims=dims) + + +def test_xarray_open_single_part(read_only_fdb_pattern_setup) -> None: + """A single-part SFC view opened as a labeled xarray DataArray. + + Verifies that: + - The DataArray dimensions match the AxisDefinition names + "values". + - Individual elements are accessible via ``.isel()`` with named dims. + - Values match the known formula sfc_value(d, t, p) = d*12 + t*3 + p. + + Array layout: + data[date, time, param] = sfc_value(d, t, p) + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE, name="date"), + AxisDefinition(["time"], Chunking.SINGLE_VALUE, name="time"), + AxisDefinition(["param"], Chunking.SINGLE_VALUE, name="param"), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + da = _open_as_dataarray(store) + + assert da.dims[:3] == ("date", "time", "param") + assert da.sizes["date"] == 3 + assert da.sizes["time"] == 4 + assert da.sizes["param"] == 3 + + # isel by named dimension — no axis indices needed + assert np.all(da.isel(date=0, time=0, param=0).values == 0) # sfc_value(0,0,0) + assert np.all(da.isel(date=2, time=3, param=2).values == 35) # sfc_value(2,3,2) + + +def test_xarray_open_fixed_size_chunking(read_only_fdb_pattern_setup) -> None: + """FixedSizeChunking on the combined date+time axis opened as a DataArray. + + FixedSizeChunk(4) groups 4 consecutive date×time combinations into one + zarr chunk. The xarray DataArray sees a flat axis of size 12. + + Array layout (combined index i = d*4 + t): + data[date_time, param] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=4)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE, name="param"), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + da = _open_as_dataarray(store) + + assert da.dims[:2] == ("date_time", "param") + assert da.sizes["date_time"] == 12 + assert da.sizes["param"] == 3 + + assert np.all(da.isel(date_time=0, param=0).values == 0) # sfc_value(0,0,0) + assert np.all(da.isel(date_time=11, param=2).values == 35) # sfc_value(2,3,2) + assert np.all(da.isel(date_time=4, param=0).values == 12) # sfc_value(1,0,0) + + +def test_xarray_multi_part_sfc_pl(read_only_fdb_pattern_setup) -> None: + """Multi-part SFC + PL view opened as a DataArray. + + Both parts share the same date×time axis (12 values). They extend along + axis 1: SFC params occupy indices 0–2, PL param×level indices 3–11. + + Array layout: + Axis 0 "date_time": 12 combined date×time values + Axis 1 "variable": 3 SFC + 9 PL = 12 combined values + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE, name="date_time"), + AxisDefinition(["param"], Chunking.SINGLE_VALUE, name="variable"), + ], + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "pl", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [131, 132, 133], + "levelist": [50, 100, 150], + }, + [ + AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE, name="date_time"), + AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE, name="variable"), + ], + ExtractorType.Grib(), + ) + builder.extend_on_axis(1) + store = builder.build() + + da = _open_as_dataarray(store) + + assert da.dims[:2] == ("date_time", "variable") + assert da.sizes["date_time"] == 12 + assert da.sizes["variable"] == 12 + + # SFC: (d=0,t=0) → date_time=0; param=165 → variable=0; sfc_value(0,0,0) = 0 + assert np.all(da.isel(date_time=0, variable=0).values == 0) + # PL: (d=0,t=0) → date_time=0; param=131,level=50 → variable=3; pl_value(0,0,0,0) = 36 + assert np.all(da.isel(date_time=0, variable=3).values == 36) + # PL corner: date_time=11; variable=11; pl_value(2,3,2,2) = 36+72+27+6+2 = 143 + assert np.all(da.isel(date_time=11, variable=11).values == 143) + + +def test_xarray_reduction_over_named_dim(read_only_fdb_pattern_setup) -> None: + """xarray reduction along a named dimension. + + Computes the mean over the time axis using the dimension label rather than + an integer axis index. The result is compared against the reference formula. + + sfc_value(d, t, p) = d*12 + t*3 + p + mean over t in [0,1,2,3]: d*12 + (0+3+6+9)/4 + p = d*12 + 4.5 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE, name="date"), + AxisDefinition(["time"], Chunking.SINGLE_VALUE, name="time"), + AxisDefinition(["param"], Chunking.SINGLE_VALUE, name="param"), + ], + ExtractorType.Grib(), + ) + store = builder.build() + + da = _open_as_dataarray(store) + + # Mean over the "time" dimension by name — no axis integer needed + time_mean = da.mean(dim="time") + + assert time_mean.dims[:2] == ("date", "param") + + for d in range(3): + for p in range(3): + expected = d * 12 + 4.5 + p + assert np.allclose(time_mean.isel(date=d, param=p).values, expected), ( + f"time mean mismatch at d={d}, p={p}: " + f"got {float(time_mean.isel(date=d, param=p).values[0]):.2f}, " + f"expected {expected:.2f}" + ) diff --git a/tests/z3fdb/zarr_interface_conformity/__init__.py b/tests/z3fdb/interface/zarr_interface_conformity/__init__.py similarity index 100% rename from tests/z3fdb/zarr_interface_conformity/__init__.py rename to tests/z3fdb/interface/zarr_interface_conformity/__init__.py diff --git a/tests/z3fdb/zarr_interface_conformity/_mocks.py b/tests/z3fdb/interface/zarr_interface_conformity/_mocks.py similarity index 99% rename from tests/z3fdb/zarr_interface_conformity/_mocks.py rename to tests/z3fdb/interface/zarr_interface_conformity/_mocks.py index 731e760fd..1ca25c739 100644 --- a/tests/z3fdb/zarr_interface_conformity/_mocks.py +++ b/tests/z3fdb/interface/zarr_interface_conformity/_mocks.py @@ -34,7 +34,7 @@ def shape(self): return self._shape @override - def chunkShape(self): + def chunk_shape(self): return self._chunk_shape @override diff --git a/tests/z3fdb/zarr_interface_conformity/conftest.py b/tests/z3fdb/interface/zarr_interface_conformity/conftest.py similarity index 100% rename from tests/z3fdb/zarr_interface_conformity/conftest.py rename to tests/z3fdb/interface/zarr_interface_conformity/conftest.py diff --git a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py similarity index 93% rename from tests/z3fdb/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py rename to tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py index 678703aec..29df03068 100644 --- a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py +++ b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_consolidated_metadata.py @@ -14,7 +14,7 @@ from zarr.core.buffer import default_buffer_prototype from zarr.core.sync import sync -from tests.z3fdb.zarr_interface_conformity._mocks import MockChunkedDataView, make_array +from tests.z3fdb.interface.zarr_interface_conformity._mocks import MockChunkedDataView, make_array from z3fdb._internal.zarr import FdbSource, FdbZarrArray, FdbZarrGroup, FdbZarrStore pytestmark = pytest.mark.offline @@ -26,11 +26,7 @@ def _flat_metadata(store): return json.loads(buf.to_bytes())["consolidated_metadata"]["metadata"] -# --------------------------------------------------------------------------- # Simple inline stores -# --------------------------------------------------------------------------- - - def test_consolidated_metadata_flat_group(): """Root group with two array children: flat metadata with keys {a, b}.""" store = FdbZarrStore(FdbZarrGroup(name="", children=[make_array("a"), make_array("b")])) @@ -90,10 +86,7 @@ def test_zmetadata_key_not_present(): assert sync(store.get(".zmetadata", prototype=default_buffer_prototype())) is None -# --------------------------------------------------------------------------- # Deep hierarchy — uses the shared deep_store fixture (see conftest.py) -# --------------------------------------------------------------------------- - _DEEP_ARRAYS = { "arr_root", "grp_a/arr_a1", @@ -147,7 +140,7 @@ def test_consolidated_metadata_deep_group_node(deep_store, path, expected_nested def test_consolidated_metadata_zarr_api(deep_store): """zarr.open_group reads consolidated metadata and navigates nested nodes.""" - root_group = zarr.open_group(deep_store, mode="r", zarr_format=3) + root_group = zarr.open_group(deep_store, mode="r") assert root_group.metadata array = root_group["grp_a/grp_a_inner/arr_ai1"] assert array.metadata diff --git a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_interface.py b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_interface.py similarity index 99% rename from tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_interface.py rename to tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_interface.py index 5fc509c5c..c00973a44 100644 --- a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_interface.py +++ b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_interface.py @@ -51,7 +51,7 @@ def store_kwargs(self, read_only_fdb_setup): @pytest.fixture def store(self, read_only_fdb_setup) -> FdbZarrStore: builder = SimpleStoreBuilder(read_only_fdb_setup) - builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.GRIB) + builder.add_part(_MARS_REQUEST, _AXES, ExtractorType.Grib()) store = builder.build() log.debug("store fixture: %s, known_paths=%d", type(store).__name__, len(store._known_paths)) return store diff --git a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_path_correctness.py b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_path_correctness.py similarity index 94% rename from tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_path_correctness.py rename to tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_path_correctness.py index 3b9ebd562..60fec03ad 100644 --- a/tests/z3fdb/zarr_interface_conformity/test_z3fdb_store_path_correctness.py +++ b/tests/z3fdb/interface/zarr_interface_conformity/test_z3fdb_store_path_correctness.py @@ -43,7 +43,7 @@ from zarr.core.buffer import default_buffer_prototype from zarr.core.sync import _collect_aiterator, sync -from tests.z3fdb.zarr_interface_conformity._mocks import ( +from tests.z3fdb.interface.zarr_interface_conformity._mocks import ( ARR_A1, ARR_AI1, ARR_AI2, @@ -64,27 +64,31 @@ def test_total_key_count(deep_store: FdbZarrStore) -> None: - """5 arrays + 6 groups verified against expected path breakdown.""" + """_known_paths holds metadata only: one zarr.json per node (5 arrays + 6 groups = 11). + + Chunk keys are computed on demand and not stored in _known_paths: + zarr.json root group + arr_root/zarr.json 1-D 4 chunks + grp_a/zarr.json + grp_a/arr_a1/zarr.json 2-D 6 chunks + grp_a/grp_a_inner/zarr.json + grp_a/grp_a_inner/arr_ai1/zarr.json 3-D 4 chunks + grp_a/grp_a_inner/arr_ai2/zarr.json 1-D 2 chunks + grp_b/zarr.json + grp_b/grp_bb/zarr.json + grp_b/grp_bb/grp_bbd/zarr.json + grp_b/grp_bb/grp_bbd/arr_bbd/zarr.json 2-D 4 chunks + """ total = len(deep_store._known_paths) - log.debug("total paths: %d", total) - # 1 root zarr.json - # 5 arr_root (1 + 4 chunks) - # 1 grp_a/zarr.json - # 7 arr_a1 (1 + 6 chunks) - # 1 grp_a_inner/zarr.json - # 5 arr_ai1 (1 + 4 chunks) - # 3 arr_ai2 (1 + 2 chunks) - # 1 grp_b/zarr.json - # 1 grp_bb/zarr.json - # 1 grp_bbd/zarr.json - # 5 arr_bbd (1 + 4 chunks) - assert total == 31 + log.debug("total metadata paths: %d", total) + assert total == 11 def test_list_returns_all_known_paths(deep_store: FdbZarrStore) -> None: + """list() returns metadata + chunk keys (11 + 20 = 31); _known_paths is a subset.""" keys = sync(_collect_aiterator(deep_store.list())) log.debug("list() returned %d keys", len(keys)) - assert set(keys) == set(deep_store._known_paths) + assert set(deep_store._known_paths).issubset(set(keys)) assert len(keys) == 31 @@ -302,7 +306,7 @@ def test_exists_unknown_path(deep_store: FdbZarrStore) -> None: def test_list_prefix_empty_returns_all(deep_store: FdbZarrStore) -> None: keys = sync(_collect_aiterator(deep_store.list_prefix(""))) log.debug("list_prefix('') -> %d keys", len(keys)) - assert set(keys) == set(deep_store._known_paths) + assert set(keys) == set(deep_store) @pytest.mark.parametrize( diff --git a/tests/z3fdb/permutation_tests/test_axis_definition_ordering.py b/tests/z3fdb/permutation_tests/test_axis_definition_ordering.py index 590d3e997..11e2da1d5 100644 --- a/tests/z3fdb/permutation_tests/test_axis_definition_ordering.py +++ b/tests/z3fdb/permutation_tests/test_axis_definition_ordering.py @@ -35,7 +35,7 @@ def _open_array(store): - return zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + return zarr.open_array(store) def _canonical_value(date, time, param, step=0): @@ -53,7 +53,7 @@ def test_canonical_axis_order_all_chunked(read_only_fdb_pattern_setup) -> None: AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -75,7 +75,7 @@ def test_canonical_axis_order_date_time_non_chunked(read_only_fdb_pattern_setup) AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -97,7 +97,7 @@ def test_swapped_time_date_axes_non_chunked(read_only_fdb_pattern_setup) -> None AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -118,7 +118,7 @@ def test_merged_date_time_axis_non_chunked(read_only_fdb_pattern_setup) -> None: AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -140,7 +140,7 @@ def test_merged_time_date_axis_non_chunked_switched(read_only_fdb_pattern_setup) AxisDefinition(["time", "date"], Chunking.WHOLE_AXIS), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data diff --git a/tests/z3fdb/permutation_tests/test_axis_definition_permutations.py b/tests/z3fdb/permutation_tests/test_axis_definition_permutations.py index bbb5f2d9d..3b5bbc4e5 100644 --- a/tests/z3fdb/permutation_tests/test_axis_definition_permutations.py +++ b/tests/z3fdb/permutation_tests/test_axis_definition_permutations.py @@ -39,14 +39,14 @@ def _open_array(store): - return zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + return zarr.open_array(store) def _canonical_value(date, time, param, step=0): return step + param * 1 + time * 3 * 1 + date * 4 * 3 * 1 -@pytest.mark.parametrize("index_permutation", permutations([0, 1, 2, 3])) +@pytest.mark.parametrize("index_permutation", list(permutations([0, 1, 2, 3]))) def test_all_four_axis_permutations_chunked(read_only_fdb_pattern_setup, index_permutation) -> None: """All 24 permutations of four individually-chunked axes produce correct data.""" axes = [ @@ -61,7 +61,7 @@ def test_all_four_axis_permutations_chunked(read_only_fdb_pattern_setup, index_p builder.add_part( CANONICAL_REQUEST, [axes[i] for i in index_permutation], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -104,7 +104,7 @@ def test_three_axis_permutations_with_merged_date_time( builder.add_part( CANONICAL_REQUEST, [axes[i] for i in index_permutation], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data diff --git a/tests/z3fdb/permutation_tests/test_builder_extension.py b/tests/z3fdb/permutation_tests/test_builder_extension.py new file mode 100644 index 000000000..729aefa77 --- /dev/null +++ b/tests/z3fdb/permutation_tests/test_builder_extension.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) +# SPDX-License-Identifier: Apache-2.0 + +"""Tests that verify extension works for SingleValueChunking and FixedSizeChunking +when parts share the same axis structure but span non-overlapping value ranges. + +These tests complement the existing test_store_multiple_parts.py which only tests +extension along axes that differ in key-structure between parts (e.g. SFC ["param"] +vs PL ["param","levelist"]). Here we split the *same* MARS key's value range across +multiple parts — e.g. dates=[2020-01-01,2020-01-02] in Part 1 and date=[2020-01-03] +in Part 2. + +Fixture data (from conftest.py build_pattern_grib_messages): + dates = [20200101, 20200102, 20200103] (3 values, 0-based index d) + times = [0, 600, 1200, 1800] (4 values, 0-based index t) + params_sfc = [165, 166, 167] (3 values, 0-based index p) + params_pl = [131, 132, 133] (3 values, 0-based index p) + levels = [50, 100, 150] (3 values, 0-based index l) + +Field-value formulas (all indices 0-based): + sfc_value(d, t, p) = d*12 + t*3 + p + pl_value(d, t, p, l) = 36 + d*36 + t*9 + p*3 + l +""" + +import numpy as np +import pytest +import zarr + +from z3fdb import ( + AxisDefinition, + Chunking, + ExtractorType, + SimpleStoreBuilder, +) + +pytestmark = pytest.mark.offline + +COMMON = { + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, +} + + +# SingleValueChunking: two date batches, extend on date axis +def test_single_value_chunking_date_extension( + read_only_fdb_pattern_setup, +) -> None: + """Extension along the date axis (axis 0) using SingleValueChunking. + + Both parts share the same axis structure: + Axis 0: ["date"] — SingleValueChunking (chunk size 1) + Axis 1: ["time"] — SingleValueChunking (chunk size 1) + Axis 2: ["param"] — SingleValueChunking (chunk size 1) + + Part 1 covers dates 2020-01-01 and 2020-01-02 (d=0,1 → 2 chunks on axis 0). + Part 2 covers date 2020-01-03 (d=2 → 1 chunk on axis 0). + + After extend_on_axis(0): 3 date chunks × 4 time chunks × 3 param chunks. + + Array layout: + data[d, t, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + _axes = [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE), + AxisDefinition(["time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + _axes, + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + _axes, + ExtractorType.Grib(), + ) + builder.extend_on_axis(0) + store = builder.build() + + data = zarr.open_array(store) + + assert data.shape[:3] == (3, 4, 3) + assert data.chunks[:3] == (1, 1, 1) + + # Spot-checks within Part 1 (d=0,1) + assert np.all(data[0, 0, 0] == 0) # sfc_value(0, 0, 0) + assert np.all(data[0, 0, 1] == 1) # sfc_value(0, 0, 1) + assert np.all(data[0, 1, 0] == 3) # sfc_value(0, 1, 0) + assert np.all(data[1, 0, 0] == 12) # sfc_value(1, 0, 0) + assert np.all(data[1, 3, 2] == 23) # sfc_value(1, 3, 2) + + # Cross-part boundary: d=2 (Part 2) + assert np.all(data[2, 0, 0] == 24) # sfc_value(2, 0, 0) = 2*12 + 0 + 0 + assert np.all(data[2, 0, 1] == 25) # sfc_value(2, 0, 1) + assert np.all(data[2, 3, 2] == 35) # sfc_value(2, 3, 2) = 24 + 9 + 2 + + +# FixedSizeChunking: two date batches, combined date+time axis +def test_fixed_size_chunking_date_extension( + read_only_fdb_pattern_setup, +) -> None: + """Extension along the combined date+time axis (axis 0) using FixedSizeChunk(4). + + Both parts share the same axis structure: + Axis 0: ["date","time"] — FixedSizeChunk(chunk_shape=4) + Axis 1: ["param"] — SingleValueChunking + + Part 1: dates=[2020-01-01, 2020-01-02], 4 times → combined size = 8 → 2 chunks of 4 + Part 2: date= [2020-01-03], 4 times → combined size = 4 → 1 chunk of 4 + + chunkSizeCheck for Part 2 axis ["date","time"] with cardinalities [1,4]: + k=0: trailing=1, card=4, d=4, 4%4==0 → valid ✓ + + After extend_on_axis(0): 12 combined-index values, 3 chunks of 4. + + Array layout (combined axis 0 index = d*4 + t): + data[d*4 + t, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + + axes = [ + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=4)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + axes, + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-03"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + axes, + ExtractorType.Grib(), + ) + builder.extend_on_axis(0) + store = builder.build() + + data = zarr.open_array(store) + + assert data.shape[:2] == (12, 3) + assert data.chunks[:2] == (4, 1) + + # Combined index i = d*4 + t; value = sfc_value(d, t, p) = d*12 + t*3 + p + # Part 1 occupies combined indices [0, 7] + assert np.all(data[0, 0] == 0) # (d=0,t=0), p=0 → 0 + assert np.all(data[0, 2] == 2) # (d=0,t=0), p=2 → 2 + assert np.all(data[1, 0] == 3) # (d=0,t=1), p=0 → 3 + assert np.all(data[3, 0] == 9) # (d=0,t=3), p=0 → 9 (last in chunk 0) + assert np.all(data[4, 0] == 12) # (d=1,t=0), p=0 → 12 (first in chunk 1) + assert np.all(data[7, 2] == 23) # (d=1,t=3), p=2 → 12+9+2 = 23 + + # Cross-part boundary: Part 2 starts at combined index 8 (d=2, t=0) + assert np.all(data[8, 0] == 24) # sfc_value(2, 0, 0) = 24 + assert np.all(data[8, 2] == 26) # sfc_value(2, 0, 2) = 26 + assert np.all(data[11, 2] == 35) # sfc_value(2, 3, 2) = 24+9+2 = 35 + + +# FixedSizeChunking: two time halves, extend on time axis +def test_fixed_size_chunking_time_extension( + read_only_fdb_pattern_setup, +) -> None: + """Extension along the time axis (axis 1) using FixedSizeChunk(2). + + Both parts share all 3 dates and all 3 SFC params, but cover different + halves of the 4-time sequence: + Part 1: times=[0, 600] (indices t=0,1) → 1 chunk of 2 + Part 2: times=[1200,1800] (indices t=2,3) → 1 chunk of 2 + + Axes on both parts: + Axis 0: ["date"] — SingleValueChunking + Axis 1: ["time"] — FixedSizeChunk(chunk_shape=2) ← extension axis + Axis 2: ["param"] — SingleValueChunking + + After extend_on_axis(1): 4 time values, 2 chunks of 2. + + Array layout: + data[d, t, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + + axes = [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE), + AxisDefinition(["time"], Chunking.FixedSizeChunk(chunk_shape=2)), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [0, 600], + "param": [165, 166, 167], + }, + axes, + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02", "2020-01-03"], + "time": [1200, 1800], + "param": [165, 166, 167], + }, + axes, + ExtractorType.Grib(), + ) + builder.extend_on_axis(1) + store = builder.build() + + data = zarr.open_array(store) + + assert data.shape[:3] == (3, 4, 3) + assert data.chunks[:3] == (1, 2, 1) + + # Part 1 — time indices t=0,1 (times 0, 600) + assert np.all(data[0, 0, 0] == 0) # sfc_value(0, 0, 0) + assert np.all(data[0, 1, 0] == 3) # sfc_value(0, 1, 0) + assert np.all(data[1, 0, 2] == 14) # sfc_value(1, 0, 2) = 12+0+2 + + # Cross-part boundary — time indices t=2,3 (times 1200, 1800), Part 2 + assert np.all(data[0, 2, 0] == 6) # sfc_value(0, 2, 0) + assert np.all(data[0, 3, 2] == 11) # sfc_value(0, 3, 2) = 0+9+2 + assert np.all(data[2, 3, 2] == 35) # sfc_value(2, 3, 2) = 24+9+2 + + +# Three-part SingleValueChunking, one date per part +def test_three_part_single_value_extension( + read_only_fdb_pattern_setup, +) -> None: + """Three parts, one per date, SingleValueChunking on all axes. + + Demonstrates that extend_on_axis works for more than two same-structure parts. + Each part has a single date: + Part 1: date=[2020-01-01] → d=0 + Part 2: date=[2020-01-02] → d=1 + Part 3: date=[2020-01-03] → d=2 + + After extend_on_axis(0): 3 date × 4 time × 3 param = (3, 4, 3) shape. + + Array layout: + data[d, t, p] = sfc_value(d, t, p) = d*12 + t*3 + p + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + + axes = [ + AxisDefinition(["date"], Chunking.SINGLE_VALUE), + AxisDefinition(["time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + for date_str in ["2020-01-01", "2020-01-02", "2020-01-03"]: + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": [date_str], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + axes, + ExtractorType.Grib(), + ) + builder.extend_on_axis(0) + store = builder.build() + + data = zarr.open_array(store) + + assert data.shape[:3] == (3, 4, 3) + assert data.chunks[:3] == (1, 1, 1) + + # First element of each date (d=0,1,2; t=0; p=0) + assert np.all(data[0, 0, 0] == 0) # sfc_value(0, 0, 0) + assert np.all(data[1, 0, 0] == 12) # sfc_value(1, 0, 0) + assert np.all(data[2, 0, 0] == 24) # sfc_value(2, 0, 0) + + # Last element of each date (t=3; p=2) + assert np.all(data[0, 3, 2] == 11) # sfc_value(0, 3, 2) = 0+9+2 + assert np.all(data[1, 3, 2] == 23) # sfc_value(1, 3, 2) = 12+9+2 + assert np.all(data[2, 3, 2] == 35) # sfc_value(2, 3, 2) = 24+9+2 + + # Mid-axis spot-check + assert np.all(data[1, 2, 1] == 19) # sfc_value(1, 2, 1) = 12+6+1 + assert np.all(data[1, 2, 2] == 20) # sfc_value(1, 2, 2) = 12+6+2 + + +# FixedSizeChunking rejects a part whose axis size isn't a multiple +def test_fixed_size_chunking_misaligned_part_boundary_rejected( + read_only_fdb_pattern_setup, +) -> None: + """Documents the current limitation: each part's axis size must independently + pass AxisMapper::chunkSizeCheck. + + Part 2 has a single date with FixedSizeChunk(2) on the date axis. + chunkSizeCheck(axis=[1 date], chunkSize=2): + trailingProduct=1, card=1, d=2, 1%2 != 0 → rejected. + + The build() call must raise RuntimeError mentioning "AxisMapper::mapAxisToChunks". + """ + builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) + + def _axes_with_fsc2(): + return [ + AxisDefinition(["date"], Chunking.FixedSizeChunk(chunk_shape=2)), + AxisDefinition(["time"], Chunking.SINGLE_VALUE), + AxisDefinition(["param"], Chunking.SINGLE_VALUE), + ] + + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-01", "2020-01-02"], + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + _axes_with_fsc2(), + ExtractorType.Grib(), + ) + builder.add_part( + { + **COMMON, + "levtype": "sfc", + "date": ["2020-01-03"], # 1 date, but chunk size = 2 -> invalid + "time": [0, 600, 1200, 1800], + "param": [165, 166, 167], + }, + _axes_with_fsc2(), + ExtractorType.Grib(), + ) + builder.extend_on_axis(0) + + with pytest.raises(RuntimeError, match="AxisMapper::mapAxisToChunks"): + builder.build() diff --git a/tests/z3fdb/permutation_tests/test_scrambled_request_ordering.py b/tests/z3fdb/permutation_tests/test_scrambled_request_ordering.py index 920c8b208..7d191c149 100644 --- a/tests/z3fdb/permutation_tests/test_scrambled_request_ordering.py +++ b/tests/z3fdb/permutation_tests/test_scrambled_request_ordering.py @@ -29,7 +29,7 @@ def _open_array(store): - return zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + return zarr.open_array(store) def _scrambled_value_at(param, time, date): @@ -61,7 +61,7 @@ def test_scrambled_request_canonical_axis_order(read_only_fdb_pattern_setup) -> AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["step"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -94,7 +94,7 @@ def test_scrambled_request_swapped_axis_definitions(read_only_fdb_pattern_setup) AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["date"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data @@ -122,7 +122,7 @@ def test_scrambled_request_all_axes_merged_non_chunked(read_only_fdb_pattern_set "param": [167, 165, 166], }, [AxisDefinition(["time", "step", "param", "date"], Chunking.WHOLE_AXIS)], - ExtractorType.GRIB, + ExtractorType.Grib(), ) data = _open_array(builder.build()) assert data diff --git a/tests/z3fdb/permutation_tests/test_store_fixed_size_chunking.py b/tests/z3fdb/permutation_tests/test_store_fixed_size_chunking.py index 0421318e6..1399fb0d6 100644 --- a/tests/z3fdb/permutation_tests/test_store_fixed_size_chunking.py +++ b/tests/z3fdb/permutation_tests/test_store_fixed_size_chunking.py @@ -26,17 +26,27 @@ # pl_value(d, t, p, l) = 36 + d*36 + t*9 + p*3 + l SFC_REQUEST = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", "stream": "oper", + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", "date": ["2020-01-01", "2020-01-02", "2020-01-03"], - "levtype": "sfc", "step": 0, + "levtype": "sfc", + "step": 0, "param": [165, 166, 167], "time": [0, 600, 1200, 1800], } PL_REQUEST = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", "stream": "oper", + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", "date": ["2020-01-01", "2020-01-02", "2020-01-03"], - "levtype": "pl", "step": 0, + "levtype": "pl", + "step": 0, "param": [131, 132, 133], "levelist": [50, 100, 150], "time": [0, 600, 1200, 1800], @@ -70,23 +80,23 @@ def test_individual_chunking_combined_datetime_axis( builder.add_part( SFC_REQUEST, [ - AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunkShape=4)), - AxisDefinition(["param"], Chunking.FixedSizeChunk(chunkShape=3)), + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=4)), + AxisDefinition(["param"], Chunking.FixedSizeChunk(chunk_shape=3)), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, [ - AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunkShape=4)), - AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunkShape=3)), + AxisDefinition(["date", "time"], Chunking.FixedSizeChunk(chunk_shape=4)), + AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunk_shape=3)), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) # Shape: combined date*time = 12; SFC param (3) + PL param*levelist (9) = 12 assert data.shape[:2] == (12, 12) @@ -152,24 +162,24 @@ def test_individual_chunking_separate_time_axis( SFC_REQUEST, [ AxisDefinition(["date"], Chunking.SINGLE_VALUE), - AxisDefinition(["time"], Chunking.FixedSizeChunk(chunkShape=2)), - AxisDefinition(["param"], Chunking.FixedSizeChunk(chunkShape=3)), + AxisDefinition(["time"], Chunking.FixedSizeChunk(chunk_shape=2)), + AxisDefinition(["param"], Chunking.FixedSizeChunk(chunk_shape=3)), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, [ AxisDefinition(["date"], Chunking.SINGLE_VALUE), - AxisDefinition(["time"], Chunking.FixedSizeChunk(chunkShape=2)), - AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunkShape=3)), + AxisDefinition(["time"], Chunking.FixedSizeChunk(chunk_shape=2)), + AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunk_shape=3)), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(2) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) # Shape: date=3, time=4; SFC param (3) + PL param*levelist (9) = 12 assert data.shape[:3] == (3, 4, 12) @@ -235,25 +245,25 @@ def test_individual_chunking_reordered_axes( builder.add_part( SFC_REQUEST, [ - AxisDefinition(["param"], Chunking.FixedSizeChunk(chunkShape=3)), - AxisDefinition(["time"], Chunking.FixedSizeChunk(chunkShape=2)), + AxisDefinition(["param"], Chunking.FixedSizeChunk(chunk_shape=3)), + AxisDefinition(["time"], Chunking.FixedSizeChunk(chunk_shape=2)), AxisDefinition(["date"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, [ - AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunkShape=3)), - AxisDefinition(["time"], Chunking.FixedSizeChunk(chunkShape=2)), + AxisDefinition(["param", "levelist"], Chunking.FixedSizeChunk(chunk_shape=3)), + AxisDefinition(["time"], Chunking.FixedSizeChunk(chunk_shape=2)), AxisDefinition(["date"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(0) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) # Shape: SFC param (3) + PL param*levelist (9) = 12; time=4; date=3 assert data.shape[:3] == (12, 4, 3) @@ -316,9 +326,7 @@ def test_individual_chunking_reordered_axes( @pytest.mark.parametrize("chunk_size", _FOUR_KEY_VALID_CHUNK_SIZES) -def test_individual_chunking_four_key_single_axis( - read_only_fdb_pattern_setup, chunk_size: int -) -> None: +def test_individual_chunking_four_key_single_axis(read_only_fdb_pattern_setup, chunk_size: int) -> None: """Mixed-levtype view: Part 1 is PL (four-key axis), Part 2 is SFC (three-key axis). Part 1 (PL): date=[2020-01-01], 1 x 4 x 3 x 3 = 36 values @@ -352,8 +360,12 @@ def test_individual_chunking_four_key_single_axis( assert PART2_SIZE % chunk_size == 0 COMMON = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "step": 0, + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, } builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) @@ -369,10 +381,10 @@ def test_individual_chunking_four_key_single_axis( [ AxisDefinition( ["date", "time", "param", "levelist"], - Chunking.FixedSizeChunk(chunkShape=chunk_size), + Chunking.FixedSizeChunk(chunk_shape=chunk_size), ) ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( { @@ -385,32 +397,32 @@ def test_individual_chunking_four_key_single_axis( [ AxisDefinition( ["date", "time", "param"], - Chunking.FixedSizeChunk(chunkShape=chunk_size), + Chunking.FixedSizeChunk(chunk_shape=chunk_size), ) ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(0) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data.shape[0] == PART1_SIZE + PART2_SIZE # = 54 assert data.chunks[0] == chunk_size # Part 1 (i in [0, 36)): i = t*9 + p*3 + l, data[i] = pl_value(0,t,p,l) = 36 + i - assert np.all(data[0] == 36) # t=0, p=131, l=50 -> pl_value(0,0,0,0) = 36 - assert np.all(data[1] == 37) # t=0, p=131, l=100 - assert np.all(data[3] == 39) # t=0, p=132, l=50 - assert np.all(data[9] == 45) # t=600, p=131, l=50 -> pl_value(0,1,0,0) = 45 + assert np.all(data[0] == 36) # t=0, p=131, l=50 -> pl_value(0,0,0,0) = 36 + assert np.all(data[1] == 37) # t=0, p=131, l=100 + assert np.all(data[3] == 39) # t=0, p=132, l=50 + assert np.all(data[9] == 45) # t=600, p=131, l=50 -> pl_value(0,1,0,0) = 45 assert np.all(data[35] == 71) # t=1800, p=133, l=150 -> pl_value(0,3,2,2) = 71 # Part 1 -> Part 2 boundary at index 36 (chunk-aligned: PART1_SIZE % chunk_size == 0) assert PART1_SIZE % chunk_size == 0 # Part 2 (i in [36, 54)): i = 36 + d*9 + t*3 + p, data[i] = sfc_value(d,t,p) = d*12 + t*3 + p - assert np.all(data[36] == 0) # d=0, t=0, p=165 -> sfc_value(0,0,0) = 0 - assert np.all(data[37] == 1) # d=0, t=0, p=166 - assert np.all(data[44] == 8) # d=0, t=1200, p=167 -> sfc_value(0,2,2) = 8 + assert np.all(data[36] == 0) # d=0, t=0, p=165 -> sfc_value(0,0,0) = 0 + assert np.all(data[37] == 1) # d=0, t=0, p=166 + assert np.all(data[44] == 8) # d=0, t=1200, p=167 -> sfc_value(0,2,2) = 8 assert np.all(data[45] == 12) # d=1, t=0, p=165 -> sfc_value(1,0,0) = 12 assert np.all(data[53] == 20) # d=1, t=1200, p=167 -> sfc_value(1,2,2) = 20 @@ -422,9 +434,7 @@ def test_individual_chunking_four_key_single_axis( @pytest.mark.parametrize("invalid_chunk_size", _FOUR_KEY_INVALID_CHUNK_SIZES) -def test_individual_chunking_rejects_invalid_chunk_size( - read_only_fdb_pattern_setup, invalid_chunk_size: int -) -> None: +def test_individual_chunking_rejects_invalid_chunk_size(read_only_fdb_pattern_setup, invalid_chunk_size: int) -> None: """build() raises when the chunk size violates the key-hierarchy alignment rule. Uses the PL axis ["date","time","param","levelist"] with cardinalities [1,4,3,3]. @@ -433,8 +443,12 @@ def test_individual_chunking_rejects_invalid_chunk_size( re-raises it as a RuntimeError whose message contains "AxisMapper::mapAxisToChunks". """ COMMON = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "step": 0, + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "step": 0, } builder = SimpleStoreBuilder(read_only_fdb_pattern_setup) builder.add_part( @@ -449,10 +463,10 @@ def test_individual_chunking_rejects_invalid_chunk_size( [ AxisDefinition( ["date", "time", "param", "levelist"], - Chunking.FixedSizeChunk(chunkShape=invalid_chunk_size), + Chunking.FixedSizeChunk(chunk_shape=invalid_chunk_size), ) ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) with pytest.raises(RuntimeError, match="AxisMapper::mapAxisToChunks"): builder.build() diff --git a/tests/z3fdb/permutation_tests/test_store_missing_values.py b/tests/z3fdb/permutation_tests/test_store_missing_values.py index 35f926157..ecc46cefc 100644 --- a/tests/z3fdb/permutation_tests/test_store_missing_values.py +++ b/tests/z3fdb/permutation_tests/test_store_missing_values.py @@ -41,9 +41,7 @@ def swh_grib_download(tmp_path_factory): - missing_mask: boolean array, True where the GRIB bitmap marks the point missing - date: the date string used in the MARS request """ - opendata = pytest.importorskip( - "ecmwf.opendata", reason="pip install ecmwf-opendata to run online tests" - ) + opendata = pytest.importorskip("ecmwf.opendata", reason="pip install ecmwf-opendata to run online tests") date = datetime.date.today() - datetime.timedelta(days=2) tmp = tmp_path_factory.mktemp("swh_online") @@ -113,7 +111,7 @@ def test_bitmap_missing_points_become_fill_value(fdb_with_swh_bitmap): "param": 140229, }, [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.fill_missing_value(-20.0) view = builder.build() @@ -144,12 +142,12 @@ def test_fill_value_propagated_to_zarr_metadata(fdb_with_swh_bitmap): "param": 140229, }, [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.fill_missing_value(-20.0) store = builder.build() - arr = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + arr = zarr.open_array(store) assert arr.fill_value == -20.0 @@ -172,7 +170,7 @@ def test_default_fill_value_replaces_bitmap_sentinel(fdb_with_swh_bitmap): "param": 140229, }, [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, + ExtractorType.Grib(), ) view = builder.build() diff --git a/tests/z3fdb/permutation_tests/test_store_multiple_parts.py b/tests/z3fdb/permutation_tests/test_store_multiple_parts.py index 4bd17dd93..1a483170d 100644 --- a/tests/z3fdb/permutation_tests/test_store_multiple_parts.py +++ b/tests/z3fdb/permutation_tests/test_store_multiple_parts.py @@ -23,34 +23,47 @@ def test_axis_check_merge(read_only_fdb_setup_for_sfc_pl_example) -> None: builder = SimpleStoreBuilder(read_only_fdb_setup_for_sfc_pl_example) builder.add_part( { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "sfc", "step": 0, "param": [165, 166], + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "sfc", + "step": 0, + "param": [165, 166], "time": "0/to/21/by/3", }, [ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "pl", "step": 0, "param": [131, 132], - "levelist": [50, 100], "time": "0/to/21/by/3", + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "pl", + "step": 0, + "param": [131, 132], + "levelist": [50, 100], + "time": "0/to/21/by/3", }, [ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -66,34 +79,47 @@ def test_axis_check_merge_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> builder = SimpleStoreBuilder(read_only_fdb_setup_for_sfc_pl_example) builder.add_part( { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "sfc", "step": 0, "param": [165, 166], + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "sfc", + "step": 0, + "param": [165, 166], "time": "0/to/21/by/3", }, [ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.WHOLE_AXIS), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "pl", "step": 0, "param": [131, 132], - "levelist": [50, 100], "time": "0/to/21/by/3", + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "pl", + "step": 0, + "param": [131, 132], + "levelist": [50, 100], + "time": "0/to/21/by/3", }, [ AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.WHOLE_AXIS), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -106,17 +132,30 @@ def test_axis_check_merge_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> SFC_REQUEST = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "sfc", "step": 0, "param": [165, 166], + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "sfc", + "step": 0, + "param": [165, 166], "time": "0/to/21/by/3", } PL_REQUEST = { - "type": "an", "class": "ea", "domain": "g", "expver": "0001", - "stream": "oper", "date": ["2020-01-01", "2020-01-02"], - "levtype": "pl", "step": 0, "param": [131, 132], - "levelist": [50, 100], "time": "0/to/21/by/3", + "type": "an", + "class": "ea", + "domain": "g", + "expver": "0001", + "stream": "oper", + "date": ["2020-01-01", "2020-01-02"], + "levtype": "pl", + "step": 0, + "param": [131, 132], + "levelist": [50, 100], + "time": "0/to/21/by/3", } @@ -129,7 +168,7 @@ def test_extend_on_axis_0(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["param"], Chunking.SINGLE_VALUE), AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, @@ -137,12 +176,12 @@ def test_extend_on_axis_0(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(0) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -164,7 +203,7 @@ def test_extend_on_axis_0_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> AxisDefinition(["param"], Chunking.WHOLE_AXIS), AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, @@ -172,12 +211,12 @@ def test_extend_on_axis_0_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> AxisDefinition(["param", "levelist"], Chunking.WHOLE_AXIS), AxisDefinition(["date", "time"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(0) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -199,7 +238,7 @@ def test_non_extension_axis_no_chunking(read_only_fdb_setup_for_sfc_pl_example) AxisDefinition(["date", "time"], Chunking.WHOLE_AXIS), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, @@ -207,12 +246,12 @@ def test_non_extension_axis_no_chunking(read_only_fdb_setup_for_sfc_pl_example) AxisDefinition(["date", "time"], Chunking.WHOLE_AXIS), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -236,7 +275,7 @@ def test_single_key_axes(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["time"], Chunking.SINGLE_VALUE), AxisDefinition(["param"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, @@ -245,12 +284,12 @@ def test_single_key_axes(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["time"], Chunking.SINGLE_VALUE), AxisDefinition(["param", "levelist"], Chunking.SINGLE_VALUE), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(2) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data @@ -274,7 +313,7 @@ def test_all_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["date", "time"], Chunking.WHOLE_AXIS), AxisDefinition(["param"], Chunking.WHOLE_AXIS), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.add_part( PL_REQUEST, @@ -282,12 +321,12 @@ def test_all_no_chunking(read_only_fdb_setup_for_sfc_pl_example) -> None: AxisDefinition(["date", "time"], Chunking.WHOLE_AXIS), AxisDefinition(["param", "levelist"], Chunking.WHOLE_AXIS), ], - ExtractorType.GRIB, + ExtractorType.Grib(), ) builder.extend_on_axis(1) store = builder.build() - data = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) + data = zarr.open_array(store) assert data diff --git a/tests/z3fdb/test_store_v3_online.py b/tests/z3fdb/test_store_v3_online.py deleted file mode 100644 index 100fb830d..000000000 --- a/tests/z3fdb/test_store_v3_online.py +++ /dev/null @@ -1,186 +0,0 @@ -# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF) -# SPDX-License-Identifier: Apache-2.0 - -"""Online z3fdb tests. - -These tests download real ECMWF open data to verify that bitmap-masked grid points -(e.g. land points in a significant-wave-height field) are returned as the configured -fill value rather than the raw eccodes missingValue sentinel. - -Requires: - pip install ecmwf-opendata -""" - -import datetime -import logging - -import eccodes as ec -import numpy as np -import pytest -import zarr - -from pychunked_data_view import ( - AxisDefinition, - ChunkedDataViewBuilder, - Chunking, - ExtractorType, -) -from pyfdb import FDB -from z3fdb import SimpleStoreBuilder - -pytestmark = pytest.mark.online - - -@pytest.fixture(scope="module") -def swh_grib_download(tmp_path_factory): - """Download a real SWH GRIB field from ECMWF open data (runs once per module). - - Returns (grib_path, expected_float32_values, missing_mask, date) where: - - grib_path: path to the downloaded GRIB file - - expected_float32_values: raw float32 values as eccodes returns them - - missing_mask: boolean array, True where the GRIB bitmap marks the point missing - - date: the date string used in the MARS request - """ - opendata = pytest.importorskip( - "ecmwf.opendata", reason="pip install ecmwf-opendata to run online tests" - ) - - date = datetime.date.today() - datetime.timedelta(days=2) - tmp = tmp_path_factory.mktemp("swh_online") - grib = tmp / "swh_opendata.grib2" - - try: - opendata.Client(source="ecmwf", model="ifs", resol="0p25").retrieve( - date=date.isoformat(), - time=0, - type="fc", - stream="wave", - step=0, - param="swh", - target=str(grib), - ) - except Exception as e: - if getattr(getattr(e, "response", None), "status_code", None) == 429: - pytest.skip("ECMWF Open Data rate limit exceeded (HTTP 429)") - raise - - with open(grib, "rb") as f: - gid = ec.codes_grib_new_from_file(f) - expected = ec.codes_get_values(gid).astype(np.float32) - missing_sentinel = float(ec.codes_get(gid, "missingValue")) - missing = expected == np.float32(missing_sentinel) - ec.codes_release(gid) - - logging.debug(f"#Values: {len(expected)}") - logging.debug(f"#Missing: {np.sum(missing)}") - logging.debug(f"#Missing Entries/#Entries: {np.sum(missing)} / {len(expected)}") - - assert 0 < missing.sum() < len(expected), ( - "SWH field must have both missing and non-missing values for this test to be meaningful" - ) - - return grib, expected, missing, date - - -@pytest.fixture -def fdb_with_swh_bitmap(empty_fdb, swh_grib_download): - """Archive the downloaded SWH field into a fresh empty FDB. - - Returns (config_path, expected_float32_values, missing_mask, date). - Uses the empty_fdb fixture so the FDB infrastructure is not duplicated here. - """ - config_path = empty_fdb - grib, expected, missing, date = swh_grib_download - - fdb = FDB(config_path.read_text()) - fdb.archive(grib.read_bytes()) - fdb.flush() - - return config_path, expected, missing, date - - -def test_bitmap_missing_points_become_fill_value(fdb_with_swh_bitmap): - """Bitmap-masked grid points must be returned as fill_value, not the eccodes sentinel.""" - config_path, expected, missing, date = fdb_with_swh_bitmap - - builder = ChunkedDataViewBuilder(config_path) - builder.add_part( - { - "class": "od", - "expver": "0001", - "stream": "wave", - "domain": "g", - "type": "fc", - "levtype": "sfc", - "date": date.strftime("%Y%m%d"), - "time": "0000", - "step": 0, - "param": 140229, - }, - [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, - ) - builder.fill_missing_value(-20.0) - view = builder.build() - - assert view.fill_missing_value() == -20.0 - - values = view.at((0, 0)) - np.testing.assert_array_equal(values[missing], view.fill_missing_value()) - np.testing.assert_array_equal(values[~missing], expected[~missing]) - - -def test_fill_value_propagated_to_zarr_metadata(fdb_with_swh_bitmap): - """fill_value set on the builder must appear in the zarr array metadata.""" - config_path, _, _, date = fdb_with_swh_bitmap - - builder = SimpleStoreBuilder(config_path) - builder.add_part( - { - "class": "od", - "expver": "0001", - "stream": "wave", - "domain": "g", - "type": "fc", - "levtype": "sfc", - "date": date.strftime("%Y%m%d"), - "time": "0000", - "step": 0, - "param": 140229, - }, - [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, - ) - builder.fill_missing_value(-20.0) - store = builder.build() - - arr = zarr.open_array(store, mode="r", zarr_format=3, use_consolidated=False) - assert arr.fill_value == -20.0 - - -def test_default_fill_value_replaces_bitmap_sentinel(fdb_with_swh_bitmap): - """Without an explicit fill_value call, the default (NaN) replaces the eccodes sentinel.""" - config_path, expected, missing, date = fdb_with_swh_bitmap - - builder = ChunkedDataViewBuilder(config_path) - builder.add_part( - { - "class": "od", - "expver": "0001", - "stream": "wave", - "domain": "g", - "type": "fc", - "levtype": "sfc", - "date": date.strftime("%Y%m%d"), - "time": "0000", - "step": 0, - "param": 140229, - }, - [AxisDefinition(["step"], Chunking.SINGLE_VALUE)], - ExtractorType.GRIB, - ) - view = builder.build() - - values = view.at((0, 0)) - assert np.all(np.isnan(values[missing])), "missing points should be NaN by default" - np.testing.assert_array_equal(values[~missing], expected[~missing])