Begin decoupling checksum logic

Author: BrianMichell

pyproject.toml

@@ -31,6 +31,8 @@ dependencies = [
     "universal-pathlib>=0.3.3",
     "xarray>=2025.10.1",
     "zarr>=3.1.3",
+    "crc32c_dist_rs",
+    "google-crc32c>=1.5.0",
 ]
 
 [project.optional-dependencies]
@@ -80,7 +82,12 @@ docs = [
 required-version = ">=0.8.17"
 
 [tool.uv.sources]
-crc32c_dist_rs = { path = "/home/brian_michell_tgs_com/source/crc32c_dist_rs/target/wheels/crc32c_dist_rs-0.1.4-cp313-cp313-manylinux_2_34_x86_64.whl" }
+crc32c_dist_rs = [
+    { path = "/home/brian_michell_tgs_com/source/crc32c_dist_rs/target/wheels/crc32c_dist_rs-0.1.4-cp311-cp311-manylinux_2_34_x86_64.whl", marker = "python_version == '3.11'" },
+    { path = "/home/brian_michell_tgs_com/source/crc32c_dist_rs/target/wheels/crc32c_dist_rs-0.1.4-cp312-cp312-manylinux_2_34_x86_64.whl", marker = "python_version == '3.12'" },
+    { path = "/home/brian_michell_tgs_com/source/crc32c_dist_rs/target/wheels/crc32c_dist_rs-0.1.4-cp313-cp313-manylinux_2_34_x86_64.whl", marker = "python_version == '3.13'" },
+]
+
 
 [tool.ruff]
 target-version = "py311"

src/mdio/segy/_workers.py

@@ -32,9 +32,8 @@ def header_scan_worker(
     segy_file_kwargs: SegyFileArguments,
     trace_range: tuple[int, int],
     subset: list[str] | None = None,
-    calculate_checksum: bool = False,
-) -> HeaderArray | tuple[HeaderArray, tuple[int, int, int]]:
-    """Header scan worker with optional checksum calculation.
+) -> HeaderArray:
+    """Header scan worker.
 
     If SegyFile is not open, it can either accept a path string or a handle that was opened in
     a different context manager.
@@ -43,11 +42,9 @@ def header_scan_worker(
         segy_file_kwargs: Arguments to open SegyFile instance.
         trace_range: Tuple consisting of the trace ranges to read.
         subset: List of header names to filter and keep.
-        calculate_checksum: If True, also calculate CRC32C for this trace range.
 
     Returns:
-        HeaderArray if calculate_checksum is False, otherwise tuple of (HeaderArray, checksum_info)
-        where checksum_info is (byte_offset, crc32c, byte_length).
+        HeaderArray parsed from SEG-Y library.
     """
     segy_file = SegyFileWrapper(**segy_file_kwargs)
 
@@ -73,33 +70,7 @@ def header_scan_worker(
     # (singleton) so we can concat and assign stuff later.
     trace_header = np.array(trace_header, dtype=new_dtype, ndmin=1)
 
-    if not calculate_checksum:
-        return HeaderArray(trace_header)  # wrap back so we can use aliases
-
-    # Verify checksum libraries are available (should have been checked earlier)
-    if not is_checksum_available():
-        logger.warning("Checksum calculation requested but libraries not available. Skipping checksum.")
-        # Return headers without checksum info - this should not happen in practice
-        # as the caller should check availability first
-        return HeaderArray(trace_header)
-
-    # Calculate checksum from the raw bytes ALREADY IN MEMORY (NO ADDITIONAL I/O!)
-    raw_bytes = traces.trace_buffer_array.tobytes()
-    partial_crc32c = calculate_bytes_crc32c(raw_bytes)
-
-    # Calculate byte offset and length
-    trace_header_size = segy_file.spec.trace.header.itemsize
-    # sample_size = segy_file.spec.trace.sample.itemsize
-    sample_size = 4  # This will always be a 4-byte float
-    num_samples = len(segy_file.sample_labels)
-    trace_size = trace_header_size + (num_samples * sample_size)
-
-    byte_offset = 3600 + start_trace * trace_size
-    byte_length = len(raw_bytes)
-
-    checksum_info = (byte_offset, partial_crc32c, byte_length)
-
-    return HeaderArray(trace_header), checksum_info
+    return HeaderArray(trace_header)  # wrap back so we can use aliases
 
 
 def trace_worker(  # noqa: PLR0913

src/mdio/segy/checksum.py

@@ -18,9 +18,13 @@
 import os
 from typing import TYPE_CHECKING
 from typing import Any
+import numpy as np
+from segy import SegyFile
+from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 
+from segy.arrays import HeaderArray
 if TYPE_CHECKING:
-    from segy.arrays import HeaderArray
+    from mdio.segy._workers import SegyFileArguments
 
 logger = logging.getLogger(__name__)
 
@@ -43,6 +47,72 @@
     DistributedCRC32C = None  # type: ignore[assignment,misc]
 
 
+def header_scan_worker(
+    segy_file_kwargs: "SegyFileArguments",
+    trace_range: tuple[int, int],
+    subset: list[str] | None = None,
+) -> HeaderArray | tuple[HeaderArray, tuple[int, int, int]]:
+    """Header scan worker with optional checksum calculation.
+
+    If SegyFile is not open, it can either accept a path string or a handle that was opened in
+    a different context manager.
+
+    Args:
+        segy_file_kwargs: Arguments to open SegyFile instance.
+        trace_range: Tuple consisting of the trace ranges to read.
+        subset: List of header names to filter and keep.
+        calculate_checksum: If True, also calculate CRC32C for this trace range.
+
+    Returns:
+        HeaderArray if calculate_checksum is False, otherwise tuple of (HeaderArray, checksum_info)
+        where checksum_info is (byte_offset, crc32c, byte_length).
+    """
+    print("Using header_scan_worker from checksum.py")
+    segy_file = SegyFile(**segy_file_kwargs)
+
+    start_trace, end_trace = trace_range
+    trace_indices = list(range(start_trace, end_trace))
+
+    # Read full trace data (header + samples) ONCE using the raw trace wrapper
+    # This avoids duplicate I/O - we get both headers and raw bytes in a single read
+    traces = SegyFileRawTraceWrapper(segy_file, trace_indices)
+
+    # Extract headers from the data we already read (no additional I/O)
+    trace_header = traces.header
+
+    if subset is not None:
+        # struct field selection needs a list, not a tuple; a subset is a tuple from the template.
+        trace_header = trace_header[list(subset)]
+
+    # Get non-void fields from dtype and copy to new array for memory efficiency
+    fields = trace_header.dtype.fields
+    non_void_fields = [(name, dtype) for name, (dtype, _) in fields.items()]
+    new_dtype = np.dtype(non_void_fields)
+
+    # Copy to non-padded memory, ndmin is to handle the case where there is 1 trace in block
+    # (singleton) so we can concat and assign stuff later.
+    trace_header = np.array(trace_header, dtype=new_dtype, ndmin=1)
+
+
+    # Calculate checksum from the raw bytes ALREADY IN MEMORY (NO ADDITIONAL I/O!)
+    raw_bytes = traces.trace_buffer_array.tobytes()
+    partial_crc32c = calculate_bytes_crc32c(raw_bytes)
+
+    # Calculate byte offset and length
+    trace_header_size = segy_file.spec.trace.header.itemsize
+    # sample_size = segy_file.spec.trace.sample.itemsize
+    sample_size = 4  # This will always be a 4-byte float
+    num_samples = len(segy_file.sample_labels)
+    trace_size = trace_header_size + (num_samples * sample_size)
+
+    byte_offset = 3600 + start_trace * trace_size
+    byte_length = len(raw_bytes)
+
+    checksum_info = (byte_offset, partial_crc32c, byte_length)
+
+    return HeaderArray(trace_header), checksum_info
+
+
 def is_checksum_available() -> bool:
     """Check if checksum libraries are available.
 

src/mdio/segy/parsers.py

@@ -16,10 +16,11 @@
 from tqdm.auto import tqdm
 from upath import UPath
 
-from mdio.segy._workers import header_scan_worker
+# All imports from the `checksum` module are slated to be removed in the future
 from mdio.segy.checksum import create_distributed_crc32c
 from mdio.segy.checksum import finalize_distributed_checksum
 from mdio.segy.checksum import is_checksum_available
+from mdio.segy.checksum import should_calculate_checksum
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
@@ -51,6 +52,14 @@ def parse_headers(  # noqa: PLR0913
         HeaderArray if calculate_checksum is False.
         Tuple of (HeaderArray, combined_crc32c) if calculate_checksum is True.
     """
+    # Dynamically import the appropriate header_scan_worker based on checksum requirement
+    if should_calculate_checksum() and is_checksum_available():
+        from mdio.segy.checksum import header_scan_worker
+    else:
+        from mdio.segy._workers import header_scan_worker
+
+    # Type hint for the dynamically imported function
+    header_scan_worker: Any
     # Initialize combiner only if checksum calculation is needed and available
     combiner: Any = None
     if calculate_checksum:
@@ -101,7 +110,6 @@ def parse_headers(  # noqa: PLR0913
             repeat(segy_file_kwargs),
             trace_ranges,
             repeat(subset),
-            repeat(calculate_checksum),
         )
 
         if progress_bar is True:
@@ -124,4 +132,58 @@ def parse_headers(  # noqa: PLR0913
         # Merge headers and return
         return np.concatenate(results)
 
-    return finalize_distributed_checksum(results, combiner)
+    return finalize_distributed_checksum(results, combiner)
+
+# def _parse_headers(
+#     segy_file_kwargs: SegyFileArguments,
+#     num_traces: int,
+#     subset: list[str] | None = None,
+#     block_size: int = 10000,
+#     progress_bar: bool = True,
+# ) -> HeaderArray:
+#     """Read and parse given `byte_locations` from SEG-Y file.
+
+#     Args:
+#         segy_file_kwargs: SEG-Y file arguments.
+#         num_traces: Total number of traces in the SEG-Y file.
+#         subset: List of header names to filter and keep.
+#         block_size: Number of traces to read for each block.
+#         progress_bar: Enable or disable progress bar. Default is True.
+
+#     Returns:
+#         HeaderArray. Keys are the index names, values are numpy arrays of parsed headers for the
+#         current block. Array is of type byte_type except IBM32 which is mapped to FLOAT32.
+#     """
+#     trace_count = num_traces
+#     n_blocks = int(ceil(trace_count / block_size))
+
+#     trace_ranges = []
+#     for idx in range(n_blocks):
+#         start, stop = idx * block_size, (idx + 1) * block_size
+#         stop = min(stop, trace_count)
+
+#         trace_ranges.append((start, stop))
+
+#     num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
+#     num_workers = min(n_blocks, num_cpus)
+
+#     tqdm_kw = {"unit": "block", "dynamic_ncols": True}
+#     # For Unix async writes with s3fs/fsspec & multiprocessing, use 'spawn' instead of default
+#     # 'fork' to avoid deadlocks on cloud stores. Slower but necessary. Default on Windows.
+#     context = mp.get_context("spawn")
+#     with ProcessPoolExecutor(num_workers, mp_context=context) as executor:
+#         lazy_work = executor.map(header_scan_worker, repeat(segy_file_kwargs), trace_ranges, repeat(subset))
+
+#         if progress_bar is True:
+#             lazy_work = tqdm(
+#                 iterable=lazy_work,
+#                 total=n_blocks,
+#                 desc="Scanning SEG-Y for geometry attributes",
+#                 **tqdm_kw,
+#             )
+
+#         # This executes the lazy work.
+#         headers: list[HeaderArray] = list(lazy_work)
+
+#     # Merge blocks before return
+#     return np.concatenate(headers)

tests/integration/test_crc32c_checksum.py

@@ -126,6 +126,7 @@ def test_corrupted_data_detects_changes(self, tmp_path: Path, monkeypatch: pytes
         """Test that checksums detect data corruption using synthetic streamer data."""
         # Enable checksum calculation via environment variable
         monkeypatch.setenv("MDIO__IMPORT__RAW_HEADERS", "true")
+        monkeypatch.setenv("MDIO__IMPORT__DO_CRC32C", "true")
 
         # Use streamer configuration for this test
         test_conf = STREAMER_3D_CONF