viirya commented on code in PR #59991:
URL: https://github.com/apache/spark/pull/59991#discussion_r4433278032
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -260,4 +261,42 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
allocator.close()
}
}
+
+ test("SPARK-60776: unloader releases the buffers of a batch it fails to
unload") {
+ var calls = 0
+ val codec = new BoundedZstdCompressionCodec(3) {
+ override protected def doCompress(allocator: BufferAllocator, input:
ArrowBuf): ArrowBuf = {
+ calls += 1
+ if (calls == 4) throw new IllegalStateException("compression failed")
+ super.doCompress(allocator, input)
+ }
+ }
+ val arrowSchema = ArrowUtils.toArrowSchema(
+ StructType(Seq(StructField("c", LongType))), "UTC", true, false)
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName,
0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ val writer = ArrowWriter.create(root)
+ val unloader = ArrowCompressionUtils.createVectorUnloader(root, codec)
+ Utils.tryWithSafeFinally {
+ writer.write(InternalRow(1L))
+ writer.finish()
+ // A batch compresses a validity buffer and a data buffer, so the first
batch succeeds and the
+ // second fails holding its compressed validity buffer and its retained
data buffer.
+ val batch = unloader.getRecordBatch()
+ Utils.tryWithSafeFinally {
+ val e = intercept[IllegalStateException](unloader.getRecordBatch())
Review Comment:
After the Arrow 20 upgrade, closing the compressed validity buffer a second
time ends up as a suppressed exception on `e`. So this `intercept` still
passes, and the test fails only later in `root.close()`. Adding
`assert(e.getSuppressed.isEmpty)` makes it fail here, with a clearer message.
--
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]
