Repository: dask/dask
Pull request #10704 (2023-12-14 13:30:55)

Add changelog entry for ``dask-expr`` switch


Patch:
diff --git a/.github/labeler.yml b/.github/labeler.yml
deleted file mode 100644
index 5247cc239..000000000
--- a/.github/labeler.yml
+++ /dev/null
@@ -1,22 +0,0 @@
-dataframe:
-  - dask/dataframe/*
-  - dask/dataframe/**/*
-
-array:
-  - dask/array/*
-  - dask/array/**/*
-
-io:
-  - dask/dataframe/io/*
-  - dask/dataframe/io/**/*
-
-documentation:
-  - docs/*
-  - docs/**/*
-
-dispatch:
-  - dask/array/backends.py
-  - dask/array/dispatch.py
-  - dask/dataframe/backends.py
-  - dask/dataframe/dispatch.py
-  - dask/dataframe/extensions.py
diff --git a/.github/workflows/label-prs.yml b/.github/workflows/label-prs.yml
deleted file mode 100644
index 834bbcd32..000000000
--- a/.github/workflows/label-prs.yml
+++ /dev/null
@@ -1,12 +0,0 @@
-name: "PR Labeler"
-on:
-- pull_request_target
-
-jobs:
-  label:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/labeler@v5.0.0
-      with:
-        repo-token: "${{ secrets.GITHUB_TOKEN }}"
-        sync-labels: false
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 82d43f315..a46530fd2 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -1412,7 +1412,15 @@ def _split_partition(df, on, nsplits):
     # add a "_partitions" column to perform the split.
     if not isinstance(on, _Frame):
         on = _select_columns_or_index(df, on)
-    partitions = partitioning_index(on, nsplits)
+
+    dtypes = {}
+    for col, dtype in on.dtypes.items():
+        if pd.api.types.is_numeric_dtype(dtype):
+            dtypes[col] = np.float64
+    if not dtypes:
+        dtypes = None
+
+    partitions = partitioning_index(on, nsplits, cast_dtype=dtypes)
     df2 = df.assign(_partitions=partitions)
     return shuffle_group(
         df2,
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index d83e4d118..d99320828 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -420,6 +420,13 @@ def shuffle(
         # selection will not match (important when merging).
         index = index.to_frame()
 
+    dtypes = {}
+    for col, dtype in index.dtypes.items():
+        if pd.api.types.is_numeric_dtype(dtype):
+            dtypes[col] = np.float64
+    if not dtypes:
+        dtypes = None
+
     meta = df._meta._constructor_sliced([0])
     # Ensure that we have the same index as before to avoid alignment
     # when calculating meta dtypes later on
@@ -429,6 +436,7 @@ def shuffle(
         npartitions=npartitions or df.npartitions,
         meta=meta,
         transform_divisions=False,
+        cast_dtype=dtypes,
     )
     df2 = df.assign(_partitions=partitions)
     df2._meta.index.name = df._meta.index.name
@@ -786,7 +794,7 @@ def rearrange_by_column_tasks(
 ########################################################
 
 
-def partitioning_index(df, npartitions):
+def partitioning_index(df, npartitions, cast_dtype=None):
     """
     Computes a deterministic index mapping each record to a partition.
 
@@ -797,12 +805,18 @@ def partitioning_index(df, npartitions):
     df : DataFrame/Series/Index
     npartitions : int
         The number of partitions to group into.
+    cast_dtype : dtype, optional
+        The dtype to cast to to avoid nullability issues
 
     Returns
     -------
     partitions : ndarray
         An array of int64 values mapping each record to a partition.
     """
+    if cast_dtype is not None:
+        # Fixme: astype raises with strings in numeric columns, but raising
+        # here might be very noisy
+        df = df.astype(cast_dtype, errors="ignore")
     return hash_object_dispatch(df, index=False) % int(npartitions)
 
 
diff --git a/dask/delayed.py b/dask/delayed.py
index c599c5cb0..7b5439909 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -716,6 +716,10 @@ class DelayedLeaf(Delayed):
     def __doc__(self):
         return self._obj.__doc__
 
+    @property
+    def __wrapped__(self):
+        return self._obj
+
 
 class DelayedAttr(Delayed):
     __slots__ = ("_obj", "_attr")
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 1970ef812..297acdf5f 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -11,6 +11,34 @@ Released on December 15, 2023
 Highlights
 ^^^^^^^^^^
 
+Logical Query Planning now available for Dask DataFrames
+""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+Dask DataFrames are now much more performant by using a logical query planner.
+This feature is currently off by default, but can be turned on with:
+
+.. code:: python
+
+    dask.config.set({"dataframe.query-planning": True})
+
+You also need to have ``dask-expr`` installed:
+
+.. code:: bash
+
+    pip install dask-expr
+
+We've seen promising performance improvements so far, see
+`this blog post <https://blog.coiled.io/blog/dask-expr-tpch-dask.html>`__
+and `these regularly updated benchmarks <https://tpch.coiled.io>`__  for more information.
+A more detailed explanation of how the query optimizer works can be found in
+`this blog post <https://blog.coiled.io/blog/dask-expr-introduction.html>`__.
+
+This feature is still under active development
+and the `API <https://github.com/dask-contrib/dask-expr#api-coverage>`__ isn't stable yet,
+so breaking changes can occur. We expect to make the query optimizer the default early next year.
+
+See :pr:`10634` from `Patrick Hoefler`_ for details.
+
 Dtype inference in ``read_parquet``
 """""""""""""""""""""""""""""""""""
 
@@ -20,6 +48,8 @@ original Arrow array, and thus avoid the conversion to NumPy object. Additionall
 ``read_parquet`` will no longer infer nested and binary types as strings, they will
 be stored in NumPy object arrays.
 
+See :pr:`10698` and :pr:`10705` from `Patrick Hoefler`_ for details.
+
 
 .. _v2023.12.0:
 
@@ -37,6 +67,15 @@ PipInstall restart and environment variables
 The ``distributed.PipInstall`` plugin now has more robust restart logic and also supports
 `environment variables <https://pip.pypa.io/en/stable/reference/requirements-file-format/#using-environment-variables>`_.
 
+Below shows how users can use the ``distributed.PipInstall`` plugin and a ``TOKEN`` environment
+variable to securely install a package from a private repository:
+
+.. code:: python
+
+  from dask.distributed import PipInstall
+  plugin = PipInstall(packages=["private_package@git+https://${TOKEN}@github.com/dask/private_package.git])
+  client.register_plugin(plugin)
+
 See :pr-distributed:`8374`, :pr-distributed:`8357`, and :pr-distributed:`8343` from `Hendrik Makait`_ for details.
 
 


Test patch:
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index 4f08b063d..cdaa3a537 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1693,3 +1693,25 @@ def test_set_index_pyarrow_dtype(data, dtype):
     pdf_result = pdf.set_index("arrow_col")
     ddf_result = ddf.set_index("arrow_col")
     assert_eq(ddf_result, pdf_result)
+
+
+def test_shuffle_nulls_introduced():
+    df1 = pd.DataFrame([[True], [False]] * 50, columns=["A"])
+    df1["B"] = list(range(100))
+
+    df2 = pd.DataFrame(
+        [[2, 3], [109, 2], [345, 3], [50, 7], [95, 1]], columns=["B", "C"]
+    )
+
+    ddf1 = dd.from_pandas(df1, npartitions=10)
+    ddf2 = dd.from_pandas(df2, npartitions=1)
+    meta = pd.Series(dtype=int, index=pd.Index([], dtype=bool, name="A"), name="A")
+    result = (
+        dd.merge(ddf1, ddf2, how="outer", on="B")
+        .groupby("A")
+        .apply(lambda df: len(df), meta=meta)
+    )
+    expected = (
+        pd.merge(df1, df2, how="outer", on="B").groupby("A").apply(lambda df: len(df))
+    )
+    assert_eq(result, expected, check_names=False)
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index fbb8dc39f..dc9b796ad 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -855,3 +855,4 @@ def test_delayed_function_attributes_forwarded():
 
     assert add.__name__ == "add"
     assert add.__doc__ == "This is a docstring"
+    assert add.__wrapped__(1, 2) == 3