You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/30 09:28:01 UTC
(datafusion-comet) 01/01: fix: release input Arrow streams that native never takes (#6372)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6372-f375c4c345a61c9a41c240f527351b076cbdfffb
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 59be74a91a5b909bdb3beceae06c6b86caab2420
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Sep 30 09:27:42 2026 +0000
fix: release input Arrow streams that native never takes (#6372)
Native releases an exported input ArrowArrayStream only after
create_plan takes it on the plan's first executePlan. The task-completion
listener that CometArrowStream.stream registers called only close(),
which frees the C struct without running the release callback. So a
stream native never took (createPlan failed, the iterator was never
polled, the task failed before the first executePlan, or planning failed
before reaching that input) kept its ArrowReader open, and arrow-java's
JNI global ref kept it reachable for the life of the executor.
Call release() before close() in that listener. It is a no-op on a
stream native took, because from_raw leaves a null release callback in
the JVM struct and arrow-java skips a null callback. With the listener
covering every path, drop the release loop in CometExecIterator's
construction failure handler, which covered only a failure after
createPlan succeeded.
---
.../scala/org/apache/comet/CometExecIterator.scala | 14 +---
.../execution/arrow/CometNativeArrowSource.scala | 10 +++
.../comet/exec/CometNativeShuffleSuite.scala | 14 ++--
.../spark/CometExecIteratorLifecycleSuite.scala | 83 +++++++++++++++++++++-
4 files changed, 104 insertions(+), 17 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index 7dab2dab4f..5c9cada361 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -26,7 +26,6 @@ import java.util.concurrent.atomic.AtomicBoolean
import scala.util.control.NonFatal
-import org.apache.arrow.c.ArrowArrayStream
import org.apache.arrow.memory.BufferAllocator
import org.apache.hadoop.conf.Configuration
import org.apache.spark._
@@ -161,16 +160,9 @@ class CometExecIterator(
case closeFailure: Throwable => failure.addSuppressed(closeFailure)
}
- // Native only takes ownership of Arrow streams during the first executePlan call.
- inputObjects.foreach {
- case stream: ArrowArrayStream =>
- try {
- stream.release()
- } catch {
- case releaseFailure: Throwable => failure.addSuppressed(releaseFailure)
- }
- case _ =>
- }
+ // Native has not taken the Arrow input streams yet (it does on the first executePlan), so
+ // the task-completion listener that `CometArrowStream.stream` registered for each one
+ // releases it.
shuffleBlockIterators.values.foreach { iterator =>
try {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index a60e2811ec..19532cefa1 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -225,6 +225,14 @@ object CometArrowStream extends Logging {
* (Spark fires listeners in reverse registration order, and the listener that drops the native
* plan is registered later by `CometExecIterator`), so `allocator.close` finds zero outstanding
* bytes.
+ *
+ * Native takes the stream on the plan's first `executePlan`. The listener releases a stream
+ * native never took (plan creation failed, the consumer was never polled, the task failed
+ * first, or planning failed before reaching this input) itself: `ArrowArrayStream.close` alone
+ * frees only the C struct, leaving the reader open and pinned for the life of the executor by
+ * the JNI global ref that arrow-java holds on the stream's private data. Taking the stream
+ * leaves a null release callback in the JVM's struct, and arrow-java's `release` skips a null
+ * callback, so a stream native took is still released only once.
*/
def stream(
name: String,
@@ -249,6 +257,8 @@ object CometArrowStream extends Logging {
if (context != null) {
val streamRef = arrowStream
context.addTaskCompletionListener[Unit] { _ =>
+ // Release before close: `close` frees the struct that `release` reads.
+ streamRef.release()
streamRef.close()
allocator.close()
}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index 66e5ccbce9..cbacf4368a 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -34,7 +34,7 @@ import org.apache.arrow.memory.ArrowBuf
import org.apache.arrow.vector.ipc.ArrowReader
import org.apache.arrow.vector.types.pojo.{Field, Schema}
import org.apache.hadoop.fs.Path
-import org.apache.spark.SparkEnv
+import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row}
import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
@@ -280,12 +280,18 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
test(
"failed RSS callback registration closes every Arrow and shuffle input and preserves its error") {
val planBytes = rssShufflePlanBytes
+ val arrowInputsClosed = spark.sparkContext.longAccumulator("arrowInputsClosed")
val results = spark.sparkContext
.parallelize(Seq(17), 1)
.mapPartitions { _ =>
var readerClosed = false
var ownedBuffer: ArrowBuf = null
+ // The task-completion listener that exports the Arrow input releases it. Registered
+ // before that one, this listener runs after it.
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+ if (readerClosed && ownedBuffer.refCnt() == 0) arrowInputsClosed.add(1)
+ }
val arrowInput = CometArrowStream
.stream(
"native-rss-registration-failure-test",
@@ -333,13 +339,12 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
shuffleBlockIterators = shuffleInputs,
shufflePartitionPusher = Some(null))
iterator.close()
- Iterator.single((false, false, false, false))
+ Iterator.single((false, false, false))
} catch {
case failure: Throwable =>
Iterator.single(
(
failure.getMessage.contains("Remote shuffle callback must not be null"),
- readerClosed && ownedBuffer.refCnt() == 0,
closedInputs.toSet == Set("failing", "remaining"),
failure.getSuppressed.exists(
_.getMessage.contains("shuffle input cleanup sentinel"))))
@@ -347,7 +352,8 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
}
.collect()
- assert(results.sameElements(Array((true, true, true, true))))
+ assert(results.sameElements(Array((true, true, true))))
+ assert(arrowInputsClosed.value == 1, "the Arrow input was not released by task end")
}
test("native shuffle plan preserves local partition writer and legacy output path") {
diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
index 59c1231ef2..3bbca8f409 100644
--- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
@@ -22,11 +22,14 @@ package org.apache.spark
import java.io.ByteArrayInputStream
import java.lang.ref.WeakReference
import java.util.Properties
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
import scala.reflect.ClassTag
+import org.apache.arrow.c.ArrowArrayStream
import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.ipc.ArrowReader
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
import org.apache.spark.executor.TaskMetrics
import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.rdd.RDD
@@ -37,7 +40,7 @@ import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
import org.apache.spark.sql.functions.{col, udf}
import org.apache.spark.sql.types.{LongType, StructField, StructType}
-import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native}
+import org.apache.comet.{CometConf, CometExecIterator, CometNativeException, CometShuffleBlockIterator, Native}
import org.apache.comet.serde.Config.ConfigMap
import org.apache.comet.serde.OperatorOuterClass
@@ -595,6 +598,82 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
assert(memoryUsageLogInterval(Some("false")) == 0L)
assert(memoryUsageLogInterval(Some("10 seconds please")) == 0L)
}
+
+ /**
+ * Exports an empty input stream the way a native plan's JVM input is exported, over a reader
+ * that counts how many times it is closed. Only the stream's release callback closes it.
+ */
+ private def closeCountingInput(closes: AtomicInteger): ArrowArrayStream =
+ CometArrowStream
+ .stream(
+ "input-release-test",
+ allocator =>
+ new ArrowReader(allocator) {
+ override protected def readSchema(): Schema =
+ new Schema(
+ java.util.Collections
+ .singletonList(Field.nullable("test", new ArrowType.Int(64, true))))
+
+ override def loadNextBatch(): Boolean = false
+
+ override def bytesRead(): Long = 0L
+
+ override protected def closeReadSource(): Unit = closes.incrementAndGet()
+ })
+ .next()
+
+ private def iteratorOver(input: ArrowArrayStream, planBytes: Array[Byte]): CometExecIterator =
+ new CometExecIterator(
+ id = CometExec.newIterId,
+ inputObjects = Array[Object](input),
+ numOutputCols = 1,
+ protobufQueryPlan = planBytes,
+ nativeMetrics = CometMetricNode(Map.empty),
+ numParts = 1,
+ partitionIndex = 0)
+
+ private def limitPlanBytes: Array[Byte] =
+ CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get.toByteArray
+
+ test("an input stream native took is released once, when its plan is released") {
+ withTaskContext(4600000L) {
+ val closes = new AtomicInteger()
+ val iter = iteratorOver(closeCountingInput(closes), limitPlanBytes)
+ // The first executePlan takes the stream, and exhausting the empty input closes the
+ // iterator, whose releasePlan drops the native reader and so releases the stream.
+ assert(!iter.hasNext)
+ assert(closes.get() == 1, s"the input reader closed ${closes.get()} times on plan release")
+ // Taking the stream left a released struct on the JVM side, so the release at task end
+ // must not close the reader a second time.
+ TaskContext.get().markTaskCompleted(None)
+ assert(closes.get() == 1, s"the input reader closed ${closes.get()} times by task end")
+ }
+ }
+
+ test("an input stream native never took is released when the task ends") {
+ withTaskContext(4700000L) {
+ val closes = new AtomicInteger()
+ // Never polled, as when the task is killed before its first hasNext, so native never takes
+ // the stream.
+ iteratorOver(closeCountingInput(closes), limitPlanBytes)
+ TaskContext.get().markTaskCompleted(None)
+ assert(closes.get() == 1, s"the input reader closed ${closes.get()} times")
+ }
+ }
+
+ test("an input stream is released when the task ends after its plan fails to be created") {
+ withTaskContext(4800000L) {
+ val closes = new AtomicInteger()
+ val input = closeCountingInput(closes)
+ val thrown = intercept[CometNativeException](iteratorOver(input, Array[Byte](1, 2, 3)))
+ // Guard against a vacuous pass: createPlan must have failed on the plan bytes.
+ assert(
+ thrown.getMessage.contains("Fail to deserialize to native operator"),
+ s"expected a plan deserialization failure, got: $thrown")
+ TaskContext.get().markTaskCompleted(None)
+ assert(closes.get() == 1, s"the input reader closed ${closes.get()} times")
+ }
+ }
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org