Multi shard consolidation

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",

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.",

src/mdio/builder/templates/base.py

@@ -46,7 +46,9 @@ 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._var_shard_shape: tuple[int, ...] = ()
         self.synthesize_missing_dims: tuple[str, ...] = ()
 
         self._builder: MDIODatasetBuilder | None = None
@@ -303,6 +305,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."""
@@ -331,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:

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")),
+        )

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

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:

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"]

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,
+    )