Resolve data serialization issues

Author: BrianMichell

src/mdio/converters/segy.py

@@ -11,13 +11,15 @@
 from segy.config import SegySettings
 from segy.standards.codes import MeasurementSystem as segy_MeasurementSystem
 from segy.standards.fields.trace import Rev0 as TraceHeaderFieldsRev0
+from xarray import Variable
 
 from mdio.constants import UINT32_MAX
 from mdio.converters.exceptions import EnvironmentFormatError
 from mdio.converters.exceptions import GridTraceCountError
 from mdio.converters.exceptions import GridTraceSparsityError
 from mdio.converters.type_converter import to_structured_type
 from mdio.core.grid import Grid
+from mdio.schemas.v1.dataset_serializer import _get_fill_value
 from mdio.schemas.v1.dataset_serializer import to_xarray_dataset
 from mdio.schemas.v1.units import AllUnits
 from mdio.schemas.v1.units import LengthUnitEnum
@@ -227,7 +229,18 @@ def populate_non_dim_coordinates(
     """Populate the xarray dataset with coordinate variables."""
     not_null = grid.map[:] != UINT32_MAX
     for c_name, c_values in coordinates.items():
-        dataset[c_name].values[not_null] = c_values
+        encodings = dataset[c_name].encoding
+        tmp_coords = np.full(
+            not_null.shape, dtype=dataset[c_name].dtype, fill_value=_get_fill_value(dataset[c_name].dtype)
+        )
+        tmp_coords[not_null] = c_values
+        dataset[c_name].values = tmp_coords
+        dataset[c_name] = Variable(
+            dataset[c_name].dims,
+            tmp_coords,
+            attrs=dataset[c_name].attrs,
+            encoding=encodings,  # Ensure we preserve all of the encodings.
+        )
         drop_vars_delayed.append(c_name)
     return dataset, drop_vars_delayed
 

src/mdio/segy/_workers.py

@@ -22,6 +22,7 @@
 from xarray import Variable
 
 from mdio.constants import UINT32_MAX
+from mdio.schemas.v1.dataset_serializer import _get_fill_value
 from mdio.schemas.v1.stats import CenteredBinHistogram
 from mdio.schemas.v1.stats import SummaryStatistics
 
@@ -125,6 +126,7 @@ def trace_worker(  # noqa: PLR0913
 
     if hdr_key in worker_variables:
         # Create temporary array for headers with the correct shape
+        # TODO(BrianMichell): Implement this better so that we can enable fill values without changing the code. #noqa: TD003
         tmp_headers = np.zeros(not_null.shape, dtype=ds_to_write[hdr_key].dtype)
         tmp_headers[not_null] = traces.header
         # Create a new Variable object to avoid copying the temporary array
@@ -136,7 +138,11 @@ def trace_worker(  # noqa: PLR0913
 
     # Get the sample dimension size from the data variable itself
     sample_dim_size = ds_to_write[data_variable_name].shape[-1]
-    tmp_samples = np.zeros(not_null.shape + (sample_dim_size,), dtype=ds_to_write[data_variable_name].dtype)
+    tmp_samples = np.full(
+        not_null.shape + (sample_dim_size,),
+        dtype=ds_to_write[data_variable_name].dtype,
+        fill_value=_get_fill_value(ds_to_write[data_variable_name].dtype),
+    )
 
     # Assign trace samples to the correct positions
     # We need to handle the fact that traces.sample is (num_traces, num_samples)