This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new e33017a9c620 [SPARK-57383][SQL][PYTHON] Honor configured Arrow zstd
compression level when writing Arrow batches
e33017a9c620 is described below
commit e33017a9c62008e007333dad8ca13ab049aaf9d8
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Jun 11 08:51:06 2026 -0700
[SPARK-57383][SQL][PYTHON] Honor configured Arrow zstd compression level
when writing Arrow batches
### What changes were proposed in this pull request?
This PR fixes a bug where the zstd compression level configured via
`spark.sql.execution.arrow.compression.zstd.level` was silently ignored
everywhere Arrow batches are compressed. Three places shared the same broken
pattern:
- `ArrowConverters.ArrowBatchIterator` (SPARK-54134)
- `PythonArrowInput` (SPARK-54226; also covers `GroupedPythonArrowInput`,
which reuses this codec via SPARK-55328)
- `CoGroupedArrowPythonRunner` (SPARK-54226)
They constructed `new ZstdCompressionCodec(level)` only to read its codec
type, then rebuilt the codec through
`CompressionCodec.Factory.INSTANCE.createCodec(codecType)`. The codec type enum
does not carry a level, so that single-argument factory overload always builds
a codec at the zstd default level (3), dropping the configured one.
The codec construction is extracted into a shared
`ArrowCompressionUtils.createCompressionCodec` helper that constructs the
level-carrying codec instance directly (the helper lives in `sql/core` because
`sql/api`, where `ArrowUtils` is, has no `arrow-compression` dependency). The
level only matters on the write side; the read side looks up the codec by the
type recorded in the IPC message, so reads are unaffected and the on-wire
format is unchanged.
The same bug class was found by dbtsai during review of #56334
(https://github.com/apache/spark/pull/56334#discussion_r3391654988); that PR
fixes the cache-side instance of the pattern, and this PR fixes the remaining
three pre-existing instances.
### Why are the changes needed?
Users tuning `spark.sql.execution.arrow.compression.zstd.level` for Python
UDF exchange or `df.toArrow()` got no effect at all: every level compressed
identically at the default level 3, with no error or warning.
### Does this PR introduce _any_ user-facing change?
Yes. The configured zstd level now actually takes effect; previously all
levels behaved like the default level 3. The bug exists in released Spark
4.1.0/4.1.1/4.1.2 (SPARK-54134 and SPARK-54226 were backported to branch-4.1)
as well as 4.2.0 RCs and master, so this fix is a candidate for backporting to
branch-4.1 and branch-4.2. Note that a branch-4.1 backport needs to fix a
fourth copy of the pattern: there `GroupedPythonArrowInput` still has its own
codec construction, since the SPA [...]
### How was this patch tested?
New `ArrowCompressionUtilsSuite`. The regression test compresses the same
compressible-but-varying batch at zstd level -5 and level 19 and asserts level
19 produces a strictly smaller payload. Against the old codec construction this
test fails with byte-identical sizes at both levels (verified locally). A
second test covers the `none` codec and the unsupported-codec error.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #56444 from viirya/fix-arrow-zstd-level.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../execution/arrow/ArrowCompressionUtils.scala | 50 ++++++++++++
.../sql/execution/arrow/ArrowConverters.scala | 21 +-----
.../python/CoGroupedArrowPythonRunner.scala | 23 +-----
.../sql/execution/python/PythonArrowInput.scala | 24 ++----
.../arrow/ArrowCompressionUtilsSuite.scala | 88 ++++++++++++++++++++++
5 files changed, 149 insertions(+), 57 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala
new file mode 100644
index 000000000000..19b84d1cd5d8
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.execution.arrow
+
+import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
+import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
+
+import org.apache.spark.SparkException
+
+private[sql] object ArrowCompressionUtils {
+
+ /**
+ * Creates the write-side Arrow [[CompressionCodec]] for the codec selected
by
+ * `spark.sql.execution.arrow.compression.codec`, honoring
+ * `spark.sql.execution.arrow.compression.zstd.level` for zstd.
+ *
+ * The codec instance must be constructed directly rather than through
+ * `CompressionCodec.Factory.INSTANCE.createCodec(codecType)`: the codec
type enum does not
+ * carry a compression level, so that factory overload always builds a codec
at the default
+ * level, silently dropping the configured one. The level only matters on
the write side; the
+ * read side looks up the codec by the type recorded in the IPC message, so
it is unaffected.
+ */
+ def createCompressionCodec(
+ codecName: String,
+ zstdCompressionLevel: Int): CompressionCodec = {
+ codecName match {
+ case "none" => NoCompressionCodec.INSTANCE
+ case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
+ case "lz4" => new Lz4CompressionCodec()
+ case other =>
+ throw SparkException.internalError(
+ s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
+ }
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
index b1e8217ff257..464cac157b25 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
@@ -24,15 +24,12 @@ import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import scala.util.control.NonFatal
-import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
import org.apache.arrow.flatbuf.MessageHeader
import org.apache.arrow.memory.BufferAllocator
import org.apache.arrow.vector._
-import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter,
ReadChannel, WriteChannel}
import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, IpcOption,
MessageSerializer}
-import org.apache.spark.SparkException
import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
import org.apache.spark.network.util.JavaUtils
@@ -99,22 +96,8 @@ private[sql] object ArrowConverters extends Logging {
protected val root = VectorSchemaRoot.create(arrowSchema, allocator)
// Create compression codec based on config
- private val compressionCodecName = SQLConf.get.arrowCompressionCodec
- private val codec = compressionCodecName match {
- case "none" => NoCompressionCodec.INSTANCE
- case "zstd" =>
- val compressionLevel = SQLConf.get.arrowZstdCompressionLevel
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new
ZstdCompressionCodec(compressionLevel).getCodecType()
- factory.createCodec(codecType)
- case "lz4" =>
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new Lz4CompressionCodec().getCodecType()
- factory.createCodec(codecType)
- case other =>
- throw SparkException.internalError(
- s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
- }
+ private val codec = ArrowCompressionUtils.createCompressionCodec(
+ SQLConf.get.arrowCompressionCodec, SQLConf.get.arrowZstdCompressionLevel)
protected val unloader = new VectorUnloader(root, true, codec, true)
protected val arrowWriter = ArrowWriter.create(root)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
index df108187c9f0..be36b368719a 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
@@ -20,14 +20,12 @@ package org.apache.spark.sql.execution.python
import java.io.DataOutputStream
import java.util
-import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
import org.apache.arrow.vector.{VectorSchemaRoot, VectorUnloader}
-import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
-import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions,
PythonWorker}
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.execution.arrow.ArrowWriterWrapper
+import org.apache.spark.sql.execution.arrow.{ArrowCompressionUtils,
ArrowWriterWrapper}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
@@ -90,21 +88,8 @@ class CoGroupedArrowPythonRunner(
// Helper method to create VectorUnloader with compression
private def createUnloader(root: VectorSchemaRoot): VectorUnloader = {
- val codec = compressionCodecName match {
- case "none" => NoCompressionCodec.INSTANCE
- case "zstd" =>
- val compressionLevel = SQLConf.get.arrowZstdCompressionLevel
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new
ZstdCompressionCodec(compressionLevel).getCodecType()
- factory.createCodec(codecType)
- case "lz4" =>
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new Lz4CompressionCodec().getCodecType()
- factory.createCodec(codecType)
- case other =>
- throw SparkException.internalError(
- s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
- }
+ val codec = ArrowCompressionUtils.createCompressionCodec(
+ compressionCodecName, SQLConf.get.arrowZstdCompressionLevel)
new VectorUnloader(root, true, codec, true)
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
index 2b200294803d..b0c457615382 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
@@ -19,19 +19,18 @@ package org.apache.spark.sql.execution.python
import java.io.DataOutputStream
import java.nio.channels.Channels
-import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
import org.apache.arrow.memory.BufferAllocator
import org.apache.arrow.vector.{VectorSchemaRoot, VectorUnloader}
-import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
+import org.apache.arrow.vector.compression.CompressionCodec
import org.apache.arrow.vector.ipc.ArrowStreamWriter
import org.apache.arrow.vector.ipc.WriteChannel
import org.apache.arrow.vector.ipc.message.MessageSerializer
-import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.api.python.{BasePythonRunner, PythonWorker}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.arrow
-import org.apache.spark.sql.execution.arrow.{ArrowWriter, ArrowWriterWrapper}
+import org.apache.spark.sql.execution.arrow.{ArrowCompressionUtils,
ArrowWriter, ArrowWriterWrapper}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
@@ -69,21 +68,8 @@ private[python] trait PythonArrowInput[IN] { self:
BasePythonRunner[IN, _] =>
}
// Create compression codec based on config
- protected def codec: CompressionCodec = SQLConf.get.arrowCompressionCodec
match {
- case "none" => NoCompressionCodec.INSTANCE
- case "zstd" =>
- val compressionLevel = SQLConf.get.arrowZstdCompressionLevel
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new ZstdCompressionCodec(compressionLevel).getCodecType()
- factory.createCodec(codecType)
- case "lz4" =>
- val factory = CompressionCodec.Factory.INSTANCE
- val codecType = new Lz4CompressionCodec().getCodecType()
- factory.createCodec(codecType)
- case other =>
- throw SparkException.internalError(
- s"Unsupported Arrow compression codec: $other. Supported values: none,
zstd, lz4")
- }
+ protected def codec: CompressionCodec =
ArrowCompressionUtils.createCompressionCodec(
+ SQLConf.get.arrowCompressionCodec, SQLConf.get.arrowZstdCompressionLevel)
protected var writer: ArrowStreamWriter = _
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala
new file mode 100644
index 000000000000..50a980dcf37a
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala
@@ -0,0 +1,88 @@
+/*
+ * 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.execution.arrow
+
+import java.io.ByteArrayOutputStream
+import java.nio.channels.Channels
+
+import org.apache.arrow.vector.{VarCharVector, VectorSchemaRoot,
VectorUnloader}
+import org.apache.arrow.vector.compression.NoCompressionCodec
+import org.apache.arrow.vector.ipc.WriteChannel
+import org.apache.arrow.vector.ipc.message.MessageSerializer
+
+import org.apache.spark.{SparkException, SparkFunSuite}
+import org.apache.spark.sql.types.{StringType, StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+
+class ArrowCompressionUtilsSuite extends SparkFunSuite {
+
+ // Serializes one Arrow record batch compressed at the given zstd level and
returns its size.
+ private def compressedSize(level: Int): Int = {
+ val sparkSchema = StructType(Seq(StructField("str_col", StringType)))
+ val arrowSchema = ArrowUtils.toArrowSchema(sparkSchema, "UTC", true, false)
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName,
0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ try {
+ root.allocateNew()
+ val strVector = root.getVector("str_col").asInstanceOf[VarCharVector]
+ // Compressible but non-trivial corpus: shared structure with per-row
variation, so
+ // different zstd levels produce measurably different output sizes.
+ (0 until 2000).foreach { i =>
+ val value =
+ s"[email protected],record-${i % 97},payload-${i *
2654435761L}".getBytes("UTF-8")
+ strVector.setSafe(i, value, 0, value.length)
+ }
+ root.setRowCount(2000)
+ val codec = ArrowCompressionUtils.createCompressionCodec("zstd", level)
+ val recordBatch = new VectorUnloader(root, true, codec,
true).getRecordBatch()
+ try {
+ val out = new ByteArrayOutputStream()
+ MessageSerializer.serialize(new
WriteChannel(Channels.newChannel(out)), recordBatch)
+ out.size()
+ } finally {
+ recordBatch.close()
+ }
+ } finally {
+ root.close()
+ allocator.close()
+ }
+ }
+
+ test("SPARK-57383: zstd compression level is honored by
createCompressionCodec") {
+ // Regression test: the codec used to be rebuilt through the
single-argument factory
+ // overload, which silently dropped the level, so every configured level
compressed at the
+ // zstd default. Compress the same batch at an ultra-fast negative level
and at a high level
+ // and assert the high level yields a strictly smaller payload.
+ val fastSize = compressedSize(-5)
+ val highSize = compressedSize(19)
+ assert(highSize < fastSize,
+ s"zstd level 19 should compress smaller than level -5, " +
+ s"got level 19 -> $highSize bytes vs level -5 -> $fastSize bytes; " +
+ "equal sizes mean the configured level is being ignored")
+ }
+
+ test("codec name 'none' maps to the no-op codec and unknown names fail") {
+ assert(ArrowCompressionUtils.createCompressionCodec("none", 3) ===
+ NoCompressionCodec.INSTANCE)
+ val e = intercept[SparkException] {
+ ArrowCompressionUtils.createCompressionCodec("snappy", 3)
+ }
+ assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
