legend-pydataobj 1.11.6__py3-none-any.whl → 1.11.8__py3-none-any.whl

lgdo/lh5/concat.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import fnmatch
+import logging
+
+from lgdo.lh5 import LH5Iterator
+
+from .. import Array, Scalar, Struct, Table, VectorOfVectors, lh5
+
+log = logging.getLogger(__name__)
+
+
+def _get_obj_list(
+    lh5_files: list, include_list: list | None = None, exclude_list: list | None = None
+) -> list[str]:
+    """Extract a list of lh5 objects to concatenate.
+
+    Parameters
+    ----------
+    lh5_files
+        list of input files to concatenate.
+    include_list
+        patterns for tables to include.
+    exclude_list
+        patterns for tables to exclude.
+
+    """
+    file0 = lh5_files[0]
+    obj_list_full = set(lh5.ls(file0, recursive=True))
+
+    # let's remove objects with nested LGDOs inside
+    to_remove = set()
+    for name in obj_list_full:
+        if len(fnmatch.filter(obj_list_full, f"{name}/*")) > 1:
+            to_remove.add(name)
+    obj_list_full -= to_remove
+
+    obj_list = set()
+    # now first remove excluded stuff
+    if exclude_list is not None:
+        for exc in exclude_list:
+            obj_list_full -= set(fnmatch.filter(obj_list_full, exc.strip("/")))
+
+    # then make list of included, based on latest list
+    if include_list is not None:
+        for inc in include_list:
+            obj_list |= set(fnmatch.filter(obj_list_full, inc.strip("/")))
+    else:
+        obj_list = obj_list_full
+
+    # sort
+    return sorted(obj_list)
+
+
+def _get_lgdos(file, obj_list):
+    """Get name of LGDO objects."""
+
+    store = lh5.LH5Store()
+    h5f0 = store.gimme_file(file)
+
+    lgdos = []
+    lgdo_structs = {}
+
+    # loop over object list in the first file
+    for name in obj_list:
+        # now loop over groups starting from root
+        current = ""
+        for item in name.split("/"):
+            current = f"{current}/{item}".strip("/")
+
+            if current in lgdos:
+                break
+
+            # not even an LGDO (i.e. a plain HDF5 group)!
+            if "datatype" not in h5f0[current].attrs:
+                continue
+
+            # read as little as possible
+            obj = store.read(current, h5f0, n_rows=1)
+            if isinstance(obj, (Table, Array, VectorOfVectors)):
+                lgdos.append(current)
+
+            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}"
+    log.info(msg)
+
+    msg = f"nested structs: {lgdo_structs}"
+    log.info(msg)
+
+    h5f0.close()
+
+    if lgdos == []:
+        msg = "did not find any field to concatenate, exit"
+        raise RuntimeError(msg)
+
+    return lgdos, lgdo_structs
+
+
+def _inplace_table_filter(name, table, obj_list):
+    """filter objects nested in this LGDO"""
+    skm = fnmatch.filter(obj_list, f"{name}/*")
+    kept = {it.removeprefix(name).strip("/").split("/")[0] for it in skm}
+
+    # now remove fields
+    for k in list(table.keys()):
+        if k not in kept:
+            table.remove_column(k)
+
+    msg = f"fields left in table '{name}': {table.keys()}"
+    log.debug(msg)
+
+    # recurse!
+    for k2, v2 in table.items():
+        if not isinstance(v2, Table):
+            continue
+
+        _inplace_table_filter(f"{name}/{k2}", v2, obj_list)
+
+
+def _remove_nested_fields(lgdos: dict, obj_list: list):
+    """Remove (nested) table fields based on obj_list."""
+
+    for key, val in lgdos.items():
+        if not isinstance(val, Table):
+            continue
+
+        _inplace_table_filter(key, val, obj_list)
+
+
+def lh5concat(
+    lh5_files: list,
+    output: str,
+    overwrite: bool = False,
+    *,
+    include_list: list | None = None,
+    exclude_list: list | None = None,
+) -> None:
+    """Concatenate LGDO Arrays, VectorOfVectors and Tables in LH5 files.
+
+    Parameters
+    ----------
+    lh5_files
+        list of input files to concatenate.
+    output
+        path to the output file
+    include_list
+        patterns for tables to include.
+    exclude_list
+        patterns for tables to exclude.
+    """
+
+    if len(lh5_files) < 2:
+        msg = "you must provide at least two input files"
+        raise RuntimeError(msg)
+
+    # determine list of objects by recursively ls'ing first file
+    obj_list = _get_obj_list(
+        lh5_files, include_list=include_list, exclude_list=exclude_list
+    )
+
+    msg = f"objects matching include patterns {include_list} in {lh5_files[0]}: {obj_list}"
+    log.info(msg)
+
+    lgdos, lgdo_structs = _get_lgdos(lh5_files[0], obj_list)
+    first_done = False
+    store = lh5.LH5Store()
+
+    # loop over lgdo objects
+    for lgdo in lgdos:
+        # iterate over the files
+        for lh5_obj in LH5Iterator(lh5_files, lgdo):
+            data = {lgdo: lh5_obj}
+
+            # remove the nested fields
+            _remove_nested_fields(data, obj_list)
+
+            if first_done is False:
+                msg = f"creating output file {output}"
+                log.info(msg)
+
+                store.write(
+                    data[lgdo],
+                    lgdo,
+                    output,
+                    wo_mode="overwrite_file"
+                    if (overwrite and not first_done)
+                    else "write_safe",
+                )
+                first_done = True
+
+            else:
+                msg = f"appending to {output}"
+                log.info(msg)
+
+                if isinstance(data[lgdo], Table):
+                    _inplace_table_filter(lgdo, data[lgdo], obj_list)
+
+                store.write(data[lgdo], lgdo, output, wo_mode="append")
+
+    if lgdo_structs != {}:
+        output_file = store.gimme_file(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()

lgdo/lh5/core.py

@@ -4,6 +4,7 @@ import bisect
 import inspect
 import sys
 from collections.abc import Mapping, Sequence
+from contextlib import suppress
 from typing import Any
 
 import h5py
@@ -92,8 +93,7 @@ def read(
         will be set to ``True``, while the rest will default to ``False``.
     obj_buf
         Read directly into memory provided in `obj_buf`. Note: the buffer
-        will be expanded to accommodate the data requested. To maintain the
-        buffer length, send in ``n_rows = len(obj_buf)``.
+        will be resized to accommodate the data retrieved.
     obj_buf_start
         Start location in ``obj_buf`` for read. For concatenating data to
         array-like objects.
@@ -106,25 +106,25 @@ def read(
 
     Returns
     -------
-    (object, n_rows_read)
-        `object` is the read-out object `n_rows_read` is the number of rows
-        successfully read out. Essential for arrays when the amount of data
-        is smaller than the object buffer.  For scalars and structs
-        `n_rows_read` will be``1``. For tables it is redundant with
-        ``table.loc``. If `obj_buf` is ``None``, only `object` is returned.
+    object
+        the read-out object
     """
     if isinstance(lh5_file, h5py.File):
         lh5_obj = lh5_file[name]
     elif isinstance(lh5_file, str):
         lh5_file = h5py.File(lh5_file, mode="r", locking=locking)
-        lh5_obj = lh5_file[name]
+        try:
+            lh5_obj = lh5_file[name]
+        except KeyError as ke:
+            err = f"Object {name} not found in file {lh5_file.filename}"
+            raise KeyError(err) from ke
     else:
-        lh5_files = list(lh5_file)
-
-        n_rows_read = 0
-        obj_buf_is_new = False
+        if obj_buf is not None:
+            obj_buf.resize(obj_buf_start)
+        else:
+            obj_buf_start = 0
 
-        for i, h5f in enumerate(lh5_files):
+        for i, h5f in enumerate(lh5_file):
             if (
                 isinstance(idx, (list, tuple))
                 and len(idx) > 0
@@ -146,33 +146,26 @@ def read(
                 idx = np.array(idx[0])[n_rows_to_read_i:] - n_rows_i
             else:
                 idx_i = None
-            n_rows_i = n_rows - n_rows_read
 
-            obj_ret = read(
+            obj_buf_start_i = len(obj_buf) if obj_buf else 0
+            n_rows_i = n_rows - (obj_buf_start_i - obj_buf_start)
+
+            obj_buf = read(
                 name,
                 h5f,
-                start_row,
+                start_row if i == 0 else 0,
                 n_rows_i,
                 idx_i,
                 use_h5idx,
                 field_mask,
                 obj_buf,
-                obj_buf_start,
+                obj_buf_start_i,
                 decompress,
             )
-            if isinstance(obj_ret, tuple):
-                obj_buf, n_rows_read_i = obj_ret
-                obj_buf_is_new = True
-            else:
-                obj_buf = obj_ret
-                n_rows_read_i = len(obj_buf)
 
-            n_rows_read += n_rows_read_i
-            if n_rows_read >= n_rows or obj_buf is None:
-                return obj_buf, n_rows_read
-            start_row = 0
-            obj_buf_start += n_rows_read_i
-        return obj_buf if obj_buf_is_new else (obj_buf, n_rows_read)
+            if obj_buf is None or (len(obj_buf) - obj_buf_start) >= n_rows:
+                return obj_buf
+        return obj_buf
 
     if isinstance(idx, (list, tuple)) and len(idx) > 0 and not np.isscalar(idx[0]):
         idx = idx[0]
@@ -192,8 +185,10 @@ def read(
         obj_buf_start=obj_buf_start,
         decompress=decompress,
     )
+    with suppress(AttributeError):
+        obj.resize(obj_buf_start + n_rows_read)
 
-    return obj if obj_buf is None else (obj, n_rows_read)
+    return obj
 
 
 def write(
@@ -273,11 +268,13 @@ def write(
           end of array is the same as ``append``.
         - ``overwrite_file`` or ``of``: delete file if present prior to
           writing to it. `write_start` should be 0 (its ignored).
-        - ``append_column`` or ``ac``: append columns from an
-          :class:`~.lgdo.table.Table` `obj` only if there is an existing
-          :class:`~.lgdo.table.Table` in the `lh5_file` with the same
-          `name` and :class:`~.lgdo.table.Table.size`. If the sizes don't
-          match, or if there are matching fields, it errors out.
+        - ``append_column`` or ``ac``: append fields/columns from an
+          :class:`~.lgdo.struct.Struct` `obj` (and derived types such as
+          :class:`~.lgdo.table.Table`) only if there is an existing
+          :class:`~.lgdo.struct.Struct` in the `lh5_file` with the same `name`.
+          If there are matching fields, it errors out. If appending to a
+          ``Table`` and the size of the new column is different from the size
+          of the existing table, it errors out.
     write_start
         row in the output file (if already existing) to start overwriting
         from.

lgdo/lh5/iterator.py

@@ -24,7 +24,8 @@ class LH5Iterator(typing.Iterator):
 
     This can be used as an iterator:
 
-    >>> for lh5_obj, i_entry, n_rows in LH5Iterator(...):
+
+    >>> for lh5_obj in LH5Iterator(...):
     >>>    # do the thing!
 
     This is intended for if you are reading a large quantity of data. This
@@ -42,6 +43,8 @@ class LH5Iterator(typing.Iterator):
     In addition to accessing requested data via ``lh5_obj``, several
     properties exist to tell you where that data came from:
 
+    - lh5_it.current_i_entry: get the index within the entry list of the
+      first entry that is currently read
     - lh5_it.current_local_entries: get the entry numbers relative to the
       file the data came from
     - lh5_it.current_global_entries: get the entry number relative to the
@@ -49,9 +52,9 @@ class LH5Iterator(typing.Iterator):
     - lh5_it.current_files: get the file name corresponding to each entry
     - lh5_it.current_groups: get the group name corresponding to each entry
 
-    This class can also be used either for random access:
+    This class can also be used for random access:
 
-    >>> lh5_obj, n_rows = lh5_it.read(i_entry)
+    >>> lh5_obj = lh5_it.read(i_entry)
 
     to read the block of entries starting at i_entry. In case of multiple files
     or the use of an event selection, i_entry refers to a global event index
@@ -65,6 +68,8 @@ class LH5Iterator(typing.Iterator):
         base_path: str = "",
         entry_list: list[int] | list[list[int]] | None = None,
         entry_mask: list[bool] | list[list[bool]] | None = None,
+        i_start: int = 0,
+        n_entries: int | None = None,
         field_mask: dict[str, bool] | list[str] | tuple[str] | None = None,
         buffer_len: int = "100*MB",
         file_cache: int = 10,
@@ -89,6 +94,10 @@ class LH5Iterator(typing.Iterator):
         entry_mask
             mask of entries to read. If a list of arrays is provided, expect
             one for each file. Ignore if a selection list is provided.
+        i_start
+            index of first entry to start at when iterating
+        n_entries
+            number of entries to read before terminating iteration
         field_mask
             mask of which fields to read. See :meth:`LH5Store.read` for
             more details.
@@ -183,7 +192,8 @@ class LH5Iterator(typing.Iterator):
             msg = f"can't open any files from {lh5_files}"
             raise RuntimeError(msg)
 
-        self.n_rows = 0
+        self.i_start = i_start
+        self.n_entries = n_entries
         self.current_i_entry = 0
         self.next_i_entry = 0
 
@@ -317,14 +327,21 @@ class LH5Iterator(typing.Iterator):
                 )
         return self.global_entry_list
 
-    def read(self, i_entry: int) -> tuple[LGDO, int]:
-        """Read the nextlocal chunk of events, starting at i_entry. Return the
-        LH5 buffer and number of rows read."""
-        self.n_rows = 0
-        i_file = np.searchsorted(self.entry_map, i_entry, "right")
+    def read(self, i_entry: int, n_entries: int | None = None) -> LGDO:
+        "Read the nextlocal chunk of events, starting at entry."
+        self.lh5_buffer.resize(0)
+
+        if n_entries is None:
+            n_entries = self.buffer_len
+        elif n_entries == 0:
+            return self.lh5_buffer
+        elif n_entries > self.buffer_len:
+            msg = "n_entries cannot be larger than buffer_len"
+            raise ValueError(msg)
 
         # if file hasn't been opened yet, search through files
         # sequentially until we find the right one
+        i_file = np.searchsorted(self.entry_map, i_entry, "right")
         if i_file < len(self.lh5_files) and self.entry_map[i_file] == np.iinfo("q").max:
             while i_file < len(self.lh5_files) and i_entry >= self._get_file_cumentries(
                 i_file
@@ -332,10 +349,10 @@ class LH5Iterator(typing.Iterator):
                 i_file += 1
 
         if i_file == len(self.lh5_files):
-            return (self.lh5_buffer, self.n_rows)
+            return self.lh5_buffer
         local_i_entry = i_entry - self._get_file_cumentries(i_file - 1)
 
-        while self.n_rows < self.buffer_len and i_file < len(self.file_map):
+        while len(self.lh5_buffer) < n_entries and i_file < len(self.file_map):
             # Loop through files
             local_idx = self.get_file_entrylist(i_file)
             if local_idx is not None and len(local_idx) == 0:
@@ -344,18 +361,17 @@ class LH5Iterator(typing.Iterator):
                 continue
 
             i_local = local_i_entry if local_idx is None else local_idx[local_i_entry]
-            self.lh5_buffer, n_rows = self.lh5_st.read(
+            self.lh5_buffer = self.lh5_st.read(
                 self.groups[i_file],
                 self.lh5_files[i_file],
                 start_row=i_local,
-                n_rows=self.buffer_len - self.n_rows,
+                n_rows=n_entries - len(self.lh5_buffer),
                 idx=local_idx,
                 field_mask=self.field_mask,
                 obj_buf=self.lh5_buffer,
-                obj_buf_start=self.n_rows,
+                obj_buf_start=len(self.lh5_buffer),
             )
 
-            self.n_rows += n_rows
             i_file += 1
             local_i_entry = 0
 
@@ -364,7 +380,7 @@ class LH5Iterator(typing.Iterator):
         if self.friend is not None:
             self.friend.read(i_entry)
 
-        return (self.lh5_buffer, self.n_rows)
+        return self.lh5_buffer
 
     def reset_field_mask(self, mask):
         """Replaces the field mask of this iterator and any friends with mask"""
@@ -375,7 +391,7 @@ class LH5Iterator(typing.Iterator):
     @property
     def current_local_entries(self) -> NDArray[int]:
         """Return list of local file entries in buffer"""
-        cur_entries = np.zeros(self.n_rows, dtype="int32")
+        cur_entries = np.zeros(len(self.lh5_buffer), dtype="int32")
         i_file = np.searchsorted(self.entry_map, self.current_i_entry, "right")
         file_start = self._get_file_cumentries(i_file - 1)
         i_local = self.current_i_entry - file_start
@@ -402,7 +418,7 @@ class LH5Iterator(typing.Iterator):
     @property
     def current_global_entries(self) -> NDArray[int]:
         """Return list of local file entries in buffer"""
-        cur_entries = np.zeros(self.n_rows, dtype="int32")
+        cur_entries = np.zeros(len(self.lh5_buffer), dtype="int32")
         i_file = np.searchsorted(self.entry_map, self.current_i_entry, "right")
         file_start = self._get_file_cumentries(i_file - 1)
         i_local = self.current_i_entry - file_start
@@ -433,7 +449,7 @@ class LH5Iterator(typing.Iterator):
     @property
     def current_files(self) -> NDArray[str]:
         """Return list of file names for entries in buffer"""
-        cur_files = np.zeros(self.n_rows, dtype=object)
+        cur_files = np.zeros(len(self.lh5_buffer), dtype=object)
         i_file = np.searchsorted(self.entry_map, self.current_i_entry, "right")
         file_start = self._get_file_cumentries(i_file - 1)
         i_local = self.current_i_entry - file_start
@@ -455,7 +471,7 @@ class LH5Iterator(typing.Iterator):
     @property
     def current_groups(self) -> NDArray[str]:
         """Return list of group names for entries in buffer"""
-        cur_groups = np.zeros(self.n_rows, dtype=object)
+        cur_groups = np.zeros(len(self.lh5_buffer), dtype=object)
         i_file = np.searchsorted(self.entry_map, self.current_i_entry, "right")
         file_start = self._get_file_cumentries(i_file - 1)
         i_local = self.current_i_entry - file_start
@@ -485,14 +501,19 @@ class LH5Iterator(typing.Iterator):
     def __iter__(self) -> typing.Iterator:
         """Loop through entries in blocks of size buffer_len."""
         self.current_i_entry = 0
-        self.next_i_entry = 0
+        self.next_i_entry = self.i_start
         return self
 
     def __next__(self) -> tuple[LGDO, int, int]:
-        """Read next buffer_len entries and return lh5_table, iterator entry
-        and n_rows read."""
-        buf, n_rows = self.read(self.next_i_entry)
-        self.next_i_entry = self.current_i_entry + n_rows
-        if n_rows == 0:
+        """Read next buffer_len entries and return lh5_table and iterator entry."""
+        n_entries = self.n_entries
+        if n_entries is not None:
+            n_entries = min(
+                self.buffer_len, n_entries + self.i_start - self.next_i_entry
+            )
+
+        buf = self.read(self.next_i_entry, n_entries)
+        if len(buf) == 0:
             raise StopIteration
-        return (buf, self.current_i_entry, n_rows)
+        self.next_i_entry = self.current_i_entry + len(buf)
+        return buf

lgdo/lh5/store.py

@@ -5,21 +5,20 @@ HDF5 files.
 
 from __future__ import annotations
 
-import bisect
 import logging
-import os
 import sys
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from inspect import signature
+from pathlib import Path
 from typing import Any
 
 import h5py
-import numpy as np
 from numpy.typing import ArrayLike
 
 from .. import types
 from . import _serializers, utils
+from .core import read
 
 log = logging.getLogger(__name__)
 
@@ -93,16 +92,16 @@ class LH5Store:
             return self.files[lh5_file]
 
         if self.base_path != "":
-            full_path = os.path.join(self.base_path, lh5_file)
+            full_path = Path(self.base_path) / lh5_file
         else:
-            full_path = lh5_file
+            full_path = Path(lh5_file)
 
-        file_exists = os.path.exists(full_path)
+        file_exists = full_path.exists()
         if mode != "r":
-            directory = os.path.dirname(full_path)
-            if directory != "" and not os.path.exists(directory):
+            directory = full_path.parent
+            if directory != "" and not full_path.parent.exists():
                 log.debug(f"making path {directory}")
-                os.makedirs(directory)
+                directory.mkdir(parents=True, exist_ok=True)
 
         if mode == "r" and not file_exists:
             msg = f"file {full_path} not found"
@@ -155,7 +154,7 @@ class LH5Store:
         """Returns an LH5 object appropriate for use as a pre-allocated buffer
         in a read loop. Sets size to `size` if object has a size.
         """
-        obj, n_rows = self.read(name, lh5_file, n_rows=0, field_mask=field_mask)
+        obj = self.read(name, lh5_file, n_rows=0, field_mask=field_mask)
         if hasattr(obj, "resize") and size is not None:
             obj.resize(new_size=size)
         return obj
@@ -182,72 +181,20 @@ class LH5Store:
         """
         # grab files from store
         if isinstance(lh5_file, (str, h5py.File)):
-            lh5_obj = self.gimme_file(lh5_file, "r", **file_kwargs)[name]
+            h5f = self.gimme_file(lh5_file, "r", **file_kwargs)
         else:
-            lh5_files = list(lh5_file)
-            n_rows_read = 0
-
-            for i, h5f in enumerate(lh5_files):
-                if (
-                    isinstance(idx, (list, tuple))
-                    and len(idx) > 0
-                    and not np.isscalar(idx[0])
-                ):
-                    # a list of lists: must be one per file
-                    idx_i = idx[i]
-                elif idx is not None:
-                    # make idx a proper tuple if it's not one already
-                    if not (isinstance(idx, tuple) and len(idx) == 1):
-                        idx = (idx,)
-                    # idx is a long continuous array
-                    n_rows_i = utils.read_n_rows(name, h5f)
-                    # find the length of the subset of idx that contains indices
-                    # that are less than n_rows_i
-                    n_rows_to_read_i = bisect.bisect_left(idx[0], n_rows_i)
-                    # now split idx into idx_i and the remainder
-                    idx_i = np.array(idx[0])[:n_rows_to_read_i]
-                    idx = np.array(idx[0])[n_rows_to_read_i:] - n_rows_i
-                else:
-                    idx_i = None
-                n_rows_i = n_rows - n_rows_read
-
-                obj_buf, n_rows_read_i = self.read(
-                    name,
-                    h5f,
-                    start_row,
-                    n_rows_i,
-                    idx_i,
-                    use_h5idx,
-                    field_mask,
-                    obj_buf,
-                    obj_buf_start,
-                    decompress,
-                )
-
-                n_rows_read += n_rows_read_i
-                if n_rows_read >= n_rows or obj_buf is None:
-                    return obj_buf, n_rows_read
-                start_row = 0
-                obj_buf_start += n_rows_read_i
-            return obj_buf, n_rows_read
-
-        if isinstance(idx, (list, tuple)) and len(idx) > 0 and not np.isscalar(idx[0]):
-            idx = idx[0]
-        if isinstance(idx, np.ndarray) and idx.dtype == np.dtype("?"):
-            idx = np.where(idx)[0]
-
-        return _serializers._h5_read_lgdo(
-            lh5_obj.id,
-            lh5_obj.file.filename,
-            lh5_obj.name,
-            start_row=start_row,
-            n_rows=n_rows,
-            idx=idx,
-            use_h5idx=use_h5idx,
-            field_mask=field_mask,
-            obj_buf=obj_buf,
-            obj_buf_start=obj_buf_start,
-            decompress=decompress,
+            h5f = [self.gimme_file(f, "r", **file_kwargs) for f in lh5_file]
+        return read(
+            name,
+            h5f,
+            start_row,
+            n_rows,
+            idx,
+            use_h5idx,
+            field_mask,
+            obj_buf,
+            obj_buf_start,
+            decompress,
         )
 
     def write(