[incubator-kyuubi] branch master updated: [KYUUBI #2664] Kyuubi Spark TPC-H Connector - SupportsReportStatistics

[ch...@apache.org]

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 dbe315e8a [KYUUBI #2664] Kyuubi Spark TPC-H Connector - SupportsReportStatistics
dbe315e8a is described below

commit dbe315e8a2ac51c2c81e88cab3d62924d3196a65
Author: jiaoqingbo <11...@qq.com>
AuthorDate: Thu May 26 17:19:10 2022 +0800

    [KYUUBI #2664] Kyuubi Spark TPC-H Connector - SupportsReportStatistics
    
    ### _Why are the changes needed?_
    
    fix #2664
    
    ### _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 #2740 from jiaoqingbo/kyuubi-2664.
    
    Closes #2664
    
    ebac8adb [jiaoqingbo] fix ut failed
    6532cc5d [jiaoqingbo] [KYUUBI #2664] Kyuubi Spark TPC-H Connector - SupportsReportStatistics
    
    Authored-by: jiaoqingbo <11...@qq.com>
    Signed-off-by: Cheng Pan <ch...@apache.org>
---
 .../kyuubi/spark/connector/tpch/SparkUtils.scala   | 78 ++++++++++++++++++++++
 .../spark/connector/tpch/TPCHBatchScan.scala       | 10 ++-
 .../spark/connector/tpch/TPCHStatisticsUtils.scala |  2 +-
 .../spark/connector/tpch/TPCHCatalogSuite.scala    | 30 ++++++++-
 4 files changed, 117 insertions(+), 3 deletions(-)

diff --git a/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/SparkUtils.scala b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/SparkUtils.scala
new file mode 100644
index 000000000..84b2835ba
--- /dev/null
+++ b/extensions/spark/kyuubi-spark-connector-tpch/src/main/scala/org/apache/kyuubi/spark/connector/tpch/SparkUtils.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 org.apache.spark.SPARK_VERSION
+
+object SparkUtils {
+
+  /**
+   * Given a Kyuubi/Spark/Hive version string,
+   * return the (major version number, minor version number).
+   * E.g., for 2.0.1-SNAPSHOT, return (2, 0).
+   */
+  def majorMinorVersion(version: String): (Int, Int) = {
+    """^(\d+)\.(\d+)(\..*)?$""".r.findFirstMatchIn(version) match {
+      case Some(m) =>
+        (m.group(1).toInt, m.group(2).toInt)
+      case None =>
+        throw new IllegalArgumentException(s"Tried to parse '$version' as a project" +
+          s" version string, but it could not find the major and minor version numbers.")
+    }
+  }
+
+  /**
+   * Given a Kyuubi/Spark/Hive version string, return the major version number.
+   * E.g., for 2.0.1-SNAPSHOT, return 2.
+   */
+  def majorVersion(version: String): Int = majorMinorVersion(version)._1
+
+  /**
+   * Given a Kyuubi/Spark/Hive version string, return the minor version number.
+   * E.g., for 2.0.1-SNAPSHOT, return 0.
+   */
+  def minorVersion(version: String): Int = majorMinorVersion(version)._2
+
+  def isSparkVersionAtMost(ver: String): Boolean = {
+    val runtimeMajor = majorVersion(SPARK_VERSION)
+    val targetMajor = majorVersion(ver)
+    (runtimeMajor < targetMajor) || {
+      val runtimeMinor = minorVersion(SPARK_VERSION)
+      val targetMinor = minorVersion(ver)
+      runtimeMajor == targetMajor && runtimeMinor <= targetMinor
+    }
+  }
+
+  def isSparkVersionAtLeast(ver: String): Boolean = {
+    val runtimeMajor = majorVersion(SPARK_VERSION)
+    val targetMajor = majorVersion(ver)
+    (runtimeMajor > targetMajor) || {
+      val runtimeMinor = minorVersion(SPARK_VERSION)
+      val targetMinor = minorVersion(ver)
+      runtimeMajor == targetMajor && runtimeMinor >= targetMinor
+    }
+  }
+
+  def isSparkVersionEqualTo(ver: String): Boolean = {
+    val runtimeMajor = majorVersion(SPARK_VERSION)
+    val targetMajor = majorVersion(ver)
+    val runtimeMinor = minorVersion(SPARK_VERSION)
+    val targetMinor = minorVersion(ver)
+    runtimeMajor == targetMajor && runtimeMinor == targetMinor
+  }
+}
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
index 4c9e8c13e..939e7c29e 100644
--- 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
@@ -19,6 +19,7 @@ package org.apache.kyuubi.spark.connector.tpch
 
 import java.time.LocalDate
 import java.time.format.DateTimeFormatter
+import java.util.OptionalLong
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -38,7 +39,9 @@ class TPCHBatchScan(
     @transient table: TpchTable[_],
     scale: Int,
     schema: StructType) extends ScanBuilder
-  with Scan with Batch with Serializable {
+  with SupportsReportStatistics with Batch with Serializable {
+
+  private val _sizeInBytes: Long = TPCHStatisticsUtils.sizeInBytes(table, scale)
 
   private val _numRows: Long = TPCHStatisticsUtils.numRows(table, scale)
 
@@ -69,6 +72,11 @@ class TPCHBatchScan(
     new TPCHPartitionReader(chuck.table, chuck.scale, chuck.parallelism, chuck.index, schema)
   }
 
+  override def estimateStatistics: Statistics = new Statistics {
+    override def sizeInBytes: OptionalLong = OptionalLong.of(_sizeInBytes)
+    override def numRows: OptionalLong = OptionalLong.of(_numRows)
+  }
+
 }
 
 class TPCHPartitionReader(
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
index 1108829d8..1137521fe 100644
--- 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
@@ -46,7 +46,7 @@ object TPCHStatisticsUtils {
       case (PART_SUPPLIER, scale) => 800000L * scale
       case (SUPPLIER, scale) => 10000L * scale
       case (NATION, _) => 25L
-      case (REGION, _) => 25L
+      case (REGION, _) => 5L
     }
   }
 
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
index 2496a0cd8..a831745b0 100644
--- 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
@@ -17,7 +17,7 @@
 
 package org.apache.kyuubi.spark.connector.tpch
 
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{AnalysisException, SparkSession}
 
 import org.apache.kyuubi.KyuubiFunSuite
 
@@ -51,6 +51,34 @@ class TPCHCatalogSuite extends KyuubiFunSuite {
     assert(spark.table("tpch.sf1.region").count === 5)
   }
 
+  test("tpch.sf1 stats") {
+    def assertStats(tableName: String, sizeInBytes: BigInt, rowCount: BigInt): Unit = {
+      val stats = spark.table(tableName).queryExecution.analyzed.stats
+      assert(stats.sizeInBytes == sizeInBytes)
+      // stats.rowCount only has value after SPARK-33954
+      if (SparkUtils.isSparkVersionAtLeast("3.2")) {
+        assert(stats.rowCount.contains(rowCount), tableName)
+      }
+    }
+
+    assertStats("tpch.sf1.customer", 26850000, 150000)
+    assertStats("tpch.sf1.orders", 156000000, 1500000)
+    assertStats("tpch.sf1.lineitem", 672136080, 6001215)
+    assertStats("tpch.sf1.part", 31000000, 200000)
+    assertStats("tpch.sf1.partsupp", 115200000, 800000)
+    assertStats("tpch.sf1.supplier", 1590000, 10000)
+    assertStats("tpch.sf1.nation", 3200, 25)
+    assertStats("tpch.sf1.region", 620, 5)
+
+  }
+
+  test("nonexistent table") {
+    val exception = intercept[AnalysisException] {
+      spark.table("tpch.sf1.nonexistent_table")
+    }
+    assert(exception.message === "Table or view not found: tpch.sf1.nonexistent_table")
+  }
+
   override def afterAll(): Unit = {
     super.afterAll()
     spark.stop()