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

lgdo/lh5/_serializers/read/vector_of_vectors.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import logging
+import sys
+
+import numba
+import numpy as np
+
+from ....types import (
+    Array,
+    VectorOfVectors,
+)
+from ... import datatype as dtypeutils
+from ...exceptions import LH5DecodeError
+from .array import (
+    _h5_read_array,
+)
+
+log = logging.getLogger(__name__)
+
+
+def _h5_read_vector_of_vectors(
+    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, VectorOfVectors):
+        msg = "object buffer is not a VectorOfVectors"
+        raise LH5DecodeError(msg, h5f, name)
+
+    # read out cumulative_length
+    cumulen_buf = None if obj_buf is None else obj_buf.cumulative_length
+    cumulative_length, n_rows_read = _h5_read_array(
+        f"{name}/cumulative_length",
+        h5f,
+        start_row=start_row,
+        n_rows=n_rows,
+        idx=idx,
+        use_h5idx=use_h5idx,
+        obj_buf=cumulen_buf,
+        obj_buf_start=obj_buf_start,
+    )
+    # get a view of just what was read out for cleaner code below
+    this_cumulen_nda = cumulative_length.nda[
+        obj_buf_start : obj_buf_start + n_rows_read
+    ]
+
+    if idx is not None and n_rows_read > 0:
+        # get the starting indices for each array in flattened data:
+        # the starting index for array[i] is cumulative_length[i-1]
+        idx2 = (np.asarray(idx[0]).copy() - 1,)
+
+        # re-read cumulative_length with these indices
+        # note this will allocate memory for fd_starts!
+        fd_start = None
+        if idx2[0][0] == -1:
+            idx2 = (idx2[0][1:],)
+            fd_start = 0  # this variable avoids an ndarray append
+
+        fd_starts, fds_n_rows_read = _h5_read_array(
+            f"{name}/cumulative_length",
+            h5f,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx2,
+            use_h5idx=use_h5idx,
+            obj_buf=None,
+        )
+        fd_starts = fd_starts.nda  # we just need the nda
+        if fd_start is None:
+            fd_start = fd_starts[0]
+
+        # compute the length that flattened_data will have after the
+        # fancy-indexed read
+        fd_n_rows = np.sum(this_cumulen_nda[-len(fd_starts) :] - fd_starts)
+        if fd_start == 0:
+            fd_n_rows += this_cumulen_nda[0]
+
+        # now make fd_idx
+        fd_idx = np.empty(fd_n_rows, dtype="int32")
+        fd_idx = _make_fd_idx(fd_starts, this_cumulen_nda, fd_idx)
+
+        # Now clean up this_cumulen_nda, to be ready
+        # to match the in-memory version of flattened_data. Note: these
+        # operations on the view change the original array because they are
+        # numpy arrays, not lists.
+        this_cumulen_nda[-len(fd_starts) :] -= fd_starts
+        np.cumsum(this_cumulen_nda, out=this_cumulen_nda)
+
+    else:
+        fd_idx = None
+
+        # determine the start_row and n_rows for the flattened_data readout
+        fd_start = 0
+        if start_row > 0 and n_rows_read > 0:
+            # need to read out the cumulen sample -before- the first sample
+            # read above in order to get the starting row of the first
+            # vector to read out in flattened_data
+            fd_start = h5f[f"{name}/cumulative_length"][start_row - 1]
+
+            # check limits for values that will be used subsequently
+            if this_cumulen_nda[-1] < fd_start:
+                log.debug(
+                    f"this_cumulen_nda[-1] = {this_cumulen_nda[-1]}, "
+                    f"fd_start = {fd_start}, "
+                    f"start_row = {start_row}, "
+                    f"n_rows_read = {n_rows_read}"
+                )
+                msg = (
+                    f"cumulative_length non-increasing between entries "
+                    f"{start_row} and {start_row+n_rows_read}"
+                )
+                raise LH5DecodeError(msg, h5f, name)
+
+        # determine the number of rows for the flattened_data readout
+        fd_n_rows = this_cumulen_nda[-1] if n_rows_read > 0 else 0
+
+        # Now done with this_cumulen_nda, so we can clean it up to be ready
+        # to match the in-memory version of flattened_data. Note: these
+        # operations on the view change the original array because they are
+        # numpy arrays, not lists.
+        #
+        # First we need to subtract off the in-file offset for the start of
+        # read for flattened_data
+        this_cumulen_nda -= fd_start
+
+    # If we started with a partially-filled buffer, add the
+    # appropriate offset for the start of the in-memory flattened
+    # data for this read.
+    fd_buf_start = np.uint32(0)
+    if obj_buf_start > 0:
+        fd_buf_start = cumulative_length.nda[obj_buf_start - 1]
+        this_cumulen_nda += fd_buf_start
+
+    # Now prepare the object buffer if necessary
+    fd_buf = None
+    if obj_buf is not None:
+        fd_buf = obj_buf.flattened_data
+        # 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.resize(fdb_size)
+
+    # now read
+    lgdotype = dtypeutils.datatype(h5f[f"{name}/flattened_data"].attrs["datatype"])
+    if lgdotype is Array:
+        _func = _h5_read_array
+    elif lgdotype is VectorOfVectors:
+        _func = _h5_read_vector_of_vectors
+    else:
+        msg = "type {lgdotype.__name__} is not supported"
+        raise LH5DecodeError(msg, h5f, f"{name}/flattened_data")
+
+    flattened_data, _ = _func(
+        f"{name}/flattened_data",
+        h5f,
+        start_row=fd_start,
+        n_rows=fd_n_rows,
+        idx=fd_idx,
+        use_h5idx=use_h5idx,
+        obj_buf=fd_buf,
+        obj_buf_start=fd_buf_start,
+    )
+
+    if obj_buf is not None:
+        # if the buffer is partially filled, cumulative_length will be invalid
+        # (i.e. non monotonically increasing). Let's fix that but filling the
+        # rest of the array with the length of flattened_data
+        end = obj_buf_start + n_rows_read
+        obj_buf.cumulative_length.nda[end:] = obj_buf.cumulative_length.nda[end - 1]
+
+        return obj_buf, n_rows_read
+
+    return (
+        VectorOfVectors(
+            flattened_data=flattened_data,
+            cumulative_length=cumulative_length,
+            attrs=h5f[name].attrs,
+        ),
+        n_rows_read,
+    )
+
+
+@numba.njit(parallel=False, fastmath=True)
+def _make_fd_idx(starts, stops, idx):
+    k = 0
+    if len(starts) < len(stops):
+        for i in range(stops[0]):
+            idx[k] = i
+            k += 1
+        stops = stops[1:]
+    for j in range(len(starts)):
+        for i in range(starts[j], stops[j]):
+            idx[k] = i
+            k += 1
+    return (idx,)

lgdo/lh5/_serializers/write/array.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import logging
+
+import h5py
+import numpy as np
+
+from .... import types
+from ...exceptions import LH5EncodeError
+
+log = logging.getLogger(__name__)
+
+DEFAULT_HDF5_SETTINGS: dict[str, ...] = {"shuffle": True, "compression": "gzip"}
+
+
+def _h5_write_array(
+    obj,
+    name,
+    lh5_file,
+    group="/",
+    start_row=0,
+    n_rows=None,
+    wo_mode="append",
+    write_start=0,
+    **h5py_kwargs,
+):
+    assert isinstance(obj, types.Array)
+
+    if n_rows is None or n_rows > obj.nda.shape[0] - start_row:
+        n_rows = obj.nda.shape[0] - start_row
+
+    nda = obj.nda[start_row : start_row + n_rows]
+
+    # hack to store bools as uint8 for c / Julia compliance
+    if nda.dtype.name == "bool":
+        nda = nda.astype(np.uint8)
+
+    # need to create dataset from ndarray the first time for speed
+    # creating an empty dataset and appending to that is super slow!
+    if (wo_mode != "a" and write_start == 0) or name not in group:
+        # this is needed in order to have a resizable (in the first
+        # axis) data set, i.e. rows can be appended later
+        # NOTE: this automatically turns chunking on!
+        maxshape = (None,) + nda.shape[1:]
+        h5py_kwargs.setdefault("maxshape", maxshape)
+
+        if wo_mode == "o" and name in group:
+            log.debug(f"overwriting {name} in {group}")
+            del group[name]
+
+        # set default compression options
+        for k, v in DEFAULT_HDF5_SETTINGS.items():
+            h5py_kwargs.setdefault(k, v)
+
+        # compress using the 'compression' LGDO attribute, if available
+        if "compression" in obj.attrs:
+            comp_algo = obj.attrs["compression"]
+            if isinstance(comp_algo, dict):
+                h5py_kwargs |= obj.attrs["compression"]
+            else:
+                h5py_kwargs["compression"] = obj.attrs["compression"]
+
+        # and even the 'hdf5_settings' one, preferred
+        if "hdf5_settings" in obj.attrs:
+            h5py_kwargs |= obj.attrs["hdf5_settings"]
+
+        # create HDF5 dataset
+        ds = group.create_dataset(name, data=nda, **h5py_kwargs)
+
+        # attach HDF5 dataset attributes, but not "compression"!
+        _attrs = obj.getattrs(datatype=True)
+        _attrs.pop("compression", None)
+        _attrs.pop("hdf5_settings", None)
+        ds.attrs.update(_attrs)
+
+        return
+
+    # Now append or overwrite
+    ds = group[name]
+    if not isinstance(ds, h5py.Dataset):
+        msg = (
+            f"existing HDF5 object '{name}' in group '{group}'"
+            " is not a dataset! Cannot overwrite or append"
+        )
+        raise LH5EncodeError(msg, lh5_file, group, name)
+
+    old_len = ds.shape[0]
+    if wo_mode == "a":
+        write_start = old_len
+
+    ds.resize(write_start + nda.shape[0], axis=0)
+    ds[write_start:] = nda

lgdo/lh5/_serializers/write/composite.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import logging
+
+import h5py
+
+from .... import compression, types
+from ... import datatype, utils
+from ...exceptions import LH5EncodeError
+from .array import _h5_write_array
+from .scalar import _h5_write_scalar
+from .vector_of_vectors import _h5_write_vector_of_vectors
+
+log = logging.getLogger(__name__)
+
+
+def _h5_write_lgdo(
+    obj,
+    name,
+    lh5_file,
+    group="/",
+    start_row=0,
+    n_rows=None,
+    wo_mode="append",
+    write_start=0,
+    **h5py_kwargs,
+):
+    assert isinstance(obj, types.LGDO)
+
+    if wo_mode == "write_safe":
+        wo_mode = "w"
+    if wo_mode == "append":
+        wo_mode = "a"
+    if wo_mode == "overwrite":
+        wo_mode = "o"
+    if wo_mode == "overwrite_file":
+        wo_mode = "of"
+        write_start = 0
+    if wo_mode == "append_column":
+        wo_mode = "ac"
+    if wo_mode not in ["w", "a", "o", "of", "ac"]:
+        msg = f"unknown wo_mode '{wo_mode}'"
+        raise LH5EncodeError(msg, lh5_file, group, name)
+
+    # "mode" is for the h5df.File and wo_mode is for this function
+    # In hdf5, 'a' is really "modify" -- in addition to appending, you can
+    # change any object in the file. So we use file:append for
+    # write_object:overwrite.
+    mode = "w" if wo_mode == "of" else "a"
+
+    if not isinstance(lh5_file, h5py.File):
+        lh5_file = h5py.File(lh5_file, mode=mode)
+
+    log.debug(
+        f"writing {obj!r}[{start_row}:{n_rows}] as "
+        f"{lh5_file.filename}:{group}/{name}[{write_start}:], "
+        f"mode = {wo_mode}, h5py_kwargs = {h5py_kwargs}"
+    )
+
+    group = utils.get_h5_group(group, lh5_file)
+
+    if wo_mode == "w" and name in group:
+        msg = f"can't overwrite '{name}' in wo_mode 'write_safe'"
+        raise LH5EncodeError(msg, lh5_file, group, name)
+
+    # struct or table or waveform table
+    if isinstance(obj, types.Struct):
+        return _h5_write_struct(
+            obj,
+            name,
+            lh5_file,
+            group=group,
+            start_row=start_row,
+            n_rows=n_rows,
+            wo_mode=wo_mode,
+            write_start=write_start,
+            **h5py_kwargs,
+        )
+
+    # scalars
+    if isinstance(obj, types.Scalar):
+        return _h5_write_scalar(obj, name, lh5_file, group, wo_mode)
+
+    # vector of encoded vectors
+    if isinstance(
+        obj, (types.VectorOfEncodedVectors, types.ArrayOfEncodedEqualSizedArrays)
+    ):
+        group = utils.get_h5_group(
+            name, group, grp_attrs=obj.attrs, overwrite=(wo_mode == "o")
+        )
+
+        # ask not to further compress flattened_data, it is already compressed!
+        obj.encoded_data.flattened_data.attrs["compression"] = None
+
+        _h5_write_vector_of_vectors(
+            obj.encoded_data,
+            "encoded_data",
+            lh5_file,
+            group=group,
+            start_row=start_row,
+            n_rows=n_rows,
+            wo_mode=wo_mode,
+            write_start=write_start,
+            **h5py_kwargs,
+        )
+
+        if isinstance(obj.decoded_size, types.Scalar):
+            _h5_write_scalar(
+                obj.decoded_size,
+                "decoded_size",
+                lh5_file,
+                group=group,
+                wo_mode=wo_mode,
+            )
+        else:
+            _h5_write_array(
+                obj.decoded_size,
+                "decoded_size",
+                lh5_file,
+                group=group,
+                start_row=start_row,
+                n_rows=n_rows,
+                wo_mode=wo_mode,
+                write_start=write_start,
+                **h5py_kwargs,
+            )
+
+        return None
+
+    # vector of vectors
+    if isinstance(obj, types.VectorOfVectors):
+        return _h5_write_vector_of_vectors(
+            obj,
+            name,
+            lh5_file,
+            group=group,
+            start_row=start_row,
+            n_rows=n_rows,
+            wo_mode=wo_mode,
+            write_start=write_start,
+            **h5py_kwargs,
+        )
+
+    # if we get this far, must be one of the Array types
+    if isinstance(obj, types.Array):
+        return _h5_write_array(
+            obj,
+            name,
+            lh5_file,
+            group=group,
+            start_row=start_row,
+            n_rows=n_rows,
+            wo_mode=wo_mode,
+            write_start=write_start,
+            **h5py_kwargs,
+        )
+
+    msg = f"do not know how to write '{name}' of type '{type(obj).__name__}'"
+    raise LH5EncodeError(msg, lh5_file, group, name)
+
+
+def _h5_write_struct(
+    obj,
+    name,
+    lh5_file,
+    group="/",
+    start_row=0,
+    n_rows=None,
+    wo_mode="append",
+    write_start=0,
+    **h5py_kwargs,
+):
+    assert isinstance(obj, types.Struct)
+
+    # In order to append a column, we need to update the
+    # `table{old_fields}` value in `group.attrs['datatype"]` to include
+    # the new fields.  One way to do this is to override
+    # `obj.attrs["datatype"]` to include old and new fields. Then we
+    # can write the fields to the table as normal.
+    if wo_mode == "ac":
+        old_group = utils.get_h5_group(name, group)
+        lgdotype = datatype.datatype(old_group.attrs["datatype"])
+        fields = datatype.get_struct_fields(old_group.attrs["datatype"])
+        if not issubclass(lgdotype, types.Struct):
+            msg = f"Trying to append columns to an object of type {lgdotype.__name__}"
+            raise LH5EncodeError(msg, lh5_file, group, name)
+
+        # If the mode is `append_column`, make sure we aren't appending
+        # a table that has a column of the same name as in the existing
+        # table. Also make sure that the field we are adding has the
+        # same size
+        if len(list(set(fields).intersection(set(obj.keys())))) != 0:
+            msg = (
+                f"Can't append {list(set(fields).intersection(set(obj.keys())))} "
+                "column(s) to a table with the same field(s)"
+            )
+            raise LH5EncodeError(msg, lh5_file, group, name)
+        # It doesn't matter what key we access, as all fields in the old table have the same size
+        if old_group[next(iter(old_group.keys()))].size != obj.size:
+            msg = (
+                f"Table sizes don't match. Trying to append column of size {obj.size} "
+                f"to a table of size {old_group[next(iter(old_group.keys()))].size}."
+            )
+            raise LH5EncodeError(msg, lh5_file, group, name)
+
+        # Now we can append the obj.keys() to the old fields, and then update obj.attrs.
+        fields.extend(list(obj.keys()))
+        obj.attrs.pop("datatype")
+        obj.attrs["datatype"] = "table" + "{" + ",".join(fields) + "}"
+
+    group = utils.get_h5_group(
+        name,
+        group,
+        grp_attrs=obj.attrs,
+        overwrite=(wo_mode in ["o", "ac"]),
+    )
+    # If the mode is overwrite, then we need to peek into the file's
+    # table's existing fields.  If we are writing a new table to the
+    # group that does not contain an old field, we should delete that
+    # old field from the file
+    if wo_mode == "o":
+        # Find the old keys in the group that are not present in the
+        # new table's keys, then delete them
+        for key in list(set(group.keys()) - set(obj.keys())):
+            log.debug(f"{key} is not present in new table, deleting field")
+            del group[key]
+
+    for field in obj:
+        # eventually compress waveform table values with LGDO's
+        # custom codecs before writing
+        # if waveformtable.values.attrs["compression"] is NOT a
+        # WaveformCodec, just leave it there
+        obj_fld = None
+        if (
+            isinstance(obj, types.WaveformTable)
+            and field == "values"
+            and not isinstance(obj.values, types.VectorOfEncodedVectors)
+            and not isinstance(obj.values, types.ArrayOfEncodedEqualSizedArrays)
+            and "compression" in obj.values.attrs
+            and isinstance(obj.values.attrs["compression"], compression.WaveformCodec)
+        ):
+            codec = obj.values.attrs["compression"]
+            obj_fld = compression.encode(obj.values, codec=codec)
+        else:
+            obj_fld = obj[field]
+
+        # Convert keys to string for dataset names
+        f = str(field)
+        _h5_write_lgdo(
+            obj_fld,
+            f,
+            lh5_file,
+            group=group,
+            start_row=start_row,
+            n_rows=n_rows,
+            wo_mode=wo_mode,
+            write_start=write_start,
+            **h5py_kwargs,
+        )

lgdo/lh5/_serializers/write/scalar.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import logging
+
+from .... import types
+from ...exceptions import LH5EncodeError
+
+log = logging.getLogger(__name__)
+
+
+def _h5_write_scalar(obj, name, lh5_file, group="/", wo_mode="append"):
+    assert isinstance(obj, types.Scalar)
+
+    if name in group:
+        if wo_mode in ["o", "a"]:
+            log.debug(f"overwriting {name} in {group}")
+            del group[name]
+        else:
+            msg = f"tried to overwrite but wo_mode is {wo_mode!r}"
+            raise LH5EncodeError(msg, lh5_file, group, name)
+
+    ds = group.create_dataset(name, shape=(), data=obj.value)
+    ds.attrs.update(obj.attrs)

lgdo/lh5/_serializers/write/vector_of_vectors.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import logging
+
+from .... import types
+from ... import utils
+from ...exceptions import LH5EncodeError
+from .array import _h5_write_array
+
+log = logging.getLogger(__name__)
+
+
+def _h5_write_vector_of_vectors(
+    obj,
+    name,
+    lh5_file,
+    group="/",
+    start_row=0,
+    n_rows=None,
+    wo_mode="append",
+    write_start=0,
+    **h5py_kwargs,
+):
+    assert isinstance(obj, types.VectorOfVectors)
+
+    group = utils.get_h5_group(
+        name, group, grp_attrs=obj.attrs, overwrite=(wo_mode == "o")
+    )
+    if n_rows is None or n_rows > obj.cumulative_length.nda.shape[0] - start_row:
+        n_rows = obj.cumulative_length.nda.shape[0] - start_row
+
+    # if appending we need to add an appropriate offset to the
+    # cumulative lengths as appropriate for the in-file object
+    offset = 0  # declare here because we have to subtract it off at the end
+    if (wo_mode in ("a", "o")) and "cumulative_length" in group:
+        len_cl = len(group["cumulative_length"])
+        if wo_mode == "a":
+            write_start = len_cl
+        if len_cl > 0:
+            offset = group["cumulative_length"][write_start - 1]
+
+    # First write flattened_data array. Only write rows with data.
+    fd_start = 0 if start_row == 0 else obj.cumulative_length.nda[start_row - 1]
+    fd_n_rows = (
+        obj.cumulative_length.nda[start_row + n_rows - 1] - fd_start
+        if len(obj.cumulative_length) > 0
+        else 0
+    )
+
+    if isinstance(obj.flattened_data, types.Array):
+        _func = _h5_write_array
+    elif isinstance(obj.flattened_data, types.VectorOfVectors):
+        _func = _h5_write_vector_of_vectors
+    else:
+        msg = (
+            "don't know how to serialize to disk flattened_data "
+            "of {type(obj.flattened_data).__name__} type"
+        )
+        raise LH5EncodeError(msg, lh5_file, group, f"{name}.flattened_data")
+
+    _func(
+        obj.flattened_data,
+        "flattened_data",
+        lh5_file,
+        group=group,
+        start_row=fd_start,
+        n_rows=fd_n_rows,
+        wo_mode=wo_mode,
+        write_start=offset,
+        **h5py_kwargs,
+    )
+
+    # now offset is used to give appropriate in-file values for
+    # cumulative_length. Need to adjust it for start_row
+    if start_row > 0:
+        offset -= obj.cumulative_length.nda[start_row - 1]
+
+    # Add offset to obj.cumulative_length itself to avoid memory allocation.
+    # Then subtract it off after writing! (otherwise it will be changed
+    # upon return)
+    cl_dtype = obj.cumulative_length.nda.dtype.type
+    obj.cumulative_length.nda += cl_dtype(offset)
+
+    _h5_write_array(
+        obj.cumulative_length,
+        "cumulative_length",
+        lh5_file,
+        group=group,
+        start_row=start_row,
+        n_rows=n_rows,
+        wo_mode=wo_mode,
+        write_start=write_start,
+        **h5py_kwargs,
+    )
+    obj.cumulative_length.nda -= cl_dtype(offset)