[superset] branch master updated: fix: missing mulitiple metrics on pivot operator (#16026)

[pk...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

pkdotson pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 3f57ea7  fix: missing mulitiple metrics on pivot operator (#16026)
3f57ea7 is described below

commit 3f57ea76db6036c52ace9c0e2b28c92e8a1a2677
Author: Yongjie Zhao <yo...@gmail.com>
AuthorDate: Mon Aug 2 19:37:10 2021 +0100

    fix: missing mulitiple metrics on pivot operator (#16026)
    
    * fix: missing mulitiple metrics on pivot operator
    
    * code smell
---
 superset/utils/pandas_postprocessing.py            |  5 ++--
 .../pandas_postprocessing_tests.py                 | 28 ++++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/superset/utils/pandas_postprocessing.py b/superset/utils/pandas_postprocessing.py
index 0d8105c..0e03dea 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 57bcc1f..1763dad 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"},