Repository: dask/dask

Issue:
Add magic partition_info keyword to dd.map_partitions function
In dask.array we allow user provided functions to get a special keyword, `block_info` with information about the array and where the block that they're operating in that array lives:

https://github.com/dask/dask/blob/5b2ecf4a5fbc38b7b1a013abd45299b8d5a235e0/dask/array/core.py#L656-L674

We might consider doing the same for `dask.dataframe.map_partitions`.  We have a bit less information to provide, but can probably include:

1.  The number of partitions in the dataframe
2.  Which partition we're on specifically

This requires using the `dask.base.has_keyword` function to inspect the user provided function, and then manipulating the task graph a bit to add the extra keyword argument.

Inspired by this stack overflow question: https://stackoverflow.com/questions/51135739/how-to-modify-values-on-specific-dask-dataframe-partiton


Discussion:
dd.read_csv(*.csv)  makes one partition one csv file.  

Most of time, csv files are raw data, when csv files have been  filled into the partitions, we need to do some data cleanning works base on  partition.

If could add  a label  for each  partition number  when call read_csv(*.csv),that would be  helpful  to judge which  csv file the partition number refers to.

@isveryok Adding this to read_csv sounds complicated. I see it interfering with things like user-provided columns. Do you have an example that would benefit from it?

Agreed that `partition_info` useful for map_partitions.
I made a start on this with https://github.com/dask/dask/compare/master...KrishanBhasin:feature/partition-info as a way of learning more about how Dask works.
If anyone has any pointers or tips on how to approach I'd apprecaite it; equally if anyone is able to beat me to a working PR please feel free.
You're more likely to get feedback as a pull request since those are easier to comment on.

`dask.array.map_blocks` does something similar. I don't know whether it's a good model to follow though (it may be more complicated than we need).
Thanks, I'll do that - I'm keen for feedback/pointers but I was worried about putting in a half-baked draft PR that I may not get around to finishing.

Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1,3 +1,4 @@
+import copy
 import operator
 import warnings
 from collections.abc import Iterator, Sequence
@@ -29,6 +30,7 @@
 from .. import threaded
 from ..context import globalmethod
 from ..utils import (
+    has_keyword,
     random_state_data,
     pseudorandom,
     derived_from,
@@ -44,7 +46,7 @@
 )
 from ..array.core import Array, normalize_arg
 from ..array.utils import zeros_like_safe
-from ..blockwise import blockwise, Blockwise
+from ..blockwise import blockwise, Blockwise, subs
 from ..base import DaskMethodsMixin, tokenize, dont_optimize, is_dask_collection
 from ..delayed import delayed, Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
@@ -582,9 +584,10 @@ def map_partitions(self, func, *args, **kwargs):
         args, kwargs :
             Arguments and keywords to pass to the function. The partition will
             be the first argument, and these will be passed *after*. Arguments
-            and keywords may contain ``Scalar``, ``Delayed`` or regular
-            python objects. DataFrame-like args (both dask and pandas) will be
-            repartitioned to align (if necessary) before applying the function.
+            and keywords may contain ``Scalar``, ``Delayed``, ``partition_info``
+            or regular python objects. DataFrame-like args (both dask and
+            pandas) will be repartitioned to align (if necessary) before
+            applying the function.
         $META
 
         Examples
@@ -648,6 +651,23 @@ def map_partitions(self, func, *args, **kwargs):
         to clear them afterwards:
 
         >>> ddf.map_partitions(func).clear_divisions()  # doctest: +SKIP
+
+        Your map function gets information about where it is in the dataframe by
+        accepting a special ``partition_info`` keyword argument.
+
+        >>> def func(partition, partition_info=None):
+        ...     pass
+
+        This will receive the following information:
+
+        >>> partition_info  # doctest: +SKIP
+        {'number': 1, 'division': 3}
+
+        For each argument and keyword arguments that are dask dataframes you will
+        receive the number (n) which represents the nth partition of the dataframe
+        and the division (the first index value in the partition). If divisions
+        are not known (for instance if the index is not sorted) then you will get
+        None as the division.
         """
         return map_partitions(func, self, *args, **kwargs)
 
@@ -5319,6 +5339,9 @@ def map_partitions(
     """
     name = kwargs.pop("token", None)
 
+    if has_keyword(func, "partition_info"):
+        kwargs["partition_info"] = {"number": -1, "divisions": None}
+
     assert callable(func)
     if name is not None:
         token = tokenize(meta, *args, **kwargs)
@@ -5341,6 +5364,9 @@ def map_partitions(
     else:
         meta = make_meta(meta, index=meta_index)
 
+    if has_keyword(func, "partition_info"):
+        kwargs["partition_info"] = "__dummy__"
+
     if all(isinstance(arg, Scalar) for arg in args):
         layer = {
             (name, 0): (apply, func, (tuple, [(arg._name, 0) for arg in args]), kwargs)
@@ -5410,6 +5436,22 @@ def map_partitions(
             if not valid_divisions(divisions):
                 divisions = [None] * (dfs[0].npartitions + 1)
 
+    if has_keyword(func, "partition_info"):
+        dsk = dict(dsk)
+
+        for k, v in dsk.items():
+            vv = v
+            v = v[0]
+            number = k[-1]
+            assert isinstance(number, int)
+            info = {"number": number, "division": divisions[number]}
+            v = copy.copy(v)  # Need to copy and unpack subgraph callable
+            v.dsk = copy.copy(v.dsk)
+            [(key, task)] = v.dsk.items()
+            task = subs(task, {"__dummy__": info})
+            v.dsk[key] = task
+            dsk[k] = (v,) + vv[1:]
+
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
     return new_dd_object(graph, name, meta, divisions)
 


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -858,6 +858,19 @@ def test_map_partitions_type():
     assert all(x == pd.DataFrame for x in result)
 
 
+def test_map_partitions_partition_info():
+    def f(df, partition_info=None):
+        assert partition_info is not None
+        assert "number" in partition_info
+        assert "division" in partition_info
+        assert dsk[("x", partition_info["number"])].equals(df)
+        assert dsk[("x", d.divisions.index(partition_info["division"]))].equals(df)
+        return df
+
+    result = d.map_partitions(f, meta=d).compute(scheduler="single-threaded")
+    assert type(result) == pd.DataFrame
+
+
 def test_map_partitions_names():
     func = lambda x: x
     assert sorted(dd.map_partitions(func, d, meta=d).dask) == sorted(