AliasDataFrame: add index-based subframe join and robust error handling

- Updated `register_subframe()` to explicitly require `index_columns` for join key(s)
- Enhanced `_prepare_subframe_joins()` to:
  - auto-materialize subframe aliases if missing
  - raise informative KeyError when column or alias does not exist
- Added logic to propagate subframe metadata (including join indices) in save/load and ROOT export/import
- Expanded test coverage:
  - Added subframe alias tests for automatic materialization and error reporting
  - Added 2D index subframe join test (e.g. using ["run", "track_id"])
  - Refactored test setup to avoid shared state interference
  - Asserted raised exceptions for missing subframe attributes
- Minor fixes to alias materialization and type assertions

Author: miranov25

UTILS/dfextensions/AliasDataFrame.py

@@ -11,17 +11,23 @@
 
 class SubframeRegistry:
     def __init__(self):
-        self.subframes = {}
+        self.subframes = {}  # name → {'frame': adf, 'index': index_columns}
 
-    def add_subframe(self, name, alias_df):
-        self.subframes[name] = alias_df
+    def add_subframe(self, name, alias_df, index_columns, pre_index=False):
+        if pre_index and not alias_df.df.index.names == index_columns:
+            alias_df.df.set_index(index_columns, inplace=True)
+        self.subframes[name] = {'frame': alias_df, 'index': index_columns}
 
     def get(self, name):
+        return self.subframes.get(name, {}).get('frame', None)
+
+    def get_entry(self, name):
         return self.subframes.get(name, None)
 
     def items(self):
         return self.subframes.items()
 
+
 def convert_expr_to_root(expr):
     class RootTransformer(ast.NodeTransformer):
         FUNC_MAP = {
@@ -77,8 +83,8 @@ def __getattr__(self, item):
             return self.df[item]
         raise AttributeError(f"'{type(self).__name__}' object has no attribute '{item}'")
 
-    def register_subframe(self, name, adf):
-        self._subframes.add_subframe(name, adf)
+    def register_subframe(self, name, adf, index_columns, pre_index=False):
+        self._subframes.add_subframe(name, adf, index_columns, pre_index=pre_index)
 
     def get_subframe(self, name):
         return self._subframes.get(name)
@@ -88,10 +94,41 @@ def _default_functions(self):
         env = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
         env.update({k: getattr(np, k) for k in dir(np) if not k.startswith("_")})
         env["np"] = np
-        for sf_name, sf in self._subframes.items():
-            env[sf_name] = sf
+        for sf_name, sf_entry in self._subframes.items():
+            env[sf_name] = sf_entry['frame']
         return env
 
+    def _prepare_subframe_joins(self, expr):
+        tokens = re.findall(r'(\b\w+)\.(\w+)', expr)
+        for sf_name, sf_col in tokens:
+            entry = self._subframes.get_entry(sf_name)
+            if not entry:
+                continue
+            sub_adf = entry['frame']
+            sub_df = sub_adf.df
+            index_cols = entry['index']
+            if isinstance(index_cols, str):
+                index_cols = [index_cols]
+            merge_cols = index_cols + [sf_col]
+            suffix = f'__{sf_name}'
+
+            try:
+                cols_to_merge = sub_df[merge_cols]
+            except KeyError:
+                if sf_col in sub_adf.aliases:
+                    sub_adf.materialize_alias(sf_col)
+                    sub_df = sub_adf.df
+                    cols_to_merge = sub_df[merge_cols]
+                else:
+                    raise KeyError(f"Subframe '{sf_name}' does not contain or define alias '{sf_col}'")
+
+            joined = self.df.merge(cols_to_merge, on=index_cols, suffixes=('', suffix))
+            col_renamed = f'{sf_col}{suffix}'
+            if col_renamed in joined.columns:
+                self.df[col_renamed] = joined[col_renamed].values
+                expr = expr.replace(f'{sf_name}.{sf_col}', col_renamed)
+        return expr
+
     def _check_for_cycles(self):
         try:
             self._topological_sort()
@@ -107,8 +144,8 @@ def add_alias(self, name, expression, dtype=None, is_constant=False):
         self._check_for_cycles()
 
     def _eval_in_namespace(self, expr):
+        expr = self._prepare_subframe_joins(expr)
         local_env = {col: self.df[col] for col in self.df.columns}
-        local_env.update({k: self.df[k] for k in self.aliases if k in self.df})
         local_env.update(self._default_functions())
         return eval(expr, {}, local_env)
 
@@ -300,8 +337,8 @@ def export_tree(self, filename_or_file, treename="tree", dropAliasColumns=True):
             self._write_metadata_to_root(filename_or_file, treename)
         else:
             self._write_to_uproot(filename_or_file, treename, dropAliasColumns)
-        for subframe_name, sub_adf in self._subframes.items():
-            sub_adf._write_metadata_to_root(filename_or_file, f"{treename}__subframe__{subframe_name}")
+        for subframe_name, entry in self._subframes.items():
+            entry["frame"]._write_metadata_to_root(filename_or_file, f"{treename}__subframe__{subframe_name}")
 
     def _write_to_uproot(self, uproot_file, treename, dropAliasColumns):
         export_cols = [col for col in self.df.columns if not dropAliasColumns or col not in self.aliases]
@@ -310,8 +347,8 @@ def _write_to_uproot(self, uproot_file, treename, dropAliasColumns):
 
         uproot_file[treename] = export_df
 
-        for subframe_name, sub_adf in self._subframes.items():
-            sub_adf.export_tree(uproot_file, f"{treename}__subframe__{subframe_name}", dropAliasColumns)
+        for subframe_name, entry in self._subframes.items():
+            entry["frame"].export_tree(uproot_file, f"{treename}__subframe__{subframe_name}", dropAliasColumns)
 
     def _write_metadata_to_root(self, filename, treename):
         f = ROOT.TFile.Open(filename, "UPDATE")
@@ -325,6 +362,7 @@ def _write_metadata_to_root(self, filename, treename):
             tree.SetAlias(alias, expr_str)
         metadata = {
             "aliases": self.aliases,
+            "subframe_indices": {k: v["index"] for k, v in self._subframes.items()},
             "dtypes": {k: v.__name__ for k, v in self.alias_dtypes.items()},
             "constants": list(self.constant_aliases),
             "subframes": list(self._subframes.subframes.keys())
@@ -334,6 +372,7 @@ def _write_metadata_to_root(self, filename, treename):
         tree.Write("", ROOT.TObject.kOverwrite)
         f.Close()
 
+    @staticmethod
     def read_tree(filename, treename="tree"):
         with uproot.open(filename) as f:
             df = f[treename].arrays(library="pd")
@@ -354,7 +393,10 @@ def read_tree(filename, treename="tree"):
                         adf.constant_aliases.update(jmeta.get("constants", []))
                         for sf_name in jmeta.get("subframes", []):
                             sf = AliasDataFrame.read_tree(filename, treename=f"{treename}__subframe__{sf_name}")
-                            adf.register_subframe(sf_name, sf)
+                            index = jmeta.get("subframe_indices", {}).get(sf_name)
+                            if index is None:
+                                raise ValueError(f"Missing index_columns for subframe '{sf_name}' in metadata")
+                            adf.register_subframe(sf_name, sf, index_columns=index)
                         break
                     except Exception:
                         pass

UTILS/dfextensions/AliasDataFrameTest.py

@@ -94,11 +94,10 @@ def test_export_import_tree_roundtrip(self):
         os.remove(tmp_path)
 
 class TestAliasDataFrameWithSubframes(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
+    def setUp(self):
         n_tracks = 1000
         n_clusters = 100
-        cls.df_tracks = pd.DataFrame({
+        df_tracks = pd.DataFrame({
             "track_index": np.arange(n_tracks),
             "mX": np.random.normal(0, 10, n_tracks),
             "mY": np.random.normal(0, 10, n_tracks),
@@ -107,65 +106,63 @@ def setUpClass(cls):
             "mEta": np.random.normal(0, 1, n_tracks),
         })
 
-        cluster_idx = np.repeat(cls.df_tracks["track_index"], n_clusters)
-        cls.df_clusters = pd.DataFrame({
+        cluster_idx = np.repeat(df_tracks["track_index"], n_clusters)
+        df_clusters = pd.DataFrame({
             "track_index": cluster_idx,
             "mX": np.random.normal(0, 10, len(cluster_idx)),
             "mY": np.random.normal(0, 10, len(cluster_idx)),
             "mZ": np.random.normal(0, 10, len(cluster_idx)),
         })
 
-        cls.adf_tracks = AliasDataFrame(cls.df_tracks)
-        cls.adf_clusters = AliasDataFrame(cls.df_clusters)
-        cls.adf_clusters.register_subframe("T", cls.adf_tracks)
-
-    def test_alias_cluster_radius(self):
-        self.adf_clusters.add_alias("mR", "sqrt(mX**2 + mY**2)")
-        self.adf_clusters.materialize_all()
-        expected = np.sqrt(self.adf_clusters.df["mX"]**2 + self.adf_clusters.df["mY"]**2)
-        pd.testing.assert_series_equal(self.adf_clusters.df["mR"], expected, check_names=False)
+        self.df_tracks = df_tracks
+        self.df_clusters = df_clusters
 
     def test_alias_cluster_track_dx(self):
-        self.adf_clusters.add_alias("mDX", "mX - T.mX")
-        self.adf_clusters.materialize_all()
-        merged = self.adf_clusters.df.merge(self.adf_tracks.df, on="track_index", suffixes=("", "_track"))
-        expected = merged["mX"] - merged["mX_track"]
-        pd.testing.assert_series_equal(self.adf_clusters.df["mDX"].reset_index(drop=True), expected.reset_index(drop=True), check_names=False)
-
-    def test_unregistered_subframe_raises_error(self):
-        adf_tmp = AliasDataFrame(self.df_clusters)
-        adf_tmp.add_alias("mDX", "mX - T.mX")
-        with self.assertRaises(NameError):
-            adf_tmp.materialize_all()
+        adf_clusters = AliasDataFrame(self.df_clusters.copy())
+        adf_tracks = AliasDataFrame(self.df_tracks.copy())
+        adf_clusters.register_subframe("T", adf_tracks, index_columns="track_index")
+        adf_clusters.add_alias("mDX", "mX - T.mX")
+        adf_clusters.materialize_all()
+        merged = adf_clusters.df.merge(adf_tracks.df, on="track_index", suffixes=("", "_trk"))
+        expected = merged["mX"] - merged["mX_trk"]
+        pd.testing.assert_series_equal(adf_clusters.df["mDX"].reset_index(drop=True), expected.reset_index(drop=True), check_names=False)
+
+    def test_subframe_invalid_alias_raises(self):
+        adf_clusters = AliasDataFrame(self.df_clusters.copy())
+        adf_tracks = AliasDataFrame(self.df_tracks.copy())
+        adf_clusters.register_subframe("T", adf_tracks, index_columns="track_index")
+        adf_clusters.add_alias("invalid", "T.nonexistent")
+
+        with self.assertRaises(KeyError) as cm:
+            adf_clusters.materialize_alias("invalid")
+
+        self.assertIn("T", str(cm.exception))
+        self.assertIn("nonexistent", str(cm.exception))
 
     def test_save_and_load_integrity(self):
-        import tempfile
+        adf_clusters = AliasDataFrame(self.df_clusters.copy())
+        adf_tracks = AliasDataFrame(self.df_tracks.copy())
+        adf_clusters.register_subframe("T", adf_tracks, index_columns="track_index")
+        adf_clusters.add_alias("mDX", "mX - T.mX")
+        adf_clusters.materialize_all()
+
         with tempfile.TemporaryDirectory() as tmpdir:
             path_clusters = os.path.join(tmpdir, "clusters.parquet")
             path_tracks = os.path.join(tmpdir, "tracks.parquet")
-            self.adf_clusters.save(path_clusters)
-            self.adf_tracks.save(path_tracks)
+            adf_clusters.save(path_clusters)
+            adf_tracks.save(path_tracks)
 
             adf_tracks_loaded = AliasDataFrame.load(path_tracks)
             adf_clusters_loaded = AliasDataFrame.load(path_clusters)
-            adf_clusters_loaded.register_subframe("T", adf_tracks_loaded)
+            adf_clusters_loaded.register_subframe("T", adf_tracks_loaded, index_columns="track_index")
             adf_clusters_loaded.add_alias("mDX", "mX - T.mX")
             adf_clusters_loaded.materialize_all()
 
-            assert "mDX" in adf_clusters_loaded.df.columns
-            mean_diff = np.mean(adf_clusters_loaded.df["mDX"] - self.adf_clusters.df["mDX"])
-            assert abs(mean_diff) < 1e-3, f"Mean difference too large: {mean_diff}"
-            self.assertDictEqual(self.adf_clusters.aliases, adf_clusters_loaded.aliases)
-
-    def test_export_tree_read_tree_with_subframe(self):
-        with tempfile.NamedTemporaryFile(suffix=".root", delete=False) as tmp:
-            self.adf_clusters.export_tree(tmp.name, treename="clusters")
-            tmp_path = tmp.name
-
-        adf_loaded = AliasDataFrame.read_tree(tmp_path, treename="clusters")
-        self.assertIn("T", adf_loaded._subframes.subframes)
-        self.assertTrue(isinstance(adf_loaded.get_subframe("T"), AliasDataFrame))
-        os.remove(tmp_path)
+            self.assertIn("mDX", adf_clusters_loaded.df.columns)
+            merged = adf_clusters_loaded.df.merge(adf_tracks_loaded.df, on="track_index", suffixes=("", "_trk"))
+            expected = merged["mX"] - merged["mX_trk"]
+            pd.testing.assert_series_equal(adf_clusters_loaded.df["mDX"].reset_index(drop=True), expected.reset_index(drop=True), check_names=False)
+            self.assertDictEqual(adf_clusters.aliases, adf_clusters_loaded.aliases)
 
 if __name__ == "__main__":
     unittest.main()