Reimplement disaster recovery logic

Author: BrianMichell

src/mdio/builder/schemas/dtype.py

@@ -32,6 +32,7 @@ class ScalarType(StrEnum):
     COMPLEX64 = "complex64"
     COMPLEX128 = "complex128"
     COMPLEX256 = "complex256"
+    HEADERS_V3 = "r1920"  # Raw number of BITS, must be a multiple of 8
 
 
 class StructuredField(CamelCaseStrictModel):

src/mdio/constants.py

@@ -64,4 +64,5 @@ class ZarrFormat(IntEnum):
     ScalarType.COMPLEX64: complex(np.nan, np.nan),
     ScalarType.COMPLEX128: complex(np.nan, np.nan),
     ScalarType.COMPLEX256: complex(np.nan, np.nan),
+    ScalarType.HEADERS_V3: b"\x00" * 240,
 }

src/mdio/converters/segy.py

@@ -22,6 +22,11 @@
 from mdio.converters.exceptions import GridTraceSparsityError
 from mdio.converters.type_converter import to_structured_type
 from mdio.core.grid import Grid
+from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.chunk_grid import RegularChunkShape
+from mdio.builder.schemas.compressors import Blosc
+from mdio.builder.schemas.compressors import BloscCname
+from mdio.builder.schemas.dtype import ScalarType
 from mdio.segy import blocked_io
 from mdio.segy.utilities import get_grid_plan
 
@@ -333,6 +338,58 @@ def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, A
         dataset.metadata.attributes["gridOverrides"] = grid_overrides
 
 
+def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> AbstractDatasetTemplate:
+    """Add raw headers capability to the MDIO template by monkey-patching its _add_variables method.
+    This function modifies the template's _add_variables method to also add a raw headers variable
+    with the following characteristics:
+    - Same rank as the Headers variable (all dimensions except vertical)
+    - Name: "RawHeaders"
+    - Type: ScalarType.HEADERS
+    - No coordinates
+    - zstd compressor
+    - No additional metadata
+    - Chunked the same as the Headers variable
+    Args:
+        mdio_template: The MDIO template to mutate
+    """
+    # Check if raw headers enhancement has already been applied to avoid duplicate additions
+    if hasattr(mdio_template, '_mdio_raw_headers_enhanced'):
+        return mdio_template
+
+    # Store the original _add_variables method
+    original_add_variables = mdio_template._add_variables
+
+    def enhanced_add_variables() -> None:
+        # Call the original method first
+        original_add_variables()
+
+        # Now add the raw headers variable
+        chunk_shape = mdio_template._var_chunk_shape[:-1]
+
+        # Create chunk grid metadata
+        chunk_metadata = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=chunk_shape))
+        from mdio.builder.schemas.v1.variable import VariableMetadata
+
+        # Add the raw headers variable using the builder's add_variable method
+        mdio_template._builder.add_variable(
+            name="raw_headers",
+            long_name="Raw Headers",
+            dimensions=mdio_template._dim_names[:-1],  # All dimensions except vertical
+            data_type=ScalarType.HEADERS_V3,
+            compressor=Blosc(cname=BloscCname.zstd),
+            coordinates=None,  # No coordinates as specified
+            metadata=VariableMetadata(chunk_grid=chunk_metadata),
+        )
+
+    # Replace the template's _add_variables method
+    mdio_template._add_variables = enhanced_add_variables
+
+    # Mark the template as enhanced to prevent duplicate monkey-patching
+    mdio_template._mdio_raw_headers_enhanced = True
+
+    return mdio_template
+
+
 def segy_to_mdio(  # noqa PLR0913
     segy_spec: SegySpec,
     mdio_template: AbstractDatasetTemplate,
@@ -372,6 +429,11 @@ def segy_to_mdio(  # noqa PLR0913
 
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
+
+    if os.getenv("MDIO__DO_RAW_HEADERS") == "1":
+        logger.warning("MDIO__DO_RAW_HEADERS is experimental and expected to change or be removed.")
+        mdio_template = _add_raw_headers_to_template(mdio_template)
+
     horizontal_unit = _get_horizontal_coordinate_unit(segy_dimensions)
     mdio_ds: Dataset = mdio_template.build_dataset(
         name=mdio_template.name,

src/mdio/converters/type_converter.py

@@ -78,6 +78,8 @@ def to_structured_type(data_type: np_dtype) -> StructuredType:
 def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
     """Get the numpy dtype for a variable."""
     if isinstance(data_type, ScalarType):
+        if data_type == ScalarType.HEADERS_V3:
+            return np_dtype("|V240")
         return np_dtype(data_type.value)
     if isinstance(data_type, StructuredType):
         return np_dtype([(f.name, f.format.value) for f in data_type.fields])

src/mdio/segy/_workers.py

@@ -124,12 +124,15 @@ def trace_worker(  # noqa: PLR0913
     traces = segy_file.trace[live_trace_indexes]
 
     header_key = "headers"
+    raw_header_key = "raw_headers"
 
     # Get subset of the dataset that has not yet been saved
     # The headers might not be present in the dataset
     worker_variables = [data_variable_name]
     if header_key in dataset.data_vars:  # Keeping the `if` here to allow for more worker configurations
         worker_variables.append(header_key)
+    if raw_header_key in dataset.data_vars:
+        worker_variables.append(raw_header_key)
 
     ds_to_write = dataset[worker_variables]
 
@@ -148,7 +151,15 @@ def trace_worker(  # noqa: PLR0913
             attrs=ds_to_write[header_key].attrs,
             encoding=ds_to_write[header_key].encoding,  # Not strictly necessary, but safer than not doing it.
         )
-
+    if raw_header_key in worker_variables:
+        tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
+        tmp_raw_headers[not_null] = traces.header.view("|V240")  # TODO: Ensure this is using the RAW view and not an interpreted view.
+        ds_to_write[raw_header_key] = Variable(
+            ds_to_write[raw_header_key].dims,
+            tmp_raw_headers,
+            attrs=ds_to_write[raw_header_key].attrs,
+            encoding=ds_to_write[raw_header_key].encoding,  # Not strictly necessary, but safer than not doing it.
+        )
     data_variable = ds_to_write[data_variable_name]
     fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))
     tmp_samples = np.full_like(data_variable, fill_value=fill_value)