This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c905f40b fix: Fix memory bloat caused by holding too many unclosed
`ArrowReaderIterator`s (#929)
c905f40b is described below
commit c905f40b13b1ebc5525719e93e462dd432c5743c
Author: Kristin Cowalcijk <[email protected]>
AuthorDate: Tue Sep 10 14:10:05 2024 +0800
fix: Fix memory bloat caused by holding too many unclosed
`ArrowReaderIterator`s (#929)
---
.../execution/shuffle/ArrowReaderIterator.scala | 17 +++++++++++----
.../shuffle/CometBlockStoreShuffleReader.scala | 24 +++++++++++-----------
2 files changed, 25 insertions(+), 16 deletions(-)
diff --git
a/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala
b/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala
index 3c0fa153..933e0b66 100644
---
a/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala
+++
b/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala
@@ -31,8 +31,12 @@ class ArrowReaderIterator(channel: ReadableByteChannel,
source: String)
private val reader = StreamReader(channel, source)
private var batch = nextBatch()
private var currentBatch: ColumnarBatch = null
+ private var isClosed: Boolean = false
override def hasNext: Boolean = {
+ if (isClosed) {
+ return false
+ }
if (batch.isDefined) {
return true
}
@@ -42,10 +46,12 @@ class ArrowReaderIterator(channel: ReadableByteChannel,
source: String)
// memory leak.
if (currentBatch != null) {
currentBatch.close()
+ currentBatch = null
}
batch = nextBatch()
if (batch.isEmpty) {
+ close()
return false
}
true
@@ -69,10 +75,13 @@ class ArrowReaderIterator(channel: ReadableByteChannel,
source: String)
def close(): Unit =
synchronized {
- if (currentBatch != null) {
- currentBatch.close()
- currentBatch = null
+ if (!isClosed) {
+ if (currentBatch != null) {
+ currentBatch.close()
+ currentBatch = null
+ }
+ reader.close()
+ isClosed = true
}
- reader.close()
}
}
diff --git
a/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala
b/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala
index 90e0bb14..e026cbeb 100644
---
a/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala
+++
b/common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala
@@ -88,20 +88,20 @@ class CometBlockStoreShuffleReader[K, C](
/** Read the combined key-values for this reduce task */
override def read(): Iterator[Product2[K, C]] = {
+ var currentReadIterator: ArrowReaderIterator = null
+
+ // Closes last read iterator after the task is finished.
+ // We need to close read iterator during iterating input streams,
+ // instead of one callback per read iterator. Otherwise if there are too
many
+ // read iterators, it may blow up the call stack and cause OOM.
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (currentReadIterator != null) {
+ currentReadIterator.close()
+ }
+ }
+
val recordIter = fetchIterator
.flatMap { case (_, inputStream) =>
- var currentReadIterator: ArrowReaderIterator = null
-
- // Closes last read iterator after the task is finished.
- // We need to close read iterator during iterating input streams,
- // instead of one callback per read iterator. Otherwise if there are
too many
- // read iterators, it may blow up the call stack and cause OOM.
- context.addTaskCompletionListener[Unit] { _ =>
- if (currentReadIterator != null) {
- currentReadIterator.close()
- }
- }
-
IpcInputStreamIterator(inputStream, decompressingNeeded = true,
context)
.flatMap { channel =>
if (currentReadIterator != null) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
