Repository: arviz-devs/arviz

Issue:
Add vars / vars_name for to_dataframe
Save some memory and make the call faster.


Discussion:
Can you please locate the file/area where you are raising the suggestion?

Resolving patch:
diff --git a/arviz/data/inference_data.py b/arviz/data/inference_data.py
--- a/arviz/data/inference_data.py
+++ b/arviz/data/inference_data.py
@@ -1,5 +1,6 @@
 # pylint: disable=too-many-lines,too-many-public-methods
 """Data structure for using netcdf groups with xarray."""
+import os
 import re
 import sys
 import uuid
@@ -23,14 +24,13 @@
     Union,
     overload,
 )
-import os
 
 import numpy as np
 import xarray as xr
 from packaging import version
 
 from ..rcparams import rcParams
-from ..utils import HtmlTemplate, _subset_list, either_dict_or_kwargs
+from ..utils import HtmlTemplate, _subset_list, _var_names, either_dict_or_kwargs
 from .base import _extend_xr_method, _make_json_serializable, dict_to_dataset
 
 if sys.version_info[:2] >= (3, 9):
@@ -620,6 +620,8 @@ def to_dataframe(
         self,
         groups=None,
         filter_groups=None,
+        var_names=None,
+        filter_vars=None,
         include_coords=True,
         include_index=True,
         index_origin=None,
@@ -635,6 +637,7 @@ def to_dataframe(
         skipped implicitly.
 
         Raises TypeError if no valid groups are found.
+        Raises ValueError if no data are selected.
 
         Parameters
         ----------
@@ -646,6 +649,15 @@ def to_dataframe(
             If "like", interpret groups as substrings of the real group or metagroup names.
             If "regex", interpret groups as regular expressions on the real group or
             metagroup names. A la `pandas.filter`.
+        var_names : str or list of str, optional
+            Variables to be extracted. Prefix the variables by `~` when you want to exclude them.
+        filter_vars: {None, "like", "regex"}, optional
+            If `None` (default), interpret var_names as the real variables names. If "like",
+            interpret var_names as substrings of the real variables names. If "regex",
+            interpret var_names as regular expressions on the real variables names. A la
+            `pandas.filter`.
+            Like with plotting, sometimes it's easier to subset saying what to exclude
+            instead of what to include
         include_coords: bool
             Add coordinate values to column name (tuple).
         include_index: bool
@@ -677,6 +689,11 @@ def to_dataframe(
         dfs = {}
         for group in group_names:
             dataset = self[group]
+            group_var_names = _var_names(var_names, dataset, filter_vars, "ignore")
+            if (group_var_names is not None) and not group_var_names:
+                continue
+            if group_var_names is not None:
+                dataset = dataset[[var_name for var_name in group_var_names if var_name in dataset]]
             df = None
             coords_to_idx = {
                 name: dict(map(reversed, enumerate(dataset.coords[name].values, index_origin)))
@@ -712,8 +729,11 @@ def to_dataframe(
                     df = dataframe
                     continue
                 df = df.join(dataframe, how="outer")
-            df = df.reset_index()
-            dfs[group] = df
+            if df is not None:
+                df = df.reset_index()
+                dfs[group] = df
+        if not dfs:
+            raise ValueError("No data selected for the dataframe.")
         if len(dfs) > 1:
             for group, df in dfs.items():
                 df.columns = [
diff --git a/arviz/utils.py b/arviz/utils.py
--- a/arviz/utils.py
+++ b/arviz/utils.py
@@ -21,7 +21,7 @@ def _check_tilde_start(x):
     return bool(isinstance(x, str) and x.startswith("~"))
 
 
-def _var_names(var_names, data, filter_vars=None):
+def _var_names(var_names, data, filter_vars=None, errors="raise"):
     """Handle var_names input across arviz.
 
     Parameters
@@ -34,6 +34,8 @@ def _var_names(var_names, data, filter_vars=None):
          interpret var_names as substrings of the real variables names. If "regex",
          interpret var_names as regular expressions on the real variables names. A la
         `pandas.filter`.
+    errors: {"raise", "ignore"}, optional, default="raise"
+        Select either to raise or ignore the invalid names.
 
     Returns
     -------
@@ -44,6 +46,9 @@ def _var_names(var_names, data, filter_vars=None):
             f"'filter_vars' can only be None, 'like', or 'regex', got: '{filter_vars}'"
         )
 
+    if errors not in {"raise", "ignore"}:
+        raise ValueError(f"'errors' can only be 'raise', or 'ignore', got: '{errors}'")
+
     if var_names is not None:
         if isinstance(data, (list, tuple)):
             all_vars = []
@@ -66,14 +71,16 @@ def _var_names(var_names, data, filter_vars=None):
             )
 
         try:
-            var_names = _subset_list(var_names, all_vars, filter_items=filter_vars, warn=False)
+            var_names = _subset_list(
+                var_names, all_vars, filter_items=filter_vars, warn=False, errors=errors
+            )
         except KeyError as err:
             msg = " ".join(("var names:", f"{err}", "in dataset"))
             raise KeyError(msg) from err
     return var_names
 
 
-def _subset_list(subset, whole_list, filter_items=None, warn=True):
+def _subset_list(subset, whole_list, filter_items=None, warn=True, errors="raise"):
     """Handle list subsetting (var_names, groups...) across arviz.
 
     Parameters
@@ -87,6 +94,8 @@ def _subset_list(subset, whole_list, filter_items=None, warn=True):
         names. If "like", interpret `subset` as substrings of the elements in
         `whole_list`. If "regex", interpret `subset` as regular expressions to match
         elements in `whole_list`. A la `pandas.filter`.
+    errors: {"raise", "ignore"}, optional, default="raise"
+        Select either to raise or ignore the invalid names.
 
     Returns
     -------
@@ -142,7 +151,7 @@ def _subset_list(subset, whole_list, filter_items=None, warn=True):
             subset = [item for item in whole_list for name in subset if re.search(name, item)]
 
         existing_items = np.isin(subset, whole_list)
-        if not np.all(existing_items):
+        if not np.all(existing_items) and (errors == "raise"):
             raise KeyError(f"{np.array(subset)[~existing_items]} are not present")
 
     return subset


Test patch:
diff --git a/arviz/tests/base_tests/test_data.py b/arviz/tests/base_tests/test_data.py
--- a/arviz/tests/base_tests/test_data.py
+++ b/arviz/tests/base_tests/test_data.py
@@ -763,6 +763,69 @@ def test_to_dataframe(self, kwargs):
                     )
             assert all(item in test_data.columns for item in ("chain", "draw"))
 
+    @pytest.mark.parametrize(
+        "kwargs",
+        (
+            {
+                "var_names": ["parameter_1", "parameter_2", "variable_1", "variable_2"],
+                "filter_vars": None,
+                "var_results": [
+                    ("posterior", "parameter_1"),
+                    ("posterior", "parameter_2"),
+                    ("prior", "parameter_1"),
+                    ("prior", "parameter_2"),
+                    ("posterior", "variable_1"),
+                    ("posterior", "variable_2"),
+                ],
+            },
+            {
+                "var_names": "parameter",
+                "filter_vars": "like",
+                "groups": "posterior",
+                "var_results": ["parameter_1", "parameter_2"],
+            },
+            {
+                "var_names": "~parameter",
+                "filter_vars": "like",
+                "groups": "posterior",
+                "var_results": ["variable_1", "variable_2", "custom_name"],
+            },
+            {
+                "var_names": [".+_2$", "custom_name"],
+                "filter_vars": "regex",
+                "groups": "posterior",
+                "var_results": ["parameter_2", "variable_2", "custom_name"],
+            },
+            {
+                "var_names": ["lp"],
+                "filter_vars": "regex",
+                "groups": "sample_stats",
+                "var_results": ["lp"],
+            },
+        ),
+    )
+    def test_to_dataframe_selection(self, kwargs):
+        results = kwargs.pop("var_results")
+        idata = from_dict(
+            posterior={
+                "parameter_1": np.random.randn(4, 100),
+                "parameter_2": np.random.randn(4, 100),
+                "variable_1": np.random.randn(4, 100),
+                "variable_2": np.random.randn(4, 100),
+                "custom_name": np.random.randn(4, 100),
+            },
+            prior={
+                "parameter_1": np.random.randn(4, 100),
+                "parameter_2": np.random.randn(4, 100),
+            },
+            sample_stats={
+                "lp": np.random.randn(4, 100),
+            },
+        )
+        test_data = idata.to_dataframe(**kwargs)
+        assert not test_data.empty
+        assert set(test_data.columns).symmetric_difference(results) == set(["chain", "draw"])
+
     def test_to_dataframe_bad(self):
         idata = from_dict(
             posterior={"a": np.random.randn(4, 100, 3, 4, 5), "b": np.random.randn(4, 100)},
@@ -781,6 +844,9 @@ def test_to_dataframe_bad(self):
         with pytest.raises(KeyError):
             idata.to_dataframe(groups=["invalid_group"])
 
+        with pytest.raises(ValueError):
+            idata.to_dataframe(var_names=["c"])
+
     @pytest.mark.parametrize("use", (None, "args", "kwargs"))
     def test_map(self, use):
         idata = load_arviz_data("centered_eight")