Repository: dask/dask
Pull request #11875 (2025-04-08 17:04:25)

Enable column projection in ``MapPartitions``
This feature was a requested at the last [Dask Community Meeting](https://docs.google.com/document/d/1UqNAP87a56ERH_xkQsS5Q_0PKYybd5Lj2WANy_hRzI0/edit?tab=t.0). Adds a `required_columns` argument to `map_partitions`. If this argument is specified, column projections are no-longer blocked by `MapPartitions` expressions. 


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._

      9 files  ± 0        9 suites  ±0   3h 25m 38s ⏱️ - 3m 35s
 18 011 tests + 2   16 797 ✅ + 2   1 214 💤 ±0  0 ❌ ±0 
161 173 runs  +18  149 068 ✅ +17  12 105 💤 +1  0 ❌ ±0 

Results for commit 5138b1c4. ± Comparison against base commit 0fa5e18d.

[test-results]:data:application/gzip;base64,H4sIAKFg9WcC/03MSQ7DIAyF4atErLvAjKaXqSiDhJqhIrCqeveSKKFZ/p/l9yExjWEl98HcBrLWVHr4mm1Jy9wSGOfYqB3LdgakAGc/1urchkob/cdXeu+fILpFm8ZmtEPIecmH5Drv0wpA86P7tDBUYcc+TeVpl+m9r8tumaZUWhAJHJ/gBAfmgMkQPTKK3lpgInItA0YjlDLk+wPeHbZ5GAEAAA==

rjzamora: >The part I didn't understand is why the user couldn't do `df[["z"]].map_partitions(func)`

The user can certainly do that if they know that they don't need any columns besides `"z"` within `func`, or in any operation after the `map_partitions` call.

I think the best way to understand the purpose if this feature is to look at `test_column_projection_map_partitions`. In that test, the user-defined function requires the `"z"` column, but the next operation is to select a completely different column (`"x"`). If we were to do `df[["z"]].map_partitions(func)`, we would get an error, because `"x"` would no-longer be in the schema.
TomAugspurger: So I think for that test example, where the `func` requires a column `z` and the final selection is for `x`,  the two different APIs would be:

```
df.map_partitions(func, required_columns=["z"])[["x"]]
```

vs.

```
df[["x", "z"]].map_partitions(func)[["x"]]
```

And the ultimate goal is that we get a `read_parquet(columns=["x", "z"])`. Does that sound right?

Not duplicating the column name `x` is nice. Decoupling the columns needed for the function from the other columns is nice (you could imagine generating the arguments for `map_partitions` dynamically). I'm still not quite convinced it's worth it but am happy to defer to you / others on this.
rjzamora: >I'm still not quite convinced it's worth it but am happy to defer to you / others on this.

I think you need to think of the problem from the perspective of someone who is using `map_partitions` to define their own function for transforming a Dask-DataFrame (e.g. `cool_func(df: dd.DataFrame, *args, **kwargs) -> dd.DataFrame`).

The author of `cool_func` may be different from the user, and the author may want the user to be able to select specific columns **after** `cool_func` is called. The user could indeed select the necessary columns both before and after they call `cool_func` to get optimal behavior. However, the only way to make sure that optimal column-projection is applied is if the author can specify something like `necessary_columns=` when they call `map_partitions` within their API.
rjzamora: >Only additional question / request would be to confirm the behavior when the user passes a required_columns that isn't in the dataframe. Do we raise a KeyError at graph construction time? Or execution time?

Ah, good question. I can add a block to raise a KeyError if any columns in `required_columns` are missing. Right now, we simply ignore projected columns that aren't in `frame`, because we are hooking into logic that doesn't know if a projected column is being added by the *current* expression (but we **do** know the `required_columns` subset must be in `frame` in this specific case).

Patch:
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index c4dbd12d8..25f0524b0 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -994,6 +994,7 @@ Expr={expr}"""
         clear_divisions=False,
         align_dataframes=False,
         parent_meta=None,
+        required_columns=None,
         **kwargs,
     ):
         """Apply a Python function to each partition
@@ -1021,6 +1022,11 @@ Expr={expr}"""
         clear_divisions : bool, default False
             Whether divisions should be cleared. If True, `transform_divisions`
             will be ignored.
+        required_columns : list or None, default None
+            List of columns that ``func`` requires for execution. These columns
+            must belong to the first DataFrame argument (in ``args``). If None
+            is specified (the default), the query optimizer will assume that
+            all input columns are required.
         $META
 
         Examples
@@ -1118,6 +1124,7 @@ Expr={expr}"""
             clear_divisions=clear_divisions,
             align_dataframes=align_dataframes,
             parent_meta=parent_meta,
+            required_columns=required_columns,
             **kwargs,
         )
 
@@ -6126,6 +6133,7 @@ def map_partitions(
     clear_divisions=False,
     align_dataframes=False,
     parent_meta=None,
+    required_columns=None,
     **kwargs,
 ):
     """Apply Python function on each DataFrame partition.
@@ -6158,6 +6166,11 @@ def map_partitions(
         If False, all inputs must have either the same number of partitions
         or a single partition. Single-partition inputs will be broadcast to
         every partition of multi-partition inputs.
+    required_columns : list or None, default None
+        List of columns that ``func`` requires for execution. These columns
+        must belong to the first DataFrame argument (in ``args``). If None
+        is specified (the default), the query optimizer will assume that
+        all input columns are required.
     $META
     """
     if align_dataframes:
@@ -6189,6 +6202,7 @@ def map_partitions(
         clear_divisions,
         align_dataframes,
         parent_meta,
+        required_columns,
         newkwargs.pop("token", None),
         Dict(newkwargs),
         len(args) - 1,
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index da88412bc..6b19514d3 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -8,6 +8,7 @@ import warnings
 import weakref
 from collections import defaultdict
 from collections.abc import Callable, Collection, Mapping
+from functools import partial
 from typing import Any as AnyType
 
 import numpy as np
@@ -654,6 +655,7 @@ class MapPartitions(Blockwise):
         "clear_divisions",
         "align_dataframes",
         "parent_meta",
+        "required_columns",
         "token",
         "kwargs",
         "nargs",
@@ -662,6 +664,7 @@ class MapPartitions(Blockwise):
         "kwargs": None,
         "align_dataframes": True,
         "parent_meta": None,
+        "required_columns": None,
         "token": None,
         "nargs": 0,
     }
@@ -756,6 +759,38 @@ class MapPartitions(Blockwise):
                 **kwargs,
             )
 
+    @staticmethod
+    def projected_operation(mapped_func, post_projection, *args, **kwargs):
+        # Apply a mapped function and then project columns.
+        # Used by `_simplify_up` to apply column projection.
+        return mapped_func(*args, **kwargs)[post_projection]
+
+    def _simplify_up(self, parent, dependents):
+        if isinstance(parent, Projection) and self.required_columns is not None:
+            if missing := set(self.required_columns) - set(self.frame.columns):
+                raise KeyError(
+                    f"Some elements of `required_columns` are missing: {missing}"
+                )
+
+            columns = determine_column_projection(
+                self, parent, dependents, additional_columns=self.required_columns
+            )
+            columns = [col for col in self.frame.columns if col in columns]
+
+            if columns == self.frame.columns:
+                # Don't add unnecessary Projections
+                return
+
+            return type(parent)(
+                type(self)(
+                    self.frame[columns],
+                    partial(self.projected_operation, self.func, parent.columns),
+                    self.meta[parent.columns],
+                    *self.operands[3:],
+                ),
+                *parent.operands[1:],
+            )
+
 
 def _get_meta_ufunc(dfs, args, func):
     dasks = [arg for arg in args if isinstance(arg, (Expr, Array))]
@@ -984,6 +1019,7 @@ class MapOverlap(MapPartitions):
             self.clear_divisions,
             self.align_dataframes,
             None,
+            None,
             self.token,
             self._kwargs,
             len(self.args[1:]),
diff --git a/dask/dataframe/dask_expr/_groupby.py b/dask/dataframe/dask_expr/_groupby.py
index d554f2411..6e42df80c 100644
--- a/dask/dataframe/dask_expr/_groupby.py
+++ b/dask/dataframe/dask_expr/_groupby.py
@@ -1390,6 +1390,7 @@ class GroupByCumulative(Expr, GroupByBase):
             True,
             None,
             None,
+            None,
             {"chunk": self.chunk, "columns": columns, **dropna, **self.numeric_only},
             len(self.by),
             *self.by,
@@ -1420,6 +1421,7 @@ class GroupByCumulative(Expr, GroupByBase):
             True,
             None,
             None,
+            None,
             {"chunk": M.last, "columns": columns, **dropna},
             len(by),
             *by,
diff --git a/dask/dataframe/dask_expr/_merge_asof.py b/dask/dataframe/dask_expr/_merge_asof.py
index eb9684cc6..852e41d8c 100644
--- a/dask/dataframe/dask_expr/_merge_asof.py
+++ b/dask/dataframe/dask_expr/_merge_asof.py
@@ -131,6 +131,7 @@ class MergeAsof(Merge):
                 True,
                 None,
                 None,
+                None,
                 self._kwargs,
                 1,
                 self.right,
@@ -154,6 +155,7 @@ class MergeAsof(Merge):
                 True,
                 None,
                 None,
+                None,
                 {"left_index": True, "right_index": True},
                 1,
                 right,


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index 1f1754199..845302072 100644
--- a/dask/dataframe/dask_expr/tests/test_collection.py
+++ b/dask/dataframe/dask_expr/tests/test_collection.py
@@ -143,6 +143,25 @@ def test_column_projection_modify_list(df, pdf):
     assert_eq(result, pdf[["x"]])
 
 
+@pytest.mark.parametrize("required_columns", [None, ["z"], ["z", "foo"]])
+def test_column_projection_map_partitions(required_columns):
+
+    pdf = pd.DataFrame({"x": [1] * 5, "y": [2] * 5, "z": range(5)})
+    df = from_pandas(pdf, npartitions=2)
+
+    def myfunc(x):
+        return x.assign(zz=x["z"] + 1)
+
+    df = df.map_partitions(myfunc, meta=myfunc(pdf), required_columns=required_columns)
+    if required_columns and "foo" in required_columns:
+        with pytest.raises(KeyError, match="foo"):
+            df["x"].simplify()
+    else:
+        assert_eq(df["x"], myfunc(pdf)["x"])
+        if required_columns:
+            assert set(df["x"].simplify().frame.frame.columns) == {"x", "z"}
+
+
 def test_setitem(pdf, df):
     pdf = pdf.copy()
     pdf["z"] = pdf.x + pdf.y