This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new c522b20ac883 [SPARK-57383][4.1][SQL][PYTHON] Honor configured Arrow
zstd compression level when writing Arrow batches
c522b20ac883 is described below
commit c522b20ac883a3e1c612e61256daf53f953b5601
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Jun 11 12:42:50 2026 -0700
[SPARK-57383][4.1][SQL][PYTHON] Honor configured Arrow zstd compression
level when writing Arrow batches
### What changes were proposed in this pull request?
Backport of #56444 (commit e33017a9c62) to branch-4.1.
This 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. The affected code 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 [...]
Compared to the master patch, this backport fixes one additional copy of
the broken pattern: on branch-4.1 `GroupedPythonArrowInput` still has its own
codec construction in `createUnloaderForGroup` (the SPARK-55328 deduplication
that makes it reuse `PythonArrowInput.codec` is master/4.2-only). The four
fixed sites are:
- `ArrowConverters.ArrowBatchIterator`
- `PythonArrowInput`
- `GroupedPythonArrowInput` (branch-4.1 only; folded into
`PythonArrowInput` on master)
- `CoGroupedArrowPythonRunner`
All four now use the shared `ArrowCompressionUtils.createCompressionCodec`
helper, which constructs the level-carrying codec instance directly. 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.
### 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. Released
4.1.0/4.1.1/4.1.2 are affected.
### 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.
### How was this patch tested?
`ArrowCompressionUtilsSuite`, brought over with the backport. 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 it fails with byte-identical sizes at both
levels. Verified the suite passes on this branch.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #56451 from viirya/SPARK-57383-branch-4.1.
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 | 41 ++--------
.../arrow/ArrowCompressionUtilsSuite.scala | 88 ++++++++++++++++++++++
5 files changed, 150 insertions(+), 73 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 3d435ac016ed..a1b51afb9ded 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 7f6efbae8881..32fd2054fdd2 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,
PythonRDD, 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
@@ -86,21 +84,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 f77b0a9342b0..a81052bbe0b2 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,18 +19,16 @@ 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.vector.{VectorSchemaRoot, VectorUnloader}
-import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
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, PythonRDD, 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
@@ -77,22 +75,8 @@ private[python] trait PythonArrowInput[IN] { self:
BasePythonRunner[IN, _] =>
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 var writer: ArrowStreamWriter = _
@@ -276,21 +260,8 @@ private[python] trait GroupedPythonArrowInput { self:
RowInputArrowPythonRunner
// Helper method to create VectorUnloader with compression for grouped
operations
private def createUnloaderForGroup(root: VectorSchemaRoot): VectorUnloader =
{
- val codec = 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")
- }
+ val codec = ArrowCompressionUtils.createCompressionCodec(
+ SQLConf.get.arrowCompressionCodec, SQLConf.get.arrowZstdCompressionLevel)
new VectorUnloader(root, true, codec, true)
}
protected override def newWriter(
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]
