byronwang123 commented on code in PR #7636:
URL: https://github.com/apache/kyuubi/pull/7636#discussion_r3804903950
##########
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:
Addressed. Removed the unused Dataset overload.
--
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]
