Begin breaking changes for v1 ingestion

Author: BrianMichell

src/mdio/converters/segy.py

@@ -26,6 +26,9 @@
 from mdio.segy.compat import mdio_segy_spec
 from mdio.segy.utilities import get_grid_plan
 
+from mdio.core.v1.builder import MDIODatasetBuilder as MDIOBuilder
+from mdio.core.utils_write import get_live_mask_chunksize as live_chunks
+
 if TYPE_CHECKING:
     from collections.abc import Sequence
     from pathlib import Path
@@ -384,6 +387,11 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915, PLR0912
     mdio_spec_grid = mdio_spec.customize(trace_header_fields=index_fields)
     segy_grid = SegyFile(url=segy_path, spec=mdio_spec_grid, settings=segy_settings)
 
+    # print(mdio_spec_grid)
+    # print(index_bytes)
+    # print(index_names)
+    # print(index_types)
+
     dimensions, chunksize, index_headers = get_grid_plan(
         segy_file=segy_grid,
         return_headers=True,
@@ -394,6 +402,71 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915, PLR0912
     grid_density_qc(grid, num_traces)
     grid.build_map(index_headers)
 
+    print(f"Dimensions: {dimensions}")
+    print(f"Chunksize: {chunksize}")
+    print(f"Index headers: {index_headers}")
+
+    builder = MDIOBuilder(name=mdio_path_or_buffer, attributes={"Description": "PH"})
+
+    for dim in dimensions:
+        builder.add_dimension(dim.name, dim.size, data_type=str(dim.coords.dtype))
+
+    # TODO: This name is bad
+    if chunksize is not None:
+        lc = list(chunksize)[:-1]
+    else:
+        lc = list(grid.shape[:-1])
+
+    builder.add_variable(
+        name="live_mask",
+        data_type="bool",
+        coordinates=["live_mask"],
+        dimensions=[dim.name for dim in dimensions[:-1]],
+        metadata={
+            "chunkGrid": {
+                "name": "regular",
+                "configuration": {
+                    "chunkShape": live_chunks(lc)
+                }
+            }
+        }
+    )
+
+    print(f"Chunksize: {chunksize}")
+
+    if chunksize is not None:
+        metadata = {
+            "chunkGrid": {
+                "name": "regular",
+                "configuration": {
+                    "chunkShape": list(chunksize),
+                }
+            }
+        }
+    else:
+        metadata = None
+
+    builder.add_variable(
+        name="seismic",
+        data_type="float32",
+        coordinates=["live_mask"],
+        metadata=metadata,
+    )
+
+    ds = builder.to_mdio(store=mdio_path_or_buffer)
+
+    import json
+    contract=json.loads(builder.build().json())
+    from rich import print as rprint
+    oc = {
+        "metadata": contract["metadata"],
+        "variables": contract["variables"],
+    }
+    rprint(oc)
+    new_coords = {dim.name: dim.coords for dim in dimensions}
+    ds = ds.assign_coords(new_coords)
+    ds.to_mdio(store=mdio_path_or_buffer, mode="r+")
+
     # Check grid validity by ensuring every trace's header-index is within dimension bounds
     valid_mask = np.ones(grid.num_traces, dtype=bool)
     for d_idx in range(len(grid.header_index_arrays)):
@@ -413,57 +486,18 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915, PLR0912
     del valid_mask
     gc.collect()
 
-    if chunksize is None:
-        dim_count = len(index_names) + 1
-        if dim_count == 2:  # noqa: PLR2004
-            chunksize = (512,) * 2
-
-        elif dim_count == 3:  # noqa: PLR2004
-            chunksize = (64,) * 3
-
-        else:
-            msg = (
-                f"Default chunking for {dim_count}-D seismic data is not implemented yet. "
-                "Please explicity define chunk sizes."
-            )
-            raise NotImplementedError(msg)
-
-        suffix = [str(x) for x in range(dim_count)]
-        suffix = "".join(suffix)
-    else:
-        suffix = [dim_chunks if dim_chunks > 0 else None for dim_chunks in chunksize]
-        suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
-        suffix = "".join(suffix)
-
-    compressors = get_compressor(lossless, compression_tolerance)
-    header_dtype = segy.spec.trace.header.dtype.newbyteorder("=")
-    var_conf = MDIOVariableConfig(
-        name=f"chunked_{suffix}",
-        dtype="float32",
-        chunks=chunksize,
-        compressors=compressors,
-        header_dtype=header_dtype,
-    )
-    config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf])
-
-    root_group = create_empty(
-        config,
-        overwrite=overwrite,
-        storage_options=storage_options_output,
-        consolidate_meta=False,
-    )
-    data_group = root_group["data"]
-    meta_group = root_group["metadata"]
-    data_array = data_group[f"chunked_{suffix}"]
-    header_array = meta_group[f"chunked_{suffix}_trace_headers"]
+    live_mask_array = ds.live_mask
+    # Cast to MDIODataArray to access the to_mdio method
+    from mdio.core.v1._overloads import MDIODataArray
+    live_mask_array.__class__ = MDIODataArray
 
-    live_mask_array = meta_group["live_mask"]
-    # 'live_mask_array' has the same first N–1 dims as 'grid.shape[:-1]'
     # Build a ChunkIterator over the live_mask (no sample axis)
     from mdio.core.indexing import ChunkIterator
 
-    chunker = ChunkIterator(live_mask_array, chunk_samples=True)
+    # chunker = ChunkIterator(live_mask_array, chunk_samples=True)
+    chunker = ChunkIterator(ds.live_mask, chunk_samples=True)
     for chunk_indices in chunker:
+        print(f"chunk_indices: {chunk_indices}")
         # chunk_indices is a tuple of N–1 slice objects
         trace_ids = grid.get_traces_for_chunk(chunk_indices)
         if trace_ids.size == 0:
@@ -502,35 +536,42 @@ def segy_to_mdio(  # noqa: PLR0913, PLR0915, PLR0912
         del local_coords
 
         # Write the entire block to Zarr at once
-        live_mask_array.set_basic_selection(selection=chunk_indices, value=block)
+        live_mask_array.loc[chunk_indices] = block
 
         # Free block immediately after writing
         del block
 
         # Force garbage collection periodically to free memory aggressively
         gc.collect()
 
+    live_mask_array.to_mdio(store=mdio_path_or_buffer, mode="r+")
+
     # Final cleanup
     del live_mask_array
     del chunker
     gc.collect()
 
-    nonzero_count = grid.num_traces
+    # nonzero_count = grid.num_traces
+
+    # write_attribute(name="trace_count", zarr_group=root_group, attribute=nonzero_count)
+    # write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n"))
+    # write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict())
 
-    write_attribute(name="trace_count", zarr_group=root_group, attribute=nonzero_count)
-    write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n"))
-    write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict())
+    da = ds.seismic
+    da.__class__ = MDIODataArray
 
     # Write traces
     stats = blocked_io.to_zarr(
         segy_file=segy,
         grid=grid,
-        data_array=data_array,
-        header_array=header_array,
+        # data_array=ds.seismic,
+        data_array=da,
+        # header_array=header_array,
+        mdio_path_or_buffer=mdio_path_or_buffer,
     )
 
     # Write actual stats
-    for key, value in stats.items():
-        write_attribute(name=key, zarr_group=root_group, attribute=value)
+    # for key, value in stats.items():
+    #     write_attribute(name=key, zarr_group=root_group, attribute=value)
 
-    zarr.consolidate_metadata(root_group.store)
+    # zarr.consolidate_metadata(root_group.store)

src/mdio/core/indexing.py

@@ -6,6 +6,8 @@
 import numpy as np
 from zarr import Array
 
+import xarray as xr
+
 
 class ChunkIterator:
     """Iterator for traversing a Zarr array in chunks.
@@ -36,9 +38,21 @@ class ChunkIterator:
         ...
     """
 
-    def __init__(self, array: Array, chunk_samples: bool = True):
-        self.arr_shape = array.shape
-        self.len_chunks = array.chunks
+    def __init__(self, array: Array | xr.DataArray, chunk_samples: bool = True):
+
+        if isinstance(array, xr.DataArray):
+            self.arr_shape = array.shape
+            self.len_chunks = array.encoding.get("chunks", self.arr_shape)  # TODO: Chunks don't appear to be present in the encoding. array.chunks is related to dask chunks.
+
+            print(f"arr_shape: {self.arr_shape}")
+            print(f"len_chunks: {self.len_chunks}")
+
+            print(f"array.encoding: {array.encoding}")
+            print(f"array.chunksizes: {array.chunksizes}")
+
+        else:
+            self.arr_shape = array.shape
+            self.len_chunks = array.chunks
 
         # If chunk_samples is False, set the last dimension's chunk size to its full extent
         if not chunk_samples:

src/mdio/core/v1/builder.py

@@ -268,6 +268,14 @@ def _generate_encodings() -> dict:
         """
         # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray
         # dimension_separator_encoding = V2ChunkKeyEncoding(separator="/").to_dict()
+        
+        # Collect dimension sizes (same approach as _construct_mdio_dataset)
+        dims: dict[str, int] = {}
+        for var in mdio_ds.variables:
+            for d in var.dimensions:
+                if isinstance(d, NamedDimension):
+                    dims[d.name] = d.size
+        
         global_encodings = {}
         for var in mdio_ds.variables:
             fill_value = 0
@@ -276,6 +284,10 @@ def _generate_encodings() -> dict:
             chunks = None
             if var.metadata is not None and var.metadata.chunk_grid is not None:
                 chunks = var.metadata.chunk_grid.configuration.chunk_shape
+            else:
+                # When no chunk_grid is provided, set chunks to shape to avoid chunking
+                dim_names = [d.name if isinstance(d, NamedDimension) else d for d in var.dimensions]
+                chunks = tuple(dims[name] for name in dim_names)
             global_encodings[var.name] = {
                 "chunks": chunks,
                 # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray

src/mdio/segy/_workers.py

@@ -54,9 +54,10 @@ def header_scan_worker(segy_file: SegyFile, trace_range: tuple[int, int]) -> Hea
 def trace_worker(
     segy_file: SegyFile,
     data_array: Array,
-    metadata_array: Array,
+    # metadata_array: Array,
     grid: Grid,
     chunk_indices: tuple[slice, ...],
+    mdio_path_or_buffer: str,
 ) -> tuple[Any, ...] | None:
     """Worker function for multi-process enabled blocked SEG-Y I/O.
 
@@ -89,8 +90,8 @@ def trace_worker(
     # Build a temporary buffer for data and metadata for this chunk
     chunk_shape = tuple(sli.stop - sli.start for sli in chunk_indices[:-1]) + (grid.shape[-1],)
     tmp_data = np.zeros(chunk_shape, dtype=data_array.dtype)
-    meta_shape = tuple(sli.stop - sli.start for sli in chunk_indices[:-1])
-    tmp_metadata = np.zeros(meta_shape, dtype=metadata_array.dtype)
+    # meta_shape = tuple(sli.stop - sli.start for sli in chunk_indices[:-1])
+    # tmp_metadata = np.zeros(meta_shape, dtype=metadata_array.dtype)
 
     # Compute local coordinates within the chunk for each trace
     local_coords: list[np.ndarray] = []
@@ -109,10 +110,10 @@ def trace_worker(
 
     # Populate the temporary buffers
     tmp_data[full_idx] = samples
-    tmp_metadata[tuple(local_coords)] = headers.view(tmp_metadata.dtype)
+    # tmp_metadata[tuple(local_coords)] = headers.view(tmp_metadata.dtype)
 
     # Flush metadata to Zarr
-    metadata_array.set_basic_selection(selection=chunk_indices[:-1], value=tmp_metadata)
+    # metadata_array.set_basic_selection(selection=chunk_indices[:-1], value=tmp_metadata)
 
     # Determine nonzero samples and early-exit if none
     nonzero_mask = samples != 0
@@ -121,7 +122,15 @@ def trace_worker(
         return None
 
     # Flush data to Zarr
-    data_array.set_basic_selection(selection=chunk_indices, value=tmp_data)
+    # data_array.set_basic_selection(selection=chunk_indices, value=tmp_data)
+
+    print(f"Writing data to the underlying array...")
+    print(f"Chunk indices: {chunk_indices}")
+    print(f"tmp_data shape: {tmp_data.shape}")
+
+    data_array.loc[chunk_indices] = tmp_data
+    data_array.to_mdio(store=mdio_path_or_buffer, mode="r+")
+
 
     # Calculate statistics
     flattened_nonzero = samples[nonzero_mask]