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/10/06 13:51:55 UTC
(datafusion-comet) 01/01: fix: rethrow the exception a JVM input throws instead of a CometNativeException (#6243)
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-6243-070951d2a77fbaa597056d8de48d29812b10f2e5
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 011cce3916cbc99556c66b6368b1ddfaae30e2fe
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Oct 6 13:51:33 2026 +0000
fix: rethrow the exception a JVM input throws instead of a CometNativeException (#6243)
* fix: rethrow the exception a JVM input throws instead of a CometNativeException
A native plan pulls each batch of a JVM input through the Arrow C stream
exported by CometArrowStream.stream. When the JVM producer throws while
building a batch, Arrow Java's exported stream catches the throwable and
hands native only its text, so the task failed with a CometNativeException
("C Data interface error: ...") instead of the exception Spark throws. For
inputs of Arrow-backed batches this happened past the first batch only,
because schema reconciliation reads the first one on the JVM. Row inputs
through CometSparkToColumnarExec hit it on every batch.
CometArrowStream.stream now wraps each exported reader to keep the first
throwable its loadNextBatch throws, keyed by the exported stream until the
task completes. When executePlan fails, CometExecIterator rethrows that
throwable if one of its input streams recorded one.
Closes #6234.
* fix: drop unneeded string interpolator flagged by scalafix
* refactor: match the recorded input failure once, share the failed-write test helper
CometExecIterator looked up the recorded input failure twice, in the guard
and in the body of the catch case. An extractor does it once, and the
comment at the catch site shrinks to a line now that
CometArrowStream.inputFailure carries the explanation.
The native-versus-JVM failed-write scaffolding in CometIcebergWriteActionSuite
was copied three times. failedWrites runs a write against a native-writer
table and a JVM-writer table, checks that the native run went through
CometIcebergWriteExec, and returns both failures.
---
docs/source/contributor-guide/ffi.md | 10 +++
.../scala/org/apache/comet/CometExecIterator.scala | 17 ++++
.../execution/arrow/CometNativeArrowSource.scala | 66 ++++++++++++++-
.../comet/CometIcebergWriteActionSuite.scala | 97 ++++++++++++++++------
.../org/apache/comet/exec/CometExecSuite.scala | 51 +++++++++++-
.../spark/CometExecIteratorLifecycleSuite.scala | 55 +++++++++++-
6 files changed, 267 insertions(+), 29 deletions(-)
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index 7027080bbf..36328468f3 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -119,6 +119,16 @@ realigns under-aligned buffers on import ([apache/arrow-rs#10030](https://github
the `realigns_under_aligned_decimal128` test in `scan.rs` guards against an arrow downgrade that would bring back the
panic ([apache/arrow-rs#10028](https://github.com/apache/arrow-rs/issues/10028)).
+### Errors from the JVM Producer
+
+Arrow Java's exported stream catches whatever the reader throws in `get_next` and hands native only its text, so
+native fails the plan with a `CometNativeException` built from that text. `CometArrowStream.stream` wraps every
+reader to keep the throwable itself until the task completes, and `CometExecIterator` rethrows it in place of the
+native error. The task then fails with the exception Spark would have thrown, such as a `SparkArithmeticException`
+from an upstream plan. For an input of Arrow-backed `ColumnarBatch`es, the first batch never takes this path: schema
+reconciliation reads it on the JVM before the stream is exported, so what it throws propagates directly. A test of
+this path therefore has to fail a later batch.
+
### Schema Reconciliation
`CometArrowStream.reconcileStreamSchema` advertises the stream's schema from the actual `CometVector` types in the
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index 22143b4f18..545c01fafd 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -26,6 +26,7 @@ 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._
@@ -33,6 +34,7 @@ import org.apache.spark.broadcast.Broadcast
import org.apache.spark.internal.Logging
import org.apache.spark.network.util.ByteUnit
import org.apache.spark.sql.comet.CometMetricNode
+import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.vectorized._
import org.apache.spark.util.SerializableConfiguration
@@ -224,6 +226,17 @@ class CometExecIterator(
CometExecIterator.startMemoryUsageLog()
+ /**
+ * Matches any throwable once a JVM input of this plan has failed, and extracts what that input
+ * threw. Native saw only its text. See `CometArrowStream.inputFailure`.
+ */
+ private object InputFailure {
+ def unapply(nativeFailure: Throwable): Option[Throwable] =
+ inputObjects.iterator
+ .collect { case stream: ArrowArrayStream => CometArrowStream.inputFailure(stream) }
+ .collectFirst { case Some(failure) => failure }
+ }
+
private def getNextBatch: Option[ColumnarBatch] = {
assert(partitionIndex >= 0 && partitionIndex < numParts)
@@ -246,6 +259,10 @@ class CometExecIterator(
result
} catch {
+ // Native saw only the text of what a JVM input threw, so rethrow the exception itself.
+ case InputFailure(failure) =>
+ throw failure
+
// Handle CometQueryExecutionException with JSON payload first
case e: CometQueryExecutionException =>
logError(s"Native execution for task $taskAttemptId failed", e)
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 19532cefa1..316afe68a2 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
@@ -19,11 +19,15 @@
package org.apache.spark.sql.comet.execution.arrow
+import java.util.concurrent.ConcurrentHashMap
+
import scala.jdk.CollectionConverters._
import scala.reflect.ClassTag
import org.apache.arrow.c.{ArrowArrayStream, Data}
import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.dictionary.Dictionary
import org.apache.arrow.vector.ipc.ArrowReader
import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema}
import org.apache.spark.TaskContext
@@ -233,16 +237,19 @@ object CometArrowStream extends Logging {
* 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.
+ *
+ * The reader is wrapped so that what it throws while native pulls a batch stays available to
+ * [[inputFailure]] until the task completes.
*/
def stream(
name: String,
readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = {
val context = TaskContext.get()
val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue)
- var reader: ArrowReader = null
+ var reader: FailureRecordingReader = null
var arrowStream: ArrowArrayStream = null
try {
- reader = readerFactory(allocator)
+ reader = new FailureRecordingReader(allocator, readerFactory(allocator))
arrowStream = ArrowArrayStream.allocateNew(allocator)
Data.exportArrayStream(allocator, reader, arrowStream)
} catch {
@@ -256,7 +263,9 @@ object CometArrowStream extends Logging {
}
if (context != null) {
val streamRef = arrowStream
+ exportedReaders.put(streamRef, reader)
context.addTaskCompletionListener[Unit] { _ =>
+ exportedReaders.remove(streamRef)
// Release before close: `close` frees the struct that `release` reads.
streamRef.release()
streamRef.close()
@@ -266,6 +275,59 @@ object CometArrowStream extends Logging {
Iterator.single(arrowStream)
}
+ /**
+ * The reader behind each stream exported by [[stream]] whose task has not completed yet, keyed
+ * by the stream object that its consumer holds as an input.
+ */
+ private val exportedReaders = new ConcurrentHashMap[ArrowArrayStream, FailureRecordingReader]()
+
+ /**
+ * What the reader behind `stream` threw while native pulled a batch from it, if anything. Arrow
+ * Java's exported stream catches that throwable and hands native only its text, so the native
+ * plan fails with a `CometNativeException` built from the text. The throwable itself is what
+ * the task would have thrown without Comet, so the plan's consumer rethrows it instead.
+ */
+ def inputFailure(stream: ArrowArrayStream): Option[Throwable] =
+ Option(exportedReaders.get(stream)).flatMap(reader => Option(reader.failure))
+
+ /**
+ * Forwards everything to `reader`, keeping the first throwable its `loadNextBatch` throws for
+ * [[inputFailure]]. Its own `ArrowReader` state is never initialized: every method that the
+ * stream exporter calls is forwarded, and so is `close`.
+ */
+ private final class FailureRecordingReader(allocator: BufferAllocator, reader: ArrowReader)
+ extends ArrowReader(allocator) {
+
+ @volatile var failure: Throwable = _
+
+ override def loadNextBatch(): Boolean =
+ try reader.loadNextBatch()
+ catch {
+ case t: Throwable =>
+ if (failure == null) failure = t
+ throw t
+ }
+
+ override def getVectorSchemaRoot(): VectorSchemaRoot = reader.getVectorSchemaRoot()
+
+ override def getDictionaryVectors(): java.util.Map[java.lang.Long, Dictionary] =
+ reader.getDictionaryVectors()
+
+ override def lookup(id: Long): Dictionary = reader.lookup(id)
+
+ override def getDictionaryIds(): java.util.Set[java.lang.Long] = reader.getDictionaryIds()
+
+ override def bytesRead(): Long = reader.bytesRead()
+
+ override def close(): Unit = reader.close()
+
+ override def close(closeReadSource: Boolean): Unit = reader.close(closeReadSource)
+
+ override protected def closeReadSource(): Unit = ()
+
+ override protected def readSchema(): Schema = reader.getVectorSchemaRoot().getSchema
+ }
+
/** Close a resource, swallowing any error, for use in rollback paths. Null-safe. */
private def closeQuietly(c: AutoCloseable): Unit = {
if (c != null) {
diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
index e13a2ec21a..53388f1b49 100644
--- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
@@ -38,7 +38,7 @@ import org.json4s.jackson.JsonMethods.parse
import org.apache.hadoop.fs.{Path => HadoopPath}
import org.apache.parquet.hadoop.ParquetFileReader
import org.apache.parquet.hadoop.util.HadoopInputFile
-import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext}
+import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, SparkThrowable, Success, TaskContext}
import org.apache.spark.rdd.RDD
import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
import org.apache.spark.sql.CometTestBase
@@ -2235,18 +2235,8 @@ class CometIcebergWriteActionSuite
"region",
"amount")
- var nativeFailure: Throwable = null
- val nativePlans = withNativeEnabled {
- capturePlans(spark, includeFailures = true) {
- nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows))
- }
- }
- val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows))
- assert(
- nativePlans.exists(p =>
- collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty),
- "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" +
- nativePlans.mkString("\n--\n"))
+ val (nativeFailure, jvmFailure) =
+ failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows))
val nativeError = clusteredWriterError(nativeFailure)
val jvmError = clusteredWriterError(jvmFailure)
@@ -2293,18 +2283,8 @@ class CometIcebergWriteActionSuite
(2, new Timestamp(0L), Array[Byte](0)),
(3, new Timestamp(-1500L), Array[Byte](0, 1, -1))).toDF("id", "ts", "bin")
- var nativeFailure: Throwable = null
- val nativePlans = withNativeEnabled {
- capturePlans(spark, includeFailures = true) {
- nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows))
- }
- }
- val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows))
- assert(
- nativePlans.exists(p =>
- collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty),
- "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" +
- nativePlans.mkString("\n--\n"))
+ val (nativeFailure, jvmFailure) =
+ failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows))
val (_, nativeMessage) = clusteredWriterError(nativeFailure)
assert(
@@ -2319,6 +2299,51 @@ class CometIcebergWriteActionSuite
}
}
+ // https://github.com/apache/datafusion-comet/issues/6234. The native writer pulls its input
+ // through an Arrow C stream, and Arrow Java hands native only the text of an exception thrown
+ // while producing a batch, so the user got a CometNativeException instead of Spark's exception.
+ // The first batch is read on the JVM to derive the stream's schema, which is why the overflow
+ // has to land past it: row 90000 of a single 100000-row data file, so one task reads it all.
+ test("native acceleration: an input error past the first batch fails like the JVM writer") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { _ =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ val (nativeTable, jvmTable) = ("overflow_native", "overflow_jvm")
+ Seq("overflow_src", nativeTable, jvmTable).foreach { table =>
+ spark.sql(s"CREATE TABLE $catalog.$ns.$table (v INT) USING iceberg")
+ }
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100000, 1, 1)
+ .selectExpr("CAST(id AS INT) AS v")
+ .writeTo(s"$catalog.$ns.overflow_src")
+ .append()
+ }
+ def insert(table: String): Unit =
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.$table " +
+ s"SELECT v + ${Int.MaxValue - 90000} FROM $catalog.$ns.overflow_src")
+
+ val (nativeFailure, jvmFailure) = failedWrites(nativeTable, jvmTable)(insert)
+
+ def overflow(failure: Throwable): (Int, Class[_], String, String) = {
+ val chain = causeChain(failure)
+ val depth = chain.indexWhere(_.isInstanceOf[ArithmeticException])
+ assert(depth >= 0, s"no ArithmeticException in the cause chain of $failure")
+ val error = chain(depth).asInstanceOf[Throwable with SparkThrowable]
+ (depth, error.getClass, error.getErrorClass, error.getSqlState)
+ }
+ assert(overflow(jvmFailure)._3 == "ARITHMETIC_OVERFLOW", s"JVM writer: $jvmFailure")
+ assert(
+ overflow(nativeFailure) == overflow(jvmFailure),
+ s"native writer: $nativeFailure\nJVM writer: $jvmFailure")
+ Seq(nativeTable, jvmTable).foreach { table =>
+ assert(spark.table(s"$catalog.$ns.$table").count() == 0, s"$table committed rows")
+ }
+ }
+ }
+ }
+
test("native acceleration: target-file-size rolls on iceberg-java's 1000-row cadence") {
assumeNativeAcceleration()
withIcebergCatalog { warehouseDir =>
@@ -3646,6 +3671,28 @@ class CometIcebergWriteActionSuite
(depth, chain(depth).getMessage)
}
+ /**
+ * Runs `write` against `nativeTable` with the native writer and against `jvmTable` with the JVM
+ * writer, expects both to fail, and returns the two failures. The native write has to go
+ * through [[CometIcebergWriteExec]], so a fallback to the JVM writer cannot pass as parity.
+ */
+ private def failedWrites(nativeTable: String, jvmTable: String)(
+ write: String => Unit): (Throwable, Throwable) = {
+ var nativeFailure: Throwable = null
+ val nativePlans = withNativeEnabled {
+ capturePlans(spark, includeFailures = true) {
+ nativeFailure = intercept[Exception](write(nativeTable))
+ }
+ }
+ val jvmFailure = intercept[Exception](write(jvmTable))
+ assert(
+ nativePlans.exists(p =>
+ collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty),
+ "expected the failed write to have gone through CometIcebergWriteExec. Plans:\n" +
+ nativePlans.mkString("\n--\n"))
+ (nativeFailure, jvmFailure)
+ }
+
private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = {
val before = countSnapshots(tableName)
val plans = capturePlans(spark)(action)
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
index 22e76a9abd..5071c40cb6 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -52,7 +52,7 @@ import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
-import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo}
+import org.apache.comet.{CometConf, CometExecIterator, CometNativeException, ExtendedExplainInfo}
import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus}
import org.apache.comet.rules.CometCoalesceShufflePartitions
import org.apache.comet.serde.Config.ConfigMap
@@ -3979,6 +3979,55 @@ class CometExecSuite extends CometTestBase {
})
}
+ // https://github.com/apache/datafusion-comet/issues/6234. The final TopK reads its child's
+ // batches through an Arrow C stream, and Arrow Java hands native only the text of an exception
+ // thrown while producing one. The first batch is read on the JVM to derive the stream's schema,
+ // so the overflow has to land past it: row 90000 of a single file, read by one task.
+ test("TakeOrderedAndProjectExec: an input error past the first batch keeps Spark's exception") {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempPath { dir =>
+ spark
+ .range(0, 100000, 1, 1)
+ .selectExpr("CAST(id AS INT) AS id")
+ .write
+ .parquet(dir.getCanonicalPath)
+ spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("overflow_src")
+ val df =
+ sql(s"SELECT id + ${Int.MaxValue - 90000} AS v FROM overflow_src ORDER BY v LIMIT 5")
+ checkSparkError(df, "ARITHMETIC_OVERFLOW")
+ }
+ }
+ }
+
+ // https://github.com/apache/datafusion-comet/issues/6234. CometSparkToColumnarExec exports its
+ // rows to native without reading a batch on the JVM first, so even an exception thrown for the
+ // first batch reached the user only as the text inside a CometNativeException.
+ test("SparkToColumnar keeps the exception its row input throws") {
+ val rows = spark.sparkContext.parallelize(1 to 10, 1).map { i =>
+ if (i == 5) throw new IllegalStateException("injected row input failure")
+ Row(i)
+ }
+ val df = spark
+ .createDataFrame(rows, StructType(Seq(StructField("a", IntegerType))))
+ .groupBy()
+ .sum("a")
+ assert(
+ collect(df.queryExecution.executedPlan) { case c: CometSparkToColumnarExec => c }.nonEmpty,
+ "expected the rows to reach native through CometSparkToColumnarExec:\n" +
+ df.queryExecution.executedPlan)
+
+ val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df)
+ def injected(error: Option[Throwable], engine: String): (Class[_], Int) = {
+ val chain = causeChain(error.getOrElse(fail(s"$engine did not fail")))
+ assert(!chain.exists(_.isInstanceOf[CometNativeException]), s"$engine: ${chain.head}")
+ val depth = chain.indexWhere(t =>
+ t.isInstanceOf[IllegalStateException] && t.getMessage == "injected row input failure")
+ assert(depth >= 0, s"$engine did not surface the injected failure: ${chain.head}")
+ (chain.head.getClass, depth)
+ }
+ assert(injected(cometError, "Comet") == injected(sparkError, "Spark"))
+ }
+
// Arrow Java ignores ArrowArray.offset on import, and a sliced boolean is the one array arrow-rs
// exports with a non-zero offset, so native has to zero it at every level, not only the top.
// https://github.com/apache/datafusion-comet/issues/6288
diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
index d6237f76b8..0485fb6181 100644
--- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala
@@ -29,6 +29,7 @@ import scala.reflect.ClassTag
import org.apache.arrow.c.ArrowArrayStream
import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.BigIntVector
import org.apache.arrow.vector.ipc.ArrowReader
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
import org.apache.spark.executor.TaskMetrics
@@ -40,11 +41,13 @@ import org.apache.spark.sql.comet.{CometExec, CometExecUtils, CometMetricNode, C
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.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.spark.util.JsonProtocol
-import org.apache.comet.{CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native}
+import org.apache.comet.{CometArrowAllocator, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native}
import org.apache.comet.serde.Config.ConfigMap
import org.apache.comet.serde.OperatorOuterClass
+import org.apache.comet.vector.CometPlainVector
/**
* Regression tests for the native plan lifecycle: every `createPlan` must be balanced by exactly
@@ -237,6 +240,56 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
}
}
+ test("an input that throws after its first batch fails the plan with its own exception") {
+ val allocator = CometArrowAllocator.newChildAllocator("input-failure", 0, Long.MaxValue)
+ try {
+ withTaskContext(4600000L) {
+ val vector = new BigIntVector("test", allocator)
+ vector.allocateNew(2)
+ vector.set(0, 1L)
+ vector.set(1, 2L)
+ vector.setValueCount(2)
+ val first = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), 2)
+ val boom = new IllegalStateException("injected input failure")
+ // The stream's schema is derived from the first batch on the JVM, so a failure there
+ // propagates unchanged. Native pulls every later batch through the exported stream, and
+ // Arrow Java hands native only the text of what that pull threw.
+ val input = Iterator.single(first) ++ new Iterator[ColumnarBatch] {
+ override def hasNext: Boolean = throw boom
+ override def next(): ColumnarBatch = throw boom
+ }
+ val schema = StructType(Seq(StructField("test", LongType, nullable = false)))
+ val stream = CometArrowStream.fromColumnarBatchIter(
+ input,
+ schema,
+ CometArrowStream.NATIVE_TIMEZONE,
+ "input-failure")
+ val limitOp =
+ CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get
+ val iter = CometExec.getCometIterator(
+ Array(stream.asInstanceOf[Object]),
+ 1,
+ limitOp,
+ CometMetricNode(Map.empty),
+ 1,
+ 0,
+ None,
+ Seq.empty)
+
+ assert(iter.hasNext && iter.next().numRows() == 2)
+ val thrown = intercept[Throwable](iter.hasNext)
+ assert(thrown eq boom, s"expected the input's own exception, got: $thrown")
+
+ TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None)
+ assert(
+ CometArrowStream.inputFailure(stream).isEmpty,
+ "a completed task must not keep its input streams' failures")
+ }
+ } finally {
+ allocator.close()
+ }
+ }
+
test("a plan fed only by native scans returns its memory when its consumer stops early") {
// See issue #2453. The sort keeps its sorted runs in memory, and the merge that reads them
// spawns a Tokio task for each run.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org