From a09415c73592519ebfa9f26490be0b0084396239 Mon Sep 17 00:00:00 2001
From: Ben Lasscock <ben.lasscock@tgs.com>
Date: Mon, 13 Jul 2026 02:01:20 +0000
Subject: [PATCH 1/3] Add multi-shard SEG-Y to single MDIO consolidation

Additive feature for consolidating many SEG-Y shards (e.g. a shot survey split
across files) into one MDIO store, cloud-to-cloud, without changing the default
single-file segy_to_mdio behavior.

- allocate_mdio_grid / append_segy_shard: two-phase region-write workflow that
  builds an empty global grid, then writes each shard into its region in place.
- plan_consolidation (+ ConsolidationPlan): deterministic planner computing per
  shard chunk ownership, per-chunk fill vs merge (read-modify-write) mode, the
  shard conflict graph, and concurrency waves. Orchestration is left to callers.
- Optional read-modify-write in trace_worker / blocked_io.to_zarr via a
  merge_chunks set, so shards sharing a boundary chunk don't clobber each other.
  Defaults to the original fast pure-write path when not provided.
- Export the new API from mdio, mdio.converters.
---
 src/mdio/__init__.py                   |   8 +
 src/mdio/converters/__init__.py        |   4 +-
 src/mdio/converters/segy.py            |  90 ++++++++
 src/mdio/ingestion/consolidation.py    | 246 ++++++++++++++++++++
 src/mdio/ingestion/segy/consolidate.py | 303 +++++++++++++++++++++++++
 src/mdio/segy/_workers.py              |  28 ++-
 src/mdio/segy/blocked_io.py            |  20 +-
 7 files changed, 689 insertions(+), 10 deletions(-)
 create mode 100644 src/mdio/ingestion/consolidation.py
 create mode 100644 src/mdio/ingestion/segy/consolidate.py

diff --git a/src/mdio/__init__.py b/src/mdio/__init__.py
index efd764927..99da183a5 100644
--- a/src/mdio/__init__.py
+++ b/src/mdio/__init__.py
@@ -6,9 +6,13 @@
 
 from mdio.api.io import open_mdio
 from mdio.api.io import to_mdio
+from mdio.converters import allocate_mdio_grid
+from mdio.converters import append_segy_shard
 from mdio.converters import mdio_to_segy
 from mdio.converters import segy_to_mdio
 from mdio.ingestion import ResolvedSchema
+from mdio.ingestion.consolidation import ConsolidationPlan
+from mdio.ingestion.consolidation import plan_consolidation
 from mdio.optimize.access_pattern import OptimizedAccessPatternConfig
 from mdio.optimize.access_pattern import optimize_access_patterns
 from mdio.segy.geometry import GridOverrides
@@ -24,6 +28,10 @@
     "GridOverrides",
     "open_mdio",
     "to_mdio",
+    "allocate_mdio_grid",
+    "append_segy_shard",
+    "plan_consolidation",
+    "ConsolidationPlan",
     "mdio_to_segy",
     "segy_to_mdio",
     "OptimizedAccessPatternConfig",
diff --git a/src/mdio/converters/__init__.py b/src/mdio/converters/__init__.py
index fd88595ff..ca9b1c11c 100644
--- a/src/mdio/converters/__init__.py
+++ b/src/mdio/converters/__init__.py
@@ -1,6 +1,8 @@
 """MDIO Data conversion API."""
 
 from mdio.converters.mdio import mdio_to_segy
+from mdio.converters.segy import allocate_mdio_grid
+from mdio.converters.segy import append_segy_shard
 from mdio.converters.segy import segy_to_mdio
 
-__all__ = ["mdio_to_segy", "segy_to_mdio"]
+__all__ = ["allocate_mdio_grid", "append_segy_shard", "mdio_to_segy", "segy_to_mdio"]
diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py
index 759db69a4..940f928d9 100644
--- a/src/mdio/converters/segy.py
+++ b/src/mdio/converters/segy.py
@@ -18,6 +18,7 @@
     from upath import UPath
 
     from mdio.builder.templates.base import AbstractDatasetTemplate
+    from mdio.core.dimension import Dimension
 
 
 def _coerce_grid_overrides(
@@ -78,3 +79,92 @@ def segy_to_mdio(  # noqa: PLR0913
         grid_overrides=typed_grid_overrides,
         segy_header_overrides=segy_header_overrides,
     )
+
+
+def allocate_mdio_grid(  # noqa: PLR0913
+    segy_spec: SegySpec,
+    mdio_template: AbstractDatasetTemplate,
+    global_dimensions: list[Dimension],
+    output_path: UPath | Path | str,
+    reference_segy_path: UPath | Path | str,
+    overwrite: bool = False,
+    grid_overrides: GridOverrides | dict[str, Any] | None = None,
+) -> UPath:
+    """Allocate an empty global MDIO store for multi-shard consolidation.
+
+    ADDITIVE: creates the dataset skeleton sized to ``global_dimensions`` (the union of
+    all shards' coordinates, including the trailing sample dimension) so shards can be
+    written into sub-regions with :func:`append_segy_shard`. Does not alter single-file
+    ``segy_to_mdio`` behavior.
+
+    Args:
+        segy_spec: SEG-Y spec shared by all shards.
+        mdio_template: MDIO template shared by all shards.
+        global_dimensions: Ordered global grid dimensions (incl. sample dimension).
+        output_path: Output MDIO store path.
+        reference_segy_path: Any one shard, used only to derive store units/metadata.
+        overwrite: Whether to overwrite an existing store.
+        grid_overrides: Optional grid overrides (dict accepted but deprecated).
+
+    Returns:
+        The normalized output path of the allocated store.
+    """
+    typed_grid_overrides = _coerce_grid_overrides(grid_overrides)
+
+    from mdio.ingestion.segy.consolidate import allocate_mdio_grid as _allocate  # noqa: PLC0415
+
+    return _allocate(
+        segy_spec=segy_spec,
+        mdio_template=mdio_template,
+        global_dimensions=global_dimensions,
+        output_path=output_path,
+        reference_segy_path=reference_segy_path,
+        overwrite=overwrite,
+        grid_overrides=typed_grid_overrides,
+    )
+
+
+def append_segy_shard(  # noqa: PLR0913
+    segy_spec: SegySpec,
+    mdio_template: AbstractDatasetTemplate,
+    global_dimensions: list[Dimension],
+    input_path: UPath | Path | str,
+    output_path: UPath | Path | str,
+    grid_overrides: GridOverrides | dict[str, Any] | None = None,
+    segy_header_overrides: SegyHeaderOverrides | None = None,
+    merge_chunks: set[tuple[int, ...]] | None = None,
+) -> dict[str, slice]:
+    """Ingest one SEG-Y shard into its region of a pre-allocated global MDIO store.
+
+    ADDITIVE: the store must already exist (see :func:`allocate_mdio_grid`). The shard's
+    traces are placed at their global grid positions via an in-place (``mode="r+"``)
+    write; other regions are untouched. Call once per shard.
+
+    Args:
+        segy_spec: SEG-Y spec (same as allocation).
+        mdio_template: MDIO template (same as allocation).
+        global_dimensions: Global grid dimensions used at allocation (incl. sample dim).
+        input_path: The shard SEG-Y path.
+        output_path: The pre-allocated global MDIO store path.
+        grid_overrides: Optional grid overrides (dict accepted but deprecated).
+        segy_header_overrides: Optional SEG-Y header overrides for this shard.
+        merge_chunks: Optional set of shared chunk-grid indices to write read-modify-write
+            (e.g. ``plan.merge_chunks_for(shard_id)`` from :func:`mdio.plan_consolidation`).
+
+    Returns:
+        The spatial region (dim name -> slice) the shard was written into.
+    """
+    typed_grid_overrides = _coerce_grid_overrides(grid_overrides)
+
+    from mdio.ingestion.segy.consolidate import append_segy_shard as _append  # noqa: PLC0415
+
+    return _append(
+        segy_spec=segy_spec,
+        mdio_template=mdio_template,
+        global_dimensions=global_dimensions,
+        input_path=input_path,
+        output_path=output_path,
+        grid_overrides=typed_grid_overrides,
+        segy_header_overrides=segy_header_overrides,
+        merge_chunks=merge_chunks,
+    )
diff --git a/src/mdio/ingestion/consolidation.py b/src/mdio/ingestion/consolidation.py
new file mode 100644
index 000000000..dbac7aa0f
--- /dev/null
+++ b/src/mdio/ingestion/consolidation.py
@@ -0,0 +1,246 @@
+"""Deterministic planning for multi-shard SEG-Y -> single MDIO consolidation.
+
+ADDITIVE and format-agnostic. Given each shard's *global* index coverage and the store's
+spatial chunk shape, this computes - with pure arithmetic, no I/O - everything needed to
+consolidate many shards into one MDIO store safely and with maximum parallelism:
+
+  * which output chunks each shard writes,
+  * per-chunk write mode: ``"fill"`` (single owner -> fast pure write) vs ``"merge"``
+    (multiple owners -> read-modify-write so shards sharing a boundary chunk don't clobber
+    each other),
+  * the shard conflict graph (two shards conflict iff they share a chunk), and
+  * a deterministic set of concurrency ``waves`` (no two conflicting shards in the same
+    wave), so an external orchestrator can parallelize non-overlapping shards and stage
+    the conflicting ones.
+
+Orchestration itself is intentionally NOT handled here (or anywhere in MDIO) - the caller
+turns this plan into an execution graph for whatever runtime it uses.
+
+Determinism: outputs are a pure function of the inputs (coverage, chunk shape, and the
+shard iteration order used for tie-breaking in wave coloring).
+"""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+from dataclasses import field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from collections.abc import Sequence
+
+# chunk index in the chunk grid (one int per spatial dim)
+ChunkIndex = tuple[int, ...]
+# a shard's global coverage: spatial dim name -> (start, stop) half-open index range
+Coverage = dict[str, tuple[int, int]]
+
+
+@dataclass
+class ConsolidationPlan:
+    """Deterministic consolidation plan (see module docstring)."""
+
+    spatial_dims: tuple[str, ...]
+    chunk_shape: tuple[int, ...]
+    chunk_owners: dict[ChunkIndex, list[str]]
+    chunk_modes: dict[ChunkIndex, str]
+    shard_chunks: dict[str, list[ChunkIndex]]
+    shard_merge_chunks: dict[str, list[ChunkIndex]]
+    conflicts: dict[str, list[str]]
+    waves: list[list[str]]
+    global_shape: tuple[int, ...] | None = None
+    warnings: list[str] = field(default_factory=list)
+
+    @property
+    def merge_chunk_count(self) -> int:
+        """Number of chunks that require read-modify-write."""
+        return sum(1 for m in self.chunk_modes.values() if m == "merge")
+
+    def merge_chunks_for(self, shard_id: str) -> set[ChunkIndex]:
+        """Set of chunks this shard must write in ``merge`` (RMW) mode."""
+        return set(self.shard_merge_chunks.get(shard_id, []))
+
+    def to_dict(self) -> dict:
+        """JSON-serializable view (tuple keys rendered as comma-joined strings)."""
+
+        def _key(idx: ChunkIndex) -> str:
+            return ",".join(map(str, idx))
+
+        return {
+            "spatial_dims": list(self.spatial_dims),
+            "chunk_shape": list(self.chunk_shape),
+            "global_shape": list(self.global_shape) if self.global_shape else None,
+            "chunk_owners": {_key(k): v for k, v in self.chunk_owners.items()},
+            "chunk_modes": {_key(k): v for k, v in self.chunk_modes.items()},
+            "shard_chunks": {s: [_key(c) for c in cs] for s, cs in self.shard_chunks.items()},
+            "shard_merge_chunks": {s: [_key(c) for c in cs] for s, cs in self.shard_merge_chunks.items()},
+            "conflicts": self.conflicts,
+            "waves": self.waves,
+            "merge_chunk_count": self.merge_chunk_count,
+            "warnings": self.warnings,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> ConsolidationPlan:
+        """Reconstruct a plan from :meth:`to_dict` output (inverse of ``to_dict``).
+
+        Chunk keys serialized as comma-joined strings (e.g. ``"1,0"``) are parsed back to
+        integer tuples. Derived-only fields (``merge_chunk_count``) are ignored.
+        """
+
+        def _idx(key: str) -> ChunkIndex:
+            return tuple(int(p) for p in key.split(",")) if key else ()
+
+        return cls(
+            spatial_dims=tuple(data["spatial_dims"]),
+            chunk_shape=tuple(int(c) for c in data["chunk_shape"]),
+            chunk_owners={_idx(k): list(v) for k, v in data.get("chunk_owners", {}).items()},
+            chunk_modes={_idx(k): v for k, v in data.get("chunk_modes", {}).items()},
+            shard_chunks={s: [_idx(c) for c in cs] for s, cs in data.get("shard_chunks", {}).items()},
+            shard_merge_chunks={s: [_idx(c) for c in cs] for s, cs in data.get("shard_merge_chunks", {}).items()},
+            conflicts={s: list(v) for s, v in data.get("conflicts", {}).items()},
+            waves=[list(w) for w in data.get("waves", [])],
+            global_shape=tuple(data["global_shape"]) if data.get("global_shape") else None,
+            warnings=list(data.get("warnings", [])),
+        )
+
+
+def coverage_from_index_arrays(
+    global_coords: Mapping[str, Sequence],
+    shard_values: Mapping[str, Sequence],
+) -> Coverage:
+    """Compute a shard's global index coverage from its coordinate values.
+
+    Args:
+        global_coords: dim name -> sorted global coordinate vector (the union across shards).
+        shard_values: dim name -> the coordinate values present in this shard.
+
+    Returns:
+        dim name -> (start, stop) half-open global index range for the shard.
+    """
+    import numpy as np
+
+    coverage: Coverage = {}
+    for dim, gcoords in global_coords.items():
+        if dim not in shard_values:
+            continue
+        g = np.asarray(gcoords)
+        idx = np.searchsorted(g, np.asarray(shard_values[dim]))
+        coverage[dim] = (int(idx.min()), int(idx.max()) + 1)
+    return coverage
+
+
+def _chunk_indices_for(coverage: Coverage, spatial_dims: Sequence[str], chunk_shape: Sequence[int]) -> list[ChunkIndex]:
+    """All chunk-grid indices a coverage box touches."""
+    per_dim_ranges = []
+    for dim, csize in zip(spatial_dims, chunk_shape, strict=True):
+        start, stop = coverage[dim]
+        first = start // csize
+        last = (stop - 1) // csize
+        per_dim_ranges.append(range(first, last + 1))
+    return list(itertools.product(*per_dim_ranges))
+
+
+def _color_waves(shard_ids: Sequence[str], conflicts: dict[str, list[str]]) -> list[list[str]]:
+    """Greedy deterministic graph coloring -> concurrency waves.
+
+    No two conflicting shards share a wave, so within a wave no two shards write the same
+    chunk. Iteration follows ``shard_ids`` order for reproducibility.
+    """
+    color: dict[str, int] = {}
+    for sid in shard_ids:
+        used = {color[n] for n in conflicts.get(sid, []) if n in color}
+        c = 0
+        while c in used:
+            c += 1
+        color[sid] = c
+
+    num_waves = (max(color.values()) + 1) if color else 0
+    waves: list[list[str]] = [[] for _ in range(num_waves)]
+    for sid in shard_ids:
+        waves[color[sid]].append(sid)
+    return waves
+
+
+def plan_consolidation(
+    shard_coverage: Mapping[str, Coverage],
+    spatial_dims: Sequence[str],
+    chunk_shape: Sequence[int],
+    global_shape: Sequence[int] | None = None,
+) -> ConsolidationPlan:
+    """Compute a deterministic consolidation plan.
+
+    Args:
+        shard_coverage: shard id -> {spatial dim -> (start, stop)} global index coverage.
+            Iteration order of this mapping is the tie-break order for wave coloring.
+        spatial_dims: Ordered spatial dimension names (must match the store's dim order,
+            excluding the trailing sample dimension).
+        chunk_shape: Spatial chunk sizes, aligned to ``spatial_dims``. MUST equal the
+            store's data-variable spatial chunk sizes for the write-mode mask to line up.
+        global_shape: Optional spatial global sizes (for validation/reporting).
+
+    Returns:
+        A :class:`ConsolidationPlan`.
+
+    Raises:
+        ValueError: If a shard is missing coverage for a spatial dimension.
+    """
+    spatial_dims = tuple(spatial_dims)
+    chunk_shape = tuple(int(c) for c in chunk_shape)
+    warnings: list[str] = []
+
+    shard_ids = list(shard_coverage.keys())
+
+    shard_chunks: dict[str, list[ChunkIndex]] = {}
+    chunk_owners: dict[ChunkIndex, list[str]] = {}
+    for sid in shard_ids:
+        cov = shard_coverage[sid]
+        missing = [d for d in spatial_dims if d not in cov]
+        if missing:
+            err = f"Shard '{sid}' missing coverage for dimensions {missing}."
+            raise ValueError(err)
+        chunks = _chunk_indices_for(cov, spatial_dims, chunk_shape)
+        shard_chunks[sid] = chunks
+        for c in chunks:
+            chunk_owners.setdefault(c, []).append(sid)
+
+    chunk_modes: dict[ChunkIndex, str] = {
+        c: ("merge" if len(owners) > 1 else "fill") for c, owners in chunk_owners.items()
+    }
+
+    shard_merge_chunks: dict[str, list[ChunkIndex]] = {
+        sid: [c for c in chunks if chunk_modes[c] == "merge"] for sid, chunks in shard_chunks.items()
+    }
+
+    # Conflict graph: shards sharing any chunk (necessarily a merge chunk).
+    conflicts: dict[str, list[str]] = {sid: [] for sid in shard_ids}
+    for owners in chunk_owners.values():
+        if len(owners) < 2:
+            continue
+        for a in owners:
+            for b in owners:
+                if a != b and b not in conflicts[a]:
+                    conflicts[a].append(b)
+
+    if any(chunk_modes[c] == "merge" for c in chunk_modes):
+        warnings.append(
+            "Some chunks are shared by multiple shards and will use read-modify-write. "
+            "Shards sharing a chunk are placed in different waves and MUST run sequentially "
+            "with respect to each other."
+        )
+
+    waves = _color_waves(shard_ids, conflicts)
+
+    return ConsolidationPlan(
+        spatial_dims=spatial_dims,
+        chunk_shape=chunk_shape,
+        chunk_owners=chunk_owners,
+        chunk_modes=chunk_modes,
+        shard_chunks=shard_chunks,
+        shard_merge_chunks=shard_merge_chunks,
+        conflicts=conflicts,
+        waves=waves,
+        global_shape=tuple(int(s) for s in global_shape) if global_shape is not None else None,
+        warnings=warnings,
+    )
diff --git a/src/mdio/ingestion/segy/consolidate.py b/src/mdio/ingestion/segy/consolidate.py
new file mode 100644
index 000000000..82752add5
--- /dev/null
+++ b/src/mdio/ingestion/segy/consolidate.py
@@ -0,0 +1,303 @@
+"""Multi-shard SEG-Y -> single MDIO consolidation via region writes.
+
+ADDITIVE feature. This module does NOT change the default single-file
+``segy_to_mdio`` behavior. It adds a two-phase workflow for the common case where
+one logical dataset (e.g. a shot survey) is split across many SEG-Y files that
+should become a single MDIO:
+
+  1. :func:`allocate_mdio_grid` - create the empty global MDIO store from a template
+     plus the *global* dimension coordinates (the union across all shards). This
+     writes the dataset skeleton (dimension coordinates + fill-valued arrays) once.
+
+  2. :func:`append_segy_shard` - ingest one shard and place its traces into their
+     region of the pre-allocated global grid using an in-place (``mode="r+"``) Zarr
+     write, leaving all other regions untouched. Call once per shard.
+
+Everything is cloud-native (reads the SEG-Y over the network, writes into the object
+store). No SEG-Y bytes are downloaded/concatenated.
+
+Caller (or an LLM agent driving consolidation) MUST guarantee:
+  * All shards share the same sample axis, data sample format, revision, and
+    header/index byte layout (so one ``segy_spec`` + template applies to all).
+  * Every shard's spatial coordinates are a subset of ``global_dimensions``.
+  * Shards are DISJOINT and, on the concatenation dimension, occupy a contiguous
+    block whose boundaries align to the store's chunk boundaries on that dimension.
+    (A boundary chunk shared by two shards would be clobbered by the second write.)
+
+These invariants are exactly what upstream inspection tooling should verify before
+consolidation.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import numpy as np
+from segy.config import SegyFileSettings
+from zarr import open_group as zarr_open_group
+
+from mdio.api.io import _normalize_path
+from mdio.api.io import _normalize_storage_options
+from mdio.api.io import to_mdio
+from mdio.builder.xarray_builder import to_xarray_dataset
+from mdio.core.grid import Grid
+from mdio.ingestion.dataset_factory import build_mdio_dataset
+from mdio.ingestion.schema.resolver import SchemaResolver
+from mdio.ingestion.segy.coordinates import get_spatial_coordinate_unit
+from mdio.ingestion.segy.coordinates import populate_coordinates
+from mdio.ingestion.segy.coordinates import resolve_units
+from mdio.ingestion.segy.index_strategies import IndexStrategyRegistry
+from mdio.ingestion.segy.raw_headers import build_raw_header_variables
+from mdio.ingestion.segy.reader import read_index_headers
+from mdio.ingestion.segy.validation import validate_spec_in_template
+from mdio.segy import blocked_io
+from mdio.segy.file import get_segy_file_info
+from mdio.segy.geometry import validate_overrides_for_template
+from mdio.segy.utilities import build_mdio_header_type
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from segy.config import SegyHeaderOverrides
+    from segy.schema import SegySpec
+    from upath import UPath
+
+    from mdio.builder.templates.base import AbstractDatasetTemplate
+    from mdio.core.dimension import Dimension
+    from mdio.segy.file import SegyFileArguments
+    from mdio.segy.geometry import GridOverrides
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_schema(mdio_template: AbstractDatasetTemplate, grid_overrides: GridOverrides | None):
+    """Resolve the (format-agnostic) schema for a template + optional grid overrides."""
+    schema_effect = IndexStrategyRegistry().schema_effect(grid_overrides)
+    return SchemaResolver().resolve(mdio_template, schema_effect)
+
+
+def allocate_mdio_grid(  # noqa: PLR0913
+    segy_spec: SegySpec,
+    mdio_template: AbstractDatasetTemplate,
+    global_dimensions: list[Dimension],
+    output_path: UPath | Path | str,
+    reference_segy_path: UPath | Path | str,
+    overwrite: bool = False,
+    grid_overrides: GridOverrides | None = None,
+) -> UPath:
+    """Allocate an empty global MDIO store sized to ``global_dimensions``.
+
+    Builds the dataset skeleton (dimension coordinates + fill-valued data/header/mask
+    arrays) so shards can later be written into sub-regions with
+    :func:`append_segy_shard`.
+
+    Args:
+        segy_spec: SEG-Y spec shared by all shards (drives the header dtype).
+        mdio_template: MDIO dataset template shared by all shards.
+        global_dimensions: Ordered dimensions of the *global* grid, INCLUDING the
+            trailing sample/vertical dimension. Coordinates are the union across shards.
+        output_path: Output MDIO store path.
+        reference_segy_path: Any one shard, used only to derive units/scalar metadata
+            for the store (header values are not written here).
+        overwrite: Whether to overwrite an existing store.
+        grid_overrides: Optional grid override configuration (must match shards).
+
+    Returns:
+        The normalized output path of the allocated store.
+
+    Raises:
+        FileExistsError: If the store exists and ``overwrite`` is False.
+    """
+    validate_overrides_for_template(grid_overrides, mdio_template)
+    validate_spec_in_template(segy_spec, mdio_template)
+
+    output_path = _normalize_path(output_path)
+    if not overwrite and output_path.exists():
+        err = f"Output location '{output_path.as_posix()}' exists. Set `overwrite=True` if intended."
+        raise FileExistsError(err)
+
+    ref_path = _normalize_path(reference_segy_path)
+    ref_kwargs: SegyFileArguments = {
+        "url": ref_path.as_posix(),
+        "spec": segy_spec,
+        "settings": SegyFileSettings(storage_options=ref_path.storage_options),
+        "header_overrides": None,
+    }
+    file_info = get_segy_file_info(ref_kwargs)
+    units = resolve_units(mdio_template, get_spatial_coordinate_unit(file_info))
+
+    schema = _resolve_schema(mdio_template, grid_overrides)
+
+    grid = Grid(dims=list(global_dimensions))
+    header_dtype = build_mdio_header_type(segy_spec)
+    extra_variables = build_raw_header_variables(schema)
+    mdio_ds = build_mdio_dataset(
+        schema=schema,
+        sizes=grid.shape,
+        header_dtype=header_dtype,
+        units=units,
+        extra_variables=extra_variables,
+    )
+
+    xr_dataset = to_xarray_dataset(mdio_ds=mdio_ds)
+
+    # Write the dimension coordinates (small, global, known up-front). Data, headers and
+    # trace_mask stay at their fill values (trace_mask fill == False => all-dead grid).
+    for dim in grid.dims:
+        xr_dataset[dim.name].values[:] = dim.coords
+
+    to_mdio(xr_dataset, output_path=output_path, mode="w", compute=False)
+    dim_names = [dim.name for dim in grid.dims]
+    to_mdio(xr_dataset[dim_names], output_path=output_path, mode="r+", compute=True)
+
+    logger.info("Allocated global MDIO grid %s at %s", grid.shape, output_path.as_posix())
+    return output_path
+
+
+def _region_slices_for(dims: tuple[str, ...], region: dict[str, slice]) -> tuple[slice, ...]:
+    """Per-variable region slices: use the region slice where the dim applies, else full."""
+    return tuple(region.get(name, slice(None)) for name in dims)
+
+
+def _write_region_vars(
+    output_path: UPath,
+    xr_dataset,
+    region: dict[str, slice],
+    var_names: list[str],
+) -> None:
+    """Write only the shard's sub-region of the given variables via in-place Zarr slicing.
+
+    Mirrors the blocked-I/O pattern (direct Zarr ``array[slices] = values``) so we avoid
+    xarray region-write constraints on dimension coordinates.
+    """
+    storage_options = _normalize_storage_options(output_path)
+    zarr_group = zarr_open_group(output_path.as_posix(), mode="r+", storage_options=storage_options)
+    for name in var_names:
+        da = xr_dataset[name]
+        slices = _region_slices_for(tuple(da.dims), region)
+        zarr_group[name][slices] = np.asarray(da.values)[slices]
+
+
+def append_segy_shard(  # noqa: PLR0913
+    segy_spec: SegySpec,
+    mdio_template: AbstractDatasetTemplate,
+    global_dimensions: list[Dimension],
+    input_path: UPath | Path | str,
+    output_path: UPath | Path | str,
+    grid_overrides: GridOverrides | None = None,
+    segy_header_overrides: SegyHeaderOverrides | None = None,
+    merge_chunks: set[tuple[int, ...]] | None = None,
+) -> dict[str, slice]:
+    """Ingest one SEG-Y shard into its region of a pre-allocated global MDIO store.
+
+    The store must already exist (see :func:`allocate_mdio_grid`). The shard's traces are
+    placed at their GLOBAL grid positions (the grid map is built by searching the shard's
+    header values into ``global_dimensions``), so nothing outside the shard is touched.
+
+    Args:
+        segy_spec: SEG-Y spec (same as used for allocation).
+        mdio_template: MDIO template (same as used for allocation).
+        global_dimensions: The global grid dimensions used at allocation (incl. sample dim).
+        input_path: The shard SEG-Y path.
+        output_path: The pre-allocated global MDIO store path.
+        grid_overrides: Optional grid overrides (same as allocation).
+        segy_header_overrides: Optional SEG-Y header overrides for this shard.
+        merge_chunks: Optional set of chunk-grid indices this shard shares with other shards;
+            these are written read-modify-write to avoid clobbering. Typically
+            ``plan.merge_chunks_for(shard_id)`` from :func:`mdio.plan_consolidation`. When
+            None, all of this shard's chunks are pure fast writes (safe only if the shard's
+            chunks are exclusive to it).
+
+    Returns:
+        The spatial region (dim name -> slice) the shard was written into.
+
+    Raises:
+        ValueError: If the shard maps to no cells within the global grid.
+    """
+    validate_overrides_for_template(grid_overrides, mdio_template)
+    validate_spec_in_template(segy_spec, mdio_template)
+
+    input_path = _normalize_path(input_path)
+    output_path = _normalize_path(output_path)
+
+    segy_file_kwargs: SegyFileArguments = {
+        "url": input_path.as_posix(),
+        "spec": segy_spec,
+        "settings": SegyFileSettings(storage_options=input_path.storage_options),
+        "header_overrides": segy_header_overrides,
+    }
+    file_info = get_segy_file_info(segy_file_kwargs)
+    units = resolve_units(mdio_template, get_spatial_coordinate_unit(file_info))
+    schema = _resolve_schema(mdio_template, grid_overrides)
+
+    indexed_headers, _shard_dims = read_index_headers(
+        segy_file_kwargs=segy_file_kwargs,
+        file_info=file_info,
+        schema=schema,
+        grid_overrides=grid_overrides,
+        synthesize_dims=mdio_template.synthesize_missing_dims,
+        template=mdio_template,
+    )
+
+    # Build the grid on the GLOBAL dimensions; build_map searchsorts the shard's header
+    # values into the global coords, so live cells land at their global positions.
+    grid = Grid(dims=list(global_dimensions))
+    grid.build_map(indexed_headers)
+
+    live_mask = np.asarray(grid.live_mask[:])
+    nonzero = np.argwhere(live_mask)
+    if nonzero.size == 0:
+        err = "Shard maps to no cells within the global grid; check global_dimensions and index map."
+        raise ValueError(err)
+
+    spatial_dim_names = grid.dim_names[:-1]
+    mins = nonzero.min(axis=0)
+    maxs = nonzero.max(axis=0) + 1
+    region = {name: slice(int(mins[i]), int(maxs[i])) for i, name in enumerate(spatial_dim_names)}
+    logger.info("Appending shard %s into region %s", input_path.as_posix(), region)
+
+    # Build a global in-memory xarray dataset (for encoding + coordinate population); do NOT
+    # re-create the store.
+    header_dtype = build_mdio_header_type(segy_spec)
+    extra_variables = build_raw_header_variables(schema)
+    mdio_ds = build_mdio_dataset(
+        schema=schema,
+        sizes=grid.shape,
+        header_dtype=header_dtype,
+        units=units,
+        extra_variables=extra_variables,
+    )
+    xr_dataset = to_xarray_dataset(mdio_ds=mdio_ds)
+
+    non_dim_coords = {}
+    for coord in schema.coordinates:
+        if coord.name in indexed_headers.dtype.names:
+            non_dim_coords[coord.name] = np.array(indexed_headers[coord.name])
+
+    xr_dataset, drop_vars_delayed = populate_coordinates(
+        dataset=xr_dataset,
+        grid=grid,
+        coords=non_dim_coords,
+        spatial_coordinate_scalar=file_info.coordinate_scalar,
+    )
+    xr_dataset.trace_mask.data[:] = grid.live_mask
+
+    # Region-scoped write of trace_mask + non-dim coordinates (dimension coordinates were
+    # written once at allocation and are excluded here).
+    non_dim_coord_names = [c.name for c in schema.coordinates if c.name in xr_dataset]
+    _write_region_vars(output_path, xr_dataset, region, var_names=non_dim_coord_names + ["trace_mask"])
+
+    # Data + structured headers: blocked-I/O writes only the shard's live chunks (mode="r+")
+    # using the global grid map, so other shards' regions are untouched.
+    xr_dataset = xr_dataset.drop_vars(drop_vars_delayed)
+    blocked_io.to_zarr(
+        segy_file_kwargs=segy_file_kwargs,
+        output_path=output_path,
+        grid_map=grid.map,
+        dataset=xr_dataset,
+        data_variable_name=schema.default_variable_name,
+        merge_chunks=merge_chunks,
+    )
+
+    return region
diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py
index a67210526..41c40d7c1 100644
--- a/src/mdio/segy/_workers.py
+++ b/src/mdio/segy/_workers.py
@@ -116,7 +116,7 @@ def trace_worker_init(  # noqa: PLR0913
     _worker_state["grid_map"] = grid_map
 
 
-def trace_worker(region: dict[str, slice]) -> SummaryStatistics | None:
+def trace_worker(region: dict[str, slice], merge: bool = False) -> SummaryStatistics | None:
     """Writes a subset of traces from a region of the dataset of Zarr file.
 
     Reads its shared inputs (SEG-Y handle, Zarr arrays, grid map) from the per-process state set up
@@ -124,6 +124,11 @@ def trace_worker(region: dict[str, slice]) -> SummaryStatistics | None:
 
     Args:
         region: Region of the dataset to write to.
+        merge: If True, read-modify-write this chunk region: start from the chunk's current
+            contents (preserving traces already written by another shard) instead of the
+            fill value, then overlay this shard's live traces. Used for multi-shard
+            consolidation where a chunk is shared across shards. Defaults to False, which
+            preserves the original fast pure-write behavior.
 
     Returns:
         SummaryStatistics object containing statistics about the written traces.
@@ -154,25 +159,34 @@ def trace_worker(region: dict[str, slice]) -> SummaryStatistics | None:
     # Compute slices once (headers exclude sample dimension)
     header_region_slices = region_slices[:-1]  # Exclude sample dimension
 
-    full_shape = tuple(s.stop - s.start for s in region_slices)
-    header_shape = tuple(s.stop - s.start for s in header_region_slices)
-
     # Write raw headers if array was provided
     # Headers only have spatial dimensions (no sample dimension)
     if raw_header_array is not None:
-        tmp_raw_headers = np.full(header_shape, raw_header_array.fill_value)
+        if merge:
+            tmp_raw_headers = np.asarray(raw_header_array[header_region_slices])
+        else:
+            header_shape = tuple(s.stop - s.start for s in header_region_slices)
+            tmp_raw_headers = np.full(header_shape, raw_header_array.fill_value)
         tmp_raw_headers[not_null] = traces.raw_header
         raw_header_array[header_region_slices] = tmp_raw_headers
 
     # Write headers if array was provided
     # Headers only have spatial dimensions (no sample dimension)
     if header_array is not None:
-        tmp_headers = np.full(header_shape, header_array.fill_value)
+        if merge:
+            tmp_headers = np.asarray(header_array[header_region_slices])
+        else:
+            header_shape = tuple(s.stop - s.start for s in header_region_slices)
+            tmp_headers = np.full(header_shape, header_array.fill_value)
         tmp_headers[not_null] = traces.header
         header_array[header_region_slices] = tmp_headers
 
     # Write the data variable
-    tmp_samples = np.full(full_shape, data_array.fill_value)
+    if merge:
+        tmp_samples = np.asarray(data_array[region_slices])
+    else:
+        full_shape = tuple(s.stop - s.start for s in region_slices)
+        tmp_samples = np.full(full_shape, data_array.fill_value)
     tmp_samples[not_null] = traces.sample
     data_array[region_slices] = tmp_samples
 
diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py
index 2eddef4cf..68a211438 100644
--- a/src/mdio/segy/blocked_io.py
+++ b/src/mdio/segy/blocked_io.py
@@ -57,6 +57,7 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     grid_map: zarr_Array,
     dataset: xr_Dataset,
     data_variable_name: str,
+    merge_chunks: set[tuple[int, ...]] | None = None,
 ) -> SummaryStatistics:
     """Blocked I/O from SEG-Y to chunked `xarray.Dataset`.
 
@@ -66,6 +67,11 @@ def to_zarr(  # noqa: PLR0913, PLR0915
         grid_map: Zarr array with grid map for the traces.
         dataset: Handle for xarray.Dataset we are writing trace data
         data_variable_name: Name of the data variable in the dataset.
+        merge_chunks: Optional set of chunk-grid indices (one int per spatial dim) that must
+            be written in read-modify-write ("merge") mode instead of a pure fill write.
+            Used for multi-shard consolidation so shards sharing a boundary chunk don't
+            clobber each other. When None (default), every chunk uses the original fast
+            pure-write path, so single-file ingestion behavior is unchanged.
 
     Returns:
         None
@@ -81,6 +87,9 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     chunk_iter = ChunkIterator(shape=data.shape, chunks=worker_chunks, dim_names=data.dims)
     num_chunks = chunk_iter.num_chunks
 
+    # Spatial chunk sizes used to map a region back to its chunk-grid index for merge lookups.
+    spatial_chunk_sizes = data_variable_chunks[:-1]
+
     zarr_format = zarr.config.get("default_zarr_format")
     use_consolidated = zarr_format == ZarrFormat.V2
 
@@ -118,8 +127,15 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     with executor:
         futures = []
         for region in chunk_iter:
-            # Only the lightweight region is pickled per block; shared inputs live in worker state.
-            future = executor.submit(trace_worker, region)
+            merge = False
+            if merge_chunks:
+                region_slices = tuple(region.values())
+                chunk_index = tuple(
+                    int(region_slices[i].start // spatial_chunk_sizes[i]) for i in range(len(spatial_chunk_sizes))
+                )
+                merge = chunk_index in merge_chunks
+            # Only the lightweight region (+ merge flag) is pickled per block; shared inputs live in worker state.
+            future = executor.submit(trace_worker, region, merge)
             futures.append(future)
 
         iterable = tqdm(

From 438f72568115c656792d3d0b1421c13a25d0cd20 Mon Sep 17 00:00:00 2001
From: Ben Lasscock <ben.lasscock@tgs.com>
Date: Thu, 23 Jul 2026 00:53:23 +0000
Subject: [PATCH 2/3] Make gun/elevation optional coordinates and prune when
 absent

A shared template's "required" set should describe the data category, not one
contributor's acquisition rig. `gun` (StreamerShotGathers3D) and the new
elevation coords (ShotReceiverLineGathers3D) are now optional coordinates:
populated when the source carries them, silently omitted otherwise.

- base: add _optional_coord_names + optional_coordinate_names property
- validation: subtract optional coords from required-field check; add
  prune_absent_optional_coordinates to drop them from the resolved schema
  when the segy_spec doesn't carry them
- pipeline + consolidate: prune optional coords after schema resolution
- shot-receiver-line: add optional source/receiver elevation for true 3D
  land/areal geometry

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/mdio/builder/templates/base.py            | 13 +++++++++
 .../seismic_3d_shot_receiver_line.py          | 27 +++++++++++++++++++
 .../templates/seismic_3d_streamer_shot.py     |  4 +++
 src/mdio/ingestion/segy/consolidate.py        | 18 +++++++++----
 src/mdio/ingestion/segy/pipeline.py           |  2 ++
 src/mdio/ingestion/segy/validation.py         | 26 ++++++++++++++++++
 .../test_seismic_3d_shot_receiver_line.py     | 12 +++++++--
 7 files changed, 95 insertions(+), 7 deletions(-)

diff --git a/src/mdio/builder/templates/base.py b/src/mdio/builder/templates/base.py
index 48d382969..b865c075d 100644
--- a/src/mdio/builder/templates/base.py
+++ b/src/mdio/builder/templates/base.py
@@ -46,6 +46,7 @@ def __init__(self, data_domain: SeismicDataDomain) -> None:
         self._calculated_dims: tuple[str, ...] = ()
         self._physical_coord_names: tuple[str, ...] = ()
         self._logical_coord_names: tuple[str, ...] = ()
+        self._optional_coord_names: tuple[str, ...] = ()
         self._var_chunk_shape: tuple[int, ...] = ()
         self.synthesize_missing_dims: tuple[str, ...] = ()
 
@@ -303,6 +304,18 @@ def coordinate_names(self) -> tuple[str, ...]:
         """Returns names of all coordinates."""
         return self._physical_coord_names + self._logical_coord_names
 
+    @property
+    def optional_coordinate_names(self) -> tuple[str, ...]:
+        """Coordinates that are populated when present in the source but are NOT required.
+
+        Unlike required coordinates, a missing optional coordinate does not fail ingestion: it
+        is simply omitted from the built dataset (see
+        :func:`mdio.ingestion.segy.validation.prune_absent_optional_coordinates`). This lets a
+        template carry a coordinate that only *some* surveys record (e.g. ``gun`` on a
+        single-source streamer acquisition) without hard-requiring it of every file.
+        """
+        return self._optional_coord_names
+
     @property
     def full_chunk_shape(self) -> tuple[int, ...]:
         """Returns the chunk shape for the variables."""
diff --git a/src/mdio/builder/templates/seismic_3d_shot_receiver_line.py b/src/mdio/builder/templates/seismic_3d_shot_receiver_line.py
index 5f2c1ad78..a4a5f01f1 100644
--- a/src/mdio/builder/templates/seismic_3d_shot_receiver_line.py
+++ b/src/mdio/builder/templates/seismic_3d_shot_receiver_line.py
@@ -23,8 +23,17 @@ def __init__(self, data_domain: SeismicDataDomain = "time"):
             "source_coord_y",
             "group_coord_x",
             "group_coord_y",
+            "source_surface_elevation",
+            "receiver_group_elevation",
         )
         self._logical_coord_names = ("orig_field_record_num",)
+        # Elevation (z) makes the source/receiver layout a true 3-D geometry rather than a flat
+        # map. Land and OBN acquisitions record it; marine and many synthetics don't - so it is
+        # OPTIONAL: populated when the SEG-Y carries it, pruned otherwise (see
+        # AbstractDatasetTemplate.optional_coordinate_names). NB: elevations are NOT in
+        # SCALE_COORDINATE_KEYS, so MDIO stores them as read (any elevation_depth_scalar is not
+        # applied); the coordinate is the raw header value.
+        self._optional_coord_names = ("source_surface_elevation", "receiver_group_elevation")
         self._var_chunk_shape = (1, 32, 1, 32, 2048)
 
     @property
@@ -44,6 +53,9 @@ def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
             CoordinateSpec(name="group_coord_x", dimensions=group_dims, dtype=ScalarType.FLOAT64),
             CoordinateSpec(name="group_coord_y", dimensions=group_dims, dtype=ScalarType.FLOAT64),
             CoordinateSpec(name="orig_field_record_num", dimensions=source_dims, dtype=ScalarType.UINT32),
+            # Optional z (see __init__); pruned by the ingestion path when the SEG-Y lacks it.
+            CoordinateSpec(name="source_surface_elevation", dimensions=source_dims, dtype=ScalarType.FLOAT64),
+            CoordinateSpec(name="receiver_group_elevation", dimensions=group_dims, dtype=ScalarType.FLOAT64),
         )
 
     def declare_dim_coordinate_types(self) -> DimCoordinateTypes:
@@ -115,3 +127,18 @@ def _add_coordinates(self) -> None:
             dimensions=("shot_line", "shot_point"),
             data_type=ScalarType.UINT32,
         )
+
+        # Optional elevation (z) coordinates - see __init__. Declared so they materialize when
+        # present; the ingestion path prunes them when the source has no such header field.
+        self._builder.add_coordinate(
+            "source_surface_elevation",
+            dimensions=("shot_line", "shot_point"),
+            data_type=ScalarType.FLOAT64,
+            metadata=CoordinateMetadata(units_v1=self.get_unit_by_key("source_surface_elevation")),
+        )
+        self._builder.add_coordinate(
+            "receiver_group_elevation",
+            dimensions=("receiver_line", "receiver"),
+            data_type=ScalarType.FLOAT64,
+            metadata=CoordinateMetadata(units_v1=self.get_unit_by_key("receiver_group_elevation")),
+        )
diff --git a/src/mdio/builder/templates/seismic_3d_streamer_shot.py b/src/mdio/builder/templates/seismic_3d_streamer_shot.py
index 744c5fe65..be58099de 100644
--- a/src/mdio/builder/templates/seismic_3d_streamer_shot.py
+++ b/src/mdio/builder/templates/seismic_3d_streamer_shot.py
@@ -19,6 +19,10 @@ def __init__(self, data_domain: SeismicDataDomain = "time"):
         self._dim_names = ("shot_point", "cable", "channel", self._data_domain)
         self._physical_coord_names = ("source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y")
         self._logical_coord_names = ("gun",)
+        # 'gun' is recorded only by some acquisitions (multi-source). Keep it as a coordinate
+        # populated when the source carries it, but don't hard-require it: single-source
+        # streamer shot data has no gun field and must still ingest as 3D.
+        self._optional_coord_names = ("gun",)
         self._var_chunk_shape = (8, 1, 128, 2048)
 
     @property
diff --git a/src/mdio/ingestion/segy/consolidate.py b/src/mdio/ingestion/segy/consolidate.py
index 82752add5..e860764b5 100644
--- a/src/mdio/ingestion/segy/consolidate.py
+++ b/src/mdio/ingestion/segy/consolidate.py
@@ -50,6 +50,7 @@
 from mdio.ingestion.segy.index_strategies import IndexStrategyRegistry
 from mdio.ingestion.segy.raw_headers import build_raw_header_variables
 from mdio.ingestion.segy.reader import read_index_headers
+from mdio.ingestion.segy.validation import prune_absent_optional_coordinates
 from mdio.ingestion.segy.validation import validate_spec_in_template
 from mdio.segy import blocked_io
 from mdio.segy.file import get_segy_file_info
@@ -71,10 +72,17 @@
 logger = logging.getLogger(__name__)
 
 
-def _resolve_schema(mdio_template: AbstractDatasetTemplate, grid_overrides: GridOverrides | None):
-    """Resolve the (format-agnostic) schema for a template + optional grid overrides."""
+def _resolve_schema(
+    mdio_template: AbstractDatasetTemplate, grid_overrides: GridOverrides | None, segy_spec: SegySpec
+):
+    """Resolve the (format-agnostic) schema for a template + optional grid overrides.
+
+    Optional coordinates absent from ``segy_spec`` are pruned so they aren't built as empty
+    vars (mirrors the single-file path in :mod:`mdio.ingestion.segy.pipeline`).
+    """
     schema_effect = IndexStrategyRegistry().schema_effect(grid_overrides)
-    return SchemaResolver().resolve(mdio_template, schema_effect)
+    schema = SchemaResolver().resolve(mdio_template, schema_effect)
+    return prune_absent_optional_coordinates(schema, segy_spec, mdio_template)
 
 
 def allocate_mdio_grid(  # noqa: PLR0913
@@ -127,7 +135,7 @@ def allocate_mdio_grid(  # noqa: PLR0913
     file_info = get_segy_file_info(ref_kwargs)
     units = resolve_units(mdio_template, get_spatial_coordinate_unit(file_info))
 
-    schema = _resolve_schema(mdio_template, grid_overrides)
+    schema = _resolve_schema(mdio_template, grid_overrides, segy_spec)
 
     grid = Grid(dims=list(global_dimensions))
     header_dtype = build_mdio_header_type(segy_spec)
@@ -229,7 +237,7 @@ def append_segy_shard(  # noqa: PLR0913
     }
     file_info = get_segy_file_info(segy_file_kwargs)
     units = resolve_units(mdio_template, get_spatial_coordinate_unit(file_info))
-    schema = _resolve_schema(mdio_template, grid_overrides)
+    schema = _resolve_schema(mdio_template, grid_overrides, segy_spec)
 
     indexed_headers, _shard_dims = read_index_headers(
         segy_file_kwargs=segy_file_kwargs,
diff --git a/src/mdio/ingestion/segy/pipeline.py b/src/mdio/ingestion/segy/pipeline.py
index b2c54fabc..129337180 100644
--- a/src/mdio/ingestion/segy/pipeline.py
+++ b/src/mdio/ingestion/segy/pipeline.py
@@ -21,6 +21,7 @@
 from mdio.ingestion.segy.raw_headers import build_raw_header_variables
 from mdio.ingestion.segy.reader import read_index_headers
 from mdio.ingestion.segy.serializer import serialize_to_mdio
+from mdio.ingestion.segy.validation import prune_absent_optional_coordinates
 from mdio.ingestion.segy.validation import validate_spec_in_template
 from mdio.segy.file import get_segy_file_info
 from mdio.segy.geometry import validate_overrides_for_template
@@ -160,6 +161,7 @@ def segy_to_mdio(  # noqa: PLR0913
     # format-agnostic resolver then applies.
     schema_effect = IndexStrategyRegistry().schema_effect(grid_overrides)
     schema = SchemaResolver().resolve(mdio_template, schema_effect)
+    schema = prune_absent_optional_coordinates(schema, segy_spec, mdio_template)
 
     indexed_headers, dimensions = read_index_headers(
         segy_file_kwargs=segy_file_kwargs,
diff --git a/src/mdio/ingestion/segy/validation.py b/src/mdio/ingestion/segy/validation.py
index 3828fda36..0424c78a5 100644
--- a/src/mdio/ingestion/segy/validation.py
+++ b/src/mdio/ingestion/segy/validation.py
@@ -10,6 +10,7 @@
     from segy.schema import SegySpec
 
     from mdio.builder.templates.base import AbstractDatasetTemplate
+    from mdio.ingestion.schema.models import ResolvedSchema
 
 
 def validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatasetTemplate) -> None:
@@ -26,6 +27,10 @@ def validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatase
     if isinstance(mdio_template, Seismic3DObnReceiverGathersTemplate):
         required_fields.discard("component")
 
+    # Optional coordinates (e.g. 'gun' on streamer shot gathers) are populated only when the
+    # source carries them; their absence must not fail ingestion.
+    required_fields -= set(mdio_template.optional_coordinate_names)
+
     if any(field in SCALE_COORDINATE_KEYS for field in required_fields):
         required_fields = required_fields | {"coordinate_scalar"}
     missing_fields = required_fields - header_fields
@@ -36,3 +41,24 @@ def validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatase
             f"not found in the provided segy_spec"
         )
         raise ValueError(err)
+
+
+def prune_absent_optional_coordinates(
+    schema: ResolvedSchema, segy_spec: SegySpec, mdio_template: AbstractDatasetTemplate
+) -> ResolvedSchema:
+    """Drop optional coordinates the SEG-Y doesn't carry, so they aren't built as empty vars.
+
+    An optional coordinate (see :attr:`AbstractDatasetTemplate.optional_coordinate_names`) is
+    declared on the template so it can be populated *when present*, but if the source lacks the
+    header field it would otherwise be materialized as an all-fill coordinate. Pruning it here
+    keeps the built dataset honest: the coordinate exists iff the data actually supplied it.
+    Returns ``schema`` unchanged when there is nothing to prune.
+    """
+    optional = set(mdio_template.optional_coordinate_names)
+    if not optional:
+        return schema
+    header_fields = {field.name for field in segy_spec.trace.header.fields}
+    kept = [c for c in schema.coordinates if c.name not in optional or c.name in header_fields]
+    if len(kept) == len(schema.coordinates):
+        return schema
+    return schema.model_copy(update={"coordinates": kept})
diff --git a/tests/unit/v1/templates/test_seismic_3d_shot_receiver_line.py b/tests/unit/v1/templates/test_seismic_3d_shot_receiver_line.py
index 34fab4948..a567d7f3c 100644
--- a/tests/unit/v1/templates/test_seismic_3d_shot_receiver_line.py
+++ b/tests/unit/v1/templates/test_seismic_3d_shot_receiver_line.py
@@ -37,14 +37,19 @@
     "group_coord_x",
     "group_coord_y",
     "orig_field_record_num",
+    # Optional 3-D geometry (z); built here because build_dataset materializes every declared
+    # coordinate. The SEG-Y ingestion path prunes these when the source lacks elevation.
+    "source_surface_elevation",
+    "receiver_group_elevation",
 ]
 
 
 def _validate_coordinates_headers_trace_mask(dataset: Dataset, headers: StructuredType, domain: str) -> None:
     """Validate the coordinate, headers, trace_mask variables in the dataset."""
     # Verify variables
-    # 5 dim coords + 5 non-dim coords + 1 data + 1 trace mask + 1 headers = 13 variables
-    assert len(dataset.variables) == 13
+    # 5 dim coords + 7 non-dim coords + 1 data + 1 trace mask + 1 headers = 15 variables
+    # (non-dim coords now include the two optional elevation coordinates)
+    assert len(dataset.variables) == 15
 
     # Verify trace headers
     validate_variable(
@@ -112,8 +117,11 @@ def test_configuration(self) -> None:
             "source_coord_y",
             "group_coord_x",
             "group_coord_y",
+            "source_surface_elevation",
+            "receiver_group_elevation",
         )
         assert t._logical_coord_names == ("orig_field_record_num",)
+        assert t.optional_coordinate_names == ("source_surface_elevation", "receiver_group_elevation")
         assert t._var_chunk_shape == (1, 32, 1, 32, 2048)
 
         assert t._builder is None

From 3e3d27d543306e34bbb13850c2a75c79d8018e88 Mon Sep 17 00:00:00 2001
From: Ben Lasscock <ben.lasscock@tgs.com>
Date: Thu, 23 Jul 2026 21:30:08 +0000
Subject: [PATCH 3/3] Add Zarr v3 sharding support to the ingestion build path

Introduce an optional shard grid so many chunks pack into one storage object
(far fewer objects / S3 GETs) while keeping the small chunk as the partial-read
unit. Adds shard_grid to VariableMetadata and shard_shape to ResolvedSchema,
a full_shard_shape property on the template base, resolver pass-through, and
shard resolution/validation (whole multiple of the chunk per dimension) in the
dataset factory. The xarray builder emits the sharding_indexed codec (v3 only)
and aligns Dask lazy blocks to shard boundaries so to_zarr's safe-chunks guard
passes; blocked_io writes whole shards to avoid read-modify-write races.

Headers are intentionally left unsharded: their structured (void) dtype trips
Zarr's sharding partial-encode path, and they are tiny relative to the sample
cube. Sharding applies to the data variable.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/mdio/builder/schemas/v1/variable.py | 10 ++++
 src/mdio/builder/templates/base.py      | 34 ++++++++++++++
 src/mdio/builder/xarray_builder.py      | 24 +++++++++-
 src/mdio/ingestion/dataset_factory.py   | 61 ++++++++++++++++++++++++-
 src/mdio/ingestion/schema/models.py     |  5 +-
 src/mdio/ingestion/schema/resolver.py   |  1 +
 src/mdio/segy/blocked_io.py             | 15 ++++--
 7 files changed, 143 insertions(+), 7 deletions(-)

diff --git a/src/mdio/builder/schemas/v1/variable.py b/src/mdio/builder/schemas/v1/variable.py
index ed317f8b4..2f8abf1ad 100644
--- a/src/mdio/builder/schemas/v1/variable.py
+++ b/src/mdio/builder/schemas/v1/variable.py
@@ -39,6 +39,16 @@ class VariableMetadata(CoordinateMetadata):
         description="Chunk grid specification for the array.",
     )
 
+    shard_grid: RegularChunkGrid | None = Field(
+        default=None,
+        description=(
+            "Optional Zarr v3 shard grid: the storage-object (write) unit. When set, the array is "
+            "stored with the sharding_indexed codec - each shard is a single object holding a grid "
+            "of chunks (the read unit given by ``chunk_grid``). The shard shape must be a whole "
+            "multiple of the chunk shape along every dimension. Ignored for Zarr v2."
+        ),
+    )
+
     stats_v1: SummaryStatistics | list[SummaryStatistics] | None = Field(
         default=None,
         description="Minimal summary statistics.",
diff --git a/src/mdio/builder/templates/base.py b/src/mdio/builder/templates/base.py
index b865c075d..e142a0410 100644
--- a/src/mdio/builder/templates/base.py
+++ b/src/mdio/builder/templates/base.py
@@ -48,6 +48,7 @@ def __init__(self, data_domain: SeismicDataDomain) -> None:
         self._logical_coord_names: tuple[str, ...] = ()
         self._optional_coord_names: tuple[str, ...] = ()
         self._var_chunk_shape: tuple[int, ...] = ()
+        self._var_shard_shape: tuple[int, ...] = ()
         self.synthesize_missing_dims: tuple[str, ...] = ()
 
         self._builder: MDIODatasetBuilder | None = None
@@ -344,6 +345,39 @@ def full_chunk_shape(self, shape: tuple[int, ...]) -> None:
 
         self._var_chunk_shape = shape
 
+    @property
+    def full_shard_shape(self) -> tuple[int, ...]:
+        """Returns the shard (storage-object) shape, or ``()`` when sharding is disabled."""
+        if not self._var_shard_shape:
+            return ()
+        if len(self._dim_sizes) != len(self._dim_names):
+            return self._var_shard_shape
+        return tuple(
+            dim_size if shard_size == -1 else shard_size
+            for shard_size, dim_size in zip(self._var_shard_shape, self._dim_sizes, strict=False)
+        )
+
+    @full_shard_shape.setter
+    def full_shard_shape(self, shape: tuple[int, ...]) -> None:
+        """Sets the shard shape for the variables (``()`` disables sharding).
+
+        The shard is the Zarr v3 storage-object/write unit; it must have the same rank as the
+        chunk shape and every entry must be a positive integer or ``-1`` (full dimension). The
+        multiple-of-chunk-shape constraint is enforced downstream at build time, once ``-1``
+        placeholders are resolved against the real dimension sizes.
+        """
+        if not shape:
+            self._var_shard_shape = ()
+            return
+        if len(shape) != len(self._dim_names):
+            msg = f"Shard shape {shape} has {len(shape)} dimensions, expected {len(self._dim_names)}"
+            raise ValueError(msg)
+        for shard_size in shape:
+            if shard_size != -1 and shard_size <= 0:
+                msg = f"Shard size must be positive integer or -1, got {shard_size}"
+                raise ValueError(msg)
+        self._var_shard_shape = shape
+
     @property
     @abstractmethod
     def _name(self) -> str:
diff --git a/src/mdio/builder/xarray_builder.py b/src/mdio/builder/xarray_builder.py
index 58e1f5834..01a4db9c2 100644
--- a/src/mdio/builder/xarray_builder.py
+++ b/src/mdio/builder/xarray_builder.py
@@ -123,6 +123,13 @@ def _get_zarr_chunks(var: Variable, all_named_dims: dict[str, NamedDimension]) -
     return _get_zarr_shape(var, all_named_dims=all_named_dims)
 
 
+def _get_zarr_shards(var: Variable) -> tuple[int, ...] | None:
+    """Get the shard (storage-object) shape for a variable, or None when it isn't sharded."""
+    if var.metadata is not None and getattr(var.metadata, "shard_grid", None) is not None:
+        return var.metadata.shard_grid.configuration.chunk_shape
+    return None
+
+
 def _compressor_to_encoding(
     compressor: mdio_Blosc | mdio_ZFP | None,
 ) -> dict[str, "zarr.codecs.Blosc | numcodecs.Blosc | numcodecs.ZFPY | zarr.codecs.ZFPY | None"] | None:
@@ -197,11 +204,17 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_Dataset:  # noqa: PLR0912
         shape = _get_zarr_shape(v, all_named_dims=all_named_dims)
         dtype = to_numpy_dtype(v.data_type)
         original_chunks = _get_zarr_chunks(v, all_named_dims=all_named_dims)
+        shards = _get_zarr_shards(v)
 
         # For efficient lazy array creation with Dask use larger chunks to minimize the task graph size
         # Initialize with original chunks for lazy array creation
         lazy_chunks = original_chunks
-        if shape != original_chunks:
+        if shards is not None:
+            # Sharded: the shard is the atomic write unit (one object). Dask blocks MUST align to
+            # shard boundaries or xarray's to_zarr safe-chunks guard rejects the parallel write
+            # ("would overlap multiple Dask chunks"). Making each lazy block one shard aligns them.
+            lazy_chunks = tuple(shards)
+        elif shape != original_chunks:
             # Compute automatic chunk sizes based on heuristics, respecting original chunks where possible
             auto_chunks = normalize_chunks("auto", shape=shape, dtype=dtype, previous_chunks=original_chunks)
 
@@ -219,7 +232,9 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_Dataset:  # noqa: PLR0912
 
         # Add array attributes
         if v.metadata is not None:
-            metadata_dict = v.metadata.model_dump(exclude_none=True, mode="json", exclude={"chunk_grid"})
+            metadata_dict = v.metadata.model_dump(
+                exclude_none=True, mode="json", exclude={"chunk_grid", "shard_grid"}
+            )
             data_array.attrs.update(metadata_dict)
         if v.long_name:
             data_array.attrs["long_name"] = v.long_name
@@ -233,6 +248,11 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_Dataset:  # noqa: PLR0912
             fill_value_key: fill_value,
         }
 
+        # Zarr v3 sharding: the shard is the storage object; chunks are the read unit within it.
+        # Sharding is a v3-only codec, so it is silently ignored for v2 stores.
+        if shards is not None and zarr_format != ZarrFormat.V2:
+            encoding["shards"] = shards
+
         compressor_encodings = _compressor_to_encoding(v.compressor)
 
         if compressor_encodings is not None:
diff --git a/src/mdio/ingestion/dataset_factory.py b/src/mdio/ingestion/dataset_factory.py
index 269185577..3398fb7d6 100644
--- a/src/mdio/ingestion/dataset_factory.py
+++ b/src/mdio/ingestion/dataset_factory.py
@@ -64,6 +64,53 @@ def _resolve_chunks(chunk_shape: tuple[int, ...], sizes: tuple[int, ...]) -> tup
     return tuple(size if chunk_size == -1 else chunk_size for chunk_size, size in zip(chunk_shape, sizes, strict=True))
 
 
+def _resolve_shards(
+    shard_shape: tuple[int, ...],
+    chunk_shape: tuple[int, ...],
+    sizes: tuple[int, ...],
+) -> tuple[int, ...]:
+    """Resolve the shard shape (``-1`` -> full size) and validate it against the chunk shape.
+
+    An empty ``shard_shape`` means sharding is disabled and returns ``()``. Otherwise the shard
+    is the Zarr v3 storage-object unit and must be a whole multiple of the (resolved) chunk shape
+    along every dimension, so an integer number of chunks packs into each shard.
+
+    Args:
+        shard_shape: Configured shard shape (may contain ``-1``); ``()`` disables sharding.
+        chunk_shape: Already-resolved chunk shape (no ``-1``).
+        sizes: Actual sizes of each dimension.
+
+    Returns:
+        The resolved shard shape, or ``()`` when sharding is disabled.
+
+    Raises:
+        ValueError: If the shard rank differs from the chunk rank, or a shard extent is not a
+            positive whole multiple of the corresponding chunk extent.
+    """
+    if not shard_shape:
+        return ()
+    if len(shard_shape) != len(chunk_shape):
+        msg = f"Shard shape {shard_shape} rank does not match chunk shape {chunk_shape}."
+        raise ValueError(msg)
+    resolved = tuple(size if s == -1 else s for s, size in zip(shard_shape, sizes, strict=True))
+    for shard_size, chunk_size in zip(resolved, chunk_shape, strict=True):
+        if shard_size <= 0 or shard_size % chunk_size != 0:
+            msg = (
+                f"Shard shape {resolved} must be a positive whole multiple of the chunk shape "
+                f"{chunk_shape} along every dimension (offending pair: shard={shard_size}, "
+                f"chunk={chunk_size})."
+            )
+            raise ValueError(msg)
+    return resolved
+
+
+def _shard_grid(shard_shape: tuple[int, ...]) -> RegularChunkGrid | None:
+    """Build a shard grid model from a resolved shard shape (``()`` -> ``None``)."""
+    if not shard_shape:
+        return None
+    return RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=shard_shape))
+
+
 def _create_dataset_builder(schema: ResolvedSchema) -> MDIODatasetBuilder:
     """Create and initialize the MDIODatasetBuilder with attributes.
 
@@ -145,6 +192,10 @@ def _add_trace_mask_and_headers(
 
     if header_dtype is not None:
         chunk_grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=resolved_chunks[:-1]))
+        # NOTE: headers are intentionally NOT sharded. They use a structured (void) dtype, and
+        # Zarr's sharding partial-encode path tries to hash the void fill value
+        # (TypeError: unhashable 'writeable void-scalar'). Headers are also tiny relative to the
+        # sample cube, so sharding them has little upside. Sharding applies to the data variable.
         builder.add_variable(
             name="headers",
             dimensions=spatial_dim_names,
@@ -159,6 +210,7 @@ def _add_main_and_extra_variables(
     builder: MDIODatasetBuilder,
     schema: ResolvedSchema,
     resolved_chunks: tuple[int, ...],
+    resolved_shards: tuple[int, ...],
     units: dict[str, AllUnitModel],
     extra_variables: list[dict[str, Any]],
 ) -> None:
@@ -168,6 +220,7 @@ def _add_main_and_extra_variables(
         builder: MDIO dataset builder.
         schema: Resolved schema.
         resolved_chunks: Resolved chunk shapes.
+        resolved_shards: Resolved shard shapes (``()`` disables sharding).
         units: Dictionary mapping coordinate/dimension names to AllUnitModel.
         extra_variables: Optional list of additional variables.
     """
@@ -181,7 +234,11 @@ def _add_main_and_extra_variables(
         data_type=ScalarType.FLOAT32,
         compressor=compressor,
         coordinates=coordinate_names,
-        metadata=VariableMetadata(chunk_grid=chunk_grid, units_v1=units.get(schema.default_variable_name)),
+        metadata=VariableMetadata(
+            chunk_grid=chunk_grid,
+            shard_grid=_shard_grid(resolved_shards),
+            units_v1=units.get(schema.default_variable_name),
+        ),
     )
 
     for var_dict in extra_variables:
@@ -223,6 +280,7 @@ def build_mdio_dataset(
     extra_variables = extra_variables or []
 
     resolved_chunks = _resolve_chunks(schema.chunk_shape, sizes)
+    resolved_shards = _resolve_shards(schema.shard_shape, resolved_chunks, sizes)
     builder = _create_dataset_builder(schema)
 
     _add_dimensions_and_coordinates(
@@ -243,6 +301,7 @@ def build_mdio_dataset(
         builder=builder,
         schema=schema,
         resolved_chunks=resolved_chunks,
+        resolved_shards=resolved_shards,
         units=units,
         extra_variables=extra_variables,
     )
diff --git a/src/mdio/ingestion/schema/models.py b/src/mdio/ingestion/schema/models.py
index ad4adaa06..0bf745db3 100644
--- a/src/mdio/ingestion/schema/models.py
+++ b/src/mdio/ingestion/schema/models.py
@@ -40,7 +40,9 @@ class ResolvedSchema(BaseModel):
         name: Name of the dataset or template.
         dimensions: Specifications for the dimensions.
         coordinates: Specifications for the coordinates.
-        chunk_shape: Chunk size for each dimension.
+        chunk_shape: Chunk (read-unit) size for each dimension.
+        shard_shape: Optional shard (storage-object/write-unit) size for each dimension. Empty
+            means no sharding. When set, must be a whole multiple of ``chunk_shape`` per dimension.
         metadata: Metadata attributes.
         default_variable_name: Name of the primary data variable.
     """
@@ -49,6 +51,7 @@ class ResolvedSchema(BaseModel):
     dimensions: list[DimensionSpec]
     coordinates: list[CoordinateSpec]
     chunk_shape: tuple[int, ...]
+    shard_shape: tuple[int, ...] = ()
     metadata: dict[str, Any] = Field(default_factory=dict)
     default_variable_name: str = "amplitude"
 
diff --git a/src/mdio/ingestion/schema/resolver.py b/src/mdio/ingestion/schema/resolver.py
index 320027dc9..08998e1c0 100644
--- a/src/mdio/ingestion/schema/resolver.py
+++ b/src/mdio/ingestion/schema/resolver.py
@@ -71,6 +71,7 @@ def _template_to_schema(self, template: AbstractDatasetTemplate) -> ResolvedSche
             dimensions=dimensions,
             coordinates=list(template.declare_coordinate_specs()),
             chunk_shape=template.full_chunk_shape,
+            shard_shape=template.full_shard_shape,
             metadata=template._load_dataset_attributes() or {},
             default_variable_name=template.default_variable_name,
         )
diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py
index 68a211438..86a63744f 100644
--- a/src/mdio/segy/blocked_io.py
+++ b/src/mdio/segy/blocked_io.py
@@ -83,12 +83,21 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     final_stats = _create_stats()
 
     data_variable_chunks = data.encoding.get("chunks")
-    worker_chunks = data_variable_chunks[:-1] + (data.shape[-1],)  # un-chunk sample axis
+
+    # Write-block granularity. With Zarr v3 sharding, a shard is a single storage object holding
+    # a grid of chunks; two workers writing different chunks of the *same* shard would perform
+    # concurrent read-modify-write on that object and clobber each other. So when the array is
+    # sharded we make the write block one whole shard (spatial), guaranteeing each worker owns a
+    # distinct set of shard objects. Unsharded arrays keep the original per-chunk write block.
+    shard_shape = data.encoding.get("shards")
+    write_block = tuple(shard_shape) if shard_shape else data_variable_chunks
+
+    worker_chunks = write_block[:-1] + (data.shape[-1],)  # un-chunk sample axis
     chunk_iter = ChunkIterator(shape=data.shape, chunks=worker_chunks, dim_names=data.dims)
     num_chunks = chunk_iter.num_chunks
 
-    # Spatial chunk sizes used to map a region back to its chunk-grid index for merge lookups.
-    spatial_chunk_sizes = data_variable_chunks[:-1]
+    # Spatial write-block sizes used to map a region back to its block-grid index for merge lookups.
+    spatial_chunk_sizes = write_block[:-1]
 
     zarr_format = zarr.config.get("default_zarr_format")
     use_consolidated = zarr_format == ZarrFormat.V2