Repository: dask/dask
Pull request #11043 (2024-04-10 18:45:31)

Add docs for query optimizer
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  +     2       15 suites  +2   3h 22m 34s :stopwatch: + 40m 48s
 13 119 tests ±     0   12 173 :white_check_mark: ±     0     942 :zzz: ±    0  4 :x: ±0 
162 443 runs  +25 964  142 264 :white_check_mark: +22 887  20 175 :zzz: +3 077  4 :x: ±0 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/23676982500).

Results for commit 0bf2918b. ± Comparison against base commit b2ec1e1a.

[test-results]:data:application/gzip;base64,H4sIAP/lFmYC/02NzQ6DIBAGX8Vw7oFdFil9GYMKCak/Dcqp6bt3qxU9znyb2bcIcfCLeFSgb5VYclxP6nNya5ynHyNoYsXruu0KwB7cLLnr9iOjTvmML5aWsKjg4sCqhBqf0pzYSDYpT1u5RiL151ImxJqK3MsowejDXdIbX8vdPI5xZRCyDWjh3hI5bbxCLQ0BBGsUeuwt//BkDDrx+QK4tWa0GQEAAA==


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index a864b9236..74c11f76d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1892,9 +1892,12 @@ class Array(DaskMethodsMixin):
         if value is np.ma.masked:
             value = np.ma.masked_all((), dtype=self.dtype)
 
-        if not is_dask_collection(value) and np.isnan(value).any():
-            if issubclass(self.dtype.type, Integral):
-                raise ValueError("cannot convert float NaN to integer")
+        if (
+            not is_dask_collection(value)
+            and issubclass(self.dtype.type, Integral)
+            and np.isnan(value).any()
+        ):
+            raise ValueError("cannot convert float NaN to integer")
 
         ## Use the "where" method for cases when key is an Array
         if isinstance(key, Array):
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
index 08949b829..237274f64 100644
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -29,6 +29,7 @@ This will raise in a future version.
 """
         if use_dask_expr is None:
             warnings.warn(msg, FutureWarning)
+            return False
         else:
             raise ImportError(msg)
     return True
@@ -126,6 +127,7 @@ if _dask_expr_enabled():
             from_legacy_dataframe,
             from_map,
             from_pandas,
+            get_collection_type,
             get_dummies,
             isna,
             map_overlap,
diff --git a/docs/source/dataframe-api.rst b/docs/source/dataframe-api.rst
index bcd1cbbe1..d95650683 100644
--- a/docs/source/dataframe-api.rst
+++ b/docs/source/dataframe-api.rst
@@ -527,6 +527,18 @@ Convert DataFrames
    DataFrame.to_dask_array
    DataFrame.to_delayed
 
+
+Convert from/to legacy DataFrames
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: dask_expr._collection
+
+.. autosummary::
+   :toctree: generated/
+
+   DataFrame.to_legacy_dataframe
+   from_legacy_dataframe
+
 Reshape DataFrames
 ~~~~~~~~~~~~~~~~~~
 
@@ -590,6 +602,19 @@ Dask Metadata
 
    make_meta
 
+
+Query Planning and Optimization
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: dask_expr._collection
+
+.. autosummary::
+   :toctree: generated/
+
+   DataFrame.explain
+   DataFrame.visualize
+   DataFrame.analyze
+
 Other functions
 ~~~~~~~~~~~~~~~
 
diff --git a/docs/source/dataframe-best-practices.rst b/docs/source/dataframe-best-practices.rst
index 74c7b9af5..483b45760 100644
--- a/docs/source/dataframe-best-practices.rst
+++ b/docs/source/dataframe-best-practices.rst
@@ -75,9 +75,7 @@ but now you want to arrange it by time, all of your partitions will have to talk
 to each other to exchange shards of data.  This can be an intensive process,
 particularly on a cluster.
 
-So, definitely set the index but try do so infrequently.  After you set the
-index, you may want to ``persist`` your data if you are on a cluster
-(see `Persist Intelligently`_):
+So, definitely set the index but try do so infrequently.
 
 .. code-block:: python
 
@@ -99,6 +97,13 @@ Persist Intelligently
 
 .. note:: This section is only relevant to users on distributed systems.
 
+.. warning::
+
+    persist has a number of drawbacks with the query optimizer. It will block all
+    optimizations and prevent us from pushing column projections or filters into
+    the IO layer. Use persist sparingly only when absolutely necessary or you need
+    the full dataset afterwards.
+
 Often DataFrame workloads look like the following:
 
 1.  Load data from files
@@ -183,12 +188,6 @@ This helps to reduce overhead and increase the effectiveness of vectorized
 Pandas operations.  You should aim for partitions that have around 100MB of
 data each.
 
-Additionally, reducing partitions is very helpful just before shuffling, which
-creates ``n log(n)`` tasks relative to the number of partitions.  DataFrames
-with less than 100 partitions are much easier to shuffle than DataFrames with
-tens of thousands.
-
-
 Joins
 -----
 
diff --git a/docs/source/dataframe-design.rst b/docs/source/dataframe-design.rst
index 95046456e..179f34058 100644
--- a/docs/source/dataframe-design.rst
+++ b/docs/source/dataframe-design.rst
@@ -62,10 +62,10 @@ several options:
 
 2. A description of the appropriate names and dtypes.  This can take several forms:
 
-    * A ``dict`` of ``{name: dtype}`` or an iterable of ``(name, dtype)``
-      specifies a DataFrame. Note that order is important: the order of the
-      names in ``meta`` should match the order of the columns
-    * A tuple of ``(name, dtype)`` specifies a series
+* A ``dict`` of ``{name: dtype}`` or an iterable of ``(name, dtype)``
+  specifies a DataFrame. Note that order is important: the order of the
+  names in ``meta`` should match the order of the columns
+* A tuple of ``(name, dtype)`` specifies a series
 
 This keyword is available on all functions/methods that take user provided
 callables (e.g. ``DataFrame.map_partitions``, ``DataFrame.apply``, etc...), as
@@ -97,6 +97,10 @@ partitions and divisions of your DataFrame with the following fields:
    >>> df.divisions
    ['2015-01-01', '2015-02-01', '2015-03-01', '2015-04-01', '2015-04-31']
 
+The number of partitions and the division values might change during optimization.
+The optimizer will try to create partitions with a sensible size to avoid straining
+the scheduler with many small partitions.
+
 Divisions includes the minimum value of every partition's index and the maximum
 value of the last partition's index.  In the example above, if the user searches
 for a specific datetime range, then we know which partitions we need to inspect
@@ -125,15 +129,28 @@ calling ``df.set_index(...)``.
 Groupby
 -------
 
-By default, groupby methods return an object with only 1 partition. This is to
-optimize performance, and assumes the groupby reduction returns an object that
-is small enough to fit into memory. If your returned object is larger than this,
-you can increase the number of output partitions using the `split_out` argument.
+By default, groupby will choose the number of output partitions based on a few
+different factors. It will look at the number of grouping keys to guess the
+cardinality of your data. It will use this information to calculate a factor
+based on the number of input partitions. You can override this behavior by
+specifying the number of output partitions using the `split_out` argument.
 
 .. code-block:: python
 
    result = df.groupby('id').value.mean()
    result.npartitions  # returns 1
 
+   result = df.groupby(['id', 'id2']).value.mean()
+   result.npartitions  # returns 5
+
    result = df.groupby('id').value.mean(split_out=8)
    result.npartitions  # returns 8
+
+Some groupby aggregation functions have a different `split_out` default value.
+`split_out=True` will keep the number of partitions constant, which is useful
+for operations that either don't reduce the number of rows very much.
+
+.. code-block:: python
+
+   result = df.groupby('id').value.nunique()
+   result.npartitions  # returns same as df.npartitions
diff --git a/docs/source/dataframe-extend.rst b/docs/source/dataframe-extend.rst
index 77257bdf0..c2b46aa6f 100644
--- a/docs/source/dataframe-extend.rst
+++ b/docs/source/dataframe-extend.rst
@@ -29,7 +29,7 @@ Extend Dispatched Methods
 
 If you are going to pass around Pandas-like objects that are not normal Pandas
 objects, then we ask you to extend a few dispatched methods: ``make_meta``,
-``get_parallel_type``, and ``concat``.
+``get_collection_type``, and ``concat``.
 
 make_meta
 """""""""
@@ -100,27 +100,27 @@ dtypes, index name, and it should return a non-empty version:
        ...
 
 
-get_parallel_type
-"""""""""""""""""
+get_collection_type
+"""""""""""""""""""
 
 Given a non-Dask DataFrame object, return the Dask equivalent:
 
 .. code-block:: python
 
-   from dask.dataframe.core import get_parallel_type
+   from dask.dataframe import get_collection_type
 
-   @get_parallel_type.register(MyDataFrame)
-   def get_parallel_type_dataframe(df):
+   @get_collection_type.register(MyDataFrame)
+   def get_collection_type_dataframe(df):
        return MyDaskDataFrame
 
 
-   @get_parallel_type.register(MySeries)
-   def get_parallel_type_series(s):
+   @get_collection_type.register(MySeries)
+   def get_collection_type_series(s):
        return MyDaskSeries
 
 
-   @get_parallel_type.register(MyIndex)
-   def get_parallel_type_index(ind):
+   @get_collection_type.register(MyIndex)
+   def get_collection_type_index(ind):
        return MyDaskIndex
 
 
diff --git a/docs/source/dataframe-groupby.rst b/docs/source/dataframe-groupby.rst
index 03261b3ad..aa977cebe 100644
--- a/docs/source/dataframe-groupby.rst
+++ b/docs/source/dataframe-groupby.rst
@@ -114,14 +114,14 @@ or as a context manager:
         ddf.groupby(...).apply(...)
 
 
-In addition, ``set_index`` also accepts a ``shuffle`` keyword argument that
+In addition, ``set_index`` also accepts a ``shuffle_method`` keyword argument that
 can be used to select either on-disk or task-based shuffling:
 
 .. code-block:: python
 
-    ddf.set_index(column, shuffle='disk')
-    ddf.set_index(column, shuffle='tasks')
-    ddf.set_index(column, shuffle='p2p')
+    ddf.set_index(column, shuffle_method='disk')
+    ddf.set_index(column, shuffle_method='tasks')
+    ddf.set_index(column, shuffle_method='p2p')
 
 
 .. _dataframe.groupby.aggregate:
diff --git a/docs/source/dataframe-joins.rst b/docs/source/dataframe-joins.rst
index 56584c226..747eb945c 100644
--- a/docs/source/dataframe-joins.rst
+++ b/docs/source/dataframe-joins.rst
@@ -75,10 +75,6 @@ maintains that index, like Parquet.
     left.to_parquet("left", overwrite=True)
     left = dd.read_parquet("left")
 
-    # If the dataframe can fit in RAM, you can also use persist
-
-    # left = left.persist()
-
     right_one = dask.datasets.timeseries(dtypes={"bar": int})
     right_two = dask.datasets.timeseries(dtypes={"baz": int})
 
diff --git a/docs/source/dataframe-optimizer.rst b/docs/source/dataframe-optimizer.rst
new file mode 100644
index 000000000..df92c597f
--- /dev/null
+++ b/docs/source/dataframe-optimizer.rst
@@ -0,0 +1,88 @@
+Optimizer
+=========
+
+.. currentmodule:: dask.dataframe
+
+.. note::
+    Dask DataFrame supports Query Planning since version 2023.03.0
+
+Optimization steps
+------------------
+
+Dask DataFrame will run several optimizations before executing the computation. These
+computations are aimed towards improving the efficiency of the query.
+
+The optimizations entail the following steps (this list is not complete):
+
+- **Projection Pushdown:**
+
+  Select only the required columns in every step. This reduces the amount of data
+  that needs to be read from storage but also the amount of data that is processed
+  along the way. Columns are dropped at the earliest stage in the query.
+
+- **Filter Pushdown:**
+
+  Push filters down as far as possible, potentially into the IO step. Filters are
+  executed in the earliest stage in the query.
+
+- **Partition Pruning:**
+
+  Partition selections are pushed down as far as possible, potentially into the IO
+  step.
+
+- **Avoiding Shuffles:**
+
+  Dask DataFrame will try to avoid shuffling data between workers as much as possible.
+  This can be achieved if the column layout is already known, i.e. if the DataFrame
+  was shuffled on the same column before. For example, executing a
+  ``df.groupby(...).apply(...)`` after a merge operation will not shuffle the data
+  again if the groupby happens on the merge columns.
+
+- **Automatically resizing partitions:**
+
+  The IO layers automatically adjust the partition count based on the column subset
+  that is selected from the dataset. Very small partitions impact the scheduler and
+  expensive operations like shuffling negatively. This is addressed by adjusting
+  the partition count automatically.
+
+Exploring the optimized query
+-----------------------------
+
+Dask will call ``df.optimize()`` before executing the computation. This method applies
+the steps mentioned above and returns a new Dask DataFrame that represents the optimized
+query.
+
+A rudimentary representation of the optimized query can be obtained by calling
+``df.pprint()``. This will print the query plan in a human-readable format to the
+command line/console. The advantage of this method is that it doesn't require any
+additional dependencies.
+
+.. code-block::
+
+    pdf = pd.DataFrame({"a": [1, 2, 3] * 5, "b": [1, 2, 3] * 5})
+    df = dd.from_pandas(pdf, npartitions=2)
+    df = df.replace(1, 5)[["a"]]
+
+    df.optimize().pprint()
+
+    Replace: to_replace=1 value=5
+        FromPandas: frame='<pandas>' npartitions=2 columns=['a'] pyarrow_strings_enabled=True
+
+A more advanced and easier to read reprepresentation can be obtained by calling
+``df.explain()``. This method requires the ``graphviz`` package to be installed.
+The method will return a graph that represents the query plan and create an image
+from it.
+
+.. code-block::
+
+    df.explain()
+
+.. figure:: images/dataframe-optimize-explain.png
+   :alt: Optimized Query
+   :align: center
+
+We can see in both representations that ``FromPandas`` consumed the column
+projection, only selecting the column ``a``.
+
+The ``explain()`` method is significantly easier to understand for more complex
+queries.
diff --git a/docs/source/dataframe-parquet.rst b/docs/source/dataframe-parquet.rst
index a271d5391..5b8646c94 100644
--- a/docs/source/dataframe-parquet.rst
+++ b/docs/source/dataframe-parquet.rst
@@ -107,7 +107,7 @@ files need to be divided into multiple row-group ranges, and the dataset
 does not contain a ``_metadata`` file, Dask will need to load all footer
 metadata up-front.
 
-We recommend aiming for 10-250 MiB in-memory size per file once loaded into
+We recommend aiming for 100-300 MiB in-memory size per file once loaded into
 pandas. Oversized partitions can lead to excessive memory usage on a single
 worker, while undersized partitions can lead to poor performance as the
 overhead of Dask dominates.
@@ -204,7 +204,7 @@ For more information on connecting to remote data, see
 
 Dask will write one file per Dask dataframe partition to this directory. To
 optimize access for downstream consumers, we recommend aiming for an in-memory
-size of 10-250 MiB per partition. This helps balance worker memory usage
+size of 100-300 MiB per partition. This helps balance worker memory usage
 against Dask overhead. You may find the
 :meth:`DataFrame.memory_usage_per_partition` method useful for determining if
 your data is partitioned optimally.
diff --git a/docs/source/dataframe-sql.rst b/docs/source/dataframe-sql.rst
index bc2a88872..4e7d9d679 100644
--- a/docs/source/dataframe-sql.rst
+++ b/docs/source/dataframe-sql.rst
@@ -33,14 +33,9 @@ The following packages may be of interest:
 - `FugueSQL`_ provides a unified interface to run SQL code on a variety of different computing frameworks.
   Specifying ``DaskExecutionEngine`` or ``DaskSQLExecutionEngine`` as the execution engine for queries allows them to be computed using Dask or dask-sql, respectively.
 
-- `pandasql`_ allows executing SQL queries on a pandas table by writing the data to
-  ``SQLite``, which may be useful for small toy examples (this package has not been
-  maintained for some time).
-
 .. _dask-sql: https://dask-sql.readthedocs.io/en/latest/
 .. _cuDF: https://docs.rapids.ai/api/cudf/stable/
 .. _FugueSQL: https://fugue-tutorials.readthedocs.io/en/latest/tutorials/fugue_sql/index.html
-.. _pandasql: https://github.com/yhat/pandasql/
 
 Database or Dask?
 -----------------
diff --git a/docs/source/dataframe.rst b/docs/source/dataframe.rst
index 7d9098d1d..96a03f077 100644
--- a/docs/source/dataframe.rst
+++ b/docs/source/dataframe.rst
@@ -10,6 +10,7 @@ Dask DataFrame
 
    Load and Save Data <dataframe-create.rst>
    Internal Design <dataframe-design.rst>
+   Optimizer <dataframe-optimizer.rst>
    Best Practices <dataframe-best-practices.rst>
    API <dataframe-api.rst>
    dataframe-extra.rst
diff --git a/docs/source/how-to/selecting-the-collection-backend.rst b/docs/source/how-to/selecting-the-collection-backend.rst
index 32b5b1485..be9eed924 100644
--- a/docs/source/how-to/selecting-the-collection-backend.rst
+++ b/docs/source/how-to/selecting-the-collection-backend.rst
@@ -133,7 +133,7 @@ Dask-Dataframe compute-based dispatch functions (as defined in ``dask.dataframe.
 
    - categorical_dtype_dispatch
    - concat_dispatch
-   - get_parallel_type
+   - get_collection_type
    - group_split_dispatch
    - grouper_dispatch
    - hash_object_dispatch


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index f4b190f2f..2e71b0a07 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -3911,6 +3911,17 @@ def test_setitem_1d():
         dx[index] = 1
 
 
+def test_setitem_masked():
+    # Test np.ma.masked assignment to object-type arrays
+    x = np.ma.array(["a", 1, 3.14], dtype=object)
+    dx = da.from_array(x.copy(), chunks=2)
+
+    x[...] = np.ma.masked
+    dx[...] = np.ma.masked
+
+    assert_eq(x.mask, da.ma.getmaskarray(dx))
+
+
 def test_setitem_hardmask():
     x = np.ma.array([1, 2, 3, 4], dtype=int)
     x.harden_mask()