rename trace wrapper and do lazy decoding

Author: tasansal

src/mdio/converters/segy.py

@@ -353,6 +353,7 @@ def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> Abst
     - zstd compressor
     - No additional metadata
     - Chunked the same as the Headers variable
+
     Args:
         mdio_template: The MDIO template to mutate
     Returns:

src/mdio/segy/_disaster_recovery_wrapper.py

@@ -0,0 +1,49 @@
+"""Consumer-side utility to get both raw and transformed header data with single filesystem read."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+    from segy import SegyFile
+
+
+class SegyFileRawTraceWrapper:
+    def __init__(self, segy_file: SegyFile, indices: int | list[int] | NDArray | slice):
+        self.segy_file = segy_file
+        self.indices = indices
+
+        self.idx = self.segy_file.trace.normalize_and_validate_query(self.indices)
+        self.trace_buffer_array = self.segy_file.trace.fetch(self.idx, raw=True)
+
+        self.trace_view = self.trace_buffer_array.view(self.segy_file.spec.trace.dtype)
+
+        self.trace_decode_pipeline = self.segy_file.accessors.trace_decode_pipeline
+        self.decoded_traces = None  # decode later when not-raw header/sample is called
+
+    def _ensure_decoded(self) -> None:
+        """Apply trace decoding pipeline if not already done."""
+        if self.decoded_traces is not None:  # already done
+            return
+        self.decoded_traces = self.trace_decode_pipeline.apply(self.trace_view.copy())
+
+    @property
+    def raw_header(self) -> NDArray:
+        """Get byte array view of the raw headers."""
+        header_itemsize = self.segy_file.spec.trace.header.itemsize  # should be 240
+        return self.trace_view.header.view(np.dtype((np.void, header_itemsize)))
+
+    @property
+    def header(self) -> NDArray:
+        """Get decoded header."""
+        self._ensure_decoded()  # decode when needed in-place to avoid copy.
+        return self.decoded_traces.header
+
+    @property
+    def sample(self) -> NDArray:
+        """Get decoded trace samples."""
+        self._ensure_decoded()  # decode when needed in-place to avoid copy.
+        return self.decoded_traces.sample

src/mdio/segy/_raw_trace_wrapper.py

@@ -12,7 +12,7 @@
 
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.dtype import ScalarType
-from mdio.segy._disaster_recovery_wrapper import SegyFileTraceDataWrapper
+from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
@@ -139,7 +139,7 @@ def trace_worker(  # noqa: PLR0913
     # For that reason, we have wrapped the accessors to provide an interface that can be removed
     # and not require additional changes to the below code.
     # NOTE: The `raw_header_key` code block should be removed in full as it will become dead code.
-    traces = SegyFileTraceDataWrapper(segy_file, live_trace_indexes)
+    traces = SegyFileRawTraceWrapper(segy_file, live_trace_indexes)
 
     ds_to_write = dataset[worker_variables]
 

src/mdio/segy/_workers.py

@@ -9,6 +9,8 @@
 
 import tempfile
 from pathlib import Path
+from typing import TYPE_CHECKING
+from typing import Any
 
 import numpy as np
 import pytest
@@ -19,7 +21,10 @@
 from segy.schema import SegySpec
 from segy.standards import get_segy_standard
 
-from mdio.segy._disaster_recovery_wrapper import SegyFileTraceDataWrapper
+from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
 
 SAMPLES_PER_TRACE = 1501
 
@@ -28,7 +33,7 @@ class TestDisasterRecoveryWrapper:
     """Test cases for disaster recovery wrapper functionality."""
 
     @pytest.fixture
-    def temp_dir(self) -> Path:
+    def temp_dir(self) -> Generator[Path, Any, None]:
         """Create a temporary directory for test files."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             yield Path(tmp_dir)
@@ -140,7 +145,7 @@ def test_wrapper_basic_functionality(self, temp_dir: Path, basic_segy_spec: Segy
 
         # Test single trace
         trace_idx = 3
-        wrapper = SegyFileTraceDataWrapper(segy_file, trace_idx)
+        wrapper = SegyFileRawTraceWrapper(segy_file, trace_idx)
 
         # Test that properties are accessible
         assert wrapper.header is not None
@@ -184,7 +189,7 @@ def test_wrapper_with_multiple_traces(self, temp_dir: Path, basic_segy_spec: Seg
 
         # Test with list of indices
         trace_indices = [0, 2, 4]
-        wrapper = SegyFileTraceDataWrapper(segy_file, trace_indices)
+        wrapper = SegyFileRawTraceWrapper(segy_file, trace_indices)
 
         # Test that properties work with multiple traces
         assert wrapper.header is not None
@@ -222,7 +227,7 @@ def test_wrapper_with_slice_indices(self, temp_dir: Path, basic_segy_spec: SegyS
         segy_file = SegyFile(segy_path, spec=spec)
 
         # Test with slice
-        wrapper = SegyFileTraceDataWrapper(segy_file, slice(5, 15))
+        wrapper = SegyFileRawTraceWrapper(segy_file, slice(5, 15))
 
         # Test that properties work with slice
         assert wrapper.header is not None
@@ -269,7 +274,7 @@ def test_different_index_types(
         segy_file = SegyFile(segy_path, spec=spec)
 
         # Create wrapper with different index types
-        wrapper = SegyFileTraceDataWrapper(segy_file, trace_indices)
+        wrapper = SegyFileRawTraceWrapper(segy_file, trace_indices)
 
         # Basic validation that we got results
         assert wrapper.header is not None