sunchao commented on code in PR #5491:
URL: https://github.com/apache/datafusion-comet/pull/5491#discussion_r3868436058
##########
spark/src/main/scala/org/apache/comet/CometExecIterator.scala:
##########
@@ -130,6 +134,37 @@ class CometExecIterator(
// worker has neither. See CometUdfBridge.evaluate.
TaskContext.get(),
Thread.currentThread().getContextClassLoader)
+
+ // Bind task-owned callbacks separately to preserve the existing
createPlan JNI signature.
+ try {
+ shufflePartitionPusher.foreach { pusher =>
+ nativeLib.setShufflePartitionPusher(createdPlan, pusher)
+ }
+ createdPlan
+ } catch {
+ case failure: Throwable =>
+ // The task-completion listener is not installed until iterator
construction succeeds.
+ try {
+ nativeUtil.close()
+ } catch {
+ case closeFailure: Throwable => failure.addSuppressed(closeFailure)
+ }
+
+ shuffleBlockIterators.values.foreach { iterator =>
Review Comment:
[P2] Release unconsumed Arrow inputs when callback registration fails
If `setShufflePartitionPusher` throws before the first `executePlan`
(including the new `Some(null)` test), `inputObjects` still contains exported
`ArrowArrayStream`s that native planning has not imported. This rollback closes
the shuffle block inputs but omits the Arrow inputs: `releasePlan` drops the
execution context's JNI references and has no native stream reader to release
at this point. The later `CometArrowStream` task listener calls
`ArrowArrayStream.close()`, which frees the wrapper without invoking the C
stream release callback, leaving the exported reader and its resources
retained. Release each unconsumed Arrow input during this rollback, preserving
other cleanup failures as suppressed exceptions, and add a regression test with
a reader that owns resources.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
