Repository: modin-project/modin

Issue:
Add range-partitioning implementation for `.pivot_table()`



Discussion:

Patch:
diff --git a/modin/core/storage_formats/pandas/groupby.py b/modin/core/storage_formats/pandas/groupby.py
--- a/modin/core/storage_formats/pandas/groupby.py
+++ b/modin/core/storage_formats/pandas/groupby.py
@@ -245,3 +245,256 @@ def mean_reduce(dfgb, **kwargs):
     "skew": GroupbyReduceImpl._build_skew_impl(),
     "sum": ("sum", "sum", lambda grp, *args, **kwargs: grp.sum(*args, **kwargs)),
 }
+
+
+class PivotTableImpl:
+    """Provide MapReduce, Range-Partitioning and Full-Column implementations for 'pivot_table()'."""
+
+    @classmethod
+    def map_reduce_impl(
+        cls, qc, unique_keys, drop_column_level, pivot_kwargs
+    ):  # noqa: PR01
+        """Compute 'pivot_table()' using MapReduce implementation."""
+        if pivot_kwargs["margins"]:
+            raise NotImplementedError(
+                "MapReduce 'pivot_table' implementation doesn't support 'margins=True' parameter"
+            )
+
+        index, columns, values = (
+            pivot_kwargs["index"],
+            pivot_kwargs["columns"],
+            pivot_kwargs["values"],
+        )
+        aggfunc = pivot_kwargs["aggfunc"]
+
+        if not GroupbyReduceImpl.has_impl_for(aggfunc):
+            raise NotImplementedError(
+                "MapReduce 'pivot_table' implementation only supports 'aggfuncs' that are implemented in 'GroupbyReduceImpl'"
+            )
+
+        if len(set(index).intersection(columns)) > 0:
+            raise NotImplementedError(
+                "MapReduce 'pivot_table' implementation doesn't support intersections of 'index' and 'columns'"
+            )
+
+        to_group, keys_columns = cls._separate_data_from_grouper(
+            qc, values, unique_keys
+        )
+        to_unstack = columns if index else None
+
+        result = GroupbyReduceImpl.build_qc_method(
+            aggfunc,
+            finalizer_fn=lambda df: cls._pivot_table_from_groupby(
+                df,
+                pivot_kwargs["dropna"],
+                drop_column_level,
+                to_unstack,
+                pivot_kwargs["fill_value"],
+            ),
+        )(
+            to_group,
+            by=keys_columns,
+            axis=0,
+            groupby_kwargs={
+                "observed": pivot_kwargs["observed"],
+                "sort": pivot_kwargs["sort"],
+            },
+            agg_args=(),
+            agg_kwargs={},
+            drop=True,
+        )
+
+        if to_unstack is None:
+            result = result.transpose()
+        return result
+
+    @classmethod
+    def full_axis_impl(
+        cls, qc, unique_keys, drop_column_level, pivot_kwargs
+    ):  # noqa: PR01
+        """Compute 'pivot_table()' using full-column-axis implementation."""
+        index, columns, values = (
+            pivot_kwargs["index"],
+            pivot_kwargs["columns"],
+            pivot_kwargs["values"],
+        )
+
+        to_group, keys_columns = cls._separate_data_from_grouper(
+            qc, values, unique_keys
+        )
+
+        def applyier(df, other):  # pragma: no cover
+            """
+            Build pivot table for a single partition.
+
+            Parameters
+            ----------
+            df : pandas.DataFrame
+                Partition of the self frame.
+            other : pandas.DataFrame
+                Broadcasted partition that contains `value` columns
+                of the self frame.
+
+            Returns
+            -------
+            pandas.DataFrame
+                Pivot table for this particular partition.
+            """
+            concated = pandas.concat([df, other], axis=1, copy=False)
+            result = pandas.pivot_table(
+                concated,
+                **pivot_kwargs,
+            )
+
+            # if only one value is specified, removing level that maps
+            # columns from `values` to the actual values
+            if drop_column_level is not None:
+                result = result.droplevel(drop_column_level, axis=1)
+
+            # in that case Pandas transposes the result of `pivot_table`,
+            # transposing it back to be consistent with column axis values along
+            # different partitions
+            if len(index) == 0 and len(columns) > 0:
+                result = result.transpose()
+
+            return result
+
+        result = qc.__constructor__(
+            to_group._modin_frame.broadcast_apply_full_axis(
+                axis=0, func=applyier, other=keys_columns._modin_frame
+            )
+        )
+
+        # transposing the result again, to be consistent with Pandas result
+        if len(index) == 0 and len(columns) > 0:
+            result = result.transpose()
+
+        return result
+
+    @classmethod
+    def range_partition_impl(
+        cls, qc, unique_keys, drop_column_level, pivot_kwargs
+    ):  # noqa: PR01
+        """Compute 'pivot_table()' using Range-Partitioning implementation."""
+        if pivot_kwargs["margins"]:
+            raise NotImplementedError(
+                "Range-partitioning 'pivot_table' implementation doesn't support 'margins=True' parameter"
+            )
+
+        index, columns, values = (
+            pivot_kwargs["index"],
+            pivot_kwargs["columns"],
+            pivot_kwargs["values"],
+        )
+
+        if len(set(index).intersection(columns)) > 0:
+            raise NotImplementedError(
+                "Range-partitioning 'pivot_table' implementation doesn't support intersections of 'index' and 'columns'"
+            )
+
+        if values is not None:
+            to_take = list(np.unique(list(index) + list(columns) + list(values)))
+            qc = qc.getitem_column_array(to_take, ignore_order=True)
+
+        to_unstack = columns if index else None
+
+        groupby_result = qc._groupby_shuffle(
+            by=list(unique_keys),
+            agg_func=pivot_kwargs["aggfunc"],
+            axis=0,
+            groupby_kwargs={
+                "observed": pivot_kwargs["observed"],
+                "sort": pivot_kwargs["sort"],
+            },
+            agg_args=(),
+            agg_kwargs={},
+            drop=True,
+        )
+
+        # the length of 'groupby_result' is typically really small here,
+        # so it's okay to call full-column function
+        result = groupby_result._modin_frame.apply_full_axis(
+            axis=0,
+            func=lambda df: cls._pivot_table_from_groupby(
+                df,
+                pivot_kwargs["dropna"],
+                drop_column_level,
+                to_unstack,
+                pivot_kwargs["fill_value"],
+                # FIXME: Range-partitioning impl has a problem with the resulting order in case of multiple grouping keys,
+                # so passing 'sort=True' explicitly in this case
+                # https://github.com/modin-project/modin/issues/6875
+                sort=pivot_kwargs["sort"] if len(unique_keys) > 1 else False,
+            ),
+        )
+
+        if to_unstack is None:
+            result = result.transpose()
+
+        return qc.__constructor__(result)
+
+    @staticmethod
+    def _pivot_table_from_groupby(
+        df, dropna, drop_column_level, to_unstack, fill_value, sort=False
+    ):
+        """
+        Convert group by aggregation result to a pivot table.
+
+        Parameters
+        ----------
+        df : pandas.DataFrame
+            Group by aggregation result.
+        dropna : bool
+            Whether to drop NaN columns.
+        drop_column_level : int or None
+            An extra columns level to drop.
+        to_unstack : list of labels or None
+            Group by keys to pass to ``.unstack()``. Reperent `columns` parameter
+            for ``.pivot_table()``.
+        fill_value : bool
+            Fill value for NaN values.
+        sort : bool, default: False
+            Whether to sort the result along index.
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        if df.index.nlevels > 1 and to_unstack is not None:
+            df = df.unstack(level=to_unstack)
+        if drop_column_level is not None:
+            df = df.droplevel(drop_column_level, axis=1)
+        if dropna:
+            df = df.dropna(axis=1, how="all")
+        if fill_value is not None:
+            df = df.fillna(fill_value, downcast="infer")
+        if sort:
+            df = df.sort_index(axis=0)
+        return df
+
+    @staticmethod
+    def _separate_data_from_grouper(qc, values, unique_keys):
+        """
+        Split `qc` for key columns to group by and values to aggregate.
+
+        Parameters
+        ----------
+        qc : PandasQueryCompiler
+        values : list of labels or None
+            List of columns to aggregate. ``None`` means all columns except 'unique_keys'.
+        unique_keys : list of labels
+            List of key columns to group by.
+
+        Returns
+        -------
+        to_aggregate : PandasQueryCompiler
+        keys_to_group : PandasQueryCompiler
+        """
+        if values is None:
+            to_aggregate = qc.drop(columns=unique_keys)
+        else:
+            to_aggregate = qc.getitem_column_array(np.unique(values), ignore_order=True)
+
+        keys_to_group = qc.getitem_column_array(unique_keys, ignore_order=True)
+
+        return to_aggregate, keys_to_group
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -75,7 +75,7 @@
 )
 
 from .aggregations import CorrCovBuilder
-from .groupby import GroupbyReduceImpl
+from .groupby import GroupbyReduceImpl, PivotTableImpl
 from .merge import MergeImpl
 from .utils import get_group_names, merge_partitioning
 
@@ -4187,62 +4187,6 @@ def __convert_by(by):
 
         return unstacked
 
-    def _pivot_table_tree_reduce(
-        self, grouper, aggfunc, drop_column_level, fill_value, dropna, to_unstack=None
-    ):
-        """
-        Build a pivot table using TreeReduce implementation.
-
-        Parameters
-        ----------
-        grouper : PandasQueryCompiler
-            QueryCompiler holding columns to group on.
-        aggfunc : str
-            Aggregation to perform against the values of the pivot table. Note that ``GroupbyReduceImpl``
-            has to be able to build implementation for this aggregation.
-        drop_column_level : bool
-            Whether to drop the top level of the columns.
-        fill_value : object
-            Fill value for None values in the result.
-        dropna : bool
-            Whether to drop NaN columns.
-        to_unstack : list, optional
-            A list of column names to pass to the `.unstack()` when building the pivot table.
-            If `None` was passed perform regular transpose instead of unstacking.
-
-        Returns
-        -------
-        PandasQueryCompiler
-            A query compiler holding a pivot table.
-        """
-
-        def make_pivot_table(df):
-            if df.index.nlevels > 1 and to_unstack is not None:
-                df = df.unstack(level=to_unstack)
-            if drop_column_level and df.columns.nlevels > 1:
-                df = df.droplevel(0, axis=1)
-            if dropna:
-                df = df.dropna(axis=1, how="all")
-            if fill_value is not None:
-                df = df.fillna(fill_value, downcast="infer")
-            return df
-
-        result = GroupbyReduceImpl.build_qc_method(
-            aggfunc, finalizer_fn=make_pivot_table
-        )(
-            self,
-            by=grouper,
-            axis=0,
-            groupby_kwargs={},
-            agg_args=(),
-            agg_kwargs={},
-            drop=True,
-        )
-
-        if to_unstack is None:
-            result = result.transpose()
-        return result
-
     def pivot_table(
         self,
         index,
@@ -4269,94 +4213,62 @@ def __convert_by(by):
                 return list(by)
             return _convert_by(by)
 
-        drop_column_level = values is not None and not is_list_like(values)
-        index, columns, values = map(__convert_by, [index, columns, values])
+        is_1d_values = values is not None and not is_list_like(values)
+        index, columns = map(__convert_by, [index, columns])
 
-        unique_keys = np.unique(index + columns)
-        unique_values = np.unique(values)
+        if len(index) + len(columns) == 0:
+            raise ValueError("No group keys passed!")
 
-        if len(values):
-            to_group = self.getitem_column_array(unique_values, ignore_order=True)
+        if is_1d_values and len(index) > 0 and len(columns) > 0:
+            drop_column_level = 1 if isinstance(aggfunc, list) else 0
         else:
-            to_group = self.drop(columns=unique_keys)
-
-        keys_columns = self.getitem_column_array(unique_keys, ignore_order=True)
-
-        # Here we can use TreeReduce implementation that tends to be more efficient rather full-axis one
-        if (
-            not margins
-            and GroupbyReduceImpl.has_impl_for(aggfunc)
-            and len(set(index).intersection(columns)) == 0
-        ):
-            return to_group._pivot_table_tree_reduce(
-                keys_columns,
-                aggfunc,
-                drop_column_level=drop_column_level,
-                fill_value=fill_value,
-                dropna=dropna,
-                to_unstack=columns if index else None,
-            )
+            drop_column_level = None
 
-        len_values = len(values)
-        if len_values == 0:
-            len_values = len(self.columns.drop(unique_keys))
+        # if the value is 'None' it will be converted to an empty list (no columns to aggregate),
+        # which is invalid for 'values', as 'None' means aggregate ALL columns instead
+        if values is not None:
+            values = __convert_by(values)
 
-        def applyier(df, other):  # pragma: no cover
-            """
-            Build pivot table for a single partition.
+        # using 'pandas.unique' instead of 'numpy' as it guarantees to not change the original order
+        unique_keys = pandas.Series(index + columns).unique()
 
-            Parameters
-            ----------
-            df : pandas.DataFrame
-                Partition of the self frame.
-            other : pandas.DataFrame
-                Broadcasted partition that contains `value` columns
-                of the self frame.
+        kwargs = {
+            "qc": self,
+            "unique_keys": unique_keys,
+            "drop_column_level": drop_column_level,
+            "pivot_kwargs": {
+                "index": index,
+                "values": values,
+                "columns": columns,
+                "aggfunc": aggfunc,
+                "fill_value": fill_value,
+                "margins": margins,
+                "dropna": dropna,
+                "margins_name": margins_name,
+                "observed": observed,
+                "sort": sort,
+            },
+        }
 
-            Returns
-            -------
-            pandas.DataFrame
-                Pivot table for this particular partition.
-            """
-            concated = pandas.concat([df, other], axis=1, copy=False)
-            result = pandas.pivot_table(
-                concated,
-                index=index,
-                values=values if len(values) > 0 else None,
-                columns=columns,
-                aggfunc=aggfunc,
-                fill_value=fill_value,
-                margins=margins,
-                dropna=dropna,
-                margins_name=margins_name,
-                observed=observed,
-                sort=sort,
+        try:
+            return PivotTableImpl.map_reduce_impl(**kwargs)
+        except NotImplementedError as e:
+            message = (
+                f"Can't use MapReduce 'pivot_table' implementation because of: {e}"
+                + "\nFalling back to a range-partitioning implementation."
             )
+            get_logger().info(message)
 
-            # if only one value is specified, removing level that maps
-            # columns from `values` to the actual values
-            if len(index) > 0 and len_values == 1 and result.columns.nlevels > 1:
-                result.columns = result.columns.droplevel(int(margins))
-
-            # in that case Pandas transposes the result of `pivot_table`,
-            # transposing it back to be consistent with column axis values along
-            # different partitions
-            if len(index) == 0 and len(columns) > 0:
-                result = result.T
-
-            return result
-
-        result = self.__constructor__(
-            to_group._modin_frame.broadcast_apply_full_axis(
-                axis=0, func=applyier, other=keys_columns._modin_frame
+        try:
+            return PivotTableImpl.range_partition_impl(**kwargs)
+        except NotImplementedError as e:
+            message = (
+                f"Can't use range-partitioning 'pivot_table' implementation because of: {e}"
+                + "\nFalling back to a full-axis implementation."
             )
-        )
-
-        # transposing the result again, to be consistent with Pandas result
-        if len(index) == 0 and len(columns) > 0:
-            result = result.transpose()
+            get_logger().info(message)
 
-        return result
+        return PivotTableImpl.full_axis_impl(**kwargs)
 
     # Get_dummies
     def get_dummies(self, columns, **kwargs):


Test patch:
diff --git a/modin/pandas/test/dataframe/test_default.py b/modin/pandas/test/dataframe/test_default.py
--- a/modin/pandas/test/dataframe/test_default.py
+++ b/modin/pandas/test/dataframe/test_default.py
@@ -701,6 +701,7 @@ def test_pivot_table_data(data, index, columns, values, aggfunc, request):
         and "int_data" in request.node.callspec.id
     ):
         raising_exceptions = TypeError("'numpy.float64' object is not callable")
+
     eval_general(
         md_df,
         pd_df,
@@ -719,6 +720,7 @@ def test_pivot_table_data(data, index, columns, values, aggfunc, request):
 @pytest.mark.parametrize(
     "index",
     [
+        pytest.param([], id="no_index_cols"),
         pytest.param(lambda df: df.columns[0], id="single_index_column"),
         pytest.param(
             lambda df: [df.columns[0], df.columns[len(df.columns) // 2 - 1]],
@@ -787,6 +789,29 @@ def test_pivot_table_margins(
     )
 
 
+@pytest.mark.parametrize(
+    "aggfunc",
+    [
+        pytest.param("sum", id="MapReduce_func"),
+        pytest.param("nunique", id="FullAxis_func"),
+    ],
+)
+@pytest.mark.parametrize("margins", [True, False])
+def test_pivot_table_fill_value(aggfunc, margins):
+    md_df, pd_df = create_test_dfs(test_data["int_data"])
+    eval_general(
+        md_df,
+        pd_df,
+        operation=lambda df, *args, **kwargs: df.pivot_table(*args, **kwargs),
+        index=md_df.columns[0],
+        columns=md_df.columns[1],
+        values=md_df.columns[2],
+        aggfunc=aggfunc,
+        margins=margins,
+        fill_value=10,
+    )
+
+
 @pytest.mark.parametrize("data", [test_data["int_data"]], ids=["int_data"])
 def test_pivot_table_dropna(data):
     eval_general(