Repository: dask/dask
Pull request #10757 (2024-01-08 14:50:56)

Handle datetimeindexes in set_index with unknown divisions
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

This is needed for dask-expr, I couldn't come up with a useful test in here because dask/dask is still broken in other places
see https://github.com/dask-contrib/dask-expr/pull/669


Patch:
diff --git a/dask/base.py b/dask/base.py
index 547bb4063..465339336 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import datetime
+import decimal
 import hashlib
 import inspect
 import os
@@ -1029,6 +1030,7 @@ normalize_token.register(
         slice,
         complex,
         type(Ellipsis),
+        decimal.Decimal,
         datetime.date,
         datetime.time,
         datetime.datetime,
diff --git a/dask/dataframe/methods.py b/dask/dataframe/methods.py
index c061571a8..a288f0bc0 100644
--- a/dask/dataframe/methods.py
+++ b/dask/dataframe/methods.py
@@ -381,7 +381,7 @@ def value_counts_combine(x, sort=True, ascending=False, **groupby_kwargs):
 
 
 def value_counts_aggregate(
-    x, sort=True, ascending=False, normalize=False, total_length=None, **groupby_kwargs
+    x, total_length=None, sort=True, ascending=False, normalize=False, **groupby_kwargs
 ):
     out = value_counts_combine(x, **groupby_kwargs)
     if normalize:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index 663e18a1a..132ddcc1e 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -916,7 +916,10 @@ def set_partitions_pre(s, divisions, ascending=True, na_position="last"):
     partitions[(partitions < 0) | (partitions >= len(divisions) - 1)] = (
         len(divisions) - 2 if ascending else 0
     )
-    partitions[s.isna().values] = len(divisions) - 2 if na_position == "last" else 0
+    nas = s.isna()
+    # We could be a ndarray already (datetime dtype)
+    nas = getattr(nas, "values", nas)
+    partitions[nas] = len(divisions) - 2 if na_position == "last" else 0
     return partitions
 
 


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index a9e39b424..449a63edc 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -5878,19 +5878,29 @@ def test_use_of_weakref_proxy():
     isinstance(res.compute(), pd.Series)
 
 
-def test_is_monotonic_numeric():
-    s = pd.Series(range(20))
-    ds = dd.from_pandas(s, npartitions=5)
-    assert_eq(s.is_monotonic_increasing, ds.is_monotonic_increasing)
-
-    s_2 = pd.Series(range(20, 0, -1))
-    ds_2 = dd.from_pandas(s_2, npartitions=5)
-    assert_eq(s_2.is_monotonic_decreasing, ds_2.is_monotonic_decreasing)
+@pytest.mark.parametrize(
+    "series",
+    [
+        [0, 1, 0, 0, 1, 0],  # not monotonic
+        [0, 1, 1, 2, 2, 3],  # monotonic increasing
+        [0, 1, 2, 3, 4, 5],  # strictly monotonic increasing
+        [0, 0, 0, 0, 0, 0],  # both monotonic increasing and monotonic decreasing
+        [0, 1, 2, 0, 1, 2],  # Partitions are individually monotonic; whole series isn't
+    ],
+)
+@pytest.mark.parametrize("reverse", [False, True])
+@pytest.mark.parametrize("cls", ["Series", "Index"])
+def test_is_monotonic_numeric(series, reverse, cls):
+    if reverse:
+        series = series[::-1]
+    pds = pd.Series(series, index=series)
+    ds = dd.from_pandas(pds, chunksize=3, sort=False)
+    if cls == "Index":
+        pds = pds.index
+        ds = ds.index
 
-    s_3 = pd.Series(list(range(0, 5)) + list(range(0, 20)))
-    ds_3 = dd.from_pandas(s_3, npartitions=5)
-    assert_eq(s_3.is_monotonic_increasing, ds_3.is_monotonic_increasing)
-    assert_eq(s_3.is_monotonic_decreasing, ds_3.is_monotonic_decreasing)
+    assert ds.is_monotonic_increasing.compute() == pds.is_monotonic_increasing
+    assert ds.is_monotonic_decreasing.compute() == pds.is_monotonic_decreasing
 
 
 @pytest.mark.skipif(PANDAS_GE_200, reason="pandas removed is_monotonic")
@@ -5913,26 +5923,25 @@ def test_is_monotonic_dt64():
     s = pd.Series(pd.date_range("20130101", periods=10))
     ds = dd.from_pandas(s, npartitions=5)
     assert_eq(s.is_monotonic_increasing, ds.is_monotonic_increasing)
+    assert_eq(s.is_monotonic_decreasing, ds.is_monotonic_decreasing)
 
     s_2 = pd.Series(list(reversed(s)))
     ds_2 = dd.from_pandas(s_2, npartitions=5)
+    assert_eq(s_2.is_monotonic_increasing, ds_2.is_monotonic_increasing)
     assert_eq(s_2.is_monotonic_decreasing, ds_2.is_monotonic_decreasing)
 
 
-def test_index_is_monotonic_numeric():
-    s = pd.Series(1, index=range(20))
+def test_index_is_monotonic_dt64():
+    s = pd.Series(1, index=pd.date_range("20130101", periods=10))
     ds = dd.from_pandas(s, npartitions=5, sort=False)
     assert_eq(s.index.is_monotonic_increasing, ds.index.is_monotonic_increasing)
+    assert_eq(s.index.is_monotonic_decreasing, ds.index.is_monotonic_decreasing)
 
-    s_2 = pd.Series(1, index=range(20, 0, -1))
+    s_2 = s[::-1]
     ds_2 = dd.from_pandas(s_2, npartitions=5, sort=False)
+    assert_eq(s_2.index.is_monotonic_increasing, ds_2.index.is_monotonic_increasing)
     assert_eq(s_2.index.is_monotonic_decreasing, ds_2.index.is_monotonic_decreasing)
 
-    s_3 = pd.Series(1, index=list(range(0, 5)) + list(range(0, 20)))
-    ds_3 = dd.from_pandas(s_3, npartitions=5, sort=False)
-    assert_eq(s_3.index.is_monotonic_increasing, ds_3.index.is_monotonic_increasing)
-    assert_eq(s_3.index.is_monotonic_decreasing, ds_3.index.is_monotonic_decreasing)
-
 
 @pytest.mark.skipif(PANDAS_GE_200, reason="pandas removed is_monotonic")
 def test_index_is_monotonic_deprecated():
@@ -5950,16 +5959,6 @@ def test_index_is_monotonic_deprecated():
     assert_eq(expected, result)
 
 
-def test_index_is_monotonic_dt64():
-    s = pd.Series(1, index=pd.date_range("20130101", periods=20))
-    ds = dd.from_pandas(s, npartitions=10, sort=False)
-    assert_eq(s.index.is_monotonic_increasing, ds.index.is_monotonic_increasing)
-
-    s_2 = s[::-1]
-    ds_2 = dd.from_pandas(s_2, npartitions=10, sort=False)
-    assert_eq(s_2.index.is_monotonic_decreasing, ds_2.index.is_monotonic_decreasing)
-
-
 def test_is_monotonic_empty_partitions():
     df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [4, 3, 2, 1]})
     ddf = dd.from_pandas(df, npartitions=2)
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index f59a8f626..9dae29229 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import datetime
+import decimal
 import inspect
 import os
 import pathlib
@@ -251,6 +252,7 @@ def test_normalize_base():
         1.1,
         "1",
         slice(1, 2, 3),
+        decimal.Decimal("1.1"),
         datetime.date(2021, 6, 25),
         pathlib.PurePath("/this/that"),
     ]: