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/05/25 03:29:15 UTC
[incubator-kyuubi] branch master updated: [KYUUBI #2663] Kyuubi Spark TPC-H Connector - Initial implementation
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 81c48b0c6 [KYUUBI #2663] Kyuubi Spark TPC-H Connector - Initial implementation
81c48b0c6 is described below
commit 81c48b0c61a1f553b3370344a164f26e6aa91698
Author: jiaoqingbo <11...@qq.com>
AuthorDate: Wed May 25 11:29:03 2022 +0800
[KYUUBI #2663] Kyuubi Spark TPC-H Connector - Initial implementation
### _Why are the changes needed?_
fix #2663
### _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
- [x] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
Closes #2731 from jiaoqingbo/kyuubi-2663.
Closes #2663
1553c59d [jiaoqingbo] code review
beb72754 [jiaoqingbo] [KYUUBI #2663] Kyuubi Spark TPC-H Connector - Initial implementation
Authored-by: jiaoqingbo <11...@qq.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
.../spark/kyuubi-spark-connector-tpch/pom.xml | 202 +++++++++++++++++++++
.../spark/connector/tpch/TPCHBatchScan.scala | 135 ++++++++++++++
.../kyuubi/spark/connector/tpch/TPCHCatalog.scala | 78 ++++++++
.../spark/connector/tpch/TPCHStatisticsUtils.scala | 65 +++++++
.../kyuubi/spark/connector/tpch/TPCHTable.scala | 78 ++++++++
.../src/test/resources/log4j.properties | 35 ++++
.../spark/connector/tpch/LocalSparkSession.scala | 39 ++++
.../spark/connector/tpch/TPCHCatalogSuite.scala | 52 ++++++
pom.xml | 14 ++
9 files changed, 698 insertions(+)
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/pom.xml b/extensions/spark/kyuubi-spark-connector-tpch/pom.xml
new file mode 100644
index 000000000..2734c91d4
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/pom.xml
@@ -0,0 +1,202 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+ -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <parent>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-parent</artifactId>
+ <version>1.6.0-SNAPSHOT</version>
+ <relativePath>../../../pom.xml</relativePath>
+ </parent>
+ <modelVersion>4.0.0</modelVersion>
+
+ <artifactId>kyuubi-spark-connector-tpch_2.12</artifactId>
+ <name>Kyuubi Spark TPC-H Connector</name>
+ <packaging>jar</packaging>
+ <url>https://kyuubi.apache.org/</url>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.scala-lang</groupId>
+ <artifactId>scala-library</artifactId>
+ <scope>provided</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-api</artifactId>
+ <scope>provided</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <scope>provided</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>io.trino.tpch</groupId>
+ <artifactId>tpch</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>com.google.guava</groupId>
+ <artifactId>guava</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.scalatestplus</groupId>
+ <artifactId>scalacheck-1-15_${scala.binary.version}</artifactId>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-client-api</artifactId>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-client-runtime</artifactId>
+ <scope>test</scope>
+ </dependency>
+
+ <!--
+ Spark requires `commons-collections` and `commons-io` but got them from transitive
+ dependencies of `hadoop-client`. As we are using Hadoop Shaded Client, we need add
+ them explicitly. See more details at SPARK-33212.
+ -->
+ <dependency>
+ <groupId>commons-collections</groupId>
+ <artifactId>commons-collections</artifactId>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>commons-io</groupId>
+ <artifactId>commons-io</artifactId>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>jakarta.xml.bind</groupId>
+ <artifactId>jakarta.xml.bind-api</artifactId>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+
+ <build>
+ <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
+ <testOutputDirectory>target/scala-${scala.binary.version}/test-classes</testOutputDirectory>
+
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-dependency-plugin</artifactId>
+ <configuration>
+ <skip>true</skip>
+ </configuration>
+ </plugin>
+
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-shade-plugin</artifactId>
+ <configuration>
+ <shadedArtifactAttached>false</shadedArtifactAttached>
+ <artifactSet>
+ <includes>
+ <include>io.trino.tpch:tpch</include>
+ <include>com.google.guava:guava</include>
+ </includes>
+ </artifactSet>
+ <relocations>
+ <relocation>
+ <pattern>com.google.common</pattern>
+ <shadedPattern>${kyuubi.shade.packageName}.com.google.common</shadedPattern>
+ <includes>
+ <include>com.google.common.**</include>
+ </includes>
+ </relocation>
+ <relocation>
+ <pattern>io.trino.tpch</pattern>
+ <shadedPattern>${kyuubi.shade.packageName}.io.trino.tpch</shadedPattern>
+ <includes>
+ <include>io.trino.tpch.**</include>
+ </includes>
+ </relocation>
+ </relocations>
+ </configuration>
+ <executions>
+ <execution>
+ <phase>package</phase>
+ <goals>
+ <goal>shade</goal>
+ </goals>
+ </execution>
+ </executions>
+ </plugin>
+
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-jar-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>prepare-test-jar</id>
+ <phase>test-compile</phase>
+ <goals>
+ <goal>test-jar</goal>
+ </goals>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+</project>
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHBatchScan.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHBatchScan.scala
new file mode 100644
index 000000000..4c9e8c13e
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHBatchScan.scala
@@ -0,0 +1,135 @@
+/*
+ * 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.tpch
+
+import java.time.LocalDate
+import java.time.format.DateTimeFormatter
+
+import scala.collection.mutable.ArrayBuffer
+
+import io.trino.tpch._
+import io.trino.tpch.GenerateUtils.formatDate
+import io.trino.tpch.TpchColumnType.Base._
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+
+case class TPCHTableChuck(table: String, scale: Int, parallelism: Int, index: Int)
+ extends InputPartition
+
+class TPCHBatchScan(
+ @transient table: TpchTable[_],
+ scale: Int,
+ schema: StructType) extends ScanBuilder
+ with Scan with Batch with Serializable {
+
+ private val _numRows: Long = TPCHStatisticsUtils.numRows(table, scale)
+
+ private val rowCountPerTask: Int = 1000000
+
+ private val parallelism: Int =
+ if (table.equals(TpchTable.NATION) || table.equals(TpchTable.REGION)) 1
+ else math.max(
+ SparkSession.active.sparkContext.defaultParallelism,
+ (_numRows / rowCountPerTask.toDouble).ceil.toInt)
+
+ override def build: Scan = this
+
+ override def toBatch: Batch = this
+
+ override def description: String =
+ s"Scan TPC-H sf$scale.${table.getTableName}, count: ${_numRows}, parallelism: $parallelism"
+
+ override def readSchema: StructType = schema
+
+ override def planInputPartitions: Array[InputPartition] =
+ (1 to parallelism).map { i =>
+ TPCHTableChuck(table.getTableName, scale, parallelism, i)
+ }.toArray
+
+ def createReaderFactory: PartitionReaderFactory = (partition: InputPartition) => {
+ val chuck = partition.asInstanceOf[TPCHTableChuck]
+ new TPCHPartitionReader(chuck.table, chuck.scale, chuck.parallelism, chuck.index, schema)
+ }
+
+}
+
+class TPCHPartitionReader(
+ table: String,
+ scale: Int,
+ parallelism: Int,
+ index: Int,
+ schema: StructType) extends PartitionReader[InternalRow] {
+
+ private val tpchTable = TpchTable.getTable(table)
+
+ private val columns = tpchTable.getColumns
+ .asInstanceOf[java.util.List[TpchColumn[TpchEntity]]]
+
+ private lazy val dateFmt: DateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd")
+
+ private val iterator = tpchTable.createGenerator(scale, index, parallelism).iterator
+
+ private var currentRow: InternalRow = _
+
+ override def next(): Boolean = {
+ val hasNext = iterator.hasNext
+ if (hasNext) currentRow = {
+ val row = iterator.next().asInstanceOf[TpchEntity]
+ val rowValue = new ArrayBuffer[String]()
+ columns.stream().forEach(column => {
+ val baseType = column.getType.getBase
+ var value: String = ""
+ baseType match {
+ case IDENTIFIER => value += column.getIdentifier(row)
+ case INTEGER => value += column.getInteger(row)
+ case DATE => value += column.getDate(row)
+ case DOUBLE => value += column.getDouble(row)
+ case VARCHAR => value += column.getString(row)
+ }
+ rowValue += value
+ })
+ val rowAny = new ArrayBuffer[Any]()
+ rowValue.zipWithIndex.map { case (value, i) =>
+ (value, schema(i).dataType) match {
+ case (null, _) => null
+ case ("", _) => null
+ case (value, IntegerType) => rowAny += value.toInt
+ case (value, LongType) => rowAny += value.toLong
+ case (value, DoubleType) => rowAny += value.toDouble
+ case (value, DateType) => rowAny += LocalDate.parse(formatDate(value.toInt), dateFmt)
+ .toEpochDay.toInt
+ case (value, StringType) => rowAny += UTF8String.fromString(value)
+ case (value, CharType(_)) => rowAny += UTF8String.fromString(value)
+ case (value, VarcharType(_)) => rowAny += UTF8String.fromString(value)
+ case (value, DecimalType()) => rowAny += Decimal(value)
+ case (value, dt) => throw new IllegalArgumentException(s"value: $value, type: $dt")
+ }
+ }
+ InternalRow.fromSeq(rowAny)
+ }
+ hasNext
+ }
+
+ override def get(): InternalRow = currentRow
+
+ override def close(): Unit = {}
+
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalog.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalog.scala
new file mode 100644
index 000000000..630c90d09
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalog.scala
@@ -0,0 +1,78 @@
+/*
+ * 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.tpch
+
+import java.util
+
+import scala.collection.JavaConverters._
+
+import io.trino.tpch.TpchTable
+import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException}
+import org.apache.spark.sql.connector.catalog.{Identifier, Table => SparkTable, TableCatalog, TableChange}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class TPCHCatalog extends TableCatalog {
+
+ val tables: Array[String] = TpchTable.getTables.asScala
+ .map(_.getTableName).toArray
+
+ val scales: Array[Int] = TPCHStatisticsUtils.SCALES
+
+ val databases: Array[String] = scales.map("sf" + _)
+
+ var options: CaseInsensitiveStringMap = _
+
+ var _name: Option[String] = None
+
+ override def name: String = _name.getOrElse("tpch")
+
+ override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {
+ this.options = options
+ this._name = Some(name)
+ }
+
+ override def listTables(namespace: Array[String]): Array[Identifier] = namespace match {
+ case Array(db) if databases contains db => tables.map(Identifier.of(namespace, _))
+ case _ => throw new NoSuchNamespaceException(namespace.mkString("."))
+ }
+
+ override def loadTable(ident: Identifier): SparkTable = (ident.namespace, ident.name) match {
+ case (Array(db), table) if (databases contains db) && tables.contains(table.toLowerCase) =>
+ new TPCHTable(table.toLowerCase, scales(databases indexOf db), options)
+ case (_, _) => throw new NoSuchTableException(ident)
+ }
+
+ override def createTable(
+ ident: Identifier,
+ schema: StructType,
+ partitions: Array[Transform],
+ properties: util.Map[String, String]): SparkTable =
+ throw new UnsupportedOperationException
+
+ override def alterTable(ident: Identifier, changes: TableChange*): SparkTable =
+ throw new UnsupportedOperationException
+
+ override def dropTable(ident: Identifier): Boolean =
+ throw new UnsupportedOperationException
+
+ override def renameTable(oldIdent: Identifier, newIdent: Identifier): Unit =
+ throw new UnsupportedOperationException
+
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHStatisticsUtils.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHStatisticsUtils.scala
new file mode 100644
index 000000000..1108829d8
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHStatisticsUtils.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.tpch
+
+import io.trino.tpch.TpchTable
+import io.trino.tpch.TpchTable._
+
+// https://www.tpc.org/tpc_documents_current_versions/pdf/tpc-h_v3.0.0.pdf
+// Page 88 Table 3: Estimated Database Size
+object TPCHStatisticsUtils {
+
+ val SCALES: Array[Int] = Array(0, 1, 10, 30, 100, 300, 1000, 3000, 10000, 30000, 100000)
+
+ def numRows(table: TpchTable[_], scale: Int): Long = {
+ require(SCALES.contains(scale), s"Unsupported scale $scale")
+ (table, scale) match {
+ case (_, 0) => 0L
+ case (CUSTOMER, scale) => 150000L * scale
+ case (ORDERS, scale) => 1500000L * scale
+ case (LINE_ITEM, 1) => 6001215L
+ case (LINE_ITEM, 10) => 59986052L
+ case (LINE_ITEM, 30) => 179998372L
+ case (LINE_ITEM, 100) => 600037902L
+ case (LINE_ITEM, 300) => 1799989091L
+ case (LINE_ITEM, 1000) => 5999989709L
+ case (LINE_ITEM, 3000) => 18000048306L
+ case (LINE_ITEM, 10000) => 59999994267L
+ case (LINE_ITEM, 30000) => 179999978268L
+ case (LINE_ITEM, 100000) => 599999969200L
+ case (PART, scale) => 200000L * scale
+ case (PART_SUPPLIER, scale) => 800000L * scale
+ case (SUPPLIER, scale) => 10000L * scale
+ case (NATION, _) => 25L
+ case (REGION, _) => 25L
+ }
+ }
+
+ def sizeInBytes(table: TpchTable[_], scale: Int): Long =
+ numRows(table, scale) * TABLE_AVG_ROW_BYTES(table)
+
+ private val TABLE_AVG_ROW_BYTES: Map[TpchTable[_], Long] = Map(
+ CUSTOMER -> 179,
+ ORDERS -> 104,
+ LINE_ITEM -> 112,
+ PART -> 155,
+ PART_SUPPLIER -> 144,
+ SUPPLIER -> 159,
+ NATION -> 128,
+ REGION -> 124)
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHTable.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHTable.scala
new file mode 100644
index 000000000..57e6c4d63
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/TPCHTable.scala
@@ -0,0 +1,78 @@
+/*
+ * 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.tpch
+
+import java.util
+import java.util.Optional
+
+import scala.collection.JavaConverters._
+import scala.collection.convert.ImplicitConversions.`list asScalaBuffer`
+
+import io.trino.tpch.{TpchColumnType, TpchEntity, TpchTable}
+import io.trino.tpch.TpchColumnType.Base._
+import org.apache.spark.sql.connector.catalog.{SupportsRead, Table => SparkTable, TableCapability}
+import org.apache.spark.sql.connector.read.ScanBuilder
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class TPCHTable(tbl: String, scale: Int, options: CaseInsensitiveStringMap)
+ extends SparkTable with SupportsRead {
+
+ // When true, use CHAR VARCHAR; otherwise use STRING
+ val useAnsiStringType: Boolean = options.getBoolean("useAnsiStringType", false)
+
+ val tpchTable: TpchTable[_] = TpchTable.getTable(tbl)
+
+ override def name: String = s"sf$scale.$tbl"
+
+ override def toString: String = s"TPCHTable($name)"
+
+ override def schema: StructType = {
+ StructType(
+ tpchTable.asInstanceOf[TpchTable[TpchEntity]].getColumns.zipWithIndex.map { case (c, _) =>
+ StructField(c.getColumnName, toSparkDataType(c.getType))
+ })
+ }
+
+ override def capabilities(): util.Set[TableCapability] =
+ Set(TableCapability.BATCH_READ).asJava
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ new TPCHBatchScan(tpchTable, scale, schema)
+ }
+
+ def toSparkDataType(tpchType: TpchColumnType): DataType = {
+ (tpchType.getBase, tpchType.getPrecision.asScala, tpchType.getScale.asScala) match {
+ case (INTEGER, None, None) => IntegerType
+ case (IDENTIFIER, None, None) => LongType
+ case (DOUBLE, None, None) => DoubleType
+ case (DATE, None, None) => DateType
+ case (VARCHAR, Some(precision), None) =>
+ if (useAnsiStringType) VarcharType(precision.toInt) else StringType
+ case (t, po, so) =>
+ throw new IllegalArgumentException(s"Unsupported TPC-H type: ($t, $po, $so)")
+ }
+ }
+
+ implicit final class RichOptional[T](val optional: Optional[T]) {
+ def asScala: Option[T] = optional match {
+ case null => null
+ case _ => if (optional.isPresent) Option(optional.get) else None
+ }
+ }
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/test/resources/log4j.properties b/extensions/spark/kyuubi-spark-connector-tpch/src/test/resources/log4j.properties
new file mode 100644
index 000000000..e3adf6f6a
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/test/resources/log4j.properties
@@ -0,0 +1,35 @@
+#
+# 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.
+#
+
+# Set everything to be logged to the file target/unit-tests.log
+log4j.rootLogger=INFO, CA, FA
+
+# Console Appender
+log4j.appender.CA=org.apache.log4j.ConsoleAppender
+log4j.appender.CA.layout=org.apache.log4j.PatternLayout
+log4j.appender.CA.layout.ConversionPattern=%d{HH:mm:ss.SSS} %p %c: %m%n
+log4j.appender.CA.Threshold = FATAL
+
+# File Appender
+log4j.appender.FA=org.apache.log4j.FileAppender
+log4j.appender.FA.append=false
+log4j.appender.FA.file=target/unit-tests.log
+log4j.appender.FA.layout=org.apache.log4j.PatternLayout
+log4j.appender.FA.layout.ConversionPattern=%d{HH:mm:ss.SSS} %t %p %c{2}: %m%n
+
+# Set the logger level of File Appender to WARN
+log4j.appender.FA.Threshold = DEBUG
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/LocalSparkSession.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/LocalSparkSession.scala
new file mode 100644
index 000000000..8b60be98c
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/LocalSparkSession.scala
@@ -0,0 +1,39 @@
+/*
+ * 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.tpch
+
+import org.apache.spark.sql.SparkSession
+
+object LocalSparkSession {
+ def stop(spark: SparkSession): Unit = {
+ if (spark != null) {
+ spark.stop()
+ }
+ // To avoid RPC rebinding to the same port, since it doesn't unbind immediately on shutdown
+ System.clearProperty("spark.driver.port")
+ }
+
+ /** Runs `f` by passing in `sc` and ensures that `sc` is stopped. */
+ def withSparkSession[T](sc: SparkSession)(f: SparkSession => T): T = {
+ try {
+ f(sc)
+ } finally {
+ stop(sc)
+ }
+ }
+}
diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalogSuite.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalogSuite.scala
new file mode 100644
index 000000000..9be0babb5
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/test/scala/org/apache/kyuubi/spark/connector/tpch/TPCHCatalogSuite.scala
@@ -0,0 +1,52 @@
+/*
+ * 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.tpch
+
+import org.apache.spark.sql.SparkSession
+
+import org.apache.kyuubi.KyuubiFunSuite
+
+class TPCHCatalogSuite extends KyuubiFunSuite {
+
+ protected lazy val spark: SparkSession = {
+ SparkSession.builder()
+ .master("local[*]")
+ .config("spark.ui.enabled", "false")
+ .config("spark.sql.catalogImplementation", "in-memory")
+ .config("spark.sql.catalog.tpch", classOf[TPCHCatalog].getName)
+ .config("spark.sql.cbo.enabled", "true")
+ .config("spark.sql.cbo.planStats.enabled", "true")
+ .getOrCreate()
+ }
+
+ test("tpch.sf1 count") {
+ assert(spark.table("tpch.sf1.customer").count === 150000)
+ assert(spark.table("tpch.sf1.orders").count === 1500000)
+ assert(spark.table("tpch.sf1.lineitem").count === 6001215)
+ assert(spark.table("tpch.sf1.part").count === 200000)
+ assert(spark.table("tpch.sf1.partsupp").count === 800000)
+ assert(spark.table("tpch.sf1.supplier").count === 10000)
+ assert(spark.table("tpch.sf1.nation").count === 25)
+ assert(spark.table("tpch.sf1.region").count === 5)
+ }
+
+ override def afterAll(): Unit = {
+ super.afterAll()
+ spark.stop()
+ }
+}
diff --git a/pom.xml b/pom.xml
index 9772b3ad9..d79a80112 100644
--- a/pom.xml
+++ b/pom.xml
@@ -69,6 +69,7 @@
<module>extensions/server/kyuubi-server-plugin</module>
<module>extensions/spark/kyuubi-spark-authz</module>
<module>extensions/spark/kyuubi-spark-connector-tpcds</module>
+ <module>extensions/spark/kyuubi-spark-connector-tpch</module>
<module>externals/kyuubi-download</module>
<module>externals/kyuubi-flink-sql-engine</module>
<module>externals/kyuubi-hive-sql-engine</module>
@@ -162,6 +163,7 @@
<thrift.version>0.9.3</thrift.version>
<trino.client.version>363</trino.client.version>
<trino.tpcds.version>1.4</trino.tpcds.version>
+ <trino.tpch.version>1.1</trino.tpch.version>
<zookeeper.version>3.4.14</zookeeper.version>
<!-- apply to kyuubi-hive-jdbc/kyuubi-hive-beeline module -->
@@ -537,6 +539,18 @@
</exclusions>
</dependency>
+ <dependency>
+ <groupId>io.trino.tpch</groupId>
+ <artifactId>tpch</artifactId>
+ <version>${trino.tpch.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>*</groupId>
+ <artifactId>*</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+
<dependency>
<groupId>com.dimafeng</groupId>
<artifactId>testcontainers-scala-scalatest_${scala.binary.version}</artifactId>