dongjoon-hyun commented on code in PR #59991:
URL: https://github.com/apache/spark/pull/59991#discussion_r4433535858
##########
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:
Thank you, @viirya. Added in 705be925fa2, with the note about the Arrow
upgrade moved next to the assertion. I ran the test with the `VectorUnloader`
and `AbstractCompressionCodec` changes of apache/arrow-java#1290 on the test
classpath. It now fails at this assertion with
`Array(java.lang.IllegalStateException: RefCnt has gone negative) was not
empty`, and it passes on Arrow 19.0.0.
--
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]
