Fix to_zarr function signature to not be changed, don't reopen the zarr store

Author: BrianMichell

src/mdio/converters/segy.py

@@ -354,7 +354,6 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915
         ...     grid_overrides={"HasDuplicates": True},
         ... )
     """
-    print("Entering segy_to_mdio")
     index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))]
     index_types = index_types or ["int32"] * len(index_bytes)
 
@@ -379,7 +378,7 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915
     num_traces = segy.num_traces
 
     # Index the dataset using a spec that interprets the user provided index headers.
-    index_fields: list[HeaderField] = []
+    index_fields = []
     for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True):
         index_fields.append(HeaderField(name=name, byte=byte, format=format_))
     mdio_spec_grid = mdio_spec.customize(trace_header_fields=index_fields)
@@ -489,16 +488,11 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915
     write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict())
 
     # Write traces
-    zarr_root = mdio_path_or_buffer  # the same path you passed earlier to create_empty
-    data_var = f"data/chunked_{suffix}"
-    header_var = f"metadata/chunked_{suffix}_trace_headers"
-
     stats = blocked_io.to_zarr(
         segy_file=segy,
         grid=grid,
-        zarr_root_path=zarr_root,
-        data_var_path=data_var,
-        header_var_path=header_var,
+        data_array=data_array,
+        header_array=header_array,
     )
 
     # Write actual stats

src/mdio/segy/blocked_io.py

@@ -1,32 +1,5 @@
 """Functions for doing blocked I/O from SEG-Y."""
 
-# from __future__ import annotations
-
-# import multiprocessing as mp
-# import os
-# from concurrent.futures import ProcessPoolExecutor
-# from itertools import repeat
-# from typing import TYPE_CHECKING, Any
-
-# import numpy as np
-# from psutil import cpu_count
-# from tqdm.auto import tqdm
-# import zarr
-
-# # from mdio.core.indexing import ChunkIterator
-# # from mdio.segy._workers import trace_worker
-
-# from mdio.core.indexing import ChunkIterator
-# from mdio.segy._workers import trace_worker
-# from mdio.segy.creation import SegyPartRecord
-# from mdio.segy.creation import concat_files
-# from mdio.segy.creation import serialize_to_segy_stack
-# from mdio.segy.utilities import find_trailing_ones_index
-
-# if TYPE_CHECKING:
-#     from segy import SegyFile
-#     from mdio.core import Grid
-
 from __future__ import annotations
 
 import multiprocessing as mp
@@ -35,6 +8,7 @@
 from itertools import repeat
 from pathlib import Path
 from typing import TYPE_CHECKING
+from typing import Any
 
 import numpy as np
 from dask.array import Array
@@ -55,61 +29,33 @@
     from numpy.typing import NDArray
     from segy import SegyFactory
     from segy import SegyFile
+    from zarr import Array as ZarrArray
 
     from mdio.core import Grid
 
 default_cpus = cpu_count(logical=True)
 
 
-def _worker_reopen(
-    zarr_root_path: str,
-    data_var_path: str,
-    header_var_path: str,
-    segy_file: SegyFile,
-    grid: Grid,
-    chunk_indices: tuple[slice, ...],
-) -> tuple[Any, ...] | None:
-    """
-    Worker function that reopens the Zarr store in this process,
-    obtains fresh array handles, and calls the real trace_worker.
-    """
-    root = zarr.open_group(zarr_root_path, mode="r+")
-    data_arr = root[data_var_path]
-    header_arr = root[header_var_path]
-    result = trace_worker(segy_file, data_arr, header_arr, grid, chunk_indices)
-    root.store.close()
-    return result
-
-
 def to_zarr(
     segy_file: SegyFile,
     grid: Grid,
-    zarr_root_path: str,
-    data_var_path: str,
-    header_var_path: str,
+    data_array: ZarrArray,
+    header_array: ZarrArray,
 ) -> dict[str, Any]:
     """Blocked I/O from SEG-Y to chunked `zarr.core.Array`.
 
-    Each worker reopens the Zarr store independently to avoid lock contention when writing.
-
     Args:
         segy_file: SEG-Y file instance.
         grid: mdio.Grid instance.
-        zarr_root_path: Filesystem path (or URI) to the root of the MDIO Zarr store.
-        data_var_path: Path within the Zarr group for the data array (e.g., "data/chunked_012").
-        header_var_path: Path within the Zarr group for the header array
-                         (e.g., "metadata/chunked_012_trace_headers").
+        data_array: Zarr array for storing trace data.
+        header_array: Zarr array for storing trace headers.
 
     Returns:
         Global statistics for the SEG-Y as a dictionary.
     """
-    # Open Zarr store only in the main process to retrieve shape/metadata
-    root = zarr.open_group(zarr_root_path, mode="r")
-    data_array_meta = root[data_var_path]  # only for shape info
     # Initialize chunk iterator (returns next chunk slice indices each iteration)
-    chunker = ChunkIterator(data_array_meta, chunk_samples=False)
+    chunker = ChunkIterator(data_array, chunk_samples=False)
     num_chunks = len(chunker)
-    root.store.close()  # close immediately
 
     # Determine number of workers
     num_cpus_env = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
@@ -125,11 +71,10 @@ def to_zarr(
     # Launch multiprocessing pool
     with ProcessPoolExecutor(max_workers=num_workers, mp_context=context) as executor:
         lazy_work = executor.map(
-            _worker_reopen,
-            repeat(zarr_root_path),
-            repeat(data_var_path),
-            repeat(header_var_path),
+            trace_worker,
             repeat(segy_file),
+            repeat(data_array),
+            repeat(header_array),
             repeat(grid),
             chunker,
             chunksize=pool_chunksize,
@@ -166,7 +111,6 @@ def to_zarr(
     return {"mean": glob_mean, "std": glob_std, "rms": glob_rms, "min": glob_min, "max": glob_max}
 
 
-
 def segy_record_concat(
     block_records: NDArray,
     file_root: str,