From 0dddffe643f0c798c63da51c16851338181177ac Mon Sep 17 00:00:00 2001 From: byronwang Date: Sat, 15 Aug 2026 04:33:03 +0800 Subject: [PATCH 1/9] [KYUUBI #7635][SPARK] Support Arrow IPC compression for query results (zstd) --- LICENSE-binary | 1 + dev/dependencyList | 4 + externals/kyuubi-spark-sql-engine/pom.xml | 6 + .../arrow/KyuubiArrowConverters.scala | 66 ++++++-- .../spark/sql/kyuubi/SparkDatasetHelper.scala | 60 +++++++- .../SparkArrowbasedOperationSuite.scala | 40 +++++ .../arrow/KyuubiArrowConvertersSuite.scala | 142 ++++++++++++++++++ .../sql/kyuubi/SparkDatasetHelperSuite.scala | 73 +++++++++ kyuubi-hive-jdbc/pom.xml | 5 + .../jdbc/hive/KyuubiArrowQueryResultSet.java | 6 +- pom.xml | 5 + 11 files changed, 391 insertions(+), 17 deletions(-) create mode 100644 externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala diff --git a/LICENSE-binary b/LICENSE-binary index df7580ad899..edf716de60b 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -320,6 +320,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 64af7549664..0d92f1a255c 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 @@ -30,6 +31,8 @@ arrow-vector/16.0.0//arrow-vector-16.0.0.jar checker-qual/3.42.0//checker-qual-3.42.0.jar classgraph/4.8.138//classgraph-4.8.138.jar commons-codec/1.17.1//commons-codec-1.17.1.jar +commons-compress/1.26.0//commons-compress-1.26.0.jar +commons-io/2.16.1//commons-io-2.16.1.jar commons-lang3/3.18.0//commons-lang3-3.18.0.jar error_prone_annotations/2.36.0//error_prone_annotations-2.36.0.jar failsafe/3.3.2//failsafe-3.3.2.jar @@ -188,3 +191,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 67308284bd6..553dfa6ac40 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -37,6 +37,12 @@ ${project.version} + + org.apache.arrow + arrow-compression + ${arrow.version} + + org.apache.kyuubi kyuubi-events_${scala.binary.version} 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 b8f590ddfbd..e8e526caa71 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 @@ -23,7 +23,9 @@ import java.nio.channels.Channels import scala.collection.JavaConverters._ import scala.collection.mutable.ArrayBuffer +import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionCodec} import org.apache.arrow.vector._ +import org.apache.arrow.vector.compression.{CompressionCodec, 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 +41,28 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) + /** + * Create the Arrow compression codec for the given codec name. Returns the no-compression codec + * when the codec name is "none" or null. + * + * The zstd codec is constructed directly with the configured compression level rather than + * through a [[CompressionCodec.Factory]]: the factory overloads built from the codec type enum + * cannot carry a compression level and silently fall back to the default one, losing the + * user-configured level. This follows the latest Spark upstream behavior + * (ArrowCompressionUtils). + */ + private def createCodec(codecName: String, zstdLevel: Int): CompressionCodec = { + if (codecName == null || codecName == "none") { + NoCompressionCodec.INSTANCE + } else { + codecName match { + case "zstd" => new ZstdCompressionCodec(zstdLevel) + case other => + throw new IllegalArgumentException(s"Unsupported arrow compression codec: $other") + } + } + } + /** * this method is to slice the input Arrow record batch byte array `bytes`, starting from `start` * and taking `length` number of elements. @@ -48,7 +72,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 +91,21 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { val recordBatch = MessageSerializer.deserializeRecordBatch( new ReadChannel(Channels.newChannel(in)), sliceAllocator) - val vectorLoader = new VectorLoader(vectorSchemaRoot) + // The decompression factory must always be provided: the codec type is read from the batch + // body compression metadata carried by the Arrow IPC message itself, so no extra hint is + // needed. Uncompressed batches do not invoke the factory at all. + val vectorLoader = new VectorLoader(vectorSchemaRoot, CommonsCompressionFactory.INSTANCE) vectorLoader.load(recordBatch) recordBatch.close() slicedVectorSchemaRoot = vectorSchemaRoot.slice(start, length) - val unloader = new VectorUnloader(slicedVectorSchemaRoot) + // The unloader and the loader above must be kept in sync: compressed batches sliced and + // re-serialized here would fail to load on the client if the compression codec was dropped. + val unloader = new VectorUnloader( + slicedVectorSchemaRoot, + true, + createCodec(codecName, zstdLevel), + true) val writeChannel = new WriteChannel(Channels.newChannel(out)) val batch = unloader.getRecordBatch() MessageSerializer.serialize(writeChannel, batch) @@ -119,7 +154,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 +202,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxRecordsPerBatch, maxEstimatedBatchSize, n, - timeZoneId) + timeZoneId, + codecName, + zstdLevel) batches.map(b => b -> batches.rowCountInLastBatch).toArray }, partsToScan) @@ -200,7 +239,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 +249,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxEstimatedBatchSize, limit, timeZoneId, - TaskContext.get) + TaskContext.get, + codecName, + zstdLevel) } /** @@ -226,7 +269,9 @@ 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) @@ -237,7 +282,10 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { Long.MaxValue) private val root = VectorSchemaRoot.create(arrowSchema, allocator) - protected val unloader = new VectorUnloader(root) + // Always use the compression-aware 4-arg constructor, matching the latest Spark upstream + // ArrowConverters. includeNullCount=true keeps the null count in the batch header, the same + // as the original 1-arg constructor used by the no-compression path. + protected val unloader = new VectorUnloader(root, true, createCodec(codecName, zstdLevel), true) 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 c503d780341..1b6ae3721fe 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, Dataset, 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} @@ -77,6 +79,28 @@ object SparkDatasetHelper extends Logging { toArrowBatchRddMethod.bind(ds).invoke() } + /** + * Read the arrow compression codec config from the session conf, reusing the Spark upstream + * configuration keys (`spark.sql.execution.arrow.compression.codec` and + * `spark.sql.execution.arrow.compression.zstd.level`, introduced in Spark 4.1). On Spark < 4.1 + * these keys are not registered as typed SQLConf entries, so they are read as raw strings: a + * session-level SET lands in the SQLConf settings map and takes precedence, otherwise the + * engine-level kyuubi-defaults.conf value applies, otherwise the default. + * Returns (codecName, zstdLevel). + */ + private def arrowCompressionConf(spark: SparkSession): (String, Int) = { + // Lower-case the codec name to match Spark upstream, whose typed entry applies + // `.transform(_.toLowerCase(Locale.ROOT))`; reading the key as a raw string here would + // otherwise make `ZSTD`/`Zstd` reject on Spark < 4.1 while accepted natively on 4.1+. + 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)]]. * Convert to an RDD of serialized ArrowRecordBatches. @@ -89,6 +113,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 +121,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 Dataset#toArrowBatchRdd (reflective) path uses the vanilla Spark ArrowConverters + // which does not apply the Kyuubi compression codec. + toArrowBatchRdd(df.queryExecution.executedPlan).toLocalIterator } } @@ -174,12 +204,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 +226,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 +245,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 8e246b3c6e6..e0ea370cacf 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,46 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp } } + test("arrow zstd compression round-trips through the JDBC client") { + // End-to-end 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") { + // The slice() path re-serializes the tail batch produced by doCollectLimit. With a batch size + // of 100 and LIMIT 150, the second batch is sliced from 100 down to 50 rows, exercising the + // compression-aware VectorLoader/VectorUnloader pairing on a compressed, sliced batch. + 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") + val resultSet = statement.executeQuery( + "select id, cast(id as string) as name from range(0, 10000) limit 150") + 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 00000000000..4a6a3a230c5 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala @@ -0,0 +1,142 @@ +/* + * 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.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) + } +} 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 791cb12b9c5..65f05179d3b 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/pom.xml b/kyuubi-hive-jdbc/pom.xml index 41d05dcc215..0032c8e620d 100644 --- a/kyuubi-hive-jdbc/pom.xml +++ b/kyuubi-hive-jdbc/pom.xml @@ -46,6 +46,11 @@ arrow-vector + + org.apache.arrow + arrow-compression + + org.apache.arrow arrow-memory-netty 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 163322ccb32..1e8ce12a14a 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.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 boolean next() throws SQLException { 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 columns = diff --git a/pom.xml b/pom.xml index a7d02e9f759..5a2be6a7ef5 100644 --- a/pom.xml +++ b/pom.xml @@ -348,6 +348,11 @@ arrow-vector ${arrow.version} + + org.apache.arrow + arrow-compression + ${arrow.version} + org.apache.arrow arrow-memory-netty From c041d1cd622fed868242036e37abd7c879ec1f20 Mon Sep 17 00:00:00 2001 From: byronwang Date: Sun, 16 Aug 2026 17:53:43 +0800 Subject: [PATCH 2/9] [KYUUBI #7635][SPARK] Bundle relocated arrow-compression into the engine jar Spark < 4.1 runtimes do not ship arrow-compression, so the compressed Arrow result path (and even the default none path when a LIMIT slices a batch, which always constructs VectorLoader with a decompression factory) would fail with NoClassDefFoundError unless users added jars manually. Bundle arrow-compression into the shaded engine jar, relocated under org.apache.kyuubi.shade.org.apache.arrow.compression, so it is self-contained on every supported Spark version and cannot conflict with the arrow-compression jar that Spark 4.1+ ships. The Arrow vector/memory/format modules and zstd-jni/commons-compress/ commons-io are excluded from shading: every supported Spark runtime provides them, and excluding them avoids classpath conflicts and native library relocation issues. --- LICENSE-binary | 1 + externals/kyuubi-spark-sql-engine/pom.xml | 30 +++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/LICENSE-binary b/LICENSE-binary index edf716de60b..c2c53faafb6 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 diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml index 553dfa6ac40..ece207871f3 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -257,6 +257,7 @@ io.netty:* io.perfmark:perfmark-api io.vertx:* + org.apache.arrow:arrow-compression org.apache.kyuubi:* org.checkerframework:checker-qual org.codehaus.mojo:animal-sniffer-annotations @@ -269,6 +270,23 @@ io.netty:netty-transport-*-kqueue io.netty:netty-transport-*-io_uring io.netty:netty-transport-native-epoll:*:linux-riscv64 + + org.apache.arrow:arrow-format + org.apache.arrow:arrow-memory-core + org.apache.arrow:arrow-memory-netty + org.apache.arrow:arrow-memory-netty-buffer-patch + org.apache.arrow:arrow-vector + com.github.luben:zstd-jni + commons-codec:commons-codec + commons-io:commons-io + org.apache.commons:commons-compress + org.apache.commons:commons-lang3 + org.immutables:value @@ -289,6 +307,7 @@ NOTICE.txt mozilla/** **/module-info.class + arrow-git.properties @@ -411,6 +430,17 @@ net.bytebuddy ${kyuubi.shade.packageName}.net.bytebuddy + + + org.apache.arrow.compression + ${kyuubi.shade.packageName}.org.apache.arrow.compression + From ed8339839c7febe317bee5a976b861316931c16d Mon Sep 17 00:00:00 2001 From: byronwang Date: Tue, 18 Aug 2026 17:11:37 +0800 Subject: [PATCH 3/9] [KYUUBI #7635][SPARK] Reject lz4 codec and drop reflective Dataset#toArrowBatchRdd The arrow-compression library is now relocated into the engine jar, so the reflective fallback that invoked the vanilla Spark Dataset#toArrowBatchRdd is no longer needed; use the plan-based helper directly so the Kyuubi zstd codec applies. Also reject lz4 explicitly instead of silently falling back to no compression, and add server-side round-trip coverage. --- externals/kyuubi-spark-sql-engine/pom.xml | 1 - .../arrow/KyuubiArrowConverters.scala | 26 ++++++++------ .../spark/sql/kyuubi/SparkDatasetHelper.scala | 36 ++++--------------- .../SparkArrowbasedOperationSuite.scala | 2 +- .../KyuubiOperationPerUserSuite.scala | 29 +++++++++++++++ 5 files changed, 53 insertions(+), 41 deletions(-) diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml index ece207871f3..0fc872c7eb2 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -307,7 +307,6 @@ NOTICE.txt mozilla/** **/module-info.class - arrow-git.properties 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 e8e526caa71..3fd148bd760 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 @@ -42,8 +42,10 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) /** - * Create the Arrow compression codec for the given codec name. Returns the no-compression codec - * when the codec name is "none" or null. + * Create the Arrow compression codec for the given codec name. Only "none" (or null, the + * no-compression codec) and "zstd" are supported; "lz4" is accepted by Spark upstream + * (SPARK-54134, 4.1.0) but not yet by Kyuubi, so it is rejected explicitly instead of + * silently falling back to no compression. * * The zstd codec is constructed directly with the configured compression level rather than * through a [[CompressionCodec.Factory]]: the factory overloads built from the codec type enum @@ -52,14 +54,18 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { * (ArrowCompressionUtils). */ private def createCodec(codecName: String, zstdLevel: Int): CompressionCodec = { - if (codecName == null || codecName == "none") { - NoCompressionCodec.INSTANCE - } else { - codecName match { - case "zstd" => new ZstdCompressionCodec(zstdLevel) - case other => - throw new IllegalArgumentException(s"Unsupported arrow compression codec: $other") - } + codecName match { + case null | "none" => + NoCompressionCodec.INSTANCE + case "zstd" => + new ZstdCompressionCodec(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") } } 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 1b6ae3721fe..4fd4ecef352 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 @@ -25,7 +25,7 @@ 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, SparkSession} +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} @@ -38,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 { @@ -65,33 +64,12 @@ 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 the arrow compression codec config from the session conf, reusing the Spark upstream - * configuration keys (`spark.sql.execution.arrow.compression.codec` and - * `spark.sql.execution.arrow.compression.zstd.level`, introduced in Spark 4.1). On Spark < 4.1 - * these keys are not registered as typed SQLConf entries, so they are read as raw strings: a - * session-level SET lands in the SQLConf settings map and takes precedence, otherwise the - * engine-level kyuubi-defaults.conf value applies, otherwise the default. - * Returns (codecName, zstdLevel). + * 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) = { - // Lower-case the codec name to match Spark upstream, whose typed entry applies - // `.transform(_.toLowerCase(Locale.ROOT))`; reading the key as a raw string here would - // otherwise make `ZSTD`/`Zstd` reject on Spark < 4.1 while accepted natively on 4.1+. + // Lowercase the codec name to match Spark upstream. val codecName = spark.sessionState.conf.getConfString( "spark.sql.execution.arrow.compression.codec", "none").toLowerCase(Locale.ROOT) @@ -102,7 +80,7 @@ object SparkDatasetHelper extends Logging { } /** - * 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]] = { @@ -130,8 +108,8 @@ object SparkDatasetHelper extends Logging { def toArrowBatchLocalIterator(df: DataFrame): Iterator[Array[Byte]] = { withNewExecutionId(df) { // use the plan-based toArrowBatchRdd so that the arrow compression codec takes effect; - // the Dataset#toArrowBatchRdd (reflective) path uses the vanilla Spark ArrowConverters - // which does not apply the Kyuubi compression codec. + // the vanilla Spark Dataset#toArrowBatchRdd uses the upstream ArrowConverters, which does + // not apply the Kyuubi compression codec. toArrowBatchRdd(df.queryExecution.executedPlan).toLocalIterator } } 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 e0ea370cacf..a626eab23e7 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 @@ -381,7 +381,7 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp } test("arrow zstd compression round-trips through the JDBC client") { - // End-to-end coverage of the compressed Arrow IPC path: the engine compresses batches with + // 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 => 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 de491e03f21..686832945e3 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 @@ -202,6 +202,35 @@ class KyuubiOperationPerUserSuite } } + test("arrow zstd compression round-trips through the Kyuubi server") { + // Full-chain coverage: JDBC client <-> Kyuubi server <-> Spark engine. The engine compresses + // Arrow IPC batches with zstd, the server forwards the payload untouched, and the JDBC client + // transparently decompresses it from the Arrow IPC body compression metadata. LIMIT 150 cuts + // through the second 100-row batch, so the tail batch is sliced and re-serialized on the + // engine side, exercising the compression-aware slice path. + withSessionConf()(Map.empty)(Map( + KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow", + "spark.sql.execution.arrow.compression.codec" -> "zstd", + "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100")) { + withJdbcStatement() { statement => + val resultSet = statement.executeQuery( + "select id, cast(id as string) as name from range(0, 10000) limit 150") + 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("scala NPE issue with hdfs jar") { val jarDir = Utils.createTempDir().toFile val udfCode = From c9afd969a2c3a71df0cb51cc29532a20a4ae0d14 Mon Sep 17 00:00:00 2001 From: byronwang Date: Tue, 18 Aug 2026 20:32:33 +0800 Subject: [PATCH 4/9] [KYUUBI #7635][SPARK][FOLLOWUP] Make LIMIT slice tests deterministic and fail fast on unsupported codec --- .../arrow/KyuubiArrowConverters.scala | 17 +++--------- .../SparkArrowbasedOperationSuite.scala | 16 ++++++++--- .../arrow/KyuubiArrowConvertersSuite.scala | 7 +++++ .../KyuubiOperationPerUserSuite.scala | 27 +++++++++++++------ 4 files changed, 42 insertions(+), 25 deletions(-) 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 3fd148bd760..b00c49dfa5a 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 @@ -41,18 +41,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) - /** - * Create the Arrow compression codec for the given codec name. Only "none" (or null, the - * no-compression codec) and "zstd" are supported; "lz4" is accepted by Spark upstream - * (SPARK-54134, 4.1.0) but not yet by Kyuubi, so it is rejected explicitly instead of - * silently falling back to no compression. - * - * The zstd codec is constructed directly with the configured compression level rather than - * through a [[CompressionCodec.Factory]]: the factory overloads built from the codec type enum - * cannot carry a compression level and silently fall back to the default one, losing the - * user-configured level. This follows the latest Spark upstream behavior - * (ArrowCompressionUtils). - */ + // Only "none" (no compression) and "zstd" codecs are supported. private def createCodec(codecName: String, zstdLevel: Int): CompressionCodec = { codecName match { case null | "none" => @@ -281,6 +270,8 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { 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 compressionCodec = createCodec(codecName, zstdLevel) private val allocator = ArrowUtils.rootAllocator.newChildAllocator( s"to${this.getClass.getSimpleName}", @@ -291,7 +282,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { // Always use the compression-aware 4-arg constructor, matching the latest Spark upstream // ArrowConverters. includeNullCount=true keeps the null count in the batch header, the same // as the original 1-arg constructor used by the no-compression path. - protected val unloader = new VectorUnloader(root, true, createCodec(codecName, zstdLevel), true) + protected val unloader = new VectorUnloader(root, true, compressionCodec, true) protected val arrowWriter = ArrowWriter.create(root) Option(context).foreach { 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 a626eab23e7..f56b414274f 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 @@ -400,16 +400,24 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp } test("arrow zstd compression round-trips when LIMIT cuts across a batch boundary") { - // The slice() path re-serializes the tail batch produced by doCollectLimit. With a batch size - // of 100 and LIMIT 150, the second batch is sliced from 100 down to 50 rows, exercising the - // compression-aware VectorLoader/VectorUnloader pairing on a compressed, sliced batch. + // 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 range(0, 10000) limit 150") + """ + |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 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 index 4a6a3a230c5..0ee62748da9 100644 --- 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 @@ -139,4 +139,11 @@ class KyuubiArrowConvertersSuite extends KyuubiFunSuite { 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")) + } } 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 686832945e3..1474afecaf9 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 @@ -26,7 +26,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} @@ -203,18 +203,29 @@ class KyuubiOperationPerUserSuite } test("arrow zstd compression round-trips through the Kyuubi server") { - // Full-chain coverage: JDBC client <-> Kyuubi server <-> Spark engine. The engine compresses - // Arrow IPC batches with zstd, the server forwards the payload untouched, and the JDBC client - // transparently decompresses it from the Arrow IPC body compression metadata. LIMIT 150 cuts - // through the second 100-row batch, so the tail batch is sliced and re-serialized on the - // engine side, exercising the compression-aware slice path. + // Verify that compressed Arrow IPC passes through Kyuubi Server unchanged. The two input + // partitions force LIMIT to slice the second 100-row batch down to 50 rows. withSessionConf()(Map.empty)(Map( KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow", "spark.sql.execution.arrow.compression.codec" -> "zstd", - "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100")) { + "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 range(0, 10000) limit 150") + """ + |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()) { From 024f3c68517efc1682c4f1cf5dc1c55b789bc4f8 Mon Sep 17 00:00:00 2001 From: byronwang Date: Tue, 18 Aug 2026 22:09:45 +0800 Subject: [PATCH 5/9] [KYUUBI #7635][SPARK][FOLLOWUP] Drop commons-compress from the arrow-compression dependency chain --- dev/dependencyList | 2 -- externals/kyuubi-spark-sql-engine/pom.xml | 9 ++++----- pom.xml | 7 +++++++ 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/dev/dependencyList b/dev/dependencyList index 0d92f1a255c..139365221fa 100644 --- a/dev/dependencyList +++ b/dev/dependencyList @@ -31,8 +31,6 @@ arrow-vector/16.0.0//arrow-vector-16.0.0.jar checker-qual/3.42.0//checker-qual-3.42.0.jar classgraph/4.8.138//classgraph-4.8.138.jar commons-codec/1.17.1//commons-codec-1.17.1.jar -commons-compress/1.26.0//commons-compress-1.26.0.jar -commons-io/2.16.1//commons-io-2.16.1.jar commons-lang3/3.18.0//commons-lang3-3.18.0.jar error_prone_annotations/2.36.0//error_prone_annotations-2.36.0.jar failsafe/3.3.2//failsafe-3.3.2.jar diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml index 0fc872c7eb2..11861fda8c1 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -271,10 +271,10 @@ io.netty:netty-transport-*-io_uring io.netty:netty-transport-native-epoll:*:linux-riscv64 org.apache.arrow:arrow-format org.apache.arrow:arrow-memory-core @@ -284,7 +284,6 @@ com.github.luben:zstd-jni commons-codec:commons-codec commons-io:commons-io - org.apache.commons:commons-compress org.apache.commons:commons-lang3 org.immutables:value diff --git a/pom.xml b/pom.xml index 5a2be6a7ef5..27f78b288c6 100644 --- a/pom.xml +++ b/pom.xml @@ -352,6 +352,13 @@ org.apache.arrow arrow-compression ${arrow.version} + + + + org.apache.commons + commons-compress + + org.apache.arrow From c761a93ac5f64a755c89f039e68e79309fe35b24 Mon Sep 17 00:00:00 2001 From: byronwang Date: Wed, 19 Aug 2026 12:02:50 +0800 Subject: [PATCH 6/9] [KYUUBI #7635][SPARK][FOLLOWUP] Make arrow-compression optional on the engine side and exclude zstd-jni from the shaded client --- externals/kyuubi-spark-sql-engine/pom.xml | 14 +--- .../arrow/ArrowCompressionSupport.scala | 33 ++++++++ .../arrow/KyuubiArrowConverters.scala | 79 +++++++++++-------- .../arrow/KyuubiArrowConvertersSuite.scala | 24 ++++++ kyuubi-hive-jdbc-shaded/pom.xml | 9 +++ .../KyuubiOperationPerUserSuite.scala | 41 +++++++++- pom.xml | 6 ++ 7 files changed, 157 insertions(+), 49 deletions(-) create mode 100644 externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml index 11861fda8c1..cddc27c1025 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -37,10 +37,12 @@ ${project.version} + org.apache.arrow arrow-compression ${arrow.version} + true @@ -257,7 +259,6 @@ io.netty:* io.perfmark:perfmark-api io.vertx:* - org.apache.arrow:arrow-compression org.apache.kyuubi:* org.checkerframework:checker-qual org.codehaus.mojo:animal-sniffer-annotations @@ -428,17 +429,6 @@ net.bytebuddy ${kyuubi.shade.packageName}.net.bytebuddy - - - org.apache.arrow.compression - ${kyuubi.shade.packageName}.org.apache.arrow.compression - diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala new file mode 100644 index 00000000000..75573e5f645 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.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 ArrowCompressionSupport { + + 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 b00c49dfa5a..d81204e6613 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 @@ -23,9 +23,8 @@ import java.nio.channels.Channels import scala.collection.JavaConverters._ import scala.collection.mutable.ArrayBuffer -import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionCodec} import org.apache.arrow.vector._ -import org.apache.arrow.vector.compression.{CompressionCodec, NoCompressionCodec} +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 @@ -41,23 +40,6 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) - // Only "none" (no compression) and "zstd" codecs are supported. - private def createCodec(codecName: String, zstdLevel: Int): CompressionCodec = { - codecName match { - case null | "none" => - NoCompressionCodec.INSTANCE - case "zstd" => - new ZstdCompressionCodec(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") - } - } - /** * this method is to slice the input Arrow record batch byte array `bytes`, starting from `start` * and taking `length` number of elements. @@ -86,21 +68,33 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { val recordBatch = MessageSerializer.deserializeRecordBatch( new ReadChannel(Channels.newChannel(in)), sliceAllocator) - // The decompression factory must always be provided: the codec type is read from the batch - // body compression metadata carried by the Arrow IPC message itself, so no extra hint is - // needed. Uncompressed batches do not invoke the factory at all. - val vectorLoader = new VectorLoader(vectorSchemaRoot, CommonsCompressionFactory.INSTANCE) + // 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) { + ArrowCompressionSupport.createLoader(vectorSchemaRoot) + } else { + new VectorLoader(vectorSchemaRoot) + } vectorLoader.load(recordBatch) recordBatch.close() slicedVectorSchemaRoot = vectorSchemaRoot.slice(start, length) - // The unloader and the loader above must be kept in sync: compressed batches sliced and - // re-serialized here would fail to load on the client if the compression codec was dropped. - val unloader = new VectorUnloader( - slicedVectorSchemaRoot, - true, - createCodec(codecName, zstdLevel), - true) + // 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" => + ArrowCompressionSupport.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) @@ -271,7 +265,19 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { protected val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true, false) // Validate the codec before allocating Arrow buffers, so an unsupported codec fails fast. - private val compressionCodec = createCodec(codecName, zstdLevel) + private val compressionEnabled = codecName match { + case null | "none" => + false + case "zstd" => + 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}", @@ -279,10 +285,13 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { Long.MaxValue) private val root = VectorSchemaRoot.create(arrowSchema, allocator) - // Always use the compression-aware 4-arg constructor, matching the latest Spark upstream - // ArrowConverters. includeNullCount=true keeps the null count in the batch header, the same - // as the original 1-arg constructor used by the no-compression path. - protected val unloader = new VectorUnloader(root, true, compressionCodec, true) + // The none path keeps the original 1-arg constructor and stays free of arrow-compression. + protected val unloader = + if (compressionEnabled) { + ArrowCompressionSupport.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/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 index 0ee62748da9..1d3e7d1f85c 100644 --- 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 @@ -26,6 +26,7 @@ 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} @@ -146,4 +147,27 @@ class KyuubiArrowConvertersSuite extends KyuubiFunSuite { } 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/kyuubi-hive-jdbc-shaded/pom.xml b/kyuubi-hive-jdbc-shaded/pom.xml index 10623940a2b..df0ff7608c1 100644 --- a/kyuubi-hive-jdbc-shaded/pom.xml +++ b/kyuubi-hive-jdbc-shaded/pom.xml @@ -42,6 +42,13 @@ jsr305 true + + + + com.github.luben + zstd-jni + true + @@ -58,6 +65,7 @@ com.google.code.findbugs:jsr305 + com.github.luben:zstd-jni @@ -97,6 +105,7 @@ ${kyuubi.shade.packageName}.io.netty + org.apache.arrow ${kyuubi.shade.packageName}.org.apache.arrow 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 1474afecaf9..02621d1ede0 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 @@ -38,6 +38,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") } @@ -203,10 +210,12 @@ class KyuubiOperationPerUserSuite } test("arrow zstd compression round-trips through the Kyuubi server") { - // Verify that compressed Arrow IPC passes through Kyuubi Server unchanged. The two input - // partitions force LIMIT to slice the second 100-row batch down to 50 rows. + // 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")) { @@ -242,6 +251,34 @@ class KyuubiOperationPerUserSuite } } + 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 27f78b288c6..866c2a1ed60 100644 --- a/pom.xml +++ b/pom.xml @@ -124,6 +124,7 @@ 2.12.0 16.0.0 + 1.5.5-11 4.9.3 4.3.4 @@ -360,6 +361,11 @@ + + com.github.luben + zstd-jni + ${zstd-jni.version} + org.apache.arrow arrow-memory-netty From 347e38aabce5aa387aa590087e8c92ccfe6b7b16 Mon Sep 17 00:00:00 2001 From: byronwang Date: Wed, 19 Aug 2026 14:48:45 +0800 Subject: [PATCH 7/9] [KYUUBI #7635][SPARK][FOLLOWUP] Detect arrow-compression capability at runtime before using zstd --- .../arrow/KyuubiArrowConverters.scala | 28 +++++++++++ .../KyuubiOperationPerUserSuite.scala | 48 +++++++++++++++++++ 2 files changed, 76 insertions(+) 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 d81204e6613..32cd80f50f9 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 @@ -40,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. @@ -73,6 +98,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { recordBatch.getBodyCompression.getCodec != NoCompressionCodec.COMPRESSION_TYPE val vectorLoader = if (compressed) { + requireArrowCompression() ArrowCompressionSupport.createLoader(vectorSchemaRoot) } else { new VectorLoader(vectorSchemaRoot) @@ -86,6 +112,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { case null | "none" => new VectorUnloader(slicedVectorSchemaRoot) case "zstd" => + requireArrowCompression() ArrowCompressionSupport.createZstdUnloader(slicedVectorSchemaRoot, zstdLevel) case "lz4" => throw new IllegalArgumentException( @@ -269,6 +296,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { case null | "none" => false case "zstd" => + requireArrowCompression() true case "lz4" => throw new IllegalArgumentException( 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 02621d1ede0..b11f6433edc 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} @@ -251,6 +252,53 @@ class KyuubiOperationPerUserSuite } } + 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 without arrow-compression fails with a clear error") { + // Enabling zstd on an engine whose classpath has no arrow-compression must fail with a clear + // dependency error, not NoClassDefFoundError. + 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 the failure surfaces at the SET + val e = intercept[SQLException] { + statement.executeQuery("set spark.sql.execution.arrow.compression.codec=zstd") + } + assert(e.getMessage.contains("arrow-compression")) + } + } + } + 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. From 34e8ac6f383b486c7349c80bf198894c921d1e68 Mon Sep 17 00:00:00 2001 From: byronwang Date: Wed, 19 Aug 2026 17:38:46 +0800 Subject: [PATCH 8/9] [KYUUBI #7635][SPARK][FOLLOWUP] Prefix Arrow compression helper with Kyuubi --- ...ionSupport.scala => KyuubiArrowCompressionSupport.scala} | 2 +- .../spark/sql/execution/arrow/KyuubiArrowConverters.scala | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) rename externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/{ArrowCompressionSupport.scala => KyuubiArrowCompressionSupport.scala} (96%) diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala similarity index 96% rename from externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala rename to externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala index 75573e5f645..478afcb6d8a 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionSupport.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala @@ -21,7 +21,7 @@ import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionC import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader} /** Isolates the optional arrow-compression dependency so the uncompressed path never loads it. */ -private[sql] object ArrowCompressionSupport { +private[sql] object KyuubiArrowCompressionSupport { def createLoader(root: VectorSchemaRoot): VectorLoader = { new VectorLoader(root, CommonsCompressionFactory.INSTANCE) 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 32cd80f50f9..9f690357b8a 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 @@ -99,7 +99,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { val vectorLoader = if (compressed) { requireArrowCompression() - ArrowCompressionSupport.createLoader(vectorSchemaRoot) + KyuubiArrowCompressionSupport.createLoader(vectorSchemaRoot) } else { new VectorLoader(vectorSchemaRoot) } @@ -113,7 +113,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { new VectorUnloader(slicedVectorSchemaRoot) case "zstd" => requireArrowCompression() - ArrowCompressionSupport.createZstdUnloader(slicedVectorSchemaRoot, zstdLevel) + KyuubiArrowCompressionSupport.createZstdUnloader(slicedVectorSchemaRoot, zstdLevel) case "lz4" => throw new IllegalArgumentException( "Arrow compression codec lz4 is not supported by Kyuubi; " + @@ -316,7 +316,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { // The none path keeps the original 1-arg constructor and stays free of arrow-compression. protected val unloader = if (compressionEnabled) { - ArrowCompressionSupport.createZstdUnloader(root, zstdLevel) + KyuubiArrowCompressionSupport.createZstdUnloader(root, zstdLevel) } else { new VectorUnloader(root) } From 8911c395f4e487092bc2c55c8619477347aeeb87 Mon Sep 17 00:00:00 2001 From: byronwang Date: Fri, 21 Aug 2026 00:44:46 +0800 Subject: [PATCH 9/9] [KYUUBI #7635][SPARK][TEST] Make session-level zstd test independent of the engine classpath Spark 4.1+ bundles arrow-compression in the Spark distribution, so the engine classpath may or may not contain it depending on the Spark version under test. The old test hard-coded the missing-dependency premise and expected the zstd SET to fail, which does not hold on Spark 4.1/4.2 and made the CI jobs fail with 'Expected exception java.sql.SQLException to be thrown, but no exception was thrown'. Assert the observable behavior of both cases instead: a clear arrow-compression dependency error when the jar is absent, and an end-to-end zstd round-trip when it is present. --- .../KyuubiOperationPerUserSuite.scala | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) 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 b11f6433edc..90f982d16b2 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 @@ -280,9 +280,11 @@ class KyuubiOperationPerUserSuite } } - test("session-level zstd without arrow-compression fails with a clear error") { - // Enabling zstd on an engine whose classpath has no arrow-compression must fail with a clear - // dependency error, not NoClassDefFoundError. + 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 => @@ -290,11 +292,20 @@ class KyuubiOperationPerUserSuite var count = 0 while (ok.next()) count += 1 assert(count === 10) - // the SET result itself is serialized as Arrow, so the failure surfaces at the SET - val e = intercept[SQLException] { + // 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}") } - assert(e.getMessage.contains("arrow-compression")) } } }