diff --git a/LICENSE-binary b/LICENSE-binary index df7580ad899..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 @@ -320,6 +321,7 @@ com.thoughtworks.paranamer:paranamer com.google.protobuf:protobuf-java-util com.google.protobuf:protobuf-java org.postgresql:postgresql +com.github.luben:zstd-jni Eclipse Distribution License - v 1.0 ------------------------------------ diff --git a/dev/dependencyList b/dev/dependencyList index 64af7549664..139365221fa 100644 --- a/dev/dependencyList +++ b/dev/dependencyList @@ -22,6 +22,7 @@ annotations/4.1.1.4//annotations-4.1.1.4.jar antlr-runtime/3.5.3//antlr-runtime-3.5.3.jar antlr4-runtime/4.9.3//antlr4-runtime-4.9.3.jar aopalliance-repackaged/2.6.1//aopalliance-repackaged-2.6.1.jar +arrow-compression/16.0.0//arrow-compression-16.0.0.jar arrow-format/16.0.0//arrow-format-16.0.0.jar arrow-memory-core/16.0.0//arrow-memory-core-16.0.0.jar arrow-memory-netty-buffer-patch/16.0.0//arrow-memory-netty-buffer-patch-16.0.0.jar @@ -188,3 +189,4 @@ units/1.7//units-1.7.jar vertx-core/4.5.3//vertx-core-4.5.3.jar vertx-grpc/4.5.3//vertx-grpc-4.5.3.jar zjsonpatch/0.3.0//zjsonpatch-0.3.0.jar +zstd-jni/1.5.5-11//zstd-jni-1.5.5-11.jar diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml index 67308284bd6..cddc27c1025 100644 --- a/externals/kyuubi-spark-sql-engine/pom.xml +++ b/externals/kyuubi-spark-sql-engine/pom.xml @@ -37,6 +37,14 @@ ${project.version} + + + org.apache.arrow + arrow-compression + ${arrow.version} + true + + org.apache.kyuubi kyuubi-events_${scala.binary.version} @@ -263,6 +271,22 @@ 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-lang3 + org.immutables:value diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala new file mode 100644 index 00000000000..478afcb6d8a --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowCompressionSupport.scala @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.arrow + +import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionCodec} +import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader} + +/** Isolates the optional arrow-compression dependency so the uncompressed path never loads it. */ +private[sql] object KyuubiArrowCompressionSupport { + + def createLoader(root: VectorSchemaRoot): VectorLoader = { + new VectorLoader(root, CommonsCompressionFactory.INSTANCE) + } + + def createZstdUnloader(root: VectorSchemaRoot, level: Int): VectorUnloader = { + new VectorUnloader(root, true, new ZstdCompressionCodec(level), true) + } +} diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala index b8f590ddfbd..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 @@ -24,6 +24,7 @@ import scala.collection.JavaConverters._ import scala.collection.mutable.ArrayBuffer import org.apache.arrow.vector._ +import org.apache.arrow.vector.compression.NoCompressionCodec import org.apache.arrow.vector.ipc.{ArrowStreamWriter, ReadChannel, WriteChannel} import org.apache.arrow.vector.ipc.message.{IpcOption, MessageSerializer} import org.apache.spark.TaskContext @@ -39,6 +40,31 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) + private val CommonsCompressionFactoryClassName = + "org.apache.arrow.compression.CommonsCompressionFactory" + private val ZstdCompressionCodecClassName = + "org.apache.arrow.compression.ZstdCompressionCodec" + + // Mirror Spark's SparkSession#enableHiveSupport: check the capability on each request instead + // of caching, because the codec is a session-level config that can change at runtime. + private def arrowCompressionAvailable: Boolean = { + try { + Utils.classForName(CommonsCompressionFactoryClassName) + Utils.classForName(ZstdCompressionCodecClassName).getConstructor(Integer.TYPE) + true + } catch { + case _: ClassNotFoundException | _: NoClassDefFoundError | _: NoSuchMethodException => + false + } + } + + private def requireArrowCompression(): Unit = { + if (!arrowCompressionAvailable) { + throw new IllegalArgumentException( + "Arrow ZSTD compression requires arrow-compression on the Spark classpath") + } + } + /** * this method is to slice the input Arrow record batch byte array `bytes`, starting from `start` * and taking `length` number of elements. @@ -48,7 +74,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { timeZoneId: String, bytes: Array[Byte], start: Int, - length: Int): Array[Byte] = { + length: Int, + codecName: String = null, + zstdLevel: Int = 3): Array[Byte] = { val in = new ByteArrayInputStream(bytes) val out = new ByteArrayOutputStream(bytes.length) @@ -65,12 +93,35 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { val recordBatch = MessageSerializer.deserializeRecordBatch( new ReadChannel(Channels.newChannel(in)), sliceAllocator) - val vectorLoader = new VectorLoader(vectorSchemaRoot) + // Only compressed batches need the factory; the none path stays free of arrow-compression. + val compressed = + recordBatch.getBodyCompression.getCodec != NoCompressionCodec.COMPRESSION_TYPE + val vectorLoader = + if (compressed) { + requireArrowCompression() + KyuubiArrowCompressionSupport.createLoader(vectorSchemaRoot) + } else { + new VectorLoader(vectorSchemaRoot) + } vectorLoader.load(recordBatch) recordBatch.close() slicedVectorSchemaRoot = vectorSchemaRoot.slice(start, length) - val unloader = new VectorUnloader(slicedVectorSchemaRoot) + // Keep the compression codec on re-serialization, or the client cannot load the batch. + val unloader = codecName match { + case null | "none" => + new VectorUnloader(slicedVectorSchemaRoot) + case "zstd" => + requireArrowCompression() + KyuubiArrowCompressionSupport.createZstdUnloader(slicedVectorSchemaRoot, zstdLevel) + case "lz4" => + throw new IllegalArgumentException( + "Arrow compression codec lz4 is not supported by Kyuubi; " + + "supported codecs: none, zstd") + case other => + throw new IllegalArgumentException( + s"Unsupported Arrow compression codec: $other; supported codecs: none, zstd") + } val writeChannel = new WriteChannel(Channels.newChannel(out)) val batch = unloader.getRecordBatch() MessageSerializer.serialize(writeChannel, batch) @@ -119,7 +170,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { collectLimitExec: CollectLimitExec, maxRecordsPerBatch: Long, maxEstimatedBatchSize: Long, - timeZoneId: String): Array[Batch] = { + timeZoneId: String, + codecName: String = null, + zstdLevel: Int = 3): Array[Batch] = { val n = collectLimitExec.limit val schema = collectLimitExec.schema if (n == 0) { @@ -165,7 +218,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxRecordsPerBatch, maxEstimatedBatchSize, n, - timeZoneId) + timeZoneId, + codecName, + zstdLevel) batches.map(b => b -> batches.rowCountInLastBatch).toArray }, partsToScan) @@ -200,7 +255,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxRecordsPerBatch: Long, maxEstimatedBatchSize: Long, limit: Long, - timeZoneId: String): ArrowBatchIterator = { + timeZoneId: String, + codecName: String = null, + zstdLevel: Int = 3): ArrowBatchIterator = { new ArrowBatchIterator( rowIter, schema, @@ -208,7 +265,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxEstimatedBatchSize, limit, timeZoneId, - TaskContext.get) + TaskContext.get, + codecName, + zstdLevel) } /** @@ -226,10 +285,27 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { maxEstimatedBatchSize: Long, limit: Long, timeZoneId: String, - context: TaskContext) + context: TaskContext, + codecName: String, + zstdLevel: Int) extends Iterator[Array[Byte]] { protected val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true, false) + // Validate the codec before allocating Arrow buffers, so an unsupported codec fails fast. + private val compressionEnabled = codecName match { + case null | "none" => + false + case "zstd" => + requireArrowCompression() + true + case "lz4" => + throw new IllegalArgumentException( + "Arrow compression codec lz4 is not supported by Kyuubi; " + + "supported codecs: none, zstd") + case other => + throw new IllegalArgumentException( + s"Unsupported Arrow compression codec: $other; supported codecs: none, zstd") + } private val allocator = ArrowUtils.rootAllocator.newChildAllocator( s"to${this.getClass.getSimpleName}", @@ -237,7 +313,13 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { Long.MaxValue) private val root = VectorSchemaRoot.create(arrowSchema, allocator) - protected val unloader = new VectorUnloader(root) + // The none path keeps the original 1-arg constructor and stays free of arrow-compression. + protected val unloader = + if (compressionEnabled) { + KyuubiArrowCompressionSupport.createZstdUnloader(root, zstdLevel) + } else { + new VectorUnloader(root) + } protected val arrowWriter = ArrowWriter.create(root) Option(context).foreach { diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala index c503d780341..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 @@ -17,13 +17,15 @@ package org.apache.spark.sql.kyuubi +import java.util.Locale + import scala.collection.mutable.ArrayBuffer import org.apache.spark.SparkContext import org.apache.spark.internal.Logging import org.apache.spark.network.util.{ByteUnit, JavaUtils} import org.apache.spark.rdd.RDD -import org.apache.spark.sql.{DataFrame, Dataset, Row} +import org.apache.spark.sql.{DataFrame, Row, SparkSession} import org.apache.spark.sql.catalyst.plans.logical.GlobalLimit import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.execution.{CollectLimitExec, CommandResultExec, HiveResult, LocalTableScanExec, QueryExecution, SparkPlan, SparkPlanHelper, SQLExecution} @@ -36,7 +38,6 @@ import org.apache.spark.sql.types._ import org.apache.kyuubi.engine.spark.KyuubiSparkUtil import org.apache.kyuubi.engine.spark.schema.RowSet import org.apache.kyuubi.engine.spark.util.SparkCatalogUtils.quoteIfNeeded -import org.apache.kyuubi.util.reflect.{DynClasses, DynMethods} object SparkDatasetHelper extends Logging { @@ -63,22 +64,23 @@ object SparkDatasetHelper extends Logging { toArrowBatchRdd(plan).collect() } - private val datasetClz = DynClasses.builder() - .impl("org.apache.spark.sql.classic.Dataset") // SPARK-49700 (4.0.0) - .impl("org.apache.spark.sql.Dataset") - .build() - - private val toArrowBatchRddMethod = - DynMethods.builder("toArrowBatchRdd") - .impl(datasetClz) - .buildChecked() - - def toArrowBatchRdd[T](ds: Dataset[T]): RDD[Array[Byte]] = { - toArrowBatchRddMethod.bind(ds).invoke() + /** + * Read session-level Arrow compression configs introduced by SPARK-54134 (4.1.0), + * with key names finalized by the SPARK-54226 follow-up. + */ + private def arrowCompressionConf(spark: SparkSession): (String, Int) = { + // Lowercase the codec name to match Spark upstream. + val codecName = spark.sessionState.conf.getConfString( + "spark.sql.execution.arrow.compression.codec", + "none").toLowerCase(Locale.ROOT) + val zstdLevel = spark.sessionState.conf.getConfString( + "spark.sql.execution.arrow.compression.zstd.level", + "3").toInt + (codecName, zstdLevel) } /** - * Forked from [[Dataset.toArrowBatchRdd(plan: SparkPlan)]]. + * Forked from [[org.apache.spark.sql.Dataset.toArrowBatchRdd(plan: SparkPlan)]]. * Convert to an RDD of serialized ArrowRecordBatches. */ def toArrowBatchRdd(plan: SparkPlan): RDD[Array[Byte]] = { @@ -89,6 +91,7 @@ object SparkDatasetHelper extends Logging { // note that, we can't pass the lazy variable `maxBatchSize` directly, this is because input // arguments are serialized and sent to the executor side for execution. val maxBatchSizePerBatch = maxBatchSize + val (codecName, zstdLevel) = arrowCompressionConf(spark) plan.execute().mapPartitionsInternal { iter => KyuubiArrowConverters.toBatchIterator( iter, @@ -96,13 +99,18 @@ object SparkDatasetHelper extends Logging { maxRecordsPerBatch, maxBatchSizePerBatch, -1, - timeZoneId) + timeZoneId, + codecName, + zstdLevel) } } def toArrowBatchLocalIterator(df: DataFrame): Iterator[Array[Byte]] = { withNewExecutionId(df) { - toArrowBatchRdd(df).toLocalIterator + // use the plan-based toArrowBatchRdd so that the arrow compression codec takes effect; + // the vanilla Spark Dataset#toArrowBatchRdd uses the upstream ArrowConverters, which does + // not apply the Kyuubi compression codec. + toArrowBatchRdd(df.queryExecution.executedPlan).toLocalIterator } } @@ -174,12 +182,15 @@ object SparkDatasetHelper extends Logging { val spark = SparkPlanHelper.sparkSession(collectLimit) val timeZoneId = spark.sessionState.conf.sessionLocalTimeZone val maxRecordsPerBatch = spark.sessionState.conf.arrowMaxRecordsPerBatch + val (codecName, zstdLevel) = arrowCompressionConf(spark) val batches = KyuubiArrowConverters.takeAsArrowBatches( collectLimit, maxRecordsPerBatch, maxBatchSize, - timeZoneId) + timeZoneId, + codecName, + zstdLevel) // note that the number of rows in the returned arrow batches may be >= `limit`, perform // the slicing operation of result @@ -193,7 +204,14 @@ object SparkDatasetHelper extends Logging { // returned ArrowRecordBatch has less than `limit` row count, safety to do conversion rest -= size.toInt } else { // size > rest - result += KyuubiArrowConverters.slice(collectLimit.schema, timeZoneId, batch, 0, rest) + result += KyuubiArrowConverters.slice( + collectLimit.schema, + timeZoneId, + batch, + 0, + rest, + codecName, + zstdLevel) rest = 0 } i += 1 @@ -205,26 +223,32 @@ object SparkDatasetHelper extends Logging { val spark = SparkPlanHelper.sparkSession(commandResult) commandResult.longMetric("numOutputRows").add(commandResult.rows.size) sendDriverMetrics(spark.sparkContext, commandResult.metrics) + val (codecName, zstdLevel) = arrowCompressionConf(spark) KyuubiArrowConverters.toBatchIterator( commandResult.rows.iterator, commandResult.schema, spark.sessionState.conf.arrowMaxRecordsPerBatch, maxBatchSize, -1, - spark.sessionState.conf.sessionLocalTimeZone).toArray + spark.sessionState.conf.sessionLocalTimeZone, + codecName, + zstdLevel).toArray } private def doLocalTableScan(localTableScan: LocalTableScanExec): Array[Array[Byte]] = { val spark = SparkPlanHelper.sparkSession(localTableScan) localTableScan.longMetric("numOutputRows").add(localTableScan.rows.size) sendDriverMetrics(spark.sparkContext, localTableScan.metrics) + val (codecName, zstdLevel) = arrowCompressionConf(spark) KyuubiArrowConverters.toBatchIterator( localTableScan.rows.iterator, localTableScan.schema, spark.sessionState.conf.arrowMaxRecordsPerBatch, maxBatchSize, -1, - spark.sessionState.conf.sessionLocalTimeZone).toArray + spark.sessionState.conf.sessionLocalTimeZone, + codecName, + zstdLevel).toArray } /** diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala index 8e246b3c6e6..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 @@ -380,6 +380,54 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp } } + test("arrow zstd compression round-trips through the JDBC client") { + // Engine-side coverage of the compressed Arrow IPC path: the engine compresses batches with + // the upstream-named codec config, and the JDBC client transparently decompresses them from + // the batch body compression metadata (no server-side codec hint on the wire). + withJdbcStatement() { statement => + statement.executeQuery( + "set spark.sql.execution.arrow.compression.codec=zstd") + val resultSet = statement.executeQuery( + "select id, cast(id as string) as name from range(0, 1000)") + var count = 0 + while (resultSet.next()) { + // per-row invariant, independent of the row order the client receives + assert(resultSet.getString(2) == resultSet.getLong(1).toString) + count += 1 + } + assert(count == 1000) + } + } + + test("arrow zstd compression round-trips when LIMIT cuts across a batch boundary") { + // Two input partitions and initialNumPartitions=1 make the LIMIT deterministic: the first + // partition yields a full 100-row batch, and the second partition yields another full 100-row + // batch with only 50 rows left, so doCollectLimit must slice() the second batch down to 50. + withJdbcStatement() { statement => + statement.executeQuery( + "set spark.sql.execution.arrow.compression.codec=zstd") + statement.executeQuery( + s"set ${SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key}=100") + statement.executeQuery( + "set spark.sql.limit.initialNumPartitions=1") + val resultSet = statement.executeQuery( + """ + |select id, cast(id as string) as name from ( + | select id from range(0, 100, 1, 1) + | union all + | select id from range(100, 10000, 1, 1) + |) t limit 150 + |""".stripMargin) + var count = 0 + while (resultSet.next()) { + // per-row invariant survives compression + slicing, independent of row order + assert(resultSet.getString(2) == resultSet.getLong(1).toString) + count += 1 + } + assert(count == 150) + } + } + private def checkResultSetFormat(statement: Statement, expectFormat: String): Unit = { val query = s""" diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala new file mode 100644 index 00000000000..1d3e7d1f85c --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConvertersSuite.scala @@ -0,0 +1,173 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.arrow + +import java.io.ByteArrayInputStream +import java.nio.channels.Channels + +import scala.collection.JavaConverters._ + +import org.apache.arrow.compression.CommonsCompressionFactory +import org.apache.arrow.flatbuf.CompressionType +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{IntVector, VectorLoader, VectorSchemaRoot} +import org.apache.arrow.vector.compression.NoCompressionCodec +import org.apache.arrow.vector.ipc.ReadChannel +import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, MessageSerializer} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.unsafe.types.UTF8String + +import org.apache.kyuubi.KyuubiFunSuite + +/** + * Test suite for the arrow compression support in [[KyuubiArrowConverters]]: + * - the configured zstd compression level must be honored when constructing the codec, i.e. + * it must not be silently dropped in favor of the default level + * - slice() must keep the decompression factory and the compression codec paired, so that + * compressed batches can be sliced when a LIMIT cuts across a batch boundary + */ +class KyuubiArrowConvertersSuite extends KyuubiFunSuite { + + private val timeZoneId = "UTC" + + private val schema = StructType(Seq( + StructField("id", IntegerType), + StructField("name", StringType))) + + // must match what ArrowUtils.toArrowSchema produces for the schema above, so that the + // decompression side can load the batches + private val arrowSchema = new Schema(List( + new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null)).asJava) + + private def rows(count: Int): Iterator[InternalRow] = + (0 until count).iterator.map(i => InternalRow(i, UTF8String.fromString(s"name_$i"))) + + private def loadRoot( + bytes: Array[Byte], + allocator: BufferAllocator): (VectorSchemaRoot, ArrowRecordBatch) = { + val root = VectorSchemaRoot.create(arrowSchema, allocator) + val recordBatch = MessageSerializer.deserializeRecordBatch( + new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))), + allocator) + new VectorLoader(root, CommonsCompressionFactory.INSTANCE).load(recordBatch) + (root, recordBatch) + } + + private def assertRoundTrip(bytes: Array[Byte], expectedRows: Int): Unit = { + val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"round-trip-$expectedRows", 0, Long.MaxValue) + try { + val (root, recordBatch) = loadRoot(bytes, allocator) + try { + assert(root.getRowCount == expectedRows) + val ids = root.getVector("id").asInstanceOf[IntVector] + assert(ids.get(0) == 0) + assert(ids.get(expectedRows - 1) == expectedRows - 1) + } finally { + root.close() + } + recordBatch.close() + } finally { + allocator.close() + } + } + + test("zstd compression level is honored and the batches round-trip") { + val noneBytes = KyuubiArrowConverters + .toBatchIterator(rows(100), schema, 1000, -1, -1, timeZoneId, null, 3) + .toArray + .head + val level1 = KyuubiArrowConverters.slice(schema, timeZoneId, noneBytes, 0, 100, "zstd", 1) + val level19 = KyuubiArrowConverters.slice(schema, timeZoneId, noneBytes, 0, 100, "zstd", 19) + + // the configured zstd level must be stored in the batch and therefore produce different + // bytes, i.e. the configured level is not silently dropped when the codec is constructed + assert(!level1.sameElements(level19)) + val allocator = + ArrowUtils.rootAllocator.newChildAllocator("compression-level", 0, Long.MaxValue) + try { + val level1Batch = MessageSerializer.deserializeRecordBatch( + new ReadChannel(Channels.newChannel(new ByteArrayInputStream(level1))), + allocator) + try { + val bodyCompression = level1Batch.getBodyCompression + assert(bodyCompression != null) + assert(bodyCompression.getCodec == CompressionType.ZSTD) + } finally { + level1Batch.close() + } + } finally { + allocator.close() + } + + assertRoundTrip(level1, 100) + assertRoundTrip(level19, 100) + } + + test("slice cuts inside the head batch of compressed results") { + val batches = KyuubiArrowConverters + .toBatchIterator(rows(150), schema, 100, -1, -1, timeZoneId, "zstd", 3) + val batch1 = batches.next() // 100 rows, compressed + + // the LIMIT cuts inside the head batch: 40 rows out of 100. + // This is the only slice() call shape that SparkDatasetHelper.doCollectLimit performs in + // production: it slices a batch only when the remaining row budget is smaller than the + // batch size (size > rest); batches that fit entirely are passed through unsliced. + val sliced = KyuubiArrowConverters.slice(schema, timeZoneId, batch1, 0, 40, "zstd", 3) + assertRoundTrip(sliced, 40) + + // exhaust the remaining batch so that the iterator closes its allocator + assert(batches.hasNext) + batches.next() + assert(!batches.hasNext) + } + + test("reject unsupported lz4 compression codec") { + val error = intercept[IllegalArgumentException] { + KyuubiArrowConverters.toBatchIterator(rows(1), schema, 100, -1, -1, timeZoneId, "lz4", 3) + } + assert(error.getMessage.contains("Arrow compression codec lz4 is not supported by Kyuubi")) + } + + test("none codec keeps the original uncompressed path") { + val bytes = KyuubiArrowConverters + .toBatchIterator(rows(100), schema, 1000, -1, -1, timeZoneId, null, 3) + .toArray + .head + val allocator = + ArrowUtils.rootAllocator.newChildAllocator("none-codec", 0, Long.MaxValue) + try { + val batch = MessageSerializer.deserializeRecordBatch( + new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))), + allocator) + try { + // the none path must produce plain IPC batches marked with NO_COMPRESSION + assert(batch.getBodyCompression.getCodec === NoCompressionCodec.COMPRESSION_TYPE) + } finally { + batch.close() + } + } finally { + allocator.close() + } + assertRoundTrip(bytes, 100) + } +} diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala index 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-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-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/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..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 @@ -17,6 +17,7 @@ package org.apache.kyuubi.operation +import java.sql.SQLException import java.util.{Properties, UUID} import org.apache.hadoop.fs.{FileSystem, FileUtil, Path} @@ -26,7 +27,7 @@ import org.apache.kyuubi.{KYUUBI_VERSION, Utils, WithKyuubiServer, WithSimpleDFS import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf.KYUUBI_ENGINE_ENV_PREFIX import org.apache.kyuubi.jdbc.KyuubiHiveDriver -import org.apache.kyuubi.jdbc.hive.{KyuubiConnection, KyuubiStatement} +import org.apache.kyuubi.jdbc.hive.{KyuubiArrowQueryResultSet, KyuubiConnection, KyuubiStatement} import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem} import org.apache.kyuubi.session.{KyuubiSessionImpl, SessionHandle} import org.apache.kyuubi.shaded.hive.service.rpc.thrift.{TExecuteStatementReq, TGetInfoReq, TGetInfoType, TStatusCode} @@ -38,6 +39,13 @@ class KyuubiOperationPerUserSuite override protected def jdbcUrl: String = getJdbcUrl + // The engine jar no longer bundles arrow-compression; the zstd test supplies it via spark.jars, + // while the other tests cover the none path without it. + private lazy val arrowCompressionJar: String = + java.nio.file.Paths.get( + Class.forName("org.apache.arrow.compression.CommonsCompressionFactory") + .getProtectionDomain.getCodeSource.getLocation.toURI).toString + override protected val conf: KyuubiConf = { KyuubiConf().set(KyuubiConf.ENGINE_SHARE_LEVEL, "user") } @@ -202,6 +210,134 @@ class KyuubiOperationPerUserSuite } } + test("arrow zstd compression round-trips through the Kyuubi server") { + // Compressed Arrow IPC must pass through the server unchanged; the two partitions force + // slice(), and the dedicated subdomain engine gets arrow-compression via spark.jars. + withSessionConf()(Map.empty)(Map( + KyuubiConf.ENGINE_SHARE_LEVEL_SUBDOMAIN.key -> "arrow-zstd", + KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow", + "spark.jars" -> arrowCompressionJar, + "spark.sql.execution.arrow.compression.codec" -> "zstd", + "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100", + "spark.sql.limit.initialNumPartitions" -> "1")) { + withJdbcStatement() { statement => + // prove the codec config reached the engine through the server + val codecResult = + statement.executeQuery("set spark.sql.execution.arrow.compression.codec") + assert(codecResult.next()) + assert(codecResult.getString("value") === "zstd") + val resultSet = statement.executeQuery( + """ + |select id, cast(id as string) as name from ( + | select id from range(0, 100, 1, 1) + | union all + | select id from range(100, 10000, 1, 1) + |) t limit 150 + |""".stripMargin) + // prove the results come back as Arrow, not silently falling back to Thrift + assert(resultSet.isInstanceOf[KyuubiArrowQueryResultSet]) + val ids = scala.collection.mutable.Set.empty[Long] + var count = 0 + while (resultSet.next()) { + // per-row invariant survives compression, transfer and slicing, independent of the row + // order the client receives + val id = resultSet.getLong(1) + assert(resultSet.getString(2) === id.toString) + ids += id + count += 1 + } + assert(count === 150) + assert(ids === (0L until 150L).toSet) + } + } + } + + test("session-level codec switches within one engine when arrow-compression is present") { + // The codec is a session-level config; each switch within one long-lived engine must apply. + withSessionConf()(Map.empty)(Map( + KyuubiConf.ENGINE_SHARE_LEVEL_SUBDOMAIN.key -> "arrow-zstd", + KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow", + "spark.jars" -> arrowCompressionJar, + "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100", + "spark.sql.limit.initialNumPartitions" -> "1")) { + withJdbcStatement() { statement => + def checkQuery(): Unit = { + val resultSet = statement.executeQuery( + "select id, cast(id as string) as name from range(0, 1000)") + var count = 0 + while (resultSet.next()) { + assert(resultSet.getString(2) === resultSet.getLong(1).toString) + count += 1 + } + assert(count === 1000) + } + checkQuery() // none + statement.executeQuery("set spark.sql.execution.arrow.compression.codec=zstd") + checkQuery() // zstd + statement.executeQuery("set spark.sql.execution.arrow.compression.codec=none") + checkQuery() // back to none + } + } + } + + test("session-level zstd either round-trips or fails with a clear arrow-compression error") { + // Spark 4.1+ bundles arrow-compression in the Spark distribution, so this suite cannot + // control whether the engine classpath contains it. Enabling zstd must either fail with a + // clear dependency error (not NoClassDefFoundError) or, when the dependency is present, + // keep Arrow results working end-to-end. + withSessionConf()(Map.empty)(Map( + KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow")) { + withJdbcStatement() { statement => + val ok = statement.executeQuery("select id from range(0, 10)") + var count = 0 + while (ok.next()) count += 1 + assert(count === 10) + // the SET result itself is serialized as Arrow, so a missing dependency surfaces at the + // SET, while a present dependency is proven by the follow-up query below + try { + statement.executeQuery("set spark.sql.execution.arrow.compression.codec=zstd") + val resultSet = statement.executeQuery("select id from range(0, 10)") + val ids = scala.collection.mutable.Set.empty[Long] + while (resultSet.next()) ids += resultSet.getLong(1) + assert(ids === (0L until 10L).toSet) + } catch { + case e: SQLException => + assert( + e.getMessage.contains("arrow-compression"), + s"Unexpected error when enabling zstd: ${e.getMessage}") + } + } + } + } + + test("arrow results work without arrow-compression on the engine classpath") { + // This engine has no arrow-compression on its classpath; the default none path, including + // the slice() path, must still work. + withSessionConf()(Map.empty)(Map( + KyuubiConf.OPERATION_RESULT_FORMAT.key -> "arrow", + "spark.sql.execution.arrow.maxRecordsPerBatch" -> "100", + "spark.sql.limit.initialNumPartitions" -> "1")) { + withJdbcStatement() { statement => + val resultSet = statement.executeQuery( + """ + |select id, cast(id as string) as name from ( + | select id from range(0, 100, 1, 1) + | union all + | select id from range(100, 10000, 1, 1) + |) t limit 150 + |""".stripMargin) + assert(resultSet.isInstanceOf[KyuubiArrowQueryResultSet]) + val ids = scala.collection.mutable.Set.empty[Long] + while (resultSet.next()) { + val id = resultSet.getLong(1) + assert(resultSet.getString(2) === id.toString) + ids += id + } + assert(ids === (0L until 150L).toSet) + } + } + } + test("scala NPE issue with hdfs jar") { val jarDir = Utils.createTempDir().toFile val udfCode = diff --git a/pom.xml b/pom.xml index a7d02e9f759..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 @@ -348,6 +349,23 @@ arrow-vector ${arrow.version} + + org.apache.arrow + arrow-compression + ${arrow.version} + + + + org.apache.commons + commons-compress + + + + + com.github.luben + zstd-jni + ${zstd-jni.version} + org.apache.arrow arrow-memory-netty