Skip to content

Commit 01791b2

Browse files
authored
Refactor trace header handling to fix memory efficiency regression (#469)
* Refactor trace header handling to fix memory efficiency regression * headerarray explicit type hint * Move HeaderArray import out of TYPE_CHECKING
1 parent f904855 commit 01791b2

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

src/mdio/segy/_workers.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,14 @@
55
import os
66
from typing import TYPE_CHECKING
77
from typing import Any
8+
from typing import cast
89

910
import numpy as np
11+
from segy.arrays import HeaderArray
1012

1113

1214
if TYPE_CHECKING:
1315
from segy import SegyFile
14-
from segy.arrays import HeaderArray
1516
from zarr import Array
1617

1718
from mdio.core import Grid
@@ -43,9 +44,20 @@ def header_scan_worker(
4344
cloud_native_mode = os.getenv("MDIO__IMPORT__CLOUD_NATIVE", default="False")
4445

4546
if cloud_native_mode.lower() in {"true", "1"}:
46-
return segy_file.trace[slice_].header
47+
trace_header = segy_file.trace[slice_].header
48+
else:
49+
trace_header = segy_file.header[slice_]
4750

48-
return segy_file.header[slice_]
51+
# Get non-void fields from dtype and copy to new array for memory efficiency
52+
fields = trace_header.dtype.fields
53+
non_void_fields = [(name, dtype) for name, (dtype, _) in fields.items()]
54+
new_dtype = np.dtype(non_void_fields)
55+
56+
# Allocate empty memory and assign non-void fields
57+
trace_header_filtered = np.empty_like(trace_header, dtype=new_dtype)
58+
trace_header_filtered[:] = trace_header
59+
60+
return cast(HeaderArray, trace_header_filtered)
4961

5062

5163
def trace_worker(

0 commit comments

Comments
 (0)