You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2022/06/16 03:31:51 UTC
[incubator-kyuubi] branch master updated: [KYUUBI #2543] Add `maxPartitionBytes` configuration for TPC-DS connecter
This is an automated email from the ASF dual-hosted git repository.
chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new 7b24ee930 [KYUUBI #2543] Add `maxPartitionBytes` configuration for TPC-DS connecter
7b24ee930 is described below
commit 7b24ee9304135a25f07913d5d7903efe55f9a26b
Author: wforget <64...@qq.com>
AuthorDate: Thu Jun 16 11:31:41 2022 +0800
[KYUUBI #2543] Add `maxPartitionBytes` configuration for TPC-DS connecter
### _Why are the changes needed?_
close #2543
### _How was this patch tested?_
- [X] Add some test cases that check the changes thoroughly including negative and positive cases if possible
- [ ] Add screenshots for manual tests if appropriate
- [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
Closes #2882 from wForget/KYUUBI-2543-2.
Closes #2543
14fa002e [Wang Zhen] fix
386c92d0 [wforget] comments
8c1fcea8 [wforget] [KYUUBI-2543] Add `maxPartitionBytes` configuration
Lead-authored-by: wforget <64...@qq.com>
Co-authored-by: Wang Zhen <wa...@qiyi.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
.../spark/connector/common/SparkConfParser.scala | 116 +++++++++++++++++++++
.../connector/common/SparkConfParserSuite.scala | 65 ++++++++++++
.../spark/connector/tpcds/TPCDSBatchScan.scala | 9 +-
.../spark/connector/tpcds/TPCDSBatchScanConf.scala | 49 +++++++++
.../kyuubi/spark/connector/tpcds/TPCDSTable.scala | 9 +-
.../spark/connector/tpcds/TPCDSTableSuite.scala | 27 +++++
6 files changed, 269 insertions(+), 6 deletions(-)
diff --git a/extensions/spark/kyuubi-spark-connector-common/src/main/scala/org/apache/kyuubi/spark/connector/common/SparkConfParser.scala b/extensions/spark/kyuubi-spark-connector-common/src/main/scala/org/apache/kyuubi/spark/connector/common/SparkConfParser.scala
new file mode 100644
index 000000000..8e1c818ae
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-common/src/main/scala/org/apache/kyuubi/spark/connector/common/SparkConfParser.scala
@@ -0,0 +1,116 @@
+/*
+ * 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.kyuubi.spark.connector.common
+
+import java.util
+
+import org.apache.spark.sql.RuntimeConfig
+
+/**
+ * Parse the value of configuration in runtime options, session configurations and table properties,
+ * the priority of parsing configuration: options > sessionConfigs > properties.
+ * @param options runtime options
+ * @param sessionConfigs spark session configurations
+ * @param properties table properties
+ */
+case class SparkConfParser(
+ options: util.Map[String, String],
+ sessionConfigs: RuntimeConfig,
+ properties: util.Map[String, String]) {
+
+ def booleanConf(): BooleanConfParser = new BooleanConfParser()
+ def intConf(): IntConfParser = new IntConfParser()
+ def longConf(): LongConfParser = new LongConfParser()
+ def doubleConf(): DoubleConfParser = new DoubleConfParser()
+ def stringConf(): StringConfParser = new StringConfParser()
+
+ class BooleanConfParser extends ConfParser[Boolean] {
+ override protected def conversion(value: String): Boolean = value.toBoolean
+ }
+
+ class IntConfParser extends ConfParser[Int] {
+ override protected def conversion(value: String): Int = value.toInt
+ }
+
+ class LongConfParser extends ConfParser[Long] {
+ override protected def conversion(value: String): Long = value.toLong
+ }
+
+ class DoubleConfParser extends ConfParser[Double] {
+ override protected def conversion(value: String): Double = value.toDouble
+ }
+
+ class StringConfParser extends ConfParser[String] {
+ override protected def conversion(value: String): String = value
+ }
+
+ abstract class ConfParser[T]() {
+ private var optionName: Option[String] = None
+ private var sessionConfName: Option[String] = None
+ private var tablePropertyName: Option[String] = None
+ private var defaultValue: Option[T] = None
+
+ def option(name: String): ConfParser[T] = {
+ this.optionName = Some(name)
+ this
+ }
+
+ def sessionConf(name: String): ConfParser[T] = {
+ this.sessionConfName = Some(name)
+ this
+ }
+
+ def tableProperty(name: String): ConfParser[T] = {
+ this.tablePropertyName = Some(name)
+ this
+ }
+
+ def defaultValue(value: T): ConfParser[T] = {
+ this.defaultValue = Some(value)
+ this
+ }
+
+ private def parse(conversion: String => T): Option[T] = {
+ var valueOpt: Option[String] = None
+ if (options != null) {
+ valueOpt = optionName.flatMap(name => Option(options.get(name)))
+ }
+ if (valueOpt.isEmpty && sessionConfigs != null) {
+ valueOpt = sessionConfName.flatMap(name => sessionConfigs.getOption(name))
+ }
+ if (valueOpt.isEmpty && properties != null) {
+ valueOpt = tablePropertyName.flatMap(name => Option(properties.get(name)))
+ }
+ valueOpt = valueOpt.filter(_ != null)
+ if (valueOpt.isDefined) {
+ valueOpt.map(conversion(_))
+ } else {
+ defaultValue
+ }
+ }
+
+ protected def conversion(value: String): T
+
+ def parse(): T = {
+ assert(defaultValue.isDefined, "Default value cannot be empty.")
+ parseOptional().get
+ }
+
+ def parseOptional(): Option[T] = parse(conversion)
+ }
+}
diff --git a/extensions/spark/kyuubi-spark-connector-common/src/test/scala/org/apache/kyuubi/spark/connector/common/SparkConfParserSuite.scala b/extensions/spark/kyuubi-spark-connector-common/src/test/scala/org/apache/kyuubi/spark/connector/common/SparkConfParserSuite.scala
new file mode 100644
index 000000000..693e8ae04
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-common/src/test/scala/org/apache/kyuubi/spark/connector/common/SparkConfParserSuite.scala
@@ -0,0 +1,65 @@
+/*
+ * 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.kyuubi.spark.connector.common
+
+// scalastyle:off anyfunsuite
+import scala.collection.JavaConverters._
+
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.funsuite.AnyFunSuite
+
+class SparkConfParserSuite extends AnyFunSuite with BeforeAndAfterAll {
+// scalastyle:on anyfunsuite
+
+ test("parse options config") {
+ assert(confParser.stringConf().option("optKey1").defaultValue("test").parse() === "optValue1")
+ assert(confParser.booleanConf().option("booleanKey").defaultValue(true).parse() === false)
+ assert(confParser.intConf().option("intKey").defaultValue(0).parse() === 10)
+ assert(confParser.longConf().option("longKey").defaultValue(0).parse() === Long.MaxValue)
+ assert(confParser.doubleConf().option("doubleKey").defaultValue(0.0).parse() === 1.1)
+ }
+
+ test("parse properties config") {
+ assert(confParser.intConf().option("key1")
+ .tableProperty("key1")
+ .defaultValue(0).parse() === 111)
+ assert(confParser.stringConf()
+ .option("propertyKey1")
+ .tableProperty("propertyKey1")
+ .defaultValue("test").parse() === "propertyValue1")
+ }
+
+ private var confParser: SparkConfParser = null
+
+ override protected def beforeAll(): Unit = {
+ super.beforeAll()
+ val options = new CaseInsensitiveStringMap(Map(
+ "key1" -> "111",
+ "optKey1" -> "optValue1",
+ "booleanKey" -> "false",
+ "intKey" -> "10",
+ "longKey" -> String.valueOf(Long.MaxValue),
+ "doubleKey" -> "1.1").asJava)
+ val properties = Map(
+ "key1" -> "333",
+ "propertyKey1" -> "propertyValue1")
+ confParser = SparkConfParser(options, null, properties.asJava)
+ }
+
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScan.scala b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScan.scala
index 577273fe7..b1f5698a5 100644
--- a/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScan.scala
+++ b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScan.scala
@@ -36,19 +36,20 @@ case class TPCDSTableChuck(table: String, scale: Double, parallelism: Int, index
class TPCDSBatchScan(
@transient table: Table,
scale: Double,
- schema: StructType) extends ScanBuilder
+ schema: StructType,
+ scanConf: TPCDSBatchScanConf) extends ScanBuilder
with SupportsReportStatistics with Batch with Serializable {
private val _sizeInBytes: Long = TPCDSStatisticsUtils.sizeInBytes(table, scale)
private val _numRows: Long = TPCDSStatisticsUtils.numRows(table, scale)
- private val rowCountPerTask: Int = 1000000
-
+ // Tables with fewer than 1000000 are not parallelized,
+ // the limit made in `io.trino.tpcds.Parallel#splitWork`.
private val parallelism: Int =
if (table.isSmall) 1
else math.max(
SparkSession.active.sparkContext.defaultParallelism,
- (_numRows / rowCountPerTask.toDouble).ceil.toInt)
+ (_sizeInBytes / scanConf.maxPartitionBytes).ceil.toInt)
override def build: Scan = this
diff --git a/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScanConf.scala b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScanConf.scala
new file mode 100644
index 000000000..cf99b6bb1
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSBatchScanConf.scala
@@ -0,0 +1,49 @@
+/*
+ * 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.kyuubi.spark.connector.tpcds
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+import org.apache.kyuubi.spark.connector.common.SparkConfParser
+import org.apache.kyuubi.spark.connector.tpcds.TPCDSBatchScanConf._
+
+case class TPCDSBatchScanConf(
+ spark: SparkSession,
+ table: Table,
+ options: CaseInsensitiveStringMap) {
+
+ private val confParser: SparkConfParser =
+ SparkConfParser(options, spark.conf, table.properties())
+
+ lazy val maxPartitionBytes: Long = confParser.longConf()
+ .option(MAX_PARTITION_BYTES_CONF)
+ .sessionConf(s"$TPCDS_CONNECTOR_READ_CONF_PREFIX.$MAX_PARTITION_BYTES_CONF")
+ .tableProperty(MAX_PARTITION_BYTES_CONF)
+ .defaultValue(MAX_PARTITION_BYTES_DEFAULT)
+ .parse()
+
+}
+
+object TPCDSBatchScanConf {
+ val TPCDS_CONNECTOR_READ_CONF_PREFIX = "spark.connector.tpcds.read"
+
+ val MAX_PARTITION_BYTES_CONF = "maxPartitionBytes"
+ val MAX_PARTITION_BYTES_DEFAULT = 128 * 1024 * 1024L
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTable.scala b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTable.scala
index 701440508..a552336ac 100644
--- a/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTable.scala
+++ b/extensions/spark/kyuubi-spark-connector-tpcds/src/main/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTable.scala
@@ -25,6 +25,7 @@ import scala.collection.JavaConverters._
import io.trino.tpcds.Table
import io.trino.tpcds.column._
import io.trino.tpcds.column.ColumnType.Base._
+import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.connector.catalog.{SupportsRead, Table => SparkTable, TableCapability}
import org.apache.spark.sql.connector.expressions.{Expressions, Transform}
import org.apache.spark.sql.connector.read.ScanBuilder
@@ -47,6 +48,8 @@ class TPCDSTable(tbl: String, scale: Double, options: CaseInsensitiveStringMap)
val tpcdsTable: Table = Table.getTable(tbl)
+ lazy val spark: SparkSession = SparkSession.active
+
override def name: String = s"${TPCDSSchemaUtils.dbName(scale)}.$tbl"
override def toString: String = s"TPCDSTable($name)"
@@ -77,8 +80,10 @@ class TPCDSTable(tbl: String, scale: Double, options: CaseInsensitiveStringMap)
override def capabilities(): util.Set[TableCapability] =
Set(TableCapability.BATCH_READ).asJava
- override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
- new TPCDSBatchScan(tpcdsTable, scale, schema)
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ val scanConf = TPCDSBatchScanConf(spark, this, options)
+ new TPCDSBatchScan(tpcdsTable, scale, schema, scanConf)
+ }
def toSparkDataType(tpcdsType: ColumnType): DataType = {
(tpcdsType.getBase, tpcdsType.getPrecision.asScala, tpcdsType.getScale.asScala) match {
diff --git a/extensions/spark/kyuubi-spark-connector-tpcds/src/test/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTableSuite.scala b/extensions/spark/kyuubi-spark-connector-tpcds/src/test/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTableSuite.scala
index ec384bfe1..af983d8c8 100644
--- a/extensions/spark/kyuubi-spark-connector-tpcds/src/test/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTableSuite.scala
+++ b/extensions/spark/kyuubi-spark-connector-tpcds/src/test/scala/org/apache/kyuubi/spark/connector/tpcds/TPCDSTableSuite.scala
@@ -21,9 +21,11 @@ import io.trino.tpcds.Table
import io.trino.tpcds.generator.CallCenterGeneratorColumn
import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.kyuubi.KyuubiFunSuite
import org.apache.kyuubi.spark.connector.common.LocalSparkSession.withSparkSession
+import org.apache.kyuubi.spark.connector.tpcds.TPCDSBatchScanConf.{MAX_PARTITION_BYTES_CONF, TPCDS_CONNECTOR_READ_CONF_PREFIX}
class TPCDSTableSuite extends KyuubiFunSuite {
@@ -120,4 +122,29 @@ class TPCDSTableSuite extends KyuubiFunSuite {
CallCenterGeneratorColumn.CC_CALL_CENTER_SK.getGlobalColumnNumber)
}
}
+
+ test("test maxPartitionBytes") {
+ val maxPartitionBytes: Long = 1 * 1024 * 1024L
+ val sparkConf = new SparkConf().setMaster("local[*]")
+ .set("spark.ui.enabled", "false")
+ .set("spark.sql.catalogImplementation", "in-memory")
+ .set("spark.sql.catalog.tpcds", classOf[TPCDSCatalog].getName)
+ .set(
+ s"$TPCDS_CONNECTOR_READ_CONF_PREFIX.$MAX_PARTITION_BYTES_CONF",
+ String.valueOf(maxPartitionBytes))
+ withSparkSession(SparkSession.builder.config(sparkConf).getOrCreate()) { spark =>
+ val tableName = "catalog_returns"
+ val table = Table.getTable(tableName)
+ val scale = 100
+ val df = spark.sql(s"select * from tpcds.sf$scale.$tableName")
+ val scan = df.queryExecution.executedPlan.collectFirst {
+ case scanExec: BatchScanExec if scanExec.scan.isInstanceOf[TPCDSBatchScan] =>
+ scanExec.scan.asInstanceOf[TPCDSBatchScan]
+ }
+ assert(scan.isDefined)
+ val expected =
+ (TPCDSStatisticsUtils.sizeInBytes(table, scale) / maxPartitionBytes).ceil.toInt
+ assert(scan.get.planInputPartitions.size == expected)
+ }
+ }
}