Merge branch 'main' into 624_prestack_schema

Author: BrianMichell

.github/workflows/constraints.txt

@@ -1,2 +1,2 @@
-bump-my-version==1.2.3
+bump-my-version==1.2.4
 nox==2025.5.1

docs/requirements.txt

@@ -1,4 +1,4 @@
-aiohttp==3.12.15
+aiohttp==3.13.0
 autodoc-pydantic==2.2.0
 furo==2025.9.25
 linkify-it-py==2.0.3

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "multidimio"
-version = "1.0.5"
+version = "1.0.6"
 description = "Cloud-native, scalable, and user-friendly multi dimensional energy data!"
 authors = [{ name = "Altay Sansal", email = "[email protected]" }]
 requires-python = ">=3.11,<3.14"
@@ -18,24 +18,24 @@ classifiers = [
 ]
 
 dependencies = [
-    "click>=8.2.1",
+    "click>=8.3.0",
     "click-params>=0.5.0",
-    "dask>=2025.9.0",
+    "dask>=2025.9.1",
     "fsspec>=2025.9.0",
     "pint>=0.25.0",
-    "psutil>=7.0.0",
-    "pydantic>=2.11.9",
+    "psutil>=7.1.0",
+    "pydantic>=2.12.0",
     "rich>=14.1.0",
-    "segy>=0.5.2",
+    "segy>=0.5.3",
     "tqdm>=4.67.1",
-    "universal-pathlib>=0.2.6",
-    "xarray>=2025.9.1",
+    "universal-pathlib>=0.3.3",
+    "xarray>=2025.10.1",
     "zarr>=3.1.3",
 ]
 
 [project.optional-dependencies]
 cloud = ["s3fs>=2025.9.0", "gcsfs>=2025.9.0", "adlfs>=2025.8.0"]
-distributed = ["distributed>=2025.9.0", "bokeh>=3.8.0"]
+distributed = ["distributed>=2025.9.1", "bokeh>=3.8.0"]
 lossy = ["zfpy>=1.0.1"]
 
 [project.urls]
@@ -48,9 +48,9 @@ mdio = "mdio.__main__:main"
 
 [dependency-groups]
 dev = [
-    "ruff>=0.13.0",
-    "coverage[toml]>=7.10.6",
-    "mypy>=1.18.1",
+    "ruff>=0.14.0",
+    "coverage[toml]>=7.10.7",
+    "mypy>=1.18.2",
     "pre-commit>=4.3.0",
     "pre-commit-hooks>=6.0.0",
     "pytest>=8.4.2",
@@ -61,9 +61,9 @@ dev = [
 ]
 
 docs = [
-    "aiohttp>=3.12.15",
+    "aiohttp>=3.13.0",
     "autodoc-pydantic>=2.2.0",
-    "furo>=2025.7.19",
+    "furo>=2025.9.25",
     "linkify-it-py>=2.0.3",
     "matplotlib>=3.10.6",
     "myst-nb>=1.3.0",
@@ -181,7 +181,7 @@ init_typed = true
 warn_required_dynamic_aliases = true
 
 [tool.bumpversion]
-current_version = "1.0.5"
+current_version = "1.0.6"
 allow_dirty = true
 commit = false
 tag = false

src/mdio/converters/segy.py

@@ -4,13 +4,13 @@
 
 import base64
 import logging
+import multiprocessing as mp
 import os
-from dataclasses import dataclass
+from concurrent.futures import ProcessPoolExecutor
 from typing import TYPE_CHECKING
 
 import numpy as np
 import zarr
-from segy import SegyFile
 from segy.config import SegyFileSettings
 from segy.config import SegyHeaderOverrides
 from segy.standards.codes import MeasurementSystem as SegyMeasurementSystem
@@ -37,9 +37,10 @@
 from mdio.core.utils_write import MAX_SIZE_LIVE_MASK
 from mdio.core.utils_write import get_constrained_chunksize
 from mdio.segy import blocked_io
+from mdio.segy._workers import SegyFileInfo
+from mdio.segy._workers import info_worker
 from mdio.segy.scalar import SCALE_COORDINATE_KEYS
 from mdio.segy.scalar import _apply_coordinate_scalar
-from mdio.segy.scalar import _get_coordinate_scalar
 from mdio.segy.utilities import get_grid_plan
 
 if TYPE_CHECKING:
@@ -54,6 +55,7 @@
     from mdio.builder.schemas import Dataset
     from mdio.builder.templates.abstract_dataset_template import AbstractDatasetTemplate
     from mdio.core.dimension import Dimension
+    from mdio.segy._workers import SegyFileArguments
 
 logger = logging.getLogger(__name__)
 
@@ -135,37 +137,9 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
         raise GridTraceSparsityError(grid.shape, num_traces, msg)
 
 
-@dataclass
-class SegyFileHeaderDump:
-    """Segy metadata information."""
-
-    text_header: str
-    binary_header_dict: dict
-    raw_binary_headers: bytes
-
-
-def _get_segy_file_header_dump(segy_file: SegyFile) -> SegyFileHeaderDump:
-    """Reads information from a SEG-Y file."""
-    text_header = segy_file.text_header
-
-    raw_binary_headers: bytes = segy_file.fs.read_block(
-        fn=segy_file.url,
-        offset=segy_file.spec.binary_header.offset,
-        length=segy_file.spec.binary_header.itemsize,
-    )
-
-    # We read here twice, but it's ok for now. Only 400-bytes.
-    binary_header_dict = segy_file.binary_header.to_dict()
-
-    return SegyFileHeaderDump(
-        text_header=text_header,
-        binary_header_dict=binary_header_dict,
-        raw_binary_headers=raw_binary_headers,
-    )
-
-
 def _scan_for_headers(
-    segy_file: SegyFile,
+    segy_file_kwargs: SegyFileArguments,
+    segy_file_info: SegyFileInfo,
     template: AbstractDatasetTemplate,
     grid_overrides: dict[str, Any] | None = None,
 ) -> tuple[list[Dimension], SegyHeaderArray]:
@@ -176,7 +150,8 @@ def _scan_for_headers(
     """
     full_chunk_size = template.full_chunk_size
     segy_dimensions, chunk_size, segy_headers = get_grid_plan(
-        segy_file=segy_file,
+        segy_file_kwargs=segy_file_kwargs,
+        segy_file_info=segy_file_info,
         return_headers=True,
         template=template,
         chunksize=full_chunk_size,
@@ -192,12 +167,29 @@ def _scan_for_headers(
     return segy_dimensions, segy_headers
 
 
-def _build_and_check_grid(segy_dimensions: list[Dimension], num_traces: int, segy_headers: SegyHeaderArray) -> Grid:
+def _read_segy_file_info(segy_file_kwargs: SegyFileArguments) -> SegyFileInfo:
+    """Read SEG-Y file in a separate process.
+
+    This is an ugly workaround for Zarr issues 3487 'Explicitly using fsspec and zarr FsspecStore causes
+    RuntimeError "Task attached to a different loop"'
+    """
+    # TODO (Dmitriy Repin): when Zarr issue 3487 is resolved, we can remove this workaround
+    # https://github.com/zarr-developers/zarr-python/issues/3487
+    with ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as executor:
+        future = executor.submit(info_worker, segy_file_kwargs)
+        return future.result()
+
+
+def _build_and_check_grid(
+    segy_dimensions: list[Dimension],
+    segy_file_info: SegyFileInfo,
+    segy_headers: SegyHeaderArray,
+) -> Grid:
     """Build and check the grid from the SEG-Y headers and dimensions.
 
     Args:
         segy_dimensions: List of of all SEG-Y dimensions to build grid from.
-        num_traces: Number of traces in the SEG-Y file.
+        segy_file_info: SegyFileInfo instance containing the SEG-Y file information.
         segy_headers: Headers read in from SEG-Y file for building the trace map.
 
     Returns:
@@ -207,6 +199,7 @@ def _build_and_check_grid(segy_dimensions: list[Dimension], num_traces: int, seg
         GridTraceCountError: If number of traces in SEG-Y file does not match the parsed grid
     """
     grid = Grid(dims=segy_dimensions)
+    num_traces = segy_file_info.num_traces
     grid_density_qc(grid, num_traces)
     grid.build_map(segy_headers)
     # Check grid validity by comparing trace numbers
@@ -303,9 +296,9 @@ def populate_non_dim_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _get_horizontal_coordinate_unit(segy_info: SegyFileHeaderDump) -> LengthUnitModel | None:
+def _get_horizontal_coordinate_unit(segy_file_info: SegyFileInfo) -> LengthUnitModel | None:
     """Get the coordinate unit from the SEG-Y headers."""
-    measurement_system_code = int(segy_info.binary_header_dict[MEASUREMENT_SYSTEM_KEY])
+    measurement_system_code = int(segy_file_info.binary_header_dict[MEASUREMENT_SYSTEM_KEY])
 
     if measurement_system_code not in (1, 2):
         logger.warning(
@@ -359,19 +352,19 @@ def _populate_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_header_dump: SegyFileHeaderDump) -> xr_Dataset:
+def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo) -> xr_Dataset:
     save_file_header = os.getenv("MDIO__IMPORT__SAVE_SEGY_FILE_HEADER", "") in ("1", "true", "yes", "on")
     if not save_file_header:
         return xr_dataset
 
     expected_rows = 40
     expected_cols = 80
 
-    text_header_rows = segy_file_header_dump.text_header.splitlines()
+    text_header_rows = segy_file_info.text_header.splitlines()
     text_header_cols_bad = [len(row) != expected_cols for row in text_header_rows]
 
     if len(text_header_rows) != expected_rows:
-        err = f"Invalid text header count: expected {expected_rows}, got {len(segy_file_header_dump.text_header)}"
+        err = f"Invalid text header count: expected {expected_rows}, got {len(segy_file_info.text_header)}"
         raise ValueError(err)
 
     if any(text_header_cols_bad):
@@ -381,12 +374,12 @@ def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_header_dump: SegyFi
     xr_dataset["segy_file_header"] = ((), "")
     xr_dataset["segy_file_header"].attrs.update(
         {
-            "textHeader": segy_file_header_dump.text_header,
-            "binaryHeader": segy_file_header_dump.binary_header_dict,
+            "textHeader": segy_file_info.text_header,
+            "binaryHeader": segy_file_info.binary_header_dict,
         }
     )
     if os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on"):
-        raw_binary_base64 = base64.b64encode(segy_file_header_dump.raw_binary_headers).decode("ascii")
+        raw_binary_base64 = base64.b64encode(segy_file_info.raw_binary_headers).decode("ascii")
         xr_dataset["segy_file_header"].attrs.update({"rawBinaryHeader": raw_binary_base64})
 
     return xr_dataset
@@ -482,6 +475,21 @@ def determine_target_size(var_type: str) -> int:
     ds.variables[index].metadata.chunk_grid = chunk_grid
 
 
+def _validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatasetTemplate) -> None:
+    """Validate that the SegySpec has all required fields in the MDIO template."""
+    header_fields = {field.name for field in segy_spec.trace.header.fields}
+
+    required_fields = set(mdio_template._dim_names[:-1]) | set(mdio_template._coord_names) | {"coordinate_scalar"}
+    missing_fields = required_fields - header_fields
+
+    if missing_fields:
+        err = (
+            f"Required fields {sorted(missing_fields)} for template {mdio_template.name} "
+            f"not found in the provided segy_spec"
+        )
+        raise ValueError(err)
+
+
 def segy_to_mdio(  # noqa PLR0913
     segy_spec: SegySpec,
     mdio_template: AbstractDatasetTemplate,
@@ -507,6 +515,8 @@ def segy_to_mdio(  # noqa PLR0913
     Raises:
         FileExistsError: If the output location already exists and overwrite is False.
     """
+    _validate_spec_in_template(segy_spec, mdio_template)
+
     input_path = _normalize_path(input_path)
     output_path = _normalize_path(output_path)
 
@@ -515,17 +525,21 @@ def segy_to_mdio(  # noqa PLR0913
         raise FileExistsError(err)
 
     segy_settings = SegyFileSettings(storage_options=input_path.storage_options)
-    segy_file = SegyFile(
-        url=input_path.as_posix(),
-        spec=segy_spec,
-        settings=segy_settings,
-        header_overrides=segy_header_overrides,
+    segy_file_kwargs: SegyFileArguments = {
+        "url": input_path.as_posix(),
+        "spec": segy_spec,
+        "settings": segy_settings,
+        "header_overrides": segy_header_overrides,
+    }
+    segy_file_info = _read_segy_file_info(segy_file_kwargs)
+
+    segy_dimensions, segy_headers = _scan_for_headers(
+        segy_file_kwargs,
+        segy_file_info,
+        template=mdio_template,
+        grid_overrides=grid_overrides,
     )
-    segy_info: SegyFileHeaderDump = _get_segy_file_header_dump(segy_file)
-
-    segy_dimensions, segy_headers = _scan_for_headers(segy_file, mdio_template, grid_overrides)
-
-    grid = _build_and_check_grid(segy_dimensions, segy_file.num_traces, segy_headers)
+    grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
@@ -537,7 +551,7 @@ def segy_to_mdio(  # noqa PLR0913
             logger.warning("MDIO__IMPORT__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_info)
+    horizontal_unit = _get_horizontal_coordinate_unit(segy_file_info)
     mdio_ds: Dataset = mdio_template.build_dataset(
         name=mdio_template.name,
         sizes=grid.shape,
@@ -554,15 +568,14 @@ def segy_to_mdio(  # noqa PLR0913
 
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
 
-    coordinate_scalar = _get_coordinate_scalar(segy_file)
     xr_dataset, drop_vars_delayed = _populate_coordinates(
         dataset=xr_dataset,
         grid=grid,
         coords=non_dim_coords,
-        horizontal_coordinate_scalar=coordinate_scalar,
+        horizontal_coordinate_scalar=segy_file_info.coordinate_scalar,
     )
 
-    xr_dataset = _add_segy_file_headers(xr_dataset, segy_info)
+    xr_dataset = _add_segy_file_headers(xr_dataset, segy_file_info)
 
     xr_dataset.trace_mask.data[:] = grid.live_mask
     # IMPORTANT: Do not drop the "trace_mask" here, as it will be used later in
@@ -583,7 +596,7 @@ def segy_to_mdio(  # noqa PLR0913
     # This is an memory-expensive and time-consuming read-write operation
     # performed in chunks to save the memory
     blocked_io.to_zarr(
-        segy_file=segy_file,
+        segy_file_kwargs=segy_file_kwargs,
         output_path=output_path,
         grid_map=grid.map,
         dataset=xr_dataset,

src/mdio/core/indexing.py

@@ -74,11 +74,13 @@ def __next__(self) -> dict[str, slice]:
             # We build slices here. It is dimension agnostic
             current_start = next(self._ranges)
 
-            # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
-            # https://github.com/TGSAI/mdio-python/issues/586
             start_indices = tuple(dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True))
 
-            stop_indices = tuple((dim + 1) * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True))
+            # Calculate stop indices, making the last slice fit the data exactly
+            stop_indices = tuple(
+                min((dim + 1) * chunk, self.arr_shape[i])
+                for i, (dim, chunk) in enumerate(zip(current_start, self.len_chunks, strict=True))
+            )
 
             slices = tuple(slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True))