Implementation of NonBinned and DuplicateIndex grid overrides (#732)

* Working example of DuplicateIndex grid override

* Remove hardcoded values + linting

* Streamline code

* Relocate template mutation to all happen in the same function -- Update logging level

* Formatting

* Clean up

* Update logging

* Remove incorrect noqa

* Checkpoint

* Fully functional demo

* pre-commit

* Revert debugging changes

* pre-commit

* Simplifications to DuplicateIndex and GridOverrides

* Add base safety for grid override template mutation

* Extract template update logic

* Resolve nondeterminstic ordering causing incorrect template mutations

* Bandaid fix to join non-binned and autoshotwrap overrides -- Should be refactored for a uniform experience

* pre-commit

* Fix memory consumption on grid calculation and not appropriately calculating trace dimension

* Apply proper dimensions to non-binned coordiante arrays -- Fix large memory spike between header scan and trace ingestion

---------

Co-authored-by: Altay Sansal <tasansal@users.noreply.github.com>

Author: BrianMichell

src/mdio/builder/dataset_builder.py

@@ -150,12 +150,32 @@ def add_coordinate(  # noqa: PLR0913
             msg = "Adding coordinate with the same name twice is not allowed"
             raise ValueError(msg)
 
-        # Validate that all referenced dimensions are already defined
+        # Resolve referenced dimensions strictly, except allow a single substitution with 'trace' if present.
         named_dimensions = []
+        trace_dim = _get_named_dimension(self._dimensions, "trace")
+        resolved_dim_names: list[str] = []
+        trace_used = False
+        missing_dims: list[str] = []
         for dim_name in dimensions:
             nd = _get_named_dimension(self._dimensions, dim_name)
+            if nd is not None:
+                if dim_name not in resolved_dim_names:
+                    resolved_dim_names.append(dim_name)
+                continue
+            if trace_dim is not None and not trace_used and "trace" not in resolved_dim_names:
+                resolved_dim_names.append("trace")
+                trace_used = True
+            else:
+                missing_dims.append(dim_name)
+
+        if missing_dims:
+            msg = f"Pre-existing dimension named {missing_dims[0]!r} is not found"
+            raise ValueError(msg)
+
+        for resolved_name in resolved_dim_names:
+            nd = _get_named_dimension(self._dimensions, resolved_name)
             if nd is None:
-                msg = f"Pre-existing dimension named {dim_name!r} is not found"
+                msg = f"Pre-existing dimension named {resolved_name!r} is not found"
                 raise ValueError(msg)
             named_dimensions.append(nd)
 
@@ -174,7 +194,7 @@ def add_coordinate(  # noqa: PLR0913
         self.add_variable(
             name=coord.name,
             long_name=coord.long_name,
-            dimensions=dimensions,  # dimension names (list[str])
+            dimensions=tuple(resolved_dim_names),  # resolved dimension names
             data_type=coord.data_type,
             compressor=compressor,
             coordinates=[name],  # Use the coordinate name as a reference

src/mdio/builder/templates/base.py

@@ -82,6 +82,9 @@ def build_dataset(
 
         Returns:
             Dataset: The constructed dataset
+
+        Raises:
+            ValueError: If coordinate already exists from subclass override.
         """
         self._dim_sizes = sizes
 
@@ -90,6 +93,20 @@ def build_dataset(
         self._builder = MDIODatasetBuilder(name=name, attributes=attributes)
         self._add_dimensions()
         self._add_coordinates()
+        # Ensure any coordinates declared on the template but not added by _add_coordinates
+        # are materialized with generic defaults. This handles coordinates added by grid overrides.
+        for coord_name in self.coordinate_names:
+            try:
+                self._builder.add_coordinate(
+                    name=coord_name,
+                    dimensions=self.spatial_dimension_names,
+                    data_type=ScalarType.FLOAT64,
+                    compressor=compressors.Blosc(cname=compressors.BloscCname.zstd),
+                    metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(coord_name)),
+                )
+            except ValueError as exc:  # coordinate may already exist
+                if "same name twice" not in str(exc):
+                    raise
         self._add_variables()
         self._add_trace_mask()
 
@@ -241,14 +258,21 @@ def _add_coordinates(self) -> None:
             )
 
         # Add non-dimension coordinates
+        # Note: coordinate_names may be modified at runtime by grid overrides,
+        # so we need to handle dynamic additions gracefully
         for name in self.coordinate_names:
-            self._builder.add_coordinate(
-                name=name,
-                dimensions=self.spatial_dimension_names,
-                data_type=ScalarType.FLOAT64,
-                compressor=compressors.Blosc(cname=compressors.BloscCname.zstd),
-                metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(name)),
-            )
+            try:
+                self._builder.add_coordinate(
+                    name=name,
+                    dimensions=self.spatial_dimension_names,
+                    data_type=ScalarType.FLOAT64,
+                    compressor=compressors.Blosc(cname=compressors.BloscCname.zstd),
+                    metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(name)),
+                )
+            except ValueError as exc:
+                # Coordinate may already exist from subclass override
+                if "same name twice" not in str(exc):
+                    raise
 
     def _add_trace_mask(self) -> None:
         """Add trace mask variables."""

src/mdio/converters/segy.py

@@ -134,6 +134,146 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
         raise GridTraceSparsityError(grid.shape, num_traces, msg)
 
 
+def _patch_add_coordinates_for_non_binned(
+    template: AbstractDatasetTemplate,
+    non_binned_dims: set[str],
+) -> None:
+    """Patch template's _add_coordinates to skip adding non-binned dims as dimension coordinates.
+
+    When NonBinned override is used, dimensions like 'offset' or 'azimuth' become coordinates
+    instead of dimensions. However, template subclasses may still try to add them as 1D
+    dimension coordinates (e.g., with dimensions=("offset",)). Since 'offset' is no longer
+    a dimension, the builder substitutes 'trace', resulting in wrong coordinate dimensions.
+
+    This function patches the template's _add_coordinates method to intercept calls to
+    builder.add_coordinate and skip adding coordinates that are non-binned dims with
+    single-element dimension tuples. These coordinates will be added later by build_dataset
+    with the correct spatial_dimension_names (e.g., (inline, crossline, trace)).
+
+    Args:
+        template: The template to patch
+        non_binned_dims: Set of dimension names that became coordinates due to NonBinned override
+    """
+    # Check if already patched to avoid duplicate patching
+    if hasattr(template, "_mdio_non_binned_patched"):
+        return
+
+    # Store the original _add_coordinates method
+    original_add_coordinates = template._add_coordinates
+
+    def patched_add_coordinates() -> None:
+        """Wrapper that intercepts builder.add_coordinate calls for non-binned dims."""
+        # Store the original add_coordinate method from the builder
+        original_builder_add_coordinate = template._builder.add_coordinate
+
+        def filtered_add_coordinate(  # noqa: ANN202
+            name: str,
+            *,
+            dimensions: tuple[str, ...],
+            **kwargs,  # noqa: ANN003
+        ):
+            """Skip adding non-binned dims as 1D dimension coordinates."""
+            # Check if this is a non-binned dim being added as a 1D dimension coordinate
+            # (i.e., the coordinate name matches a non-binned dim and has only 1 dimension)
+            if name in non_binned_dims and len(dimensions) == 1:
+                logger.debug(
+                    "Skipping 1D coordinate '%s' with dims %s - will be added with full spatial dims",
+                    name,
+                    dimensions,
+                )
+                return template._builder  # Return builder for chaining, but don't add
+
+            # Otherwise, call the original method
+            return original_builder_add_coordinate(name, dimensions=dimensions, **kwargs)
+
+        # Temporarily replace builder's add_coordinate
+        template._builder.add_coordinate = filtered_add_coordinate
+
+        try:
+            # Call the original _add_coordinates
+            original_add_coordinates()
+        finally:
+            # Restore the original add_coordinate method
+            template._builder.add_coordinate = original_builder_add_coordinate
+
+    # Replace the template's _add_coordinates method
+    template._add_coordinates = patched_add_coordinates
+
+    # Mark as patched to prevent duplicate patching
+    template._mdio_non_binned_patched = True
+
+
+def _update_template_from_grid_overrides(
+    template: AbstractDatasetTemplate,
+    grid_overrides: dict[str, Any] | None,
+    segy_dimensions: list[Dimension],
+    full_chunk_shape: tuple[int, ...],
+    chunk_size: tuple[int, ...],
+) -> None:
+    """Update template attributes to match grid plan results after grid overrides.
+
+    This function modifies the template in-place to reflect changes from grid overrides:
+    - Updates chunk shape if it changed due to overrides
+    - Updates dimension names if they changed due to overrides
+    - Adds non-binned dimensions as coordinates for NonBinned override
+    - Patches _add_coordinates to skip adding non-binned dims as dimension coordinates
+
+    Args:
+        template: The template to update
+        grid_overrides: Grid override configuration
+        segy_dimensions: Dimensions returned from grid planning
+        full_chunk_shape: Original template chunk shape
+        chunk_size: Chunk size returned from grid planning
+    """
+    # Update template to match grid_plan results after grid overrides
+    # Extract actual spatial dimensions from segy_dimensions (excluding vertical dimension)
+    actual_spatial_dims = tuple(dim.name for dim in segy_dimensions[:-1])
+
+    # Align chunk_size with actual dimensions - truncate if dimensions were filtered out
+    num_actual_spatial = len(actual_spatial_dims)
+    num_chunk_spatial = len(chunk_size) - 1  # Exclude vertical dimension chunk
+    if num_actual_spatial != num_chunk_spatial:
+        # Truncate chunk_size to match actual dimensions
+        chunk_size = chunk_size[:num_actual_spatial] + (chunk_size[-1],)
+
+    if full_chunk_shape != chunk_size:
+        logger.debug(
+            "Adjusting template chunk shape from %s to %s to match grid after overrides",
+            full_chunk_shape,
+            chunk_size,
+        )
+        template._var_chunk_shape = chunk_size
+
+    # Update dimensions if they don't match grid_plan results
+    if template.spatial_dimension_names != actual_spatial_dims:
+        logger.debug(
+            "Adjusting template dimensions from %s to %s to match grid after overrides",
+            template.spatial_dimension_names,
+            actual_spatial_dims,
+        )
+        template._dim_names = actual_spatial_dims + (template.trace_domain,)
+
+    # If using NonBinned override, expose non-binned dims as logical coordinates on the template instance
+    # and patch _add_coordinates to skip adding them as 1D dimension coordinates
+    if grid_overrides and "NonBinned" in grid_overrides and "non_binned_dims" in grid_overrides:
+        non_binned_dims = tuple(grid_overrides["non_binned_dims"])
+        if non_binned_dims:
+            logger.debug(
+                "NonBinned grid override: exposing non-binned dims as coordinates: %s",
+                non_binned_dims,
+            )
+            # Append any missing names; keep existing order and avoid duplicates
+            existing = set(template.coordinate_names)
+            to_add = tuple(n for n in non_binned_dims if n not in existing)
+            if to_add:
+                template._logical_coord_names = template._logical_coord_names + to_add
+
+            # Patch _add_coordinates to skip adding non-binned dims as 1D dimension coordinates
+            # This prevents them from being added with wrong dimensions (e.g., just "trace")
+            # They will be added later by build_dataset with full spatial_dimension_names
+            _patch_add_coordinates_for_non_binned(template, set(non_binned_dims))
+
+
 def _scan_for_headers(
     segy_file_kwargs: SegyFileArguments,
     segy_file_info: SegyFileInfo,
@@ -143,7 +283,11 @@ def _scan_for_headers(
     """Extract trace dimensions and index headers from the SEG-Y file.
 
     This is an expensive operation.
-    It scans the SEG-Y file in chunks by using ProcessPoolExecutor
+    It scans the SEG-Y file in chunks by using ProcessPoolExecutor.
+
+    Note:
+        If grid_overrides are applied to the template before calling this function,
+        the chunk_size returned from get_grid_plan should match the template's chunk shape.
     """
     full_chunk_shape = template.full_chunk_shape
     segy_dimensions, chunk_size, segy_headers = get_grid_plan(
@@ -154,13 +298,15 @@ def _scan_for_headers(
         chunksize=full_chunk_shape,
         grid_overrides=grid_overrides,
     )
-    if full_chunk_shape != chunk_size:
-        # TODO(Dmitriy): implement grid overrides
-        # https://github.com/TGSAI/mdio-python/issues/585
-        # The returned 'chunksize' is used only for grid_overrides. We will need to use it when full
-        # support for grid overrides is implemented
-        err = "Support for changing full_chunk_shape in grid overrides is not yet implemented"
-        raise NotImplementedError(err)
+
+    _update_template_from_grid_overrides(
+        template=template,
+        grid_overrides=grid_overrides,
+        segy_dimensions=segy_dimensions,
+        full_chunk_shape=full_chunk_shape,
+        chunk_size=chunk_size,
+    )
+
     return segy_dimensions, segy_headers
 
 
@@ -233,7 +379,8 @@ def _get_coordinates(
         if coord_name not in segy_headers.dtype.names:
             err = f"Coordinate '{coord_name}' not found in SEG-Y dimensions."
             raise ValueError(err)
-        non_dim_coords[coord_name] = segy_headers[coord_name]
+        # Copy the data to allow segy_headers to be garbage collected
+        non_dim_coords[coord_name] = np.array(segy_headers[coord_name])
 
     return dimensions_coords, non_dim_coords
 
@@ -255,25 +402,54 @@ def populate_non_dim_coordinates(
     drop_vars_delayed: list[str],
     spatial_coordinate_scalar: int,
 ) -> tuple[xr_Dataset, list[str]]:
-    """Populate the xarray dataset with coordinate variables."""
+    """Populate the xarray dataset with coordinate variables.
+
+    Memory optimization: Processes coordinates one at a time and explicitly
+    releases intermediate arrays to reduce peak memory usage.
+    """
     non_data_domain_dims = grid.dim_names[:-1]  # minus the data domain dimension
-    for coord_name, coord_values in coordinates.items():
+
+    # Process coordinates one at a time to minimize peak memory
+    coord_names = list(coordinates.keys())
+    for coord_name in coord_names:
+        coord_values = coordinates.pop(coord_name)  # Remove from dict to free memory
         da_coord = dataset[coord_name]
-        tmp_coord_values = dataset[coord_name].values
 
+        # Get coordinate shape from dataset (uses dask shape, no memory allocation)
+        coord_shape = da_coord.shape
+
+        # Create output array with fill value
+        fill_value = da_coord.encoding.get("_FillValue") or da_coord.encoding.get("fill_value")
+        if fill_value is None:
+            fill_value = np.nan
+        tmp_coord_values = np.full(coord_shape, fill_value, dtype=da_coord.dtype)
+
+        # Compute slices for this coordinate's dimensions
         coord_axes = tuple(non_data_domain_dims.index(coord_dim) for coord_dim in da_coord.dims)
         coord_slices = tuple(slice(None) if idx in coord_axes else 0 for idx in range(len(non_data_domain_dims)))
-        coord_trace_indices = grid.map[coord_slices]
 
+        # Read only the required slice from grid map
+        coord_trace_indices = np.asarray(grid.map[coord_slices])
+
+        # Find valid (non-null) indices
         not_null = coord_trace_indices != grid.map.fill_value
-        tmp_coord_values[not_null] = coord_values[coord_trace_indices[not_null]]
 
+        # Populate values efficiently
+        if not_null.any():
+            valid_indices = coord_trace_indices[not_null]
+            tmp_coord_values[not_null] = coord_values[valid_indices]
+
+        # Apply scalar if needed
         if coord_name in SCALE_COORDINATE_KEYS:
             tmp_coord_values = _apply_coordinate_scalar(tmp_coord_values, spatial_coordinate_scalar)
 
+        # Assign to dataset
         dataset[coord_name][:] = tmp_coord_values
         drop_vars_delayed.append(coord_name)
 
+        # Explicitly release intermediate arrays
+        del tmp_coord_values, coord_trace_indices, not_null, coord_values
+
         # TODO(Altay): Add verification of reduced coordinates being the same as the first
         # https://github.com/TGSAI/mdio-python/issues/645
 
@@ -554,6 +730,10 @@ def segy_to_mdio(  # noqa PLR0913
     grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
+
+    # Explicitly delete segy_headers to free memory - coordinate values have been copied
+    del segy_headers
+
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
 
     if settings.raw_headers: