Repository: dask/dask

Issue:
Add first and last aggregate functions to `dask.dataframe.pivot_table`
This question/feature request started on the [Dask Discourse](https://dask.discourse.group/t/using-pivot-table-with-non-numerical-data/267/2).

Currently `dask.dataframe.pivot_table` only supports `mean`, `sum` and `count` as aggregate functions. This makes `pivot_table` inconvenient for working with non-numeric data. `first` and `last` are supported aggregate functions in Pandas' `pivot_table`.

Typical data would look like:
```
   ID Col_ID    value
0   1      A    Alice
1   1      B      Bob
2   1      A  Charlie
```
Desired output in case of `aggfunc='first'`
```
Col_ID      A    B
ID
1       Alice  Bob
```
Desired output in case of `aggfunc='last'`
```
Col_ID        A    B
ID
1       Charlie  Bob
```

Example code
```
import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame({"ID": [1, 1, 1], "Col_ID": ["A", "B", "A"], "value": ["Alice", "Bob", "Charlie"]})
df_pivot = df.pivot_table(index="ID", columns="Col_ID", values="value", aggfunc="first")

ddf = dd.from_pandas(df, npartitions=2)
# count is the only aggregate functions that could be useful for non-numeric data
ddf_pivot = ddf.categorize(columns=["Col_ID"]).pivot_table(
    index="ID", columns="Col_ID", values="value", aggfunc="count"
)

```




Discussion:
Thanks for raising this issue @nordange! It seems like adding these aggregations to pivot_table would be relatively straightforward. It would just require changes here: 

https://github.com/dask/dask/blob/e44f78009d56332b0153828adb0eaf0e5d01582f/dask/dataframe/reshape.py#L186

and new `pivot_last` and `pivot_first` functions here:

https://github.com/dask/dask/blob/e44f78009d56332b0153828adb0eaf0e5d01582f/dask/dataframe/methods.py#L417

Great! I'll give it a go, @jsignell. 

Resolving patch:
diff --git a/dask/dataframe/methods.py b/dask/dataframe/methods.py
--- a/dask/dataframe/methods.py
+++ b/dask/dataframe/methods.py
@@ -402,6 +402,14 @@ def pivot_agg(df):
     return df.groupby(level=0).sum()
 
 
+def pivot_agg_first(df):
+    return df.groupby(level=0).first()
+
+
+def pivot_agg_last(df):
+    return df.groupby(level=0).last()
+
+
 def pivot_sum(df, index, columns, values):
     return pd.pivot_table(
         df, index=index, columns=columns, values=values, aggfunc="sum", dropna=False
@@ -416,6 +424,18 @@ def pivot_count(df, index, columns, values):
     ).astype(np.float64)
 
 
+def pivot_first(df, index, columns, values):
+    return pd.pivot_table(
+        df, index=index, columns=columns, values=values, aggfunc="first", dropna=False
+    )
+
+
+def pivot_last(df, index, columns, values):
+    return pd.pivot_table(
+        df, index=index, columns=columns, values=values, aggfunc="last", dropna=False
+    )
+
+
 def assign_index(df, ind):
     df = df.copy()
     df.index = ind
diff --git a/dask/dataframe/reshape.py b/dask/dataframe/reshape.py
--- a/dask/dataframe/reshape.py
+++ b/dask/dataframe/reshape.py
@@ -199,7 +199,7 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
         column to be columns
     values : scalar or list(scalar)
         column(s) to aggregate
-    aggfunc : {'mean', 'sum', 'count'}, default 'mean'
+    aggfunc : {'mean', 'sum', 'count', 'first', 'last'}, default 'mean'
 
     Returns
     -------
@@ -228,8 +228,13 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
         or is_scalar(values)
     ):
         raise ValueError("'values' must refer to an existing column or columns")
-    if not is_scalar(aggfunc) or aggfunc not in ("mean", "sum", "count"):
-        raise ValueError("aggfunc must be either 'mean', 'sum' or 'count'")
+
+    available_aggfuncs = ["mean", "sum", "count", "first", "last"]
+
+    if not is_scalar(aggfunc) or aggfunc not in available_aggfuncs:
+        raise ValueError(
+            "aggfunc must be either " + ", ".join(f"'{x}'" for x in available_aggfuncs)
+        )
 
     # _emulate can't work for empty data
     # the result must have CategoricalIndex columns
@@ -242,9 +247,26 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
             (sorted(values), columns_contents), names=[None, columns]
         )
 
-    meta = pd.DataFrame(
-        columns=new_columns, dtype=np.float64, index=pd.Index(df._meta[index])
-    )
+    if aggfunc in ["first", "last"]:
+        # Infer datatype as non-numeric values are allowed
+        if is_scalar(values):
+            meta = pd.DataFrame(
+                columns=new_columns,
+                dtype=df[values].dtype,
+                index=pd.Index(df._meta[index]),
+            )
+        else:
+            meta = pd.DataFrame(
+                columns=new_columns,
+                index=pd.Index(df._meta[index]),
+            )
+            for value_col in values:
+                meta[value_col] = meta[value_col].astype(df[values].dtypes[value_col])
+    else:
+        # Use float64 as other aggregate functions require numerical data
+        meta = pd.DataFrame(
+            columns=new_columns, dtype=np.float64, index=pd.Index(df._meta[index])
+        )
 
     kwargs = {"index": index, "columns": columns, "values": values}
 
@@ -274,6 +296,24 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
         return pv_count
     elif aggfunc == "mean":
         return pv_sum / pv_count
+    elif aggfunc == "first":
+        return apply_concat_apply(
+            [df],
+            chunk=methods.pivot_first,
+            aggregate=methods.pivot_agg_first,
+            meta=meta,
+            token="pivot_table_first",
+            chunk_kwargs=kwargs,
+        )
+    elif aggfunc == "last":
+        return apply_concat_apply(
+            [df],
+            chunk=methods.pivot_last,
+            aggregate=methods.pivot_agg_last,
+            meta=meta,
+            token="pivot_table_last",
+            chunk_kwargs=kwargs,
+        )
     else:
         raise ValueError
 


Test patch:
diff --git a/dask/dataframe/tests/test_reshape.py b/dask/dataframe/tests/test_reshape.py
--- a/dask/dataframe/tests/test_reshape.py
+++ b/dask/dataframe/tests/test_reshape.py
@@ -181,7 +181,7 @@ def test_get_dummies_errors():
 
 
 @pytest.mark.parametrize("values", ["B", ["B"], ["B", "D"]])
-@pytest.mark.parametrize("aggfunc", ["mean", "sum", "count"])
+@pytest.mark.parametrize("aggfunc", ["mean", "sum", "count", "first", "last"])
 def test_pivot_table(values, aggfunc):
     df = pd.DataFrame(
         {
@@ -210,6 +210,32 @@ def test_pivot_table(values, aggfunc):
     assert_eq(res, exp)
 
 
+@pytest.mark.parametrize("values", ["B", ["D"], ["B", "D"]])
+@pytest.mark.parametrize("aggfunc", ["first", "last"])
+def test_pivot_table_firstlast(values, aggfunc):
+
+    df = pd.DataFrame(
+        {
+            "A": np.random.choice(list("XYZ"), size=100),
+            "B": np.random.randn(100),
+            "C": pd.Categorical(np.random.choice(list("abc"), size=100)),
+            "D": np.random.choice(list("abc"), size=100),
+        }
+    )
+    ddf = dd.from_pandas(df, 5).repartition((0, 20, 40, 60, 80, 98, 99))
+
+    res = dd.pivot_table(ddf, index="A", columns="C", values=values, aggfunc=aggfunc)
+    exp = pd.pivot_table(df, index="A", columns="C", values=values, aggfunc=aggfunc)
+
+    assert_eq(exp, res)
+
+    # method
+    res = ddf.pivot_table(index="A", columns="C", values=values, aggfunc=aggfunc)
+    exp = df.pivot_table(index="A", columns="C", values=values, aggfunc=aggfunc)
+
+    assert_eq(exp, res)
+
+
 def test_pivot_table_dtype():
 
     df = pd.DataFrame(
@@ -266,7 +292,7 @@ def test_pivot_table_errors():
         dd.pivot_table(ddf, index="A", columns="C", values=[["B"]])
     assert msg in str(err.value)
 
-    msg = "aggfunc must be either 'mean', 'sum' or 'count'"
+    msg = "aggfunc must be either 'mean', 'sum', 'count', 'first', 'last'"
     with pytest.raises(ValueError) as err:
         dd.pivot_table(ddf, index="A", columns="C", values="B", aggfunc=["sum"])
     assert msg in str(err.value)