Repository: legend-exp/pygama (language: python)

Issue:
Multiple store.write_object calls
When the `write_object` function of the `Store` class is run multiple times, I find that the "datatype" attribute of the HDF5 group is not updated after the first function call.  This causes a subsequent `read_object` call to read out only what was written with the first `write_object` call.  Here's an example:

```
>>> import pygama.lh5 as lh5
>>> import h5py
>>> import numpy as np
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset2': lh5.Array(np.zeros(10))})
>>> store = lh5.Store()
>>> store.write_object(tb1, 'my_group', 'my_output.lh5')
>>> store.write_object(tb2, 'my_group', 'my_output.lh5')
>>> tb_dat, _ = store.read_object('my_group', 'my_output.lh5')
>>> tb_dat.keys()
dict_keys(['dset1'])
>>> f = h5py.File('my_output.lh5', 'r')
>>> f.keys()
<KeysViewHDF5 ['my_group']>
>>> f['my_group'].attrs['datatype']
'table{dset1}'
```

Here only the `dset1` dataset was read with `read_object` call, and, by opening the file directly, we can see that this is because only `dset1` is in "datatype" attribute for the HDF5 group.  I would have expected that `dset2` also be in the "datatype" attribute (and that `dset2` be also read in with a `read_object` call).  I would make a pull request to create such behavior in store.py, but I do not currently understand enough of the class' code to do so.

Now, in the above example, I would get the desired result if I add `dset2` to the `col_dict` argument in `lh5.Table` when creating the `tb1` object.  However, my use case requires me to add new datasets to the same group in separate steps, so such a fix is not possible.  I can also imagine a use case where new datasets are added to the same group using `write_object` using different scripts, which runs into this same issue when `read_object` is run with the name of the group.


Discussion / hints:
Jason take a look in refactor
 I've tried taking a look at this, especially after the `lgdo` refactor, to see if this is still an issue. Running the same block of code as above results in the following: 
``` 
>>> import pygama.lgdo as lh5 
>>> import h5py
>>> import numpy as np 
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset2': lh5.Array(np.zeros(10))})
>>> store = lh5.Store()
>>> store.write_object(tb1, 'my_group', 'my_output.lh5')
>>> store.write_object(tb2, 'my_group', 'my_output.lh5')
RuntimeError: grp_attrs != group.attrs but overwrite not set
```
where the `RuntimeError` is raised in a `gimme_group` function call in the second `write_object`, so `pygama` doesn't repeat the bug and instead errors out -- nice!  The question is still open as to if @sweigart 's use case can be implemented currently. Can new datasets be written to the same group? 

```
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset2': lh5.Array(np.ones(90))})
>>> store = lh5.LH5Store()
>>> store.write_object(tb1, name='dset1', group='my_group', lh5_file="my_output.lh5")
>>> store.write_object(tb2, name='dset2', group='my_group', lh5_file="my_output.lh5")
>>> tb_dat1, _ = store.read_object('my_group/dset1',  "my_output.lh5")
>>> tb_dat1.keys()
dict_keys(['dset1'])
>>> len(tb_dat1['dset1'])
10
>>> tb_dat2, _ = store.read_object('my_group/dset2',  "my_output.lh5")
>>> tb_dat2.keys()
dict_keys(['dset2'])
>>> len(tb_dat2['dset2'])
90
```
Ok, so we can definitely add new dataset subgroups to the same group, provided we give the datasets different names. 

But what if the use case we are imagining is more like trying to add a new column to an existing table in an `lh5` file? I tried running the following code to test: 

```
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset2': lh5.Array(np.ones(10))})  # second table with new field name
>>> store = lh5.LH5Store()
>>> store.write_object(tb1, name='table', lh5_file="my_output.lh5")
>>> store.write_object(tb2, name='table', lh5_file="my_output.lh5")
RuntimeError: grp_attrs != group.attrs but overwrite not set
```
The ideal output of this use case would be that the `table` dataset in the `lh5` file has the fields `dset1` and `dset2`. 
I think that this issue boils down to this: do we want to implement a way to add a column to an existing `lgdo.table` in a file, or if this behavior would be unpygamic? Adding this functionality would open a whole new set of issues, such as how should the table be resized appropriately, what happens if we try to write a second table with a field that has the same name as an existing field, etc. Does anyone have any thoughts?
Also, I found this interesting behavior with the `overwrite` option of the `wo_mode` in `write_object`: 

```
>>> import pygama.lgdo as lh5 
>>> import h5py
>>> import numpy as np
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset2': lh5.Array(np.ones(10))})  # Different field name from tb1
>>> store = lh5.LH5Store()
>>> store.write_object(tb1, 'my_table', "my_output.lh5")
>>> store.write_object(tb2, 'my_table', "my_output.lh5", wo_mode="overwrite") # Try overwriting the table
>>> tb_dat, _ = store.read_object('my_table',  "my_output.lh5")
>>> tb_dat.keys()
dict_keys(['dset2'])
>>> tb_dat['dset2']
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
``` 
This looks fine: I would expect that the `overwrite` mode would remove `tb1` and its `dset1` field from `my_table` and then `tb2`'s `dset2` column would be written in its place. However, the following happens when I use `h5py` to look at the file:

```
>>> f = h5py.File("my_output.lh5", 'r')
>>> f.keys()
<KeysViewHDF5 ['my_table']>
>>> f['my_table'].attrs['datatype']
table{dset2}
>>> f['my_table']['dset2']
<HDF5 dataset "dset2": shape (10,), type "<f8">
>>> f['my_table'].keys()
<KeysViewHDF5 ['dset1', 'dset2']>
>>> f['my_table']['dset1']
<HDF5 dataset "dset1": shape (10,), type "<f8">
>>> f['my_table']['dset1'][:]
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
```
`dset1` is still accessible in the file even though it isn't listed in `my_table`'s `datatype` attribute! I suspect that the `overwrite` option only works as intended when the field names in the original tables are identical. Let's see if this is true:

```
>>> tb1 = lh5.Table(col_dict={'dset1': lh5.Array(np.zeros(10))})
>>> tb2 = lh5.Table(col_dict={'dset1': lh5.Array(np.ones(10))})  # Same field name as tb1, but filled with ones
>>> store = lh5.LH5Store()
>>> store.write_object(tb1, 'my_table', "my_output.lh5")
>>> store.write_object(tb2, 'my_table', "my_output.lh5", wo_mode="overwrite")
>>> f = h5py.File("my_output.lh5", 'r')
>>> f['my_table'].attrs['datatype']
table{dset1}
>>> f['my_table'].keys()
<KeysViewHDF5 ['dset1']>
>>> f['my_table']['dset1']
<HDF5 dataset "dset1": shape (10,), type "<f8">
>>> f['my_table']['dset1'][:]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
```
Ok, so `overwrite` does indeed work when both tables have the same field names. 

Should we add functionality to the `overwrite` mode to allow for tables with differently named fields to replace a table in an `lh5` file? Or should we instead raise an error that the user is trying to overwrite a table that contains different fields?
> do we want to implement a way to add a column to an existing lgdo.table in a file, or if this behavior would be unpygamic? Adding this functionality would open a whole new set of issues, such as how should the table be resized appropriately, what happens if we try to write a second table with a field that has the same name as an existing field, etc.

Hi @SamuelBorden! The first important thing, as you already mentioned, is that adding this would be maybe tricky and requires a careful implementation. I think adding columns to an existing table on disk is a nice feature in general -- the question is whether it's critical for us now. If not, we can keep this open and implement it at a later point.

> dset1 is still accessible in the file even though it isn't listed in my_table's datatype attribute! 

This is a bug in my opinion and should be fixed.

Fix patch:
diff --git a/src/pygama/lgdo/lh5_store.py b/src/pygama/lgdo/lh5_store.py
index cdd04a364..4ba63be0f 100644
--- a/src/pygama/lgdo/lh5_store.py
+++ b/src/pygama/lgdo/lh5_store.py
@@ -97,7 +97,7 @@ def gimme_file(self, lh5_file: str | h5py.File, mode: str = "r") -> h5py.File:
 
     def gimme_group(
         self,
-        group: str,
+        group: str | h5py.Group,
         base_group: h5py.Group,
         grp_attrs: dict[str, Any] = None,
         overwrite: bool = False,
@@ -727,6 +727,14 @@ def write_object(
             group = self.gimme_group(
                 name, group, grp_attrs=obj.attrs, overwrite=(wo_mode == "o")
             )
+            # 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.keys():
                 self.write_object(
                     obj[field],


Test patch:
diff --git a/tests/lgdo/test_lh5_store.py b/tests/lgdo/test_lh5_store.py
index 1ff6b712e..bac00f8bf 100644
--- a/tests/lgdo/test_lh5_store.py
+++ b/tests/lgdo/test_lh5_store.py
@@ -1,3 +1,6 @@
+import logging
+import os
+
 import h5py
 import numpy as np
 import pandas as pd
@@ -318,3 +321,170 @@ def test_read_lgnd_waveform_table_fancy_idx(lgnd_file):
     )
     assert isinstance(lh5_obj, lgdo.WaveformTable)
     assert len(lh5_obj) == 19
+
+
+# First test that we can overwrite a table with the same name without deleting the original field
+def test_write_object_overwrite_table_no_deletion(caplog):
+    caplog.set_level(logging.DEBUG)
+    caplog.clear()
+
+    if os.path.exists("/tmp/write_object_overwrite_test.lh5"):
+        os.remove("/tmp/write_object_overwrite_test.lh5")
+
+    tb1 = lh5.Table(col_dict={"dset1": lh5.Array(np.zeros(10))})
+    tb2 = lh5.Table(
+        col_dict={"dset1": lh5.Array(np.ones(10))}
+    )  # Same field name, different values
+    store = LH5Store()
+    store.write_object(tb1, "my_group", "/tmp/write_object_overwrite_test.lh5")
+    store.write_object(
+        tb2, "my_group", "/tmp/write_object_overwrite_test.lh5", wo_mode="overwrite"
+    )  # Now, try to overwrite the same field
+
+    # If the old field is deleted from the file before writing the new field, then we would get an extra debug statement
+    assert "dset1 is not present in new table, deleting field" not in [
+        rec.message for rec in caplog.records
+    ]
+
+    # Now, check that the data were overwritten
+    tb_dat, _ = store.read_object("my_group", "/tmp/write_object_overwrite_test.lh5")
+    assert np.array_equal(tb_dat["dset1"].nda, np.ones(10))
+
+
+# Second: test that when we overwrite a table with fields with a different name, we delete the original field
+def test_write_object_overwrite_table_with_deletion(caplog):
+    caplog.set_level(logging.DEBUG)
+    caplog.clear()
+
+    if os.path.exists("/tmp/write_object_overwrite_test.lh5"):
+        os.remove("/tmp/write_object_overwrite_test.lh5")
+
+    tb1 = lh5.Table(col_dict={"dset1": lh5.Array(np.zeros(10))})
+    tb2 = lh5.Table(
+        col_dict={"dset2": lh5.Array(np.ones(10))}
+    )  # Same field name, different values
+    store = LH5Store()
+    store.write_object(tb1, "my_group", "/tmp/write_object_overwrite_test.lh5")
+    store.write_object(
+        tb2, "my_group", "/tmp/write_object_overwrite_test.lh5", wo_mode="overwrite"
+    )  # Now, try to overwrite with a different field
+
+    # If the old field is deleted from the file before writing the new field, then we would get a debug statement
+    assert [rec.message for rec in caplog.records][
+        2
+    ] == "dset1 is not present in new table, deleting field"
+
+    # Now, check that the data were overwritten
+    tb_dat, _ = store.read_object("my_group", "/tmp/write_object_overwrite_test.lh5")
+    assert np.array_equal(tb_dat["dset2"].nda, np.ones(10))
+
+    # Also make sure that the first table's fields aren't lurking around the lh5 file!
+    with h5py.File("/tmp/write_object_overwrite_test.lh5", "r") as lh5file:
+        assert "dset1" not in list(lh5file["my_group"].keys())
+
+
+# Third: test that when we overwrite other LGDO classes
+def test_write_object_overwrite_lgdo(caplog):
+    caplog.set_level(logging.DEBUG)
+    caplog.clear()
+
+    # Start with an lgdo.WaveformTable
+    if os.path.exists("/tmp/write_object_overwrite_test.lh5"):
+        os.remove("/tmp/write_object_overwrite_test.lh5")
+
+    tb1 = lh5.WaveformTable(
+        t0=np.zeros(10),
+        t0_units="ns",
+        dt=np.zeros(10),
+        dt_units="ns",
+        values=np.zeros((10, 10)),
+        values_units="ADC",
+    )
+    tb2 = lh5.WaveformTable(
+        t0=np.ones(10),
+        t0_units="ns",
+        dt=np.ones(10),
+        dt_units="ns",
+        values=np.ones((10, 10)),
+        values_units="ADC",
+    )  # Same field name, different values
+    store = LH5Store()
+    store.write_object(
+        tb1, "my_table", "/tmp/write_object_overwrite_test.lh5", group="my_group"
+    )
+    store.write_object(
+        tb2,
+        "my_table",
+        "/tmp/write_object_overwrite_test.lh5",
+        wo_mode="overwrite",
+        group="my_group",
+    )
+
+    # If the old field is deleted from the file before writing the new field, then we would get a debug statement
+    assert "my_table is not present in new table, deleting field" not in [
+        rec.message for rec in caplog.records
+    ]
+
+    # Now, check that the data were overwritten
+    tb_dat, _ = store.read_object(
+        "my_group/my_table", "/tmp/write_object_overwrite_test.lh5"
+    )
+    assert np.array_equal(tb_dat["values"].nda, np.ones((10, 10)))
+
+    # Now try overwriting an array, and test the write_start argument
+    array1 = lh5.Array(nda=np.zeros(10))
+    array2 = lh5.Array(nda=np.ones(20))
+    store = LH5Store()
+    store.write_object(array1, "my_array", "/tmp/write_object_overwrite_test.lh5")
+    store.write_object(
+        array2,
+        "my_array",
+        "/tmp/write_object_overwrite_test.lh5",
+        wo_mode="overwrite",
+        write_start=5,
+    )
+
+    # Now, check that the data were overwritten
+    array_dat, _ = store.read_object("my_array", "/tmp/write_object_overwrite_test.lh5")
+    expected_out_array = np.append(np.zeros(5), np.ones(20))
+
+    assert np.array_equal(array_dat.nda, expected_out_array)
+
+    # Now try overwriting a scalar
+    scalar1 = lh5.Scalar(0)
+    scalar2 = lh5.Scalar(1)
+    store = LH5Store()
+    store.write_object(scalar1, "my_scalar", "/tmp/write_object_overwrite_test.lh5")
+    store.write_object(
+        scalar2,
+        "my_scalar",
+        "/tmp/write_object_overwrite_test.lh5",
+        wo_mode="overwrite",
+    )
+
+    # Now, check that the data were overwritten
+    scalar_dat, _ = store.read_object(
+        "my_scalar", "/tmp/write_object_overwrite_test.lh5"
+    )
+
+    assert scalar_dat.value == 1
+
+    # Finally, try overwriting a vector of vectors
+    vov1 = lh5.VectorOfVectors(listoflists=[np.zeros(1), np.ones(2), np.zeros(3)])
+    vov2 = lh5.VectorOfVectors(listoflists=[np.ones(1), np.zeros(2), np.ones(3)])
+    store = LH5Store()
+    store.write_object(vov1, "my_vector", "/tmp/write_object_overwrite_test.lh5")
+    store.write_object(
+        vov2,
+        "my_vector",
+        "/tmp/write_object_overwrite_test.lh5",
+        wo_mode="overwrite",
+        write_start=1,
+    )  # start overwriting the second list of lists
+
+    vector_dat, _ = store.read_object(
+        "my_vector", "/tmp/write_object_overwrite_test.lh5"
+    )
+
+    assert np.array_equal(vector_dat.cumulative_length.nda, [1, 2, 4, 7])
+    assert np.array_equal(vector_dat.flattened_data.nda, [0, 1, 0, 0, 1, 1, 1])