You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2026/08/21 10:28:35 UTC
(kyuubi) branch master updated: [KYUUBI #7635][SPARK] Support Zstd compression for Arrow IPC query results
This is an automated email from the ASF dual-hosted git repository.
pan3793 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new cfbbeec429 [KYUUBI #7635][SPARK] Support Zstd compression for Arrow IPC query results
cfbbeec429 is described below
commit cfbbeec429db0115ce42fbaf91f13c6eb9092762
Author: byron <by...@outlook.com>
AuthorDate: Fri Aug 21 18:28:24 2026 +0800
[KYUUBI #7635][SPARK] Support Zstd compression for Arrow IPC query results
### Why are the changes needed?
Kyuubi's Arrow result format reduces serialization/deserialization overhead, but Arrow IPC results are currently transferred without compression. For large result sets, network transfer can become a significant part of fetch latency.
This follows up on #3877 (related issue #3865). Spark 4.1 introduced the upstream Arrow compression configurations (`SPARK-54134`), so this PR follows Spark's configuration and Arrow IPC compression semantics instead of introducing a Kyuubi-specific protocol.
The main changes are:
- Support `none` (default) and `zstd` for Arrow IPC query results, including ZSTD compression-level configuration.
- Preserve compression when Arrow batches are sliced, e.g. when `LIMIT` crosses a batch boundary.
- Decode compressed batches in the JDBC client using standard Arrow IPC compression metadata.
- Keep `arrow-compression` optional on the Spark engine side. The default `none` path requires no additional dependency; when zstd is enabled, Kyuubi checks the required Arrow compression capability at runtime and reports a clear error if it is unavailable.
- Update `kyuubi-hive-jdbc-shaded` to include/relocate `arrow-compression`, while excluding `zstd-jni` and declaring it optional.
- Exclude `commons-compress`, since LZ4 is not supported.
For Spark runtimes that do not provide a compatible `arrow-compression` implementation, users enabling zstd need to add one to the Spark classpath. The default `none` behavior remains unchanged.
A cluster benchmark with Kyuubi 1.13.0-SNAPSHOT + Spark 3.5.9 showed 78.8%-81.8% fewer Arrow IPC bytes, with fetch time reduced from 19.3s to 7.5s in the 8-column case and from 205.8s to 94.1s in the 34-column case.
### How was this patch tested?
Added/updated coverage for:
- ZSTD compression/decompression and compression-level behavior.
- Compressed batch slicing across `LIMIT` boundaries.
- Full JDBC client <=> Kyuubi server <=> Spark engine ZSTD round trip.
- The default `none` path with no `arrow-compression` on the engine classpath.
- Runtime capability checking when zstd is enabled without the required dependency.
- Session-level `none -> zstd -> none` switching within the same engine.
- Unsupported LZ4 configuration.
Targeted engine tests and `KyuubiOperationPerUserSuite` pass on JDK 17 / Spark 3.5.8. The shaded artifacts were also checked to verify the intended `arrow-compression` / `zstd-jni` packaging.
### Was this patch authored or co-authored using generative AI tooling?
Assisted-by: DeepSeek Harness:deepseek-v4-pro
Closes #7636 from byronwang123/kyuubi-7635-arrow-ipc-compression.
Closes #7635
8911c395f [byronwang] [KYUUBI #7635][SPARK][TEST] Make session-level zstd test independent of the engine classpath
34e8ac6f3 [byronwang] [KYUUBI #7635][SPARK][FOLLOWUP] Prefix Arrow compression helper with Kyuubi
347e38aab [byronwang] [KYUUBI #7635][SPARK][FOLLOWUP] Detect arrow-compression capability at runtime before using zstd
c761a93ac [byronwang] [KYUUBI #7635][SPARK][FOLLOWUP] Make arrow-compression optional on the engine side and exclude zstd-jni from the shaded client
024f3c685 [byronwang] [KYUUBI #7635][SPARK][FOLLOWUP] Drop commons-compress from the arrow-compression dependency chain
c9afd969a [byronwang] [KYUUBI #7635][SPARK][FOLLOWUP] Make LIMIT slice tests deterministic and fail fast on unsupported codec
ed8339839 [byronwang] [KYUUBI #7635][SPARK] Reject lz4 codec and drop reflective Dataset#toArrowBatchRdd
c041d1cd6 [byronwang] [KYUUBI #7635][SPARK] Bundle relocated arrow-compression into the engine jar
0dddffe64 [byronwang] [KYUUBI #7635][SPARK] Support Arrow IPC compression for query results (zstd)
Authored-by: byron <by...@outlook.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
LICENSE-binary | 2 +
dev/dependencyList | 2 +
externals/kyuubi-spark-sql-engine/pom.xml | 24 +++
.../arrow/KyuubiArrowCompressionSupport.scala | 33 ++++
.../execution/arrow/KyuubiArrowConverters.scala | 100 ++++++++++--
.../spark/sql/kyuubi/SparkDatasetHelper.scala | 66 +++++---
.../operation/SparkArrowbasedOperationSuite.scala | 48 ++++++
.../arrow/KyuubiArrowConvertersSuite.scala | 173 +++++++++++++++++++++
.../spark/sql/kyuubi/SparkDatasetHelperSuite.scala | 73 +++++++++
kyuubi-hive-jdbc-shaded/pom.xml | 9 ++
kyuubi-hive-jdbc/pom.xml | 5 +
.../jdbc/hive/KyuubiArrowQueryResultSet.java | 6 +-
.../operation/KyuubiOperationPerUserSuite.scala | 138 +++++++++++++++-
pom.xml | 18 +++
14 files changed, 665 insertions(+), 32 deletions(-)
diff --git a/LICENSE-binary b/LICENSE-binary
index df7580ad89..c2c53faafb 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -304,6 +304,7 @@ io.swagger.core.v3:swagger-jaxrs2
io.swagger.core.v3:swagger-models
io.vertx:vertx-core
io.vertx:vertx-grpc
+org.apache.arrow:arrow-compression
org.apache.kafka:kafka-clients
org.xerial:sqlite-jdbc
com.openai:openai-java
@@ -320,6 +321,7 @@ com.thoughtworks.paranamer:paranamer
com.google.protobuf:protobuf-java-util
com.google.protobuf:protobuf-java
org.postgresql:postgresql
+com.github.luben:zstd-jni
Eclipse Distribution License - v 1.0
------------------------------------
diff --git a/dev/dependencyList b/dev/dependencyList
index 64af754966..139365221f 100644
--- a/dev/dependencyList
+++ b/dev/dependencyList
@@ -22,6 +22,7 @@ annotations/4.1.1.4//annotations-4.1.1.4.jar
antlr-runtime/3.5.3//antlr-runtime-3.5.3.jar
antlr4-runtime/4.9.3//antlr4-runtime-4.9.3.jar
aopalliance-repackaged/2.6.1//aopalliance-repackaged-2.6.1.jar
+arrow-compression/16.0.0//arrow-compression-16.0.0.jar
arrow-format/16.0.0//arrow-format-16.0.0.jar
arrow-memory-core/16.0.0//arrow-memory-core-16.0.0.jar
arrow-memory-netty-buffer-patch/16.0.0//arrow-memory-netty-buffer-patch-16.0.0.jar
@@ -188,3 +189,4 @@ units/1.7//units-1.7.jar
vertx-core/4.5.3//vertx-core-4.5.3.jar
vertx-grpc/4.5.3//vertx-grpc-4.5.3.jar
zjsonpatch/0.3.0//zjsonpatch-0.3.0.jar
+zstd-jni/1.5.5-11//zstd-jni-1.5.5-11.jar
diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml
index 67308284bd..cddc27c102 100644
--- a/externals/kyuubi-spark-sql-engine/pom.xml
+++ b/externals/kyuubi-spark-sql-engine/pom.xml
@@ -37,6 +37,14 @@
<version>${project.version}</version>
</dependency>
+ <!-- Compile-time only: not bundled into the engine jar; provided by Spark 4.1+ or the user. -->
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ <version>${arrow.version}</version>
+ <optional>true</optional>
+ </dependency>
+
<dependency>
<groupId>org.apache.kyuubi</groupId>
<artifactId>kyuubi-events_${scala.binary.version}</artifactId>
@@ -263,6 +271,22 @@
<exclude>io.netty:netty-transport-*-kqueue</exclude>
<exclude>io.netty:netty-transport-*-io_uring</exclude>
<exclude>io.netty:netty-transport-native-epoll:*:linux-riscv64</exclude>
+ <!--
+ The Arrow vector/memory/format modules and the zstd library used by
+ arrow-compression are all provided by the Spark runtime on every
+ supported Spark version, so they are excluded from the engine jar
+ to avoid classpath conflicts and native library relocation issues.
+ -->
+ <exclude>org.apache.arrow:arrow-format</exclude>
+ <exclude>org.apache.arrow:arrow-memory-core</exclude>
+ <exclude>org.apache.arrow:arrow-memory-netty</exclude>
+ <exclude>org.apache.arrow:arrow-memory-netty-buffer-patch</exclude>
+ <exclude>org.apache.arrow:arrow-vector</exclude>
+ <exclude>com.github.luben:zstd-jni</exclude>
+ <exclude>commons-codec:commons-codec</exclude>
+ <exclude>commons-io:commons-io</exclude>
+ <exclude>org.apache.commons:commons-lang3</exclude>
+ <exclude>org.immutables:value</exclude>
</excludes>
</artifactSet>
<filters>
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala
new file mode 100644
index 0000000000..478afcb6d8
--- /dev/null
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala
@@ -0,0 +1,33 @@
+/*
+ * 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.{CommonsCompressionFactory, ZstdCompressionCodec}
+import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader}
+
+/** Isolates the optional arrow-compression dependency so the uncompressed path never loads it. */
+private[sql] object KyuubiArrowCompressionSupport {
+
+ def createLoader(root: VectorSchemaRoot): VectorLoader = {
+ new VectorLoader(root, CommonsCompressionFactory.INSTANCE)
+ }
+
+ def createZstdUnloader(root: VectorSchemaRoot, level: Int): VectorUnloader = {
+ new VectorUnloader(root, true, new ZstdCompressionCodec(level), true)
+ }
+}
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala
index b8f590ddfb..9f690357b8 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala
@@ -24,6 +24,7 @@ import scala.collection.JavaConverters._
import scala.collection.mutable.ArrayBuffer
import org.apache.arrow.vector._
+import org.apache.arrow.vector.compression.NoCompressionCodec
import org.apache.arrow.vector.ipc.{ArrowStreamWriter, ReadChannel, WriteChannel}
import org.apache.arrow.vector.ipc.message.{IpcOption, MessageSerializer}
import org.apache.spark.TaskContext
@@ -39,6 +40,31 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
type Batch = (Array[Byte], Long)
+ private val CommonsCompressionFactoryClassName =
+ "org.apache.arrow.compression.CommonsCompressionFactory"
+ private val ZstdCompressionCodecClassName =
+ "org.apache.arrow.compression.ZstdCompressionCodec"
+
+ // Mirror Spark's SparkSession#enableHiveSupport: check the capability on each request instead
+ // of caching, because the codec is a session-level config that can change at runtime.
+ private def arrowCompressionAvailable: Boolean = {
+ try {
+ Utils.classForName(CommonsCompressionFactoryClassName)
+ Utils.classForName(ZstdCompressionCodecClassName).getConstructor(Integer.TYPE)
+ true
+ } catch {
+ case _: ClassNotFoundException | _: NoClassDefFoundError | _: NoSuchMethodException =>
+ false
+ }
+ }
+
+ private def requireArrowCompression(): Unit = {
+ if (!arrowCompressionAvailable) {
+ throw new IllegalArgumentException(
+ "Arrow ZSTD compression requires arrow-compression on the Spark classpath")
+ }
+ }
+
/**
* this method is to slice the input Arrow record batch byte array `bytes`, starting from `start`
* and taking `length` number of elements.
@@ -48,7 +74,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
timeZoneId: String,
bytes: Array[Byte],
start: Int,
- length: Int): Array[Byte] = {
+ length: Int,
+ codecName: String = null,
+ zstdLevel: Int = 3): Array[Byte] = {
val in = new ByteArrayInputStream(bytes)
val out = new ByteArrayOutputStream(bytes.length)
@@ -65,12 +93,35 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
val recordBatch = MessageSerializer.deserializeRecordBatch(
new ReadChannel(Channels.newChannel(in)),
sliceAllocator)
- val vectorLoader = new VectorLoader(vectorSchemaRoot)
+ // Only compressed batches need the factory; the none path stays free of arrow-compression.
+ val compressed =
+ recordBatch.getBodyCompression.getCodec != NoCompressionCodec.COMPRESSION_TYPE
+ val vectorLoader =
+ if (compressed) {
+ requireArrowCompression()
+ KyuubiArrowCompressionSupport.createLoader(vectorSchemaRoot)
+ } else {
+ new VectorLoader(vectorSchemaRoot)
+ }
vectorLoader.load(recordBatch)
recordBatch.close()
slicedVectorSchemaRoot = vectorSchemaRoot.slice(start, length)
- val unloader = new VectorUnloader(slicedVectorSchemaRoot)
+ // Keep the compression codec on re-serialization, or the client cannot load the batch.
+ val unloader = codecName match {
+ case null | "none" =>
+ new VectorUnloader(slicedVectorSchemaRoot)
+ case "zstd" =>
+ requireArrowCompression()
+ KyuubiArrowCompressionSupport.createZstdUnloader(slicedVectorSchemaRoot, zstdLevel)
+ case "lz4" =>
+ throw new IllegalArgumentException(
+ "Arrow compression codec lz4 is not supported by Kyuubi; " +
+ "supported codecs: none, zstd")
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unsupported Arrow compression codec: $other; supported codecs: none, zstd")
+ }
val writeChannel = new WriteChannel(Channels.newChannel(out))
val batch = unloader.getRecordBatch()
MessageSerializer.serialize(writeChannel, batch)
@@ -119,7 +170,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
collectLimitExec: CollectLimitExec,
maxRecordsPerBatch: Long,
maxEstimatedBatchSize: Long,
- timeZoneId: String): Array[Batch] = {
+ timeZoneId: String,
+ codecName: String = null,
+ zstdLevel: Int = 3): Array[Batch] = {
val n = collectLimitExec.limit
val schema = collectLimitExec.schema
if (n == 0) {
@@ -165,7 +218,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
maxRecordsPerBatch,
maxEstimatedBatchSize,
n,
- timeZoneId)
+ timeZoneId,
+ codecName,
+ zstdLevel)
batches.map(b => b -> batches.rowCountInLastBatch).toArray
},
partsToScan)
@@ -200,7 +255,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
maxRecordsPerBatch: Long,
maxEstimatedBatchSize: Long,
limit: Long,
- timeZoneId: String): ArrowBatchIterator = {
+ timeZoneId: String,
+ codecName: String = null,
+ zstdLevel: Int = 3): ArrowBatchIterator = {
new ArrowBatchIterator(
rowIter,
schema,
@@ -208,7 +265,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
maxEstimatedBatchSize,
limit,
timeZoneId,
- TaskContext.get)
+ TaskContext.get,
+ codecName,
+ zstdLevel)
}
/**
@@ -226,10 +285,27 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
maxEstimatedBatchSize: Long,
limit: Long,
timeZoneId: String,
- context: TaskContext)
+ context: TaskContext,
+ codecName: String,
+ zstdLevel: Int)
extends Iterator[Array[Byte]] {
protected val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true, false)
+ // Validate the codec before allocating Arrow buffers, so an unsupported codec fails fast.
+ private val compressionEnabled = codecName match {
+ case null | "none" =>
+ false
+ case "zstd" =>
+ requireArrowCompression()
+ true
+ case "lz4" =>
+ throw new IllegalArgumentException(
+ "Arrow compression codec lz4 is not supported by Kyuubi; " +
+ "supported codecs: none, zstd")
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unsupported Arrow compression codec: $other; supported codecs: none, zstd")
+ }
private val allocator =
ArrowUtils.rootAllocator.newChildAllocator(
s"to${this.getClass.getSimpleName}",
@@ -237,7 +313,13 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging {
Long.MaxValue)
private val root = VectorSchemaRoot.create(arrowSchema, allocator)
- protected val unloader = new VectorUnloader(root)
+ // The none path keeps the original 1-arg constructor and stays free of arrow-compression.
+ protected val unloader =
+ if (compressionEnabled) {
+ KyuubiArrowCompressionSupport.createZstdUnloader(root, zstdLevel)
+ } else {
+ new VectorUnloader(root)
+ }
protected val arrowWriter = ArrowWriter.create(root)
Option(context).foreach {
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala
index c503d78034..4fd4ecef35 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala
@@ -17,13 +17,15 @@
package org.apache.spark.sql.kyuubi
+import java.util.Locale
+
import scala.collection.mutable.ArrayBuffer
import org.apache.spark.SparkContext
import org.apache.spark.internal.Logging
import org.apache.spark.network.util.{ByteUnit, JavaUtils}
import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{DataFrame, Dataset, Row}
+import org.apache.spark.sql.{DataFrame, Row, SparkSession}
import org.apache.spark.sql.catalyst.plans.logical.GlobalLimit
import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
import org.apache.spark.sql.execution.{CollectLimitExec, CommandResultExec, HiveResult, LocalTableScanExec, QueryExecution, SparkPlan, SparkPlanHelper, SQLExecution}
@@ -36,7 +38,6 @@ import org.apache.spark.sql.types._
import org.apache.kyuubi.engine.spark.KyuubiSparkUtil
import org.apache.kyuubi.engine.spark.schema.RowSet
import org.apache.kyuubi.engine.spark.util.SparkCatalogUtils.quoteIfNeeded
-import org.apache.kyuubi.util.reflect.{DynClasses, DynMethods}
object SparkDatasetHelper extends Logging {
@@ -63,22 +64,23 @@ object SparkDatasetHelper extends Logging {
toArrowBatchRdd(plan).collect()
}
- private val datasetClz = DynClasses.builder()
- .impl("org.apache.spark.sql.classic.Dataset") // SPARK-49700 (4.0.0)
- .impl("org.apache.spark.sql.Dataset")
- .build()
-
- private val toArrowBatchRddMethod =
- DynMethods.builder("toArrowBatchRdd")
- .impl(datasetClz)
- .buildChecked()
-
- def toArrowBatchRdd[T](ds: Dataset[T]): RDD[Array[Byte]] = {
- toArrowBatchRddMethod.bind(ds).invoke()
+ /**
+ * Read session-level Arrow compression configs introduced by SPARK-54134 (4.1.0),
+ * with key names finalized by the SPARK-54226 follow-up.
+ */
+ private def arrowCompressionConf(spark: SparkSession): (String, Int) = {
+ // Lowercase the codec name to match Spark upstream.
+ val codecName = spark.sessionState.conf.getConfString(
+ "spark.sql.execution.arrow.compression.codec",
+ "none").toLowerCase(Locale.ROOT)
+ val zstdLevel = spark.sessionState.conf.getConfString(
+ "spark.sql.execution.arrow.compression.zstd.level",
+ "3").toInt
+ (codecName, zstdLevel)
}
/**
- * Forked from [[Dataset.toArrowBatchRdd(plan: SparkPlan)]].
+ * Forked from [[org.apache.spark.sql.Dataset.toArrowBatchRdd(plan: SparkPlan)]].
* Convert to an RDD of serialized ArrowRecordBatches.
*/
def toArrowBatchRdd(plan: SparkPlan): RDD[Array[Byte]] = {
@@ -89,6 +91,7 @@ object SparkDatasetHelper extends Logging {
// note that, we can't pass the lazy variable `maxBatchSize` directly, this is because input
// arguments are serialized and sent to the executor side for execution.
val maxBatchSizePerBatch = maxBatchSize
+ val (codecName, zstdLevel) = arrowCompressionConf(spark)
plan.execute().mapPartitionsInternal { iter =>
KyuubiArrowConverters.toBatchIterator(
iter,
@@ -96,13 +99,18 @@ object SparkDatasetHelper extends Logging {
maxRecordsPerBatch,
maxBatchSizePerBatch,
-1,
- timeZoneId)
+ timeZoneId,
+ codecName,
+ zstdLevel)
}
}
def toArrowBatchLocalIterator(df: DataFrame): Iterator[Array[Byte]] = {
withNewExecutionId(df) {
- toArrowBatchRdd(df).toLocalIterator
+ // use the plan-based toArrowBatchRdd so that the arrow compression codec takes effect;
+ // the vanilla Spark Dataset#toArrowBatchRdd uses the upstream ArrowConverters, which does
+ // not apply the Kyuubi compression codec.
+ toArrowBatchRdd(df.queryExecution.executedPlan).toLocalIterator
}
}
@@ -174,12 +182,15 @@ object SparkDatasetHelper extends Logging {
val spark = SparkPlanHelper.sparkSession(collectLimit)
val timeZoneId = spark.sessionState.conf.sessionLocalTimeZone
val maxRecordsPerBatch = spark.sessionState.conf.arrowMaxRecordsPerBatch
+ val (codecName, zstdLevel) = arrowCompressionConf(spark)
val batches = KyuubiArrowConverters.takeAsArrowBatches(
collectLimit,
maxRecordsPerBatch,
maxBatchSize,
- timeZoneId)
+ timeZoneId,
+ codecName,
+ zstdLevel)
// note that the number of rows in the returned arrow batches may be >= `limit`, perform
// the slicing operation of result
@@ -193,7 +204,14 @@ object SparkDatasetHelper extends Logging {
// returned ArrowRecordBatch has less than `limit` row count, safety to do conversion
rest -= size.toInt
} else { // size > rest
- result += KyuubiArrowConverters.slice(collectLimit.schema, timeZoneId, batch, 0, rest)
+ result += KyuubiArrowConverters.slice(
+ collectLimit.schema,
+ timeZoneId,
+ batch,
+ 0,
+ rest,
+ codecName,
+ zstdLevel)
rest = 0
}
i += 1
@@ -205,26 +223,32 @@ object SparkDatasetHelper extends Logging {
val spark = SparkPlanHelper.sparkSession(commandResult)
commandResult.longMetric("numOutputRows").add(commandResult.rows.size)
sendDriverMetrics(spark.sparkContext, commandResult.metrics)
+ val (codecName, zstdLevel) = arrowCompressionConf(spark)
KyuubiArrowConverters.toBatchIterator(
commandResult.rows.iterator,
commandResult.schema,
spark.sessionState.conf.arrowMaxRecordsPerBatch,
maxBatchSize,
-1,
- spark.sessionState.conf.sessionLocalTimeZone).toArray
+ spark.sessionState.conf.sessionLocalTimeZone,
+ codecName,
+ zstdLevel).toArray
}
private def doLocalTableScan(localTableScan: LocalTableScanExec): Array[Array[Byte]] = {
val spark = SparkPlanHelper.sparkSession(localTableScan)
localTableScan.longMetric("numOutputRows").add(localTableScan.rows.size)
sendDriverMetrics(spark.sparkContext, localTableScan.metrics)
+ val (codecName, zstdLevel) = arrowCompressionConf(spark)
KyuubiArrowConverters.toBatchIterator(
localTableScan.rows.iterator,
localTableScan.schema,
spark.sessionState.conf.arrowMaxRecordsPerBatch,
maxBatchSize,
-1,
- spark.sessionState.conf.sessionLocalTimeZone).toArray
+ spark.sessionState.conf.sessionLocalTimeZone,
+ codecName,
+ zstdLevel).toArray
}
/**
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala
index 8e246b3c6e..f56b414274 100644
--- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala
@@ -380,6 +380,54 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp
}
}
+ test("arrow zstd compression round-trips through the JDBC client") {
+ // Engine-side coverage of the compressed Arrow IPC path: the engine compresses batches with
+ // the upstream-named codec config, and the JDBC client transparently decompresses them from
+ // the batch body compression metadata (no server-side codec hint on the wire).
+ withJdbcStatement() { statement =>
+ statement.executeQuery(
+ "set spark.sql.execution.arrow.compression.codec=zstd")
+ val resultSet = statement.executeQuery(
+ "select id, cast(id as string) as name from range(0, 1000)")
+ var count = 0
+ while (resultSet.next()) {
+ // per-row invariant, independent of the row order the client receives
+ assert(resultSet.getString(2) == resultSet.getLong(1).toString)
+ count += 1
+ }
+ assert(count == 1000)
+ }
+ }
+
+ test("arrow zstd compression round-trips when LIMIT cuts across a batch boundary") {
+ // Two input partitions and initialNumPartitions=1 make the LIMIT deterministic: the first
+ // partition yields a full 100-row batch, and the second partition yields another full 100-row
+ // batch with only 50 rows left, so doCollectLimit must slice() the second batch down to 50.
+ withJdbcStatement() { statement =>
+ statement.executeQuery(
+ "set spark.sql.execution.arrow.compression.codec=zstd")
+ statement.executeQuery(
+ s"set ${SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key}=100")
+ statement.executeQuery(
+ "set spark.sql.limit.initialNumPartitions=1")
+ val resultSet = statement.executeQuery(
+ """
+ |select id, cast(id as string) as name from (
+ | select id from range(0, 100, 1, 1)
+ | union all
+ | select id from range(100, 10000, 1, 1)
+ |) t limit 150
+ |""".stripMargin)
+ var count = 0
+ while (resultSet.next()) {
+ // per-row invariant survives compression + slicing, independent of row order
+ assert(resultSet.getString(2) == resultSet.getLong(1).toString)
+ count += 1
+ }
+ assert(count == 150)
+ }
+ }
+
private def checkResultSetFormat(statement: Statement, expectFormat: String): Unit = {
val query =
s"""
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala
new file mode 100644
index 0000000000..1d3e7d1f85
--- /dev/null
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala
@@ -0,0 +1,173 @@
+/*
+ * 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.ByteArrayInputStream
+import java.nio.channels.Channels
+
+import scala.collection.JavaConverters._
+
+import org.apache.arrow.compression.CommonsCompressionFactory
+import org.apache.arrow.flatbuf.CompressionType
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{IntVector, VectorLoader, VectorSchemaRoot}
+import org.apache.arrow.vector.compression.NoCompressionCodec
+import org.apache.arrow.vector.ipc.ReadChannel
+import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, MessageSerializer}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.unsafe.types.UTF8String
+
+import org.apache.kyuubi.KyuubiFunSuite
+
+/**
+ * Test suite for the arrow compression support in [[KyuubiArrowConverters]]:
+ * - the configured zstd compression level must be honored when constructing the codec, i.e.
+ * it must not be silently dropped in favor of the default level
+ * - slice() must keep the decompression factory and the compression codec paired, so that
+ * compressed batches can be sliced when a LIMIT cuts across a batch boundary
+ */
+class KyuubiArrowConvertersSuite extends KyuubiFunSuite {
+
+ private val timeZoneId = "UTC"
+
+ private val schema = StructType(Seq(
+ StructField("id", IntegerType),
+ StructField("name", StringType)))
+
+ // must match what ArrowUtils.toArrowSchema produces for the schema above, so that the
+ // decompression side can load the batches
+ private val arrowSchema = new Schema(List(
+ new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null),
+ new Field("name", FieldType.nullable(new ArrowType.Utf8()), null)).asJava)
+
+ private def rows(count: Int): Iterator[InternalRow] =
+ (0 until count).iterator.map(i => InternalRow(i, UTF8String.fromString(s"name_$i")))
+
+ private def loadRoot(
+ bytes: Array[Byte],
+ allocator: BufferAllocator): (VectorSchemaRoot, ArrowRecordBatch) = {
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ val recordBatch = MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))),
+ allocator)
+ new VectorLoader(root, CommonsCompressionFactory.INSTANCE).load(recordBatch)
+ (root, recordBatch)
+ }
+
+ private def assertRoundTrip(bytes: Array[Byte], expectedRows: Int): Unit = {
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator(s"round-trip-$expectedRows", 0, Long.MaxValue)
+ try {
+ val (root, recordBatch) = loadRoot(bytes, allocator)
+ try {
+ assert(root.getRowCount == expectedRows)
+ val ids = root.getVector("id").asInstanceOf[IntVector]
+ assert(ids.get(0) == 0)
+ assert(ids.get(expectedRows - 1) == expectedRows - 1)
+ } finally {
+ root.close()
+ }
+ recordBatch.close()
+ } finally {
+ allocator.close()
+ }
+ }
+
+ test("zstd compression level is honored and the batches round-trip") {
+ val noneBytes = KyuubiArrowConverters
+ .toBatchIterator(rows(100), schema, 1000, -1, -1, timeZoneId, null, 3)
+ .toArray
+ .head
+ val level1 = KyuubiArrowConverters.slice(schema, timeZoneId, noneBytes, 0, 100, "zstd", 1)
+ val level19 = KyuubiArrowConverters.slice(schema, timeZoneId, noneBytes, 0, 100, "zstd", 19)
+
+ // the configured zstd level must be stored in the batch and therefore produce different
+ // bytes, i.e. the configured level is not silently dropped when the codec is constructed
+ assert(!level1.sameElements(level19))
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator("compression-level", 0, Long.MaxValue)
+ try {
+ val level1Batch = MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(new ByteArrayInputStream(level1))),
+ allocator)
+ try {
+ val bodyCompression = level1Batch.getBodyCompression
+ assert(bodyCompression != null)
+ assert(bodyCompression.getCodec == CompressionType.ZSTD)
+ } finally {
+ level1Batch.close()
+ }
+ } finally {
+ allocator.close()
+ }
+
+ assertRoundTrip(level1, 100)
+ assertRoundTrip(level19, 100)
+ }
+
+ test("slice cuts inside the head batch of compressed results") {
+ val batches = KyuubiArrowConverters
+ .toBatchIterator(rows(150), schema, 100, -1, -1, timeZoneId, "zstd", 3)
+ val batch1 = batches.next() // 100 rows, compressed
+
+ // the LIMIT cuts inside the head batch: 40 rows out of 100.
+ // This is the only slice() call shape that SparkDatasetHelper.doCollectLimit performs in
+ // production: it slices a batch only when the remaining row budget is smaller than the
+ // batch size (size > rest); batches that fit entirely are passed through unsliced.
+ val sliced = KyuubiArrowConverters.slice(schema, timeZoneId, batch1, 0, 40, "zstd", 3)
+ assertRoundTrip(sliced, 40)
+
+ // exhaust the remaining batch so that the iterator closes its allocator
+ assert(batches.hasNext)
+ batches.next()
+ assert(!batches.hasNext)
+ }
+
+ test("reject unsupported lz4 compression codec") {
+ val error = intercept[IllegalArgumentException] {
+ KyuubiArrowConverters.toBatchIterator(rows(1), schema, 100, -1, -1, timeZoneId, "lz4", 3)
+ }
+ assert(error.getMessage.contains("Arrow compression codec lz4 is not supported by Kyuubi"))
+ }
+
+ test("none codec keeps the original uncompressed path") {
+ val bytes = KyuubiArrowConverters
+ .toBatchIterator(rows(100), schema, 1000, -1, -1, timeZoneId, null, 3)
+ .toArray
+ .head
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator("none-codec", 0, Long.MaxValue)
+ try {
+ val batch = MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))),
+ allocator)
+ try {
+ // the none path must produce plain IPC batches marked with NO_COMPRESSION
+ assert(batch.getBodyCompression.getCodec === NoCompressionCodec.COMPRESSION_TYPE)
+ } finally {
+ batch.close()
+ }
+ } finally {
+ allocator.close()
+ }
+ assertRoundTrip(bytes, 100)
+ }
+}
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala
index 791cb12b9c..65f05179d3 100644
--- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala
@@ -17,7 +17,16 @@
package org.apache.spark.sql.kyuubi
+import java.io.ByteArrayInputStream
+import java.nio.channels.Channels
+
+import org.apache.arrow.compression.CommonsCompressionFactory
+import org.apache.arrow.flatbuf.CompressionType
+import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot}
+import org.apache.arrow.vector.ipc.ReadChannel
+import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, MessageSerializer}
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.ArrowUtils
import org.apache.kyuubi.engine.spark.WithSparkSQLEngine
@@ -53,4 +62,68 @@ class SparkDatasetHelperSuite extends WithSparkSQLEngine {
query = "select * from VALUES(1),(2),(3),(4) AS t(id)"
assert(!SparkDatasetHelper.isCommandExec(spark.sql(query)))
}
+
+ test("arrow compression config flows from the session conf into the produced IPC batches") {
+ val codecKey = "spark.sql.execution.arrow.compression.codec"
+ val levelKey = "spark.sql.execution.arrow.compression.zstd.level"
+ // restore the shared engine session conf, otherwise the codec leaks into later suites/tests.
+ // getOption/set/unset are all public RuntimeConfig APIs across Spark 3.3+/4.x.
+ val restoreCodec = spark.conf.getOption(codecKey)
+ val restoreLevel = spark.conf.getOption(levelKey)
+ try {
+ spark.conf.set(codecKey, "zstd")
+ spark.conf.set(levelKey, "1")
+
+ val plan = spark.range(0, 100).queryExecution.executedPlan
+ val level1 = SparkDatasetHelper.toArrowBatchRdd(plan).collect().head
+ spark.conf.set(levelKey, "19")
+ val level19 = SparkDatasetHelper.toArrowBatchRdd(plan).collect().head
+
+ // the configured zstd level must be honored, not silently dropped when the codec is
+ // constructed
+ assert(!level1.sameElements(level19))
+
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("session-conf", 0, Long.MaxValue)
+ try {
+ val batch = deserialize(level19, allocator)
+ try {
+ val compression = batch.getBodyCompression
+ assert(compression != null)
+ assert(compression.getCodec == CompressionType.ZSTD)
+ assert(batch.getLength == 100)
+
+ // the compressed batch must be loadable through the same decompression factory the
+ // client uses
+ val arrowSchema = ArrowUtils.toArrowSchema(plan.schema, "UTC", true, false)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ try {
+ new VectorLoader(root, CommonsCompressionFactory.INSTANCE).load(batch)
+ assert(root.getRowCount == 100)
+ } finally {
+ root.close()
+ }
+ } finally {
+ batch.close()
+ }
+ } finally {
+ allocator.close()
+ }
+ } finally {
+ restoreCodec match {
+ case Some(value) => spark.conf.set(codecKey, value)
+ case None => spark.conf.unset(codecKey)
+ }
+ restoreLevel match {
+ case Some(value) => spark.conf.set(levelKey, value)
+ case None => spark.conf.unset(levelKey)
+ }
+ }
+ }
+
+ private def deserialize(bytes: Array[Byte], allocator: org.apache.arrow.memory.BufferAllocator)
+ : ArrowRecordBatch = {
+ MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))),
+ allocator)
+ }
}
diff --git a/kyuubi-hive-jdbc-shaded/pom.xml b/kyuubi-hive-jdbc-shaded/pom.xml
index 10623940a2..df0ff7608c 100644
--- a/kyuubi-hive-jdbc-shaded/pom.xml
+++ b/kyuubi-hive-jdbc-shaded/pom.xml
@@ -42,6 +42,13 @@
<artifactId>jsr305</artifactId>
<optional>true</optional>
</dependency>
+
+ <!-- zstd-jni is excluded from the shaded jar and optional: zstd consumers provide it. -->
+ <dependency>
+ <groupId>com.github.luben</groupId>
+ <artifactId>zstd-jni</artifactId>
+ <optional>true</optional>
+ </dependency>
</dependencies>
<build>
@@ -58,6 +65,7 @@
<artifactSet>
<excludes>
<exclude>com.google.code.findbugs:jsr305</exclude>
+ <exclude>com.github.luben:zstd-jni</exclude>
</excludes>
</artifactSet>
<filters>
@@ -97,6 +105,7 @@
<shadedPattern>${kyuubi.shade.packageName}.io.netty</shadedPattern>
</relocation>
<relocation>
+ <!-- arrow-compression is bundled and relocated; zstd-jni is not. -->
<pattern>org.apache.arrow</pattern>
<shadedPattern>${kyuubi.shade.packageName}.org.apache.arrow</shadedPattern>
</relocation>
diff --git a/kyuubi-hive-jdbc/pom.xml b/kyuubi-hive-jdbc/pom.xml
index 41d05dcc21..0032c8e620 100644
--- a/kyuubi-hive-jdbc/pom.xml
+++ b/kyuubi-hive-jdbc/pom.xml
@@ -46,6 +46,11 @@
<artifactId>arrow-vector</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ </dependency>
+
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-netty</artifactId>
diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiArrowQueryResultSet.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiArrowQueryResultSet.java
index 163322ccb3..1e8ce12a14 100644
--- a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiArrowQueryResultSet.java
+++ b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiArrowQueryResultSet.java
@@ -26,6 +26,7 @@ import java.util.Iterator;
import java.util.List;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Collectors;
+import org.apache.arrow.compression.CommonsCompressionFactory;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.ipc.ReadChannel;
@@ -368,7 +369,10 @@ public class KyuubiArrowQueryResultSet extends KyuubiArrowBasedResultSet {
TColumn arrowColumn = results.getColumns().get(0);
byte[] batchBytes = arrowColumn.getBinaryVal().getValues().get(0).array();
ArrowRecordBatch recordBatch = loadArrowBatch(batchBytes, allocator);
- VectorLoader vectorLoader = new VectorLoader(root);
+ // The compression codec is detected from the batch body compression metadata carried by
+ // the Arrow IPC message itself, so a single decompression factory works for both
+ // uncompressed and compressed batches -- no server hint about the codec is needed.
+ VectorLoader vectorLoader = new VectorLoader(root, CommonsCompressionFactory.INSTANCE);
vectorLoader.load(recordBatch);
recordBatch.close();
java.util.List<ArrowColumnVector> columns =
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
index 767a8e8397..40b76c12b6 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
@@ -17,6 +17,7 @@
package org.apache.kyuubi.operation
+import java.sql.SQLException
import java.util.{Properties, UUID}
import org.apache.hadoop.fs.{FileSystem, FileUtil, Path}
@@ -26,7 +27,7 @@ import org.apache.kyuubi.{KYUUBI_VERSION, Utils, WithKyuubiServer, WithSimpleDFS
import org.apache.kyuubi.config.KyuubiConf
import org.apache.kyuubi.config.KyuubiConf.KYUUBI_ENGINE_ENV_PREFIX
import org.apache.kyuubi.jdbc.KyuubiHiveDriver
-import org.apache.kyuubi.jdbc.hive.{KyuubiConnection, KyuubiStatement}
+import org.apache.kyuubi.jdbc.hive.{KyuubiArrowQueryResultSet, KyuubiConnection, KyuubiStatement}
import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
import org.apache.kyuubi.session.{KyuubiSessionImpl, SessionHandle}
import org.apache.kyuubi.shaded.hive.service.rpc.thrift.{TExecuteStatementReq, TGetInfoReq, TGetInfoType, TStatusCode}
@@ -38,6 +39,13 @@ class KyuubiOperationPerUserSuite
override protected def jdbcUrl: String = getJdbcUrl
+ // The engine jar no longer bundles arrow-compression; the zstd test supplies it via spark.jars,
+ // while the other tests cover the none path without it.
+ private lazy val arrowCompressionJar: String =
+ java.nio.file.Paths.get(
+ Class.forName("org.apache.arrow.compression.CommonsCompressionFactory")
+ .getProtectionDomain.getCodeSource.getLocation.toURI).toString
+
override protected val conf: KyuubiConf = {
KyuubiConf().set(KyuubiConf.ENGINE_SHARE_LEVEL, "user")
}
@@ -202,6 +210,134 @@ class KyuubiOperationPerUserSuite
}
}
+ test("arrow zstd compression round-trips through the Kyuubi server") {
+ // Compressed Arrow IPC must pass through the server unchanged; the two partitions force
+ // slice(), and the dedicated subdomain engine gets arrow-compression via spark.jars.
+ withSessionConf()(Map.empty)(Map(
+ KyuubiConf.ENGINE_SHARE_LEVEL_SUBDOMAIN.key -> "arrow-zstd",
+ KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow",
+ "spark.jars" -> arrowCompressionJar,
+ "spark.sql.execution.arrow.compression.codec" -> "zstd",
+ "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100",
+ "spark.sql.limit.initialNumPartitions" -> "1")) {
+ withJdbcStatement() { statement =>
+ // prove the codec config reached the engine through the server
+ val codecResult =
+ statement.executeQuery("set spark.sql.execution.arrow.compression.codec")
+ assert(codecResult.next())
+ assert(codecResult.getString("value") === "zstd")
+ val resultSet = statement.executeQuery(
+ """
+ |select id, cast(id as string) as name from (
+ | select id from range(0, 100, 1, 1)
+ | union all
+ | select id from range(100, 10000, 1, 1)
+ |) t limit 150
+ |""".stripMargin)
+ // prove the results come back as Arrow, not silently falling back to Thrift
+ assert(resultSet.isInstanceOf[KyuubiArrowQueryResultSet])
+ val ids = scala.collection.mutable.Set.empty[Long]
+ var count = 0
+ while (resultSet.next()) {
+ // per-row invariant survives compression, transfer and slicing, independent of the row
+ // order the client receives
+ val id = resultSet.getLong(1)
+ assert(resultSet.getString(2) === id.toString)
+ ids += id
+ count += 1
+ }
+ assert(count === 150)
+ assert(ids === (0L until 150L).toSet)
+ }
+ }
+ }
+
+ test("session-level codec switches within one engine when arrow-compression is present") {
+ // The codec is a session-level config; each switch within one long-lived engine must apply.
+ withSessionConf()(Map.empty)(Map(
+ KyuubiConf.ENGINE_SHARE_LEVEL_SUBDOMAIN.key -> "arrow-zstd",
+ KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow",
+ "spark.jars" -> arrowCompressionJar,
+ "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100",
+ "spark.sql.limit.initialNumPartitions" -> "1")) {
+ withJdbcStatement() { statement =>
+ def checkQuery(): Unit = {
+ val resultSet = statement.executeQuery(
+ "select id, cast(id as string) as name from range(0, 1000)")
+ var count = 0
+ while (resultSet.next()) {
+ assert(resultSet.getString(2) === resultSet.getLong(1).toString)
+ count += 1
+ }
+ assert(count === 1000)
+ }
+ checkQuery() // none
+ statement.executeQuery("set spark.sql.execution.arrow.compression.codec=zstd")
+ checkQuery() // zstd
+ statement.executeQuery("set spark.sql.execution.arrow.compression.codec=none")
+ checkQuery() // back to none
+ }
+ }
+ }
+
+ test("session-level zstd either round-trips or fails with a clear arrow-compression error") {
+ // Spark 4.1+ bundles arrow-compression in the Spark distribution, so this suite cannot
+ // control whether the engine classpath contains it. Enabling zstd must either fail with a
+ // clear dependency error (not NoClassDefFoundError) or, when the dependency is present,
+ // keep Arrow results working end-to-end.
+ withSessionConf()(Map.empty)(Map(
+ KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow")) {
+ withJdbcStatement() { statement =>
+ val ok = statement.executeQuery("select id from range(0, 10)")
+ var count = 0
+ while (ok.next()) count += 1
+ assert(count === 10)
+ // the SET result itself is serialized as Arrow, so a missing dependency surfaces at the
+ // SET, while a present dependency is proven by the follow-up query below
+ try {
+ statement.executeQuery("set spark.sql.execution.arrow.compression.codec=zstd")
+ val resultSet = statement.executeQuery("select id from range(0, 10)")
+ val ids = scala.collection.mutable.Set.empty[Long]
+ while (resultSet.next()) ids += resultSet.getLong(1)
+ assert(ids === (0L until 10L).toSet)
+ } catch {
+ case e: SQLException =>
+ assert(
+ e.getMessage.contains("arrow-compression"),
+ s"Unexpected error when enabling zstd: ${e.getMessage}")
+ }
+ }
+ }
+ }
+
+ test("arrow results work without arrow-compression on the engine classpath") {
+ // This engine has no arrow-compression on its classpath; the default none path, including
+ // the slice() path, must still work.
+ withSessionConf()(Map.empty)(Map(
+ KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow",
+ "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100",
+ "spark.sql.limit.initialNumPartitions" -> "1")) {
+ withJdbcStatement() { statement =>
+ val resultSet = statement.executeQuery(
+ """
+ |select id, cast(id as string) as name from (
+ | select id from range(0, 100, 1, 1)
+ | union all
+ | select id from range(100, 10000, 1, 1)
+ |) t limit 150
+ |""".stripMargin)
+ assert(resultSet.isInstanceOf[KyuubiArrowQueryResultSet])
+ val ids = scala.collection.mutable.Set.empty[Long]
+ while (resultSet.next()) {
+ val id = resultSet.getLong(1)
+ assert(resultSet.getString(2) === id.toString)
+ ids += id
+ }
+ assert(ids === (0L until 150L).toSet)
+ }
+ }
+ }
+
test("scala NPE issue with hdfs jar") {
val jarDir = Utils.createTempDir().toFile
val udfCode =
diff --git a/pom.xml b/pom.xml
index 8e9a00db08..80dba97898 100644
--- a/pom.xml
+++ b/pom.xml
@@ -124,6 +124,7 @@
<scala-collection-compat.version>2.12.0</scala-collection-compat.version>
<arrow.version>16.0.0</arrow.version>
+ <zstd-jni.version>1.5.5-11</zstd-jni.version>
<!-- Please don't upgrade the version to 4.10+, it depends on JDK 11 -->
<antlr4.version>4.9.3</antlr4.version>
<antlr.st4.version>4.3.4</antlr.st4.version>
@@ -348,6 +349,23 @@
<artifactId>arrow-vector</artifactId>
<version>${arrow.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ <version>${arrow.version}</version>
+ <!-- Kyuubi rejects lz4, so the commons-compress-backed LZ4 codec is not needed. -->
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.commons</groupId>
+ <artifactId>commons-compress</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>com.github.luben</groupId>
+ <artifactId>zstd-jni</artifactId>
+ <version>${zstd-jni.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-netty</artifactId>