Spark: Doing a Coalesce and foreachpartitions in spark directly on an iceberg table is leaking memory heavy iterators

[jkolash]

Apache Iceberg version

1.5.0

Query engine

Spark

Please describe the bug 🐞

Summary

Doing the following should not leak any significant amount of memory.

sparkSession.sql("select * from icebergcatalog.db.table").coalesce(4).foreachPartition( (iterator) -> {
 while (iterator.hasNext()) iterator.next();
});

A workaround is to use repartition() instead however this requires more resources to handle spilling shuffling etc..

Spark version: Spark 3.4.X

Details

The Below code can be run on a sufficiently large iceberg table.

  static AtomicInteger partitionCounter = new AtomicInteger(0);

    static void reproduceBug(SparkSession sparkSession, String table) {

        sparkSession.sql("select * from "+table).coalesce(4).foreachPartition( (iterator) -> {
            int partition = partitionCounter.getAndIncrement();
            AtomicLong rowCounter = ThreadLocal.withInitial(() -> new AtomicLong(0)).get();

            while (iterator.hasNext()) {
                iterator.next();
                if (rowCounter.getAndIncrement() % 100000 == 0) {
                    System.out.println(partition + " " + rowCounter.get());
                }
            }
        });
    }

The following image is me running the reproduceBug method over sufficiently large table that we have in our environment with ~500 columns.

The following image shows the "Dominators" report in VisualVM org.apache.spark.TaskContextImpl

Digging deeper we see that the onCallbacks is keeping an anonymous class inside org.apache.spark.sql.execution.datasources.v2.DataSourceRDD and that is holding a reference to org.apache.iceberg.spark.source.RowDataReader

I believe this callback is added here https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala#L90
We also see the org.apache.iceberg.util.Filter iterator holding a heavy reference.

Exploring the problem

Is this inherently a bug in org.apache.spark.sql.execution.datasources.v2.DataSourceRDD? Or should iterators not hold onto state no longer needed once advanced to the end? Is the iterator even exhausted? Once an iterator is exhausted there is no longer a need for referencing. However this kind of breaks the concept of a CloseableIterator which has an explicit close vs an implicit close where you could detect hasNext() is false and auto-close. then even ignore a duplicate close() as it was handled by an implicit close() of iterator exhaustion. I believe an iterator accumulating hundreds of megabytes of state kind of breaks the implicit "expected" contract of an iterator being a streaming set of objects. There might even be a distinction between closing and simply holding onto large object references.

Digging deeper I see items = org.apache.iceberg.parquet.ParquetReader$FileIterator#6 holding onto a model reference. It might be possible to null out the model references when hasNext() is false.

   @Override
    public boolean hasNext() {
      boolean hasNext = valuesRead < totalValues; 
      if (!hasNext) {
        this.model = null;
      }
      return hasNext;
    }

Willingness to contribute

• I can contribute a fix for this bug independently
• I would be willing to contribute a fix for this bug with guidance from the Iceberg community
• I cannot contribute a fix for this bug at this time

[jkolash]

I was able to make some hacky changes that reduced memory usage in this draft PR. #13298 mainly to show that these were the critical objects that needed to be GC'd that were no longer needed.

I also tested

   String[] paths = sparkSession.sql("select file_path from "+table+".files").collectAsList().stream().map( row -> row.getString(0)).toArray(String[]::new);
        System.out.println(Arrays.asList(paths));
        sparkSession.read().load(paths).coalesce(4).foreachPartition((iterator) -> {
                    int partition = partitionCounter.getAndIncrement();
                    AtomicLong rowCounter = ThreadLocal.withInitial(() -> new AtomicLong(0)).get();

                    while (iterator.hasNext()) {
                        iterator.next();
                        if (rowCounter.getAndIncrement() % 100000 == 0) {
                            System.out.println(partition + " " + rowCounter.get());
                        }
                    }
                }
        );

To show that it is the new read path through iceberg that is the issue and when loading/processing the parquet files directly the issue doesn't manifest.

The TaskContextImpl still has the onCompleteCalllbacks but the FileScanRDD is much smaller

[RussellSpitzer]

Can your force V2 Sources and try the Parquet version again (FileScanRDD is a different code path)? It feels odd to me that the iterator we are making should null itself out when fully iterated. Feels like that should be on the last "next" call as is currently in the code or in the close method?

Or is the issue here that we are saving this task context for the UI so we don't actually ever drop the iterator reference and close is never called?

[jkolash]

I'm not sure what you mean by force V2 sources and try the Parquet version again. I was comparing both the v1 and v2 memory usage and the v2 path using iceberg consumes significantly more memory.

Is there a v2 equivalent to

sparkSession.read().load(paths)

Speaking of should be handled on the last next() call. I don't see that in my experimental changes which did reduce memory usage.
https://github.com/apache/iceberg/pull/13298/files#diff-d80c15b3e5376265436aeab8b79d5a92fb629c6b81f58ad10a11b9b9d3bfcffcR134

[jkolash]

Or is the issue here that we are saving this task context for the UI so we don't actually ever drop the iterator reference and close is never called?

The close() will not likely be called until the job completes. the call back being held onto is what will do the close() operation and that wont likely happen until the job completes.

          context.addTaskCompletionListener[Unit] { _ =>
            // In case of early stopping before consuming the entire iterator,
            // we need to do one more metric update at the end of the task.
            CustomMetrics
              .updateMetrics(reader.currentMetricsValues.toImmutableArraySeq, customMetrics)
            iter.forceUpdateMetrics()
            reader.close()
          }

https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala#L90

[jkolash]

Let me know if you want me to write something to produce synthetic data for this so it can be reproduced without using real data. That way you can reproduce it on your environment.

[RussellSpitzer]

I'm not sure what you mean by force V2 sources and try the Parquet version again. I was comparing both the v1 and v2 memory usage and the v2 path using iceberg consumes significantly more memory.

Is there a v2 equivalent to

sparkSession.read().load(paths)
Speaking of should be handled on the last next() call. I don't see that in my experimental changes which did reduce memory usage. https://github.com/apache/iceberg/pull/13298/files#diff-d80c15b3e5376265436aeab8b79d5a92fb629c6b81f58ad10a11b9b9d3bfcffcR134

https://github.com/apache/iceberg/pull/13298/files#diff-7f0d08eb9c160db9e4fefa082993c58f2fe749f9be4f108c26cb4201547f2521L139

[RussellSpitzer]

Ok - for my own recap here

1. Spark is holding onto task contexts in order to invoke callbacks to get metric information till the end of the job
2. Task Context through this callback holds onto the iterator the DSV2 source makes
3. Our iterator ends up being slightly heavy, it holds onto a reader for every column in the table which means for a very wide table we end up holding onto a parquet reader for every column in the table. This ends up adding up to a significant amount of memory.

Now I think there are a few avenues to make this better

1. As @jkolash suggests we can preemptively null out the reference to the readers when the iterator has been exhausted. One issue here is that a partially consumed iterator will still hold onto the model until the job has completed.

2. We can work on the Spark side to make sure that the iterator is actually released when exhausted or the task has completed.

Thoughts @szehon-ho + @aokolnychyi ?

[szehon-ho]

I dont see immediately a spark api to easily remove the callback once an iter is exhausted.

Maybe the iceberg change is easier. but if you have some idea on spark side, it make sense to explore?

[jkolash]

I wanted to bring up another point that @RussellSpitzer helped me identify which was this issue reproduces in a similar way with the v2 Datasources

when I Ran #13297 (comment)
but set the spark option spark.sql.sources.useV1SourceList to an empty string it reproduced via the new v2 datasource code, but not the v1 datasource.

So this issue also needs to be addressed in spark as well.

On the spark side I think the close() call likely should be done on task completion not via a callback on job completion?, and the metrics can be collected before that close?

[jkolash]

The issue is that coalesce is combining N tasks into 1 task but really there are N tasks, so coalesce should be more subtask aware?

[jkolash]

Reading more deeply in DataSourceRDD.scala it seems the callback is there due to a partially consumed iterator, but spark wraps the raw iterator using the MetricsBatchIterator that could detect whether the iterator was exhausted via hasNext() == false and remove the onComplete callback?

[jkolash]

https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/TaskContextImpl.scala#L71 is currently a stack so that would need to change to something more like a linked list to alllow removing arbitrary listeners.

[jkolash]

Actually I don't think any spark api changes are needed the callback and the MetricsBatchIterator are both created in the DataSourceRDD

In java you could write something like this.

   static class HeavyCallback implements Runnable{
        byte[] heavyReference = new byte[1000000];
        @Override
        public void run() {
            
        }
    }
    
    static class InvokeOnceCallback implements Runnable {
        final AtomicReference<Runnable> callbackHandle = new AtomicReference<>();
        
        InvokeOnceCallback(Runnable target) {
            callbackHandle.set(target);
        }
        
        @Override
        public void run() {
            if (callbackHandle.get() != null) {
                callbackHandle.get().run();
                callbackHandle.set(null);
            }
        }
    }
    
    static void showcaseIndirectCallback() {
        InvokeOnceCallback callback = new InvokeOnceCallback(new HeavyCallback());

        //this can be called from either the iterable exhaustion or the original callback.
        //Once invoked the reference is removed and can be GC'd
        callback.run(); 
    }

[emmanuel099]

Hi, my team and I had a similar problem while querying and optimizing tables with a huge number of data and delete files. Spark always ran OOM, because resources (Parquet readers, ...) weren't freed properly. I did some debugging and identified CloseableIterable as the culprit. Unfortunately, I never had time to upstream it properly.

CloseableIterable has the following two problems:

1. When turning a CloseableIterable into an iterator, the returned CloseableIterator doesn't properly close the iterable, thereby causing a resource leak. By returning a custom instance of CloseableIterator which keeps track of the closable iterable, we properly close the iterable when the iterator gets closed.
2. When combining an iterable with a closeable using CloseableIterable.combine we have a resource leak if the iterable implements CloseableIterable because we never call close on the iterable. By overloading combine for CloseableIterable we can ensure that we close the iterable in addition to the closeable.

With the following patch applied, the queries and table optimizations ran without OOM and much lower memory limits.

From 6e2cfd71c2dcecf432f2db2119ac530674af5042 Mon Sep 17 00:00:00 2001
From: Emmanuel Pescosta <emmanuelpescosta099@gmail.com>
Date: Thu, 16 Jan 2025 12:19:46 +0100
Subject: [PATCH] fix(iceberg.io): Properly close CloseableIterable when
 turning it into an iterator

A CloseableIterator created from a CloseableIterable didn't close
the iterable when the iterator is closed. This may cause a resource
leak.
---
 .../apache/iceberg/io/CloseableIterable.java  | 40 ++++++++++++++++++-
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java b/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
index 06323612a..240c1ef8f 100644
--- a/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
+++ b/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
@@ -51,7 +51,7 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
 
       @Override
       public CloseableIterator<E> iterator() {
-        return CloseableIterator.withClose(iterable.iterator());
+        return closeableIteratorOf(iterable, this);
       }
     };
   }
@@ -69,7 +69,43 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
 
       @Override
       public CloseableIterator<E> iterator() {
-        return CloseableIterator.withClose(iterable.iterator());
+        return closeableIteratorOf(iterable, this);
+      }
+    };
+  }
+
+  static <E> CloseableIterable<E> combine(CloseableIterable<E> iterable, Closeable closeable) {
+    return new CloseableIterable<E>() {
+      @Override
+      public void close() throws IOException {
+        closeable.close();
+        iterable.close();
+      }
+
+      @Override
+      public CloseableIterator<E> iterator() {
+        return closeableIteratorOf(iterable, this);
+      }
+    };
+  }
+
+  private static <E> CloseableIterator<E> closeableIteratorOf(
+      Iterable<E> iterable, Closeable closeable) {
+    Iterator<E> iterator = iterable.iterator();
+    return new CloseableIterator<E>() {
+      @Override
+      public void close() throws IOException {
+        closeable.close();
+      }
+
+      @Override
+      public boolean hasNext() {
+        return iterator.hasNext();
+      }
+
+      @Override
+      public E next() {
+        return iterator.next();
       }
     };
   }
-- 
2.49.0

@jkolash Can you please check if this would fix your problem?