fix: big number to handle NULL as it did in the past (#9314)

apache · Mar 17, 2020 · 6cf36c9 · 6cf36c9
1 parent 85e9a4f
commit 6cf36c9
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/superset/viz.py b/superset/viz.py
@@ -1126,8 +1126,8 @@ def get_data(self, df: pd.DataFrame) -> VizData:
             index=DTTM_ALIAS,
             columns=[],
             values=self.metric_labels,
-            fill_value=0,
-            aggfunc=sum,
+            dropna=False,
+            aggfunc=np.min,  # looking for any (only) value, preserving `None`
         )
         df = self.apply_rolling(df)
         df[DTTM_ALIAS] = df.index

diff --git a/tests/viz_tests.py b/tests/viz_tests.py
@@ -1243,3 +1243,31 @@ def test_apply_rolling(self):
             .tolist(),
             [1.0, 1.5, 2.0, 2.5],
         )
+
+
+class BigNumberVizTestCase(SupersetTestCase):
+    def test_get_data(self):
+        datasource = self.get_datasource_mock()
+        df = pd.DataFrame(
+            data={
+                DTTM_ALIAS: pd.to_datetime(
+                    ["2019-01-01", "2019-01-02", "2019-01-05", "2019-01-07"]
+                ),
+                "y": [1.0, 2.0, 3.0, 4.0],
+            }
+        )
+        data = viz.BigNumberViz(datasource, {"metrics": ["y"]}).get_data(df)
+        self.assertEqual(data[2], {DTTM_ALIAS: pd.Timestamp("2019-01-05"), "y": 3})
+
+    def test_get_data_with_none(self):
+        datasource = self.get_datasource_mock()
+        df = pd.DataFrame(
+            data={
+                DTTM_ALIAS: pd.to_datetime(
+                    ["2019-01-01", "2019-01-02", "2019-01-05", "2019-01-07"]
+                ),
+                "y": [1.0, 2.0, None, 4.0],
+            }
+        )
+        data = viz.BigNumberViz(datasource, {"metrics": ["y"]}).get_data(df)
+        assert np.isnan(data[2]["y"])