comphead commented on code in PR #7307: URL: https://github.com/apache/datafusion-comet/pull/7307#discussion_r4405722878
########## spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala: ########## @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +import com.github.luben.zstd.Zstd + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. Review Comment: The overrun is explained in full in the `CometZstdCompressionCodec` scaladoc, again here, and a third time on `exposedLengths`. A one-line pointer to the codec's scaladoc would keep a single copy. ########## spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala: ########## @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +import com.github.luben.zstd.Zstd + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. + * + * The buffers come from [[GuardedAllocator]], which hands the codec exactly the capacity it asked + * for and keeps a sentinel-filled guard region right behind it. An overrun of the output buffer + * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed + * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: + * bytes that look structured enough for zstd to try entropy coding but do not compress, in the + * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + */ +class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { + + // zstd-jni unpacks its native library into java.io.tmpdir, which the test JVM points at + // target/tmp. Nothing else creates that directory in a fresh checkout. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + /** + * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 + * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less + * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding + * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is + * 509, allocated as 512. + */ + private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => + val request = Prefix + Zstd.compressBound(n) + val rounded = java.lang.Long.highestOneBit(request - 1) << 1 + rounded - request < Prefix + } + + // Protobuf-like bytes: random payload with a scatter of repeated field tags. The tags make the + // block pass zstd's "worth trying" heuristic, the payload makes the attempt not pay off, so + // zstd runs its entropy coder to the end of the capacity it was given before storing the block. + private def protoLike(len: Int, seed: Int): Array[Byte] = { + val rnd = new Random(seed * 7919L + len) + val out = new Array[Byte](len) + rnd.nextBytes(out) + (0 until 8 + rnd.nextInt(24)).foreach { _ => + out(rnd.nextInt(len)) = (0x08 + 8 * rnd.nextInt(4)).toByte + } + out + } + + private def compressible(len: Int, seed: Int): Array[Byte] = + Array.tabulate(len)(i => ((i / 7 + seed) % 5).toByte) + + private def roundTrip( + codec: ZstdCompressionCodec, + allocator: BufferAllocator, + guarded: GuardedAllocator, + bytes: Array[Byte]): Boolean = { + val input = allocator.buffer(bytes.length) + input.setBytes(0, bytes) + input.writerIndex(bytes.length) + // compress consumes the input reference and allocates its output from the guarded allocator. + val compressed = codec.compress(guarded, input) + val storedRaw = compressed.getLong(0) == CompressionUtil.NO_COMPRESSION_LENGTH + val plain = codec.decompress(allocator, compressed) + try { + plain.writerIndex() shouldBe bytes.length + val actual = new Array[Byte](bytes.length) + plain.getBytes(0, actual) + actual shouldBe bytes + } finally { + plain.close() + } + storedRaw + } + + private def withAllocators(f: (BufferAllocator, GuardedAllocator) => Unit): Unit = { + val root = new RootAllocator() + val guarded = new GuardedAllocator(root) + try f(root, guarded) + finally { + guarded.close() + root.close() + } + } + + test("Arrow's zstd codec writes past the output buffer it allocates") { + // Documents the defect this suite guards against (apache/arrow-java#1119, fixed in Arrow + // 20.0.0). Once Comet moves to an Arrow with the fix, CometZstdCompressionCodec and this test + // go together. + withAllocators { (root, guarded) => + val codec = new ZstdCompressionCodec(1) + exposedLengths.filter(_ < 1024).foreach { len => + (0 until 50).foreach(seed => roundTrip(codec, root, guarded, protoLike(len, seed))) + } + guarded.overruns should be > 0 + } + } + + test("Comet's zstd codec stays within the output buffer it allocates") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(1) + var raw = 0 + var packed = 0 + exposedLengths.foreach { len => + (0 until 50).foreach { seed => + if (roundTrip(codec, root, guarded, protoLike(len, seed))) raw += 1 else packed += 1 + } + } + guarded.overruns shouldBe 0 + // The inputs are meant to defeat compression, which is the path that overruns. + raw should be > 0 + } + } + + test("Comet's zstd codec round-trips compressible and empty input") { Review Comment: This mostly repeats existing coverage. Empty input never reaches `doCompress`, because `AbstractCompressionCodec.compress` returns early when `writerIndex()` is 0. The compressed round trip through the config is covered by "Comet in-memory cache round-trips under every compression codec" in `CometInMemoryCacheSuite`, which now goes through this codec. The only new check is the guard on a compressed outcome, which fits in the previous test. ########## spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala: ########## @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +import com.github.luben.zstd.Zstd + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. + * + * The buffers come from [[GuardedAllocator]], which hands the codec exactly the capacity it asked + * for and keeps a sentinel-filled guard region right behind it. An overrun of the output buffer + * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed + * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: + * bytes that look structured enough for zstd to try entropy coding but do not compress, in the + * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + */ +class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { + + // zstd-jni unpacks its native library into java.io.tmpdir, which the test JVM points at + // target/tmp. Nothing else creates that directory in a fresh checkout. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + /** + * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 + * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less + * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding + * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is + * 509, allocated as 512. + */ + private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => + val request = Prefix + Zstd.compressBound(n) + val rounded = java.lang.Long.highestOneBit(request - 1) << 1 + rounded - request < Prefix + } + + // Protobuf-like bytes: random payload with a scatter of repeated field tags. The tags make the + // block pass zstd's "worth trying" heuristic, the payload makes the attempt not pay off, so + // zstd runs its entropy coder to the end of the capacity it was given before storing the block. + private def protoLike(len: Int, seed: Int): Array[Byte] = { + val rnd = new Random(seed * 7919L + len) + val out = new Array[Byte](len) + rnd.nextBytes(out) + (0 until 8 + rnd.nextInt(24)).foreach { _ => + out(rnd.nextInt(len)) = (0x08 + 8 * rnd.nextInt(4)).toByte + } + out + } + + private def compressible(len: Int, seed: Int): Array[Byte] = + Array.tabulate(len)(i => ((i / 7 + seed) % 5).toByte) + + private def roundTrip( + codec: ZstdCompressionCodec, + allocator: BufferAllocator, + guarded: GuardedAllocator, + bytes: Array[Byte]): Boolean = { + val input = allocator.buffer(bytes.length) + input.setBytes(0, bytes) + input.writerIndex(bytes.length) + // compress consumes the input reference and allocates its output from the guarded allocator. + val compressed = codec.compress(guarded, input) + val storedRaw = compressed.getLong(0) == CompressionUtil.NO_COMPRESSION_LENGTH + val plain = codec.decompress(allocator, compressed) + try { + plain.writerIndex() shouldBe bytes.length + val actual = new Array[Byte](bytes.length) + plain.getBytes(0, actual) + actual shouldBe bytes + } finally { + plain.close() + } + storedRaw + } + + private def withAllocators(f: (BufferAllocator, GuardedAllocator) => Unit): Unit = { + val root = new RootAllocator() + val guarded = new GuardedAllocator(root) + try f(root, guarded) + finally { + guarded.close() + root.close() + } + } + + test("Arrow's zstd codec writes past the output buffer it allocates") { + // Documents the defect this suite guards against (apache/arrow-java#1119, fixed in Arrow + // 20.0.0). Once Comet moves to an Arrow with the fix, CometZstdCompressionCodec and this test + // go together. + withAllocators { (root, guarded) => + val codec = new ZstdCompressionCodec(1) + exposedLengths.filter(_ < 1024).foreach { len => + (0 until 50).foreach(seed => roundTrip(codec, root, guarded, protoLike(len, seed))) + } + guarded.overruns should be > 0 + } + } + + test("Comet's zstd codec stays within the output buffer it allocates") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(1) + var raw = 0 + var packed = 0 Review Comment: `packed` is counted but never read. Adding a few `compressible` inputs to this loop and asserting `packed should be > 0` would cover the compressed outcome under the guard and let the next test go. Otherwise it can be dropped. ########## spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala: ########## @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +import com.github.luben.zstd.Zstd + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. + * + * The buffers come from [[GuardedAllocator]], which hands the codec exactly the capacity it asked + * for and keeps a sentinel-filled guard region right behind it. An overrun of the output buffer + * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed + * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: + * bytes that look structured enough for zstd to try entropy coding but do not compress, in the + * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + */ +class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { + + // zstd-jni unpacks its native library into java.io.tmpdir, which the test JVM points at + // target/tmp. Nothing else creates that directory in a fresh checkout. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + /** + * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 + * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less + * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding + * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is + * 509, allocated as 512. + */ + private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => + val request = Prefix + Zstd.compressBound(n) + val rounded = java.lang.Long.highestOneBit(request - 1) << 1 + rounded - request < Prefix + } + + // Protobuf-like bytes: random payload with a scatter of repeated field tags. The tags make the + // block pass zstd's "worth trying" heuristic, the payload makes the attempt not pay off, so + // zstd runs its entropy coder to the end of the capacity it was given before storing the block. + private def protoLike(len: Int, seed: Int): Array[Byte] = { + val rnd = new Random(seed * 7919L + len) + val out = new Array[Byte](len) + rnd.nextBytes(out) + (0 until 8 + rnd.nextInt(24)).foreach { _ => + out(rnd.nextInt(len)) = (0x08 + 8 * rnd.nextInt(4)).toByte + } + out + } + + private def compressible(len: Int, seed: Int): Array[Byte] = + Array.tabulate(len)(i => ((i / 7 + seed) % 5).toByte) + + private def roundTrip( + codec: ZstdCompressionCodec, + allocator: BufferAllocator, + guarded: GuardedAllocator, + bytes: Array[Byte]): Boolean = { + val input = allocator.buffer(bytes.length) + input.setBytes(0, bytes) + input.writerIndex(bytes.length) + // compress consumes the input reference and allocates its output from the guarded allocator. + val compressed = codec.compress(guarded, input) + val storedRaw = compressed.getLong(0) == CompressionUtil.NO_COMPRESSION_LENGTH + val plain = codec.decompress(allocator, compressed) + try { + plain.writerIndex() shouldBe bytes.length + val actual = new Array[Byte](bytes.length) + plain.getBytes(0, actual) + actual shouldBe bytes + } finally { + plain.close() + } + storedRaw + } + + private def withAllocators(f: (BufferAllocator, GuardedAllocator) => Unit): Unit = { + val root = new RootAllocator() + val guarded = new GuardedAllocator(root) + try f(root, guarded) + finally { + guarded.close() + root.close() + } + } + + test("Arrow's zstd codec writes past the output buffer it allocates") { + // Documents the defect this suite guards against (apache/arrow-java#1119, fixed in Arrow + // 20.0.0). Once Comet moves to an Arrow with the fix, CometZstdCompressionCodec and this test + // go together. + withAllocators { (root, guarded) => + val codec = new ZstdCompressionCodec(1) + exposedLengths.filter(_ < 1024).foreach { len => + (0 until 50).foreach(seed => roundTrip(codec, root, guarded, protoLike(len, seed))) + } + guarded.overruns should be > 0 + } + } + + test("Comet's zstd codec stays within the output buffer it allocates") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(1) + var raw = 0 + var packed = 0 + exposedLengths.foreach { len => + (0 until 50).foreach { seed => + if (roundTrip(codec, root, guarded, protoLike(len, seed))) raw += 1 else packed += 1 + } + } + guarded.overruns shouldBe 0 + // The inputs are meant to defeat compression, which is the path that overruns. + raw should be > 0 + } + } + + test("Comet's zstd codec round-trips compressible and empty input") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(3) + Seq(0, 1, 7, 437, 4096, 70000).foreach { len => + val storedRaw = roundTrip(codec, root, guarded, compressible(len, len)) + if (len > 16) storedRaw shouldBe false + } + guarded.overruns shouldBe 0 + } + } + + test("the guard notices a write past the buffer") { Review Comment: The first test already shows that the guard detects an overrun, since it asserts `overruns > 0` through the same `GuardedAllocator`. Once Arrow is bumped, that test, the codec and this suite go together, so this test never adds coverage beyond it. ########## spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala: ########## @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +import com.github.luben.zstd.Zstd + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. + * + * The buffers come from [[GuardedAllocator]], which hands the codec exactly the capacity it asked + * for and keeps a sentinel-filled guard region right behind it. An overrun of the output buffer + * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed + * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: + * bytes that look structured enough for zstd to try entropy coding but do not compress, in the + * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + */ +class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { + + // zstd-jni unpacks its native library into java.io.tmpdir, which the test JVM points at + // target/tmp. Nothing else creates that directory in a fresh checkout. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + /** + * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 + * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less + * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding + * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is + * 509, allocated as 512. + */ + private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => Review Comment: `GuardedAllocator` cuts every slice to the exact requested size, so the guard catches the overrun at any length, not only where power-of-two rounding leaves less than 8 bytes of slack. In the same local probe, Arrow's capacity overran on 19177 of 49550 inputs at the other lengths below 1024. A plain range such as `400 to 600` would drive both tests and drop the rounding math and its doc comment. -- 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]
