Top-K query optimization in sort uses substantial memory

[gruuya]

Describe the bug

It seems like the Top-K query optimization is somehow conditional on the usage of a custom allocator (mimalloc/snmalloc), while in principle that shouldn't be the case?

To Reproduce

1. Grab and build bytehound: https://github.com/koute/bytehound
2. Prepare some large-ish Parquet file, e.g. https://seafowl-public.s3.eu-west-1.amazonaws.com/tutorial/trase-supply-chains.parquet:

$ du -h ~/supply-chains.parquet 
146M /home/ubuntu/supply-chains.parquet

3. Remove the custom allocator and build

diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs
index aea499d60..a92957730 100644
--- a/datafusion-cli/src/main.rs
+++ b/datafusion-cli/src/main.rs
@@ -24,13 +24,13 @@ use datafusion_cli::catalog::DynamicFileCatalog;
 use datafusion_cli::{
     exec, print_format::PrintFormat, print_options::PrintOptions, DATAFUSION_CLI_VERSION,
 };
-use mimalloc::MiMalloc;
+// use mimalloc::MiMalloc;
 use std::env;
 use std::path::Path;
 use std::sync::Arc;

-#[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
+// #[global_allocator]
+// static GLOBAL: MiMalloc = MiMalloc;

 #[derive(Debug, Parser, PartialEq)]
 #[clap(author, version, about, long_about= None)]

4. Profile a Top-K query

$ LD_PRELOAD=~/bytehound/target/release/libbytehound.so ./target/debug/datafusion-cli
DataFusion CLI v28.0.0
❯ CREATE EXTERNAL TABLE supply_chains STORED AS PARQUET LOCATION '/home/ubuntu/supply-chains.parquet';
0 rows in set. Query took 0.445 seconds.
❯ SELECT * FROM supply_chains ORDER BY flow_id DESC LIMIT 1;
...

The profile I get is

Expected behavior

With the custom allocator present the memory profile I see is like this

Additional context

No response

[alamb]

This definitely seems non ideal -- @Dandandan do you have any insights on this?

[Dandandan]

Hi @gruuya thanks for the report. I don't have a direct answer

Two things:

• It looks like you're testing in debug mode? What happens if you build in release mode?
• The sorting / top-k optimizations are not designed with a certain allocator in mind, but I think we've seen some interesting things with some allocators not actively freeing memory (e.g. when memory is fragmented). It could also be that different allocators report memory usage differently. I think you did some experiments as well at influxDB IOX with some allocators @alamb @mkmik

[gruuya]

It looks like you're testing in debug mode? What happens if you build in release mode?

True, but the discrepancy doesn't go away in the release mode, only the run time gets shorter:

• W/O the allocator (~8.5 seconds)
  

• W/ allocator (~4 seconds)
  

One potential explanation for the discrepancy could be that bytehound doesn't properly account for memory allocations under mimalloc, which I think is unlikely; even then the question remains why, apparently, the entire file gets loaded into memory instead of using the top-k flow.

[Dandandan]

Hm, it might not support mimalloc (according to readme it only supports jemalloc).

Something else: the memory usage using the current optimizations will only go down when the total allocated memory exceeds the configured memory limit (the memory usage will not be roughly constant but will go up until the limit has been reached).
Did you configure a memory limit?

[gruuya]

Hm, it might not support mimalloc (according to readme it only supports jemalloc).

Good point; to clarify I've been profiling and optimizing some stuff where we don't configure a custom allocator at all, so in the base case bytehound appears to be working correctly. The surges definitely seem to be real, as they lead to OOMs in memory-constrained envs.

Did you configure a memory limit?

I haven't been setting a memory limit (neither in the above repro using datafusion-cli, nor in seafowl itself), because I was actually investigating memory issues. Naively it seems to me like top-k algorithm shouldn't be affected either way?

[gruuya]

Hm, it might not support mimalloc (according to readme it only supports jemalloc).

Indeed, trying out jemalloc (static GLOBAL: Jemalloc = Jemalloc;) leads to the same memory profile as W/O the allocator above (debug build)

This means that the profile discrepancy was actually an artifact of the bytehound's lack of support for mimalloc/snmalloc (maybe I should revise the issue title?).

The underlying problem still remains though; fwiw here are top 3 allocations grouped by backtrace

The above backtraces could also suggest that the underlying problem is not with the top-k flow itself, but rather with some other adjacent logic which is manifested in that case (e.g. something in parquet crate).

[gruuya]

Oh I think I see the problem: the fetch count is never utilized while the sort is accumulating batches: https://github.com/apache/arrow-datafusion/blob/main/datafusion/core/src/physical_plan/sorts/sort.rs#L628-L632

Consequently, all the batches get loaded into memory inside of the ExternalSorter:

diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs
index f660f0acf..6eb63e39a 100644
--- a/datafusion/core/src/physical_plan/sorts/sort.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -159,6 +159,7 @@ impl ExternalSorter {
             }
         }

+        debug!("Inserted batch of size {} for a total of {} in-memory batches", input.num_rows(), self.in_mem_batches.len() + 1);
         self.in_mem_batches.push(input);
         self.in_mem_batches_sorted = false;
         Ok(())

[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::limit] Start GlobalLimitExec::execute for partition: 0
[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::sorts::sort] Start SortExec::execute for partition 0 of context session_id 8855be36-07c1-495b-b6fd-9cb9501cb46d and task_id None
[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::sorts::sort] End SortExec's input.execute for partition: 0
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 1 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 2 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 3 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 4 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 5 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 6 in-memory batches
...
[2023-08-01T08:14:12Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 710 in-memory batches
[2023-08-01T08:14:12Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 6836 for a total of 711 in-memory batches

This raises a (maybe naive) related question: why bother with materializing the record batches from the input stream and potentially spilling them to the disk at this point, only for them to be streamed back into streaming_merge function again?

EDIT: Alternatively, the following diff fixes the issue:

diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs
index f660f0acf..49e44a834 100644
--- a/datafusion/core/src/physical_plan/sorts/sort.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -133,11 +133,17 @@ impl ExternalSorter {
     /// Appends an unsorted [`RecordBatch`] to `in_mem_batches`
     ///
     /// Updates memory usage metrics, and possibly triggers spilling to disk
-    async fn insert_batch(&mut self, input: RecordBatch) -> Result<()> {
+    async fn insert_batch(&mut self, mut input: RecordBatch) -> Result<()> {
         if input.num_rows() == 0 {
             return Ok(());
         }

+        if self.fetch.map_or(false, |f| f < input.num_rows()) {
+            // Sort the batch pre-emptively to reduce the number of rows
+            // since the fetch count is smaller than the batch size
+            input = sort_batch(&input, &self.expr, self.fetch)?;
+        }
+
         let size = batch_byte_size(&input);
         if self.reservation.try_grow(size).is_err() {
             let before = self.reservation.size();

Not sure about the overall performance implications of this tweak however (i.e. it may impact the runtime for some cases I guess). Also the value of this tweak is progressively reduced the larger the fetch size is, and when fetch size reaches the batch size, we're back to the underlying problem.

[alamb]

cc @NGA-TRAN -- we were also seeing some unexpected peformance behavior with a plan that had SortExec(fetch=1) -- it was much slower than we expected. Perhaps @gruuya 's observation in #7149 (comment) might explain it

[gruuya]

why bother with materializing the record batches from the input stream and potentially spilling them to the disk at this point, only for them to be streamed back into streaming_merge function again?

Oh, to answer my own question it's because streaming_merge assumes the input is sorted, and to do that we need to materialize the original stream, which also involves disk spill over, got it.

One partial improvement would be something like concatenating and sorting the incoming batches into a single batch inside insert_batch, and then optionally spilling that over to disk; still, this results in a trade-off between time and memory efficiency of such cases, since it would require sorting of each intermediate batch (albeit partially sorted already).

I guess the real solution would be something like an external top-k algorithm which would work on the original input stream and use a min/max heap with spillover in case fetch is some value (actually it was my impression this was already implemented). Not sure how viable is that though.

[Dandandan]

why bother with materializing the record batches from the input stream and potentially spilling them to the disk at this point, only for them to be streamed back into streaming_merge function again?

Oh, to answer my own question it's because streaming_merge assumes the input is sorted, and to do that we need to materialize the original stream, which also involves disk spill over, got it.

One partial improvement would be something like concatenating and sorting the incoming batches into a single batch inside insert_batch, and then optionally spilling that over to disk; still, this results in a trade-off between time and memory efficiency of such cases, since it would require sorting of each intermediate batch (albeit partially sorted already).

I guess the real solution would be something like an external top-k algorithm which would work on the original input stream and use a min/max heap with spillover in case fetch is some value (actually it was my impression this was already implemented). Not sure how viable is that though.

Yes it only concatenates them once the limit has been (almost) reached, but avoids spilling the data.

Potentially this could be done at an earlier stage (say after processing 1/10 of the memory limit) to keep memory usage down during processing (while also avoiding too much overhead in terms of copying).

As well there is indeed no real TopK implementation yet, so as @alamb suggests a SortExec(fetch=1) could also be faster / more efficient with a real TopK (min/max heap) implementation.

[Dandandan]

Oh I think I see the problem: the fetch count is never utilized while the sort is accumulating batches: https://github.com/apache/arrow-datafusion/blob/main/datafusion/core/src/physical_plan/sorts/sort.rs#L628-L632

Consequently, all the batches get loaded into memory inside of the ExternalSorter:

diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs
index f660f0acf..6eb63e39a 100644
--- a/datafusion/core/src/physical_plan/sorts/sort.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -159,6 +159,7 @@ impl ExternalSorter {
             }
         }

+        debug!("Inserted batch of size {} for a total of {} in-memory batches", input.num_rows(), self.in_mem_batches.len() + 1);
         self.in_mem_batches.push(input);
         self.in_mem_batches_sorted = false;
         Ok(())

[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::limit] Start GlobalLimitExec::execute for partition: 0
[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::sorts::sort] Start SortExec::execute for partition 0 of context session_id 8855be36-07c1-495b-b6fd-9cb9501cb46d and task_id None
[2023-08-01T08:13:27Z TRACE datafusion::physical_plan::sorts::sort] End SortExec's input.execute for partition: 0
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 1 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 2 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 3 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 4 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 8192 for a total of 5 in-memory batches
[2023-08-01T08:13:28Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 6 in-memory batches
...
[2023-08-01T08:14:12Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 209 for a total of 710 in-memory batches
[2023-08-01T08:14:12Z DEBUG datafusion::physical_plan::sorts::sort] Inserted batch of size 6836 for a total of 711 in-memory batches

This raises a (maybe naive) related question: why bother with materializing the record batches from the input stream and potentially spilling them to the disk at this point, only for them to be streamed back into streaming_merge function again?

EDIT: Alternatively, the following diff fixes the issue:

diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs
index f660f0acf..49e44a834 100644
--- a/datafusion/core/src/physical_plan/sorts/sort.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -133,11 +133,17 @@ impl ExternalSorter {
     /// Appends an unsorted [`RecordBatch`] to `in_mem_batches`
     ///
     /// Updates memory usage metrics, and possibly triggers spilling to disk
-    async fn insert_batch(&mut self, input: RecordBatch) -> Result<()> {
+    async fn insert_batch(&mut self, mut input: RecordBatch) -> Result<()> {
         if input.num_rows() == 0 {
             return Ok(());
         }

+        if self.fetch.map_or(false, |f| f < input.num_rows()) {
+            // Sort the batch pre-emptively to reduce the number of rows
+            // since the fetch count is smaller than the batch size
+            input = sort_batch(&input, &self.expr, self.fetch)?;
+        }
+
         let size = batch_byte_size(&input);
         if self.reservation.try_grow(size).is_err() {
             let before = self.reservation.size();

Not sure about the overall performance implications of this tweak however (i.e. it may impact the runtime for some cases I guess). Also the value of this tweak is progressively reduced the larger the fetch size is, and when fetch size reaches the batch size, we're back to the underlying problem.

Ah, this might be a good improvement. I thought it did it like this before, but this might have been changed more recently. I think doing this should be as fast (as long as we avoid to perform sorting twice).

[gruuya]

Ah, this might be a good improvement. I thought it did it like this before, but this might have been changed more recently. I think doing this should be as fast (as long as we avoid to perform sorting twice).

So I tried to do a first pass along these lines, you can see the results here: #7180

Besides the memory the query times are actually slightly better than before, however it's only draft atm since I think there are some sub-optimal cases. For instance, previously the sorting was done in parallel via #6308 which is not the case here, and relatedly in the absence of a fetch parameter the algorithm should probably default to the previous approach.

[alamb]

Thanks @gruuya -- I'll take a look hopefully later today.

In general, perhaps it is time to make a real TopK operator in DataFusion (we are also interested in it at InfluxData)

Interestingly, I looked through the tickets and I didn't see any that specifically described TopK

https://github.com/apache/arrow-datafusion/issues?q=is%3Aissue+is%3Aopen+topk

There is a more general "TopK per partition" described in #6899

I will try and file some thoughts on TopK later this week (perhaps tomorrow)

cc @avantgardnerio who I think may also be interested in this discussion