Reduce amount of remote calls for square-like dataframes

[dchigarev]

What's the problem?

There are certain scenarios where some patterns of distributed executions work better than others for different types of datasets even if they're implementing the same function. Consider this example from Modin's optimization notes:

When we're applying a Map/MapReduce operator to a frame having a square (MxM) or close to square shape we would have a NUM_CPUS^2 partitions for this frame and so NUM_CPUS^2 remote calls for this Map function. It's known that this amount of remote calls per core results in a really poor performance with Ray (and probably Dask):

	Modin with default partitioning	Pandas
Map (.abs())	3.88s	0.09s
MapReduce (.sum())	4.42s	0.1s

p.s. all of the measurements were conducted on 2x 28 Cores (56 threads) Xeon Platinum-8276L @ 2.20 with the MODIN_NPARTITIONS=MODIN_NCPUS=112 being set.

How to run this

import modin.pandas as pd
import numpy as np
from timeit import default_timer as timer
from modin.core.execution.ray.implementations.pandas_on_ray.partitioning.partition_manager import PandasOnRayDataframePartitionManager

M = 5000

md_df = pd.DataFrame({f"col{i}": np.arange(M) for i in range(M)})
pd_df = md_df._to_pandas()

### Trying a 'Map' function
t1 = timer()
res = md_df.abs()
PandasOnRayDataframePartitionManager.wait_partitions(res._query_compiler._modin_frame._partitions.flatten())
print(f"Modin .abs: {timer() - t1}")


t1 = timer()
pd_df.abs()
print(f"Pandas .abs: {timer() - t1}")


### Trying a 'MapReduce' function
t1 = timer()
res = md_df.sum()
PandasOnRayDataframePartitionManager.wait_partitions(res._query_compiler._modin_frame._partitions.flatten())
print(f"Modin .sum: {timer() - t1}")

t1 = timer()
pd_df.sum()
print(f"Pandas .sum: {timer() - t1}")

What to do?

There are several options of how we can resolve this:

1. (Complex) Check for partitioning before every Map call and rebalance them so there'll be fewer remote calls.
2. (Easy) Check for partitioning before every Map call and if there're too many partitions then call the function across row/column axis so the number of remote calls would equal to the number of row/column partitions (fewer than the total amount of partitions).

I've made a draft implementation for each of these approaches and made measurements with the same dataset of 5000x5000 ('Rebalanced partitioning' also includes time for actual rebalancing):

	Default partitoning	Rebalanced partitioning	FullAxis approach
Map	3.88s	1.55s	0.87s
MapReduce	4.42s	1.23s	0.89s

How to run this

1. Apply the following patch with a draft implementation for rebalancing across the row axis:

From f2402765f5a55b61a430ea9ec91325a6e9856894 Mon Sep 17 00:00:00 2001
From: Dmitry Chigarev <[email protected]>
Date: Thu, 8 Dec 2022 10:02:10 -0600
Subject: [PATCH] Draft rebalancing for axis=1

Signed-off-by: Dmitry Chigarev <[email protected]>
---
 .../pandas/partitioning/partition_manager.py  | 51 ++++++++++++-------
 1 file changed, 33 insertions(+), 18 deletions(-)

diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
index c0f41740..d71dd618 100644
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -152,7 +152,7 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
         ]
 
     @classmethod
-    def row_partitions(cls, partitions):
+    def row_partitions(cls, partitions, full_axis=True):
         """
         List of `BaseDataframeAxisPartition` objects representing row-wise partitions.
 
@@ -173,7 +173,7 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
         """
         if not isinstance(partitions, list):
             partitions = [partitions]
-        return [cls._row_partition_class(row) for frame in partitions for row in frame]
+        return [cls._row_partition_class(row, full_axis=full_axis) for frame in partitions for row in frame]
 
     @classmethod
     def axis_partition(cls, partitions, axis, full_axis: bool = True):
@@ -1358,7 +1358,7 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
         [part.drain_call_queue() for row in partitions for part in row]
 
     @classmethod
-    def rebalance_partitions(cls, partitions):
+    def rebalance_partitions(cls, partitions, axis=0):
         """
         Rebalance a 2-d array of partitions if we are using ``PandasOnRay`` or ``PandasOnDask`` executions.
 
@@ -1389,7 +1389,7 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
             # the ideal number of partitions is larger than this threshold. The
             # threshold is a heuristic that may need to be tuned for performance.
             max_excess_of_num_partitions = 1.5
-            num_existing_partitions = partitions.shape[0]
+            num_existing_partitions = partitions.shape[axis]
             ideal_num_new_partitions = NPartitions.get()
             if (
                 num_existing_partitions
@@ -1399,7 +1399,7 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
             # If any partition has an unknown length, give each axis partition
             # roughly the same number of row partitions. We use `_length_cache` here
             # to avoid materializing any unmaterialized lengths.
-            if any(
+            if axis == 1 or any(
                 partition._length_cache is None
                 for row in partitions
                 for partition in row
@@ -1410,19 +1410,34 @@ class PandasDataframePartitionManager(ClassLogger, ABC):
                 chunk_size = compute_chunksize(
                     num_existing_partitions, ideal_num_new_partitions, min_block_size=1
                 )
-                new_partitions = np.array(
-                    [
-                        cls.column_partitions(
-                            partitions[i : i + chunk_size],
-                            full_axis=False,
-                        )
-                        for i in range(
-                            0,
-                            num_existing_partitions,
-                            chunk_size,
-                        )
-                    ]
-                )
+                if axis == 0:
+                    new_partitions = np.array(
+                        [
+                            cls.column_partitions(
+                                partitions[i : i + chunk_size],
+                                full_axis=False,
+                            )
+                            for i in range(
+                                0,
+                                num_existing_partitions,
+                                chunk_size,
+                            )
+                        ]
+                    )
+                else:
+                    new_partitions = np.array(
+                        [
+                            cls.row_partitions(
+                                partitions[:, i : i + chunk_size],
+                                full_axis=False
+                            )
+                            for i in range(
+                                0,
+                                num_existing_partitions,
+                                chunk_size,
+                            )
+                        ]
+                    )
                 return new_partitions, None
 
             # If we know the number of rows in every partition, then we should try
-- 
2.25.1

2. Run the following code:

import modin.pandas as pd
import numpy as np
from timeit import default_timer as timer
from modin.core.execution.ray.implementations.pandas_on_ray.partitioning.partition_manager import PandasOnRayDataframePartitionManager
import modin.config as cfg

M = 5000

def rebalance_partitions(md_df):
    npartitions_old = cfg.NPartitions.get()
    npartitions_new = int(npartitions_old ** 0.5)

    cfg.NPartitions.put(npartitions_new)

    parts, _ = PandasOnRayDataframePartitionManager.rebalance_partitions(md_df._partitions, axis=0)
    parts, _ = PandasOnRayDataframePartitionManager.rebalance_partitions(parts, axis=1)
    new_md_df = md_df.__constructor__(
        parts,
        index=md_df._index_cache,
        columns=md_df._columns_cache,
        row_lengths=[len(md_df.index) // npartitions_new] * npartitions_new,
        column_widths=[len(md_df.columns) // npartitions_new] * npartitions_new,
        dtypes=md_df._dtypes
    )

    cfg.NPartitions.put(npartitions_old)
    return new_md_df

md_df = pd.DataFrame({f"col{i}": np.arange(M) for i in range(M)})


t1 = timer()
res = md_df._query_compiler._modin_frame.map(func=lambda df: df.abs())
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"Map abs: {timer() - t1}")


t1 = timer()
new_md_df = rebalance_partitions(md_df._query_compiler._modin_frame)
res = new_md_df.map(func=lambda df: df.abs())
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"Map abs with rebalanced partitions: {timer() - t1}")

t1 = timer()
res = md_df._query_compiler._modin_frame.fold(
    func=lambda df: df.abs(),
    axis=0,
)
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"FullAxis abs: {timer() - t1}")


### Trying a 'MapReduce' function
t1 = timer()
res = md_df._query_compiler._modin_frame.tree_reduce(
    map_func=lambda df: df.sum(),
    reduce_func=lambda df: df.sum(),
    axis=0
)
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"MapReduce sum: {timer() - t1}")

t1 = timer()
new_md_df = rebalance_partitions(md_df._query_compiler._modin_frame)
res = new_md_df.tree_reduce(
    map_func=lambda df: df.sum(),
    reduce_func=lambda df: df.sum(),
    axis=0
)
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"MapReduce sum with rebalanced partitions: {timer() - t1}")

t1 = timer()
res = md_df._query_compiler._modin_frame.reduce(function=lambda df: df.sum(), axis=0)
PandasOnRayDataframePartitionManager.wait_partitions(res._partitions.flatten())
print(f"FullAxis sum: {timer() - t1}")

Why can't we just change the default partitioning scheme so square-like frames would not have NCORES^2 partitions by default?

Well, there are not only Map and MapReduce operators existing in Modin's DataFrame algebra. There's still a mass usage of full-axis functions that would be hurt by decreasing the total amount of partitions.

Would like to hear more opinions/suggestions from @modin-project/modin-core of what we should do with it.

[YarShev]

I think the first approach ("complex") is clear and reasonable. I wonder why "easy" approach is faster than "complex` in your measurements?

[vnlitvinov]

@dchigarev could you please elaborate what exactly is "rebalance" in complex option?

[dchigarev]

I think the first approach ("complex") is clear and reasonable. I wonder why "easy" approach is faster than "complex` in your measurements?

I guess it's because of the "rebalancing" implementation. The whole point of both approaches is to reduce the number of partitions to NCORES when an initial partitioning is NCORES * NCORES.

The "complex" approach does so by making each axis have sqrt(NCORES) splits, thus requiring repartitioning along both axes sequentially.

On the other hand, the "easy" approach doesn't even require repartitioning, it just builds NCORES amount of column partitions.

It appears that in a competition of "repartition along rows + repartition along cols + apply a function to 'NCORES' amount of partitions" vs "build column partitions + apply a function to 'NCORES' amount of partitions" the second approach wins.

Again, probably we could advance the repartitioning logic somehow for the first approach but I really doubt that we'll make it faster than the second one just because the first one requires manipulations with both of the axes and the second only with a single one.

[dchigarev]

@vnlitvinov

could you please elaborate what exactly is "rebalance" in complex option?

"Rebalance" means repartitioning so the frame would have fewer total amount of partitions. An exact way of doing so is implementation details I would say.

In the measurements above I used the .rebalance_partitions method along both axes sequentially in order to have a "perfect" partitioning (check the spoiler in the issue's head for the exact code). This may not be the most optimal way of rebalancing though.

[vnlitvinov]

I'm guessing "easy" is faster because it runs NCORES jobs in parallel and (probably) does not put intermediate partitions in the object store but concatenates them, runs the map function and puts the result back.

"Complex" solution could be running in sqrt(NCORES) jobs, plus it might need to put intermediate partitions in the object store.

[vnlitvinov]

That said, I propose to measure distributed workloads (on a cluster) as well, as they may behave very differently due to worker-driver exchange being much more expensieve.

[YarShev]

Ah, I see why the "easy" approach is faster. I originally thought that you proposed to first make a repartitioning in one remote call and then perform "map" in other remote call.

The "complex" approach would complicate a virtual partition because it would consist of both a row part and a col part.

"Complex" solution could be running in sqrt(NCORES) jobs, plus it might need to put intermediate partitions in the object store.

Yes, I believe "Complex" solution could be running in sqrt(NCORES) jobs, and we could avoid the need to put intermediate partitions in the object store. However, this would complicate the virtual partition implementation. Also, I am not sure what is better in terms of performance.

I think we could move forward with "easy" approach and run "map" either row-wise or col-wise in depend on the number of partitions along the axes.

[vnlitvinov]

From my point of view, if "easy" approach does not exhibit a slowdown on a real cluster (i.e. in a a non-single-node scenario), then let's implement it as it should be only benefits. We can always build a more advanced solution later on.

[YarShev]

Reopening since #7136 covered Map operator only.