[Python] Only convert in parallel for the ConsolidatedBlockCreator class for large data

[anjakefala]

Describe the enhancement requested

The Consolidated Block Creator runs the column conversion in parallel, creating a Scalar Memo Table for each column up until pa.cpu_count().

For performance reasons, jemalloc and mimalloc maintain allocations on a per-memory segment level to reduce contention between threads.

What this means is that if a user calls table.to_pandas(split_blocks=False) for a small table, a disproportionately large amount of memory gets allocated to build the Scalar Memo Table. Both jemalloc and mimalloc will essentially allocate a chunk of memory per column.

Here is some code:

import pyarrow as pa

def test_memory_usage():
    table = pa.table({'A': 'a', 'B': 'b', 'C': 'c', 'D': 'd', 'E': 'e', 'F': 'f', 'G': 'g'}) #'H':'h', 'I': 'i', 'J': 'j', 'K':'k', 'L':'l', 'M':'m', 'N':'n', 'O':'o', 'P':'p', 'Q':'q', 'R':'r', 'S':'s', 'T':'t', 'U':'u', 'V':'v', 'W':'w', 'X':'x','Y':'y','Z':'z'})
    df = table.to_pandas()

if __name__ == '__main__':
    test_memory_usage()

Here are the resulting memory allocations summarised with memray:

jemalloc with 7 columns:

📦 Total memory allocated:
        178.127MB

📊 Histogram of allocation size:
        min: 1.000B
        --------------------------------------------
        < 4.000B   :  3548 ▇▇
        < 24.000B  :  1253 ▇
        < 119.000B : 14686 ▇▇▇▇▇▇▇
        < 588.000B :  3746 ▇▇
        < 2.827KB  : 58959 ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇
        < 13.924KB :  2550 ▇▇
        < 68.569KB :   403 ▇
        < 337.661KB:    84 ▇
        < 1.624MB  :    24 ▇
        <=7.996MB  :    20 ▇
        --------------------------------------------
        max: 7.996MB

jemalloc with 26 columns:

📦 Total memory allocated:
        238.229MB

📊 Histogram of allocation size:
        min: 1.000B
        --------------------------------------------
        < 4.000B   :  3545 ▇▇
        < 24.000B  :  1627 ▇
        < 119.000B : 15086 ▇▇▇▇▇▇▇
        < 588.000B :  3882 ▇▇
        < 2.828KB  : 58971 ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇
        < 13.929KB :  2552 ▇▇
        < 68.593KB :   403 ▇
        < 337.794KB:    84 ▇
        < 1.625MB  :    24 ▇
        <=8.000MB  :    47 ▇
        --------------------------------------------
        max: 8.000MB

mimalloc with 7 columns:

📦 Total memory allocated:
        380.166MB

📊 Histogram of allocation size:
        min: 1.000B
        --------------------------------------------
        < 6.000B   :  3548 ▇▇
        < 36.000B  :  7470 ▇▇▇▇
        < 222.000B :  9524 ▇▇▇▇▇
        < 1.319KB  : 57271 ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇
        < 7.999KB  :  6492 ▇▇▇
        < 48.503KB :   775 ▇
        < 294.066KB:   150 ▇
        < 1.741MB  :    26 ▇
        < 10.556MB :     1 ▇
        <=64.000MB :     4 ▇
        --------------------------------------------
        max: 64.000MB

mimalloc with 26 columns:

📦 Total memory allocated:
        1.434GB

📊 Histogram of allocation size:
        min: 1.000B
        --------------------------------------------
        < 6.000B   :  3545 ▇▇
        < 36.000B  :  7845 ▇▇▇▇
        < 222.000B : 10001 ▇▇▇▇▇
        < 1.319KB  : 57332 ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇
        < 7.999KB  :  6501 ▇▇▇
        < 48.503KB :   794 ▇
        < 294.066KB:   150 ▇
        < 1.741MB  :    26 ▇
        < 10.556MB :     1 ▇
        <=64.000MB :    21 ▇
        --------------------------------------------
        max: 64.000MB

You can see how dramatically the memory increases even for a very small table.

My proposal is that we only do the conversion in parallel when it might make a substantial performance difference for a table of a certain size. I'm not quite sure which size, but once the code has been refactored, we can run experiments to come to a data-informed decision.

Component(s)

C++

[felipecrv]

I’ve taken a look at the code.

OptionalParallelFor is used when PandasOptions::use_threads is set. Internally, it uses the default CPU thread pool which can allocate CPU-count threads at a time. That’s why it allocates CPU-count threads as @anjakefala described.

~My first solution/idea: add another option to PandasOptions — int threads. Then if that’s different than 0, create a different ThreadPool with that as capacity when calling the parallel-for.

Easy, but it’s an extra option that users have to remember to set.~

EDIT 1: I had a better idea.

Alternative

Add more parameters to parallel-for(columns) that modify the loop that submits tasks to the thread-pool:

Start with a low minimum number of submitted tasks (eg 2. each column is a task). After that, tasks are added with a small delay (microseconds) as to give the chance of the initially submitted tasks to finish (because cols are small) and make these initially used threads to be re-used.

This might require internal ThreadPool changes to avoid waits that wouldn’t lead to a global reduction in time taken to convert all columns.

EDIT 2: the better solution is too allow the same thread to continue converting more columns if available with a loop in the main thread increasing the number of threads used if necessary

[anjakefala]

@felipecrv Is modifying ThreadPool better than an option where we use an approach similar to the SplitBlockCreator class for tables under a certain size? That's more along the line of what I was thinking of.

However, if you think work-stealing would be the most robust solution, that other functions would benefit from, I'd be game for approaching this.

I prefer the work-stealing approach because, ideally, we wouldn't require the user to know about the existence of an option to set. Folks might not know that the memory usage has to do with the spawning of individual threads. They might not even know why to_pandas spawns multiple threads.

[felipecrv]

Sharing some progress here in the open:

@anjakefala will test a version of this that uses a version of ParallelFor [1] that doesn't create one Future per task/column — the function keeps going by popping the next task (column index) from a shared std::atomic<int> called next_task. This atomic integer plays the same role as the multiple task queues in a generic work-stealing algorithm implementation.

In ParallelForWithBackoff [1], each worker starts from a task/column they should complete and when done it can keep going while next_task.fetch_add(1, std::memory_order_acq_rel) returns a task smaller than num_tasks. Getting that atomic and incrementing it is how it "steals" the task from all the other workers. This maximizes throughput and reduces thread-local allocations because the same thread can very quickly go through many columns.

[1] https://github.com/felipecrv/arrow/blob/b54cb408a8c73af0c4fc76cdceadca0c0b240cb0/cpp/src/arrow/util/parallel.h#L104-L229

[pitrou]

I'm honestly not sure what this has to do with work stealing.

If I read correctly:

• a new ScalarMemoTable is created when converting each column (either eagerly or lazily, this might or might not matter)
• columns are converted in parallel for smaller execution times
• a lot of memory allocations may occur because of this

What I would like to know is:

1. What is the concern? The total number of allocations? The max allocated working set during the workload?
2. Is it a practical concern?

[jorisvandenbossche]

Is it a practical concern?

I think the original reported case (converting a tiny table of a few kilobytes to pandas can give a spike of several hundred MBs in memory usage) is something people can certainly run into. And although it will often not be concerning (typically when working with smaller data, memory usage is not an issue, and when actually working with larger tables and memory usage becomes relevant, this overhead will disappear), it is definitely surprising and can lead to confusion. So I think it is worth "fixing".

But the potential fix I was thinking of could also be something much simpler, like with some heuristic decide to just not do the conversion in parallel for smaller data.
For the conversion the other way around (pandas -> pyarrow), we actually have some heuristic currently (in python):

arrow/python/pyarrow/pandas_compat.py

Lines 573 to 581 in a6e577d

	# NOTE(wesm): If nthreads=None, then we use a heuristic to decide whether
	# using a thread pool is worth it. Currently the heuristic is whether the
	# nrows > 100 * ncols and ncols > 1.
	if nthreads is None:
	nrows, ncols = len(df), len(df.columns)
	if nrows > ncols * 100 and ncols > 1:
	nthreads = pa.cpu_count()
	else:
	nthreads = 1

[pitrou]

Yes, I agree this could just be fixed with a heuristic, especially as parallelizing will not be performant on very small data.

[felipecrv]

I'm honestly not sure what this has to do with work stealing.

It's the inspiration for this, but if it seems to make the understanding harder, just ignore it.

What is the concern? The total number of allocations? The max allocated working set during the workload?

The practical concern is that jemalloc and mimalloc allocate significant amounts of thread-local data:

For performance reasons, jemalloc and microsoft/mimalloc#351 maintain allocations on a per-memory segment level to reduce contention between threads.

---

Yes, I agree this could just be fixed with a heuristic, especially as parallelizing will not be performant on very small data.

And that is exactly what I'm doing, but instead of looking at number of rows I look at how long it takes to convert the columns on average, while doing the conversions. Never blocking [1] or sleeping.

It's an adaptive solution: you pass a duration, if it takes more than that duration, a few more threads are allowed to start.

      // If the average wall-clock time per task is greater than the
      // small_task_duration_secs, grow the number of workers geometrically.
      if (static_cast<double>(elapsed_time.count() * steady_clock::period::num) >
          (small_task_duration_secs * steady_clock::period::den) * num_tasks_completed) {
        num_workers = std::min(
            ((num_workers * GrowthRatio::num - 1) / GrowthRatio::den) + 1, max_workers);
      }

And instead of having these issues for every kind of parallelization we do in the future, this function can be used on any problem where the size of the tasks can range from trivial to very expensive.

This what the fix looks like at the callsite:

-    return OptionalParallelFor(options_.use_threads, num_columns_, WriteColumn);
+    if (options_.use_threads) {
+      const double kSmallTaskDurationSecs = 0.008;  // 8ms
+      return ::arrow::internal::ParallelForWithBackOff(num_columns_, WriteColumn,
+                                                       kSmallTaskDurationSecs);
+    } else {
+      for (int i = 0; i < num_columns_; ++i) {
+        RETURN_NOT_OK(WriteColumn(i));
+      }
+      return Status::OK();
+    }

[1] except for the lock-free fetch/add of the next_task atomic