[KYUUBI #7635][SPARK] Support Zstd compression for Arrow IPC query results

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
 ------------------------------------

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

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>

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)
+  }
+}

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,15 +255,19 @@ 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,
       maxRecordsPerBatch,
       maxEstimatedBatchSize,
       limit,
       timeZoneId,
-      TaskContext.get)
+      TaskContext.get,
+      codecName,
+      zstdLevel)
   }
 
   /**
@@ -226,18 +285,41 @@ 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}",
         0,
         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 {