fix: missing mulitiple metrics on pivot operator (apache#16026)

* fix: missing mulitiple metrics on pivot operator * code smell
opus-42 · Nov 14, 2021 · 27f14a5 · 27f14a5
1 parent a18f7c9
commit 27f14a5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git 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
@@ -258,6 +258,7 @@ def test_pivot_exceptions(self):
         )
 
     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 @@ def test_pivot_eliminate_cartesian_product_columns(self):
         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"},