pan3793 commented on code in PR #7636:
URL: https://github.com/apache/kyuubi/pull/7636#discussion_r3794657791
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala:
##########
@@ -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
Review Comment:
Add the JIRA ticket when referring to upstream Spark changes, e.g.,
`SPARK-XXXXX (4.1.0) do sth`
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala:
##########
@@ -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+.
Review Comment:
```suggestion
// Lowercase the codec name to match Spark upstream.
```
##########
kyuubi-hive-jdbc/pom.xml:
##########
@@ -46,6 +46,11 @@
<artifactId>arrow-vector</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ </dependency>
Review Comment:
we should tune the `kyuubi-hive-jdbc-shaded`'s `pom.xml` too, include and
relocate `arrow-compression`, but exclude `zstd-jni` and declare it as an
optional dependency
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala:
##########
@@ -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 =>
Review Comment:
Why not merge the `null` and `"none"` branches into the pattern-match body?
For `lz4`, we need to either fail with a clear error message, or fall back to
null with a warning message
##########
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).
Review Comment:
the "end-to-end coverage" is inaccurate, it's just the engine side coverage,
you still need to add the server side coverage to verify the: jdbc client <==>
kyuubi server <==> spark engine
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala:
##########
@@ -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.
Review Comment:
that's too verbose, you can just say: Read the session-level arrow
compression codec configs ...
##########
externals/kyuubi-spark-sql-engine/pom.xml:
##########
@@ -405,6 +430,17 @@
<pattern>net.bytebuddy</pattern>
<shadedPattern>${kyuubi.shade.packageName}.net.bytebuddy</shadedPattern>
</relocation>
+ <!--
+ Relocate the bundled arrow-compression classes so
that Kyuubi's copy
+ never conflicts with the arrow-compression jar
shipped by Spark 4.1+.
+ The relocated codec classes still resolve the
unrelocated
+ org.apache.arrow.vector.* and
com.github.luben.zstd.* types, which are
Review Comment:
I'm worried about the half-shading decision. Some questions:
- if we don't enable this feature, and don't bundle this depenedency, does
it work fine or throw a runtime class not found error?
- if the answer to the above question is "work fine", we can leave it and
ask the user to put the `arrow-compression` dependency in their Spark classpath
manually
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala:
##########
@@ -89,20 +113,26 @@ 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,
schemaCaptured,
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
Review Comment:
do we still need `def toArrowBatchRdd[T](ds: Dataset[T])` then?
##########
pom.xml:
##########
@@ -348,6 +348,11 @@
<artifactId>arrow-vector</artifactId>
<version>${arrow.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ <version>${arrow.version}</version>
Review Comment:
actually, we can exclude `commons-compress` since we don't support lz4.
note: things have changed since I suggested `none`, `gzip`, `lz4`, and
`zstd` options, lz4-java was abandoned by the original author, and another guy
took over but focused only on security patches, so I would suggest `none` and
`zstd` only now. (arrow compression does not support `gzip` so exclude it too)
##########
externals/kyuubi-spark-sql-engine/pom.xml:
##########
@@ -283,6 +307,7 @@
<exclude>NOTICE.txt</exclude>
<exclude>mozilla/**</exclude>
<exclude>**/module-info.class</exclude>
+ <exclude>arrow-git.properties</exclude>
Review Comment:
will this cause any side-effect it is unavailable at runtime?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
