diff --git a/superset/utils/pandas_postprocessing.py b/superset/utils/pandas_postprocessing.py index 0d8105c756..0e03dea61c 100644 --- a/superset/utils/pandas_postprocessing.py +++ b/superset/utils/pandas_postprocessing.py @@ -271,8 +271,9 @@ def pivot( # pylint: disable=too-many-arguments series_set = set() if not drop_missing_columns and columns: for row in df[columns].itertuples(): - metrics_and_series = tuple(aggfunc.keys()) + tuple(row[1:]) - series_set.add(str(metrics_and_series)) + for metric in aggfunc.keys(): + series_set.add(str(tuple([metric]) + tuple(row[1:]))) + df = df.pivot_table( values=aggfunc.keys(), index=index, diff --git a/tests/integration_tests/pandas_postprocessing_tests.py b/tests/integration_tests/pandas_postprocessing_tests.py index 57bcc1fb92..1763dad336 100644 --- a/tests/integration_tests/pandas_postprocessing_tests.py +++ b/tests/integration_tests/pandas_postprocessing_tests.py @@ -258,6 +258,7 @@ class TestPostProcessing(SupersetTestCase): ) def test_pivot_eliminate_cartesian_product_columns(self): + # single metric mock_df = DataFrame( { "dttm": to_datetime(["2019-01-01", "2019-01-01"]), @@ -277,6 +278,33 @@ class TestPostProcessing(SupersetTestCase): self.assertEqual(list(df.columns), ["dttm", "0, 0", "1, 1"]) self.assertTrue(np.isnan(df["1, 1"][0])) + # multiple metrics + mock_df = DataFrame( + { + "dttm": to_datetime(["2019-01-01", "2019-01-01"]), + "a": [0, 1], + "b": [0, 1], + "metric": [9, np.NAN], + "metric2": [10, 11], + } + ) + + df = proc.pivot( + df=mock_df, + index=["dttm"], + columns=["a", "b"], + aggregates={ + "metric": {"operator": "mean"}, + "metric2": {"operator": "mean"}, + }, + drop_missing_columns=False, + ) + self.assertEqual( + list(df.columns), + ["dttm", "metric, 0, 0", "metric, 1, 1", "metric2, 0, 0", "metric2, 1, 1"], + ) + self.assertTrue(np.isnan(df["metric, 1, 1"][0])) + def test_aggregate(self): aggregates = { "asc sum": {"column": "asc_idx", "operator": "sum"},