[fbprophet] Fix frequencies

apache · Jun 14, 2022 · 34043c4 · 34043c4
1 parent fd12987
commit 34043c4
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 4 deletions.
diff --git a/superset/utils/pandas_postprocessing/utils.py b/superset/utils/pandas_postprocessing/utils.py
@@ -86,10 +86,10 @@
     "P1M": "M",
     "P3M": "Q",
     "P1Y": "A",
-    "1969-12-28T00:00:00Z/P1W": "W",
-    "1969-12-29T00:00:00Z/P1W": "W",
-    "P1W/1970-01-03T00:00:00Z": "W",
-    "P1W/1970-01-04T00:00:00Z": "W",
+    "1969-12-28T00:00:00Z/P1W": "W-SUN",
+    "1969-12-29T00:00:00Z/P1W": "W-MON",
+    "P1W/1970-01-03T00:00:00Z": "W-SAT",
+    "P1W/1970-01-04T00:00:00Z": "W-SUN",
 }
 
 RESAMPLE_METHOD = ("asfreq", "bfill", "ffill", "linear", "median", "mean", "sum")

diff --git a/tests/unit_tests/pandas_postprocessing/test_prophet.py b/tests/unit_tests/pandas_postprocessing/test_prophet.py
@@ -17,6 +17,7 @@
 from datetime import datetime
 from importlib.util import find_spec
 
+import pandas as pd
 import pytest
 
 from superset.exceptions import InvalidPostProcessingError
@@ -50,6 +51,66 @@ def test_prophet_valid():
     assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 5, 31)
     assert len(df) == 9
 
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 2), datetime(2022, 1, 9)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 16)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 2), datetime(2022, 1, 9)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="1969-12-28T00:00:00Z/P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 16)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 3), datetime(2022, 1, 10)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="1969-12-29T00:00:00Z/P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 17)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 8), datetime(2022, 1, 15)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="P1W/1970-01-03T00:00:00Z",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 22)
+    assert len(df) == 3
+
 
 def test_prophet_valid_zero_periods():
     pytest.importorskip("prophet")