legend-pydataobj 1.5.0a5__py3-none-any.whl → 1.6.0__py3-none-any.whl

lgdo/types/table.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing a special struct of arrays of
 equal length and corresponding utilities.
 """
+
 from __future__ import annotations
 
 import logging
@@ -199,6 +200,11 @@ class Table(Struct):
     ) -> pd.DataFrame:
         """Get a :class:`pandas.DataFrame` from the data in the table.
 
+        Warning
+        -------
+        This method is deprecated. Use :meth:`.view_as` to view the table as a
+        Pandas dataframe.
+
         Notes
         -----
         The requested data must be array-like, with the ``nda`` attribute.
@@ -224,6 +230,29 @@ class Table(Struct):
         )
         return self.view_as(library="pd", cols=cols, prefix=prefix)
 
+    def flatten(self, _prefix="") -> Table:
+        """Flatten the table, if nested.
+
+        Returns a new :class:`Table` (that references, not copies, the existing
+        columns) with columns in nested tables being moved to the first level
+        (and renamed appropriately).
+
+        Examples
+        --------
+        >>> repr(tbl)
+        "Table(dict={'a': Array([1 2 3], attrs={'datatype': 'array<1>{real}'}), 'tbl': Table(dict={'b': Array([4 5 6], attrs={'datatype': 'array<1>{real}'}), 'tbl1': Table(dict={'z': Array([9 9 9], attrs={'datatype': 'array<1>{real}'})}, attrs={'datatype': 'table{z}'})}, attrs={'datatype': 'table{b,tbl1}'})}, attrs={'datatype': 'table{a,tbl}'})"
+        >>> tbl.flatten().keys()
+        dict_keys(['a', 'tbl__b', 'tbl__tbl1__z'])
+        """
+        flat_table = Table(size=self.size)
+        for key, obj in self.items():
+            if isinstance(obj, Table):
+                flat_table.join(obj.flatten(_prefix=f"{_prefix}{key}__"))
+            else:
+                flat_table.add_column(_prefix + key, obj)
+
+        return flat_table
+
     def eval(
         self,
         expr: str,
@@ -237,6 +266,13 @@ class Table(Struct):
         columns are viewed as :class:`ak.Array` and the respective routines are
         therefore available.
 
+        To columns nested in subtables can be accessed by scoping with two
+        underscores (``__``). For example: ::
+
+          tbl.eval("a + tbl2__b")
+
+        computes the sum of column `a` and column `b` in the subtable `tbl2`.
+
         Parameters
         ----------
         expr
@@ -280,15 +316,16 @@ class Table(Struct):
 
         # make a dictionary of low-level objects (numpy or awkward)
         # for later computation
+        flat_self = self.flatten()
         self_unwrap = {}
         has_ak = False
         for obj in c.co_names:
-            if obj in self.keys():
-                if isinstance(self[obj], VectorOfVectors):
-                    self_unwrap[obj] = self[obj].view_as("ak", with_units=False)
+            if obj in flat_self:
+                if isinstance(flat_self[obj], VectorOfVectors):
+                    self_unwrap[obj] = flat_self[obj].view_as("ak", with_units=False)
                     has_ak = True
                 else:
-                    self_unwrap[obj] = self[obj].view_as("np", with_units=False)
+                    self_unwrap[obj] = flat_self[obj].view_as("np", with_units=False)
 
         # use numexpr if we are only dealing with numpy data types
         if not has_ak:
@@ -314,7 +351,7 @@ class Table(Struct):
 
         # resort to good ol' eval()
         globs = {"ak": ak, "np": np}
-        out_data = eval(expr, globs, (self_unwrap | parameters))  # noqa: PGH001
+        out_data = eval(expr, globs, (self_unwrap | parameters))
 
         # need to convert back to LGDO
         if isinstance(out_data, ak.Array):
@@ -381,41 +418,44 @@ class Table(Struct):
             forward physical units to the output data.
         cols
             a list of column names specifying the subset of the table's columns
-            to be added to the dataframe.
+            to be added to the data view structure.
         prefix
-            The prefix to be added to the column names. Used when recursively getting the
-            dataframe of a table inside this table.
+            The prefix to be added to the column names. Used when recursively
+            getting the dataframe of a :class:`Table` inside this
+            :class:`Table`.
 
         See Also
         --------
         .LGDO.view_as
         """
+        if cols is None:
+            cols = self.keys()
+
         if library == "pd":
             df = pd.DataFrame()
-            if cols is None:
-                cols = self.keys()
+
             for col in cols:
-                column = self[col]
-                if isinstance(column, (Array, VectorOfVectors)):
-                    tmp_ser = column.view_as("pd", with_units=with_units).rename(
-                        prefix + str(col)
-                    )
-                    df = pd.DataFrame(tmp_ser) if df.empty else df.join(tmp_ser)
-                elif isinstance(column, Table):
-                    tmp_df = column.view_as(
+                data = self[col]
+
+                if isinstance(data, Table):
+                    log.debug(f"viewing Table {col=!r} recursively")
+
+                    tmp_df = data.view_as(
                         "pd", with_units=with_units, prefix=f"{prefix}{col}_"
                     )
-                    df = tmp_df if df.empty else df.join(tmp_df)
-                elif df.empty:
-                    df[prefix + str(col)] = column.view_as("pd", with_units=with_units)
+                    for k, v in tmp_df.items():
+                        df[k] = v
+
                 else:
-                    df[prefix + str(col)] = df.join(
-                        column.view_as("pd", with_units=with_units)
+                    log.debug(
+                        f"viewing {type(data).__name__} column {col!r} as Pandas Series"
                     )
+                    df[f"{prefix}{col}"] = data.view_as("pd", with_units=with_units)
+
             return df
 
         if library == "np":
-            msg = f"Format {library} is not supported for Tables."
+            msg = f"Format {library!r} is not supported for Tables."
             raise TypeError(msg)
 
         if library == "ak":
@@ -423,7 +463,10 @@ class Table(Struct):
                 msg = "Pint does not support Awkward yet, you must view the data with_units=False"
                 raise ValueError(msg)
 
-            return ak.Array(self)
+            # NOTE: passing the Table directly (which inherits from a dict)
+            # makes it somehow really slow. Not sure why, but this could be due
+            # to extra LGDO fields (like "attrs")
+            return ak.Array({col: self[col].view_as("ak") for col in cols})
 
-        msg = f"{library} is not a supported third-party format."
+        msg = f"{library!r} is not a supported third-party format."
         raise TypeError(msg)