Re-implement base checksum requirements

Author: BrianMichell

pyproject.toml

@@ -20,8 +20,10 @@ classifiers = [
 dependencies = [
     "click>=8.3.0",
     "click-params>=0.5.0",
+    "crc32c_dist_rs",
     "dask>=2025.9.1",
     "fsspec>=2025.9.0",
+    "google-crc32c>=1.5.0",
     "pint>=0.25.0",
     "psutil>=7.1.0",
     "pydantic>=2.12.0",
@@ -37,6 +39,7 @@ dependencies = [
 cloud = ["s3fs>=2025.9.0", "gcsfs>=2025.9.0", "adlfs>=2025.8.0"]
 distributed = ["distributed>=2025.9.1", "bokeh>=3.8.0"]
 lossy = ["zfpy>=1.0.1"]
+performance = ["numba>=0.60.0"]
 
 [project.urls]
 homepage = "https://mdio.dev/"
@@ -77,6 +80,9 @@ docs = [
 [tool.uv]
 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" }
+
 [tool.ruff]
 target-version = "py311"
 src = ["src"]
@@ -196,4 +202,4 @@ module-name = "mdio"
 
 [build-system]
 requires = ["uv_build>=0.8.17,<0.9.0"]
-build-backend = "uv_build"
+build-backend = "uv_build"

src/mdio/converters/segy.py

@@ -13,8 +13,10 @@
 from segy.config import SegyHeaderOverrides
 from segy.standards.codes import MeasurementSystem as SegyMeasurementSystem
 from segy.standards.fields import binary as binary_header_fields
+from zarr import open_group as zarr_open_group
 
 from mdio.api.io import _normalize_path
+from mdio.api.io import _normalize_storage_options
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.chunk_grid import RegularChunkGrid
 from mdio.builder.schemas.chunk_grid import RegularChunkShape
@@ -140,29 +142,38 @@ def _scan_for_headers(
     segy_file_info: SegyFileInfo,
     template: AbstractDatasetTemplate,
     grid_overrides: dict[str, Any] | None = None,
-) -> tuple[list[Dimension], SegyHeaderArray]:
+    calculate_checksum: bool = False,
+) -> tuple[list[Dimension], SegyHeaderArray, int | None]:
     """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
+    Optionally calculates CRC32C checksum for all trace data during scanning.
     """
     full_chunk_size = template.full_chunk_size
-    segy_dimensions, chunk_size, segy_headers = get_grid_plan(
+    grid_plan_result = get_grid_plan(
         segy_file_kwargs=segy_file_kwargs,
         segy_file_info=segy_file_info,
         return_headers=True,
         template=template,
         chunksize=full_chunk_size,
         grid_overrides=grid_overrides,
+        calculate_checksum=calculate_checksum,
     )
+
+    if calculate_checksum:
+        segy_dimensions, chunk_size, segy_headers, trace_data_crc32c = grid_plan_result
+    else:
+        segy_dimensions, chunk_size, segy_headers = grid_plan_result
+        trace_data_crc32c = None
     if full_chunk_size != 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_size in grid overrides is not yet implemented"
         raise NotImplementedError(err)
-    return segy_dimensions, segy_headers
+    return segy_dimensions, segy_headers, trace_data_crc32c
 
 
 def _build_and_check_grid(
@@ -518,11 +529,15 @@ def segy_to_mdio(  # noqa PLR0913
     }
     segy_file_info = get_segy_file_info(segy_file_kwargs)
 
-    segy_dimensions, segy_headers = _scan_for_headers(
+    # Only calculate CRC32C when raw headers are enabled
+    calculate_checksum = os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on")
+
+    segy_dimensions, segy_headers, trace_data_crc32c = _scan_for_headers(
         segy_file_kwargs,
         segy_file_info,
         template=mdio_template,
         grid_overrides=grid_overrides,
+        calculate_checksum=calculate_checksum,
     )
     grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
@@ -580,10 +595,27 @@ def segy_to_mdio(  # noqa PLR0913
     default_variable_name = mdio_template.default_variable_name
     # This is an memory-expensive and time-consuming read-write operation
     # performed in chunks to save the memory
+    # NOTE: trace_data_crc32c was already calculated during header scanning phase
     blocked_io.to_zarr(
         segy_file_kwargs=segy_file_kwargs,
         output_path=output_path,
         grid_map=grid.map,
         dataset=xr_dataset,
         data_variable_name=default_variable_name,
     )
+
+    # Store the final CRC32C checksum in the Zarr store attributes only when calculated
+    if trace_data_crc32c is not None:
+        # The trace_data_crc32c is now the full file CRC32C from DistributedCRC32C
+        final_crc32c = trace_data_crc32c
+
+        storage_options = _normalize_storage_options(output_path)
+        zarr_group = zarr_open_group(output_path.as_posix(), mode="a", storage_options=storage_options)
+        zarr_group.attrs.update(
+            {
+                "segy_input_crc32c": final_crc32c,  # Store as integer, not hex string
+                "crc32c_algorithm": "CRC32C",
+                "checksum_scope": "full_file",
+                "checksum_library": "google-crc32c",
+            }
+        )

src/mdio/segy/_workers.py

@@ -32,8 +32,9 @@ def header_scan_worker(
     segy_file_kwargs: SegyFileArguments,
     trace_range: tuple[int, int],
     subset: list[str] | None = None,
-) -> HeaderArray:
-    """Header scan worker.
+    calculate_checksum: bool = False,
+) -> 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.
@@ -42,9 +43,11 @@ 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 parsed from SEG-Y library.
+        HeaderArray if calculate_checksum is False, otherwise tuple of (HeaderArray, checksum_info)
+        where checksum_info is (byte_offset, crc32c, byte_length).
     """
     segy_file = SegyFileWrapper(**segy_file_kwargs)
 
@@ -70,7 +73,27 @@ def header_scan_worker(
     # (singleton) so we can concat and assign stuff later.
     trace_header = np.array(trace_header, dtype=new_dtype, ndmin=1)
 
-    return HeaderArray(trace_header)  # wrap back so we can use aliases
+    if not calculate_checksum:
+        return HeaderArray(trace_header)  # wrap back so we can use aliases
+
+    # Calculate checksum from the raw bytes ALREADY IN MEMORY (NO ADDITIONAL I/O!)
+    raw_bytes = traces.trace_buffer_array.tobytes()
+    crc = google_crc32c.Checksum(raw_bytes)
+    partial_crc32c = int.from_bytes(crc.digest(), byteorder="big")
+
+    # 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 trace_worker(  # noqa: PLR0913

src/mdio/segy/parsers.py

@@ -10,8 +10,11 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from crc32c_dist_rs import DistributedCRC32C
 from psutil import cpu_count
+from segy import SegyFile
 from tqdm.auto import tqdm
+from upath import UPath
 
 from mdio.segy._workers import header_scan_worker
 
@@ -23,13 +26,32 @@
 default_cpus = cpu_count(logical=True)
 
 
-def parse_headers(
+def _finalize_checksum(
+    results: list[tuple[HeaderArray, tuple[int, int, int]]], combiner: DistributedCRC32C
+) -> tuple[HeaderArray, int]:
+    """Finalize the checksum from the results of the header scan."""
+    headers: list[HeaderArray] = []
+    for result in results:
+        headers.append(result[0])
+        byte_offset, partial_crc, byte_length = result[1]
+        combiner.add_fragment(byte_offset, byte_length, partial_crc)
+    combined_crc = combiner.try_finalize()
+    if combined_crc is None:
+        msg = "Failed to finalize CRC32C - file may not be fully covered"
+        raise ValueError(msg)
+
+    # Merge headers and return with checksum
+    return np.concatenate(headers), combined_crc
+
+
+def parse_headers(  # noqa: PLR0913
     segy_file_kwargs: SegyFileArguments,
     num_traces: int,
     subset: list[str] | None = None,
     block_size: int = 10000,
     progress_bar: bool = True,
-) -> HeaderArray:
+    calculate_checksum: bool = True,
+) -> HeaderArray | tuple[HeaderArray, int]:
     """Read and parse given `byte_locations` from SEG-Y file.
 
     Args:
@@ -38,11 +60,32 @@ def parse_headers(
         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.
+        calculate_checksum: If True, also calculate CRC32C checksum for all trace data.
 
     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.
+        HeaderArray if calculate_checksum is False.
+        Tuple of (HeaderArray, combined_crc32c) if calculate_checksum is True.
     """
+    # Use UPath for cloud/filesystem agnostic reading
+    path = UPath(segy_file_kwargs["url"])
+    raw_bytes = path.fs.read_block(
+        fn=str(path),
+        offset=0,
+        length=3600,
+    )
+
+    # Calculate trace size to determine total data length
+    # We need to open the file briefly to get the spec
+    sf = SegyFile(**segy_file_kwargs)
+    trace_header_size = sf.spec.trace.header.itemsize
+    sample_size = 4  # This will always be a 4-byte float
+    num_samples = len(sf.sample_labels)
+    trace_size = trace_header_size + (num_samples * sample_size)
+
+    # Total length is header (3600) + trace data
+    total_len = 3600 + num_traces * trace_size
+    combiner = DistributedCRC32C(raw_bytes, total_len)
+
     trace_count = num_traces
     n_blocks = int(ceil(trace_count / block_size))
 
@@ -61,18 +104,32 @@ def parse_headers(
     # '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))
+        lazy_work = executor.map(
+            header_scan_worker,
+            repeat(segy_file_kwargs),
+            trace_ranges,
+            repeat(subset),
+            repeat(calculate_checksum),
+        )
 
         if progress_bar is True:
+            desc = (
+                "Scanning SEG-Y & calculating checksum"
+                if calculate_checksum
+                else "Scanning SEG-Y for geometry attributes"
+            )
             lazy_work = tqdm(
                 iterable=lazy_work,
                 total=n_blocks,
-                desc="Scanning SEG-Y for geometry attributes",
+                desc=desc,
                 **tqdm_kw,
             )
 
         # This executes the lazy work.
-        headers: list[HeaderArray] = list(lazy_work)
+        results = list(lazy_work)
+
+    if not calculate_checksum:
+        # Merge headers and return
+        return np.concatenate(results)
 
-    # Merge blocks before return
-    return np.concatenate(headers)
+    return _finalize_checksum(results, combiner)

src/mdio/segy/utilities.py

@@ -32,7 +32,12 @@ def get_grid_plan(  # noqa:  C901, PLR0913
     template: AbstractDatasetTemplate,
     return_headers: bool = False,
     grid_overrides: dict[str, Any] | None = None,
-) -> tuple[list[Dimension], tuple[int, ...]] | tuple[list[Dimension], tuple[int, ...], HeaderArray]:
+    calculate_checksum: bool = False,
+) -> (
+    tuple[list[Dimension], tuple[int, ...]]
+    | tuple[list[Dimension], tuple[int, ...], HeaderArray]
+    | tuple[list[Dimension], tuple[int, ...], HeaderArray, int]
+):
     """Infer dimension ranges, and increments.
 
     Generates multiple dimensions with the following steps:
@@ -48,22 +53,31 @@ def get_grid_plan(  # noqa:  C901, PLR0913
         template: MDIO template where coordinate names and domain will be taken.
         return_headers: Option to return parsed headers with `Dimension` objects. Default is False.
         grid_overrides: Option to add grid overrides. See main documentation.
+        calculate_checksum: Option to calculate CRC32C checksum during scanning. Default is False.
 
     Returns:
-        All index dimensions and chunksize or dimensions and chunksize together with header values.
+        All index dimensions and chunksize, optionally with header values and/or checksum.
     """
     if grid_overrides is None:
         grid_overrides = {}
 
     # Keep only dimension and non-dimension coordinates excluding the vertical axis
     horizontal_dimensions = template.dimension_names[:-1]
     horizontal_coordinates = horizontal_dimensions + template.coordinate_names
-    headers_subset = parse_headers(
+    parse_result = parse_headers(
         segy_file_kwargs=segy_file_kwargs,
         num_traces=segy_file_info.num_traces,
         subset=horizontal_coordinates,
+        calculate_checksum=calculate_checksum,
     )
 
+    # Unpack result based on whether checksum was calculated
+    if calculate_checksum:
+        headers_subset, trace_data_crc32c = parse_result
+    else:
+        headers_subset = parse_result
+        trace_data_crc32c = None
+
     # Handle grid overrides.
     override_handler = GridOverrider()
     headers_subset, horizontal_coordinates, chunksize = override_handler.run(
@@ -86,8 +100,12 @@ def get_grid_plan(  # noqa:  C901, PLR0913
     vertical_dim = Dimension(coords=sample_labels, name=template.trace_domain)
     dimensions.append(vertical_dim)
 
+    if return_headers and calculate_checksum:
+        return dimensions, chunksize, headers_subset, trace_data_crc32c
     if return_headers:
         return dimensions, chunksize, headers_subset
+    if calculate_checksum:
+        return dimensions, chunksize, None, trace_data_crc32c
 
     return dimensions, chunksize