Begin isolating checksum logic

Author: BrianMichell

pyproject.toml

@@ -20,10 +20,8 @@ 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",
@@ -36,6 +34,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+checksum = ["google-crc32c>=1.5.0", "crc32c_dist_rs"]
 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"]

src/mdio/converters/segy.py

@@ -529,8 +529,9 @@ def segy_to_mdio(  # noqa PLR0913
     }
     segy_file_info = get_segy_file_info(segy_file_kwargs)
 
-    # Only calculate CRC32C when raw headers are enabled
-    calculate_checksum = os.getenv("MDIO__IMPORT__DO_CRC32C") in ("1", "true", "yes", "on")
+    # Check if checksum calculation should be performed
+    # This checks both the env var and library availability
+    calculate_checksum = should_calculate_checksum()
 
     segy_dimensions, segy_headers, trace_data_crc32c = _scan_for_headers(
         segy_file_kwargs,

src/mdio/segy/_workers.py

@@ -76,10 +76,16 @@ def header_scan_worker(
     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()
-    crc = google_crc32c.Checksum(raw_bytes)
-    partial_crc32c = int.from_bytes(crc.digest(), byteorder="big")
+    partial_crc32c = calculate_bytes_crc32c(raw_bytes)
 
     # Calculate byte offset and length
     trace_header_size = segy_file.spec.trace.header.itemsize

src/mdio/segy/checksum.py

@@ -0,0 +1,165 @@
+"""Checksum utilities for SEG-Y ingestion.
+
+This module provides optional CRC32C checksum functionality for SEG-Y ingestion.
+All checksum-related imports and logic are isolated here to maintain loose coupling
+with the core functionality.
+
+Checksum calculation is controlled by the MDIO__IMPORT__DO_CRC32C environment variable
+and requires optional dependencies (google-crc32c, crc32c_dist_rs).
+
+IMPORTANT: This module uses distributed CRC32C calculation to avoid reading entire files
+into memory. Never add functions that read entire files - this has severe performance
+and cost implications, especially for cloud storage.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+from typing import Any
+
+if TYPE_CHECKING:
+    from segy.arrays import HeaderArray
+
+logger = logging.getLogger(__name__)
+
+# Try to import CRC32C libraries, but gracefully handle if they're not installed
+_CHECKSUM_AVAILABLE = False
+_IMPORT_ERROR_MSG = ""
+
+try:
+    import google_crc32c
+    from crc32c_dist_rs import DistributedCRC32C
+
+    _CHECKSUM_AVAILABLE = True
+except ImportError as e:
+    _IMPORT_ERROR_MSG = (
+        f"CRC32C checksum libraries not available: {e}. "
+        "Install with: pip install multidimio[checksum] or pip install google-crc32c crc32c_dist_rs"
+    )
+    # Define placeholder types for type checking
+    google_crc32c = None  # type: ignore[assignment]
+    DistributedCRC32C = None  # type: ignore[assignment,misc]
+
+
+def is_checksum_available() -> bool:
+    """Check if checksum libraries are available.
+
+    Returns:
+        True if google-crc32c and crc32c_dist_rs are installed, False otherwise.
+    """
+    return _CHECKSUM_AVAILABLE
+
+
+def should_calculate_checksum() -> bool:
+    """Determine if checksum calculation should be performed.
+
+    Checks both the environment variable and library availability.
+
+    Returns:
+        True if MDIO__IMPORT__DO_CRC32C is enabled and libraries are available.
+    """
+    env_enabled = os.getenv("MDIO__IMPORT__DO_CRC32C", "false").lower() in ("1", "true", "yes", "on")
+
+    if env_enabled and not _CHECKSUM_AVAILABLE:
+        logger.warning(
+            "MDIO__IMPORT__DO_CRC32C is enabled but checksum libraries are not available. "
+            "Checksum calculation will be skipped. %s",
+            _IMPORT_ERROR_MSG,
+        )
+        return False
+
+    return env_enabled
+
+
+def require_checksum_libraries() -> None:
+    """Raise an error if checksum libraries are not available.
+
+    Raises:
+        ImportError: If checksum libraries are not installed.
+    """
+    if not _CHECKSUM_AVAILABLE:
+        raise ImportError(_IMPORT_ERROR_MSG)
+
+
+def calculate_bytes_crc32c(data: bytes) -> int:
+    """Calculate CRC32C checksum for a byte array.
+
+    Args:
+        data: Byte array to checksum.
+
+    Returns:
+        CRC32C checksum as integer.
+
+    Raises:
+        ImportError: If checksum libraries are not available.
+    """
+    require_checksum_libraries()
+
+    crc = google_crc32c.Checksum(data)
+    return int.from_bytes(crc.digest(), byteorder="big")
+
+
+def create_distributed_crc32c(initial_bytes: bytes, total_length: int) -> Any:
+    """Create a distributed CRC32C combiner instance.
+
+    Args:
+        initial_bytes: Initial bytes (e.g., file header).
+        total_length: Total expected file length in bytes.
+
+    Returns:
+        DistributedCRC32C instance.
+
+    Raises:
+        ImportError: If checksum libraries are not available.
+    """
+    require_checksum_libraries()
+
+    return DistributedCRC32C(initial_bytes, total_length)
+
+
+def finalize_distributed_checksum(
+    results: list[tuple[HeaderArray, tuple[int, int, int]]], combiner: Any
+) -> tuple[HeaderArray, int]:
+    """Finalize a distributed CRC32C checksum from scan results.
+
+    Args:
+        results: List of (HeaderArray, (byte_offset, partial_crc, byte_length)) tuples.
+        combiner: DistributedCRC32C instance.
+
+    Returns:
+        Tuple of (concatenated HeaderArray, final CRC32C checksum).
+
+    Raises:
+        ValueError: If checksum finalization fails.
+        ImportError: If checksum libraries are not available.
+    """
+    require_checksum_libraries()
+
+    import numpy as np
+
+    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
+
+
+__all__ = [
+    "is_checksum_available",
+    "should_calculate_checksum",
+    "require_checksum_libraries",
+    "calculate_bytes_crc32c",
+    "create_distributed_crc32c",
+    "finalize_distributed_checksum",
+]
+

src/mdio/segy/parsers.py

@@ -8,15 +8,18 @@
 from itertools import repeat
 from math import ceil
 from typing import TYPE_CHECKING
+from typing import Any
 
 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
+from mdio.segy.checksum import create_distributed_crc32c
+from mdio.segy.checksum import finalize_distributed_checksum
+from mdio.segy.checksum import is_checksum_available
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
@@ -26,24 +29,6 @@
 default_cpus = cpu_count(logical=True)
 
 
-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,
@@ -66,25 +51,32 @@ def parse_headers(  # noqa: PLR0913
         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)
+    # Initialize combiner only if checksum calculation is needed and available
+    combiner: Any = None
+    if calculate_checksum:
+        if not is_checksum_available():
+            # Checksum was requested but libraries not available - disable it
+            calculate_checksum = False
+        else:
+            # 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 = create_distributed_crc32c(raw_bytes, total_len)
 
     trace_count = num_traces
     n_blocks = int(ceil(trace_count / block_size))
@@ -132,4 +124,4 @@ def parse_headers(  # noqa: PLR0913
         # Merge headers and return
         return np.concatenate(results)
 
-    return _finalize_checksum(results, combiner)
+    return finalize_distributed_checksum(results, combiner)

tests/integration/test_crc32c_checksum.py

@@ -6,7 +6,6 @@
 
 from pathlib import Path
 
-import google_crc32c
 import pytest
 import zarr
 from tests.integration.test_segy_import_export_masked import COCA_3D_CONF
@@ -22,21 +21,34 @@
 from mdio import segy_to_mdio
 from mdio.builder.template_registry import TemplateRegistry
 from mdio.segy._workers import info_worker
+from mdio.segy.checksum import is_checksum_available
 from mdio.segy.parsers import parse_headers
 
+# Skip all tests in this module if checksum libraries are not available
+pytestmark = pytest.mark.skipif(
+    not is_checksum_available(), reason="CRC32C checksum libraries (google-crc32c, crc32c_dist_rs) not installed"
+)
+
 
 def get_expected_crc32c(segy_path: Path) -> int:
     """Calculate the expected CRC32C checksum using Google crc32c library.
 
     This calculates CRC32C over the entire file (headers + trace data),
     matching how MDIO calculates the checksum.
 
+    NOTE: This is TEST-ONLY code. Never use this in production as it reads
+    the entire file into memory which has significant performance and cost
+    penalties, especially for cloud storage.
+
     Args:
         segy_path: Path to the SEG-Y file
 
     Returns:
         CRC32C checksum as integer
     """
+    # Import here to keep it test-only
+    import google_crc32c
+
     crc = google_crc32c.Checksum()
 
     with segy_path.open("rb") as f: