Handle Day / Tick change

From pandas-dev/pandas#61985

Author: TomAugspurger

dask/dataframe/dask_expr/_util.py

@@ -43,10 +43,12 @@ def _calc_maybe_new_divisions(df, periods, freq):
 
     is_offset = isinstance(freq, pd.DateOffset)
     if is_offset:
-        if not isinstance(freq, pd.offsets.Tick):
+        if not isinstance(freq, (pd.offsets.Tick, pd.offsets.Day)):
             # Can't infer divisions on relative or anchored offsets, as
             # divisions may now split identical index value.
             # (e.g. index_partitions = [[1, 2, 3], [3, 4, 5]])
+            # As of https://github.com/pandas-dev/pandas/pull/61985
+            # Day doesn't subclass Tick.
             return None  # Would need to clear divisions
     if df.known_divisions:
         divs = pd.Series(range(len(df.divisions)), index=df.divisions)

dask/dataframe/dask_expr/tests/test_collection.py

@@ -2461,14 +2461,25 @@ def test_scalar_to_series():
 
 
 @pytest.mark.parametrize("data_freq, divs1", [("B", False), ("D", True), ("h", True)])
-def test_shift_with_freq_datetime(pdf, data_freq, divs1):
+@pytest.mark.parametrize(
+    "freq_divs2", [("s", True), ("W", False), (pd.Timedelta(10, unit="h"), True)]
+)
+@pytest.mark.parametrize("as_series", [True, False])
+def test_shift_with_freq_datetime(pdf, data_freq, divs1, freq_divs2, as_series):
     pdf.index = pd.date_range(start="2020-01-01", periods=len(pdf), freq=data_freq)
     df = from_pandas(pdf, npartitions=4)
-    for freq, divs2 in [("s", True), ("W", False), (pd.Timedelta(10, unit="h"), True)]:
-        for d, p in [(df, pdf), (df.x, pdf.x)]:
-            res = d.shift(2, freq=freq)
-            assert_eq(res, p.shift(2, freq=freq))
-            assert res.known_divisions == divs2
+
+    freq, divs2 = freq_divs2
+    if as_series:
+        pd_obj = pdf["x"]
+        dask_obj = df["x"]
+    else:
+        pd_obj = pdf
+        dask_obj = df
+
+    res = dask_obj.shift(2, freq=freq)
+    assert_eq(res, pd_obj.shift(2, freq=freq))
+    assert res.known_divisions == divs2
 
     res = df.index.shift(2)
     assert_eq(res, df.index.shift(2))