legend-pydataobj 1.10.2__tar.gz → 1.11.0__tar.gz

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: legend_pydataobj
-Version: 1.10.2
+Version: 1.11.0
 Summary: LEGEND Python Data Objects
 Author: The LEGEND Collaboration
 Maintainer: The LEGEND Collaboration

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/legend_pydataobj.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: legend_pydataobj
-Version: 1.10.2
+Version: 1.11.0
 Summary: LEGEND Python Data Objects
 Author: The LEGEND Collaboration
 Maintainer: The LEGEND Collaboration

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/lgdo/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.10.2'
-__version_tuple__ = version_tuple = (1, 10, 2)
+__version__ = version = '1.11.0'
+__version_tuple__ = version_tuple = (1, 11, 0)

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/lgdo/cli.py

@@ -7,7 +7,7 @@ import fnmatch
 import logging
 import sys
 
-from . import Array, Table, VectorOfVectors, __version__, lh5
+from . import Array, Scalar, Struct, Table, VectorOfVectors, __version__, lh5
 from . import logging as lgdogging  # eheheh
 
 log = logging.getLogger(__name__)
@@ -212,6 +212,7 @@ Exclude the /data/table1/col1 Table column:
     store = lh5.LH5Store()
     h5f0 = store.gimme_file(file0)
     lgdos = {}
+    lgdo_structs = {}
     # loop over object list in the first file
     for name in obj_list:
         # now loop over groups starting from root
@@ -222,7 +223,7 @@ Exclude the /data/table1/col1 Table column:
             if current in lgdos:
                 break
 
-            # not even an LGDO!
+            # not even an LGDO (i.e. a plain HDF5 group)!
             if "datatype" not in h5f0[current].attrs:
                 continue
 
@@ -232,14 +233,30 @@ Exclude the /data/table1/col1 Table column:
                 # read all!
                 obj, _ = store.read(current, h5f0)
                 lgdos[current] = obj
+            elif isinstance(obj, Struct):
+                # structs might be used in a "group-like" fashion (i.e. they might only
+                # contain array-like objects).
+                # note: handle after handling tables, as tables also satisfy this check.
+                lgdo_structs[current] = obj.attrs["datatype"]
+                continue
+            elif isinstance(obj, Scalar):
+                msg = f"cannot concat scalar field {current}"
+                log.warning(msg)
 
             break
 
     msg = f"first-level, array-like objects: {lgdos.keys()}"
     log.debug(msg)
+    msg = f"nested structs: {lgdo_structs.keys()}"
+    log.debug(msg)
 
     h5f0.close()
 
+    if lgdos == {}:
+        msg = "did not find any field to concatenate, exit"
+        log.error(msg)
+        return
+
     # 2. remove (nested) table fields based on obj_list
 
     def _inplace_table_filter(name, table, obj_list):
@@ -298,3 +315,14 @@ Exclude the /data/table1/col1 Table column:
                 _inplace_table_filter(name, obj, obj_list)
 
             store.write(obj, name, args.output, wo_mode="append")
+
+    # 5. reset datatypes of the "group-like" structs
+
+    if lgdo_structs != {}:
+        output_file = store.gimme_file(args.output, mode="a")
+        for struct, struct_dtype in lgdo_structs.items():
+            msg = f"reset datatype of struct {struct} to {struct_dtype}"
+            log.debug(msg)
+
+            output_file[struct].attrs["datatype"] = struct_dtype
+        output_file.close()

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/lgdo/lh5/_serializers/read/composite.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import bisect
 import logging
 import sys
-from collections import defaultdict
 
 import h5py
 import numpy as np
@@ -72,19 +71,8 @@ def _h5_read_lgdo(
             obj_buf=obj_buf,
         )
 
-    # check field_mask and make it a default dict
-    if field_mask is None:
-        field_mask = defaultdict(lambda: True)
-    elif isinstance(field_mask, dict):
-        default = True
-        if len(field_mask) > 0:
-            default = not field_mask[next(iter(field_mask.keys()))]
-        field_mask = defaultdict(lambda: default, field_mask)
-    elif isinstance(field_mask, (list, tuple, set)):
-        field_mask = defaultdict(bool, {field: True for field in field_mask})
-    elif not isinstance(field_mask, defaultdict):
-        msg = "bad field_mask type"
-        raise ValueError(msg, type(field_mask).__name__)
+    # Convert whatever we input into a defaultdict
+    field_mask = utils.build_field_mask(field_mask)
 
     if lgdotype is Struct:
         return _h5_read_struct(
@@ -246,18 +234,16 @@ def _h5_read_struct(
 
     # determine fields to be read out
     all_fields = dtypeutils.get_struct_fields(attrs["datatype"])
-    selected_fields = (
-        [field for field in all_fields if field_mask[field]]
-        if field_mask is not None
-        else all_fields
-    )
+    selected_fields = utils.eval_field_mask(field_mask, all_fields)
 
     # modify datatype in attrs if a field_mask was used
-    attrs["datatype"] = "struct{" + ",".join(selected_fields) + "}"
+    attrs["datatype"] = (
+        "struct{" + ",".join(field for field, _ in selected_fields) + "}"
+    )
 
     # loop over fields and read
     obj_dict = {}
-    for field in selected_fields:
+    for field, submask in selected_fields:
         # support for integer keys
         field_key = int(field) if attrs.get("int_keys") else str(field)
         h5o = h5py.h5o.open(h5g, field.encode("utf-8"))
@@ -269,6 +255,7 @@ def _h5_read_struct(
             n_rows=n_rows,
             idx=idx,
             use_h5idx=use_h5idx,
+            field_mask=submask,
             decompress=decompress,
         )
         h5o.close()
@@ -297,19 +284,15 @@ def _h5_read_table(
 
     # determine fields to be read out
     all_fields = dtypeutils.get_struct_fields(attrs["datatype"])
-    selected_fields = (
-        [field for field in all_fields if field_mask[field]]
-        if field_mask is not None
-        else all_fields
-    )
+    selected_fields = utils.eval_field_mask(field_mask, all_fields)
 
     # modify datatype in attrs if a field_mask was used
-    attrs["datatype"] = "table{" + ",".join(selected_fields) + "}"
+    attrs["datatype"] = "table{" + ",".join(field for field, _ in selected_fields) + "}"
 
     # read out each of the fields
     col_dict = {}
     rows_read = []
-    for field in selected_fields:
+    for field, submask in selected_fields:
         fld_buf = None
         if obj_buf is not None:
             if not isinstance(obj_buf, Table) or field not in obj_buf:
@@ -329,6 +312,7 @@ def _h5_read_table(
             use_h5idx=use_h5idx,
             obj_buf=fld_buf,
             obj_buf_start=obj_buf_start,
+            field_mask=submask,
             decompress=decompress,
         )
         h5o.close()

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/lgdo/lh5/_serializers/read/utils.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import logging
+from collections import defaultdict
+from collections.abc import Collection, Mapping
 
 import h5py
 import numpy as np
@@ -22,6 +24,56 @@ def check_obj_buf_attrs(attrs, new_attrs, fname, oname):
         raise LH5DecodeError(msg, fname, oname)
 
 
+def build_field_mask(field_mask: Mapping[str, bool] | Collection[str]) -> defaultdict:
+    # check field_mask and make it a default dict
+    if field_mask is None:
+        return defaultdict(lambda: True)
+    if isinstance(field_mask, dict):
+        default = True
+        if len(field_mask) > 0:
+            default = not field_mask[next(iter(field_mask.keys()))]
+        return defaultdict(lambda: default, field_mask)
+    if isinstance(field_mask, (list, tuple, set)):
+        return defaultdict(bool, {field: True for field in field_mask})
+    if isinstance(field_mask, defaultdict):
+        return field_mask
+    msg = "bad field_mask type"
+    raise ValueError(msg, type(field_mask).__name__)
+
+
+def eval_field_mask(
+    field_mask: defaultdict, all_fields: list[str]
+) -> list[tuple(str, defaultdict)]:
+    """Get list of fields that need to be loaded along with a sub-field-mask
+    in case we have a nested Table"""
+
+    if field_mask is None:
+        return all_fields
+
+    this_field_mask = defaultdict(field_mask.default_factory)
+    sub_field_masks = {}
+
+    for key, val in field_mask.items():
+        field = key.strip("/")
+        pos = field.find("/")
+        if pos < 0:
+            this_field_mask[field] = val
+        else:
+            sub_field = field[pos + 1 :]
+            field = field[:pos]
+            this_field_mask[field] = True
+            sub_mask = sub_field_masks.setdefault(
+                field, defaultdict(field_mask.default_factory)
+            )
+            sub_mask[sub_field] = val
+
+    return [
+        (field, sub_field_masks.get(field))
+        for field in all_fields
+        if this_field_mask[field]
+    ]
+
+
 def read_attrs(h5o, fname, oname):
     """Read all attributes for an hdf5 dataset or group using low level API
     and return them as a dict. Assume all are strings or scalar types."""
@@ -114,6 +166,7 @@ def read_size_in_bytes(h5o, fname, oname, field_mask=None):
     h5a.read(type_attr)
     type_attr = type_attr.item().decode()
     lgdotype = datatype.datatype(type_attr)
+    field_mask = build_field_mask(field_mask)
 
     # scalars are dim-0 datasets
     if lgdotype in (
@@ -124,24 +177,21 @@ def read_size_in_bytes(h5o, fname, oname, field_mask=None):
     ):
         return int(np.prod(h5o.shape) * h5o.dtype.itemsize)
 
-    # structs don't have rows
-    if lgdotype in (types.Struct, types.Histogram, types.Histogram.Axis):
-        size = 0
-        for key in h5o:
-            obj = h5py.h5o.open(h5o, key)
-            size += read_size_in_bytes(obj, fname, oname, field_mask)
-            obj.close()
-        return size
-
     # tables should have elements with all the same length
-    if lgdotype in (types.Table, types.WaveformTable):
+    if lgdotype in (
+        types.Struct,
+        types.Histogram,
+        types.Histogram.Axis,
+        types.Table,
+        types.WaveformTable,
+    ):
         # read out each of the fields
         size = 0
-        if not field_mask:
-            field_mask = datatype.get_struct_fields(type_attr)
-        for field in field_mask:
+        all_fields = datatype.get_struct_fields(type_attr)
+        selected_fields = eval_field_mask(field_mask, all_fields)
+        for field, submask in selected_fields:
             obj = h5py.h5o.open(h5o, field.encode())
-            size += read_size_in_bytes(obj, fname, field)
+            size += read_size_in_bytes(obj, fname, field, submask)
             obj.close()
         return size
 

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/src/lgdo/lh5/_serializers/read/vector_of_vectors.py

@@ -156,7 +156,7 @@ def _h5_read_vector_of_vectors(
         # grow fd_buf if necessary to hold the data
         fdb_size = fd_buf_start + fd_n_rows
         if len(fd_buf) < fdb_size:
-            fd_buf.nda.resize(fdb_size, refcheck=False)
+            fd_buf.resize(fdb_size)
 
     # now read
     h5o = h5py.h5o.open(h5g, b"flattened_data")

{legend_pydataobj-1.10.2 → legend_pydataobj-1.11.0}/tests/test_cli.py

@@ -149,3 +149,21 @@ def test_lh5concat(lgnd_test_data, tmptestdir):
         assert tbl.packet_id[i] == tbl2.packet_id[i - 10]
         assert np.array_equal(tbl.tracelist[i], tbl2.tracelist[i - 10])
         assert np.array_equal(tbl.waveform.values[i], tbl2.waveform.values[i - 10])
+
+    # test concatenating arrays in structs.
+    infile1 = f"{tmptestdir}/concat_test_struct_0.lh5"
+    tb1 = types.Table(col_dict={"col": types.Array(np.zeros(4))})
+    struct1 = types.Struct({"x": tb1})
+    store.write(struct1, "stp", infile1, wo_mode="overwrite_file")
+
+    infile2 = f"{tmptestdir}/concat_test_struct_1.lh5"
+    tb2 = types.Table(col_dict={"col": types.Array(np.ones(7))})
+    struct2 = types.Struct({"x": tb2})
+    store.write(struct2, "stp", infile2, wo_mode="overwrite_file")
+
+    outfile = f"{tmptestdir}/concat_test_struct_out.lh5"
+    cli.lh5concat(["--output", outfile, "--", infile1, infile2])
+
+    out_stp = store.read("stp", outfile)[0]
+    assert out_stp.attrs["datatype"] == "struct{x}"
+    assert np.all(out_stp.x["col"].nda == np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]))