Refactor coordinate specification handling in dataset templates

BrianMichell · BrianMichell · commit 2e67c0846179 · 2026-05-14T18:06:53.000Z
- Simplified the `declare_coordinate_specs` method in `AbstractDatasetTemplate` to use list comprehensions for better readability and performance.
- Updated coordinate specifications in `Seismic2DCdpGathersTemplate`, `Seismic3DCdpGathersTemplate`, `Seismic3DCocaGathersTemplate`, and `Seismic3DStreamerFieldRecordsTemplate` to remove redundant parameters.
- Introduced a new experimental module for handling raw headers, including functions to determine inclusion and attach binary headers, enhancing the ingestion pipeline's flexibility.
- Modified metadata handling in `metadata.py` to utilize the new raw headers functionality, streamlining the process of adding SEG-Y file headers.
diff --git a/src/mdio/builder/templates/base.py b/src/mdio/builder/templates/base.py
@@ -72,30 +72,25 @@ def _repr_html_(self) -> str:
         return template_repr_html(self)
 
     def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
-        """Return per-coordinate (name, dims, dtype, source, header_key, chunk_policy)."""
+        """Return the non-dimension coordinate specs (name, dims, dtype) for this template."""
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
-        specs = []
-        for coord_name in self.physical_coordinate_names:
-            specs.append(
-                CoordinateSpec(
-                    name=coord_name,
-                    dimensions=self.spatial_dimension_names,
-                    dtype=ScalarType.FLOAT64,
-                    source="header",
-                    header_key=coord_name,
-                )
+        specs = [
+            CoordinateSpec(
+                name=coord_name,
+                dimensions=self.spatial_dimension_names,
+                dtype=ScalarType.FLOAT64,
             )
-        for coord_name in self.logical_coordinate_names:
-            specs.append(
-                CoordinateSpec(
-                    name=coord_name,
-                    dimensions=self.spatial_dimension_names,
-                    dtype=ScalarType.UINT8 if coord_name == "gun" else ScalarType.INT32,
-                    source="header",
-                    header_key=coord_name,
-                )
+            for coord_name in self.physical_coordinate_names
+        ]
+        specs.extend(
+            CoordinateSpec(
+                name=coord_name,
+                dimensions=self.spatial_dimension_names,
+                dtype=ScalarType.UINT8 if coord_name == "gun" else ScalarType.INT32,
             )
+            for coord_name in self.logical_coordinate_names
+        )
         return tuple(specs)
 
     def build_dataset(
diff --git a/src/mdio/builder/templates/seismic_2d_cdp.py b/src/mdio/builder/templates/seismic_2d_cdp.py
@@ -43,20 +43,8 @@ def declare_coordinate_specs(self) -> tuple[Any, ...]:
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
         return (
-            CoordinateSpec(
-                name="cdp_x",
-                dimensions=("cdp",),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_x",
-            ),
-            CoordinateSpec(
-                name="cdp_y",
-                dimensions=("cdp",),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_y",
-            ),
+            CoordinateSpec(name="cdp_x", dimensions=("cdp",), dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="cdp_y", dimensions=("cdp",), dtype=ScalarType.FLOAT64),
         )
 
     def _add_coordinates(self) -> None:
diff --git a/src/mdio/builder/templates/seismic_3d_cdp.py b/src/mdio/builder/templates/seismic_3d_cdp.py
@@ -43,20 +43,8 @@ def declare_coordinate_specs(self) -> tuple[Any, ...]:
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
         return (
-            CoordinateSpec(
-                name="cdp_x",
-                dimensions=("inline", "crossline"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_x",
-            ),
-            CoordinateSpec(
-                name="cdp_y",
-                dimensions=("inline", "crossline"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_y",
-            ),
+            CoordinateSpec(name="cdp_x", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="cdp_y", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
         )
 
     def _add_coordinates(self) -> None:
diff --git a/src/mdio/builder/templates/seismic_3d_coca.py b/src/mdio/builder/templates/seismic_3d_coca.py
@@ -34,20 +34,8 @@ def declare_coordinate_specs(self) -> tuple[Any, ...]:
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
         return (
-            CoordinateSpec(
-                name="cdp_x",
-                dimensions=("inline", "crossline"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_x",
-            ),
-            CoordinateSpec(
-                name="cdp_y",
-                dimensions=("inline", "crossline"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="cdp_y",
-            ),
+            CoordinateSpec(name="cdp_x", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="cdp_y", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
         )
 
     def _add_coordinates(self) -> None:
diff --git a/src/mdio/builder/templates/seismic_3d_obn.py b/src/mdio/builder/templates/seismic_3d_obn.py
@@ -59,49 +59,15 @@ def _load_dataset_attributes(self) -> dict[str, Any]:
     def declare_coordinate_specs(self) -> tuple[Any, ...]:
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
+        receiver_dim = ("receiver",)
+        shot_dims = ("shot_line", "gun", "shot_index")
         return (
-            CoordinateSpec(
-                name="group_coord_x",
-                dimensions=("receiver",),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="group_coord_x",
-            ),
-            CoordinateSpec(
-                name="group_coord_y",
-                dimensions=("receiver",),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="group_coord_y",
-            ),
-            CoordinateSpec(
-                name="shot_point",
-                dimensions=("shot_line", "gun", "shot_index"),
-                dtype=ScalarType.UINT32,
-                source="header",
-                header_key="shot_point",
-            ),
-            CoordinateSpec(
-                name="orig_field_record_num",
-                dimensions=("shot_line", "gun", "shot_index"),
-                dtype=ScalarType.UINT32,
-                source="header",
-                header_key="orig_field_record_num",
-            ),
-            CoordinateSpec(
-                name="source_coord_x",
-                dimensions=("shot_line", "gun", "shot_index"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="source_coord_x",
-            ),
-            CoordinateSpec(
-                name="source_coord_y",
-                dimensions=("shot_line", "gun", "shot_index"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="source_coord_y",
-            ),
+            CoordinateSpec(name="group_coord_x", dimensions=receiver_dim, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="group_coord_y", dimensions=receiver_dim, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="shot_point", dimensions=shot_dims, dtype=ScalarType.UINT32),
+            CoordinateSpec(name="orig_field_record_num", dimensions=shot_dims, dtype=ScalarType.UINT32),
+            CoordinateSpec(name="source_coord_x", dimensions=shot_dims, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="source_coord_y", dimensions=shot_dims, dtype=ScalarType.FLOAT64),
         )
 
     def _add_coordinates(self) -> None:
diff --git a/src/mdio/builder/templates/seismic_3d_streamer_field.py b/src/mdio/builder/templates/seismic_3d_streamer_field.py
@@ -47,49 +47,15 @@ def _load_dataset_attributes(self) -> dict[str, Any]:
     def declare_coordinate_specs(self) -> tuple[Any, ...]:
         from mdio.ingestion.schema_resolver import CoordinateSpec
 
+        shot_dims = ("sail_line", "gun", "shot_index")
+        receiver_dims = ("sail_line", "gun", "shot_index", "cable", "channel")
         return (
-            CoordinateSpec(
-                name="orig_field_record_num",
-                dimensions=("sail_line", "gun", "shot_index"),
-                dtype=ScalarType.UINT32,
-                source="header",
-                header_key="orig_field_record_num",
-            ),
-            CoordinateSpec(
-                name="shot_point",
-                dimensions=("sail_line", "gun", "shot_index"),
-                dtype=ScalarType.UINT32,
-                source="header",
-                header_key="shot_point",
-            ),
-            CoordinateSpec(
-                name="source_coord_x",
-                dimensions=("sail_line", "gun", "shot_index"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="source_coord_x",
-            ),
-            CoordinateSpec(
-                name="source_coord_y",
-                dimensions=("sail_line", "gun", "shot_index"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="source_coord_y",
-            ),
-            CoordinateSpec(
-                name="group_coord_x",
-                dimensions=("sail_line", "gun", "shot_index", "cable", "channel"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="group_coord_x",
-            ),
-            CoordinateSpec(
-                name="group_coord_y",
-                dimensions=("sail_line", "gun", "shot_index", "cable", "channel"),
-                dtype=ScalarType.FLOAT64,
-                source="header",
-                header_key="group_coord_y",
-            ),
+            CoordinateSpec(name="orig_field_record_num", dimensions=shot_dims, dtype=ScalarType.UINT32),
+            CoordinateSpec(name="shot_point", dimensions=shot_dims, dtype=ScalarType.UINT32),
+            CoordinateSpec(name="source_coord_x", dimensions=shot_dims, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="source_coord_y", dimensions=shot_dims, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="group_coord_x", dimensions=receiver_dims, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="group_coord_y", dimensions=receiver_dims, dtype=ScalarType.FLOAT64),
         )
 
     def _add_coordinates(self) -> None:
diff --git a/src/mdio/ingestion/_raw_headers_experimental.py b/src/mdio/ingestion/_raw_headers_experimental.py
@@ -0,0 +1,58 @@
+"""EXPERIMENTAL: gating for the ``MDIO__IMPORT__RAW_HEADERS`` feature.
+
+This module concentrates all checks tied to the experimental raw-headers feature so
+that removing the feature is a single-file delete plus the (small) handful of call
+sites that import from here. The behaviour:
+
+* ``should_include_raw_headers()`` decides whether the ``raw_headers`` variable
+  should be added to the dataset for this run. Honours the env-driven
+  :class:`mdio.core.config.MDIOSettings` and the active Zarr format.
+* ``attach_raw_binary_header()`` adds the base64-encoded binary file header to a
+  segy_file_header attribute dict, if (and only if) the feature is enabled.
+
+Removal plan
+------------
+Delete this module and:
+
+#. Drop the call to :func:`should_include_raw_headers` in
+   :mod:`mdio.ingestion.pipeline` (and the resulting ``include_raw_headers``
+   kwarg threaded through :class:`~mdio.ingestion.dataset_factory.DatasetFactory`
+   / :meth:`AbstractDatasetTemplate.build_dataset` / ``_add_raw_headers``).
+#. Drop the call to :func:`attach_raw_binary_header` in
+   :mod:`mdio.ingestion.metadata`.
+#. Drop ``MDIOSettings.raw_headers``.
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+from typing import TYPE_CHECKING
+
+import zarr
+
+from mdio.constants import ZarrFormat
+from mdio.core.config import MDIOSettings
+
+if TYPE_CHECKING:
+    from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def should_include_raw_headers() -> bool:
+    """True iff the experimental flag is set and the active Zarr format supports it."""
+    if not MDIOSettings().raw_headers:
+        return False
+    if zarr.config.get("default_zarr_format") == ZarrFormat.V2:
+        logger.warning("Raw headers are only supported for Zarr v3. Skipping raw headers.")
+        return False
+    logger.warning("MDIO__IMPORT__RAW_HEADERS is experimental and expected to change or be removed.")
+    return True
+
+
+def attach_raw_binary_header(attrs: dict[str, Any], raw_binary_headers: bytes) -> None:
+    """Attach a base64-encoded raw binary header to ``attrs`` if the feature is enabled."""
+    if not MDIOSettings().raw_headers:
+        return
+    attrs["rawBinaryHeader"] = base64.b64encode(raw_binary_headers).decode("ascii")
diff --git a/src/mdio/ingestion/metadata.py b/src/mdio/ingestion/metadata.py
@@ -1,16 +1,12 @@
-"""Metadata utilities for MDIO ingestion.
-
-This module contains functions for adding metadata to MDIO datasets
-during ingestion.
-"""
+"""Metadata utilities for MDIO ingestion."""
 
 from __future__ import annotations
 
-import base64
 from typing import TYPE_CHECKING
 from typing import Any
 
 from mdio.core.config import MDIOSettings
+from mdio.ingestion._raw_headers_experimental import attach_raw_binary_header
 
 if TYPE_CHECKING:
     from xarray import Dataset as xr_Dataset
@@ -30,9 +26,7 @@ def add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, An
 
 def add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo) -> xr_Dataset:
     """Add SEG-Y file headers to the dataset as metadata."""
-    settings = MDIOSettings()
-
-    if not settings.save_segy_file_header:
+    if not MDIOSettings().save_segy_file_header:
         return xr_dataset
 
     expected_rows = 40
@@ -50,14 +44,13 @@ def add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo)
         raise ValueError(err)
 
     xr_dataset["segy_file_header"] = ((), "")
-    xr_dataset["segy_file_header"].attrs.update(
+    file_header_attrs = xr_dataset["segy_file_header"].attrs
+    file_header_attrs.update(
         {
             "textHeader": segy_file_info.text_header,
             "binaryHeader": segy_file_info.binary_header_dict,
         }
     )
-    if settings.raw_headers:
-        raw_binary_base64 = base64.b64encode(segy_file_info.raw_binary_headers).decode("ascii")
-        xr_dataset["segy_file_header"].attrs.update({"rawBinaryHeader": raw_binary_base64})
+    attach_raw_binary_header(file_header_attrs, segy_file_info.raw_binary_headers)
 
     return xr_dataset
diff --git a/src/mdio/ingestion/pipeline.py b/src/mdio/ingestion/pipeline.py
diff --git a/src/mdio/ingestion/schema_resolver.py b/src/mdio/ingestion/schema_resolver.py
diff --git a/tests/unit/test_ingestion_schema_resolver.py b/tests/unit/test_ingestion_schema_resolver.py
