legend-pydataobj 1.5.1__py3-none-any.whl → 1.6.1__py3-none-any.whl

lgdo/lh5/_serializers/read/composite.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+import bisect
+import logging
+import sys
+from collections import defaultdict
+
+import h5py
+import numpy as np
+
+from ....types import (
+    Array,
+    ArrayOfEncodedEqualSizedArrays,
+    ArrayOfEqualSizedArrays,
+    FixedSizeArray,
+    Scalar,
+    Struct,
+    Table,
+    VectorOfEncodedVectors,
+    VectorOfVectors,
+    WaveformTable,
+)
+from ... import datatype as dtypeutils
+from ...exceptions import LH5DecodeError
+from ...utils import read_n_rows
+from . import utils
+from .array import (
+    _h5_read_array,
+    _h5_read_array_of_equalsized_arrays,
+    _h5_read_fixedsize_array,
+)
+from .encoded import (
+    _h5_read_array_of_encoded_equalsized_arrays,
+    _h5_read_vector_of_encoded_vectors,
+)
+from .scalar import _h5_read_scalar
+from .vector_of_vectors import _h5_read_vector_of_vectors
+
+log = logging.getLogger(__name__)
+
+
+def _h5_read_lgdo(
+    name,
+    h5f,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    field_mask=None,
+    obj_buf=None,
+    obj_buf_start=0,
+    decompress=True,
+):
+    # Handle list-of-files recursively
+    if not isinstance(h5f, (str, h5py.File)):
+        lh5_file = list(h5f)
+        n_rows_read = 0
+
+        for i, h5f in enumerate(lh5_file):
+            if isinstance(idx, list) 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 = 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 = (idx[0][:n_rows_to_read_i],)
+                idx = (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 = _h5_read_lgdo(
+                name,
+                h5f,
+                start_row=start_row,
+                n_rows=n_rows_i,
+                idx=idx_i,
+                use_h5idx=use_h5idx,
+                field_mask=field_mask,
+                obj_buf=obj_buf,
+                obj_buf_start=obj_buf_start,
+                decompress=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 not isinstance(h5f, h5py.File):
+        h5f = h5py.File(h5f, mode="r")
+
+    log.debug(
+        f"reading {h5f.filename}:{name}[{start_row}:{n_rows}], decompress = {decompress}, "
+        + (f" with field mask {field_mask}" if field_mask else "")
+    )
+
+    # make idx a proper tuple if it's not one already
+    if not (isinstance(idx, tuple) and len(idx) == 1) and idx is not None:
+        idx = (idx,)
+
+    try:
+        lgdotype = dtypeutils.datatype(h5f[name].attrs["datatype"])
+    except KeyError as e:
+        msg = "dataset not in file or missing 'datatype' attribute"
+        raise LH5DecodeError(msg, h5f, name) from e
+
+    if lgdotype is Scalar:
+        return _h5_read_scalar(
+            name,
+            h5f,
+            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)):
+        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__)
+
+    if lgdotype is Struct:
+        return _h5_read_struct(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            field_mask=field_mask,
+            decompress=decompress,
+        )
+
+    # Below here is all array-like types. So trim idx if needed
+    if idx is not None:
+        # check if idx is just an ordered list of the integers if so can ignore
+        if (idx[0] == np.arange(0, len(idx[0]), 1)).all():
+            if n_rows > len(idx[0]):
+                n_rows = len(idx[0])
+            idx = None
+        else:
+            # chop off indices < start_row
+            i_first_valid = bisect.bisect_left(idx[0], start_row)
+            idxa = idx[0][i_first_valid:]
+            # don't readout more than n_rows indices
+            idx = (idxa[:n_rows],)  # works even if n_rows > len(idxa)
+
+    if lgdotype is Table:
+        return _h5_read_table(
+            name,
+            h5f,
+            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,
+        )
+
+    if lgdotype is ArrayOfEncodedEqualSizedArrays:
+        return _h5_read_array_of_encoded_equalsized_arrays(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+            decompress=decompress,
+        )
+
+    if lgdotype is VectorOfEncodedVectors:
+        return _h5_read_vector_of_encoded_vectors(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+            decompress=decompress,
+        )
+
+    if lgdotype is VectorOfVectors:
+        return _h5_read_vector_of_vectors(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+        )
+
+    if lgdotype is FixedSizeArray:
+        return _h5_read_fixedsize_array(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+        )
+
+    if lgdotype is ArrayOfEqualSizedArrays:
+        return _h5_read_array_of_equalsized_arrays(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+        )
+
+    if lgdotype is Array:
+        return _h5_read_array(
+            name,
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+        )
+
+    msg = f"no rule to decode {lgdotype.__name__} from LH5"
+    raise LH5DecodeError(msg, h5f, name)
+
+
+def _h5_read_struct(
+    name,
+    h5f,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    field_mask=None,
+    decompress=True,
+):
+    # TODO: it's strange to pass start_row, n_rows, idx to struct
+    # fields. If they all had shared indexing, they should be in a
+    # table... Maybe should emit a warning? Or allow them to be
+    # dicts keyed by field name?
+
+    attrs = dict(h5f[name].attrs)
+
+    # 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
+    )
+
+    # modify datatype in attrs if a field_mask was used
+    attrs["datatype"] = "struct{" + ",".join(selected_fields) + "}"
+
+    # loop over fields and read
+    obj_dict = {}
+    for field in selected_fields:
+        # support for integer keys
+        field_key = int(field) if attrs.get("int_keys") else str(field)
+        obj_dict[field_key], _ = _h5_read_lgdo(
+            f"{name}/{field}",
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            decompress=decompress,
+        )
+
+    return Struct(obj_dict=obj_dict, attrs=attrs), 1
+
+
+def _h5_read_table(
+    name,
+    h5f,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    field_mask=None,
+    obj_buf=None,
+    obj_buf_start=0,
+    decompress=True,
+):
+    if obj_buf is not None and not isinstance(obj_buf, Table):
+        msg = "provided object buffer is not a Table"
+        raise LH5DecodeError(msg, h5f, name)
+
+    attrs = dict(h5f[name].attrs)
+
+    # 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
+    )
+
+    # modify datatype in attrs if a field_mask was used
+    attrs["datatype"] = "table{" + ",".join(selected_fields) + "}"
+
+    # read out each of the fields
+    col_dict = {}
+    rows_read = []
+    for field in selected_fields:
+        fld_buf = None
+        if obj_buf is not None:
+            if not isinstance(obj_buf, Table) or field not in obj_buf:
+                msg = "provided object buffer is not a Table or columns are missing"
+                raise LH5DecodeError(msg, h5f, name)
+
+            fld_buf = obj_buf[field]
+
+        col_dict[field], n_rows_read = _h5_read_lgdo(
+            f"{name}/{field}",
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=fld_buf,
+            obj_buf_start=obj_buf_start,
+            decompress=decompress,
+        )
+
+        if obj_buf is not None and obj_buf_start + n_rows_read > len(obj_buf):
+            obj_buf.resize(obj_buf_start + n_rows_read)
+
+        rows_read.append(n_rows_read)
+
+    # warn if all columns don't read in the same number of rows
+    if len(rows_read) > 0:
+        n_rows_read = rows_read[0]
+    else:
+        n_rows_read = 0
+        log.warning(f"Table '{name}' has no fields specified by {field_mask=}")
+
+    for n in rows_read[1:]:
+        if n != n_rows_read:
+            log.warning(
+                f"Table '{name}' got strange n_rows_read = {n}, "
+                "{n_rows_read} was expected ({rows_read})"
+            )
+
+    # fields have been read out, now return a table
+    if obj_buf is None:
+        # if col_dict contains just 3 objects called t0, dt, and values,
+        # return a WaveformTable
+        if (
+            len(col_dict) == 3
+            and "t0" in col_dict
+            and "dt" in col_dict
+            and "values" in col_dict
+        ):
+            table = WaveformTable(
+                t0=col_dict["t0"], dt=col_dict["dt"], values=col_dict["values"]
+            )
+        else:
+            table = Table(col_dict=col_dict, attrs=attrs)
+
+        # set (write) loc to end of tree
+        table.loc = n_rows_read
+        return table, n_rows_read
+
+    # We have read all fields into the object buffer. Run
+    # checks: All columns should be the same size. So update
+    # table's size as necessary, warn if any mismatches are found
+    obj_buf.resize(do_warn=True)
+    # set (write) loc to end of tree
+    obj_buf.loc = obj_buf_start + n_rows_read
+
+    # check attributes
+    utils.check_obj_buf_attrs(obj_buf.attrs, attrs, h5f, name)
+
+    return obj_buf, n_rows_read

lgdo/lh5/_serializers/read/encoded.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import logging
+import sys
+
+from .... import compression as compress
+from ....types import (
+    ArrayOfEncodedEqualSizedArrays,
+    VectorOfEncodedVectors,
+)
+from ...exceptions import LH5DecodeError
+from .array import (
+    _h5_read_array,
+)
+from .scalar import _h5_read_scalar
+from .vector_of_vectors import _h5_read_vector_of_vectors
+
+log = logging.getLogger(__name__)
+
+
+def _h5_read_array_of_encoded_equalsized_arrays(
+    name,
+    h5f,
+    **kwargs,
+):
+    return _h5_read_encoded_array(ArrayOfEncodedEqualSizedArrays, name, h5f, **kwargs)
+
+
+def _h5_read_vector_of_encoded_vectors(
+    name,
+    h5f,
+    **kwargs,
+):
+    return _h5_read_encoded_array(VectorOfEncodedVectors, name, h5f, **kwargs)
+
+
+def _h5_read_encoded_array(
+    lgdotype,
+    name,
+    h5f,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    obj_buf=None,
+    obj_buf_start=0,
+    decompress=True,
+):
+    if lgdotype not in (ArrayOfEncodedEqualSizedArrays, VectorOfEncodedVectors):
+        msg = f"unsupported read of encoded type {lgdotype.__name__}"
+        raise LH5DecodeError(msg, h5f, name)
+
+    if not decompress and obj_buf is not None and not isinstance(obj_buf, lgdotype):
+        msg = f"object buffer is not a {lgdotype.__name__}"
+        raise LH5DecodeError(msg, h5f, name)
+
+    # read out decoded_size, either a Scalar or an Array
+    decoded_size_buf = encoded_data_buf = None
+    if obj_buf is not None and not decompress:
+        decoded_size_buf = obj_buf.decoded_size
+        encoded_data_buf = obj_buf.encoded_data
+
+    if lgdotype is VectorOfEncodedVectors:
+        decoded_size, _ = _h5_read_array(
+            f"{name}/decoded_size",
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            obj_buf=None if decompress else decoded_size_buf,
+            obj_buf_start=0 if decompress else obj_buf_start,
+        )
+
+    else:
+        decoded_size, _ = _h5_read_scalar(
+            f"{name}/decoded_size",
+            h5f,
+            obj_buf=None if decompress else decoded_size_buf,
+        )
+
+    # read out encoded_data, a VectorOfVectors
+    encoded_data, n_rows_read = _h5_read_vector_of_vectors(
+        f"{name}/encoded_data",
+        h5f,
+        start_row=start_row,
+        n_rows=n_rows,
+        idx=idx,
+        use_h5idx=use_h5idx,
+        obj_buf=None if decompress else encoded_data_buf,
+        obj_buf_start=0 if decompress else obj_buf_start,
+    )
+
+    # return the still encoded data in the buffer object, if there
+    if obj_buf is not None and not decompress:
+        return obj_buf, n_rows_read
+
+    # otherwise re-create the encoded LGDO
+    rawdata = lgdotype(
+        encoded_data=encoded_data,
+        decoded_size=decoded_size,
+        attrs=h5f[name].attrs,
+    )
+
+    # already return if no decompression is requested
+    if not decompress:
+        return rawdata, n_rows_read
+
+    # if no buffer, decode and return
+    if obj_buf is None and decompress:
+        return compress.decode(rawdata), n_rows_read
+
+    # eventually expand provided obj_buf, if too short
+    buf_size = obj_buf_start + n_rows_read
+    if len(obj_buf) < buf_size:
+        obj_buf.resize(buf_size)
+
+    # use the (decoded object type) buffer otherwise
+    if lgdotype is ArrayOfEncodedEqualSizedArrays:
+        compress.decode(rawdata, obj_buf[obj_buf_start:buf_size])
+
+    elif lgdotype is VectorOfEncodedVectors:
+        # FIXME: not a good idea. an in place decoding version
+        # of decode would be needed to avoid extra memory
+        # allocations
+        for i, wf in enumerate(compress.decode(rawdata)):
+            obj_buf[obj_buf_start + i] = wf
+
+    return obj_buf, n_rows_read

lgdo/lh5/_serializers/read/ndarray.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import logging
+import sys
+from bisect import bisect_left
+
+import numpy as np
+
+from ....types import Array
+from ... import datatype
+from ...exceptions import LH5DecodeError
+
+log = logging.getLogger(__name__)
+
+
+def _h5_read_ndarray(
+    name,
+    h5f,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    obj_buf=None,
+    obj_buf_start=0,
+):
+    if obj_buf is not None and not isinstance(obj_buf, Array):
+        msg = "object buffer is not an Array"
+        raise LH5DecodeError(msg, h5f, name)
+
+    # compute the number of rows to read
+    # we culled idx above for start_row and n_rows, now we have to apply
+    # the constraint of the length of the dataset
+    try:
+        ds_n_rows = h5f[name].shape[0]
+    except AttributeError as e:
+        msg = "does not seem to be an HDF5 dataset"
+        raise LH5DecodeError(msg, h5f, name) from e
+
+    if idx is not None:
+        if len(idx[0]) > 0 and idx[0][-1] >= ds_n_rows:
+            log.warning("idx indexed past the end of the array in the file. Culling...")
+            n_rows_to_read = bisect_left(idx[0], ds_n_rows)
+            idx = (idx[0][:n_rows_to_read],)
+            if len(idx[0]) == 0:
+                log.warning("idx empty after culling.")
+        n_rows_to_read = len(idx[0])
+    else:
+        n_rows_to_read = ds_n_rows - start_row
+    if n_rows_to_read > n_rows:
+        n_rows_to_read = n_rows
+
+    # if idx is passed, check if we can make it a slice instead (faster)
+    change_idx_to_slice = False
+
+    # prepare the selection for the read. Use idx if available
+    if idx is not None:
+        # check if idx is empty and convert to slice instead
+        if len(idx[0]) == 0:
+            source_sel = np.s_[0:0]
+            change_idx_to_slice = True
+        # check if idx is contiguous and increasing
+        # if so, convert it to a slice instead (faster)
+        elif np.all(np.diff(idx[0]) == 1):
+            source_sel = np.s_[idx[0][0] : idx[0][-1] + 1]
+            change_idx_to_slice = True
+        else:
+            source_sel = idx
+    else:
+        source_sel = np.s_[start_row : start_row + n_rows_to_read]
+
+    # Now read the array
+    if obj_buf is not None and n_rows_to_read > 0:
+        buf_size = obj_buf_start + n_rows_to_read
+        if len(obj_buf) < buf_size:
+            obj_buf.resize(buf_size)
+        dest_sel = np.s_[obj_buf_start:buf_size]
+
+        # this is required to make the read of multiple files faster
+        # until a better solution found.
+        if change_idx_to_slice or idx is None or use_h5idx:
+            h5f[name].read_direct(obj_buf.nda, source_sel, dest_sel)
+        else:
+            # it is faster to read the whole object and then do fancy indexing
+            obj_buf.nda[dest_sel] = h5f[name][...][source_sel]
+
+        nda = obj_buf.nda
+    elif n_rows == 0:
+        tmp_shape = (0,) + h5f[name].shape[1:]
+        nda = np.empty(tmp_shape, h5f[name].dtype)
+    elif change_idx_to_slice or idx is None or use_h5idx:
+        nda = h5f[name][source_sel]
+    else:
+        # it is faster to read the whole object and then do fancy indexing
+        nda = h5f[name][...][source_sel]
+
+    # Finally, set attributes and return objects
+    attrs = h5f[name].attrs
+
+    # special handling for bools
+    # (c and Julia store as uint8 so cast to bool)
+    if datatype.get_nested_datatype_string(attrs["datatype"]) == "bool":
+        nda = nda.astype(np.bool_)
+
+    return (nda, attrs, n_rows_to_read)

lgdo/lh5/_serializers/read/scalar.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+from ....types import Scalar
+from ...exceptions import LH5DecodeError
+
+log = logging.getLogger(__name__)
+
+
+def _h5_read_scalar(
+    name,
+    h5f,
+    obj_buf=None,
+):
+    value = h5f[name][()]
+
+    # special handling for bools
+    # (c and Julia store as uint8 so cast to bool)
+    if h5f[name].attrs["datatype"] == "bool":
+        value = np.bool_(value)
+
+    if obj_buf is not None:
+        if not isinstance(obj_buf, Scalar):
+            msg = "object buffer a Scalar"
+            raise LH5DecodeError(msg, h5f, name)
+
+        obj_buf.value = value
+        obj_buf.attrs.update(h5f[name].attrs)
+        return obj_buf, 1
+
+    return Scalar(value=value, attrs=h5f[name].attrs), 1

lgdo/lh5/_serializers/read/utils.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from ...exceptions import LH5DecodeError
+
+
+def check_obj_buf_attrs(attrs, new_attrs, file, name):
+    if set(attrs.keys()) != set(new_attrs.keys()):
+        msg = (
+            f"existing buffer and new data chunk have different attributes: "
+            f"obj_buf.attrs={attrs} != {file.filename}[{name}].attrs={new_attrs}"
+        )
+        raise LH5DecodeError(msg, file, name)