Begin investigation for sharded ingestion

Author: BrianMichell

docs/data_models/chunk_grids.md

@@ -17,13 +17,14 @@
 
 The variables in MDIO data model can represent different types of chunk grids.
 These grids are essential for managing multi-dimensional data arrays efficiently.
-In this breakdown, we will explore four distinct data models within the MDIO schema,
+In this breakdown, we will explore the chunk grid data models within the MDIO schema,
 each serving a specific purpose in data handling and organization.
 
 MDIO implements data models following the guidelines of the Zarr v3 spec and ZEPs:
 
 - [Zarr core specification (version 3)](https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html)
 - [ZEP 1 — Zarr specification version 3](https://zarr.dev/zeps/accepted/ZEP0001.html)
+- [ZEP 2 — Sharding codec](https://zarr.dev/zeps/accepted/ZEP0002.html)
 - [ZEP 3 — Variable chunking](https://zarr.dev/zeps/draft/ZEP0003.html)
 
 ## Regular Grid
@@ -126,6 +127,100 @@ conceptual and visually not to scale.
 └──────────┴─────┴───────┴───┘
 ```
 
+## Sharded Grid
+
+The [ShardedChunkGrid](ShardedChunkGrid) model implements Zarr v3 sharding, which
+organizes data in a two-level hierarchy for improved I/O performance:
+
+- **Shards**: Outer containers stored as individual files/objects
+- **Chunks**: Inner units within each shard for compression and fine-grained access
+
+Sharding is particularly beneficial for cloud storage where reducing the number of
+objects improves performance while still maintaining efficient partial reads.
+
+```{eval-rst}
+.. autosummary::
+   ShardedChunkGrid
+   ShardedChunkShape
+```
+
+:::{note}
+Sharding is only supported with Zarr v3 format. The `shardShape` must be evenly
+divisible by `chunkShape` along all dimensions.
+:::
+
+:::{important}
+**Non-shardable dtype limitation**: Sharding is only applied to simple scalar-type
+variables (e.g., `amplitude`, `cdp_x`, `cdp_y`). Variables with structured dtypes
+(e.g., `headers`) or void/bytes types (e.g., `raw_headers`) will automatically use
+regular chunking instead, as Zarr's sharding codec does not support these dtypes.
+However, these variables will use the **shard shape as their chunk shape** to maintain
+consistent I/O patterns across all variables.
+:::
+
+For a 2D array with shape `rows, cols = (256, 512)`{l=python}, we can configure
+sharding with shard size `(128, 128)` and chunk size `(32, 32)`:
+
+`{ "name": "sharded", "configuration": { "shardShape": [128, 128], "chunkShape": [32, 32] } }`{l=json}
+
+This creates a two-level structure where each shard contains multiple chunks:
+
+```bash
+ ←────────── 128 ──────────→ ←────────── 128 ──────────→ ...
+┌─────────────────────────────────────────────────────────┐
+│ ┌─────┬─────┬─────┬─────┐ ┌─────┬─────┬─────┬─────┐     │  ↑
+│ │ 32  │ 32  │ 32  │ 32  │ │ 32  │ 32  │ 32  │ 32  │     │  │
+│ ├─────┼─────┼─────┼─────┤ ├─────┼─────┼─────┼─────┤     │  │
+│ │ 32  │ 32  │ 32  │ 32  │ │ 32  │ 32  │ 32  │ 32  │     │  128
+│ ├─────┼─────┼─────┼─────┤ ├─────┼─────┼─────┼─────┤     │  │
+│ │ 32  │ 32  │ 32  │ 32  │ │ 32  │ 32  │ 32  │ 32  │     │  │
+│ ├─────┼─────┼─────┼─────┤ ├─────┼─────┼─────┼─────┤     │  │
+│ │ 32  │ 32  │ 32  │ 32  │ │ 32  │ 32  │ 32  │ 32  │     │  ↓
+│ └─────┴─────┴─────┴─────┘ └─────┴─────┴─────┴─────┘     │
+│        Shard (0,0)               Shard (0,1)            │
+├─────────────────────────────────────────────────────────┤
+│        Shard (1,0)               Shard (1,1)       ...  │
+└─────────────────────────────────────────────────────────┘
+```
+
+### Using Sharding with Templates
+
+By default, MDIO templates do not use sharding. To enable sharding on a template,
+set both `full_chunk_shape` and `full_shard_shape`:
+
+```python
+from mdio.builder.templates.seismic_3d_poststack import Seismic3DPostStackTemplate
+
+# Create template with default chunking (no sharding)
+template = Seismic3DPostStackTemplate(data_domain='depth')
+print(template.full_shard_shape)  # None - sharding disabled
+
+# Enable sharding by setting both chunk and shard shapes
+template.full_chunk_shape = (32, 32, 128)    # Inner chunk sizes
+template.full_shard_shape = (128, 128, 512)  # Outer shard sizes
+
+# Build dataset - will use ShardedChunkGrid for scalar variables (amplitude, coordinates)
+# Note: headers array will use RegularChunkGrid (structured dtypes don't support sharding)
+dataset = template.build_dataset('Seismic3D', sizes=(256, 512, 1024))
+```
+
+### Ingestion with Sharding
+
+When sharding is enabled, the SEG-Y ingestion process automatically iterates at the
+shard level instead of the chunk level. This reduces the number of I/O operations
+and improves performance, especially for cloud storage backends.
+
+- **With sharding**: Ingestion iterates over shards (larger units)
+- **Without sharding**: Ingestion iterates over chunks (default behavior)
+
+During ingestion, sharding is applied only to simple scalar-type arrays (e.g., `amplitude`).
+Arrays with non-shardable dtypes use regular chunking with the shard shape as their chunk
+size. For example, with `full_chunk_shape=(32, 32, 32)` and `full_shard_shape=(256, 256, 256)`:
+
+- `amplitude` array: chunks=(32, 32, 32), shards=(256, 256, 256)
+- `headers` array (structured dtype): chunks=(256, 256), no sharding
+- `raw_headers` array (void/bytes dtype): chunks=(256, 256), no sharding
+
 ## Model Reference
 
 :::{dropdown} RegularChunkGrid
@@ -152,3 +247,15 @@ conceptual and visually not to scale.
 ```
 
 :::
+:::{dropdown} ShardedChunkGrid
+:animate: fade-in-slide-down
+
+```{eval-rst}
+.. autopydantic_model:: ShardedChunkGrid
+
+----------
+
+.. autopydantic_model:: ShardedChunkShape
+```
+
+:::

src/mdio/builder/schemas/chunk_grid.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pydantic import Field
+from pydantic import model_validator
 
 from mdio.builder.schemas.core import CamelCaseStrictModel
 
@@ -22,6 +23,53 @@ class RectilinearChunkShape(CamelCaseStrictModel):
     )
 
 
+class ShardedChunkShape(CamelCaseStrictModel):
+    """Represents sharded chunk configuration for Zarr v3.
+
+    In Zarr v3 sharding, data is organized in a two-level hierarchy:
+    - Shards: Outer containers that are stored as individual files/objects
+    - Chunks: Inner units within each shard for compression and access
+
+    The shard_shape must be evenly divisible by chunk_shape along each dimension.
+    """
+
+    shard_shape: tuple[int, ...] = Field(
+        ...,
+        description="Lengths of the shard (outer container) along each dimension of the array.",
+    )
+
+    chunk_shape: tuple[int, ...] = Field(
+        ...,
+        description="Lengths of the chunk (inner unit) along each dimension within each shard.",
+    )
+
+    @model_validator(mode="after")
+    def validate_shapes(self) -> ShardedChunkShape:
+        """Validate that shard_shape is divisible by chunk_shape."""
+        if len(self.shard_shape) != len(self.chunk_shape):
+            msg = (
+                f"shard_shape and chunk_shape must have the same number of dimensions. "
+                f"Got shard_shape={self.shard_shape} and chunk_shape={self.chunk_shape}"
+            )
+            raise ValueError(msg)
+
+        for i, (shard_dim, chunk_dim) in enumerate(zip(self.shard_shape, self.chunk_shape, strict=True)):
+            if shard_dim % chunk_dim != 0:
+                msg = (
+                    f"shard_shape must be evenly divisible by chunk_shape along all dimensions. "
+                    f"Dimension {i}: shard_shape[{i}]={shard_dim} is not divisible by chunk_shape[{i}]={chunk_dim}"
+                )
+                raise ValueError(msg)
+            if shard_dim < chunk_dim:
+                msg = (
+                    f"shard_shape must be >= chunk_shape along all dimensions. "
+                    f"Dimension {i}: shard_shape[{i}]={shard_dim} < chunk_shape[{i}]={chunk_dim}"
+                )
+                raise ValueError(msg)
+
+        return self
+
+
 class RegularChunkGrid(CamelCaseStrictModel):
     """Represents a rectangular and regularly spaced chunk grid."""
 
@@ -36,3 +84,16 @@ class RectilinearChunkGrid(CamelCaseStrictModel):
     name: str = Field(default="rectilinear", description="The name of the chunk grid.")
 
     configuration: RectilinearChunkShape = Field(..., description="Configuration of the irregular chunk grid.")
+
+
+class ShardedChunkGrid(CamelCaseStrictModel):
+    """Represents a sharded chunk grid for Zarr v3.
+
+    Sharding enables efficient storage and access patterns by grouping multiple
+    chunks into larger shards. Each shard is stored as a single object/file,
+    reducing the number of I/O operations for sequential access patterns.
+    """
+
+    name: str = Field(default="sharded", description="The name of the chunk grid.")
+
+    configuration: ShardedChunkShape = Field(..., description="Configuration of the sharded chunk grid.")

src/mdio/builder/schemas/v1/variable.py

@@ -18,6 +18,7 @@
 from mdio.builder.schemas.base import NamedArray
 from mdio.builder.schemas.chunk_grid import RectilinearChunkGrid
 from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.chunk_grid import ShardedChunkGrid
 from mdio.builder.schemas.core import CamelCaseStrictModel
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.v1.stats import SummaryStatistics
@@ -34,7 +35,7 @@ class CoordinateMetadata(CamelCaseStrictModel):
 class VariableMetadata(CoordinateMetadata):
     """Complete Metadata for Variables and complex or large Coordinates."""
 
-    chunk_grid: RegularChunkGrid | RectilinearChunkGrid | None = Field(
+    chunk_grid: RegularChunkGrid | RectilinearChunkGrid | ShardedChunkGrid | None = Field(
         default=None,
         description="Chunk grid specification for the array.",
     )

src/mdio/builder/templates/base.py

@@ -13,6 +13,8 @@
 from mdio.builder.schemas import compressors
 from mdio.builder.schemas.chunk_grid import RegularChunkGrid
 from mdio.builder.schemas.chunk_grid import RegularChunkShape
+from mdio.builder.schemas.chunk_grid import ShardedChunkGrid
+from mdio.builder.schemas.chunk_grid import ShardedChunkShape
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.dtype import StructuredType
 from mdio.builder.schemas.v1.units import AllUnitModel
@@ -43,6 +45,7 @@ def __init__(self, data_domain: SeismicDataDomain) -> None:
         self._physical_coord_names: tuple[str, ...] = ()
         self._logical_coord_names: tuple[str, ...] = ()
         self._var_chunk_shape: tuple[int, ...] = ()
+        self._var_shard_shape: tuple[int, ...] | None = None  # Default to no sharding
 
         self._builder: MDIODatasetBuilder | None = None
         self._dim_sizes: tuple[int, ...] = ()
@@ -59,6 +62,7 @@ def __repr__(self) -> str:
             f"physical_coord_names={self._physical_coord_names}, "
             f"logical_coord_names={self._logical_coord_names}, "
             f"var_chunk_shape={self._var_chunk_shape}, "
+            f"var_shard_shape={self._var_shard_shape}, "
             f"dim_sizes={self._dim_sizes}, "
             f"units={self._units})"
         )
@@ -196,6 +200,41 @@ def full_chunk_shape(self, shape: tuple[int, ...]) -> None:
 
         self._var_chunk_shape = shape
 
+    @property
+    def full_shard_shape(self) -> tuple[int, ...] | None:
+        """Returns the shard shape for the variables, or None if sharding is disabled."""
+        if self._var_shard_shape is None:
+            return None
+
+        # If dimension sizes are not set yet, return the stored shape as-is
+        if len(self._dim_sizes) != len(self._dim_names):
+            return self._var_shard_shape
+
+        # Expand -1 values to full dimension sizes
+        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) -> None:
+        """Sets the shard shape for the variables, or None to disable sharding."""
+        if shape is None:
+            self._var_shard_shape = None
+            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)
+
+        # Validate that all values are positive integers or -1
+        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:
@@ -274,6 +313,27 @@ def _add_coordinates(self) -> None:
                 if "same name twice" not in str(exc):
                     raise
 
+    def _create_chunk_grid(self, exclude_vertical: bool = False) -> RegularChunkGrid | ShardedChunkGrid:
+        """Create the appropriate chunk grid based on sharding configuration.
+
+        Args:
+            exclude_vertical: If True, excludes the vertical (last) dimension from shapes.
+                            Used for headers and other non-sample variables.
+
+        Returns:
+            RegularChunkGrid if sharding is disabled, ShardedChunkGrid if enabled.
+        """
+        chunk_shape = self.full_chunk_shape[:-1] if exclude_vertical else self.full_chunk_shape
+        shard_shape = self.full_shard_shape
+
+        if shard_shape is None:
+            # No sharding - use regular chunk grid
+            return RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=chunk_shape))
+
+        # Sharding enabled - create sharded chunk grid
+        shard_shape = shard_shape[:-1] if exclude_vertical else shard_shape
+        return ShardedChunkGrid(configuration=ShardedChunkShape(shard_shape=shard_shape, chunk_shape=chunk_shape))
+
     def _add_trace_mask(self) -> None:
         """Add trace mask variables."""
         self._builder.add_variable(
@@ -286,7 +346,7 @@ def _add_trace_mask(self) -> None:
 
     def _add_trace_headers(self, header_dtype: StructuredType) -> None:
         """Add trace mask variables."""
-        chunk_grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=self.full_chunk_shape[:-1]))
+        chunk_grid = self._create_chunk_grid(exclude_vertical=True)
         self._builder.add_variable(
             name="headers",
             dimensions=self.spatial_dimension_names,
@@ -302,7 +362,7 @@ def _add_variables(self) -> None:
         A virtual method that can be overwritten by subclasses to add custom variables.
         Uses the class field 'builder' to add variables to the dataset.
         """
-        chunk_grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=self.full_chunk_shape))
+        chunk_grid = self._create_chunk_grid(exclude_vertical=False)
         unit = self.get_unit_by_key(self._default_variable_name)
         self._builder.add_variable(
             name=self.default_variable_name,