diff --git a/superset/utils/pandas_postprocessing/cum.py b/superset/utils/pandas_postprocessing/cum.py index 128fa970f5f79..d3eb969f79fc4 100644 --- a/superset/utils/pandas_postprocessing/cum.py +++ b/superset/utils/pandas_postprocessing/cum.py @@ -46,6 +46,7 @@ def cum( """ columns = columns or {} df_cum = df.loc[:, columns.keys()] + df_cum = df_cum.fillna(0) operation = "cum" + operator if operation not in ALLOWLIST_CUMULATIVE_FUNCTIONS or not hasattr( df_cum, operation diff --git a/tests/unit_tests/fixtures/dataframes.py b/tests/unit_tests/fixtures/dataframes.py index 31a275b735ac7..e1499792cba07 100644 --- a/tests/unit_tests/fixtures/dataframes.py +++ b/tests/unit_tests/fixtures/dataframes.py @@ -130,6 +130,11 @@ data={"label": ["x", "y", "z", "q"], "y": [1.0, 2.0, 3.0, 4.0]}, ) +timeseries_with_gap_df = DataFrame( + index=to_datetime(["2019-01-01", "2019-01-02", "2019-01-05", "2019-01-07"]), + data={"label": ["x", "y", "z", "q"], "y": [1.0, 2.0, None, 4.0]}, +) + timeseries_df2 = DataFrame( index=to_datetime(["2019-01-01", "2019-01-02", "2019-01-05", "2019-01-07"]), data={ diff --git a/tests/unit_tests/pandas_postprocessing/test_cum.py b/tests/unit_tests/pandas_postprocessing/test_cum.py index 130e0602520a1..25d7fd045f57b 100644 --- a/tests/unit_tests/pandas_postprocessing/test_cum.py +++ b/tests/unit_tests/pandas_postprocessing/test_cum.py @@ -24,6 +24,7 @@ multiple_metrics_df, single_metric_df, timeseries_df, + timeseries_with_gap_df, ) from tests.unit_tests.pandas_postprocessing.utils import series_to_list @@ -77,6 +78,19 @@ def test_cum(): ) +def test_cum_with_gap(): + # create new column (cumsum) + post_df = pp.cum( + df=timeseries_with_gap_df, + columns={"y": "y2"}, + operator="sum", + ) + assert post_df.columns.tolist() == ["label", "y", "y2"] + assert series_to_list(post_df["label"]) == ["x", "y", "z", "q"] + assert series_to_list(post_df["y"]) == [1.0, 2.0, None, 4.0] + assert series_to_list(post_df["y2"]) == [1.0, 3.0, 3.0, 7.0] + + def test_cum_after_pivot_with_single_metric(): pivot_df = pp.pivot( df=single_metric_df,