`SortMergeJoinExec` fails to allocate memory but should spill instead

[andygrove]

Describe the bug

In some configurations/environments, I see queries fail due to memory pool requests being rejected, but I would expect Comet to spill to disk instead.

In one example, I am running TPC-H @ SF=1000 (1TB) in k8s. I am specifying spark.comet.exec.replaceSortMergeJoin=false to force the use of CometSortMergeJoinExec.

    --conf spark.executor.instances=4 \
    --conf spark.executor.cores=8 \
    --conf spark.executor.memory=8G \
    --conf spark.memory.offHeap.enabled=true \
    --conf spark.memory.offHeap.size=4g \

I allocated 4 GB of off-heap memory, which equates to 512 MB per core.

I saw memory requests fail with the memory pool limit at ~512MB.

I then doubled the off-heap memory, and I still see the same issue; however, the pool is now ~1GB. I would expect spilling to kick in instead.

org.apache.comet.CometNativeException: Additional allocation failed with top memory consumers (across reservations) as:
  ExternalSorter[107]#2991(can spill: true) consumed 1024.2 MB,
  ExternalSorterMerge[107]#2990(can spill: false) consumed 16.7 MB,
  GroupedHashAggregateStream[107] ()#2994(can spill: true) consumed 0.0 B,
  GroupedHashAggregateStream[107] ()#2995(can spill: true) consumed 0.0 B,
  ExternalSorterMerge[107]#2992(can spill: false) consumed 0.0 B,

I also see pods being killed due to OOM:

NAME                                                        READY   STATUS              RESTARTS   AGE
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-1   0/1     OOMKilled           0          11m
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-3   0/1     OOMKilled           0          11m
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-4   0/1     OOMKilled           0          11m
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-5   1/1     Running             0          4s
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-6   1/1     Running             0          4s
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-7   1/1     Running             0          3s
comet-benchmark-derived-from-tpch-a1133f997d91851b-exec-8   0/1     ContainerCreating   0          1s

I also see errors in the executor logs:

25/09/24 21:24:57 WARN ExecutionMemoryPool: Internal error: release called on 917504 bytes but task only has 0 bytes of memory from the off-heap execution pool
25/09/24 21:24:57 WARN ExecutionMemoryPool: Internal error: release called on 839664 bytes but task only has 0 bytes of memory from the off-heap execution pool
25/09/24 21:24:57 WARN ExecutionMemoryPool: Internal error: release called on 917504 bytes but task only has 0 bytes of memory from the off-heap execution pool

Some related Spark logging:

25/09/24 21:27:20 INFO TaskMemoryManager: Memory used in task 20632
25/09/24 21:27:20 INFO TaskMemoryManager: 1073741824 bytes of memory were used by task 20632 but are not associated with specific consumers
25/09/24 21:27:20 INFO TaskMemoryManager: 5677121232 bytes of memory are used for execution and 962961 bytes of memory are used for storage

Steps to reproduce

No response

Expected behavior

No response

Additional context

No response

[andygrove]

Adding notes for my own reference later. Spark run completes in 3486 seconds:

--conf spark.driver.memory=8G \
--conf spark.executor.instances=6 \
--conf spark.executor.cores=8 \
--conf spark.executor.memory=16G \

The following Comet run (replaceSortMergeJoin=false) started failing with OOM on q9 after 16 mins.

--conf spark.driver.memory=8G \
--conf spark.executor.instances=4 \
--conf spark.executor.cores=8 \
--conf spark.executor.memory=8G \
--conf spark.memory.offHeap.enabled=true \
--conf spark.memory.offHeap.size=32g \
--conf spark.comet.exec.replaceSortMergeJoin=false

The following Comet run (replaceSortMergeJoin=true) completed in 2390 seconds. There may have been some executor restarts though.

--conf spark.executor.instances=4 \
--conf spark.executor.cores=8 \
--conf spark.executor.memory=16G \
--conf spark.memory.offHeap.enabled=true \
--conf spark.memory.offHeap.size=24g \
--conf spark.comet.exec.replaceSortMergeJoin=true

[comphead]

Thanks @andygrove might be related to apache/datafusion#11218

Some ongoing spilling work also apache/datafusion#17593

[andygrove]

I improved the error message from DataFusion in apache/datafusion#17870 and now see that the failed allocation is from ExternalSorterMerge.

try_grow failed for ExternalSorterMerge[53]: Additional allocation failed with top memory consumers (across reservations) as:
  ExternalSorter[53]#1569(can spill: true) consumed 262.8 MB, peak 271.9 MB,
  ExternalSorterMerge[53]#1570(can spill: false) consumed 5.5 MB, peak 10.0 MB,
  ExternalSorterMerge[53]#1568(can spill: false) consumed 5.0 MB, peak 10.0 MB,
  SMJStream[53]#1571(can spill: false) consumed 0.0 B, peak 0.0 B,
  GroupedHashAggregateStream[53] ()#1572(can spill: true) consumed 0.0 B, peak 0.0 B,
  GroupedHashAggregateStream[53] ()#1573(can spill: true) consumed 0.0 B, peak 0.0 B.

[andygrove]

org.apache.comet.CometNativeException: Additional allocation failed for ExternalSorterMerge[9] with top memory consumers (across reservations) as:
  ExternalSorter[9]#1170(can spill: true) consumed 250.2 MB, peak 266.0 MB,
  ExternalSorterMerge[9]#1169(can spill: false) consumed 5.0 MB, peak 10.0 MB,
  ExternalSorterMerge[9]#1171(can spill: false) consumed 866.9 KB, peak 10.0 MB,
  SMJStream[9]#1172(can spill: false) consumed 0.0 B, peak 0.0 B,
  GroupedHashAggregateStream[9] ()#1173(can spill: true) consumed 0.0 B, peak 0.0 B,
  GroupedHashAggregateStream[9] ()#1174(can spill: true) consumed 0.0 B, peak 0.0 B.
Error: Task 1008 failed to acquire 327680 bytes, only got 536. Reserved: 268466040

backtrace:    0: <comet::execution::memory_pools::unified_pool::CometUnifiedMemoryPool as datafusion_execution::memory_pool::MemoryPool>::try_grow
   1: <datafusion_execution::memory_pool::pool::TrackConsumersPool<I> as datafusion_execution::memory_pool::MemoryPool>::try_grow
   2: datafusion_execution::memory_pool::MemoryReservation::try_grow
   3: datafusion_physical_plan::sorts::builder::BatchBuilder::push_batch
   4: datafusion_physical_plan::sorts::merge::SortPreservingMergeStream<C>::maybe_poll_stream
   5: <datafusion_physical_plan::sorts::merge::SortPreservingMergeStream<C> as futures_core::stream::Stream>::poll_next
   6: datafusion_physical_plan::sorts::sort::ExternalSorter::sort_and_spill_in_mem_batches::{{closure}}
   7: <datafusion_physical_plan::stream::RecordBatchStreamAdapter<S> as futures_core::stream::Stream>::poll_next
   8: <comet::execution::operators::copy::CopyStream as futures_core::stream::Stream>::poll_next
   9: <datafusion_physical_plan::joins::sort_merge_join::stream::SortMergeJoinStream as futures_core::stream::Stream>::poll_next
  10: <datafusion_physical_plan::projection::ProjectionStream as futures_core::stream::Stream>::poll_next
  11: <datafusion_physical_plan::aggregates::row_hash::GroupedHashAggregateStream as futures_core::stream::Stream>::poll_next
  12: <datafusion_physical_plan::aggregates::row_hash::GroupedHashAggregateStream as futures_core::stream::Stream>::poll_next
  13: <datafusion_physical_plan::projection::ProjectionStream as futures_core::stream::Stream>::poll_next
  14: comet::execution::jni_api::Java_org_apache_comet_Native_executePlan::{{closure}}::{{closure}}
  15: Java_org_apache_comet_Native_executePlan

[comphead]

  ExternalSorterMerge[9]#1169(can spill: false) consumed 5.0 MB, peak 10.0 MB,
  ExternalSorterMerge[9]#1171(can spill: false) consumed 866.9 KB, peak 10.0 MB,
  SMJStream[9]#1172(can spill: false) consumed 0.0 B, peak 0.0 B,

spill - false 🤔 to enable spilling it should be passed the DiskManager which supports spilling

[parthchandra]

The peak memory consumption reported for the top 6 operators is nowhere close to the amount of configured memory.

[andygrove]

The peak memory consumption reported for the top 6 operators is nowhere close to the amount of configured memory.

The recent stack traces are using some intentionally low memory settings, so that I can reproduce more conveniently.

$SPARK_HOME/bin/spark-submit \
    --master $SPARK_MASTER \
    --jars $COMET_JAR \
    --driver-class-path $COMET_JAR \
    --conf spark.driver.memory=8G \
    --conf spark.executor.instances=1 \
    --conf spark.executor.cores=8 \
    --conf spark.cores.max=8 \
    --conf spark.executor.memory=16g \
    --conf spark.memory.offHeap.enabled=true \
    --conf spark.memory.offHeap.size=2g \
    --conf spark.eventLog.enabled=true \
    --conf spark.driver.extraClassPath=$COMET_JAR \
    --conf spark.executor.extraClassPath=$COMET_JAR \
    --conf spark.plugins=org.apache.spark.CometPlugin \
    --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
    --conf spark.comet.exec.replaceSortMergeJoin=false \
    --conf spark.comet.expression.allowIncompatible=true \
    --conf spark.hadoop.fs.s3a.impl=org.apache.hadoop.fs.s3a.S3AFileSystem \
    --conf spark.hadoop.fs.s3a.aws.credentials.provider=com.amazonaws.auth.DefaultAWSCredentialsProviderChain \
    tpcbench.py \
    --name comet \
    --benchmark tpch \
    --data $TPCH_DATA \
    --queries $TPCH_QUERIES \
    --output . \
    --iterations 1

[andygrove]

@comphead @parthchandra Here is what I know so far (which is not very much. I am just getting started with understanding DataFusion's code in this area).

sort_and_spill_in_mem_batches is being called, presumably as a result of not being able to allocate extra memory.

sort_and_spill_in_mem_batches then has a loop which polls for batches and then also tries to allocate memory but has fallbacks if it cannot get more memory:

        while let Some(batch) = sorted_stream.next().await {
            let batch = batch?;
            let sorted_size = get_reserved_byte_for_record_batch(&batch);
            if self.reservation.try_grow(sorted_size).is_err() {
                // Although the reservation is not enough, the batch is
                // already in memory, so it's okay to combine it with previously
                // sorted batches, and spill together.
                globally_sorted_batches.push(batch);
                self.consume_and_spill_append(&mut globally_sorted_batches)
                    .await?; // reservation is freed in spill()
            } else {
                globally_sorted_batches.push(batch);
            }
        }

Unfortunately, the call to sorted_stream.next().await also allocates memory, because it is calling SortPreservingMergeStream and its poll method will call BatchBuilder::push_batch:

Poll::Ready(self.in_progress.push_batch(idx, batch))

push_batch fails to allocate memory here:

    pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> {
        self.reservation
            .try_grow(get_record_batch_memory_size(&batch))?;

[comphead]

Yes, you are right @andygrove the way spilling done in sorter and SMJ is:

• try reserve memory
• if failed and spilling supported then spill

However we need to keep in mind that it makes sense for the specific memory pool.
So the large memory is already in the memory and spilling works later and for pool only and if the pool == available memory the spilling doesn't make sense IMO.

One idea is to configure the pool individually so spilling would start sooner

[parthchandra]

One idea is to configure the pool individually

There is more than one memory pool? Are these blocks of memory from the memory pool that are reserved for an operator?

so spilling would start sooner

That seems to be the way to do it.

From what I understand, the issue is the spilling itself requires some memory and we have already used all the available memory when we start/try to spill?