[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

14fa002 [Wang Zhen] fix
386c92d [wforget] comments
8c1fcea [wforget] [KYUUBI-2543] Add `maxPartitionBytes` configuration

Lead-authored-by: wforget <643348094@qq.com>
Co-authored-by: Wang Zhen <wangzhen07@qiyi.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>

Author: wForget

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

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

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
 

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

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 {

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