Repository: dask/dask

Issue:
API for filtering dataframes by partition
I just came across a need for a way to filter a Dask DataFrame based on some partition-level characteristics. In this case it was a Dask Geopandas dataframe and I was checking whether the partition's spatial footprint overlaps with some shape (dask-geopandas adds a `spatial_extent` attribute to each Dask GeoDataFrame). In general, this feature would be helpful whenever you have a boolean mask *on the partitions*.

The solution was to use a mixture of `dd.concat` and `dd.get_partition`


Here's a simple example:


```python
>>> import dask.dataframe as dd
>>> df = dask.dataframe.demo.make_timeseries()
>>> mask = df.id.map_partitions(lambda p: len(p.unique()) > 260).compute()
>>> keep = dd.concat([df.get_partition(i) for i, m in enumerate(mask) if m])
>>> keep
Dask DataFrame Structure:
                 name     id        x        y
npartitions=6                                 
               object  int64  float64  float64
                  ...    ...      ...      ...
...               ...    ...      ...      ...
                  ...    ...      ...      ...
                  ...    ...      ...      ...
Dask Name: concat, 8 graph layers
```

(Note that that has 6 partitions, instead of the normal 11).

It's not too difficult of a workaround, but I think a moderate expansion API would be acceptable. It doesn't quite fit `get_partition`, since it's getting multiple partitions. Perhaps a `.filter_partitions` that takes a boolean sequence the same length as the number of partitions? Or a callable that's provided each partition and should return a boolean?

I think this will become especially useful after https://github.com/dask/dask/pull/9473, since users will have more partition-level properties to filter on *without* a compute (cc @rjzamora).


Discussion:
Don't forget the partitions indexer

On Mon, Feb 13, 2023, 3:45 PM Tom Augspurger ***@***.***>
wrote:

> I just came across a need for a way to filter a Dask DataFrame based on
> some partition-level characteristics. In this case it was a Dask Geopandas
> dataframe and I was checking whether the partition's spatial footprint
> overlaps with some shape (dask-geopandas adds a spatial_extent attribute
> to each Dask GeoDataFrame). In general, this feature would be helpful
> whenever you have a boolean mask *on the partitions*.
>
> The solution was to use a mixture of dd.concat and dd.get_partition
>
> Here's a simple example:
>
> >>> import dask>>> df = dask.dataframe.demo.make_timeseries()>>> mask = df.id.map_partitions(lambda p: len(p.unique()) > 260).compute()>>> keep = dd.concat([df.get_partition(i) for i, m in enumerate(mask) if m])>>> keepDask DataFrame Structure:
>                  name     id        x        ynpartitions=6
>                object  int64  float64  float64
>                   ...    ...      ...      ...
> ...               ...    ...      ...      ...
>                   ...    ...      ...      ...
>                   ...    ...      ...      ...Dask Name: concat, 8 graph layers
>
> (Note that that has 6 partitions, instead of the normal 11).
>
> It's not too difficult of a workaround, but I think a moderate expansion
> API would be acceptable. It doesn't quite fit get_partition, since it's
> getting multiple partitions. Perhaps a .filter_partitions that takes a
> boolean sequence the same length as the number of partitions? Or a callable
> that's provided each partition and should return a boolean?
>
> I think this will become especially useful after #9473
> <https://github.com/dask/dask/pull/9473>, since users will have more
> partition-level properties to filter on *without* a compute (cc @rjzamora
> <https://github.com/rjzamora>).
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/9951>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTGFTO75TVCPVEYTOALWXKTO3ANCNFSM6AAAAAAU2ZIRFU>
> .
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>

@TomAugspurger thoughts on using the existing `df.partitions` indexer? 

```python
In [1]: import dask.dataframe

In [2]: df = dask.dataframe.demo.make_timeseries()

In [3]: mask = df.id.map_partitions(lambda p: len(p.unique()) > 260).compute()

In [4]: df.npartitions
Out[4]: 11

In [5]: mask
Out[5]:
0      True
1      True
2      True
3      True
4     False
5      True
6      True
7      True
8      True
9      True
10     True
dtype: bool

In [6]: df.partitions[mask]
Out[6]:
Dask DataFrame Structure:
                  name     id        x        y
npartitions=10
2000-01-31      object  int64  float64  float64
2000-02-29         ...    ...      ...      ...
...                ...    ...      ...      ...
2000-11-30         ...    ...      ...      ...
2000-12-31         ...    ...      ...      ...
Dask Name: blocks, 2 graph layers
```

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -675,7 +675,35 @@ def compute_current_divisions(self, col=None):
         return compute_divisions(self, col=col)
 
     def get_partition(self, n):
-        """Get a dask DataFrame/Series representing the `nth` partition."""
+        """
+        Get a dask DataFrame/Series representing the `nth` partition.
+
+        Parameters
+        ----------
+        n : int
+            The 0-indexed partition number to select.
+
+        Returns
+        -------
+        Dask DataFrame or Series
+            The same type as the original object.
+
+        Examples
+        --------
+        >>> import dask
+        >>> ddf = dask.datasets.timeseries(start="2021-01-01", end="2021-01-07", freq="1H")
+        >>> ddf.get_partition(0)  # doctest: +NORMALIZE_WHITESPACE
+        Dask DataFrame Structure:
+                         name     id        x        y
+        npartitions=1
+        2021-01-01     object  int64  float64  float64
+        2021-01-02        ...    ...      ...      ...
+        Dask Name: get-partition, 2 graph layers
+
+        See Also
+        --------
+        DataFrame.partitions
+        """
         if 0 <= n < self.npartitions:
             name = f"get-partition-{str(n)}-{self._name}"
             divisions = self.divisions[n : n + 2]
@@ -1356,9 +1384,10 @@ def partitions(self):
         """Slice dataframe by partitions
 
         This allows partitionwise slicing of a Dask Dataframe.  You can perform normal
-        Numpy-style slicing but now rather than slice elements of the array you
+        Numpy-style slicing, but now rather than slice elements of the array you
         slice along partitions so, for example, ``df.partitions[:5]`` produces a new
-        Dask Dataframe of the first five partitions.
+        Dask Dataframe of the first five partitions. Valid indexers are integers, sequences
+        of integers, slices, or boolean masks.
 
         Examples
         --------


Test patch: