Implement in-place aoesa decompression

legend-exp · Oct 28, 2023 · f47a13b · f47a13b
1 parent c1d2b7b
commit f47a13b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 25 deletions.
diff --git a/src/lgdo/compression/radware.py b/src/lgdo/compression/radware.py
@@ -251,12 +251,6 @@ def decode(
                 dtype=int32,
                 attrs=sig_in.getattrs(),
             )
-        else:
-            if not (
-                isinstance(sig_out, lgdo.ArrayOfEqualSizedArrays)
-                and sig_out.nda.shape == (len(sig_in), sig_in.decoded_size.value)
-            ):
-                raise ValueError("sig_out is of the wrong format")
 
         siglen = np.empty(len(sig_in), dtype=uint32)
         # save original encoded vector lengths
@@ -269,7 +263,7 @@ def decode(
         # can now decode on the 2D matrix together with number of bytes to read per row
         _, siglen = decode(
             (sig_in.encoded_data.to_aoesa(preserve_dtype=True).nda, nbytes),
-            sig_out.nda,
+            sig_out if isinstance(sig_out, np.ndarray) else sig_out.nda,
             shift=shift,
         )
 

diff --git a/src/lgdo/lh5_store.py b/src/lgdo/lh5_store.py
@@ -531,26 +531,31 @@ def read_object(
                 elif 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 enc_lgdo == VectorOfEncodedVectors and not isinstance(
-                    obj_buf, VectorOfVectors
-                ):
-                    raise ValueError(
-                        f"obj_buf for decoded '{name}' not a VectorOfVectors"
-                    )
-                elif enc_lgdo == ArrayOfEncodedEqualSizedArrays and not isinstance(
-                    obj_buf, ArrayOfEqualSizedArrays
-                ):
-                    raise ValueError(
-                        f"obj_buf for decoded '{name}' not an ArrayOfEqualSizedArrays"
-                    )
+                if enc_lgdo == ArrayOfEncodedEqualSizedArrays:
+                    if not isinstance(obj_buf, ArrayOfEqualSizedArrays):
+                        raise ValueError(
+                            f"obj_buf for decoded '{name}' not an ArrayOfEqualSizedArrays"
+                        )
+
+                    compress.decode(rawdata, obj_buf[obj_buf_start:buf_size])
+
+                elif enc_lgdo == VectorOfEncodedVectors:
+                    if not isinstance(obj_buf, VectorOfVectors):
+                        raise ValueError(
+                            f"obj_buf for decoded '{name}' not a VectorOfVectors"
+                        )
 
-                # FIXME: not a good idea. an in place decoding version
-                # of decode would be needed to avoid extra memory
-                # allocations
-                # FIXME: obj_buf_start??? Write a unit test
-                for i, wf in enumerate(compress.decode(rawdata)):
-                    obj_buf[i] = wf
+                    # 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