Repository: mwaskom/seaborn

Issue:
`FacetGrid` data no longer contains all columns from original input dataframe
After recently upgrading from seaborn 0.9.1 to 0.11.1, I noticed that after using `relplot` on a dataframe, the returned `FacetGrid` object only contains a transformed dataframe which makes `FacetGrid.facet_data()` return a lot less useful for additional plot manipulation using information from the original dataframe. A minimal repro:

```
>>> import seaborn as sns
>>> import pandas as pd
>>> import numpy as np
>>> data = pd.DataFrame(np.random.randint(0,3,size=(15, 7)), columns=list('ABCDEFG'))
>>> data.columns
Index(['A', 'B', 'C', 'D', 'E', 'F', 'G'], dtype='object')
>>> g = sns.relplot(x='A', y='B', row='C', col='D', hue='E', data=data)
>>> g.data.columns
Index(['x', 'y', 'hue', 'size', 'style', 'units', 'C', 'D'], dtype='object')
>>> for (i, j, k), data_ijk in g.facet_data():
...     data_ijk.F  # this no longer works
```

Previously, `g.data` would contain all of the original dataframe columns, and I was relying on that so that I could use `g.facet_data()` to augment plots with additional information from the original dataframe.

Is there a way to get back that functionality?

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom) Report
> Merging [#2623](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom) (f9016cf) into [master](https://codecov.io/gh/mwaskom/seaborn/commit/9c3dba6eb0193552123eaa14fac1a90974fa24f9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom) (9c3dba6) will **decrease** coverage by `0.46%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/mwaskom/seaborn/pull/2623/graphs/tree.svg?width=650&height=150&src=pr&token=SLuzHnA9Bl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom)](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom)

```diff
@@            Coverage Diff             @@
##           master    #2623      +/-   ##
==========================================
- Coverage   97.45%   96.99%   -0.47%     
==========================================
  Files          17       17              
  Lines        6374     6350      -24     
==========================================
- Hits         6212     6159      -53     
- Misses        162      191      +29     
```


| [Impacted Files](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom) | Coverage Δ | |
|---|---|---|
| [seaborn/distributions.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9kaXN0cmlidXRpb25zLnB5) | `96.05% <100.00%> (-0.33%)` | :arrow_down: |
| [seaborn/relational.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9yZWxhdGlvbmFsLnB5) | `99.69% <100.00%> (+<0.01%)` | :arrow_up: |
| [seaborn/\_docstrings.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9fZG9jc3RyaW5ncy5weQ==) | `89.18% <0.00%> (-10.82%)` | :arrow_down: |
| [seaborn/\_statistics.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9fc3RhdGlzdGljcy5weQ==) | `98.46% <0.00%> (-1.54%)` | :arrow_down: |
| [seaborn/regression.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9yZWdyZXNzaW9uLnB5) | `96.65% <0.00%> (-1.23%)` | :arrow_down: |
| [seaborn/matrix.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9tYXRyaXgucHk=) | `94.96% <0.00%> (-0.92%)` | :arrow_down: |
| [seaborn/categorical.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9jYXRlZ29yaWNhbC5weQ==) | `98.01% <0.00%> (-0.53%)` | :arrow_down: |
| [seaborn/utils.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi91dGlscy5weQ==) | `94.27% <0.00%> (-0.34%)` | :arrow_down: |
| [seaborn/\_core.py](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom#diff-c2VhYm9ybi9fY29yZS5weQ==) | `98.33% <0.00%> (-0.15%)` | :arrow_down: |
| ... and [3 more](https://codecov.io/gh/mwaskom/seaborn/pull/2623/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom). Last update [9c3dba6...f9016cf](https://codecov.io/gh/mwaskom/seaborn/pull/2623?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Michael+Waskom).


Patch:
diff --git a/doc/releases/v0.12.0.txt b/doc/releases/v0.12.0.txt
index bd4bc7ce..c3ff1fbe 100644
--- a/doc/releases/v0.12.0.txt
+++ b/doc/releases/v0.12.0.txt
@@ -32,7 +32,7 @@ A paper describing seaborn was published in the `Journal of Open Source Software
 
 - |Feature| Added a ``refline`` method to :class:`FacetGrid` and :class:`JointGrid` for including horizontal and/or vertical reference lines using :meth:`matplotlib.axes.Axes.axhline`/:meth:`matplotlib.axes.Axes.axvline` (:pr:`2620`).
 
-- |API| In :func:`lmplot`, the `sharex`, `sharey`, and `legend_out` parameters have been deprecated from the function signature, but they can be passed using the new `facet_kws` parameter (:pr:`2576`).
+- |API| |Enhancement| In :func:`lmplot`, added a new `facet_kws` parameter and deprecated the `sharex`, `sharey`, and `legend_out` parameters from the function signature; pass them in the `facet_kws` dictionary (:pr:`2576`).
 
 - |Fix| In :func:`lineplot, allowed the `dashes` keyword to set the style of a line without mapping a `style` variable (:pr:`2449`).
 
@@ -44,6 +44,8 @@ A paper describing seaborn was published in the `Journal of Open Source Software
 
 - |Fix| In the relational plots, fixed a bug where legend entries for the `size` semantic were incorrect when `size_norm` extrapolated beyond the range of the data (:pr:`2580`).
 
+- |Fix| In :func:`relplot` and :func:`displot`, fixed a bug where the dataframe attached to the returned `FacetGrid` object dropped columns that were not used in the plot (:pr:`2623`).
+
 - |Fix| In :func:`relplot`, fixed an error that would be raised when a column used to facet shared a name with one of the plot variables (:pr:`2581`).
 
 - |Fix| In :func:`histplot` and :func:`kdeplot`, fixed a bug where the `alpha` parameter was ignored when `fill=False` (:pr:`2460`).
diff --git a/seaborn/distributions.py b/seaborn/distributions.py
index abcf1aa6..1b0210b7 100644
--- a/seaborn/distributions.py
+++ b/seaborn/distributions.py
@@ -2177,8 +2177,8 @@ def displot(
             p.variables[var] = f"_{var}_"
 
     # Adapt the plot_data dataframe for use with FacetGrid
-    data = p.plot_data.rename(columns=p.variables)
-    data = data.loc[:, ~data.columns.duplicated()]
+    grid_data = p.plot_data.rename(columns=p.variables)
+    grid_data = grid_data.loc[:, ~grid_data.columns.duplicated()]
 
     col_name = p.variables.get("col", None)
     row_name = p.variables.get("row", None)
@@ -2187,7 +2187,7 @@ def displot(
         facet_kws = {}
 
     g = FacetGrid(
-        data=data, row=row_name, col=col_name,
+        data=grid_data, row=row_name, col=col_name,
         col_wrap=col_wrap, row_order=row_order,
         col_order=col_order, height=height,
         aspect=aspect,
@@ -2311,6 +2311,21 @@ def displot(
     g.set_titles()
     g.tight_layout()
 
+    if data is not None and (x is not None or y is not None):
+        if not isinstance(data, pd.DataFrame):
+            data = pd.DataFrame(data)
+        g.data = pd.merge(
+            data,
+            g.data[g.data.columns.difference(data.columns)],
+            left_index=True,
+            right_index=True,
+        )
+    else:
+        wide_cols = {
+            k: f"_{k}_" if v is None else v for k, v in p.variables.items()
+        }
+        g.data = p.plot_data.rename(columns=wide_cols)
+
     return g
 
 
diff --git a/seaborn/relational.py b/seaborn/relational.py
index 86cf692b..6721cdb2 100644
--- a/seaborn/relational.py
+++ b/seaborn/relational.py
@@ -1,6 +1,7 @@
 import warnings
 
 import numpy as np
+import pandas as pd
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 
@@ -1002,7 +1003,18 @@ def relplot(
     orig_cols = {
         f"_{k}": f"_{k}_" if v is None else v for k, v in variables.items()
     }
-    g.data = g.data.rename(columns=orig_cols)
+    grid_data = g.data.rename(columns=orig_cols)
+    if data is not None and (x is not None or y is not None):
+        if not isinstance(data, pd.DataFrame):
+            data = pd.DataFrame(data)
+        g.data = pd.merge(
+            data,
+            grid_data[grid_data.columns.difference(data.columns)],
+            left_index=True,
+            right_index=True,
+        )
+    else:
+        g.data = grid_data
 
     return g
 


Test patch:
diff --git a/seaborn/tests/test_distributions.py b/seaborn/tests/test_distributions.py
index b26191fe..37088227 100644
--- a/seaborn/tests/test_distributions.py
+++ b/seaborn/tests/test_distributions.py
@@ -2320,6 +2320,19 @@ class TestDisPlot:
         clim2 = g.axes.flat[1].collections[0].get_clim()
         assert clim1[1] > clim2[1]
 
+    def test_facetgrid_data(self, long_df):
+
+        g = displot(
+            data=long_df.to_dict(orient="list"),
+            x="z",
+            hue=long_df["a"].rename("hue_var"),
+            col=long_df["c"].to_numpy(),
+        )
+        expected_cols = set(long_df.columns.to_list() + ["hue_var", "_col_"])
+        assert set(g.data.columns) == expected_cols
+        assert_array_equal(g.data["hue_var"], long_df["a"])
+        assert_array_equal(g.data["_col_"], long_df["c"])
+
 
 def integrate(y, x):
     """"Simple numerical integration for testing KDE code."""
diff --git a/seaborn/tests/test_relational.py b/seaborn/tests/test_relational.py
index fadbc911..4ffea504 100644
--- a/seaborn/tests/test_relational.py
+++ b/seaborn/tests/test_relational.py
@@ -617,15 +617,19 @@ class TestRelationalPlotter(Helpers):
         for line, color in zip(lines, palette):
             assert line.get_color() == color
 
-    def test_relplot_data_columns(self, long_df):
+    def test_relplot_data(self, long_df):
 
-        long_df = long_df.assign(x_var=long_df["x"], y_var=long_df["y"])
         g = relplot(
-            data=long_df,
-            x="x_var", y="y_var",
-            hue=long_df["a"].to_numpy(), col="c"
+            data=long_df.to_dict(orient="list"),
+            x="x",
+            y=long_df["y"].rename("y_var"),
+            hue=long_df["a"].to_numpy(),
+            col="c",
         )
-        assert g.data.columns.to_list() == ["x_var", "y_var", "_hue_", "c"]
+        expected_cols = set(long_df.columns.to_list() + ["_hue_", "y_var"])
+        assert set(g.data.columns) == expected_cols
+        assert_array_equal(g.data["y_var"], long_df["y"])
+        assert_array_equal(g.data["_hue_"], long_df["a"])
 
     def test_facet_variable_collision(self, long_df):