diff --git a/.github/workflows/license.yml b/.github/workflows/license.yml index fb74c19d172..0c3c960f1a1 100644 --- a/.github/workflows/license.yml +++ b/.github/workflows/license.yml @@ -44,7 +44,7 @@ jobs: - run: | build/mvn org.apache.rat:apache-rat-plugin:check \ -Ptpcds -Pkubernetes-it \ - -Pspark-3.3 -Pspark-3.4 -Pspark-3.5 -Pspark-4.0 -Pspark-4.1 -Pspark-4.2 + -Pspark-3.5 -Pspark-4.0 -Pspark-4.1 -Pspark-4.2 - name: Upload rat report if: failure() uses: actions/upload-artifact@v7 diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml index f3cd8cdf510..999f84400e2 100644 --- a/.github/workflows/master.yml +++ b/.github/workflows/master.yml @@ -55,8 +55,6 @@ jobs: python: - '3.9' spark: - - '3.3' - - '3.4' - '3.5' spark-archive: [""] exclude-tags: [""] @@ -80,18 +78,6 @@ jobs: spark-archive: '-Pscala-2.13' exclude-tags: '-Dmaven.plugin.scalatest.exclude.tags=org.scalatest.tags.Slow,org.apache.kyuubi.tags.DeltaTest,org.apache.kyuubi.tags.IcebergTest,org.apache.kyuubi.tags.PaimonTest,org.apache.kyuubi.tags.HudiTest' comment: 'normal' - - java: 8 - python: '3.9' - spark: '3.5' - spark-archive: '-Dspark.archive.mirror=https://www.apache.org/dyn/closer.lua/spark/spark-3.3.3 -Dspark.archive.name=spark-3.3.3-bin-hadoop3.tgz -Pzookeeper-3.6' - exclude-tags: '-Dmaven.plugin.scalatest.exclude.tags=org.scalatest.tags.Slow,org.apache.kyuubi.tags.DeltaTest,org.apache.kyuubi.tags.IcebergTest,org.apache.kyuubi.tags.PaimonTest,org.apache.kyuubi.tags.HudiTest,org.apache.kyuubi.tags.SparkLocalClusterTest' - comment: 'verify-on-spark-3.3-binary' - - java: 8 - python: '3.9' - spark: '3.5' - spark-archive: '-Dspark.archive.mirror=https://www.apache.org/dyn/closer.lua/spark/spark-3.4.4 -Dspark.archive.name=spark-3.4.4-bin-hadoop3.tgz -Pzookeeper-3.6' - exclude-tags: '-Dmaven.plugin.scalatest.exclude.tags=org.scalatest.tags.Slow,org.apache.kyuubi.tags.DeltaTest,org.apache.kyuubi.tags.IcebergTest,org.apache.kyuubi.tags.PaimonTest,org.apache.kyuubi.tags.SparkLocalClusterTest' - comment: 'verify-on-spark-3.4-binary' - java: 17 python: '3.11' spark: '3.5' @@ -248,17 +234,6 @@ jobs: fail-fast: false max-parallel: 1 matrix: - java: - - 17 - scala: - - '2.12' - - '2.13' - spark-compile: - - '3.5' - spark-runtime: - - '3.4' - - '3.3' - comment: [ "normal" ] include: - java: 17 scala: "2.13" diff --git a/.github/workflows/publish-snapshot-nexus.yml b/.github/workflows/publish-snapshot-nexus.yml index dd4ca79f5c4..e174e2f1114 100644 --- a/.github/workflows/publish-snapshot-nexus.yml +++ b/.github/workflows/publish-snapshot-nexus.yml @@ -31,27 +31,23 @@ jobs: matrix: branch: - master - - branch-1.8 - - branch-1.9 + - branch-1.12 profiles: - - -Pflink-provided,spark-provided,hive-provided,spark-3.3 - - -Pflink-provided,spark-provided,hive-provided,spark-3.4,tpcds - include: - - branch: master - profiles: -Pflink-provided,spark-provided,hive-provided,spark-3.5 - - branch: branch-1.9 - profiles: -Pflink-provided,spark-provided,hive-provided,spark-3.5 + - -Pflink-provided,spark-provided,hive-provided,spark-3.5,tpcds + - -Pflink-provided,spark-provided,hive-provided,scala-2.13,spark-4.0 + - -Pflink-provided,spark-provided,hive-provided,scala-2.13,spark-4.1 + - -Pflink-provided,spark-provided,hive-provided,scala-2.13,spark-4.2,tpcds steps: - uses: actions/checkout@v7 with: ref: ${{ matrix.branch }} - name: Free up disk space run: ./.github/scripts/free_disk_space.sh - - name: Setup JDK 8 + - name: Setup JDK 17 uses: actions/setup-java@v5 with: distribution: temurin - java-version: 8 + java-version: 17 cache: 'maven' - name: Publish Snapshot Jar to Nexus - ${{ matrix.branch }} env: diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 9c4cadd7d3f..e5d486ae8cc 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -44,7 +44,7 @@ jobs: fail-fast: false max-parallel: 15 matrix: - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.9", "3.10"] env: PYTHONHASHSEED: random steps: diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index b51bcbc5b86..aeefd89837e 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -35,7 +35,7 @@ jobs: max-parallel: 15 matrix: profiles: - - '-Pflink-provided,hive-provided,spark-provided,spark-3.5,spark-3.4,spark-3.3,tpcds,kubernetes-it' + - '-Pflink-provided,hive-provided,spark-provided,spark-3.5,tpcds,kubernetes-it' steps: - uses: actions/checkout@v7 @@ -66,8 +66,6 @@ jobs: run: | export MAVEN_ARGS="-DskipTests -Dorg.slf4j.simpleLogger.defaultLogLevel=warn -Drat.skip=true -Dmaven.scalastyle.skip=true -Dspotless.check.skip=true -Pflink-provided,hive-provided,spark-provided" build/mvn clean install -Ptpcds - build/mvn clean install -pl extensions/spark/kyuubi-extension-spark-3-3 -am -Pspark-3.3 - build/mvn clean install -pl extensions/spark/kyuubi-extension-spark-3-4 -am -Pspark-3.4 build/mvn clean install -pl extensions/spark/kyuubi-extension-spark-3-5,extensions/spark/kyuubi-spark-connector-hive -am -Pspark-3.5 build/mvn clean install -pl extensions/spark/kyuubi-extension-spark-4-0 -am -Pspark-4.0 -Pscala-2.13 build/mvn clean install -pl extensions/spark/kyuubi-extension-spark-4-1 -am -Pspark-4.1 -Pscala-2.13 diff --git a/AGENTS.md b/AGENTS.md index 2ca87fd03b1..3e6901a6259 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -67,14 +67,14 @@ Use `-am` (also-make) when building or testing a single module — without it, M ### Engine profile matrix -| Profile | Notes | -|---|---| -| `-Pspark-3.5` (default), `-Pspark-{3.3,3.4,4.0,4.1,master}` | Spark version | -| `-Pflink-1.20` (default), `-Pflink-{1.17,1.18,1.19}` | Flink version | -| `-Pscala-2.13` | Scala 2.13 (default is 2.12) | -| `-P{spark,flink,hive}-provided` | skip bundled engine downloads | -| `-Pmirror-cdn` | use Apache mirror CDN for engine archives | -| `-Pfast` | skip tests/style/docs/enforcer/RAT/downloads | +| Profile | Notes | +|---------------------------------------------------------|----------------------------------------------| +| `-Pspark-3.5` (default), `-Pspark-{4.0,4.1,4.2,master}` | Spark version | +| `-Pflink-1.20` (default), `-Pflink-{1.17,1.18,1.19}` | Flink version | +| `-Pscala-2.13` | Scala 2.13 (default is 2.12) | +| `-P{spark,flink,hive}-provided` | skip bundled engine downloads | +| `-Pmirror-cdn` | use Apache mirror CDN for engine archives | +| `-Pfast` | skip tests/style/docs/enforcer/RAT/downloads | When engine code varies across versions, gate source/binary differences by Maven profile and runtime capability differences by feature detection — not by parsing version strings. diff --git a/build/release/release.sh b/build/release/release.sh index 09b5a18b9e1..e80088d28d8 100755 --- a/build/release/release.sh +++ b/build/release/release.sh @@ -110,16 +110,6 @@ upload_svn_staging() { } upload_nexus_staging() { - # Spark Extension Plugin for Spark 3.3 and Scala 2.12 - ${KYUUBI_DIR}/build/mvn clean deploy -DskipTests -Papache-release,flink-provided,spark-provided,hive-provided,spark-3.3 \ - -s "${KYUUBI_DIR}/build/release/asf-settings.xml" \ - -pl extensions/spark/kyuubi-extension-spark-3-3 -am - - # Spark Extension Plugin for Spark 3.4 and Scala 2.12 - ${KYUUBI_DIR}/build/mvn clean deploy -DskipTests -Papache-release,flink-provided,spark-provided,hive-provided,spark-3.4 \ - -s "${KYUUBI_DIR}/build/release/asf-settings.xml" \ - -pl extensions/spark/kyuubi-extension-spark-3-4 -am - # Spark Extension Plugin for Spark 4.0 and Scala 2.13 ${KYUUBI_DIR}/build/mvn clean deploy -DskipTests -Papache-release,flink-provided,spark-provided,hive-provided,spark-4.0,scala-2.13 \ -s "${KYUUBI_DIR}/build/release/asf-settings.xml" \ diff --git a/dev/kyuubi-codecov/pom.xml b/dev/kyuubi-codecov/pom.xml index 0dfe6842e4c..4d922e2f07d 100644 --- a/dev/kyuubi-codecov/pom.xml +++ b/dev/kyuubi-codecov/pom.xml @@ -155,56 +155,6 @@ - - spark-3.3 - - - org.apache.kyuubi - kyuubi-extension-spark-3-3_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-connector-hive_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-authz_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-lineage_${scala.binary.version} - ${project.version} - - - - - spark-3.4 - - - org.apache.kyuubi - kyuubi-extension-spark-3-4_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-connector-hive_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-authz_${scala.binary.version} - ${project.version} - - - org.apache.kyuubi - kyuubi-spark-lineage_${scala.binary.version} - ${project.version} - - - spark-3.5 diff --git a/dev/reformat b/dev/reformat index 28ce515f5f1..eed4fc1b40e 100755 --- a/dev/reformat +++ b/dev/reformat @@ -20,7 +20,7 @@ set -x KYUUBI_HOME="$(cd "`dirname "$0"`/.."; pwd)" -PROFILES="-Pflink-provided,hive-provided,spark-provided,spark-4.2,spark-4.1,spark-4.0,spark-3.5,spark-3.4,spark-3.3,tpcds,kubernetes-it" +PROFILES="-Pflink-provided,hive-provided,spark-provided,spark-4.2,spark-4.1,spark-4.0,spark-3.5,tpcds,kubernetes-it" # python style checks rely on `black` in path if ! command -v black &> /dev/null diff --git a/docs/contributing/code/building.md b/docs/contributing/code/building.md index 689b705f7bb..8621bbacc88 100644 --- a/docs/contributing/code/building.md +++ b/docs/contributing/code/building.md @@ -61,11 +61,17 @@ mvn clean install -pl '!dev/kyuubi-codecov,!kyuubi-assembly' -DskipTests Since v1.1.0, Kyuubi support building with different Spark profiles, -| Profile | Default | Since | -|-------------|---------|-------| -| -Pspark-3.3 | | 1.6.0 | -| -Pspark-3.4 | | 1.8.0 | -| -Pspark-3.5 | ✓ | 1.8.0 | +| Profile | Default | Since | Removed | +|-------------|---------|--------|---------| +| -Pspark-3.0 | | 1.1.0 | 1.6.0 | +| -Pspark-3.1 | | 1.1.0 | 1.9.0 | +| -Pspark-3.2 | | 1.4.0 | 1.10.0 | +| -Pspark-3.3 | | 1.6.0 | 1.13.0 | +| -Pspark-3.4 | | 1.8.0 | 1.13.0 | +| -Pspark-3.5 | ✓ | 1.8.0 | | +| -Pspark-4.0 | | 1.10.0 | | +| -Pspark-4.1 | | 1.11.0 | | +| -Pspark-4.2 | | 1.12.0 | | ## Building Kyuubi Against Different Scala Versions @@ -73,7 +79,7 @@ Since v1.8.0, Kyuubi support building with different Scala profile. Currently, K | Profile | Default | Since | |--------------|---------|-------| -| (Scala 2.12) | ✓ | - | +| -Pscala-2.12 | ✓ | - | | -Pscala-2.13 | | 1.8.0 | Please activate `scala-2.13` profile when Scala 2.13 support is needed. The GA tests have covered integration test with the Kyuubi server, engines and related plugins, while the Flink engine and it's integration tests are not included for the reason that Flink does not support Scala 2.13 yet and will pull out client support for Scala. diff --git a/docs/deployment/migration-guide.md b/docs/deployment/migration-guide.md index 7bf37c6d2ce..bdcadd87609 100644 --- a/docs/deployment/migration-guide.md +++ b/docs/deployment/migration-guide.md @@ -17,6 +17,10 @@ # Kyuubi Migration Guide +## Upgrading from Kyuubi 1.12 to 1.13 + +* Since Kyuubi 1.13, the support of Spark engine for Spark 3.3 and 3.4 is removed. + ## Upgrading from Kyuubi 1.11 to 1.12 * Since Kyuubi 1.12, the support of variable `` substitution in config `kyuubi.metadata.store.jdbc.url` is deprecated, use `{{KYUUBI_HOME}}` instead. diff --git a/docs/extensions/engines/spark/jvm-quake.md b/docs/extensions/engines/spark/jvm-quake.md index 5f350be0a2b..b69ab0ef44a 100644 --- a/docs/extensions/engines/spark/jvm-quake.md +++ b/docs/extensions/engines/spark/jvm-quake.md @@ -23,7 +23,7 @@ So introducing JVMQuake provides more granular monitoring of GC behavior, enabli ## Usage -JVM Quake is implemented through Spark plugins, This plugin technically supports Spark 3.0 onwards, but was only verified with Spark 3.3 to 4.0 in CI. +JVM Quake is implemented through Spark plugins, This plugin technically supports Spark 3.0 onwards, but was only verified with Spark 3.5 to 4.2 in CI. ### Build with Apache Maven diff --git a/docs/extensions/engines/spark/lineage.md b/docs/extensions/engines/spark/lineage.md index 73b9f6af059..a5863eb04f0 100644 --- a/docs/extensions/engines/spark/lineage.md +++ b/docs/extensions/engines/spark/lineage.md @@ -122,18 +122,18 @@ build/mvn clean package -pl :kyuubi-spark-lineage_2.12 -am -DskipTests -Dspark.v The available `spark.version`s are shown in the following table. -| Spark Version | Supported | Remark | -|:-------------:|:---------:|:------:| -| master | √ | - | -| 3.5.x | √ | - | -| 3.4.x | √ | - | -| 3.3.x | √ | - | -| 3.2.x | √ | - | -| 3.1.x | x | - | -| 3.0.x | x | - | -| 2.4.x | x | - | - -Currently, Spark released with Scala 2.12 are supported. +| Spark Version | Supported | +|:-------------:|:---------:| +| 4.2.x | √ | +| 4.1.x | √ | +| 4.0.x | √ | +| 3.5.x | √ | +| 3.4.x | x | +| 3.3.x | x | +| 3.2.x | x | +| 3.1.x | x | +| 3.0.x | x | +| 2.4.x | x | ### Test with ScalaTest Maven plugin diff --git a/docs/extensions/engines/spark/rules.md b/docs/extensions/engines/spark/rules.md index 4b8c52ebe59..5dd337d8382 100644 --- a/docs/extensions/engines/spark/rules.md +++ b/docs/extensions/engines/spark/rules.md @@ -17,7 +17,7 @@ # Auxiliary Optimization Rules -Kyuubi provides SQL extension out of box. Due to the version compatibility with Apache Spark, currently we support Apache Spark branch-3.1 and later. +Kyuubi provides SQL extension out of box. Due to the version compatibility with Apache Spark, currently we support Apache Spark 3.5 and later. And don't worry, Kyuubi will support the new Apache Spark version in the future. Thanks to the adaptive query execution framework (AQE), Kyuubi can do these optimizations. ## Features @@ -57,20 +57,20 @@ And don't worry, Kyuubi will support the new Apache Spark version in the future. ## Usage -| Kyuubi Spark SQL extension | Supported Spark version(s) | Available since | EOL | Bundled in Binary release tarball | Maven profile | -|----------------------------|----------------------------|------------------|-------|-----------------------------------|---------------| -| kyuubi-extension-spark-3-1 | 3.1.x | 1.3.0-incubating | 1.8.0 | 1.3.0-incubating | spark-3.1 | -| kyuubi-extension-spark-3-2 | 3.2.x | 1.4.0-incubating | 1.9.0 | 1.4.0-incubating | spark-3.2 | -| kyuubi-extension-spark-3-3 | 3.3.x | 1.6.0-incubating | N/A | 1.6.0-incubating | spark-3.3 | -| kyuubi-extension-spark-3-4 | 3.4.x | 1.8.0 | N/A | 1.8.0 | spark-3.4 | -| kyuubi-extension-spark-3-5 | 3.5.x | 1.8.0 | N/A | 1.9.0 | spark-3.5 | -| kyuubi-extension-spark-4-0 | 4.0.x | 1.10.0 | N/A | 1.10.0 | spark-4.0 | -| kyuubi-extension-spark-4-1 | 4.1.x | 1.11.0 | N/A | 1.11.0 | spark-4.1 | -| kyuubi-extension-spark-4-2 | 4.2.x | 1.12.0 | N/A | 1.12.0 | spark-4.2 | +| Kyuubi Spark SQL extension | Supported Spark version(s) | Available since | EOL | Bundled in Binary release tarball | Maven profile | +|----------------------------|----------------------------|------------------|--------|-----------------------------------|---------------| +| kyuubi-extension-spark-3-1 | 3.1.x | 1.3.0-incubating | 1.8.0 | 1.3.0-incubating - 1.4.1 | spark-3.1 | +| kyuubi-extension-spark-3-2 | 3.2.x | 1.4.0-incubating | 1.9.0 | 1.5.0 - 1.6.1 | spark-3.2 | +| kyuubi-extension-spark-3-3 | 3.3.x | 1.6.0-incubating | 1.12.0 | N/A | spark-3.3 | +| kyuubi-extension-spark-3-4 | 3.4.x | 1.8.0 | 1.12.0 | N/A | spark-3.4 | +| kyuubi-extension-spark-3-5 | 3.5.x | 1.8.0 | N/A | N/A | spark-3.5 | +| kyuubi-extension-spark-4-0 | 4.0.x | 1.10.0 | N/A | N/A | spark-4.0 | +| kyuubi-extension-spark-4-1 | 4.1.x | 1.11.0 | N/A | N/A | spark-4.1 | +| kyuubi-extension-spark-4-2 | 4.2.x | 1.12.0 | N/A | N/A | spark-4.2 | 1. Check the matrix that if you are using the supported Spark version, and find the corresponding Kyuubi Spark SQL Extension jar 2. Get the Kyuubi Spark SQL Extension jar - 1. Each Kyuubi binary release tarball only contains one default version of Kyuubi Spark SQL Extension jar, if you are looking for such version, you can find it under `$KYUUBI_HOME/extension` + 1. Kyuubi binary release tarballs bundled one default version of the Kyuubi Spark SQL Extension jar until 1.6.1, and it is no longer bundled since 1.7.0 2. All supported versions of Kyuubi Spark SQL Extension jar will be deployed to [Maven Central](https://search.maven.org/search?q=kyuubi-extension-spark) 3. If you like, you can compile Kyuubi Spark SQL Extension jar by yourself, please activate the corresponding Maven's profile on you compile command, i.e. you can get Kyuubi Spark SQL Extension jar for Spark 3.5 under `extensions/spark/kyuubi-extension-spark-3-5/target` when compile with `-Pspark-3.5` 3. Put the Kyuubi Spark SQL extension jar `kyuubi-extension-spark-*.jar` into `$SPARK_HOME/jars` diff --git a/docs/quick_start/quick_start.rst b/docs/quick_start/quick_start.rst index bd97cfec562..eacc6f18673 100644 --- a/docs/quick_start/quick_start.rst +++ b/docs/quick_start/quick_start.rst @@ -43,7 +43,7 @@ pre-installed and the ``JAVA_HOME`` is correctly set to each component. **Kyuubi** Gateway \ |release| \ - Kyuubi Server Engine lib - Kyuubi Engine Beeline - Kyuubi Beeline - **Spark** Engine 3.3 to 3.5, 4.0 to 4.2 A Spark distribution + **Spark** Engine 3.5, 4.0 to 4.2 A Spark distribution **Flink** Engine 1.17 to 1.20, 2.0 to 2.3 A Flink distribution **Trino** Engine N/A A Trino cluster allows to access via trino-client v411 **Doris** Engine N/A A Doris cluster @@ -146,7 +146,7 @@ To install Spark, you need to unpack the tarball. For example, .. code-block:: - $ tar zxf spark-3.4.2-bin-hadoop3.tgz + $ tar zxf spark-3.5.8-bin-hadoop3.tgz Configuration ~~~~~~~~~~~~~ diff --git a/docs/security/authorization/spark/build.md b/docs/security/authorization/spark/build.md index 9d6a18a5ef6..abfebdb471b 100644 --- a/docs/security/authorization/spark/build.md +++ b/docs/security/authorization/spark/build.md @@ -53,23 +53,22 @@ Sometimes, it may be incompatible with other Spark distributions, then you may n For example, ```shell -build/mvn clean package -pl :kyuubi-spark-authz_2.12 -am -DskipTests -Pspark-3.4 -Dspark.version=3.4.1 +build/mvn clean package -pl :kyuubi-spark-authz_2.12 -am -DskipTests -Pspark-3.5 -Dspark.version=3.5.8 ``` The available `spark.version`s are shown in the following table. -| Spark Version | Supported | Remark | -|:-----------------:|:---------:|:----------------------------------------------------------------------------------------------------------------------:| -| master | √ | - | -| 3.5.x | √ | - | -| 3.4.x | √ | - | -| 3.3.x | √ | - | -| 3.2.x | x | EOL since v1.11.0 | -| 3.1.x | x | EOL since v1.10.0 | -| 3.0.x | x | EOL since v1.9.0 | -| 2.4.x and earlier | × | [PR 2367](https://github.com/apache/kyuubi/pull/2367) is used to track how we work with older releases with scala 2.11 | - -Currently, Spark released with Scala 2.12 are supported. +| Spark Version | Supported | Removed | +|:-------------:|:---------:|:-------:| +| 4.2.x | √ | - | +| 4.1.x | √ | - | +| 4.0.x | √ | - | +| 3.5.x | √ | - | +| 3.4.x | x | 1.13.0 | +| 3.3.x | x | 1.13.0 | +| 3.2.x | x | 1.11.0 | +| 3.1.x | x | 1.10.0 | +| 3.0.x | x | 1.9.0 | ### Build against Different Apache Ranger Versions @@ -83,20 +82,14 @@ build/mvn clean package -pl :kyuubi-spark-authz_2.12 -am -DskipTests -Dranger.ve The available `ranger.version`s are shown in the following table. -| Ranger Version | Supported | Remark | -|:--------------:|:---------:|:-----------------------------------------------------------------------------------------:| -| 2.6.x | √ | - | -| 2.5.x | √ | - | -| 2.4.x | √ | - | -| 2.3.x | √ | - | -| 2.2.x | √ | - | -| 2.1.x | √ | - | -| 2.0.x | √ | - | -| 1.2.x | √ | - | -| 1.1.x | √ | - | -| 1.0.x | √ | - | -| 0.7.x | √ | - | -| 0.6.x | X | [KYUUBI-4672](https://github.com/apache/kyuubi/issues/4672) reported unresolved failures. | +| Ranger Version | Supported | Remark | +|:--------------:|:---------:|:------:| +| 2.6.x | √ | - | +| 2.5.x | √ | - | +| 2.4.x | √ | - | +| 2.3.x | √ | - | +| 2.2.x | √ | - | +| 2.1.x | √ | - | Currently, all ranger releases are supported. diff --git a/extensions/spark/kyuubi-extension-spark-3-3/benchmarks/ZorderCoreBenchmark-results.txt b/extensions/spark/kyuubi-extension-spark-3-3/benchmarks/ZorderCoreBenchmark-results.txt deleted file mode 100644 index 0ed1cbf7b73..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/benchmarks/ZorderCoreBenchmark-results.txt +++ /dev/null @@ -1,18 +0,0 @@ -Java HotSpot(TM) 64-Bit Server VM 1.8.0_271-b09 on Mac OS X 10.16 -Intel(R) Core(TM) i7-4770HQ CPU @ 2.20GHz -1000000 rows zorder core benchmark: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------- -2 int columns benchmark 191 201 12 5.2 191.1 1.0X -3 int columns benchmark 262 304 72 3.8 261.5 0.7X -4 int columns benchmark 302 310 9 3.3 302.1 0.6X -2 long columns benchmark 185 188 3 5.4 185.4 1.0X -3 long columns benchmark 241 243 3 4.2 240.6 0.8X -4 long columns benchmark 291 335 69 3.4 290.6 0.7X - -Java HotSpot(TM) 64-Bit Server VM 1.8.0_271-b09 on Mac OS X 10.16 -Intel(R) Core(TM) i7-4770HQ CPU @ 2.20GHz -10000000 iterations paddingTo8Byte benchmark: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ----------------------------------------------------------------------------------------------------------------------------- -2 length benchmark 163 168 5 61.2 16.3 1.0X -16 length benchmark 154 155 2 65.1 15.4 1.1X - diff --git a/extensions/spark/kyuubi-extension-spark-3-3/pom.xml b/extensions/spark/kyuubi-extension-spark-3-3/pom.xml deleted file mode 100644 index 4273add03cd..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/pom.xml +++ /dev/null @@ -1,222 +0,0 @@ - - - - 4.0.0 - - org.apache.kyuubi - kyuubi-parent - 1.13.0-SNAPSHOT - ../../../pom.xml - - - kyuubi-extension-spark-3-3_${scala.binary.version} - jar - Kyuubi Dev Spark Extensions (for Spark 3.3) - https://kyuubi.apache.org/ - - - - org.scala-lang - scala-library - provided - - - - org.apache.spark - spark-sql_${scala.binary.version} - provided - - - - org.apache.spark - spark-hive_${scala.binary.version} - provided - - - - org.apache.hadoop - hadoop-client-api - provided - - - - org.apache.kyuubi - kyuubi-download - ${project.version} - pom - test - - - - org.apache.kyuubi - kyuubi-util-scala_${scala.binary.version} - ${project.version} - test-jar - test - - - - org.apache.spark - spark-core_${scala.binary.version} - test-jar - test - - - - org.apache.spark - spark-catalyst_${scala.binary.version} - test-jar - test - - - - org.scalatestplus - scalacheck-1-17_${scala.binary.version} - test - - - - org.apache.spark - spark-sql_${scala.binary.version} - ${spark.version} - test-jar - test - - - - org.apache.hadoop - hadoop-client-runtime - test - - - - javax.servlet - javax.servlet-api - test - - - - jakarta.servlet - jakarta.servlet-api - test - - - - org.apache.logging.log4j - log4j-slf4j-impl - test - - - - - - - - org.codehaus.mojo - build-helper-maven-plugin - - - regex-property - - regex-property - - - spark.home - ${project.basedir}/../../../externals/kyuubi-download/target/${spark.archive.name} - (.+)\.tgz - $1 - - - - - - org.antlr - antlr4-maven-plugin - - true - ${project.basedir}/src/main/antlr4 - - - - - org.apache.maven.plugins - maven-shade-plugin - - false - - - org.apache.kyuubi:* - - - - - - - shade - - package - - - - - org.scalatest - scalatest-maven-plugin - - - - ${scala.binary.version} - - - - - target/scala-${scala.binary.version}/classes - target/scala-${scala.binary.version}/test-classes - - - - - spark-home-from-archive - - - !env.SPARK_HOME - - - - - - org.scalatest - scalatest-maven-plugin - - - ${spark.home} - - - - - - - - diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 b/extensions/spark/kyuubi-extension-spark-3-3/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 deleted file mode 100644 index e52b7f5cfeb..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 +++ /dev/null @@ -1,191 +0,0 @@ -/* - * 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. - */ - -grammar KyuubiSparkSQL; - -@members { - /** - * Verify whether current token is a valid decimal token (which contains dot). - * Returns true if the character that follows the token is not a digit or letter or underscore. - * - * For example: - * For char stream "2.3", "2." is not a valid decimal token, because it is followed by digit '3'. - * For char stream "2.3_", "2.3" is not a valid decimal token, because it is followed by '_'. - * For char stream "2.3W", "2.3" is not a valid decimal token, because it is followed by 'W'. - * For char stream "12.0D 34.E2+0.12 " 12.0D is a valid decimal token because it is followed - * by a space. 34.E2 is a valid decimal token because it is followed by symbol '+' - * which is not a digit or letter or underscore. - */ - public boolean isValidDecimal() { - int nextChar = _input.LA(1); - if (nextChar >= 'A' && nextChar <= 'Z' || nextChar >= '0' && nextChar <= '9' || - nextChar == '_') { - return false; - } else { - return true; - } - } - } - -tokens { - DELIMITER -} - -singleStatement - : statement EOF - ; - -statement - : OPTIMIZE multipartIdentifier whereClause? zorderClause #optimizeZorder - | .*? #passThrough - ; - -whereClause - : WHERE partitionPredicate = predicateToken - ; - -zorderClause - : ZORDER BY order+=multipartIdentifier (',' order+=multipartIdentifier)* - ; - -// We don't have an expression rule in our grammar here, so we just grab the tokens and defer -// parsing them to later. -predicateToken - : .+? - ; - -multipartIdentifier - : parts+=identifier ('.' parts+=identifier)* - ; - -identifier - : strictIdentifier - ; - -strictIdentifier - : IDENTIFIER #unquotedIdentifier - | quotedIdentifier #quotedIdentifierAlternative - | nonReserved #unquotedIdentifier - ; - -quotedIdentifier - : BACKQUOTED_IDENTIFIER - ; - -nonReserved - : AND - | BY - | FALSE - | DATE - | INTERVAL - | OPTIMIZE - | OR - | TABLE - | TIMESTAMP - | TRUE - | WHERE - | ZORDER - ; - -AND: 'AND'; -BY: 'BY'; -FALSE: 'FALSE'; -DATE: 'DATE'; -INTERVAL: 'INTERVAL'; -OPTIMIZE: 'OPTIMIZE'; -OR: 'OR'; -TABLE: 'TABLE'; -TIMESTAMP: 'TIMESTAMP'; -TRUE: 'TRUE'; -WHERE: 'WHERE'; -ZORDER: 'ZORDER'; - -MINUS: '-'; - -BIGINT_LITERAL - : DIGIT+ 'L' - ; - -SMALLINT_LITERAL - : DIGIT+ 'S' - ; - -TINYINT_LITERAL - : DIGIT+ 'Y' - ; - -INTEGER_VALUE - : DIGIT+ - ; - -DECIMAL_VALUE - : DIGIT+ EXPONENT - | DECIMAL_DIGITS EXPONENT? {isValidDecimal()}? - ; - -DOUBLE_LITERAL - : DIGIT+ EXPONENT? 'D' - | DECIMAL_DIGITS EXPONENT? 'D' {isValidDecimal()}? - ; - -BIGDECIMAL_LITERAL - : DIGIT+ EXPONENT? 'BD' - | DECIMAL_DIGITS EXPONENT? 'BD' {isValidDecimal()}? - ; - -BACKQUOTED_IDENTIFIER - : '`' ( ~'`' | '``' )* '`' - ; - -IDENTIFIER - : (LETTER | DIGIT | '_')+ - ; - -fragment DECIMAL_DIGITS - : DIGIT+ '.' DIGIT* - | '.' DIGIT+ - ; - -fragment EXPONENT - : 'E' [+-]? DIGIT+ - ; - -fragment DIGIT - : [0-9] - ; - -fragment LETTER - : [A-Z] - ; - -SIMPLE_COMMENT - : '--' ~[\r\n]* '\r'? '\n'? -> channel(HIDDEN) - ; - -BRACKETED_COMMENT - : '/*' .*? '*/' -> channel(HIDDEN) - ; - -WS : [ \r\n\t]+ -> channel(HIDDEN) - ; - -// Catch-all for anything we can't recognize. -// We use this to be able to ignore and recover all the text -// when splitting statements with DelimiterLexer -UNRECOGNIZED - : . - ; diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala deleted file mode 100644 index e740554d443..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala +++ /dev/null @@ -1,54 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunc, UnresolvedRelation, UnresolvedTableOrView, UnresolvedView} -import org.apache.spark.sql.catalyst.plans.logical.{DropFunction, DropNamespace, DropTable, DropView, LogicalPlan, NoopCommand, UncacheTable} -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.command.{AlterTableDropPartitionCommand, DropTableCommand} - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -case class DropIgnoreNonexistent(session: SparkSession) extends Rule[LogicalPlan] { - - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(DROP_IGNORE_NONEXISTENT)) { - plan match { - case i @ AlterTableDropPartitionCommand(_, _, false, _, _) => - i.copy(ifExists = true) - case i @ DropTableCommand(_, false, _, _) => - i.copy(ifExists = true) - case i @ DropNamespace(_, false, _) => - i.copy(ifExists = true) - // like: org.apache.spark.sql.catalyst.analysis.ResolveCommandsWithIfExists - case DropTable(u: UnresolvedTableOrView, false, _) => - NoopCommand("DROP TABLE", u.multipartIdentifier) - case DropView(u: UnresolvedView, false) => - NoopCommand("DROP VIEW", u.multipartIdentifier) - case UncacheTable(u: UnresolvedRelation, false, _) => - NoopCommand("UNCACHE TABLE", u.multipartIdentifier) - case DropFunction(u: UnresolvedFunc, false) => - NoopCommand("DROP FUNCTION", u.multipartIdentifier) - case _ => plan - } - } else { - plan - } - } - -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala deleted file mode 100644 index fcbf5c0a122..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala +++ /dev/null @@ -1,110 +0,0 @@ -/* - * 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.sql - -import scala.annotation.tailrec - -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Expression, NamedExpression, UnaryExpression} -import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys -import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftAnti, LeftOuter, LeftSemi, RightOuter} -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, Project, Sort, SubqueryAlias, View} - -/** - * Infer the columns for Rebalance and Sort to improve the compression ratio. - * - * For example - * {{{ - * INSERT INTO TABLE t PARTITION(p='a') - * SELECT * FROM t1 JOIN t2 on t1.c1 = t2.c1 - * }}} - * the inferred columns are: t1.c1 - */ -object InferRebalanceAndSortOrders { - - type PartitioningAndOrdering = (Seq[Expression], Seq[Expression]) - - private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = { - @tailrec - def throughUnary(e: Expression): Expression = e match { - case u: UnaryExpression if u.deterministic => - throughUnary(u.child) - case _ => e - } - - named.flatMap { - case a @ Alias(child, _) => - Some((throughUnary(child).canonicalized, a.toAttribute)) - case _ => None - }.toMap - } - - def infer(plan: LogicalPlan): Option[PartitioningAndOrdering] = { - def candidateKeys( - input: LogicalPlan, - output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = { - input match { - case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, _, _, _, _, _) => - joinType match { - case LeftSemi | LeftAnti | LeftOuter => Some((leftKeys, leftKeys)) - case RightOuter => Some((rightKeys, rightKeys)) - case Inner | FullOuter => - if (output.isEmpty) { - Some((leftKeys ++ rightKeys, leftKeys ++ rightKeys)) - } else { - assert(leftKeys.length == rightKeys.length) - val keys = leftKeys.zip(rightKeys).flatMap { case (left, right) => - if (left.references.subsetOf(output)) { - Some(left) - } else if (right.references.subsetOf(output)) { - Some(right) - } else { - None - } - } - Some((keys, keys)) - } - case _ => None - } - case agg: Aggregate => - val aliasMap = getAliasMap(agg.aggregateExpressions) - Some(( - agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)), - agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o)))) - case s: Sort => Some((s.order.map(_.child), s.order.map(_.child))) - case p: Project => - val aliasMap = getAliasMap(p.projectList) - candidateKeys(p.child, p.references).map { case (partitioning, ordering) => - ( - partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)), - ordering.map(o => aliasMap.getOrElse(o.canonicalized, o))) - } - case f: Filter => candidateKeys(f.child, output) - case s: SubqueryAlias => candidateKeys(s.child, output) - case v: View => candidateKeys(v.child, output) - - case _ => None - } - } - - candidateKeys(plan).map { case (partitioning, ordering) => - ( - partitioning.filter(_.references.subsetOf(plan.outputSet)), - ordering.filter(_.references.subsetOf(plan.outputSet))) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala deleted file mode 100644 index 92626f02745..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala +++ /dev/null @@ -1,93 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.plans.physical.Distribution -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.QueryStageExec -import org.apache.spark.sql.execution.aggregate.BaseAggregateExec -import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * Insert shuffle node before join if it doesn't exist to make `OptimizeSkewedJoin` works. - */ -object InsertShuffleNodeBeforeJoin extends Rule[SparkPlan] { - - override def apply(plan: SparkPlan): SparkPlan = { - // this rule has no meaning without AQE - if (!conf.getConf(FORCE_SHUFFLE_BEFORE_JOIN) || - !conf.getConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)) { - return plan - } - - val newPlan = insertShuffleBeforeJoin(plan) - if (plan.fastEquals(newPlan)) { - plan - } else { - // make sure the output partitioning and ordering will not be broken. - KyuubiEnsureRequirements.apply(newPlan) - } - } - - // SPARK-33832 (Spark 3.3) moves the rule OptimizeSkewedJoin from queryStageOptimizerRules - // to queryStagePreparationRules, injecting shuffle after OptimizeSkewedJoin may produce - // invalid query plan. - private def insertShuffleBeforeJoin(plan: SparkPlan): SparkPlan = plan transformUp { - case smj @ SortMergeJoinExec(_, _, _, _, l, r, isSkewJoin) if !isSkewJoin => - smj.withNewChildren(checkAndInsertShuffle(smj.requiredChildDistribution.head, l) :: - checkAndInsertShuffle(smj.requiredChildDistribution(1), r) :: Nil) - - case shj: ShuffledHashJoinExec if !shj.isSkewJoin => - if (!shj.left.isInstanceOf[Exchange] && !shj.right.isInstanceOf[Exchange]) { - shj.withNewChildren(withShuffleExec(shj.requiredChildDistribution.head, shj.left) :: - withShuffleExec(shj.requiredChildDistribution(1), shj.right) :: Nil) - } else if (!shj.left.isInstanceOf[Exchange]) { - shj.withNewChildren( - withShuffleExec(shj.requiredChildDistribution.head, shj.left) :: shj.right :: Nil) - } else if (!shj.right.isInstanceOf[Exchange]) { - shj.withNewChildren( - shj.left :: withShuffleExec(shj.requiredChildDistribution(1), shj.right) :: Nil) - } else { - shj - } - } - - private def checkAndInsertShuffle( - distribution: Distribution, - child: SparkPlan): SparkPlan = child match { - case SortExec(_, _, _: Exchange, _) => - child - case SortExec(_, _, _: QueryStageExec, _) => - child - case sort @ SortExec(_, _, agg: BaseAggregateExec, _) => - sort.withNewChildren(withShuffleExec(distribution, agg) :: Nil) - case _ => - withShuffleExec(distribution, child) - } - - private def withShuffleExec(distribution: Distribution, child: SparkPlan): SparkPlan = { - val numPartitions = distribution.requiredNumPartitions - .getOrElse(conf.numShufflePartitions) - ShuffleExchangeExec(distribution.createPartitioning(numPartitions), child) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala deleted file mode 100644 index a17e0a4652b..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala +++ /dev/null @@ -1,133 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.expressions.SortOrder -import org.apache.spark.sql.catalyst.plans.physical.{BroadcastDistribution, Distribution, UnspecifiedDistribution} -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} - -/** - * Copy from Apache Spark `EnsureRequirements` - * 1. remove reorder join predicates - * 2. remove shuffle pruning - */ -object KyuubiEnsureRequirements extends Rule[SparkPlan] { - private def ensureDistributionAndOrdering(operator: SparkPlan): SparkPlan = { - val requiredChildDistributions: Seq[Distribution] = operator.requiredChildDistribution - val requiredChildOrderings: Seq[Seq[SortOrder]] = operator.requiredChildOrdering - var children: Seq[SparkPlan] = operator.children - assert(requiredChildDistributions.length == children.length) - assert(requiredChildOrderings.length == children.length) - - // Ensure that the operator's children satisfy their output distribution requirements. - children = children.zip(requiredChildDistributions).map { - case (child, distribution) if child.outputPartitioning.satisfies(distribution) => - child - case (child, BroadcastDistribution(mode)) => - BroadcastExchangeExec(mode, child) - case (child, distribution) => - val numPartitions = distribution.requiredNumPartitions - .getOrElse(conf.numShufflePartitions) - ShuffleExchangeExec(distribution.createPartitioning(numPartitions), child) - } - - // Get the indexes of children which have specified distribution requirements and need to have - // same number of partitions. - val childrenIndexes = requiredChildDistributions.zipWithIndex.filter { - case (UnspecifiedDistribution, _) => false - case (_: BroadcastDistribution, _) => false - case _ => true - }.map(_._2) - - val childrenNumPartitions = - childrenIndexes.map(children(_).outputPartitioning.numPartitions).toSet - - if (childrenNumPartitions.size > 1) { - // Get the number of partitions which is explicitly required by the distributions. - val requiredNumPartitions = { - val numPartitionsSet = childrenIndexes.flatMap { - index => requiredChildDistributions(index).requiredNumPartitions - }.toSet - assert( - numPartitionsSet.size <= 1, - s"$operator have incompatible requirements of the number of partitions for its children") - numPartitionsSet.headOption - } - - // If there are non-shuffle children that satisfy the required distribution, we have - // some tradeoffs when picking the expected number of shuffle partitions: - // 1. We should avoid shuffling these children. - // 2. We should have a reasonable parallelism. - val nonShuffleChildrenNumPartitions = - childrenIndexes.map(children).filterNot(_.isInstanceOf[ShuffleExchangeExec]) - .map(_.outputPartitioning.numPartitions) - val expectedChildrenNumPartitions = - if (nonShuffleChildrenNumPartitions.nonEmpty) { - if (nonShuffleChildrenNumPartitions.length == childrenIndexes.length) { - // Here we pick the max number of partitions among these non-shuffle children. - nonShuffleChildrenNumPartitions.max - } else { - // Here we pick the max number of partitions among these non-shuffle children as the - // expected number of shuffle partitions. However, if it's smaller than - // `conf.numShufflePartitions`, we pick `conf.numShufflePartitions` as the - // expected number of shuffle partitions. - math.max(nonShuffleChildrenNumPartitions.max, conf.defaultNumShufflePartitions) - } - } else { - childrenNumPartitions.max - } - - val targetNumPartitions = requiredNumPartitions.getOrElse(expectedChildrenNumPartitions) - - children = children.zip(requiredChildDistributions).zipWithIndex.map { - case ((child, distribution), index) if childrenIndexes.contains(index) => - if (child.outputPartitioning.numPartitions == targetNumPartitions) { - child - } else { - val defaultPartitioning = distribution.createPartitioning(targetNumPartitions) - child match { - // If child is an exchange, we replace it with a new one having defaultPartitioning. - case ShuffleExchangeExec(_, c, _) => ShuffleExchangeExec(defaultPartitioning, c) - case _ => ShuffleExchangeExec(defaultPartitioning, child) - } - } - - case ((child, _), _) => child - } - } - - // Now that we've performed any necessary shuffles, add sorts to guarantee output orderings: - children = children.zip(requiredChildOrderings).map { case (child, requiredOrdering) => - // If child.outputOrdering already satisfies the requiredOrdering, we do not need to sort. - if (SortOrder.orderingSatisfies(child.outputOrdering, requiredOrdering)) { - child - } else { - SortExec(requiredOrdering, global = false, child = child) - } - } - - operator.withNewChildren(children) - } - - def apply(plan: SparkPlan): SparkPlan = plan.transformUp { - case operator: SparkPlan => - ensureDistributionAndOrdering(operator) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala deleted file mode 100644 index 483d3ed73ab..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala +++ /dev/null @@ -1,191 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec -import org.apache.spark.sql.execution.command.{ResetCommand, SetCommand} -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * This rule split stage into two parts: - * 1. previous stage - * 2. final stage - * For final stage, we can inject extra config. It's useful if we use repartition to optimize - * small files that needs bigger shuffle partition size than previous. - * - * Let's say we have a query with 3 stages, then the logical machine like: - * - * Set/Reset Command -> cleanup previousStage config if user set the spark config. - * Query -> AQE -> stage1 -> preparation (use previousStage to overwrite spark config) - * -> AQE -> stage2 -> preparation (use spark config) - * -> AQE -> stage3 -> preparation (use finalStage config to overwrite spark config, - * store spark config to previousStage.) - * - * An example of the new finalStage config: - * `spark.sql.adaptive.advisoryPartitionSizeInBytes` -> - * `spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes` - */ -case class FinalStageConfigIsolation(session: SparkSession) extends Rule[SparkPlan] { - import FinalStageConfigIsolation._ - - override def apply(plan: SparkPlan): SparkPlan = { - // this rule has no meaning without AQE - if (!conf.getConf(FINAL_STAGE_CONFIG_ISOLATION) || - !conf.getConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)) { - return plan - } - - if (isFinalStage(plan)) { - // We can not get the whole plan at query preparation phase to detect if current plan is - // for writing, so we depend on a tag which is been injected at post resolution phase. - // Note: we should still do clean up previous config for non-final stage to avoid such case: - // the first statement is write, but the second statement is query. - if (conf.getConf(FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY) && - !MarkNumOutputColumnsRule.isWrite(session, plan)) { - return plan - } - - // set config for final stage - session.conf.getAll.filter(_._1.startsWith(FINAL_STAGE_CONFIG_PREFIX)).foreach { - case (k, v) => - val sparkConfigKey = s"spark.sql.${k.substring(FINAL_STAGE_CONFIG_PREFIX.length)}" - val previousStageConfigKey = - s"$PREVIOUS_STAGE_CONFIG_PREFIX${k.substring(FINAL_STAGE_CONFIG_PREFIX.length)}" - // store the previous config only if we have not stored, to avoid some query only - // have one stage that will overwrite real config. - if (!session.sessionState.conf.contains(previousStageConfigKey)) { - val originalValue = - if (session.conf.getOption(sparkConfigKey).isDefined) { - session.sessionState.conf.getConfString(sparkConfigKey) - } else { - // the default value of config is None, so we need to use a internal tag - INTERNAL_UNSET_CONFIG_TAG - } - logInfo(s"Store config: $sparkConfigKey to previousStage, " + - s"original value: $originalValue ") - session.sessionState.conf.setConfString(previousStageConfigKey, originalValue) - } - logInfo(s"For final stage: set $sparkConfigKey = $v.") - session.conf.set(sparkConfigKey, v) - } - } else { - // reset config for previous stage - session.conf.getAll.filter(_._1.startsWith(PREVIOUS_STAGE_CONFIG_PREFIX)).foreach { - case (k, v) => - val sparkConfigKey = s"spark.sql.${k.substring(PREVIOUS_STAGE_CONFIG_PREFIX.length)}" - logInfo(s"For previous stage: set $sparkConfigKey = $v.") - if (v == INTERNAL_UNSET_CONFIG_TAG) { - session.conf.unset(sparkConfigKey) - } else { - session.conf.set(sparkConfigKey, v) - } - // unset config so that we do not need to reset configs for every previous stage - session.conf.unset(k) - } - } - - plan - } - - private def isFinalStage(plan: SparkPlan): Boolean = { - var shuffleNum = 0 - var broadcastNum = 0 - var reusedNum = 0 - var queryStageNum = 0 - - def collectNumber(p: SparkPlan): SparkPlan = { - p transform { - case shuffle: ShuffleExchangeLike => - shuffleNum += 1 - shuffle - - case broadcast: BroadcastExchangeLike => - broadcastNum += 1 - broadcast - - case reusedExchangeExec: ReusedExchangeExec => - reusedNum += 1 - reusedExchangeExec - - // query stage is leaf node so we need to transform it manually - // compatible with Spark 3.5: - // SPARK-42101: table cache is a independent query stage, so do not need include it. - case queryStage: QueryStageExec if queryStage.nodeName != "TableCacheQueryStage" => - queryStageNum += 1 - collectNumber(queryStage.plan) - queryStage - } - } - collectNumber(plan) - - if (shuffleNum == 0) { - // we don not care about broadcast stage here since it won't change partition number. - true - } else if (shuffleNum + broadcastNum + reusedNum == queryStageNum) { - true - } else { - false - } - } -} -object FinalStageConfigIsolation { - final val SQL_PREFIX = "spark.sql." - final val FINAL_STAGE_CONFIG_PREFIX = "spark.sql.finalStage." - final val PREVIOUS_STAGE_CONFIG_PREFIX = "spark.sql.previousStage." - final val INTERNAL_UNSET_CONFIG_TAG = "__INTERNAL_UNSET_CONFIG_TAG__" - - def getPreviousStageConfigKey(configKey: String): Option[String] = { - if (configKey.startsWith(SQL_PREFIX)) { - Some(s"$PREVIOUS_STAGE_CONFIG_PREFIX${configKey.substring(SQL_PREFIX.length)}") - } else { - None - } - } -} - -case class FinalStageConfigIsolationCleanRule(session: SparkSession) extends Rule[LogicalPlan] { - import FinalStageConfigIsolation._ - - override def apply(plan: LogicalPlan): LogicalPlan = plan match { - case set @ SetCommand(Some((k, Some(_)))) if k.startsWith(SQL_PREFIX) => - checkAndUnsetPreviousStageConfig(k) - set - - case reset @ ResetCommand(Some(k)) if k.startsWith(SQL_PREFIX) => - checkAndUnsetPreviousStageConfig(k) - reset - - case other => other - } - - private def checkAndUnsetPreviousStageConfig(configKey: String): Unit = { - getPreviousStageConfigKey(configKey).foreach { previousStageConfigKey => - if (session.sessionState.conf.contains(previousStageConfigKey)) { - logInfo(s"For previous stage: unset $previousStageConfigKey") - session.conf.unset(previousStageConfigKey) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala deleted file mode 100644 index e72a6c07354..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala +++ /dev/null @@ -1,241 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.network.util.ByteUnit -import org.apache.spark.sql.internal.SQLConf._ - -object KyuubiSQLConf { - - val INSERT_REPARTITION_BEFORE_WRITE = - buildConf("spark.sql.optimizer.insertRepartitionBeforeWrite.enabled") - .doc("Add repartition node at the top of query plan. An approach of merging small files.") - .version("1.2.0") - .booleanConf - .createWithDefault(true) - - val FORCE_SHUFFLE_BEFORE_JOIN = - buildConf("spark.sql.optimizer.forceShuffleBeforeJoin.enabled") - .doc("Ensure shuffle node exists before shuffled join (shj and smj) to make AQE " + - "`OptimizeSkewedJoin` works (complex scenario join, multi table join).") - .version("1.2.0") - .booleanConf - .createWithDefault(false) - - val FINAL_STAGE_CONFIG_ISOLATION = - buildConf("spark.sql.optimizer.finalStageConfigIsolation.enabled") - .doc("If true, the final stage support use different config with previous stage. " + - "The prefix of final stage config key should be `spark.sql.finalStage.`." + - "For example, the raw spark config: `spark.sql.adaptive.advisoryPartitionSizeInBytes`, " + - "then the final stage config should be: " + - "`spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes`.") - .version("1.2.0") - .booleanConf - .createWithDefault(false) - - val INSERT_ZORDER_BEFORE_WRITING = - buildConf("spark.sql.optimizer.insertZorderBeforeWriting.enabled") - .doc("When true, we will follow target table properties to insert zorder or not. " + - "The key properties are: 1) kyuubi.zorder.enabled; if this property is true, we will " + - "insert zorder before writing data. 2) kyuubi.zorder.cols; string split by comma, we " + - "will zorder by these cols.") - .version("1.4.0") - .booleanConf - .createWithDefault(true) - - val ZORDER_GLOBAL_SORT_ENABLED = - buildConf("spark.sql.optimizer.zorderGlobalSort.enabled") - .doc("When true, we do a global sort using zorder. Note that, it can cause data skew " + - "issue if the zorder columns have less cardinality. When false, we only do local sort " + - "using zorder.") - .version("1.4.0") - .booleanConf - .createWithDefault(true) - - val REBALANCE_BEFORE_ZORDER = - buildConf("spark.sql.optimizer.rebalanceBeforeZorder.enabled") - .doc("When true, we do a rebalance before zorder in case data skew. " + - "Note that, if the insertion is dynamic partition we will use the partition " + - "columns to rebalance.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val REBALANCE_ZORDER_COLUMNS_ENABLED = - buildConf("spark.sql.optimizer.rebalanceZorderColumns.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do rebalance before " + - s"Z-Order. If it's dynamic partition insert, the rebalance expression will include " + - s"both partition columns and Z-Order columns.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val TWO_PHASE_REBALANCE_BEFORE_ZORDER = - buildConf("spark.sql.optimizer.twoPhaseRebalanceBeforeZorder.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do two phase rebalance " + - s"before Z-Order for the dynamic partition write. The first phase rebalance using " + - s"dynamic partition column; The second phase rebalance using dynamic partition column + " + - s"Z-Order columns.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val ZORDER_USING_ORIGINAL_ORDERING_ENABLED = - buildConf("spark.sql.optimizer.zorderUsingOriginalOrdering.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do sort by " + - s"the original ordering i.e. lexicographical order.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val WATCHDOG_MAX_PARTITIONS = - buildConf("spark.sql.watchdog.maxPartitions") - .doc("Set the max partition number when spark scans a data source. " + - "Enable maxPartitions Strategy by specifying this configuration. " + - "Add maxPartitions Strategy to avoid scan excessive partitions " + - "on partitioned table, it's optional that works with defined") - .version("1.4.0") - .intConf - .createOptional - - val WATCHDOG_MAX_FILE_SIZE = - buildConf("spark.sql.watchdog.maxFileSize") - .doc("Set the maximum size in bytes of files when spark scans a data source. " + - "Enable maxFileSize Strategy by specifying this configuration. " + - "Add maxFileSize Strategy to avoid scan excessive size of files," + - " it's optional that works with defined") - .version("1.8.0") - .bytesConf(ByteUnit.BYTE) - .createOptional - - val DROP_IGNORE_NONEXISTENT = - buildConf("spark.sql.optimizer.dropIgnoreNonExistent") - .doc("Do not report an error if DROP DATABASE/TABLE/VIEW/FUNCTION/PARTITION specifies " + - "a non-existent database/table/view/function/partition") - .version("1.5.0") - .booleanConf - .createWithDefault(false) - - val INFER_REBALANCE_AND_SORT_ORDERS = - buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.enabled") - .doc("When ture, infer columns for rebalance and sort orders from original query, " + - "e.g. the join keys from join. It can avoid compression ratio regression.") - .version("1.7.0") - .booleanConf - .createWithDefault(false) - - val INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS = - buildConf("spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns") - .doc("The max columns of inferred columns.") - .version("1.7.0") - .intConf - .checkValue(_ > 0, "must be positive number") - .createWithDefault(3) - - val INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE = - buildConf("spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled") - .doc("When true, add repartition even if the original plan does not have shuffle.") - .version("1.7.0") - .booleanConf - .createWithDefault(false) - - val FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY = - buildConf("spark.sql.optimizer.finalStageConfigIsolationWriteOnly.enabled") - .doc("When true, only enable final stage isolation for writing.") - .version("1.7.0") - .booleanConf - .createWithDefault(true) - - val FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED = - buildConf("spark.sql.finalWriteStage.eagerlyKillExecutors.enabled") - .doc("When true, eagerly kill redundant executors before running final write stage.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL = - buildConf("spark.sql.finalWriteStage.eagerlyKillExecutors.killAll") - .doc("When true, eagerly kill all executors before running final write stage. " + - "Mainly for test.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_SKIP_KILLING_EXECUTORS_FOR_TABLE_CACHE = - buildConf("spark.sql.finalWriteStage.skipKillingExecutorsForTableCache") - .doc("When true, skip killing executors if the plan has table caches.") - .version("1.8.0") - .booleanConf - .createWithDefault(true) - - val FINAL_WRITE_STAGE_PARTITION_FACTOR = - buildConf("spark.sql.finalWriteStage.retainExecutorsFactor") - .doc("If the target executors * factor < active executors, and " + - "target executors * factor > min executors, then kill redundant executors.") - .version("1.8.0") - .doubleConf - .checkValue(_ >= 1, "must be bigger than or equal to 1") - .createWithDefault(1.2) - - val FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED = - buildConf("spark.sql.finalWriteStage.resourceIsolation.enabled") - .doc( - "When true, make final write stage resource isolation using custom RDD resource profile.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_EXECUTOR_CORES = - buildConf("spark.sql.finalWriteStage.executorCores") - .doc("Specify the executor core request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .intConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_MEMORY = - buildConf("spark.sql.finalWriteStage.executorMemory") - .doc("Specify the executor on heap memory request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD = - buildConf("spark.sql.finalWriteStage.executorMemoryOverhead") - .doc("Specify the executor memory overhead request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY = - buildConf("spark.sql.finalWriteStage.executorOffHeapMemory") - .doc("Specify the executor off heap memory request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val SCRIPT_TRANSFORMATION_ENABLED = - buildConf("spark.sql.execution.scriptTransformation.enabled") - .doc("When false, script transformation is not allowed.") - .version("1.9.0") - .booleanConf - .createWithDefault(true) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala deleted file mode 100644 index 88c5a988fd9..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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.sql - -import java.sql.SQLException - -class KyuubiSQLExtensionException(reason: String, cause: Throwable) - extends SQLException(reason, cause) { - - def this(reason: String) = { - this(reason, null) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala deleted file mode 100644 index 083623d4edf..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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.sql - -import scala.collection.JavaConverters.asScalaBufferConverter -import scala.collection.mutable.ListBuffer - -import org.antlr.v4.runtime.ParserRuleContext -import org.antlr.v4.runtime.misc.Interval -import org.antlr.v4.runtime.tree.ParseTree -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedRelation, UnresolvedStar} -import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.parser.ParserUtils.withOrigin -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project, Sort} - -import org.apache.kyuubi.sql.KyuubiSparkSQLParser._ -import org.apache.kyuubi.sql.zorder.{OptimizeZorderStatement, Zorder} - -class KyuubiSparkSQLAstBuilder extends KyuubiSparkSQLBaseVisitor[AnyRef] with SQLConfHelper { - - def buildOptimizeStatement( - unparsedPredicateOptimize: UnparsedPredicateOptimize, - parseExpression: String => Expression): LogicalPlan = { - - val UnparsedPredicateOptimize(tableIdent, tablePredicate, orderExpr) = - unparsedPredicateOptimize - - val predicate = tablePredicate.map(parseExpression) - verifyPartitionPredicates(predicate) - val table = UnresolvedRelation(tableIdent) - val tableWithFilter = predicate match { - case Some(expr) => Filter(expr, table) - case None => table - } - val query = - Sort( - SortOrder(orderExpr, Ascending, NullsLast, Seq.empty) :: Nil, - conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED), - Project(Seq(UnresolvedStar(None)), tableWithFilter)) - OptimizeZorderStatement(tableIdent, query) - } - - private def verifyPartitionPredicates(predicates: Option[Expression]): Unit = { - predicates.foreach { - case p if !isLikelySelective(p) => - throw new KyuubiSQLExtensionException(s"unsupported partition predicates: ${p.sql}") - case _ => - } - } - - /** - * Forked from Apache Spark's org.apache.spark.sql.catalyst.expressions.PredicateHelper - * The `PredicateHelper.isLikelySelective()` is available since Spark 3.3, forked for Spark - * that is lower than 3.3. - * - * Returns whether an expression is likely to be selective - */ - private def isLikelySelective(e: Expression): Boolean = e match { - case Not(expr) => isLikelySelective(expr) - case And(l, r) => isLikelySelective(l) || isLikelySelective(r) - case Or(l, r) => isLikelySelective(l) && isLikelySelective(r) - case _: StringRegexExpression => true - case _: BinaryComparison => true - case _: In | _: InSet => true - case _: StringPredicate => true - case BinaryPredicate(_) => true - case _: MultiLikeBase => true - case _ => false - } - - private object BinaryPredicate { - def unapply(expr: Expression): Option[Expression] = expr match { - case _: Contains => Option(expr) - case _: StartsWith => Option(expr) - case _: EndsWith => Option(expr) - case _ => None - } - } - - /** - * Create an expression from the given context. This method just passes the context on to the - * visitor and only takes care of typing (We assume that the visitor returns an Expression here). - */ - protected def expression(ctx: ParserRuleContext): Expression = typedVisit(ctx) - - protected def multiPart(ctx: ParserRuleContext): Seq[String] = typedVisit(ctx) - - override def visitSingleStatement(ctx: SingleStatementContext): LogicalPlan = { - visit(ctx.statement()).asInstanceOf[LogicalPlan] - } - - override def visitOptimizeZorder( - ctx: OptimizeZorderContext): UnparsedPredicateOptimize = withOrigin(ctx) { - val tableIdent = multiPart(ctx.multipartIdentifier()) - - val predicate = Option(ctx.whereClause()) - .map(_.partitionPredicate) - .map(extractRawText(_)) - - val zorderCols = ctx.zorderClause().order.asScala - .map(visitMultipartIdentifier) - .map(UnresolvedAttribute(_)) - .toSeq - - val orderExpr = - if (zorderCols.length == 1) { - zorderCols.head - } else { - Zorder(zorderCols) - } - UnparsedPredicateOptimize(tableIdent, predicate, orderExpr) - } - - override def visitPassThrough(ctx: PassThroughContext): LogicalPlan = null - - override def visitMultipartIdentifier(ctx: MultipartIdentifierContext): Seq[String] = - withOrigin(ctx) { - ctx.parts.asScala.map(typedVisit[String]).toSeq - } - - override def visitIdentifier(ctx: IdentifierContext): String = { - withOrigin(ctx) { - ctx.strictIdentifier() match { - case quotedContext: QuotedIdentifierAlternativeContext => - typedVisit[String](quotedContext) - case _ => ctx.getText - } - } - } - - override def visitQuotedIdentifier(ctx: QuotedIdentifierContext): String = { - withOrigin(ctx) { - ctx.BACKQUOTED_IDENTIFIER().getText.stripPrefix("`").stripSuffix("`").replace("``", "`") - } - } - - override def visitZorderClause(ctx: ZorderClauseContext): Seq[UnresolvedAttribute] = - withOrigin(ctx) { - val res = ListBuffer[UnresolvedAttribute]() - ctx.multipartIdentifier().forEach { identifier => - res += UnresolvedAttribute(identifier.parts.asScala.map(typedVisit[String]).toSeq) - } - res.toSeq - } - - private def typedVisit[T](ctx: ParseTree): T = { - ctx.accept(this).asInstanceOf[T] - } - - private def extractRawText(exprContext: ParserRuleContext): String = { - // Extract the raw expression which will be parsed later - exprContext.getStart.getInputStream.getText(new Interval( - exprContext.getStart.getStartIndex, - exprContext.getStop.getStopIndex)) - } -} - -/** - * a logical plan contains an unparsed expression that will be parsed by spark. - */ -trait UnparsedExpressionLogicalPlan extends LogicalPlan { - override def output: Seq[Attribute] = throw new UnsupportedOperationException() - - override def children: Seq[LogicalPlan] = throw new UnsupportedOperationException() - - protected def withNewChildrenInternal( - newChildren: IndexedSeq[LogicalPlan]): LogicalPlan = - throw new UnsupportedOperationException() -} - -case class UnparsedPredicateOptimize( - tableIdent: Seq[String], - tablePredicate: Option[String], - orderExpr: Expression) extends UnparsedExpressionLogicalPlan {} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala deleted file mode 100644 index 3dda669a8a3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala +++ /dev/null @@ -1,50 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSessionExtensions - -import org.apache.kyuubi.sql.zorder.{InsertZorderBeforeWritingDatasource, InsertZorderBeforeWritingHive, ResolveZorder} - -class KyuubiSparkSQLCommonExtension extends (SparkSessionExtensions => Unit) { - override def apply(extensions: SparkSessionExtensions): Unit = { - KyuubiSparkSQLCommonExtension.injectCommonExtensions(extensions) - } -} - -object KyuubiSparkSQLCommonExtension { - def injectCommonExtensions(extensions: SparkSessionExtensions): Unit = { - // inject zorder parser and related rules - extensions.injectParser { case (_, parser) => new SparkKyuubiSparkSQLParser(parser) } - extensions.injectResolutionRule(ResolveZorder) - - // Note that: - // InsertZorderBeforeWritingDatasource and InsertZorderBeforeWritingHive - // should be applied before - // RepartitionBeforeWriting and RebalanceBeforeWriting - // because we can only apply one of them (i.e. Global Sort or Repartition/Rebalance) - extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingDatasource) - extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingHive) - extensions.injectPostHocResolutionRule(FinalStageConfigIsolationCleanRule) - - extensions.injectQueryStagePrepRule(_ => InsertShuffleNodeBeforeJoin) - - extensions.injectPostHocResolutionRule(MarkNumOutputColumnsRule(_)) - extensions.injectQueryStagePrepRule(FinalStageConfigIsolation(_)) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala deleted file mode 100644 index 33ff3e3177a..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.{FinalStageResourceManager, InjectCustomResourceProfile, SparkSessionExtensions} - -import org.apache.kyuubi.sql.watchdog.{KyuubiUnsupportedOperationsCheck, MaxScanStrategy} - -// scalastyle:off line.size.limit -/** - * Depend on Spark SQL Extension framework, we can use this extension follow steps - * 1. move this jar into $SPARK_HOME/jars - * 2. add config into `spark-defaults.conf`: `spark.sql.extensions=org.apache.kyuubi.sql.KyuubiSparkSQLExtension` - */ -// scalastyle:on line.size.limit -class KyuubiSparkSQLExtension extends (SparkSessionExtensions => Unit) { - override def apply(extensions: SparkSessionExtensions): Unit = { - KyuubiSparkSQLCommonExtension.injectCommonExtensions(extensions) - - extensions.injectPostHocResolutionRule(RebalanceBeforeWritingDatasource) - extensions.injectPostHocResolutionRule(RebalanceBeforeWritingHive) - extensions.injectPostHocResolutionRule(DropIgnoreNonexistent) - - // watchdog extension - extensions.injectCheckRule(_ => KyuubiUnsupportedOperationsCheck) - extensions.injectPlannerStrategy(MaxScanStrategy) - - extensions.injectQueryStagePrepRule(FinalStageResourceManager(_)) - extensions.injectQueryStagePrepRule(InjectCustomResourceProfile) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala deleted file mode 100644 index c4418c33c44..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala +++ /dev/null @@ -1,140 +0,0 @@ -/* - * 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.sql - -import org.antlr.v4.runtime._ -import org.antlr.v4.runtime.atn.PredictionMode -import org.antlr.v4.runtime.misc.{Interval, ParseCancellationException} -import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.{FunctionIdentifier, SQLConfHelper, TableIdentifier} -import org.apache.spark.sql.catalyst.expressions.Expression -import org.apache.spark.sql.catalyst.parser.{ParseErrorListener, ParseException, ParserInterface, PostProcessor} -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.catalyst.trees.Origin -import org.apache.spark.sql.types.{DataType, StructType} - -abstract class KyuubiSparkSQLParserBase extends ParserInterface with SQLConfHelper { - def delegate: ParserInterface - def astBuilder: KyuubiSparkSQLAstBuilder - - override def parsePlan(sqlText: String): LogicalPlan = parse(sqlText) { parser => - astBuilder.visit(parser.singleStatement()) match { - case optimize: UnparsedPredicateOptimize => - astBuilder.buildOptimizeStatement(optimize, delegate.parseExpression) - case plan: LogicalPlan => plan - case _ => delegate.parsePlan(sqlText) - } - } - - protected def parse[T](command: String)(toResult: KyuubiSparkSQLParser => T): T = { - val lexer = new KyuubiSparkSQLLexer( - new UpperCaseCharStream(CharStreams.fromString(command))) - lexer.removeErrorListeners() - lexer.addErrorListener(ParseErrorListener) - - val tokenStream = new CommonTokenStream(lexer) - val parser = new KyuubiSparkSQLParser(tokenStream) - parser.addParseListener(PostProcessor) - parser.removeErrorListeners() - parser.addErrorListener(ParseErrorListener) - - try { - try { - // first, try parsing with potentially faster SLL mode - parser.getInterpreter.setPredictionMode(PredictionMode.SLL) - toResult(parser) - } catch { - case _: ParseCancellationException => - // if we fail, parse with LL mode - tokenStream.seek(0) // rewind input stream - parser.reset() - - // Try Again. - parser.getInterpreter.setPredictionMode(PredictionMode.LL) - toResult(parser) - } - } catch { - case e: ParseException if e.command.isDefined => - throw e - case e: ParseException => - throw e.withCommand(command) - case e: AnalysisException => - val position = Origin(e.line, e.startPosition) - throw new ParseException(Option(command), e.message, position, position) - } - } - - override def parseExpression(sqlText: String): Expression = { - delegate.parseExpression(sqlText) - } - - override def parseTableIdentifier(sqlText: String): TableIdentifier = { - delegate.parseTableIdentifier(sqlText) - } - - override def parseFunctionIdentifier(sqlText: String): FunctionIdentifier = { - delegate.parseFunctionIdentifier(sqlText) - } - - override def parseMultipartIdentifier(sqlText: String): Seq[String] = { - delegate.parseMultipartIdentifier(sqlText) - } - - override def parseTableSchema(sqlText: String): StructType = { - delegate.parseTableSchema(sqlText) - } - - override def parseDataType(sqlText: String): DataType = { - delegate.parseDataType(sqlText) - } - - /** - * This functions was introduced since spark-3.3, for more details, please see - * https://github.com/apache/spark/pull/34543 - */ - override def parseQuery(sqlText: String): LogicalPlan = { - delegate.parseQuery(sqlText) - } -} - -class SparkKyuubiSparkSQLParser( - override val delegate: ParserInterface) - extends KyuubiSparkSQLParserBase { - def astBuilder: KyuubiSparkSQLAstBuilder = new KyuubiSparkSQLAstBuilder -} - -/* Copied from Apache Spark's to avoid dependency on Spark Internals */ -class UpperCaseCharStream(wrapped: CodePointCharStream) extends CharStream { - override def consume(): Unit = wrapped.consume() - override def getSourceName(): String = wrapped.getSourceName - override def index(): Int = wrapped.index - override def mark(): Int = wrapped.mark - override def release(marker: Int): Unit = wrapped.release(marker) - override def seek(where: Int): Unit = wrapped.seek(where) - override def size(): Int = wrapped.size - - override def getText(interval: Interval): String = wrapped.getText(interval) - - // scalastyle:off - override def LA(i: Int): Int = { - val la = wrapped.LA(i) - if (la == 0 || la == IntStream.EOF) la - else Character.toUpperCase(la) - } - // scalastyle:on -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/MarkNumOutputColumnsRule.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/MarkNumOutputColumnsRule.scala deleted file mode 100644 index 5312eefb887..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/MarkNumOutputColumnsRule.scala +++ /dev/null @@ -1,68 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Union, V2WriteCommand} -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.command.DataWritingCommand - -case class MarkNumOutputColumnsRule(session: SparkSession) - extends Rule[LogicalPlan] { - import MarkNumOutputColumnsRule._ - - override def apply(plan: LogicalPlan): LogicalPlan = { - if (!conf.getConf(KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION)) { - return plan - } - if (session.conf.getOption(OUTPUT_NUM_COLUMNS).isDefined) { - session.conf.unset(OUTPUT_NUM_COLUMNS) - } - - def numOutputColumns(p: LogicalPlan): Option[Int] = p match { - case w: DataWritingCommand => Some(w.outputColumnNames.size) - case w: V2WriteCommand => Some(w.query.output.size) - case u: Union if u.children.nonEmpty => numOutputColumns(u.children.head) - case _ => None - } - val numCols = numOutputColumns(plan) - numCols.foreach { n => - session.conf.set(OUTPUT_NUM_COLUMNS, s"$n") - } - plan - } -} - -object MarkNumOutputColumnsRule { - val OUTPUT_NUM_COLUMNS = "spark.sql.internal.numOutputColumns" - - def numOutputColumns(session: SparkSession): Option[String] = { - session.conf.getOption(OUTPUT_NUM_COLUMNS) - } - - def isWrite(session: SparkSession, plan: SparkPlan): Boolean = { - numOutputColumns(session).exists { num => - try { - num.toInt == plan.output.size - } catch { - case _: Throwable => false - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala deleted file mode 100644 index 2c6980fdb6e..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala +++ /dev/null @@ -1,68 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, SortOrder} -import org.apache.spark.sql.catalyst.plans.logical._ - -trait RepartitionBuilderWithRebalance extends RepartitionBuilder { - override def buildRepartition( - dynamicPartitionColumns: Seq[Attribute], - query: LogicalPlan): LogicalPlan = { - if (!conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS) || - dynamicPartitionColumns.nonEmpty) { - RebalancePartitions(dynamicPartitionColumns, query) - } else { - val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS) - val inferred = InferRebalanceAndSortOrders.infer(query) - if (inferred.isDefined) { - val (partitioning, ordering) = inferred.get - val rebalance = RebalancePartitions(partitioning.take(maxColumns), query) - if (ordering.nonEmpty) { - val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending)) - Sort(sortOrders, false, rebalance) - } else { - rebalance - } - } else { - RebalancePartitions(dynamicPartitionColumns, query) - } - } - } -} - -/** - * For datasource table, there two commands can write data to table - * 1. InsertIntoHadoopFsRelationCommand - * 2. CreateDataSourceTableAsSelectCommand - * This rule add a RebalancePartitions node between write and query - */ -case class RebalanceBeforeWritingDatasource(session: SparkSession) - extends RepartitionBeforeWritingDatasourceBase - with RepartitionBuilderWithRebalance {} - -/** - * For Hive table, there two commands can write data to table - * 1. InsertIntoHiveTable - * 2. CreateHiveTableAsSelectCommand - * This rule add a RebalancePartitions node between write and query - */ -case class RebalanceBeforeWritingHive(session: SparkSession) - extends RepartitionBeforeWritingHiveBase - with RepartitionBuilderWithRebalance {} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala deleted file mode 100644 index 0e2e901265e..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala +++ /dev/null @@ -1,153 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.command.{CreateDataSourceTableAsSelectCommand, InsertIntoDataSourceDirCommand} -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.hive.execution.{CreateHiveTableAsSelectCommand, InsertIntoHiveDirCommand, InsertIntoHiveTable, OptimizedCreateHiveTableAsSelectCommand} -import org.apache.spark.sql.internal.StaticSQLConf - -trait RepartitionBuilder extends Rule[LogicalPlan] with RepartitionBeforeWriteHelper { - def buildRepartition( - dynamicPartitionColumns: Seq[Attribute], - query: LogicalPlan): LogicalPlan -} - -/** - * For datasource table, there two commands can write data to table - * 1. InsertIntoHadoopFsRelationCommand - * 2. CreateDataSourceTableAsSelectCommand - * This rule add a repartition node between write and query - */ -abstract class RepartitionBeforeWritingDatasourceBase extends RepartitionBuilder { - - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE)) { - addRepartition(plan) - } else { - plan - } - } - - private def addRepartition(plan: LogicalPlan): LogicalPlan = plan match { - case i @ InsertIntoHadoopFsRelationCommand(_, sp, _, pc, bucket, _, _, query, _, _, _, _) - if query.resolved && bucket.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = pc.filterNot(attr => sp.contains(attr.name)) - i.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case c @ CreateDataSourceTableAsSelectCommand(table, _, query, _) - if query.resolved && table.bucketSpec.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = - query.output.filter(attr => table.partitionColumnNames.contains(attr.name)) - c.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case i @ InsertIntoDataSourceDirCommand(_, _, query, _) - if query.resolved && canInsertRepartitionByExpression(query) => - i.copy(query = buildRepartition(Seq.empty, query)) - - case u @ Union(children, _, _) => - u.copy(children = children.map(addRepartition)) - - case _ => plan - } -} - -/** - * For Hive table, there two commands can write data to table - * 1. InsertIntoHiveTable - * 2. CreateHiveTableAsSelectCommand - * This rule add a repartition node between write and query - */ -abstract class RepartitionBeforeWritingHiveBase extends RepartitionBuilder { - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(StaticSQLConf.CATALOG_IMPLEMENTATION) == "hive" && - conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE)) { - addRepartition(plan) - } else { - plan - } - } - - def addRepartition(plan: LogicalPlan): LogicalPlan = plan match { - case i @ InsertIntoHiveTable(table, partition, query, _, _, _) - if query.resolved && table.bucketSpec.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = partition.filter(_._2.isEmpty).keys - .flatMap(name => query.output.find(_.name == name)).toSeq - i.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case c @ CreateHiveTableAsSelectCommand(table, query, _, _) - if query.resolved && table.bucketSpec.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = - query.output.filter(attr => table.partitionColumnNames.contains(attr.name)) - c.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case c @ OptimizedCreateHiveTableAsSelectCommand(table, query, _, _) - if query.resolved && table.bucketSpec.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = - query.output.filter(attr => table.partitionColumnNames.contains(attr.name)) - c.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case c @ InsertIntoHiveDirCommand(_, _, query, _, _) - if query.resolved && canInsertRepartitionByExpression(query) => - c.copy(query = buildRepartition(Seq.empty, query)) - - case u @ Union(children, _, _) => - u.copy(children = children.map(addRepartition)) - - case _ => plan - } -} - -trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] { - private def hasBenefit(plan: LogicalPlan): Boolean = { - def probablyHasShuffle: Boolean = plan.find { - case _: Join => true - case _: Aggregate => true - case _: Distinct => true - case _: Deduplicate => true - case _: Window => true - case s: Sort if s.global => true - case _: RepartitionOperation => true - case _: RebalancePartitions => true - case _: GlobalLimit => true - case _ => false - }.isDefined - - conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE) || probablyHasShuffle - } - - def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = { - def canInsert(p: LogicalPlan): Boolean = p match { - case Project(_, child) => canInsert(child) - case SubqueryAlias(_, child) => canInsert(child) - case Limit(_, _) => false - case _: Sort => false - case _: RepartitionOperation => false - case _: RebalancePartitions => false - case _ => true - } - - // 1. make sure AQE is enabled, otherwise it is no meaning to add a shuffle - // 2. make sure it does not break the semantics of original plan - // 3. try to avoid adding a shuffle if it has potential performance regression - conf.adaptiveExecutionEnabled && canInsert(plan) && hasBenefit(plan) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala deleted file mode 100644 index 2b4d3940ada..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala +++ /dev/null @@ -1,35 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ScriptTransformation} - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} - -object KyuubiUnsupportedOperationsCheck extends (LogicalPlan => Unit) with SQLConfHelper { - override def apply(plan: LogicalPlan): Unit = - conf.getConf(KyuubiSQLConf.SCRIPT_TRANSFORMATION_ENABLED) match { - case false => plan foreach { - case _: ScriptTransformation => - throw new KyuubiSQLExtensionException("Script transformation is not allowed") - case _ => - } - case true => - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala deleted file mode 100644 index e44309192a9..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala +++ /dev/null @@ -1,30 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -final class MaxPartitionExceedException( - private val reason: String = "", - private val cause: Throwable = None.orNull) - extends KyuubiSQLExtensionException(reason, cause) - -final class MaxFileSizeExceedException( - private val reason: String = "", - private val cause: Throwable = None.orNull) - extends KyuubiSQLExtensionException(reason, cause) diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala deleted file mode 100644 index a85b34035ee..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala +++ /dev/null @@ -1,337 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.hadoop.fs.Path -import org.apache.spark.sql.{PruneFileSourcePartitionHelper, SparkSession, Strategy} -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} -import org.apache.spark.sql.catalyst.planning.ScanOperation -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.datasources.{CatalogFileIndex, HadoopFsRelation, InMemoryFileIndex, LogicalRelation} -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation -import org.apache.spark.sql.types.StructType - -import org.apache.kyuubi.sql.KyuubiSQLConf - -/** - * Add MaxScanStrategy to avoid scan excessive partitions or files - * 1. Check if scan exceed maxPartition of partitioned table - * 2. Check if scan exceed maxFileSize (calculated by hive table and partition statistics) - * This Strategy Add Planner Strategy after LogicalOptimizer - * @param session - */ -case class MaxScanStrategy(session: SparkSession) - extends Strategy - with SQLConfHelper - with PruneFileSourcePartitionHelper { - override def apply(plan: LogicalPlan): Seq[SparkPlan] = { - val maxScanPartitionsOpt = conf.getConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS) - val maxFileSizeOpt = conf.getConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE) - if (maxScanPartitionsOpt.isDefined || maxFileSizeOpt.isDefined) { - checkScan(plan, maxScanPartitionsOpt, maxFileSizeOpt) - } - Nil - } - - private def checkScan( - plan: LogicalPlan, - maxScanPartitionsOpt: Option[Int], - maxFileSizeOpt: Option[Long]): Unit = { - plan match { - case ScanOperation(_, _, relation: HiveTableRelation) => - if (relation.isPartitioned) { - relation.prunedPartitions match { - case Some(prunedPartitions) => - if (maxScanPartitionsOpt.exists(_ < prunedPartitions.size)) { - throw new MaxPartitionExceedException( - s""" - |SQL job scan hive partition: ${prunedPartitions.size} - |exceed restrict of hive scan maxPartition ${maxScanPartitionsOpt.get} - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - lazy val scanFileSize = prunedPartitions.flatMap(_.stats).map(_.sizeInBytes).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - Some(relation.tableMeta), - prunedPartitions.flatMap(_.storage.locationUri).map(_.toString), - relation.partitionCols.map(_.name)) - } - case _ => - lazy val scanPartitions: Int = session - .sessionState.catalog.externalCatalog.listPartitionNames( - relation.tableMeta.database, - relation.tableMeta.identifier.table).size - if (maxScanPartitionsOpt.exists(_ < scanPartitions)) { - throw new MaxPartitionExceedException( - s""" - |Your SQL job scan a whole huge table without any partition filter, - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - - lazy val scanFileSize: BigInt = - relation.tableMeta.stats.map(_.sizeInBytes).getOrElse { - session - .sessionState.catalog.externalCatalog.listPartitions( - relation.tableMeta.database, - relation.tableMeta.identifier.table).flatMap(_.stats).map(_.sizeInBytes).sum - } - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |Your SQL job scan a whole huge table without any partition filter, - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - } - } else { - lazy val scanFileSize = relation.tableMeta.stats.map(_.sizeInBytes).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - Some(relation.tableMeta)) - } - } - case ScanOperation( - _, - filters, - relation @ LogicalRelation( - fsRelation @ HadoopFsRelation( - fileIndex: InMemoryFileIndex, - partitionSchema, - _, - _, - _, - _), - _, - _, - _)) => - if (fsRelation.partitionSchema.nonEmpty) { - val (partitionKeyFilters, dataFilter) = - getPartitionKeyFiltersAndDataFilters( - fsRelation.sparkSession, - relation, - partitionSchema, - filters, - relation.output) - val prunedPartitions = fileIndex.listFiles( - partitionKeyFilters.toSeq, - dataFilter) - if (maxScanPartitionsOpt.exists(_ < prunedPartitions.size)) { - throw maxPartitionExceedError( - prunedPartitions.size, - maxScanPartitionsOpt.get, - relation.catalogTable, - fileIndex.rootPaths, - fsRelation.partitionSchema) - } - lazy val scanFileSize = prunedPartitions.flatMap(_.files).map(_.getLen).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - relation.catalogTable, - fileIndex.rootPaths.map(_.toString), - fsRelation.partitionSchema.map(_.name)) - } - } else { - lazy val scanFileSize = fileIndex.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - relation.catalogTable) - } - } - case ScanOperation( - _, - filters, - logicalRelation @ LogicalRelation( - fsRelation @ HadoopFsRelation( - catalogFileIndex: CatalogFileIndex, - partitionSchema, - _, - _, - _, - _), - _, - _, - _)) => - if (fsRelation.partitionSchema.nonEmpty) { - val (partitionKeyFilters, _) = - getPartitionKeyFiltersAndDataFilters( - fsRelation.sparkSession, - logicalRelation, - partitionSchema, - filters, - logicalRelation.output) - - val fileIndex = catalogFileIndex.filterPartitions( - partitionKeyFilters.toSeq) - - lazy val prunedPartitionSize = fileIndex.partitionSpec().partitions.size - if (maxScanPartitionsOpt.exists(_ < prunedPartitionSize)) { - throw maxPartitionExceedError( - prunedPartitionSize, - maxScanPartitionsOpt.get, - logicalRelation.catalogTable, - catalogFileIndex.rootPaths, - fsRelation.partitionSchema) - } - - lazy val scanFileSize = fileIndex - .listFiles(Nil, Nil).flatMap(_.files).map(_.getLen).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - logicalRelation.catalogTable, - catalogFileIndex.rootPaths.map(_.toString), - fsRelation.partitionSchema.map(_.name)) - } - } else { - lazy val scanFileSize = catalogFileIndex.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - logicalRelation.catalogTable) - } - } - case ScanOperation( - _, - _, - relation @ DataSourceV2ScanRelation(_, _, _, _)) => - val table = relation.relation.table - if (table.partitioning().nonEmpty) { - val partitionColumnNames = table.partitioning().map(_.describe()) - val stats = relation.computeStats() - lazy val scanFileSize = stats.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize ${maxFileSizeOpt.get} - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${table.name()} - |Partition Structure: ${partitionColumnNames.mkString(",")} - |""".stripMargin) - } - } else { - val stats = relation.computeStats() - lazy val scanFileSize = stats.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize ${maxFileSizeOpt.get} - |detail as below: - |Table: ${table.name()} - |""".stripMargin) - } - } - case _ => - } - } - - def maxPartitionExceedError( - prunedPartitionSize: Int, - maxPartitionSize: Int, - tableMeta: Option[CatalogTable], - rootPaths: Seq[Path], - partitionSchema: StructType): Throwable = { - val truncatedPaths = - if (rootPaths.length > 5) { - rootPaths.slice(0, 5).mkString(",") + """... """ + (rootPaths.length - 5) + " more paths" - } else { - rootPaths.mkString(",") - } - - new MaxPartitionExceedException( - s""" - |SQL job scan data source partition: $prunedPartitionSize - |exceed restrict of data source scan maxPartition $maxPartitionSize - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |RootPaths: $truncatedPaths - |Partition Structure: ${partitionSchema.map(_.name).mkString(", ")} - |""".stripMargin) - } - - private def partTableMaxFileExceedError( - scanFileSize: Number, - maxFileSize: Long, - tableMeta: Option[CatalogTable], - rootPaths: Seq[String], - partitions: Seq[String]): Throwable = { - val truncatedPaths = - if (rootPaths.length > 5) { - rootPaths.slice(0, 5).mkString(",") + """... """ + (rootPaths.length - 5) + " more paths" - } else { - rootPaths.mkString(",") - } - - new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize $maxFileSize - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |RootPaths: $truncatedPaths - |Partition Structure: ${partitions.mkString(", ")} - |""".stripMargin) - } - - private def nonPartTableMaxFileExceedError( - scanFileSize: Number, - maxFileSize: Long, - tableMeta: Option[CatalogTable]): Throwable = { - new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize $maxFileSize - |detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |Location: ${tableMeta.map(_.location).getOrElse("")} - |""".stripMargin) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala deleted file mode 100644 index e03e49b7d87..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala +++ /dev/null @@ -1,220 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.catalog.CatalogTable -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Expression, NullsLast, SortOrder} -import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.command.CreateDataSourceTableAsSelectCommand -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.hive.execution.{CreateHiveTableAsSelectCommand, InsertIntoHiveTable, OptimizedCreateHiveTableAsSelectCommand} - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} - -trait ZorderBuilder { - def buildZorder(children: Seq[Expression]): ZorderBase -} - -trait InsertZorderHelper extends Rule[LogicalPlan] with ZorderBuilder { - private val KYUUBI_ZORDER_ENABLED = "kyuubi.zorder.enabled" - private val KYUUBI_ZORDER_COLS = "kyuubi.zorder.cols" - - def isZorderEnabled(props: Map[String, String]): Boolean = { - props.contains(KYUUBI_ZORDER_ENABLED) && - "true".equalsIgnoreCase(props(KYUUBI_ZORDER_ENABLED)) && - props.contains(KYUUBI_ZORDER_COLS) - } - - def getZorderColumns(props: Map[String, String]): Seq[String] = { - val cols = props.get(KYUUBI_ZORDER_COLS) - assert(cols.isDefined) - cols.get.split(",").map(_.trim) - } - - def canInsertZorder(query: LogicalPlan): Boolean = query match { - case Project(_, child) => canInsertZorder(child) - case _: RepartitionOperation | _: RebalancePartitions - if !conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) => true - // TODO: actually, we can force zorder even if existed some shuffle - case _: Sort => false - case _: RepartitionOperation => false - case _: RebalancePartitions => false - case _ => true - } - - def insertZorder( - catalogTable: CatalogTable, - plan: LogicalPlan, - dynamicPartitionColumns: Seq[Attribute]): LogicalPlan = { - if (!canInsertZorder(plan)) { - return plan - } - val cols = getZorderColumns(catalogTable.properties) - val resolver = session.sessionState.conf.resolver - val output = plan.output - val bound = cols.flatMap(col => output.find(attr => resolver(attr.name, col))) - if (bound.size < cols.size) { - logWarning(s"target table does not contain all zorder cols: ${cols.mkString(",")}, " + - s"please check your table properties ${KYUUBI_ZORDER_COLS}.") - plan - } else { - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) && - conf.getConf(KyuubiSQLConf.REBALANCE_BEFORE_ZORDER)) { - throw new KyuubiSQLExtensionException(s"${KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key} " + - s"and ${KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key} can not be enabled together.") - } - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) && - dynamicPartitionColumns.nonEmpty) { - logWarning(s"Dynamic partition insertion with global sort may produce small files.") - } - - val zorderExpr = - if (bound.length == 1) { - bound - } else if (conf.getConf(KyuubiSQLConf.ZORDER_USING_ORIGINAL_ORDERING_ENABLED)) { - bound.asInstanceOf[Seq[Expression]] - } else { - buildZorder(bound) :: Nil - } - val (global, orderExprs, child) = - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED)) { - (true, zorderExpr, plan) - } else if (conf.getConf(KyuubiSQLConf.REBALANCE_BEFORE_ZORDER)) { - val rebalanceExpr = - if (dynamicPartitionColumns.isEmpty) { - // static partition insert - bound - } else if (conf.getConf(KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED)) { - // improve data compression ratio - dynamicPartitionColumns.asInstanceOf[Seq[Expression]] ++ bound - } else { - dynamicPartitionColumns.asInstanceOf[Seq[Expression]] - } - // for dynamic partition insert, Spark always sort the partition columns, - // so here we sort partition columns + zorder. - val rebalance = - if (dynamicPartitionColumns.nonEmpty && - conf.getConf(KyuubiSQLConf.TWO_PHASE_REBALANCE_BEFORE_ZORDER)) { - // improve compression ratio - RebalancePartitions( - rebalanceExpr, - RebalancePartitions(dynamicPartitionColumns, plan)) - } else { - RebalancePartitions(rebalanceExpr, plan) - } - (false, dynamicPartitionColumns.asInstanceOf[Seq[Expression]] ++ zorderExpr, rebalance) - } else { - (false, zorderExpr, plan) - } - val order = orderExprs.map { expr => - SortOrder(expr, Ascending, NullsLast, Seq.empty) - } - Sort(order, global, child) - } - } - - override def buildZorder(children: Seq[Expression]): ZorderBase = Zorder(children) - - def session: SparkSession - def applyInternal(plan: LogicalPlan): LogicalPlan - - final override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING)) { - applyInternal(plan) - } else { - plan - } - } -} - -case class InsertZorderBeforeWritingDatasource(session: SparkSession) - extends InsertZorderHelper { - override def applyInternal(plan: LogicalPlan): LogicalPlan = plan match { - case insert: InsertIntoHadoopFsRelationCommand - if insert.query.resolved && - insert.bucketSpec.isEmpty && insert.catalogTable.isDefined && - isZorderEnabled(insert.catalogTable.get.properties) => - val dynamicPartition = - insert.partitionColumns.filterNot(attr => insert.staticPartitions.contains(attr.name)) - val newQuery = insertZorder(insert.catalogTable.get, insert.query, dynamicPartition) - if (newQuery.eq(insert.query)) { - insert - } else { - insert.copy(query = newQuery) - } - - case ctas: CreateDataSourceTableAsSelectCommand - if ctas.query.resolved && - ctas.table.bucketSpec.isEmpty && isZorderEnabled(ctas.table.properties) => - val dynamicPartition = - ctas.query.output.filter(attr => ctas.table.partitionColumnNames.contains(attr.name)) - val newQuery = insertZorder(ctas.table, ctas.query, dynamicPartition) - if (newQuery.eq(ctas.query)) { - ctas - } else { - ctas.copy(query = newQuery) - } - - case _ => plan - } -} - -case class InsertZorderBeforeWritingHive(session: SparkSession) - extends InsertZorderHelper { - override def applyInternal(plan: LogicalPlan): LogicalPlan = plan match { - case insert: InsertIntoHiveTable - if insert.query.resolved && - insert.table.bucketSpec.isEmpty && isZorderEnabled(insert.table.properties) => - val dynamicPartition = insert.partition.filter(_._2.isEmpty).keys - .flatMap(name => insert.query.output.find(_.name == name)).toSeq - val newQuery = insertZorder(insert.table, insert.query, dynamicPartition) - if (newQuery.eq(insert.query)) { - insert - } else { - insert.copy(query = newQuery) - } - - case ctas: CreateHiveTableAsSelectCommand - if ctas.query.resolved && - ctas.tableDesc.bucketSpec.isEmpty && isZorderEnabled(ctas.tableDesc.properties) => - val dynamicPartition = - ctas.query.output.filter(attr => ctas.tableDesc.partitionColumnNames.contains(attr.name)) - val newQuery = insertZorder(ctas.tableDesc, ctas.query, dynamicPartition) - if (newQuery.eq(ctas.query)) { - ctas - } else { - ctas.copy(query = newQuery) - } - - case octas: OptimizedCreateHiveTableAsSelectCommand - if octas.query.resolved && - octas.tableDesc.bucketSpec.isEmpty && isZorderEnabled(octas.tableDesc.properties) => - val dynamicPartition = - octas.query.output.filter(attr => octas.tableDesc.partitionColumnNames.contains(attr.name)) - val newQuery = insertZorder(octas.tableDesc, octas.query, dynamicPartition) - if (newQuery.eq(octas.query)) { - octas - } else { - octas.copy(query = newQuery) - } - - case _ => plan - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala deleted file mode 100644 index 21d1cf2a25b..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala +++ /dev/null @@ -1,78 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.{Row, SparkSession} -import org.apache.spark.sql.catalyst.catalog.CatalogTable -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.command.DataWritingCommand -import org.apache.spark.sql.hive.execution.InsertIntoHiveTable - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -/** - * A runnable command for zorder, we delegate to real command to execute - */ -abstract class OptimizeZorderCommandBase extends DataWritingCommand { - def catalogTable: CatalogTable - - override def outputColumnNames: Seq[String] = query.output.map(_.name) - - private def isHiveTable: Boolean = { - catalogTable.provider.isEmpty || - (catalogTable.provider.isDefined && "hive".equalsIgnoreCase(catalogTable.provider.get)) - } - - private def getWritingCommand(session: SparkSession): DataWritingCommand = { - // TODO: Support convert hive relation to datasource relation, can see - // [[org.apache.spark.sql.hive.RelationConversions]] - InsertIntoHiveTable( - catalogTable, - catalogTable.partitionColumnNames.map(p => (p, None)).toMap, - query, - overwrite = true, - ifPartitionNotExists = false, - outputColumnNames) - } - - override def run(session: SparkSession, child: SparkPlan): Seq[Row] = { - // TODO: Support datasource relation - // TODO: Support read and insert overwrite the same table for some table format - if (!isHiveTable) { - throw new KyuubiSQLExtensionException("only support hive table") - } - - val command = getWritingCommand(session) - command.run(session, child) - DataWritingCommand.propogateMetrics(session.sparkContext, command, metrics) - Seq.empty - } -} - -/** - * A runnable command for zorder, we delegate to real command to execute - */ -case class OptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan) - extends OptimizeZorderCommandBase { - protected def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = { - copy(query = newChild) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala deleted file mode 100644 index 895f9e24be3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} - -/** - * A zorder statement that contains we parsed from SQL. - * We should convert this plan to certain command at Analyzer. - */ -case class OptimizeZorderStatement( - tableIdentifier: Seq[String], - query: LogicalPlan) extends UnaryNode { - override def child: LogicalPlan = query - override def output: Seq[Attribute] = child.output - protected def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = - copy(query = newChild) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala deleted file mode 100644 index 9f735caa7a7..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.TableIdentifier -import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} -import org.apache.spark.sql.catalyst.expressions.AttributeSet -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, SubqueryAlias} -import org.apache.spark.sql.catalyst.rules.Rule - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -/** - * Resolve `OptimizeZorderStatement` to `OptimizeZorderCommand` - */ -abstract class ResolveZorderBase extends Rule[LogicalPlan] { - def session: SparkSession - def buildOptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan): OptimizeZorderCommandBase - - protected def checkQueryAllowed(query: LogicalPlan): Unit = query foreach { - case Filter(condition, SubqueryAlias(_, tableRelation: HiveTableRelation)) => - if (tableRelation.partitionCols.isEmpty) { - throw new KyuubiSQLExtensionException("Filters are only supported for partitioned table") - } - - val partitionKeyIds = AttributeSet(tableRelation.partitionCols) - if (condition.references.isEmpty || !condition.references.subsetOf(partitionKeyIds)) { - throw new KyuubiSQLExtensionException("Only partition column filters are allowed") - } - - case _ => - } - - protected def getTableIdentifier(tableIdent: Seq[String]): TableIdentifier = tableIdent match { - case Seq(tbl) => TableIdentifier.apply(tbl) - case Seq(db, tbl) => TableIdentifier.apply(tbl, Some(db)) - case _ => throw new KyuubiSQLExtensionException( - "only support session catalog table, please use db.table instead") - } - - override def apply(plan: LogicalPlan): LogicalPlan = plan match { - case statement: OptimizeZorderStatement if statement.query.resolved => - checkQueryAllowed(statement.query) - val tableIdentifier = getTableIdentifier(statement.tableIdentifier) - val catalogTable = session.sessionState.catalog.getTableMetadata(tableIdentifier) - buildOptimizeZorderCommand(catalogTable, statement.query) - - case _ => plan - } -} - -/** - * Resolve `OptimizeZorderStatement` to `OptimizeZorderCommand` - */ -case class ResolveZorder(session: SparkSession) extends ResolveZorderBase { - override def buildOptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan): OptimizeZorderCommandBase = { - OptimizeZorderCommand(catalogTable, query) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala deleted file mode 100644 index e4d98ccbe84..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala +++ /dev/null @@ -1,95 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.analysis.TypeCheckResult -import org.apache.spark.sql.catalyst.expressions.Expression -import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, FalseLiteral} -import org.apache.spark.sql.catalyst.expressions.codegen.Block._ -import org.apache.spark.sql.types.{BinaryType, DataType} - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -abstract class ZorderBase extends Expression { - override def foldable: Boolean = children.forall(_.foldable) - override def nullable: Boolean = false - override def dataType: DataType = BinaryType - override def prettyName: String = "zorder" - - override def checkInputDataTypes(): TypeCheckResult = { - try { - defaultNullValues - TypeCheckResult.TypeCheckSuccess - } catch { - case e: KyuubiSQLExtensionException => - TypeCheckResult.TypeCheckFailure(e.getMessage) - } - } - - @transient - private[this] lazy val defaultNullValues: Array[Any] = - children.map(_.dataType) - .map(ZorderBytesUtils.defaultValue) - .toArray - - override def eval(input: InternalRow): Any = { - val childrenValues = children.zipWithIndex.map { - case (child: Expression, index) => - val v = child.eval(input) - if (v == null) { - defaultNullValues(index) - } else { - v - } - } - ZorderBytesUtils.interleaveBits(childrenValues.toArray) - } - - override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { - val evals = children.map(_.genCode(ctx)) - val defaultValues = ctx.addReferenceObj("defaultValues", defaultNullValues) - val values = ctx.freshName("values") - val util = ZorderBytesUtils.getClass.getName.stripSuffix("$") - val inputs = evals.zipWithIndex.map { - case (eval, index) => - s""" - |${eval.code} - |if (${eval.isNull}) { - | $values[$index] = $defaultValues[$index]; - |} else { - | $values[$index] = ${eval.value}; - |} - |""".stripMargin - } - ev.copy( - code = - code""" - |byte[] ${ev.value} = null; - |Object[] $values = new Object[${evals.length}]; - |${inputs.mkString("\n")} - |${ev.value} = $util.interleaveBits($values); - |""".stripMargin, - isNull = FalseLiteral) - } -} - -case class Zorder(children: Seq[Expression]) extends ZorderBase { - protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = - copy(children = newChildren) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala deleted file mode 100644 index d7034e6804d..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala +++ /dev/null @@ -1,517 +0,0 @@ -/* - * 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.sql.zorder - -import java.lang.{Double => jDouble, Float => jFloat} - -import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DateType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampType} -import org.apache.spark.unsafe.types.UTF8String - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -object ZorderBytesUtils { - final private val BIT_8_MASK = 1 << 7 - final private val BIT_16_MASK = 1 << 15 - final private val BIT_32_MASK = 1 << 31 - final private val BIT_64_MASK = 1L << 63 - - def interleaveBits(inputs: Array[Any]): Array[Byte] = { - inputs.length match { - // it's a more fast approach, use O(8 * 8) - // can see http://graphics.stanford.edu/~seander/bithacks.html#InterleaveTableObvious - case 1 => longToByte(toLong(inputs(0))) - case 2 => interleave2Longs(toLong(inputs(0)), toLong(inputs(1))) - case 3 => interleave3Longs(toLong(inputs(0)), toLong(inputs(1)), toLong(inputs(2))) - case 4 => - interleave4Longs(toLong(inputs(0)), toLong(inputs(1)), toLong(inputs(2)), toLong(inputs(3))) - case 5 => interleave5Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4))) - case 6 => interleave6Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5))) - case 7 => interleave7Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5)), - toLong(inputs(6))) - case 8 => interleave8Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5)), - toLong(inputs(6)), - toLong(inputs(7))) - - case _ => - // it's the default approach, use O(64 * n), n is the length of inputs - interleaveBitsDefault(inputs.map(toByteArray)) - } - } - - private def interleave2Longs(l1: Long, l2: Long): Array[Byte] = { - // output 8 * 16 bits - val result = new Array[Byte](16) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toShort - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toShort - - var z = 0 - var j = 0 - while (j < 8) { - val x_masked = tmp1 & (1 << j) - val y_masked = tmp2 & (1 << j) - z |= (x_masked << j) - z |= (y_masked << (j + 1)) - j = j + 1 - } - result((7 - i) * 2 + 1) = (z & 0xFF).toByte - result((7 - i) * 2) = ((z >> 8) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave3Longs(l1: Long, l2: Long, l3: Long): Array[Byte] = { - // output 8 * 24 bits - val result = new Array[Byte](24) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toInt - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toInt - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toInt - - var z = 0 - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - z |= (r1_mask << (2 * j)) | (r2_mask << (2 * j + 1)) | (r3_mask << (2 * j + 2)) - j = j + 1 - } - result((7 - i) * 3 + 2) = (z & 0xFF).toByte - result((7 - i) * 3 + 1) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 3) = ((z >> 16) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave4Longs(l1: Long, l2: Long, l3: Long, l4: Long): Array[Byte] = { - // output 8 * 32 bits - val result = new Array[Byte](32) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toInt - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toInt - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toInt - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toInt - - var z = 0 - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - z |= (r1_mask << (3 * j)) | (r2_mask << (3 * j + 1)) | (r3_mask << (3 * j + 2)) | - (r4_mask << (3 * j + 3)) - j = j + 1 - } - result((7 - i) * 4 + 3) = (z & 0xFF).toByte - result((7 - i) * 4 + 2) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 4 + 1) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 4) = ((z >> 24) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave5Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long): Array[Byte] = { - // output 8 * 40 bits - val result = new Array[Byte](40) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - z |= (r1_mask << (4 * j)) | (r2_mask << (4 * j + 1)) | (r3_mask << (4 * j + 2)) | - (r4_mask << (4 * j + 3)) | (r5_mask << (4 * j + 4)) - j = j + 1 - } - result((7 - i) * 5 + 4) = (z & 0xFF).toByte - result((7 - i) * 5 + 3) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 5 + 2) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 5 + 1) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 5) = ((z >> 32) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave6Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long): Array[Byte] = { - // output 8 * 48 bits - val result = new Array[Byte](48) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - z |= (r1_mask << (5 * j)) | (r2_mask << (5 * j + 1)) | (r3_mask << (5 * j + 2)) | - (r4_mask << (5 * j + 3)) | (r5_mask << (5 * j + 4)) | (r6_mask << (5 * j + 5)) - j = j + 1 - } - result((7 - i) * 6 + 5) = (z & 0xFF).toByte - result((7 - i) * 6 + 4) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 6 + 3) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 6 + 2) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 6 + 1) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 6) = ((z >> 40) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave7Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long, - l7: Long): Array[Byte] = { - // output 8 * 56 bits - val result = new Array[Byte](56) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - val tmp7 = ((l7 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - val r7_mask = tmp7 & (1 << j) - z |= (r1_mask << (6 * j)) | (r2_mask << (6 * j + 1)) | (r3_mask << (6 * j + 2)) | - (r4_mask << (6 * j + 3)) | (r5_mask << (6 * j + 4)) | (r6_mask << (6 * j + 5)) | - (r7_mask << (6 * j + 6)) - j = j + 1 - } - result((7 - i) * 7 + 6) = (z & 0xFF).toByte - result((7 - i) * 7 + 5) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 7 + 4) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 7 + 3) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 7 + 2) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 7 + 1) = ((z >> 40) & 0xFF).toByte - result((7 - i) * 7) = ((z >> 48) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave8Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long, - l7: Long, - l8: Long): Array[Byte] = { - // output 8 * 64 bits - val result = new Array[Byte](64) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - val tmp7 = ((l7 >> (i * 8)) & 0xFF).toLong - val tmp8 = ((l8 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - val r7_mask = tmp7 & (1 << j) - val r8_mask = tmp8 & (1 << j) - z |= (r1_mask << (7 * j)) | (r2_mask << (7 * j + 1)) | (r3_mask << (7 * j + 2)) | - (r4_mask << (7 * j + 3)) | (r5_mask << (7 * j + 4)) | (r6_mask << (7 * j + 5)) | - (r7_mask << (7 * j + 6)) | (r8_mask << (7 * j + 7)) - j = j + 1 - } - result((7 - i) * 8 + 7) = (z & 0xFF).toByte - result((7 - i) * 8 + 6) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 8 + 5) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 8 + 4) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 8 + 3) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 8 + 2) = ((z >> 40) & 0xFF).toByte - result((7 - i) * 8 + 1) = ((z >> 48) & 0xFF).toByte - result((7 - i) * 8) = ((z >> 56) & 0xFF).toByte - i = i + 1 - } - result - } - - def interleaveBitsDefault(arrays: Array[Array[Byte]]): Array[Byte] = { - var totalLength = 0 - var maxLength = 0 - arrays.foreach { array => - totalLength += array.length - maxLength = maxLength.max(array.length * 8) - } - val result = new Array[Byte](totalLength) - var resultBit = 0 - - var bit = 0 - while (bit < maxLength) { - val bytePos = bit / 8 - val bitPos = bit % 8 - - for (arr <- arrays) { - val len = arr.length - if (bytePos < len) { - val resultBytePos = totalLength - 1 - resultBit / 8 - val resultBitPos = resultBit % 8 - result(resultBytePos) = - updatePos(result(resultBytePos), resultBitPos, arr(len - 1 - bytePos), bitPos) - resultBit += 1 - } - } - bit += 1 - } - result - } - - def updatePos(a: Byte, apos: Int, b: Byte, bpos: Int): Byte = { - var temp = (b & (1 << bpos)).toByte - if (apos > bpos) { - temp = (temp << (apos - bpos)).toByte - } else if (apos < bpos) { - temp = (temp >> (bpos - apos)).toByte - } - val atemp = (a & (1 << apos)).toByte - if (atemp == temp) { - return a - } - (a ^ (1 << apos)).toByte - } - - def toLong(a: Any): Long = { - a match { - case b: Boolean => (if (b) 1 else 0).toLong ^ BIT_64_MASK - case b: Byte => b.toLong ^ BIT_64_MASK - case s: Short => s.toLong ^ BIT_64_MASK - case i: Int => i.toLong ^ BIT_64_MASK - case l: Long => l ^ BIT_64_MASK - case f: Float => java.lang.Float.floatToRawIntBits(f).toLong ^ BIT_64_MASK - case d: Double => java.lang.Double.doubleToRawLongBits(d) ^ BIT_64_MASK - case str: UTF8String => str.getPrefix - case dec: Decimal => dec.toLong ^ BIT_64_MASK - case other: Any => - throw new KyuubiSQLExtensionException("Unsupported z-order type: " + other.getClass) - } - } - - def toByteArray(a: Any): Array[Byte] = { - a match { - case bo: Boolean => - booleanToByte(bo) - case b: Byte => - byteToByte(b) - case s: Short => - shortToByte(s) - case i: Int => - intToByte(i) - case l: Long => - longToByte(l) - case f: Float => - floatToByte(f) - case d: Double => - doubleToByte(d) - case str: UTF8String => - // truncate or padding str to 8 byte - paddingTo8Byte(str.getBytes) - case dec: Decimal => - longToByte(dec.toLong) - case other: Any => - throw new KyuubiSQLExtensionException("Unsupported z-order type: " + other.getClass) - } - } - - def booleanToByte(a: Boolean): Array[Byte] = { - if (a) { - byteToByte(1.toByte) - } else { - byteToByte(0.toByte) - } - } - - def byteToByte(a: Byte): Array[Byte] = { - val tmp = (a ^ BIT_8_MASK).toByte - Array(tmp) - } - - def shortToByte(a: Short): Array[Byte] = { - val tmp = a ^ BIT_16_MASK - Array(((tmp >> 8) & 0xFF).toByte, (tmp & 0xFF).toByte) - } - - def intToByte(a: Int): Array[Byte] = { - val result = new Array[Byte](4) - var i = 0 - val tmp = a ^ BIT_32_MASK - while (i <= 3) { - val offset = i * 8 - result(3 - i) = ((tmp >> offset) & 0xFF).toByte - i += 1 - } - result - } - - def longToByte(a: Long): Array[Byte] = { - val result = new Array[Byte](8) - var i = 0 - val tmp = a ^ BIT_64_MASK - while (i <= 7) { - val offset = i * 8 - result(7 - i) = ((tmp >> offset) & 0xFF).toByte - i += 1 - } - result - } - - def floatToByte(a: Float): Array[Byte] = { - val fi = jFloat.floatToRawIntBits(a) - intToByte(fi) - } - - def doubleToByte(a: Double): Array[Byte] = { - val dl = jDouble.doubleToRawLongBits(a) - longToByte(dl) - } - - def paddingTo8Byte(a: Array[Byte]): Array[Byte] = { - val len = a.length - if (len == 8) { - a - } else if (len > 8) { - val result = new Array[Byte](8) - System.arraycopy(a, 0, result, 0, 8) - result - } else { - val result = new Array[Byte](8) - System.arraycopy(a, 0, result, 8 - len, len) - result - } - } - - def defaultByteArrayValue(dataType: DataType): Array[Byte] = toByteArray { - defaultValue(dataType) - } - - def defaultValue(dataType: DataType): Any = { - dataType match { - case BooleanType => - true - case ByteType => - Byte.MaxValue - case ShortType => - Short.MaxValue - case IntegerType | DateType => - Int.MaxValue - case LongType | TimestampType | _: DecimalType => - Long.MaxValue - case FloatType => - Float.MaxValue - case DoubleType => - Double.MaxValue - case StringType => - // we pad string to 8 bytes so it's equal to long - UTF8String.fromBytes(longToByte(Long.MaxValue)) - case other: Any => - throw new KyuubiSQLExtensionException(s"Unsupported z-order type: ${other.catalogString}") - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala deleted file mode 100644 index 32fb9f5ce84..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala +++ /dev/null @@ -1,283 +0,0 @@ -/* - * 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.spark.sql - -import scala.annotation.tailrec -import scala.collection.mutable -import scala.collection.mutable.ArrayBuffer - -import org.apache.spark.{ExecutorAllocationClient, MapOutputTrackerMaster, SparkContext, SparkEnv} -import org.apache.spark.internal.Logging -import org.apache.spark.resource.ResourceProfile -import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SortExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive._ -import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec -import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec} - -import org.apache.kyuubi.sql.{KyuubiSQLConf, MarkNumOutputColumnsRule} - -/** - * This rule assumes the final write stage has less cores requirement than previous, otherwise - * this rule would take no effect. - * - * It provide a feature: - * 1. Kill redundant executors before running final write stage - */ -case class FinalStageResourceManager(session: SparkSession) - extends Rule[SparkPlan] with FinalRebalanceStageHelper { - override def apply(plan: SparkPlan): SparkPlan = { - if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED)) { - return plan - } - - if (!MarkNumOutputColumnsRule.isWrite(session, plan)) { - return plan - } - - val sc = session.sparkContext - val dra = sc.getConf.getBoolean("spark.dynamicAllocation.enabled", false) - val coresPerExecutor = sc.getConf.getInt("spark.executor.cores", 1) - val minExecutors = sc.getConf.getInt("spark.dynamicAllocation.minExecutors", 0) - val maxExecutors = sc.getConf.getInt("spark.dynamicAllocation.maxExecutors", Int.MaxValue) - val factor = conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_PARTITION_FACTOR) - val hasImprovementRoom = maxExecutors - 1 > minExecutors * factor - // Fast fail if: - // 1. DRA off - // 2. only work with yarn and k8s - // 3. maxExecutors is not bigger than minExecutors * factor - if (!dra || !sc.schedulerBackend.isInstanceOf[CoarseGrainedSchedulerBackend] || - !hasImprovementRoom) { - return plan - } - - val stageOpt = findFinalRebalanceStage(plan) - if (stageOpt.isEmpty) { - return plan - } - - // It's not safe to kill executors if this plan contains table cache. - // If the executor loses then the rdd would re-compute those partition. - if (hasTableCache(plan) && - conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_SKIP_KILLING_EXECUTORS_FOR_TABLE_CACHE)) { - return plan - } - - // TODO: move this to query stage optimizer when updating Spark to 3.5.x - // Since we are in `prepareQueryStage`, the AQE shuffle read has not been applied. - // So we need to apply it by self. - val shuffleRead = queryStageOptimizerRules.foldLeft(stageOpt.get.asInstanceOf[SparkPlan]) { - case (latest, rule) => rule.apply(latest) - } - val (targetCores, stage) = shuffleRead match { - case AQEShuffleReadExec(stage: ShuffleQueryStageExec, partitionSpecs) => - (partitionSpecs.length, stage) - case stage: ShuffleQueryStageExec => - // we can still kill executors if no AQE shuffle read, e.g., `.repartition(2)` - (stage.shuffle.numPartitions, stage) - case _ => - // it should never happen in current Spark, but to be safe do nothing if happens - logWarning("BUG, Please report to Apache Kyuubi community") - return plan - } - // The condition whether inject custom resource profile: - // - target executors < active executors - // - active executors - target executors > min executors - val numActiveExecutors = sc.getExecutorIds().length - val targetExecutors = (math.ceil(targetCores.toFloat / coresPerExecutor) * factor).toInt - .max(1) - val hasBenefits = targetExecutors < numActiveExecutors && - (numActiveExecutors - targetExecutors) > minExecutors - logInfo(s"The snapshot of current executors view, " + - s"active executors: $numActiveExecutors, min executor: $minExecutors, " + - s"target executors: $targetExecutors, has benefits: $hasBenefits") - if (hasBenefits) { - val shuffleId = stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleDependency.shuffleId - val numReduce = stage.plan.asInstanceOf[ShuffleExchangeExec].numPartitions - // Now, there is only a final rebalance stage waiting to execute and all tasks of previous - // stage are finished. Kill redundant existed executors eagerly so the tasks of final - // stage can be centralized scheduled. - killExecutors(sc, targetExecutors, shuffleId, numReduce) - } - - plan - } - - /** - * The priority of kill executors follow: - * 1. kill executor who is younger than other (The older the JIT works better) - * 2. kill executor who produces less shuffle data first - */ - private def findExecutorToKill( - sc: SparkContext, - targetExecutors: Int, - shuffleId: Int, - numReduce: Int): Seq[String] = { - val tracker = SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] - val shuffleStatusOpt = tracker.shuffleStatuses.get(shuffleId) - if (shuffleStatusOpt.isEmpty) { - return Seq.empty - } - val shuffleStatus = shuffleStatusOpt.get - val executorToBlockSize = new mutable.HashMap[String, Long] - shuffleStatus.withMapStatuses { mapStatus => - mapStatus.foreach { status => - var i = 0 - var sum = 0L - while (i < numReduce) { - sum += status.getSizeForBlock(i) - i += 1 - } - executorToBlockSize.getOrElseUpdate(status.location.executorId, sum) - } - } - - val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] - val executorsWithRegistrationTs = backend.getExecutorsWithRegistrationTs() - val existedExecutors = executorsWithRegistrationTs.keys.toSet - val expectedNumExecutorToKill = existedExecutors.size - targetExecutors - if (expectedNumExecutorToKill < 1) { - return Seq.empty - } - - val executorIdsToKill = new ArrayBuffer[String]() - // We first kill executor who does not hold shuffle block. It would happen because - // the last stage is running fast and finished in a short time. The existed executors are - // from previous stages that have not been killed by DRA, so we can not find it by tracking - // shuffle status. - // We should evict executors by their alive time first and retain all of executors which - // have better locality for shuffle block. - executorsWithRegistrationTs.toSeq.sortBy(_._2).foreach { case (id, _) => - if (executorIdsToKill.length < expectedNumExecutorToKill && - !executorToBlockSize.contains(id)) { - executorIdsToKill.append(id) - } - } - - // Evict the rest executors according to the shuffle block size - executorToBlockSize.toSeq.sortBy(_._2).foreach { case (id, _) => - if (executorIdsToKill.length < expectedNumExecutorToKill && existedExecutors.contains(id)) { - executorIdsToKill.append(id) - } - } - - executorIdsToKill.toSeq - } - - private def killExecutors( - sc: SparkContext, - targetExecutors: Int, - shuffleId: Int, - numReduce: Int): Unit = { - val executorAllocationClient = sc.schedulerBackend.asInstanceOf[ExecutorAllocationClient] - - val executorsToKill = - if (conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL)) { - executorAllocationClient.getExecutorIds() - } else { - findExecutorToKill(sc, targetExecutors, shuffleId, numReduce) - } - logInfo(s"Request to kill executors, total count ${executorsToKill.size}, " + - s"[${executorsToKill.mkString(", ")}].") - if (executorsToKill.isEmpty) { - return - } - - // Note, `SparkContext#killExecutors` does not allow with DRA enabled, - // see `https://github.com/apache/spark/pull/20604`. - // It may cause the status in `ExecutorAllocationManager` inconsistent with - // `CoarseGrainedSchedulerBackend` for a while. But it should be synchronous finally. - // - // We should adjust target num executors, otherwise `YarnAllocator` might re-request original - // target executors if DRA has not updated target executors yet. - // Note, DRA would re-adjust executors if there are more tasks to be executed, so we are safe. - // - // * We kill executor - // * YarnAllocator re-request target executors - // * DRA can not release executors since they are new added - // ----------------------------------------------------------------> timeline - executorAllocationClient.killExecutors( - executorIds = executorsToKill, - adjustTargetNumExecutors = true, - countFailures = false, - force = false) - - FinalStageResourceManager.getAdjustedTargetExecutors(sc) - .filter(_ < targetExecutors).foreach { adjustedExecutors => - val delta = targetExecutors - adjustedExecutors - logInfo(s"Target executors after kill ($adjustedExecutors) is lower than required " + - s"($targetExecutors). Requesting $delta additional executor(s).") - executorAllocationClient.requestExecutors(delta) - } - } - - @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq( - OptimizeSkewInRebalancePartitions, - CoalesceShufflePartitions(session), - OptimizeShuffleWithLocalRead) -} - -object FinalStageResourceManager extends Logging { - - private[sql] def getAdjustedTargetExecutors(sc: SparkContext): Option[Int] = { - sc.schedulerBackend match { - case schedulerBackend: CoarseGrainedSchedulerBackend => - try { - val field = classOf[CoarseGrainedSchedulerBackend] - .getDeclaredField("requestedTotalExecutorsPerResourceProfile") - field.setAccessible(true) - schedulerBackend.synchronized { - val requestedTotalExecutorsPerResourceProfile = - field.get(schedulerBackend).asInstanceOf[mutable.HashMap[ResourceProfile, Int]] - val defaultRp = sc.resourceProfileManager.defaultResourceProfile - requestedTotalExecutorsPerResourceProfile.get(defaultRp) - } - } catch { - case e: Exception => - logWarning("Failed to get requestedTotalExecutors of Default ResourceProfile", e) - None - } - case _ => None - } - } -} - -trait FinalRebalanceStageHelper extends AdaptiveSparkPlanHelper { - @tailrec - final protected def findFinalRebalanceStage(plan: SparkPlan): Option[ShuffleQueryStageExec] = { - plan match { - case p: ProjectExec => findFinalRebalanceStage(p.child) - case f: FilterExec => findFinalRebalanceStage(f.child) - case s: SortExec if !s.global => findFinalRebalanceStage(s.child) - case stage: ShuffleQueryStageExec - if stage.isMaterialized && stage.mapStats.isDefined && - stage.plan.isInstanceOf[ShuffleExchangeExec] && - stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleOrigin != ENSURE_REQUIREMENTS => - Some(stage) - case _ => None - } - } - - final protected def hasTableCache(plan: SparkPlan): Boolean = { - find(plan) { - case _: InMemoryTableScanExec => true - case _ => false - }.isDefined - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala deleted file mode 100644 index 30c042b2a2c..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{CustomResourceProfileExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive._ - -import org.apache.kyuubi.sql.{KyuubiSQLConf, MarkNumOutputColumnsRule} - -/** - * Inject custom resource profile for final write stage, so we can specify custom - * executor resource configs. - */ -case class InjectCustomResourceProfile(session: SparkSession) - extends Rule[SparkPlan] with FinalRebalanceStageHelper { - override def apply(plan: SparkPlan): SparkPlan = { - if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED)) { - return plan - } - - if (!MarkNumOutputColumnsRule.isWrite(session, plan)) { - return plan - } - - val stage = findFinalRebalanceStage(plan) - if (stage.isEmpty) { - return plan - } - - // TODO: Ideally, We can call `CoarseGrainedSchedulerBackend.requestTotalExecutors` eagerly - // to reduce the task submit pending time, but it may lose task locality. - // - // By default, it would request executors when catch stage submit event. - injectCustomResourceProfile(plan, stage.get.id) - } - - private def injectCustomResourceProfile(plan: SparkPlan, id: Int): SparkPlan = { - plan match { - case stage: ShuffleQueryStageExec if stage.id == id => - CustomResourceProfileExec(stage) - case _ => plan.mapChildren(child => injectCustomResourceProfile(child, id)) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala deleted file mode 100644 index ce496eb474c..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExpressionSet, PredicateHelper, SubqueryExpression} -import org.apache.spark.sql.catalyst.plans.logical.LeafNode -import org.apache.spark.sql.execution.datasources.DataSourceStrategy -import org.apache.spark.sql.types.StructType - -trait PruneFileSourcePartitionHelper extends PredicateHelper { - - def getPartitionKeyFiltersAndDataFilters( - sparkSession: SparkSession, - relation: LeafNode, - partitionSchema: StructType, - filters: Seq[Expression], - output: Seq[AttributeReference]): (ExpressionSet, Seq[Expression]) = { - val normalizedFilters = DataSourceStrategy.normalizeExprs( - filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f)), - output) - val partitionColumns = - relation.resolve(partitionSchema, sparkSession.sessionState.analyzer.resolver) - val partitionSet = AttributeSet(partitionColumns) - val (partitionFilters, dataFilters) = normalizedFilters.partition(f => - f.references.subsetOf(partitionSet)) - val extraPartitionFilter = - dataFilters.flatMap(extractPredicatesWithinOutputSet(_, partitionSet)) - - (ExpressionSet(partitionFilters ++ extraPartitionFilter), dataFilters) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala deleted file mode 100644 index 043d6496b22..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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.spark.sql.execution - -import org.apache.spark.network.util.{ByteUnit, JavaUtils} -import org.apache.spark.rdd.RDD -import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfileBuilder} -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} -import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * This node wraps the final executed plan and inject custom resource profile to the RDD. - * It assumes that, the produced RDD would create the `ResultStage` in `DAGScheduler`, - * so it makes resource isolation between previous and final stage. - * - * Note that, Spark does not support config `minExecutors` for each resource profile. - * Which means, it would retain `minExecutors` for each resource profile. - * So, suggest set `spark.dynamicAllocation.minExecutors` to 0 if enable this feature. - */ -case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode { - override def output: Seq[Attribute] = child.output - override def outputPartitioning: Partitioning = child.outputPartitioning - override def outputOrdering: Seq[SortOrder] = child.outputOrdering - override def supportsColumnar: Boolean = child.supportsColumnar - override def supportsRowBased: Boolean = child.supportsRowBased - override protected def doCanonicalize(): SparkPlan = child.canonicalized - - private val executorCores = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_CORES).getOrElse( - sparkContext.getConf.getInt("spark.executor.cores", 1)) - private val executorMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY).getOrElse( - sparkContext.getConf.get("spark.executor.memory", "2G")) - private val executorMemoryOverhead = - conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD) - .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G")) - private val executorOffHeapMemory = - if (sparkContext.getConf.getBoolean("spark.memory.offHeap.enabled", false)) { - conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY) - } else { - None - } - - override lazy val metrics: Map[String, SQLMetric] = { - val base = Map( - "executorCores" -> SQLMetrics.createMetric(sparkContext, "executor cores"), - "executorMemory" -> SQLMetrics.createMetric(sparkContext, "executor memory (MiB)"), - "executorMemoryOverhead" -> SQLMetrics.createMetric( - sparkContext, - "executor memory overhead (MiB)")) - val addition = executorOffHeapMemory.map(_ => - "executorOffHeapMemory" -> - SQLMetrics.createMetric(sparkContext, "executor off heap memory (MiB)")).toMap - base ++ addition - } - - private def wrapResourceProfile[T](rdd: RDD[T]): RDD[T] = { - if (Utils.isTesting) { - // do nothing for local testing - return rdd - } - - metrics("executorCores") += executorCores - metrics("executorMemory") += JavaUtils.byteStringAs(executorMemory, ByteUnit.MiB) - metrics("executorMemoryOverhead") += JavaUtils.byteStringAs( - executorMemoryOverhead, - ByteUnit.MiB) - executorOffHeapMemory.foreach(m => - metrics("executorOffHeapMemory") += JavaUtils.byteStringAs(m, ByteUnit.MiB)) - - val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) - SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) - - val resourceProfileBuilder = new ResourceProfileBuilder() - val executorResourceRequests = new ExecutorResourceRequests() - executorResourceRequests.cores(executorCores) - executorResourceRequests.memory(executorMemory) - executorResourceRequests.memoryOverhead(executorMemoryOverhead) - executorOffHeapMemory.foreach(executorResourceRequests.offHeapMemory) - resourceProfileBuilder.require(executorResourceRequests) - rdd.withResources(resourceProfileBuilder.build()) - rdd - } - - override protected def doExecute(): RDD[InternalRow] = { - val rdd = child.execute() - wrapResourceProfile(rdd) - } - - override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { - val rdd = child.executeColumnar() - wrapResourceProfile(rdd) - } - - override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = { - this.copy(child = newChild) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/resources/log4j2-test.xml b/extensions/spark/kyuubi-extension-spark-3-3/src/test/resources/log4j2-test.xml deleted file mode 100644 index 3110216c17c..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/resources/log4j2-test.xml +++ /dev/null @@ -1,43 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala deleted file mode 100644 index ff13268948f..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala +++ /dev/null @@ -1,53 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.plans.logical.{DropNamespace, NoopCommand} -import org.apache.spark.sql.execution.command._ - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class DropIgnoreNonexistentSuite extends KyuubiSparkSQLExtensionTest { - - test("drop ignore nonexistent") { - withSQLConf(KyuubiSQLConf.DROP_IGNORE_NONEXISTENT.key -> "true") { - // drop nonexistent database - val df1 = sql("DROP DATABASE nonexistent_database") - assert(df1.queryExecution.analyzed.asInstanceOf[DropNamespace].ifExists == true) - - // drop nonexistent table - val df2 = sql("DROP TABLE nonexistent_table") - assert(df2.queryExecution.analyzed.isInstanceOf[NoopCommand]) - - // drop nonexistent view - val df3 = sql("DROP VIEW nonexistent_view") - assert(df3.queryExecution.analyzed.isInstanceOf[NoopCommand]) - - // drop nonexistent function - val df4 = sql("DROP FUNCTION nonexistent_function") - assert(df4.queryExecution.analyzed.isInstanceOf[NoopCommand]) - - // drop nonexistent PARTITION - withTable("test") { - sql("CREATE TABLE IF NOT EXISTS test(i int) PARTITIONED BY (p int)") - val df5 = sql("ALTER TABLE test DROP PARTITION (p = 1)") - assert(df5.queryExecution.analyzed - .asInstanceOf[AlterTableDropPartitionCommand].ifExists == true) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala deleted file mode 100644 index e1ad1b2cab6..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala +++ /dev/null @@ -1,207 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.{FinalStageConfigIsolation, KyuubiSQLConf, MarkNumOutputColumnsRule} - -class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - test("final stage config set reset check") { - withSQLConf( - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "false", - "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum" -> "1", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "100") { - // use loop to double check final stage config doesn't affect the sql query each other - (1 to 3).foreach { _ => - sql("SELECT COUNT(*) FROM VALUES(1) as t(c)").collect() - assert(spark.sessionState.conf.getConfString( - "spark.sql.previousStage.adaptive.coalescePartitions.minPartitionNum") === - FinalStageConfigIsolation.INTERNAL_UNSET_CONFIG_TAG) - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.coalescePartitions.minPartitionNum") === - "1") - assert(spark.sessionState.conf.getConfString( - "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum") === - "1") - - // 64MB - assert(spark.sessionState.conf.getConfString( - "spark.sql.previousStage.adaptive.advisoryPartitionSizeInBytes") === - "67108864b") - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.advisoryPartitionSizeInBytes") === - "100") - assert(spark.sessionState.conf.getConfString( - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes") === - "100") - } - - sql("SET spark.sql.adaptive.advisoryPartitionSizeInBytes=1") - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.advisoryPartitionSizeInBytes") === - "1") - assert(!spark.sessionState.conf.contains( - "spark.sql.previousStage.adaptive.advisoryPartitionSizeInBytes")) - - sql("SET a=1") - assert(spark.sessionState.conf.getConfString("a") === "1") - - sql("RESET spark.sql.adaptive.coalescePartitions.minPartitionNum") - assert(!spark.sessionState.conf.contains( - "spark.sql.adaptive.coalescePartitions.minPartitionNum")) - assert(!spark.sessionState.conf.contains( - "spark.sql.previousStage.adaptive.coalescePartitions.minPartitionNum")) - - sql("RESET a") - assert(!spark.sessionState.conf.contains("a")) - } - } - - test("final stage config isolation") { - def checkPartitionNum( - sqlString: String, - previousPartitionNum: Int, - finalPartitionNum: Int): Unit = { - val df = sql(sqlString) - df.collect() - val shuffleReaders = collect(df.queryExecution.executedPlan) { - case customShuffleReader: AQEShuffleReadExec => customShuffleReader - } - assert(shuffleReaders.nonEmpty) - // reorder stage by stage id to ensure we get the right stage - val sortedShuffleReaders = shuffleReaders.sortWith { - case (s1, s2) => - s1.child.asInstanceOf[QueryStageExec].id < s2.child.asInstanceOf[QueryStageExec].id - } - if (sortedShuffleReaders.length > 1) { - assert(sortedShuffleReaders.head.partitionSpecs.length === previousPartitionNum) - } - assert(sortedShuffleReaders.last.partitionSpecs.length === finalPartitionNum) - assert(df.rdd.partitions.length === finalPartitionNum) - assert(MarkNumOutputColumnsRule.numOutputColumns(spark).isEmpty) - } - - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", - SQLConf.SHUFFLE_PARTITIONS.key -> "3", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "false", - "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "1", - "spark.sql.adaptive.coalescePartitions.minPartitionSize" -> "1", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "10000000") { - - // use loop to double check final stage config doesn't affect the sql query each other - (1 to 3).foreach { _ => - checkPartitionNum( - "SELECT c1, count(*) FROM t1 GROUP BY c1", - 1, - 1) - - checkPartitionNum( - "SELECT c2, count(*) FROM (SELECT c1, count(*) as c2 FROM t1 GROUP BY c1) GROUP BY c2", - 3, - 1) - - checkPartitionNum( - "SELECT t1.c1, count(*) FROM t1 JOIN t2 ON t1.c2 = t2.c2 GROUP BY t1.c1", - 3, - 1) - - checkPartitionNum( - """ - | SELECT /*+ REPARTITION */ - | t1.c1, count(*) FROM t1 - | JOIN t2 ON t1.c2 = t2.c2 - | JOIN t3 ON t1.c1 = t3.c1 - | GROUP BY t1.c1 - |""".stripMargin, - 3, - 1) - - // one shuffle reader - checkPartitionNum( - """ - | SELECT /*+ BROADCAST(t1) */ - | t1.c1, t2.c2 FROM t1 - | JOIN t2 ON t1.c2 = t2.c2 - | DISTRIBUTE BY c1 - |""".stripMargin, - 1, - 1) - - // test ReusedExchange - checkPartitionNum( - """ - |SELECT /*+ REPARTITION */ t0.c2 FROM ( - |SELECT t1.c1, (count(*) + c1) as c2 FROM t1 GROUP BY t1.c1 - |) t0 JOIN ( - |SELECT t1.c1, (count(*) + c1) as c2 FROM t1 GROUP BY t1.c1 - |) t1 ON t0.c2 = t1.c2 - |""".stripMargin, - 3, - 1) - - // one shuffle reader - checkPartitionNum( - """ - |SELECT t0.c1 FROM ( - |SELECT t1.c1 FROM t1 GROUP BY t1.c1 - |) t0 JOIN ( - |SELECT t1.c1 FROM t1 GROUP BY t1.c1 - |) t1 ON t0.c1 = t1.c1 - |""".stripMargin, - 1, - 1) - } - } - } - - test("final stage config isolation write only") { - withSQLConf( - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "true", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "7") { - sql("set spark.sql.adaptive.advisoryPartitionSizeInBytes=5") - sql("SELECT * FROM t1").count() - assert(MarkNumOutputColumnsRule.numOutputColumns(spark).isEmpty) - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("5")) - - withTable("tmp") { - sql("CREATE TABLE t1 USING PARQUET SELECT /*+ repartition */ 1 AS c1, 'a' AS c2") - assert(MarkNumOutputColumnsRule.numOutputColumns(spark).contains("2")) - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("7")) - } - - sql("SELECT * FROM t1").count() - assert(MarkNumOutputColumnsRule.numOutputColumns(spark).isEmpty) - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("5")) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala deleted file mode 100644 index 4b9991ef6f2..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala +++ /dev/null @@ -1,62 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.scalatest.time.{Minutes, Span} - -import org.apache.kyuubi.sql.KyuubiSQLConf -import org.apache.kyuubi.tags.SparkLocalClusterTest - -@SparkLocalClusterTest -class FinalStageResourceManagerSuite extends KyuubiSparkSQLExtensionTest { - - override def sparkConf(): SparkConf = { - // It is difficult to run spark in local-cluster mode when spark.testing is set. - sys.props.remove("spark.testing") - - super.sparkConf().set("spark.master", "local-cluster[3, 1, 1024]") - .set("spark.dynamicAllocation.enabled", "true") - .set("spark.dynamicAllocation.initialExecutors", "3") - .set("spark.dynamicAllocation.minExecutors", "1") - .set("spark.dynamicAllocation.shuffleTracking.enabled", "true") - .set(KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key, "true") - .set(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED.key, "true") - } - - test("[KYUUBI #5136][Bug] Final Stage hangs forever") { - // Prerequisite to reproduce the bug: - // 1. Dynamic allocation is enabled. - // 2. Dynamic allocation min executors is 1. - // 3. target executors < active executors. - // 4. No active executor is left after FinalStageResourceManager killed executors. - // This is possible because FinalStageResourceManager retained executors may already be - // requested to be killed but not died yet. - // 5. Final Stage required executors is 1. - withSQLConf( - (KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL.key, "true")) { - withTable("final_stage") { - eventually(timeout(Span(10, Minutes))) { - sql( - "CREATE TABLE final_stage AS SELECT id, count(*) as num FROM (SELECT 0 id) GROUP BY id") - } - assert(FinalStageResourceManager.getAdjustedTargetExecutors(spark.sparkContext).get == 1) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala deleted file mode 100644 index b0767b18708..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} -import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class InjectResourceProfileSuite extends KyuubiSparkSQLExtensionTest { - private def checkCustomResourceProfile(sqlString: String, exists: Boolean): Unit = { - @volatile var lastEvent: SparkListenerSQLAdaptiveExecutionUpdate = null - val listener = new SparkListener { - override def onOtherEvent(event: SparkListenerEvent): Unit = { - event match { - case e: SparkListenerSQLAdaptiveExecutionUpdate => lastEvent = e - case _ => - } - } - } - - spark.sparkContext.addSparkListener(listener) - try { - sql(sqlString).collect() - spark.sparkContext.listenerBus.waitUntilEmpty() - assert(lastEvent != null) - var current = lastEvent.sparkPlanInfo - var shouldStop = false - while (!shouldStop) { - if (current.nodeName != "CustomResourceProfile") { - if (current.children.isEmpty) { - assert(!exists) - shouldStop = true - } else { - current = current.children.head - } - } else { - assert(exists) - shouldStop = true - } - } - } finally { - spark.sparkContext.removeSparkListener(listener) - } - } - - test("Inject resource profile") { - withTable("t") { - withSQLConf( - "spark.sql.adaptive.forceApply" -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED.key -> "true") { - - sql("CREATE TABLE t (c1 int, c2 string) USING PARQUET") - - checkCustomResourceProfile("INSERT INTO TABLE t VALUES(1, 'a')", false) - checkCustomResourceProfile("SELECT 1", false) - checkCustomResourceProfile( - "INSERT INTO TABLE t SELECT /*+ rebalance */ * FROM VALUES(1, 'a')", - true) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala deleted file mode 100644 index f0d38465734..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala +++ /dev/null @@ -1,19 +0,0 @@ -/* - * 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.spark.sql - -class InsertShuffleNodeBeforeJoinSuite extends InsertShuffleNodeBeforeJoinSuiteBase diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala deleted file mode 100644 index c657dee49f3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala +++ /dev/null @@ -1,98 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeLike} -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} - -import org.apache.kyuubi.sql.KyuubiSQLConf - -trait InsertShuffleNodeBeforeJoinSuiteBase extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - override def sparkConf(): SparkConf = { - super.sparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLCommonExtension") - } - - test("force shuffle before join") { - def checkShuffleNodeNum(sqlString: String, num: Int): Unit = { - var expectedResult: Seq[Row] = Seq.empty - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - expectedResult = sql(sqlString).collect() - } - val df = sql(sqlString) - checkAnswer(df, expectedResult) - assert( - collect(df.queryExecution.executedPlan) { - case shuffle: ShuffleExchangeLike if shuffle.shuffleOrigin == ENSURE_REQUIREMENTS => - shuffle - }.size == num) - } - - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - KyuubiSQLConf.FORCE_SHUFFLE_BEFORE_JOIN.key -> "true") { - Seq("SHUFFLE_HASH", "MERGE").foreach { joinHint => - // positive case - checkShuffleNodeNum( - s""" - |SELECT /*+ $joinHint(t2, t3) */ t1.c1, t1.c2, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN t3 ON t1.c1 = t3.c1 - | """.stripMargin, - 4) - - // negative case - checkShuffleNodeNum( - s""" - |SELECT /*+ $joinHint(t2, t3) */ t1.c1, t1.c2, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN t3 ON t1.c2 = t3.c2 - | """.stripMargin, - 4) - } - - checkShuffleNodeNum( - """ - |SELECT t1.c1, t2.c1, t3.c2 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN ( - | SELECT c2, count(*) FROM t1 GROUP BY c2 - | ) t3 ON t1.c1 = t3.c2 - | """.stripMargin, - 5) - - checkShuffleNodeNum( - """ - |SELECT t1.c1, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN ( - | SELECT c1, count(*) FROM t1 GROUP BY c1 - | ) t3 ON t1.c1 = t3.c1 - | """.stripMargin, - 5) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala deleted file mode 100644 index 1fa136e60b9..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala +++ /dev/null @@ -1,94 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.hadoop.hive.conf.HiveConf.ConfVars -import org.apache.spark.SparkConf -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.test.SQLTestData.TestData -import org.apache.spark.sql.test.SQLTestUtils -import org.apache.spark.util.Utils - -trait KyuubiSparkSQLExtensionTest extends QueryTest - with SQLTestUtils - with AdaptiveSparkPlanHelper { - sys.props.put("spark.testing", "1") - - private var _spark: Option[SparkSession] = None - protected def spark: SparkSession = _spark.getOrElse { - throw new RuntimeException("test spark session don't initial before using it.") - } - - override protected def beforeAll(): Unit = { - if (_spark.isEmpty) { - _spark = Option(SparkSession.builder() - .master("local[1]") - .config(sparkConf) - .enableHiveSupport() - .getOrCreate()) - } - super.beforeAll() - } - - override protected def afterAll(): Unit = { - super.afterAll() - cleanupData() - _spark.foreach(_.stop) - } - - protected def setupData(): Unit = { - val self = spark - import self.implicits._ - spark.sparkContext.parallelize( - (1 to 100).map(i => TestData(i, i.toString)), - 10) - .toDF("c1", "c2").createOrReplaceTempView("t1") - spark.sparkContext.parallelize( - (1 to 10).map(i => TestData(i, i.toString)), - 5) - .toDF("c1", "c2").createOrReplaceTempView("t2") - spark.sparkContext.parallelize( - (1 to 50).map(i => TestData(i, i.toString)), - 2) - .toDF("c1", "c2").createOrReplaceTempView("t3") - } - - private def cleanupData(): Unit = { - spark.sql("DROP VIEW IF EXISTS t1") - spark.sql("DROP VIEW IF EXISTS t2") - spark.sql("DROP VIEW IF EXISTS t3") - } - - def sparkConf(): SparkConf = { - val basePath = Utils.createTempDir() + "/" + getClass.getCanonicalName - val metastorePath = basePath + "/metastore_db" - val warehousePath = basePath + "/warehouse" - new SparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLExtension") - .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true") - .set("spark.hadoop.hive.exec.dynamic.partition.mode", "nonstrict") - .set("spark.hadoop.hive.metastore.client.capability.check", "false") - .set( - ConfVars.METASTORECONNECTURLKEY.varname, - s"jdbc:derby:;databaseName=$metastorePath;create=true") - .set(StaticSQLConf.WAREHOUSE_PATH, warehousePath) - .set("spark.ui.enabled", "false") - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala deleted file mode 100644 index 1a776123461..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala +++ /dev/null @@ -1,324 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, RebalancePartitions, Sort} -import org.apache.spark.sql.execution.command.{DataWritingCommand, InsertIntoDataSourceDirCommand} -import org.apache.spark.sql.hive.HiveUtils -import org.apache.spark.sql.hive.execution.{InsertIntoHiveDirCommand, OptimizedCreateHiveTableAsSelectCommand} - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest { - - test("check rebalance exists") { - def check( - df: => DataFrame, - expectedRebalanceNumEnabled: Int = 1, - expectedRebalanceNumDisabled: Int = 0): Unit = { - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - assert( - df.queryExecution.analyzed.collect { - case r: RebalancePartitions => r - }.size == expectedRebalanceNumEnabled) - } - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "false") { - assert( - df.queryExecution.analyzed.collect { - case r: RebalancePartitions => r - }.size == expectedRebalanceNumDisabled) - } - } - - // It's better to set config explicitly in case of we change the default value. - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2='a') " + - "SELECT * FROM VALUES(1),(2) AS t(c1)")) - } - - withTable("tmp1", "tmp2") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - sql(s"CREATE TABLE tmp2 (c1 int) $storage PARTITIONED BY (c2 string)") - check( - sql( - """FROM VALUES(1),(2) - |INSERT INTO TABLE tmp1 PARTITION(c2='a') SELECT * - |INSERT INTO TABLE tmp2 PARTITION(c2='a') SELECT * - |""".stripMargin), - 2) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check(sql("INSERT INTO TABLE tmp1 SELECT * FROM VALUES(1),(2),(3) AS t(c1)")) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check( - sql("INSERT INTO TABLE tmp1 SELECT /*+ REBALANCE */ * FROM VALUES(1),(2),(3) AS t(c1)"), - 1, - 1) - } - - withTable("tmp1", "tmp2") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - sql(s"CREATE TABLE tmp2 (c1 int) $storage") - check( - sql( - """FROM VALUES(1),(2),(3) - |INSERT INTO TABLE tmp1 SELECT * - |INSERT INTO TABLE tmp2 SELECT * - |""".stripMargin), - 2) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage AS SELECT * FROM VALUES(1),(2),(3) AS t(c1)") - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage PARTITIONED BY(c2) AS " + - s"SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)") - } - } - } - } - - test("check rebalance does not exists") { - def check(df: DataFrame): Unit = { - assert( - df.queryExecution.analyzed.collect { - case r: RebalancePartitions => r - }.isEmpty) - } - - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - // test no write command - check(sql("SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - check(sql("SELECT count(*) FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - - // test not supported plan - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT /*+ repartition(10) */ * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2) ORDER BY c1")) - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2) LIMIT 10")) - } - } - - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "false") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check(sql("INSERT INTO TABLE tmp1 SELECT * FROM VALUES(1),(2),(3) AS t(c1)")) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage AS SELECT * FROM VALUES(1),(2),(3) AS t(c1)") - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage PARTITIONED BY(c2) AS " + - s"SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)") - } - } - } - } - - test("test dynamic partition write") { - def checkRepartitionExpression(df: DataFrame): Unit = { - assert(df.queryExecution.analyzed.collect { - case r: RebalancePartitions if r.partitionExpressions.size == 1 => - assert(r.partitionExpressions.head.asInstanceOf[Attribute].name === "c2") - r - }.size == 1) - } - - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - checkRepartitionExpression(sql("INSERT INTO TABLE tmp1 SELECT 1 as c1, 'a' as c2 ")) - } - - withTable("tmp1") { - checkRepartitionExpression( - sql("CREATE TABLE tmp1 PARTITIONED BY(C2) SELECT 1 as c1, 'a' as c2 ")) - } - } - } - } - - test("OptimizedCreateHiveTableAsSelectCommand") { - withSQLConf( - HiveUtils.CONVERT_METASTORE_PARQUET.key -> "true", - HiveUtils.CONVERT_METASTORE_CTAS.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTable("t") { - val df = sql(s"CREATE TABLE t STORED AS parquet AS SELECT 1 as a") - val ctas = df.queryExecution.analyzed.collect { - case _: OptimizedCreateHiveTableAsSelectCommand => true - } - assert(ctas.size == 1) - val repartition = df.queryExecution.analyzed.collect { - case _: RebalancePartitions => true - } - assert(repartition.size == 1) - } - } - } - - test("Infer rebalance and sorder orders") { - def checkShuffleAndSort(dataWritingCommand: LogicalPlan, sSize: Int, rSize: Int): Unit = { - assert(dataWritingCommand.isInstanceOf[DataWritingCommand]) - val plan = dataWritingCommand.asInstanceOf[DataWritingCommand].query - assert(plan.collect { - case s: Sort => s - }.size == sSize) - assert(plan.collect { - case r: RebalancePartitions if r.partitionExpressions.size == rSize => r - }.nonEmpty || rSize == 0) - } - - withView("v") { - withTable("t", "input1", "input2") { - withSQLConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS.key -> "true") { - sql(s"CREATE TABLE t (c1 int, c2 long) USING PARQUET PARTITIONED BY (p string)") - sql(s"CREATE TABLE input1 USING PARQUET AS SELECT * FROM VALUES(1,2),(1,3)") - sql(s"CREATE TABLE input2 USING PARQUET AS SELECT * FROM VALUES(1,3),(1,3)") - sql(s"CREATE VIEW v as SELECT col1, count(*) as col2 FROM input1 GROUP BY col1") - - val df0 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT /*+ broadcast(input2) */ input1.col1, input2.col1 - |FROM input1 - |JOIN input2 - |ON input1.col1 = input2.col1 - |""".stripMargin) - checkShuffleAndSort(df0.queryExecution.analyzed, 1, 1) - - val df1 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT /*+ broadcast(input2) */ input1.col1, input1.col2 - |FROM input1 - |LEFT JOIN input2 - |ON input1.col1 = input2.col1 and input1.col2 = input2.col2 - |""".stripMargin) - checkShuffleAndSort(df1.queryExecution.analyzed, 1, 2) - - val df2 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT col1 as c1, count(*) as c2 - |FROM input1 - |GROUP BY col1 - |HAVING count(*) > 0 - |""".stripMargin) - checkShuffleAndSort(df2.queryExecution.analyzed, 1, 1) - - // dynamic partition - val df3 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p) - |SELECT /*+ broadcast(input2) */ input1.col1, input1.col2, input1.col2 - |FROM input1 - |JOIN input2 - |ON input1.col1 = input2.col1 - |""".stripMargin) - checkShuffleAndSort(df3.queryExecution.analyzed, 0, 1) - - // non-deterministic - val df4 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT col1 + rand(), count(*) as c2 - |FROM input1 - |GROUP BY col1 - |""".stripMargin) - checkShuffleAndSort(df4.queryExecution.analyzed, 0, 0) - - // view - val df5 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT * FROM v - |""".stripMargin) - checkShuffleAndSort(df5.queryExecution.analyzed, 1, 1) - } - } - } - } - - test("Test rebalance in InsertIntoHiveDirCommand") { - withSQLConf( - HiveUtils.CONVERT_METASTORE_PARQUET.key -> "false", - HiveUtils.CONVERT_METASTORE_CTAS.key -> "false", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTempDir(tmpDir => { - spark.range(0, 1000, 1, 10).createOrReplaceTempView("tmp_table") - val df = sql(s"INSERT OVERWRITE DIRECTORY '${tmpDir.getPath}' " + - s"STORED AS PARQUET SELECT * FROM tmp_table") - val insertHiveDirCommand = df.queryExecution.analyzed.collect { - case _: InsertIntoHiveDirCommand => true - } - assert(insertHiveDirCommand.size == 1) - val repartition = df.queryExecution.analyzed.collect { - case _: RebalancePartitions => true - } - assert(repartition.size == 1) - }) - } - } - - test("Test rebalance in InsertIntoDataSourceDirCommand") { - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTempDir(tmpDir => { - spark.range(0, 1000, 1, 10).createOrReplaceTempView("tmp_table") - val df = sql(s"INSERT OVERWRITE DIRECTORY '${tmpDir.getPath}' " + - s"USING PARQUET SELECT * FROM tmp_table") - assert(df.queryExecution.analyzed.isInstanceOf[InsertIntoDataSourceDirCommand]) - val repartition = - df.queryExecution.analyzed.asInstanceOf[InsertIntoDataSourceDirCommand].query.collect { - case _: RebalancePartitions => true - } - assert(repartition.size == 1) - }) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala deleted file mode 100644 index 670d9ce7e43..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.spark.sql - -import java.util.OptionalLong - -import org.apache.spark.sql.connector.{RangeInputPartition, SimpleBatchTable, SimpleScanBuilder, SimpleWritableDataSource} -import org.apache.spark.sql.connector.catalog.Table -import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, Transform} -import org.apache.spark.sql.connector.read._ -import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, Partitioning} -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -class ReportStatisticsAndPartitionAwareDataSource extends SimpleWritableDataSource { - - class MyScanBuilder( - val partitionKeys: Seq[String]) extends SimpleScanBuilder - with SupportsReportStatistics with SupportsReportPartitioning { - - override def estimateStatistics(): Statistics = { - new Statistics { - override def sizeInBytes(): OptionalLong = OptionalLong.of(80) - - override def numRows(): OptionalLong = OptionalLong.of(10) - - } - } - - override def planInputPartitions(): Array[InputPartition] = { - Array(RangeInputPartition(0, 5), RangeInputPartition(5, 10)) - } - - override def outputPartitioning(): Partitioning = { - new KeyGroupedPartitioning(partitionKeys.map(FieldReference(_)).toArray, 10) - } - } - - override def getTable(options: CaseInsensitiveStringMap): Table = { - new SimpleBatchTable { - override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new MyScanBuilder(Seq("i")) - } - - override def partitioning(): Array[Transform] = { - Array(Expressions.identity("i")) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala deleted file mode 100644 index 2035d352562..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala +++ /dev/null @@ -1,53 +0,0 @@ -/* - * 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.spark.sql - -import java.util.OptionalLong - -import org.apache.spark.sql.connector._ -import org.apache.spark.sql.connector.catalog.Table -import org.apache.spark.sql.connector.read._ -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -class ReportStatisticsDataSource extends SimpleWritableDataSource { - - class MyScanBuilder extends SimpleScanBuilder - with SupportsReportStatistics { - - override def estimateStatistics(): Statistics = { - new Statistics { - override def sizeInBytes(): OptionalLong = OptionalLong.of(80) - - override def numRows(): OptionalLong = OptionalLong.of(10) - } - } - - override def planInputPartitions(): Array[InputPartition] = { - Array(RangeInputPartition(0, 5), RangeInputPartition(5, 10)) - } - - } - - override def getTable(options: CaseInsensitiveStringMap): Table = { - new SimpleBatchTable { - override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new MyScanBuilder - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala deleted file mode 100644 index 957089340ca..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala +++ /dev/null @@ -1,20 +0,0 @@ -/* - * 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.spark.sql - -class WatchDogSuite extends WatchDogSuiteBase {} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala deleted file mode 100644 index b5f8c67261f..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala +++ /dev/null @@ -1,254 +0,0 @@ -/* - * 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.spark.sql - -import java.io.File - -import scala.collection.JavaConverters._ - -import org.apache.commons.io.FileUtils -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} -import org.apache.kyuubi.sql.watchdog.{MaxFileSizeExceedException, MaxPartitionExceedException} - -trait WatchDogSuiteBase extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - case class LimitAndExpected(limit: Int, expected: Int) - - val limitAndExpecteds = List(LimitAndExpected(1, 1), LimitAndExpected(11, 10)) - - private def checkMaxPartition: Unit = { - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS.key -> "100") { - checkAnswer(sql("SELECT count(distinct(p)) FROM test"), Row(10) :: Nil) - } - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS.key -> "5") { - sql("SELECT * FROM test where p=1").queryExecution.sparkPlan - - sql(s"SELECT * FROM test WHERE p in (${Range(0, 5).toList.mkString(",")})") - .queryExecution.sparkPlan - - intercept[MaxPartitionExceedException]( - sql("SELECT * FROM test where p != 1").queryExecution.sparkPlan) - - intercept[MaxPartitionExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - - intercept[MaxPartitionExceedException](sql( - s"SELECT * FROM test WHERE p in (${Range(0, 6).toList.mkString(",")})") - .queryExecution.sparkPlan) - } - } - - test("watchdog with scan maxPartitions -- hive") { - Seq("textfile", "parquet").foreach { format => - withTable("test", "temp") { - sql( - s""" - |CREATE TABLE test(i int) - |PARTITIONED BY (p int) - |STORED AS $format""".stripMargin) - spark.range(0, 10, 1).selectExpr("id as col") - .createOrReplaceTempView("temp") - - for (part <- Range(0, 10)) { - sql( - s""" - |INSERT OVERWRITE TABLE test PARTITION (p='$part') - |select col from temp""".stripMargin) - } - checkMaxPartition - } - } - } - - test("watchdog with scan maxPartitions -- data source") { - withTempDir { dir => - withTempView("test") { - spark.range(10).selectExpr("id", "id as p") - .write - .partitionBy("p") - .mode("overwrite") - .save(dir.getCanonicalPath) - spark.read.load(dir.getCanonicalPath).createOrReplaceTempView("test") - checkMaxPartition - } - } - } - - private def checkMaxFileSize(tableSize: Long, nonPartTableSize: Long): Unit = { - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> tableSize.toString) { - checkAnswer(sql("SELECT count(distinct(p)) FROM test"), Row(10) :: Nil) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (tableSize / 2).toString) { - sql("SELECT * FROM test where p=1").queryExecution.sparkPlan - - sql(s"SELECT * FROM test WHERE p in (${Range(0, 3).toList.mkString(",")})") - .queryExecution.sparkPlan - - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test where p != 1").queryExecution.sparkPlan) - - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - - intercept[MaxFileSizeExceedException](sql( - s"SELECT * FROM test WHERE p in (${Range(0, 6).toList.mkString(",")})") - .queryExecution.sparkPlan) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> nonPartTableSize.toString) { - checkAnswer(sql("SELECT count(*) FROM test_non_part"), Row(10000) :: Nil) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (nonPartTableSize - 1).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan) - } - } - - test("watchdog with scan maxFileSize -- hive") { - Seq(false).foreach { convertMetastoreParquet => - withTable("test", "test_non_part", "temp") { - spark.range(10000).selectExpr("id as col") - .createOrReplaceTempView("temp") - - // partitioned table - sql( - s""" - |CREATE TABLE test(i int) - |PARTITIONED BY (p int) - |STORED AS parquet""".stripMargin) - for (part <- Range(0, 10)) { - sql( - s""" - |INSERT OVERWRITE TABLE test PARTITION (p='$part') - |select col from temp""".stripMargin) - } - - val tablePath = new File(spark.sessionState.catalog.externalCatalog - .getTable("default", "test").location) - val tableSize = FileUtils.listFiles(tablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // non-partitioned table - sql( - s""" - |CREATE TABLE test_non_part(i int) - |STORED AS parquet""".stripMargin) - sql( - s""" - |INSERT OVERWRITE TABLE test_non_part - |select col from temp""".stripMargin) - sql("ANALYZE TABLE test_non_part COMPUTE STATISTICS") - - val nonPartTablePath = new File(spark.sessionState.catalog.externalCatalog - .getTable("default", "test_non_part").location) - val nonPartTableSize = FileUtils.listFiles(nonPartTablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(nonPartTableSize > 0) - - // check - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> convertMetastoreParquet.toString) { - checkMaxFileSize(tableSize, nonPartTableSize) - } - } - } - } - - test("watchdog with scan maxFileSize -- data source") { - withTempDir { dir => - withTempView("test", "test_non_part") { - // partitioned table - val tablePath = new File(dir, "test") - spark.range(10).selectExpr("id", "id as p") - .write - .partitionBy("p") - .mode("overwrite") - .parquet(tablePath.getCanonicalPath) - spark.read.load(tablePath.getCanonicalPath).createOrReplaceTempView("test") - - val tableSize = FileUtils.listFiles(tablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // non-partitioned table - val nonPartTablePath = new File(dir, "test_non_part") - spark.range(10000).selectExpr("id", "id as p") - .write - .mode("overwrite") - .parquet(nonPartTablePath.getCanonicalPath) - spark.read.load(nonPartTablePath.getCanonicalPath).createOrReplaceTempView("test_non_part") - - val nonPartTableSize = FileUtils.listFiles(nonPartTablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // check - checkMaxFileSize(tableSize, nonPartTableSize) - } - } - } - - test("disable script transformation") { - withSQLConf(KyuubiSQLConf.SCRIPT_TRANSFORMATION_ENABLED.key -> "false") { - val e = intercept[KyuubiSQLExtensionException] { - sql("SELECT TRANSFORM('') USING 'ls /'") - } - assert(e.getMessage == "Script transformation is not allowed") - } - } - - test("watchdog with scan maxFileSize -- data source v2") { - val df = spark.read.format(classOf[ReportStatisticsAndPartitionAwareDataSource].getName).load() - df.createOrReplaceTempView("test") - val logical = df.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - val tableSize = logical.computeStats().sizeInBytes.toLong - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> tableSize.toString) { - sql("SELECT * FROM test").queryExecution.sparkPlan - } - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (tableSize / 2).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - } - - val nonPartDf = spark.read.format(classOf[ReportStatisticsDataSource].getName).load() - nonPartDf.createOrReplaceTempView("test_non_part") - val nonPartLogical = nonPartDf.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - val nonPartTableSize = nonPartLogical.computeStats().sizeInBytes.toLong - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> nonPartTableSize.toString) { - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (nonPartTableSize / 2).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala deleted file mode 100644 index d6c52b3a5af..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.benchmark.Benchmark -import org.apache.spark.sql.benchmark.KyuubiBenchmarkBase -import org.apache.spark.sql.internal.StaticSQLConf - -import org.apache.kyuubi.sql.zorder.ZorderBytesUtils - -/** - * Benchmark to measure performance with zorder core. - * - * {{{ - * RUN_BENCHMARK=1 ./build/mvn clean test \ - * -pl extensions/spark/kyuubi-extension-spark-3-3 -am \ - * -Pspark-3.3,kyuubi-extension-spark-3-3 \ - * -Dtest=none -DwildcardSuites=org.apache.spark.sql.ZorderCoreBenchmark - * }}} - */ -class ZorderCoreBenchmark extends KyuubiSparkSQLExtensionTest with KyuubiBenchmarkBase { - private val runBenchmark = sys.env.contains("RUN_BENCHMARK") - private val numRows = 1 * 1000 * 1000 - - private def randomInt(numColumns: Int): Seq[Array[Any]] = { - (1 to numRows).map { l => - val arr = new Array[Any](numColumns) - (0 until numColumns).foreach(col => arr(col) = l) - arr - } - } - - private def randomLong(numColumns: Int): Seq[Array[Any]] = { - (1 to numRows).map { l => - val arr = new Array[Any](numColumns) - (0 until numColumns).foreach(col => arr(col) = l.toLong) - arr - } - } - - private def interleaveMultiByteArrayBenchmark(): Unit = { - val benchmark = - new Benchmark(s"$numRows rows zorder core benchmark", numRows, output = output) - benchmark.addCase("2 int columns benchmark", 3) { _ => - randomInt(2).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("3 int columns benchmark", 3) { _ => - randomInt(3).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("4 int columns benchmark", 3) { _ => - randomInt(4).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("2 long columns benchmark", 3) { _ => - randomLong(2).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("3 long columns benchmark", 3) { _ => - randomLong(3).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("4 long columns benchmark", 3) { _ => - randomLong(4).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.run() - } - - private def paddingTo8ByteBenchmark() { - val iterations = 10 * 1000 * 1000 - - val b2 = Array('a'.toByte, 'b'.toByte) - val benchmark = - new Benchmark(s"$iterations iterations paddingTo8Byte benchmark", iterations, output = output) - benchmark.addCase("2 length benchmark", 3) { _ => - (1 to iterations).foreach(_ => ZorderBytesUtils.paddingTo8Byte(b2)) - } - - val b16 = Array.tabulate(16) { i => i.toByte } - benchmark.addCase("16 length benchmark", 3) { _ => - (1 to iterations).foreach(_ => ZorderBytesUtils.paddingTo8Byte(b16)) - } - - benchmark.run() - } - - test("zorder core benchmark") { - assume(runBenchmark) - - withHeader { - interleaveMultiByteArrayBenchmark() - paddingTo8ByteBenchmark() - } - } - - override def sparkConf(): SparkConf = { - super.sparkConf().remove(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuite.scala deleted file mode 100644 index d18e3035945..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuite.scala +++ /dev/null @@ -1,124 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.parser.ParserInterface -import org.apache.spark.sql.catalyst.plans.logical.{RebalancePartitions, Sort} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.{KyuubiSQLConf, SparkKyuubiSparkSQLParser} -import org.apache.kyuubi.sql.zorder.Zorder - -trait ZorderSuiteSpark extends ZorderSuiteBase { - - test("Add rebalance before zorder") { - Seq("true" -> false, "false" -> true).foreach { case (useOriginalOrdering, zorder) => - withSQLConf( - KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false", - KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED.key -> "true", - KyuubiSQLConf.ZORDER_USING_ORIGINAL_ORDERING_ENABLED.key -> useOriginalOrdering) { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) PARTITIONED BY (d string) - | TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,C2') - |""".stripMargin) - val p = sql("INSERT INTO TABLE t PARTITION(d='a') SELECT * FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p.collect { - case sort: Sort - if !sort.global && - ((sort.order.exists(_.child.isInstanceOf[Zorder]) && zorder) || - (!sort.order.exists(_.child.isInstanceOf[Zorder]) && !zorder)) => sort - }.size == 1) - assert(p.collect { - case rebalance: RebalancePartitions - if rebalance.references.map(_.name).exists(_.equals("c1")) => rebalance - }.size == 1) - - val p2 = sql("INSERT INTO TABLE t PARTITION(d) SELECT * FROM VALUES(1,'a','b')") - .queryExecution.analyzed - assert(p2.collect { - case sort: Sort - if (!sort.global && Seq("c1", "c2", "d").forall(x => - sort.references.map(_.name).exists(_.equals(x)))) && - ((sort.order.exists(_.child.isInstanceOf[Zorder]) && zorder) || - (!sort.order.exists(_.child.isInstanceOf[Zorder]) && !zorder)) => sort - }.size == 1) - assert(p2.collect { - case rebalance: RebalancePartitions - if Seq("c1", "c2", "d").forall(x => - rebalance.references.map(_.name).exists(_.equals(x))) => rebalance - }.size == 1) - } - } - } - } - - test("Two phase rebalance before Z-Order") { - withSQLConf( - SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> - "org.apache.spark.sql.catalyst.optimizer.CollapseRepartition", - KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false", - KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.TWO_PHASE_REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED.key -> "true") { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int) PARTITIONED BY (d string) - | TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1') - |""".stripMargin) - val p = sql("INSERT INTO TABLE t PARTITION(d) SELECT * FROM VALUES(1,'a')") - val rebalance = p.queryExecution.optimizedPlan.innerChildren - .flatMap(_.collect { case r: RebalancePartitions => r }) - assert(rebalance.size == 2) - assert(rebalance.head.partitionExpressions.flatMap(_.references.map(_.name)) - .contains("d")) - assert(rebalance.head.partitionExpressions.flatMap(_.references.map(_.name)) - .contains("c1")) - - assert(rebalance(1).partitionExpressions.flatMap(_.references.map(_.name)) - .contains("d")) - assert(!rebalance(1).partitionExpressions.flatMap(_.references.map(_.name)) - .contains("c1")) - } - } - } -} - -trait ParserSuite { self: ZorderSuiteBase => - override def createParser: ParserInterface = { - new SparkKyuubiSparkSQLParser(spark.sessionState.sqlParser) - } -} - -class ZorderWithCodegenEnabledSuite - extends ZorderWithCodegenEnabledSuiteBase - with ZorderSuiteSpark - with ParserSuite {} - -class ZorderWithCodegenDisabledSuite - extends ZorderWithCodegenDisabledSuiteBase - with ZorderSuiteSpark - with ParserSuite {} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala deleted file mode 100644 index 99f2c2cafe8..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala +++ /dev/null @@ -1,834 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} -import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar} -import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, EqualTo, Expression, ExpressionEvalHelper, Literal, NullsLast, SortOrder} -import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, OneRowRelation, Project, Sort} -import org.apache.spark.sql.execution.command.CreateDataSourceTableAsSelectCommand -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.hive.execution.{CreateHiveTableAsSelectCommand, InsertIntoHiveTable, OptimizedCreateHiveTableAsSelectCommand} -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.types._ - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} -import org.apache.kyuubi.sql.zorder.{OptimizeZorderCommandBase, OptimizeZorderStatement, Zorder, ZorderBytesUtils} - -trait ZorderSuiteBase extends KyuubiSparkSQLExtensionTest with ExpressionEvalHelper { - override def sparkConf(): SparkConf = { - super.sparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLCommonExtension") - } - - test("optimize unpartitioned table") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("up") { - sql(s"DROP TABLE IF EXISTS up") - - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - sql(s"CREATE TABLE up (c1 INT, c2 INT, c3 INT)") - sql(s"INSERT INTO TABLE up VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - val e = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE up WHERE c1 > 1 ZORDER BY c1, c2") - } - assert(e.getMessage == "Filters are only supported for partitioned table") - - sql("OPTIMIZE up ZORDER BY c1, c2") - val res = sql("SELECT c1, c2 FROM up").collect() - - assert(res.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r = res(i) - assert(t(0) == r.getInt(0)) - assert(t(1) == r.getInt(1)) - } - } - } - } - - test("optimize partitioned table") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("p") { - sql("DROP TABLE IF EXISTS p") - - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - - sql(s"CREATE TABLE p (c1 INT, c2 INT, c3 INT) PARTITIONED BY (id INT)") - sql(s"ALTER TABLE p ADD PARTITION (id = 1)") - sql(s"ALTER TABLE p ADD PARTITION (id = 2)") - sql(s"INSERT INTO TABLE p PARTITION (id = 1) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - sql(s"INSERT INTO TABLE p PARTITION (id = 2) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - sql(s"OPTIMIZE p ZORDER BY c1, c2") - - val res1 = sql(s"SELECT c1, c2 FROM p WHERE id = 1").collect() - val res2 = sql(s"SELECT c1, c2 FROM p WHERE id = 2").collect() - - assert(res1.length == 16) - assert(res2.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r1 = res1(i) - assert(t(0) == r1.getInt(0)) - assert(t(1) == r1.getInt(1)) - - val r2 = res2(i) - assert(t(0) == r2.getInt(0)) - assert(t(1) == r2.getInt(1)) - } - } - } - } - - test("optimize partitioned table with filters") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("p") { - sql("DROP TABLE IF EXISTS p") - - val target1 = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - val target2 = Seq( - Seq(0, 0), - Seq(0, 1), - Seq(0, 2), - Seq(0, 3), - Seq(1, 0), - Seq(1, 1), - Seq(1, 2), - Seq(1, 3), - Seq(2, 0), - Seq(2, 1), - Seq(2, 2), - Seq(2, 3), - Seq(3, 0), - Seq(3, 1), - Seq(3, 2), - Seq(3, 3)) - sql(s"CREATE TABLE p (c1 INT, c2 INT, c3 INT) PARTITIONED BY (id INT)") - sql(s"ALTER TABLE p ADD PARTITION (id = 1)") - sql(s"ALTER TABLE p ADD PARTITION (id = 2)") - sql(s"INSERT INTO TABLE p PARTITION (id = 1) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - sql(s"INSERT INTO TABLE p PARTITION (id = 2) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - val e = intercept[KyuubiSQLExtensionException]( - sql(s"OPTIMIZE p WHERE id = 1 AND c1 > 1 ZORDER BY c1, c2")) - assert(e.getMessage == "Only partition column filters are allowed") - - sql(s"OPTIMIZE p WHERE id = 1 ZORDER BY c1, c2") - - val res1 = sql(s"SELECT c1, c2 FROM p WHERE id = 1").collect() - val res2 = sql(s"SELECT c1, c2 FROM p WHERE id = 2").collect() - - assert(res1.length == 16) - assert(res2.length == 16) - - for (i <- target1.indices) { - val t1 = target1(i) - val r1 = res1(i) - assert(t1(0) == r1.getInt(0)) - assert(t1(1) == r1.getInt(1)) - - val t2 = target2(i) - val r2 = res2(i) - assert(t2(0) == r2.getInt(0)) - assert(t2(1) == r2.getInt(1)) - } - } - } - } - - test("optimize zorder with datasource table") { - // TODO remove this if we support datasource table - withTable("t") { - sql("CREATE TABLE t (c1 int, c2 int) USING PARQUET") - val msg = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE t ZORDER BY c1, c2") - }.getMessage - assert(msg.contains("only support hive table")) - } - } - - private def checkZorderTable( - enabled: Boolean, - cols: String, - planHasRepartition: Boolean, - resHasSort: Boolean): Unit = { - def checkSort(plan: LogicalPlan): Unit = { - assert(plan.isInstanceOf[Sort] === resHasSort) - plan match { - case sort: Sort => - val colArr = cols.split(",") - val refs = - if (colArr.length == 1) { - sort.order.head - .child.asInstanceOf[AttributeReference] :: Nil - } else { - sort.order.head - .child.asInstanceOf[Zorder].children.map(_.references.head) - } - assert(refs.size === colArr.size) - refs.zip(colArr).foreach { case (ref, col) => - assert(ref.name === col.trim) - } - case _ => - } - } - - val repartition = - if (planHasRepartition) { - "/*+ repartition */" - } else { - "" - } - withSQLConf("spark.sql.shuffle.partitions" -> "1") { - // hive - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> "false") { - withTable("zorder_t1", "zorder_t2_true", "zorder_t2_false") { - sql( - s""" - |CREATE TABLE zorder_t1 (c1 int, c2 string, c3 long, c4 double) STORED AS PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - |""".stripMargin) - val df1 = sql(s""" - |INSERT INTO TABLE zorder_t1 - |SELECT $repartition * FROM VALUES(1,'a',2,4D),(2,'b',3,6D) - |""".stripMargin) - assert(df1.queryExecution.analyzed.isInstanceOf[InsertIntoHiveTable]) - checkSort(df1.queryExecution.analyzed.children.head) - - Seq("true", "false").foreach { optimized => - withSQLConf( - "spark.sql.hive.convertMetastoreCtas" -> optimized, - "spark.sql.hive.convertMetastoreParquet" -> optimized) { - val df2 = - sql( - s""" - |CREATE TABLE zorder_t2_$optimized STORED AS PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - | - |SELECT $repartition * FROM - |VALUES(1,'a',2,4D),(2,'b',3,6D) AS t(c1 ,c2 , c3, c4) - |""".stripMargin) - if (optimized.toBoolean) { - assert(df2.queryExecution.analyzed - .isInstanceOf[OptimizedCreateHiveTableAsSelectCommand]) - } else { - assert(df2.queryExecution.analyzed.isInstanceOf[CreateHiveTableAsSelectCommand]) - } - checkSort(df2.queryExecution.analyzed.children.head) - } - } - } - } - - // datasource - withTable("zorder_t3", "zorder_t4") { - sql( - s""" - |CREATE TABLE zorder_t3 (c1 int, c2 string, c3 long, c4 double) USING PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - |""".stripMargin) - val df1 = sql(s""" - |INSERT INTO TABLE zorder_t3 - |SELECT $repartition * FROM VALUES(1,'a',2,4D),(2,'b',3,6D) - |""".stripMargin) - assert(df1.queryExecution.analyzed.isInstanceOf[InsertIntoHadoopFsRelationCommand]) - checkSort(df1.queryExecution.analyzed.children.head) - - val df2 = - sql( - s""" - |CREATE TABLE zorder_t4 USING PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - | - |SELECT $repartition * FROM - |VALUES(1,'a',2,4D),(2,'b',3,6D) AS t(c1 ,c2 , c3, c4) - |""".stripMargin) - assert(df2.queryExecution.analyzed.isInstanceOf[CreateDataSourceTableAsSelectCommand]) - checkSort(df2.queryExecution.analyzed.children.head) - } - } - } - - test("Support insert zorder by table properties") { - withSQLConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING.key -> "false") { - checkZorderTable(true, "c1", false, false) - checkZorderTable(false, "c1", false, false) - } - withSQLConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING.key -> "true") { - checkZorderTable(true, "", false, false) - checkZorderTable(true, "c5", false, false) - checkZorderTable(true, "c1,c5", false, false) - checkZorderTable(false, "c3", false, false) - checkZorderTable(true, "c3", true, false) - checkZorderTable(true, "c3", false, true) - checkZorderTable(true, "c2,c4", false, true) - checkZorderTable(true, "c4, c2, c1, c3", false, true) - } - } - - test("zorder: check unsupported data type") { - def checkZorderPlan(zorder: Expression): Unit = { - val msg = intercept[AnalysisException] { - val plan = Project(Seq(Alias(zorder, "c")()), OneRowRelation()) - spark.sessionState.analyzer.checkAnalysis(plan) - }.getMessage - assert(msg.contains("Unsupported z-order type: void")) - } - - checkZorderPlan(Zorder(Seq(Literal(null, NullType)))) - checkZorderPlan(Zorder(Seq(Literal(1, IntegerType), Literal(null, NullType)))) - } - - test("zorder: check supported data type") { - val children = Seq( - Literal.create(false, BooleanType), - Literal.create(null, BooleanType), - Literal.create(1.toByte, ByteType), - Literal.create(null, ByteType), - Literal.create(1.toShort, ShortType), - Literal.create(null, ShortType), - Literal.create(1, IntegerType), - Literal.create(null, IntegerType), - Literal.create(1L, LongType), - Literal.create(null, LongType), - Literal.create(1f, FloatType), - Literal.create(null, FloatType), - Literal.create(1d, DoubleType), - Literal.create(null, DoubleType), - Literal.create("1", StringType), - Literal.create(null, StringType), - Literal.create(1L, TimestampType), - Literal.create(null, TimestampType), - Literal.create(1, DateType), - Literal.create(null, DateType), - Literal.create(BigDecimal(1, 1), DecimalType(1, 1)), - Literal.create(null, DecimalType(1, 1))) - val zorder = Zorder(children) - val plan = Project(Seq(Alias(zorder, "c")()), OneRowRelation()) - spark.sessionState.analyzer.checkAnalysis(plan) - assert(zorder.foldable) - -// // scalastyle:off -// val resultGen = org.apache.commons.codec.binary.Hex.encodeHex( -// zorder.eval(InternalRow.fromSeq(children)).asInstanceOf[Array[Byte]], false) -// resultGen.grouped(2).zipWithIndex.foreach { case (char, i) => -// print("0x" + char(0) + char(1) + ", ") -// if ((i + 1) % 10 == 0) { -// println() -// } -// } -// // scalastyle:on - - val expected = Array( - 0xFB, 0xEA, 0xAA, 0xBA, 0xAE, 0xAB, 0xAA, 0xEA, 0xBA, 0xAE, 0xAB, 0xAA, 0xEA, 0xBA, 0xA6, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xBA, 0xBB, 0xAA, 0xAA, 0xAA, - 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0x9A, 0xAA, 0xAA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xEA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, - 0xAA, 0xAA, 0xBE, 0xAA, 0xAA, 0x8A, 0xBA, 0xAA, 0x2A, 0xEA, 0xA8, 0xAA, 0xAA, 0xA2, 0xAA, - 0xAA, 0x8A, 0xAA, 0xAA, 0x2F, 0xEB, 0xFE) - .map(_.toByte) - checkEvaluation(zorder, expected, InternalRow.fromSeq(children)) - } - - private def checkSort(input: DataFrame, expected: Seq[Row], dataType: Array[DataType]): Unit = { - withTempDir { dir => - input.repartition(3).write.mode("overwrite").format("parquet").save(dir.getCanonicalPath) - val df = spark.read.format("parquet") - .load(dir.getCanonicalPath) - .repartition(1) - assert(df.schema.fields.map(_.dataType).sameElements(dataType)) - val exprs = Seq("c1", "c2").map(col).map(_.expr) - val sortOrder = SortOrder(Zorder(exprs), Ascending, NullsLast, Seq.empty) - val zorderSort = Sort(Seq(sortOrder), true, df.logicalPlan) - val result = Dataset.ofRows(spark, zorderSort) - checkAnswer(result, expected) - } - } - - test("sort with zorder -- boolean column") { - val schema = StructType(StructField("c1", BooleanType) :: StructField("c2", BooleanType) :: Nil) - val nonNullDF = spark.createDataFrame( - spark.sparkContext.parallelize( - Seq(Row(false, false), Row(false, true), Row(true, false), Row(true, true))), - schema) - val expected = - Row(false, false) :: Row(true, false) :: Row(false, true) :: Row(true, true) :: Nil - checkSort(nonNullDF, expected, Array(BooleanType, BooleanType)) - val df = spark.createDataFrame( - spark.sparkContext.parallelize( - Seq(Row(false, false), Row(false, null), Row(null, false), Row(null, null))), - schema) - val expected2 = - Row(false, false) :: Row(null, false) :: Row(false, null) :: Row(null, null) :: Nil - checkSort(df, expected2, Array(BooleanType, BooleanType)) - } - - test("sort with zorder -- int column") { - // TODO: add more datatype unit test - val session = spark - import session.implicits._ - // generate 4 * 4 matrix - val len = 3 - val input = spark.range(len + 1).selectExpr("cast(id as int) as c1") - .select($"c1", explode(sequence(lit(0), lit(len))) as "c2") - val expected = - Row(0, 0) :: Row(1, 0) :: Row(0, 1) :: Row(1, 1) :: - Row(2, 0) :: Row(3, 0) :: Row(2, 1) :: Row(3, 1) :: - Row(0, 2) :: Row(1, 2) :: Row(0, 3) :: Row(1, 3) :: - Row(2, 2) :: Row(3, 2) :: Row(2, 3) :: Row(3, 3) :: Nil - checkSort(input, expected, Array(IntegerType, IntegerType)) - - // contains null value case. - val nullDF = spark.range(1).selectExpr("cast(null as int) as c1") - val input2 = spark.range(len).selectExpr("cast(id as int) as c1") - .union(nullDF) - .select( - $"c1", - explode(concat(sequence(lit(0), lit(len - 1)), array(lit(null)))) as "c2") - val expected2 = Row(0, 0) :: Row(1, 0) :: Row(0, 1) :: Row(1, 1) :: - Row(2, 0) :: Row(2, 1) :: Row(0, 2) :: Row(1, 2) :: - Row(2, 2) :: Row(null, 0) :: Row(null, 1) :: Row(null, 2) :: - Row(0, null) :: Row(1, null) :: Row(2, null) :: Row(null, null) :: Nil - checkSort(input2, expected2, Array(IntegerType, IntegerType)) - } - - test("sort with zorder -- string column") { - val schema = StructType(StructField("c1", StringType) :: StructField("c2", StringType) :: Nil) - val rdd = spark.sparkContext.parallelize(Seq( - Row("a", "a"), - Row("a", "b"), - Row("a", "c"), - Row("a", "d"), - Row("b", "a"), - Row("b", "b"), - Row("b", "c"), - Row("b", "d"), - Row("c", "a"), - Row("c", "b"), - Row("c", "c"), - Row("c", "d"), - Row("d", "a"), - Row("d", "b"), - Row("d", "c"), - Row("d", "d"))) - val input = spark.createDataFrame(rdd, schema) - val expected = Row("a", "a") :: Row("b", "a") :: Row("c", "a") :: Row("a", "b") :: - Row("a", "c") :: Row("b", "b") :: Row("c", "b") :: Row("b", "c") :: - Row("c", "c") :: Row("d", "a") :: Row("d", "b") :: Row("d", "c") :: - Row("a", "d") :: Row("b", "d") :: Row("c", "d") :: Row("d", "d") :: Nil - checkSort(input, expected, Array(StringType, StringType)) - - val rdd2 = spark.sparkContext.parallelize(Seq( - Row(null, "a"), - Row("a", "b"), - Row("a", "c"), - Row("a", null), - Row("b", "a"), - Row(null, "b"), - Row("b", null), - Row("b", "d"), - Row("c", "a"), - Row("c", null), - Row(null, "c"), - Row("c", "d"), - Row("d", null), - Row("d", "b"), - Row("d", "c"), - Row(null, "d"), - Row(null, null))) - val input2 = spark.createDataFrame(rdd2, schema) - val expected2 = Row("b", "a") :: Row("c", "a") :: Row("a", "b") :: Row("a", "c") :: - Row("d", "b") :: Row("d", "c") :: Row("b", "d") :: Row("c", "d") :: - Row(null, "a") :: Row(null, "b") :: Row(null, "c") :: Row(null, "d") :: - Row("a", null) :: Row("b", null) :: Row("c", null) :: Row("d", null) :: - Row(null, null) :: Nil - checkSort(input2, expected2, Array(StringType, StringType)) - } - - test("test special value of short int long type") { - val df1 = spark.createDataFrame(Seq( - (-1, -1L), - (Int.MinValue, Int.MinValue.toLong), - (1, 1L), - (Int.MaxValue - 1, Int.MaxValue.toLong), - (Int.MaxValue - 1, Int.MaxValue.toLong - 1), - (Int.MaxValue, Int.MaxValue.toLong + 1), - (Int.MaxValue, Int.MaxValue.toLong))).toDF("c1", "c2") - val expected1 = - Row(Int.MinValue, Int.MinValue.toLong) :: - Row(-1, -1L) :: - Row(1, 1L) :: - Row(Int.MaxValue - 1, Int.MaxValue.toLong - 1) :: - Row(Int.MaxValue - 1, Int.MaxValue.toLong) :: - Row(Int.MaxValue, Int.MaxValue.toLong) :: - Row(Int.MaxValue, Int.MaxValue.toLong + 1) :: Nil - checkSort(df1, expected1, Array(IntegerType, LongType)) - - val df2 = spark.createDataFrame(Seq( - (-1, -1.toShort), - (Short.MinValue.toInt, Short.MinValue), - (1, 1.toShort), - (Short.MaxValue.toInt, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toInt + 1, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toInt, Short.MaxValue), - (Short.MaxValue.toInt + 1, Short.MaxValue))).toDF("c1", "c2") - val expected2 = - Row(Short.MinValue.toInt, Short.MinValue) :: - Row(-1, -1.toShort) :: - Row(1, 1.toShort) :: - Row(Short.MaxValue.toInt, Short.MaxValue - 1) :: - Row(Short.MaxValue.toInt, Short.MaxValue) :: - Row(Short.MaxValue.toInt + 1, Short.MaxValue - 1) :: - Row(Short.MaxValue.toInt + 1, Short.MaxValue) :: Nil - checkSort(df2, expected2, Array(IntegerType, ShortType)) - - val df3 = spark.createDataFrame(Seq( - (-1L, -1.toShort), - (Short.MinValue.toLong, Short.MinValue), - (1L, 1.toShort), - (Short.MaxValue.toLong, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toLong + 1, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toLong, Short.MaxValue), - (Short.MaxValue.toLong + 1, Short.MaxValue))).toDF("c1", "c2") - val expected3 = - Row(Short.MinValue.toLong, Short.MinValue) :: - Row(-1L, -1.toShort) :: - Row(1L, 1.toShort) :: - Row(Short.MaxValue.toLong, Short.MaxValue - 1) :: - Row(Short.MaxValue.toLong, Short.MaxValue) :: - Row(Short.MaxValue.toLong + 1, Short.MaxValue - 1) :: - Row(Short.MaxValue.toLong + 1, Short.MaxValue) :: Nil - checkSort(df3, expected3, Array(LongType, ShortType)) - } - - test("skip zorder if only requires one column") { - withTable("t") { - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> "false") { - sql("CREATE TABLE t (c1 int, c2 string) stored as parquet") - val order1 = sql("OPTIMIZE t ZORDER BY c1").queryExecution.analyzed - .asInstanceOf[OptimizeZorderCommandBase].query.asInstanceOf[Sort].order.head.child - assert(!order1.isInstanceOf[Zorder]) - assert(order1.isInstanceOf[AttributeReference]) - } - } - } - - test("Add config to control if zorder using global sort") { - withTable("t") { - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,c2') - |""".stripMargin) - val p1 = sql("OPTIMIZE t ZORDER BY c1, c2").queryExecution.analyzed - assert(p1.collect { - case shuffle: Sort if !shuffle.global => shuffle - }.size == 1) - - val p2 = sql("INSERT INTO TABLE t SELECT * FROM VALUES(1,'a')").queryExecution.analyzed - assert(p2.collect { - case shuffle: Sort if !shuffle.global => shuffle - }.size == 1) - } - } - } - - test("Allow insert zorder after repartition if zorder using local sort") { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,c2') - |""".stripMargin) - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") { - val p1 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p1.collect { - case sort: Sort if !sort.global => sort - }.size == 1) - } - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "true") { - val p2 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p2.collect { - case sort: Sort if !sort.global => sort - }.size == 0) - } - } - } - - test("fast approach test") { - Seq[Seq[Any]]( - Seq(1L, 2L), - Seq(1L, 2L, 3L), - Seq(1L, 2L, 3L, 4L), - Seq(1L, 2L, 3L, 4L, 5L), - Seq(1L, 2L, 3L, 4L, 5L, 6L), - Seq(1L, 2L, 3L, 4L, 5L, 6L, 7L), - Seq(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L)) - .foreach { inputs => - assert(java.util.Arrays.equals( - ZorderBytesUtils.interleaveBits(inputs.toArray), - ZorderBytesUtils.interleaveBitsDefault(inputs.map(ZorderBytesUtils.toByteArray).toArray))) - } - } - - test("OPTIMIZE command is parsed as expected") { - val parser = createParser - val globalSort = spark.conf.get(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) - - assert(parser.parsePlan("OPTIMIZE p zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project(Seq(UnresolvedStar(None)), UnresolvedRelation(TableIdentifier("p")))))) - - assert(parser.parsePlan("OPTIMIZE p zorder by c1, c2") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder( - Zorder(Seq(UnresolvedAttribute("c1"), UnresolvedAttribute("c2"))), - Ascending, - NullsLast, - Seq.empty) :: Nil, - globalSort, - Project(Seq(UnresolvedStar(None)), UnresolvedRelation(TableIdentifier("p")))))) - - assert(parser.parsePlan("OPTIMIZE p where id = 1 zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo(UnresolvedAttribute("id"), Literal(1)), - UnresolvedRelation(TableIdentifier("p"))))))) - - assert(parser.parsePlan("OPTIMIZE p where id = 1 zorder by c1, c2") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder( - Zorder(Seq(UnresolvedAttribute("c1"), UnresolvedAttribute("c2"))), - Ascending, - NullsLast, - Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo(UnresolvedAttribute("id"), Literal(1)), - UnresolvedRelation(TableIdentifier("p"))))))) - - assert(parser.parsePlan("OPTIMIZE p where id = current_date() zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo( - UnresolvedAttribute("id"), - UnresolvedFunction("current_date", Seq.empty, false)), - UnresolvedRelation(TableIdentifier("p"))))))) - - // TODO: add following case support - intercept[ParseException] { - parser.parsePlan("OPTIMIZE p zorder by (c1)") - } - - intercept[ParseException] { - parser.parsePlan("OPTIMIZE p zorder by (c1, c2)") - } - } - - test("OPTIMIZE partition predicates constraint") { - withTable("p") { - sql("CREATE TABLE p (c1 INT, c2 INT) PARTITIONED BY (event_date DATE)") - val e1 = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE p WHERE event_date = current_date as c ZORDER BY c1, c2") - } - assert(e1.getMessage.contains("unsupported partition predicates")) - - val e2 = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE p WHERE c1 = 1 ZORDER BY c1, c2") - } - assert(e2.getMessage == "Only partition column filters are allowed") - } - } - - test("optimize sort by backquoted column name") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("up") { - sql(s"DROP TABLE IF EXISTS up") - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - sql(s"CREATE TABLE up (c1 INT, `@c2` INT, c3 INT)") - sql(s"INSERT INTO TABLE up VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - sql("OPTIMIZE up ZORDER BY c1, `@c2`") - val res = sql("SELECT c1, `@c2` FROM up").collect() - - assert(res.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r = res(i) - assert(t(0) == r.getInt(0)) - assert(t(1) == r.getInt(1)) - } - } - } - } - - def createParser: ParserInterface -} - -trait ZorderWithCodegenEnabledSuiteBase extends ZorderSuiteBase { - override def sparkConf(): SparkConf = { - val conf = super.sparkConf - conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") - conf - } -} - -trait ZorderWithCodegenDisabledSuiteBase extends ZorderSuiteBase { - override def sparkConf(): SparkConf = { - val conf = super.sparkConf - conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "false") - conf.set(SQLConf.CODEGEN_FACTORY_MODE.key, "NO_CODEGEN") - conf - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala deleted file mode 100644 index b891a7224a0..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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.spark.sql.benchmark - -import java.io.{File, FileOutputStream, OutputStream} - -import scala.collection.JavaConverters._ - -import com.google.common.reflect.ClassPath -import org.scalatest.Assertions._ - -trait KyuubiBenchmarkBase { - var output: Option[OutputStream] = None - - private val prefix = { - val benchmarkClasses = ClassPath.from(Thread.currentThread.getContextClassLoader) - .getTopLevelClassesRecursive("org.apache.spark.sql").asScala.toArray - assert(benchmarkClasses.nonEmpty) - val benchmark = benchmarkClasses.find(_.load().getName.endsWith("Benchmark")) - val targetDirOrProjDir = - new File(benchmark.get.load().getProtectionDomain.getCodeSource.getLocation.toURI) - .getParentFile.getParentFile - if (targetDirOrProjDir.getName == "target") { - targetDirOrProjDir.getParentFile.getCanonicalPath + "/" - } else { - targetDirOrProjDir.getCanonicalPath + "/" - } - } - - def withHeader(func: => Unit): Unit = { - val version = System.getProperty("java.version").split("\\D+")(0).toInt - val jdkString = if (version > 8) s"-jdk$version" else "" - val resultFileName = - s"${this.getClass.getSimpleName.replace("$", "")}$jdkString-results.txt" - val dir = new File(s"${prefix}benchmarks/") - if (!dir.exists()) { - // scalastyle:off println - println(s"Creating ${dir.getAbsolutePath} for benchmark results.") - // scalastyle:on println - dir.mkdirs() - } - val file = new File(dir, resultFileName) - if (!file.exists()) { - file.createNewFile() - } - output = Some(new FileOutputStream(file)) - - func - - output.foreach { o => - if (o != null) { - o.close() - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/pom.xml b/extensions/spark/kyuubi-extension-spark-3-4/pom.xml deleted file mode 100644 index d14ce39bbb3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/pom.xml +++ /dev/null @@ -1,222 +0,0 @@ - - - - 4.0.0 - - org.apache.kyuubi - kyuubi-parent - 1.13.0-SNAPSHOT - ../../../pom.xml - - - kyuubi-extension-spark-3-4_${scala.binary.version} - jar - Kyuubi Dev Spark Extensions (for Spark 3.4) - https://kyuubi.apache.org/ - - - - org.scala-lang - scala-library - provided - - - - org.apache.spark - spark-sql_${scala.binary.version} - provided - - - - org.apache.spark - spark-hive_${scala.binary.version} - provided - - - - org.apache.hadoop - hadoop-client-api - provided - - - - org.apache.kyuubi - kyuubi-download - ${project.version} - pom - test - - - - org.apache.kyuubi - kyuubi-util-scala_${scala.binary.version} - ${project.version} - test-jar - test - - - - org.apache.spark - spark-core_${scala.binary.version} - test-jar - test - - - - org.apache.spark - spark-catalyst_${scala.binary.version} - test-jar - test - - - - org.scalatestplus - scalacheck-1-17_${scala.binary.version} - test - - - - org.apache.spark - spark-sql_${scala.binary.version} - ${spark.version} - test-jar - test - - - - org.apache.hadoop - hadoop-client-runtime - test - - - - javax.servlet - javax.servlet-api - test - - - - jakarta.servlet - jakarta.servlet-api - test - - - - org.apache.logging.log4j - log4j-slf4j-impl - test - - - - - - - - org.codehaus.mojo - build-helper-maven-plugin - - - regex-property - - regex-property - - - spark.home - ${project.basedir}/../../../externals/kyuubi-download/target/${spark.archive.name} - (.+)\.tgz - $1 - - - - - - org.antlr - antlr4-maven-plugin - - true - ${project.basedir}/src/main/antlr4 - - - - - org.apache.maven.plugins - maven-shade-plugin - - false - - - org.apache.kyuubi:* - - - - - - - shade - - package - - - - - org.scalatest - scalatest-maven-plugin - - - - ${scala.binary.version} - - - - - target/scala-${scala.binary.version}/classes - target/scala-${scala.binary.version}/test-classes - - - - - spark-home-from-archive - - - !env.SPARK_HOME - - - - - - org.scalatest - scalatest-maven-plugin - - - ${spark.home} - - - - - - - - diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 b/extensions/spark/kyuubi-extension-spark-3-4/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 deleted file mode 100644 index e52b7f5cfeb..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSparkSQL.g4 +++ /dev/null @@ -1,191 +0,0 @@ -/* - * 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. - */ - -grammar KyuubiSparkSQL; - -@members { - /** - * Verify whether current token is a valid decimal token (which contains dot). - * Returns true if the character that follows the token is not a digit or letter or underscore. - * - * For example: - * For char stream "2.3", "2." is not a valid decimal token, because it is followed by digit '3'. - * For char stream "2.3_", "2.3" is not a valid decimal token, because it is followed by '_'. - * For char stream "2.3W", "2.3" is not a valid decimal token, because it is followed by 'W'. - * For char stream "12.0D 34.E2+0.12 " 12.0D is a valid decimal token because it is followed - * by a space. 34.E2 is a valid decimal token because it is followed by symbol '+' - * which is not a digit or letter or underscore. - */ - public boolean isValidDecimal() { - int nextChar = _input.LA(1); - if (nextChar >= 'A' && nextChar <= 'Z' || nextChar >= '0' && nextChar <= '9' || - nextChar == '_') { - return false; - } else { - return true; - } - } - } - -tokens { - DELIMITER -} - -singleStatement - : statement EOF - ; - -statement - : OPTIMIZE multipartIdentifier whereClause? zorderClause #optimizeZorder - | .*? #passThrough - ; - -whereClause - : WHERE partitionPredicate = predicateToken - ; - -zorderClause - : ZORDER BY order+=multipartIdentifier (',' order+=multipartIdentifier)* - ; - -// We don't have an expression rule in our grammar here, so we just grab the tokens and defer -// parsing them to later. -predicateToken - : .+? - ; - -multipartIdentifier - : parts+=identifier ('.' parts+=identifier)* - ; - -identifier - : strictIdentifier - ; - -strictIdentifier - : IDENTIFIER #unquotedIdentifier - | quotedIdentifier #quotedIdentifierAlternative - | nonReserved #unquotedIdentifier - ; - -quotedIdentifier - : BACKQUOTED_IDENTIFIER - ; - -nonReserved - : AND - | BY - | FALSE - | DATE - | INTERVAL - | OPTIMIZE - | OR - | TABLE - | TIMESTAMP - | TRUE - | WHERE - | ZORDER - ; - -AND: 'AND'; -BY: 'BY'; -FALSE: 'FALSE'; -DATE: 'DATE'; -INTERVAL: 'INTERVAL'; -OPTIMIZE: 'OPTIMIZE'; -OR: 'OR'; -TABLE: 'TABLE'; -TIMESTAMP: 'TIMESTAMP'; -TRUE: 'TRUE'; -WHERE: 'WHERE'; -ZORDER: 'ZORDER'; - -MINUS: '-'; - -BIGINT_LITERAL - : DIGIT+ 'L' - ; - -SMALLINT_LITERAL - : DIGIT+ 'S' - ; - -TINYINT_LITERAL - : DIGIT+ 'Y' - ; - -INTEGER_VALUE - : DIGIT+ - ; - -DECIMAL_VALUE - : DIGIT+ EXPONENT - | DECIMAL_DIGITS EXPONENT? {isValidDecimal()}? - ; - -DOUBLE_LITERAL - : DIGIT+ EXPONENT? 'D' - | DECIMAL_DIGITS EXPONENT? 'D' {isValidDecimal()}? - ; - -BIGDECIMAL_LITERAL - : DIGIT+ EXPONENT? 'BD' - | DECIMAL_DIGITS EXPONENT? 'BD' {isValidDecimal()}? - ; - -BACKQUOTED_IDENTIFIER - : '`' ( ~'`' | '``' )* '`' - ; - -IDENTIFIER - : (LETTER | DIGIT | '_')+ - ; - -fragment DECIMAL_DIGITS - : DIGIT+ '.' DIGIT* - | '.' DIGIT+ - ; - -fragment EXPONENT - : 'E' [+-]? DIGIT+ - ; - -fragment DIGIT - : [0-9] - ; - -fragment LETTER - : [A-Z] - ; - -SIMPLE_COMMENT - : '--' ~[\r\n]* '\r'? '\n'? -> channel(HIDDEN) - ; - -BRACKETED_COMMENT - : '/*' .*? '*/' -> channel(HIDDEN) - ; - -WS : [ \r\n\t]+ -> channel(HIDDEN) - ; - -// Catch-all for anything we can't recognize. -// We use this to be able to ignore and recover all the text -// when splitting statements with DelimiterLexer -UNRECOGNIZED - : . - ; diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala deleted file mode 100644 index e33632b8b30..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/DropIgnoreNonexistent.scala +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunctionName, UnresolvedRelation} -import org.apache.spark.sql.catalyst.plans.logical.{DropFunction, DropNamespace, LogicalPlan, NoopCommand, UncacheTable} -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.command.{AlterTableDropPartitionCommand, DropTableCommand} - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -case class DropIgnoreNonexistent(session: SparkSession) extends Rule[LogicalPlan] { - - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(DROP_IGNORE_NONEXISTENT)) { - plan match { - case i @ AlterTableDropPartitionCommand(_, _, false, _, _) => - i.copy(ifExists = true) - case i @ DropTableCommand(_, false, _, _) => - i.copy(ifExists = true) - case i @ DropNamespace(_, false, _) => - i.copy(ifExists = true) - case UncacheTable(u: UnresolvedRelation, false, _) => - NoopCommand("UNCACHE TABLE", u.multipartIdentifier) - case DropFunction(u: UnresolvedFunctionName, false) => - NoopCommand("DROP FUNCTION", u.multipartIdentifier) - case _ => plan - } - } else { - plan - } - } - -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala deleted file mode 100644 index fcbf5c0a122..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala +++ /dev/null @@ -1,110 +0,0 @@ -/* - * 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.sql - -import scala.annotation.tailrec - -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Expression, NamedExpression, UnaryExpression} -import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys -import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftAnti, LeftOuter, LeftSemi, RightOuter} -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, Project, Sort, SubqueryAlias, View} - -/** - * Infer the columns for Rebalance and Sort to improve the compression ratio. - * - * For example - * {{{ - * INSERT INTO TABLE t PARTITION(p='a') - * SELECT * FROM t1 JOIN t2 on t1.c1 = t2.c1 - * }}} - * the inferred columns are: t1.c1 - */ -object InferRebalanceAndSortOrders { - - type PartitioningAndOrdering = (Seq[Expression], Seq[Expression]) - - private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = { - @tailrec - def throughUnary(e: Expression): Expression = e match { - case u: UnaryExpression if u.deterministic => - throughUnary(u.child) - case _ => e - } - - named.flatMap { - case a @ Alias(child, _) => - Some((throughUnary(child).canonicalized, a.toAttribute)) - case _ => None - }.toMap - } - - def infer(plan: LogicalPlan): Option[PartitioningAndOrdering] = { - def candidateKeys( - input: LogicalPlan, - output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = { - input match { - case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, _, _, _, _, _) => - joinType match { - case LeftSemi | LeftAnti | LeftOuter => Some((leftKeys, leftKeys)) - case RightOuter => Some((rightKeys, rightKeys)) - case Inner | FullOuter => - if (output.isEmpty) { - Some((leftKeys ++ rightKeys, leftKeys ++ rightKeys)) - } else { - assert(leftKeys.length == rightKeys.length) - val keys = leftKeys.zip(rightKeys).flatMap { case (left, right) => - if (left.references.subsetOf(output)) { - Some(left) - } else if (right.references.subsetOf(output)) { - Some(right) - } else { - None - } - } - Some((keys, keys)) - } - case _ => None - } - case agg: Aggregate => - val aliasMap = getAliasMap(agg.aggregateExpressions) - Some(( - agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)), - agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o)))) - case s: Sort => Some((s.order.map(_.child), s.order.map(_.child))) - case p: Project => - val aliasMap = getAliasMap(p.projectList) - candidateKeys(p.child, p.references).map { case (partitioning, ordering) => - ( - partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)), - ordering.map(o => aliasMap.getOrElse(o.canonicalized, o))) - } - case f: Filter => candidateKeys(f.child, output) - case s: SubqueryAlias => candidateKeys(s.child, output) - case v: View => candidateKeys(v.child, output) - - case _ => None - } - } - - candidateKeys(plan).map { case (partitioning, ordering) => - ( - partitioning.filter(_.references.subsetOf(plan.outputSet)), - ordering.filter(_.references.subsetOf(plan.outputSet))) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala deleted file mode 100644 index 92626f02745..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/InsertShuffleNodeBeforeJoin.scala +++ /dev/null @@ -1,93 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.plans.physical.Distribution -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.QueryStageExec -import org.apache.spark.sql.execution.aggregate.BaseAggregateExec -import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * Insert shuffle node before join if it doesn't exist to make `OptimizeSkewedJoin` works. - */ -object InsertShuffleNodeBeforeJoin extends Rule[SparkPlan] { - - override def apply(plan: SparkPlan): SparkPlan = { - // this rule has no meaning without AQE - if (!conf.getConf(FORCE_SHUFFLE_BEFORE_JOIN) || - !conf.getConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)) { - return plan - } - - val newPlan = insertShuffleBeforeJoin(plan) - if (plan.fastEquals(newPlan)) { - plan - } else { - // make sure the output partitioning and ordering will not be broken. - KyuubiEnsureRequirements.apply(newPlan) - } - } - - // SPARK-33832 (Spark 3.3) moves the rule OptimizeSkewedJoin from queryStageOptimizerRules - // to queryStagePreparationRules, injecting shuffle after OptimizeSkewedJoin may produce - // invalid query plan. - private def insertShuffleBeforeJoin(plan: SparkPlan): SparkPlan = plan transformUp { - case smj @ SortMergeJoinExec(_, _, _, _, l, r, isSkewJoin) if !isSkewJoin => - smj.withNewChildren(checkAndInsertShuffle(smj.requiredChildDistribution.head, l) :: - checkAndInsertShuffle(smj.requiredChildDistribution(1), r) :: Nil) - - case shj: ShuffledHashJoinExec if !shj.isSkewJoin => - if (!shj.left.isInstanceOf[Exchange] && !shj.right.isInstanceOf[Exchange]) { - shj.withNewChildren(withShuffleExec(shj.requiredChildDistribution.head, shj.left) :: - withShuffleExec(shj.requiredChildDistribution(1), shj.right) :: Nil) - } else if (!shj.left.isInstanceOf[Exchange]) { - shj.withNewChildren( - withShuffleExec(shj.requiredChildDistribution.head, shj.left) :: shj.right :: Nil) - } else if (!shj.right.isInstanceOf[Exchange]) { - shj.withNewChildren( - shj.left :: withShuffleExec(shj.requiredChildDistribution(1), shj.right) :: Nil) - } else { - shj - } - } - - private def checkAndInsertShuffle( - distribution: Distribution, - child: SparkPlan): SparkPlan = child match { - case SortExec(_, _, _: Exchange, _) => - child - case SortExec(_, _, _: QueryStageExec, _) => - child - case sort @ SortExec(_, _, agg: BaseAggregateExec, _) => - sort.withNewChildren(withShuffleExec(distribution, agg) :: Nil) - case _ => - withShuffleExec(distribution, child) - } - - private def withShuffleExec(distribution: Distribution, child: SparkPlan): SparkPlan = { - val numPartitions = distribution.requiredNumPartitions - .getOrElse(conf.numShufflePartitions) - ShuffleExchangeExec(distribution.createPartitioning(numPartitions), child) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala deleted file mode 100644 index a17e0a4652b..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiEnsureRequirements.scala +++ /dev/null @@ -1,133 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.expressions.SortOrder -import org.apache.spark.sql.catalyst.plans.physical.{BroadcastDistribution, Distribution, UnspecifiedDistribution} -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} - -/** - * Copy from Apache Spark `EnsureRequirements` - * 1. remove reorder join predicates - * 2. remove shuffle pruning - */ -object KyuubiEnsureRequirements extends Rule[SparkPlan] { - private def ensureDistributionAndOrdering(operator: SparkPlan): SparkPlan = { - val requiredChildDistributions: Seq[Distribution] = operator.requiredChildDistribution - val requiredChildOrderings: Seq[Seq[SortOrder]] = operator.requiredChildOrdering - var children: Seq[SparkPlan] = operator.children - assert(requiredChildDistributions.length == children.length) - assert(requiredChildOrderings.length == children.length) - - // Ensure that the operator's children satisfy their output distribution requirements. - children = children.zip(requiredChildDistributions).map { - case (child, distribution) if child.outputPartitioning.satisfies(distribution) => - child - case (child, BroadcastDistribution(mode)) => - BroadcastExchangeExec(mode, child) - case (child, distribution) => - val numPartitions = distribution.requiredNumPartitions - .getOrElse(conf.numShufflePartitions) - ShuffleExchangeExec(distribution.createPartitioning(numPartitions), child) - } - - // Get the indexes of children which have specified distribution requirements and need to have - // same number of partitions. - val childrenIndexes = requiredChildDistributions.zipWithIndex.filter { - case (UnspecifiedDistribution, _) => false - case (_: BroadcastDistribution, _) => false - case _ => true - }.map(_._2) - - val childrenNumPartitions = - childrenIndexes.map(children(_).outputPartitioning.numPartitions).toSet - - if (childrenNumPartitions.size > 1) { - // Get the number of partitions which is explicitly required by the distributions. - val requiredNumPartitions = { - val numPartitionsSet = childrenIndexes.flatMap { - index => requiredChildDistributions(index).requiredNumPartitions - }.toSet - assert( - numPartitionsSet.size <= 1, - s"$operator have incompatible requirements of the number of partitions for its children") - numPartitionsSet.headOption - } - - // If there are non-shuffle children that satisfy the required distribution, we have - // some tradeoffs when picking the expected number of shuffle partitions: - // 1. We should avoid shuffling these children. - // 2. We should have a reasonable parallelism. - val nonShuffleChildrenNumPartitions = - childrenIndexes.map(children).filterNot(_.isInstanceOf[ShuffleExchangeExec]) - .map(_.outputPartitioning.numPartitions) - val expectedChildrenNumPartitions = - if (nonShuffleChildrenNumPartitions.nonEmpty) { - if (nonShuffleChildrenNumPartitions.length == childrenIndexes.length) { - // Here we pick the max number of partitions among these non-shuffle children. - nonShuffleChildrenNumPartitions.max - } else { - // Here we pick the max number of partitions among these non-shuffle children as the - // expected number of shuffle partitions. However, if it's smaller than - // `conf.numShufflePartitions`, we pick `conf.numShufflePartitions` as the - // expected number of shuffle partitions. - math.max(nonShuffleChildrenNumPartitions.max, conf.defaultNumShufflePartitions) - } - } else { - childrenNumPartitions.max - } - - val targetNumPartitions = requiredNumPartitions.getOrElse(expectedChildrenNumPartitions) - - children = children.zip(requiredChildDistributions).zipWithIndex.map { - case ((child, distribution), index) if childrenIndexes.contains(index) => - if (child.outputPartitioning.numPartitions == targetNumPartitions) { - child - } else { - val defaultPartitioning = distribution.createPartitioning(targetNumPartitions) - child match { - // If child is an exchange, we replace it with a new one having defaultPartitioning. - case ShuffleExchangeExec(_, c, _) => ShuffleExchangeExec(defaultPartitioning, c) - case _ => ShuffleExchangeExec(defaultPartitioning, child) - } - } - - case ((child, _), _) => child - } - } - - // Now that we've performed any necessary shuffles, add sorts to guarantee output orderings: - children = children.zip(requiredChildOrderings).map { case (child, requiredOrdering) => - // If child.outputOrdering already satisfies the requiredOrdering, we do not need to sort. - if (SortOrder.orderingSatisfies(child.outputOrdering, requiredOrdering)) { - child - } else { - SortExec(requiredOrdering, global = false, child = child) - } - } - - operator.withNewChildren(children) - } - - def apply(plan: SparkPlan): SparkPlan = plan.transformUp { - case operator: SparkPlan => - ensureDistributionAndOrdering(operator) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala deleted file mode 100644 index a7fcbecd422..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiQueryStagePreparation.scala +++ /dev/null @@ -1,194 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec -import org.apache.spark.sql.execution.command.{ResetCommand, SetCommand} -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * This rule split stage into two parts: - * 1. previous stage - * 2. final stage - * For final stage, we can inject extra config. It's useful if we use repartition to optimize - * small files that needs bigger shuffle partition size than previous. - * - * Let's say we have a query with 3 stages, then the logical machine like: - * - * Set/Reset Command -> cleanup previousStage config if user set the spark config. - * Query -> AQE -> stage1 -> preparation (use previousStage to overwrite spark config) - * -> AQE -> stage2 -> preparation (use spark config) - * -> AQE -> stage3 -> preparation (use finalStage config to overwrite spark config, - * store spark config to previousStage.) - * - * An example of the new finalStage config: - * `spark.sql.adaptive.advisoryPartitionSizeInBytes` -> - * `spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes` - */ -case class FinalStageConfigIsolation(session: SparkSession) extends Rule[SparkPlan] { - import FinalStageConfigIsolation._ - - override def apply(plan: SparkPlan): SparkPlan = { - // this rule has no meaning without AQE - if (!conf.getConf(FINAL_STAGE_CONFIG_ISOLATION) || - !conf.getConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)) { - return plan - } - - if (isFinalStage(plan)) { - // We can not get the whole plan at query preparation phase to detect if current plan is - // for writing, so we depend on a tag which is been injected at post resolution phase. - // Note: we should still do clean up previous config for non-final stage to avoid such case: - // the first statement is write, but the second statement is query. - if (conf.getConf(FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY) && - !WriteUtils.isWrite(session, plan)) { - return plan - } - - // set config for final stage - session.conf.getAll.filter(_._1.startsWith(FINAL_STAGE_CONFIG_PREFIX)).foreach { - case (k, v) => - val sparkConfigKey = s"spark.sql.${k.substring(FINAL_STAGE_CONFIG_PREFIX.length)}" - val previousStageConfigKey = - s"$PREVIOUS_STAGE_CONFIG_PREFIX${k.substring(FINAL_STAGE_CONFIG_PREFIX.length)}" - // store the previous config only if we have not stored, to avoid some query only - // have one stage that will overwrite real config. - if (!session.sessionState.conf.contains(previousStageConfigKey)) { - val originalValue = - if (session.conf.getOption(sparkConfigKey).isDefined) { - session.sessionState.conf.getConfString(sparkConfigKey) - } else { - // the default value of config is None, so we need to use a internal tag - INTERNAL_UNSET_CONFIG_TAG - } - logInfo(s"Store config: $sparkConfigKey to previousStage, " + - s"original value: $originalValue ") - session.sessionState.conf.setConfString(previousStageConfigKey, originalValue) - } - logInfo(s"For final stage: set $sparkConfigKey = $v.") - session.conf.set(sparkConfigKey, v) - } - } else { - // reset config for previous stage - session.conf.getAll.filter(_._1.startsWith(PREVIOUS_STAGE_CONFIG_PREFIX)).foreach { - case (k, v) => - val sparkConfigKey = s"spark.sql.${k.substring(PREVIOUS_STAGE_CONFIG_PREFIX.length)}" - logInfo(s"For previous stage: set $sparkConfigKey = $v.") - if (v == INTERNAL_UNSET_CONFIG_TAG) { - session.conf.unset(sparkConfigKey) - } else { - session.conf.set(sparkConfigKey, v) - } - // unset config so that we do not need to reset configs for every previous stage - session.conf.unset(k) - } - } - - plan - } - - /** - * Currently formula depend on AQE in Spark 3.1.1, not sure it can work in future. - */ - private def isFinalStage(plan: SparkPlan): Boolean = { - var shuffleNum = 0 - var broadcastNum = 0 - var reusedNum = 0 - var queryStageNum = 0 - - def collectNumber(p: SparkPlan): SparkPlan = { - p transform { - case shuffle: ShuffleExchangeLike => - shuffleNum += 1 - shuffle - - case broadcast: BroadcastExchangeLike => - broadcastNum += 1 - broadcast - - case reusedExchangeExec: ReusedExchangeExec => - reusedNum += 1 - reusedExchangeExec - - // query stage is leaf node so we need to transform it manually - // compatible with Spark 3.5: - // SPARK-42101: table cache is a independent query stage, so do not need include it. - case queryStage: QueryStageExec if queryStage.nodeName != "TableCacheQueryStage" => - queryStageNum += 1 - collectNumber(queryStage.plan) - queryStage - } - } - collectNumber(plan) - - if (shuffleNum == 0) { - // we don not care about broadcast stage here since it won't change partition number. - true - } else if (shuffleNum + broadcastNum + reusedNum == queryStageNum) { - true - } else { - false - } - } -} -object FinalStageConfigIsolation { - final val SQL_PREFIX = "spark.sql." - final val FINAL_STAGE_CONFIG_PREFIX = "spark.sql.finalStage." - final val PREVIOUS_STAGE_CONFIG_PREFIX = "spark.sql.previousStage." - final val INTERNAL_UNSET_CONFIG_TAG = "__INTERNAL_UNSET_CONFIG_TAG__" - - def getPreviousStageConfigKey(configKey: String): Option[String] = { - if (configKey.startsWith(SQL_PREFIX)) { - Some(s"$PREVIOUS_STAGE_CONFIG_PREFIX${configKey.substring(SQL_PREFIX.length)}") - } else { - None - } - } -} - -case class FinalStageConfigIsolationCleanRule(session: SparkSession) extends Rule[LogicalPlan] { - import FinalStageConfigIsolation._ - - override def apply(plan: LogicalPlan): LogicalPlan = plan match { - case set @ SetCommand(Some((k, Some(_)))) if k.startsWith(SQL_PREFIX) => - checkAndUnsetPreviousStageConfig(k) - set - - case reset @ ResetCommand(Some(k)) if k.startsWith(SQL_PREFIX) => - checkAndUnsetPreviousStageConfig(k) - reset - - case other => other - } - - private def checkAndUnsetPreviousStageConfig(configKey: String): Unit = { - getPreviousStageConfigKey(configKey).foreach { previousStageConfigKey => - if (session.sessionState.conf.contains(previousStageConfigKey)) { - logInfo(s"For previous stage: unset $previousStageConfigKey") - session.conf.unset(previousStageConfigKey) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala deleted file mode 100644 index e72a6c07354..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala +++ /dev/null @@ -1,241 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.network.util.ByteUnit -import org.apache.spark.sql.internal.SQLConf._ - -object KyuubiSQLConf { - - val INSERT_REPARTITION_BEFORE_WRITE = - buildConf("spark.sql.optimizer.insertRepartitionBeforeWrite.enabled") - .doc("Add repartition node at the top of query plan. An approach of merging small files.") - .version("1.2.0") - .booleanConf - .createWithDefault(true) - - val FORCE_SHUFFLE_BEFORE_JOIN = - buildConf("spark.sql.optimizer.forceShuffleBeforeJoin.enabled") - .doc("Ensure shuffle node exists before shuffled join (shj and smj) to make AQE " + - "`OptimizeSkewedJoin` works (complex scenario join, multi table join).") - .version("1.2.0") - .booleanConf - .createWithDefault(false) - - val FINAL_STAGE_CONFIG_ISOLATION = - buildConf("spark.sql.optimizer.finalStageConfigIsolation.enabled") - .doc("If true, the final stage support use different config with previous stage. " + - "The prefix of final stage config key should be `spark.sql.finalStage.`." + - "For example, the raw spark config: `spark.sql.adaptive.advisoryPartitionSizeInBytes`, " + - "then the final stage config should be: " + - "`spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes`.") - .version("1.2.0") - .booleanConf - .createWithDefault(false) - - val INSERT_ZORDER_BEFORE_WRITING = - buildConf("spark.sql.optimizer.insertZorderBeforeWriting.enabled") - .doc("When true, we will follow target table properties to insert zorder or not. " + - "The key properties are: 1) kyuubi.zorder.enabled; if this property is true, we will " + - "insert zorder before writing data. 2) kyuubi.zorder.cols; string split by comma, we " + - "will zorder by these cols.") - .version("1.4.0") - .booleanConf - .createWithDefault(true) - - val ZORDER_GLOBAL_SORT_ENABLED = - buildConf("spark.sql.optimizer.zorderGlobalSort.enabled") - .doc("When true, we do a global sort using zorder. Note that, it can cause data skew " + - "issue if the zorder columns have less cardinality. When false, we only do local sort " + - "using zorder.") - .version("1.4.0") - .booleanConf - .createWithDefault(true) - - val REBALANCE_BEFORE_ZORDER = - buildConf("spark.sql.optimizer.rebalanceBeforeZorder.enabled") - .doc("When true, we do a rebalance before zorder in case data skew. " + - "Note that, if the insertion is dynamic partition we will use the partition " + - "columns to rebalance.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val REBALANCE_ZORDER_COLUMNS_ENABLED = - buildConf("spark.sql.optimizer.rebalanceZorderColumns.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do rebalance before " + - s"Z-Order. If it's dynamic partition insert, the rebalance expression will include " + - s"both partition columns and Z-Order columns.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val TWO_PHASE_REBALANCE_BEFORE_ZORDER = - buildConf("spark.sql.optimizer.twoPhaseRebalanceBeforeZorder.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do two phase rebalance " + - s"before Z-Order for the dynamic partition write. The first phase rebalance using " + - s"dynamic partition column; The second phase rebalance using dynamic partition column + " + - s"Z-Order columns.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val ZORDER_USING_ORIGINAL_ORDERING_ENABLED = - buildConf("spark.sql.optimizer.zorderUsingOriginalOrdering.enabled") - .doc(s"When true and ${REBALANCE_BEFORE_ZORDER.key} is true, we do sort by " + - s"the original ordering i.e. lexicographical order.") - .version("1.6.0") - .booleanConf - .createWithDefault(false) - - val WATCHDOG_MAX_PARTITIONS = - buildConf("spark.sql.watchdog.maxPartitions") - .doc("Set the max partition number when spark scans a data source. " + - "Enable maxPartitions Strategy by specifying this configuration. " + - "Add maxPartitions Strategy to avoid scan excessive partitions " + - "on partitioned table, it's optional that works with defined") - .version("1.4.0") - .intConf - .createOptional - - val WATCHDOG_MAX_FILE_SIZE = - buildConf("spark.sql.watchdog.maxFileSize") - .doc("Set the maximum size in bytes of files when spark scans a data source. " + - "Enable maxFileSize Strategy by specifying this configuration. " + - "Add maxFileSize Strategy to avoid scan excessive size of files," + - " it's optional that works with defined") - .version("1.8.0") - .bytesConf(ByteUnit.BYTE) - .createOptional - - val DROP_IGNORE_NONEXISTENT = - buildConf("spark.sql.optimizer.dropIgnoreNonExistent") - .doc("Do not report an error if DROP DATABASE/TABLE/VIEW/FUNCTION/PARTITION specifies " + - "a non-existent database/table/view/function/partition") - .version("1.5.0") - .booleanConf - .createWithDefault(false) - - val INFER_REBALANCE_AND_SORT_ORDERS = - buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.enabled") - .doc("When ture, infer columns for rebalance and sort orders from original query, " + - "e.g. the join keys from join. It can avoid compression ratio regression.") - .version("1.7.0") - .booleanConf - .createWithDefault(false) - - val INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS = - buildConf("spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns") - .doc("The max columns of inferred columns.") - .version("1.7.0") - .intConf - .checkValue(_ > 0, "must be positive number") - .createWithDefault(3) - - val INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE = - buildConf("spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled") - .doc("When true, add repartition even if the original plan does not have shuffle.") - .version("1.7.0") - .booleanConf - .createWithDefault(false) - - val FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY = - buildConf("spark.sql.optimizer.finalStageConfigIsolationWriteOnly.enabled") - .doc("When true, only enable final stage isolation for writing.") - .version("1.7.0") - .booleanConf - .createWithDefault(true) - - val FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED = - buildConf("spark.sql.finalWriteStage.eagerlyKillExecutors.enabled") - .doc("When true, eagerly kill redundant executors before running final write stage.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL = - buildConf("spark.sql.finalWriteStage.eagerlyKillExecutors.killAll") - .doc("When true, eagerly kill all executors before running final write stage. " + - "Mainly for test.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_SKIP_KILLING_EXECUTORS_FOR_TABLE_CACHE = - buildConf("spark.sql.finalWriteStage.skipKillingExecutorsForTableCache") - .doc("When true, skip killing executors if the plan has table caches.") - .version("1.8.0") - .booleanConf - .createWithDefault(true) - - val FINAL_WRITE_STAGE_PARTITION_FACTOR = - buildConf("spark.sql.finalWriteStage.retainExecutorsFactor") - .doc("If the target executors * factor < active executors, and " + - "target executors * factor > min executors, then kill redundant executors.") - .version("1.8.0") - .doubleConf - .checkValue(_ >= 1, "must be bigger than or equal to 1") - .createWithDefault(1.2) - - val FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED = - buildConf("spark.sql.finalWriteStage.resourceIsolation.enabled") - .doc( - "When true, make final write stage resource isolation using custom RDD resource profile.") - .version("1.8.0") - .booleanConf - .createWithDefault(false) - - val FINAL_WRITE_STAGE_EXECUTOR_CORES = - buildConf("spark.sql.finalWriteStage.executorCores") - .doc("Specify the executor core request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .intConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_MEMORY = - buildConf("spark.sql.finalWriteStage.executorMemory") - .doc("Specify the executor on heap memory request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD = - buildConf("spark.sql.finalWriteStage.executorMemoryOverhead") - .doc("Specify the executor memory overhead request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY = - buildConf("spark.sql.finalWriteStage.executorOffHeapMemory") - .doc("Specify the executor off heap memory request for final write stage. " + - "It would be passed to the RDD resource profile.") - .version("1.8.0") - .stringConf - .createOptional - - val SCRIPT_TRANSFORMATION_ENABLED = - buildConf("spark.sql.execution.scriptTransformation.enabled") - .doc("When false, script transformation is not allowed.") - .version("1.9.0") - .booleanConf - .createWithDefault(true) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala deleted file mode 100644 index 88c5a988fd9..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLExtensionException.scala +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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.sql - -import java.sql.SQLException - -class KyuubiSQLExtensionException(reason: String, cause: Throwable) - extends SQLException(reason, cause) { - - def this(reason: String) = { - this(reason, null) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala deleted file mode 100644 index 7ee439a4399..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLAstBuilder.scala +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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.sql - -import scala.collection.JavaConverters.asScalaBufferConverter -import scala.collection.mutable.ListBuffer - -import org.antlr.v4.runtime.ParserRuleContext -import org.antlr.v4.runtime.misc.Interval -import org.antlr.v4.runtime.tree.ParseTree -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedRelation, UnresolvedStar} -import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.parser.ParserUtils.withOrigin -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project, Sort} - -import org.apache.kyuubi.sql.KyuubiSparkSQLParser._ -import org.apache.kyuubi.sql.zorder.{OptimizeZorderStatement, Zorder} - -class KyuubiSparkSQLAstBuilder extends KyuubiSparkSQLBaseVisitor[AnyRef] with SQLConfHelper { - - def buildOptimizeStatement( - unparsedPredicateOptimize: UnparsedPredicateOptimize, - parseExpression: String => Expression): LogicalPlan = { - - val UnparsedPredicateOptimize(tableIdent, tablePredicate, orderExpr) = - unparsedPredicateOptimize - - val predicate = tablePredicate.map(parseExpression) - verifyPartitionPredicates(predicate) - val table = UnresolvedRelation(tableIdent) - val tableWithFilter = predicate match { - case Some(expr) => Filter(expr, table) - case None => table - } - val query = - Sort( - SortOrder(orderExpr, Ascending, NullsLast, Seq.empty) :: Nil, - conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED), - Project(Seq(UnresolvedStar(None)), tableWithFilter)) - OptimizeZorderStatement(tableIdent, query) - } - - private def verifyPartitionPredicates(predicates: Option[Expression]): Unit = { - predicates.foreach { - case p if !isLikelySelective(p) => - throw new KyuubiSQLExtensionException(s"unsupported partition predicates: ${p.sql}") - case _ => - } - } - - /** - * Forked from Apache Spark's org.apache.spark.sql.catalyst.expressions.PredicateHelper - * The `PredicateHelper.isLikelySelective()` is available since Spark-3.3, forked for Spark - * that is lower than 3.3. - * - * Returns whether an expression is likely to be selective - */ - private def isLikelySelective(e: Expression): Boolean = e match { - case Not(expr) => isLikelySelective(expr) - case And(l, r) => isLikelySelective(l) || isLikelySelective(r) - case Or(l, r) => isLikelySelective(l) && isLikelySelective(r) - case _: StringRegexExpression => true - case _: BinaryComparison => true - case _: In | _: InSet => true - case _: StringPredicate => true - case BinaryPredicate(_) => true - case _: MultiLikeBase => true - case _ => false - } - - private object BinaryPredicate { - def unapply(expr: Expression): Option[Expression] = expr match { - case _: Contains => Option(expr) - case _: StartsWith => Option(expr) - case _: EndsWith => Option(expr) - case _ => None - } - } - - /** - * Create an expression from the given context. This method just passes the context on to the - * visitor and only takes care of typing (We assume that the visitor returns an Expression here). - */ - protected def expression(ctx: ParserRuleContext): Expression = typedVisit(ctx) - - protected def multiPart(ctx: ParserRuleContext): Seq[String] = typedVisit(ctx) - - override def visitSingleStatement(ctx: SingleStatementContext): LogicalPlan = { - visit(ctx.statement()).asInstanceOf[LogicalPlan] - } - - override def visitOptimizeZorder( - ctx: OptimizeZorderContext): UnparsedPredicateOptimize = withOrigin(ctx) { - val tableIdent = multiPart(ctx.multipartIdentifier()) - - val predicate = Option(ctx.whereClause()) - .map(_.partitionPredicate) - .map(extractRawText(_)) - - val zorderCols = ctx.zorderClause().order.asScala - .map(visitMultipartIdentifier) - .map(UnresolvedAttribute(_)) - .toSeq - - val orderExpr = - if (zorderCols.length == 1) { - zorderCols.head - } else { - Zorder(zorderCols) - } - UnparsedPredicateOptimize(tableIdent, predicate, orderExpr) - } - - override def visitPassThrough(ctx: PassThroughContext): LogicalPlan = null - - override def visitMultipartIdentifier(ctx: MultipartIdentifierContext): Seq[String] = - withOrigin(ctx) { - ctx.parts.asScala.map(typedVisit[String]).toSeq - } - - override def visitIdentifier(ctx: IdentifierContext): String = { - withOrigin(ctx) { - ctx.strictIdentifier() match { - case quotedContext: QuotedIdentifierAlternativeContext => - typedVisit[String](quotedContext) - case _ => ctx.getText - } - } - } - - override def visitQuotedIdentifier(ctx: QuotedIdentifierContext): String = { - withOrigin(ctx) { - ctx.BACKQUOTED_IDENTIFIER().getText.stripPrefix("`").stripSuffix("`").replace("``", "`") - } - } - - override def visitZorderClause(ctx: ZorderClauseContext): Seq[UnresolvedAttribute] = - withOrigin(ctx) { - val res = ListBuffer[UnresolvedAttribute]() - ctx.multipartIdentifier().forEach { identifier => - res += UnresolvedAttribute(identifier.parts.asScala.map(typedVisit[String]).toSeq) - } - res.toSeq - } - - private def typedVisit[T](ctx: ParseTree): T = { - ctx.accept(this).asInstanceOf[T] - } - - private def extractRawText(exprContext: ParserRuleContext): String = { - // Extract the raw expression which will be parsed later - exprContext.getStart.getInputStream.getText(new Interval( - exprContext.getStart.getStartIndex, - exprContext.getStop.getStopIndex)) - } -} - -/** - * a logical plan contains an unparsed expression that will be parsed by spark. - */ -trait UnparsedExpressionLogicalPlan extends LogicalPlan { - override def output: Seq[Attribute] = throw new UnsupportedOperationException() - - override def children: Seq[LogicalPlan] = throw new UnsupportedOperationException() - - protected def withNewChildrenInternal( - newChildren: IndexedSeq[LogicalPlan]): LogicalPlan = - throw new UnsupportedOperationException() -} - -case class UnparsedPredicateOptimize( - tableIdent: Seq[String], - tablePredicate: Option[String], - orderExpr: Expression) extends UnparsedExpressionLogicalPlan {} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala deleted file mode 100644 index c4ddcef2b13..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSessionExtensions - -import org.apache.kyuubi.sql.zorder.{InsertZorderBeforeWritingDatasource, InsertZorderBeforeWritingHive, ResolveZorder} - -class KyuubiSparkSQLCommonExtension extends (SparkSessionExtensions => Unit) { - override def apply(extensions: SparkSessionExtensions): Unit = { - KyuubiSparkSQLCommonExtension.injectCommonExtensions(extensions) - } -} - -object KyuubiSparkSQLCommonExtension { - def injectCommonExtensions(extensions: SparkSessionExtensions): Unit = { - // inject zorder parser and related rules - extensions.injectParser { case (_, parser) => new SparkKyuubiSparkSQLParser(parser) } - extensions.injectResolutionRule(ResolveZorder) - - // Note that: - // InsertZorderBeforeWritingDatasource and InsertZorderBeforeWritingHive - // should be applied before - // RepartitionBeforeWriting and RebalanceBeforeWriting - // because we can only apply one of them (i.e. Global Sort or Repartition/Rebalance) - extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingDatasource) - extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingHive) - extensions.injectPostHocResolutionRule(FinalStageConfigIsolationCleanRule) - - extensions.injectQueryStagePrepRule(_ => InsertShuffleNodeBeforeJoin) - - extensions.injectQueryStagePrepRule(FinalStageConfigIsolation(_)) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala deleted file mode 100644 index 33ff3e3177a..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.{FinalStageResourceManager, InjectCustomResourceProfile, SparkSessionExtensions} - -import org.apache.kyuubi.sql.watchdog.{KyuubiUnsupportedOperationsCheck, MaxScanStrategy} - -// scalastyle:off line.size.limit -/** - * Depend on Spark SQL Extension framework, we can use this extension follow steps - * 1. move this jar into $SPARK_HOME/jars - * 2. add config into `spark-defaults.conf`: `spark.sql.extensions=org.apache.kyuubi.sql.KyuubiSparkSQLExtension` - */ -// scalastyle:on line.size.limit -class KyuubiSparkSQLExtension extends (SparkSessionExtensions => Unit) { - override def apply(extensions: SparkSessionExtensions): Unit = { - KyuubiSparkSQLCommonExtension.injectCommonExtensions(extensions) - - extensions.injectPostHocResolutionRule(RebalanceBeforeWritingDatasource) - extensions.injectPostHocResolutionRule(RebalanceBeforeWritingHive) - extensions.injectPostHocResolutionRule(DropIgnoreNonexistent) - - // watchdog extension - extensions.injectCheckRule(_ => KyuubiUnsupportedOperationsCheck) - extensions.injectPlannerStrategy(MaxScanStrategy) - - extensions.injectQueryStagePrepRule(FinalStageResourceManager(_)) - extensions.injectQueryStagePrepRule(InjectCustomResourceProfile) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala deleted file mode 100644 index c4418c33c44..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLParser.scala +++ /dev/null @@ -1,140 +0,0 @@ -/* - * 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.sql - -import org.antlr.v4.runtime._ -import org.antlr.v4.runtime.atn.PredictionMode -import org.antlr.v4.runtime.misc.{Interval, ParseCancellationException} -import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.{FunctionIdentifier, SQLConfHelper, TableIdentifier} -import org.apache.spark.sql.catalyst.expressions.Expression -import org.apache.spark.sql.catalyst.parser.{ParseErrorListener, ParseException, ParserInterface, PostProcessor} -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.catalyst.trees.Origin -import org.apache.spark.sql.types.{DataType, StructType} - -abstract class KyuubiSparkSQLParserBase extends ParserInterface with SQLConfHelper { - def delegate: ParserInterface - def astBuilder: KyuubiSparkSQLAstBuilder - - override def parsePlan(sqlText: String): LogicalPlan = parse(sqlText) { parser => - astBuilder.visit(parser.singleStatement()) match { - case optimize: UnparsedPredicateOptimize => - astBuilder.buildOptimizeStatement(optimize, delegate.parseExpression) - case plan: LogicalPlan => plan - case _ => delegate.parsePlan(sqlText) - } - } - - protected def parse[T](command: String)(toResult: KyuubiSparkSQLParser => T): T = { - val lexer = new KyuubiSparkSQLLexer( - new UpperCaseCharStream(CharStreams.fromString(command))) - lexer.removeErrorListeners() - lexer.addErrorListener(ParseErrorListener) - - val tokenStream = new CommonTokenStream(lexer) - val parser = new KyuubiSparkSQLParser(tokenStream) - parser.addParseListener(PostProcessor) - parser.removeErrorListeners() - parser.addErrorListener(ParseErrorListener) - - try { - try { - // first, try parsing with potentially faster SLL mode - parser.getInterpreter.setPredictionMode(PredictionMode.SLL) - toResult(parser) - } catch { - case _: ParseCancellationException => - // if we fail, parse with LL mode - tokenStream.seek(0) // rewind input stream - parser.reset() - - // Try Again. - parser.getInterpreter.setPredictionMode(PredictionMode.LL) - toResult(parser) - } - } catch { - case e: ParseException if e.command.isDefined => - throw e - case e: ParseException => - throw e.withCommand(command) - case e: AnalysisException => - val position = Origin(e.line, e.startPosition) - throw new ParseException(Option(command), e.message, position, position) - } - } - - override def parseExpression(sqlText: String): Expression = { - delegate.parseExpression(sqlText) - } - - override def parseTableIdentifier(sqlText: String): TableIdentifier = { - delegate.parseTableIdentifier(sqlText) - } - - override def parseFunctionIdentifier(sqlText: String): FunctionIdentifier = { - delegate.parseFunctionIdentifier(sqlText) - } - - override def parseMultipartIdentifier(sqlText: String): Seq[String] = { - delegate.parseMultipartIdentifier(sqlText) - } - - override def parseTableSchema(sqlText: String): StructType = { - delegate.parseTableSchema(sqlText) - } - - override def parseDataType(sqlText: String): DataType = { - delegate.parseDataType(sqlText) - } - - /** - * This functions was introduced since spark-3.3, for more details, please see - * https://github.com/apache/spark/pull/34543 - */ - override def parseQuery(sqlText: String): LogicalPlan = { - delegate.parseQuery(sqlText) - } -} - -class SparkKyuubiSparkSQLParser( - override val delegate: ParserInterface) - extends KyuubiSparkSQLParserBase { - def astBuilder: KyuubiSparkSQLAstBuilder = new KyuubiSparkSQLAstBuilder -} - -/* Copied from Apache Spark's to avoid dependency on Spark Internals */ -class UpperCaseCharStream(wrapped: CodePointCharStream) extends CharStream { - override def consume(): Unit = wrapped.consume() - override def getSourceName(): String = wrapped.getSourceName - override def index(): Int = wrapped.index - override def mark(): Int = wrapped.mark - override def release(marker: Int): Unit = wrapped.release(marker) - override def seek(where: Int): Unit = wrapped.seek(where) - override def size(): Int = wrapped.size - - override def getText(interval: Interval): String = wrapped.getText(interval) - - // scalastyle:off - override def LA(i: Int): Int = { - val la = wrapped.LA(i) - if (la == 0 || la == IntStream.EOF) la - else Character.toUpperCase(la) - } - // scalastyle:on -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala deleted file mode 100644 index 2c6980fdb6e..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala +++ /dev/null @@ -1,68 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, SortOrder} -import org.apache.spark.sql.catalyst.plans.logical._ - -trait RepartitionBuilderWithRebalance extends RepartitionBuilder { - override def buildRepartition( - dynamicPartitionColumns: Seq[Attribute], - query: LogicalPlan): LogicalPlan = { - if (!conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS) || - dynamicPartitionColumns.nonEmpty) { - RebalancePartitions(dynamicPartitionColumns, query) - } else { - val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS) - val inferred = InferRebalanceAndSortOrders.infer(query) - if (inferred.isDefined) { - val (partitioning, ordering) = inferred.get - val rebalance = RebalancePartitions(partitioning.take(maxColumns), query) - if (ordering.nonEmpty) { - val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending)) - Sort(sortOrders, false, rebalance) - } else { - rebalance - } - } else { - RebalancePartitions(dynamicPartitionColumns, query) - } - } - } -} - -/** - * For datasource table, there two commands can write data to table - * 1. InsertIntoHadoopFsRelationCommand - * 2. CreateDataSourceTableAsSelectCommand - * This rule add a RebalancePartitions node between write and query - */ -case class RebalanceBeforeWritingDatasource(session: SparkSession) - extends RepartitionBeforeWritingDatasourceBase - with RepartitionBuilderWithRebalance {} - -/** - * For Hive table, there two commands can write data to table - * 1. InsertIntoHiveTable - * 2. CreateHiveTableAsSelectCommand - * This rule add a RebalancePartitions node between write and query - */ -case class RebalanceBeforeWritingHive(session: SparkSession) - extends RepartitionBeforeWritingHiveBase - with RepartitionBuilderWithRebalance {} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala deleted file mode 100644 index f631c880cd1..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala +++ /dev/null @@ -1,135 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.command.InsertIntoDataSourceDirCommand -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.hive.execution.{InsertIntoHiveDirCommand, InsertIntoHiveTable} -import org.apache.spark.sql.internal.StaticSQLConf - -trait RepartitionBuilder extends Rule[LogicalPlan] with RepartitionBeforeWriteHelper { - def buildRepartition( - dynamicPartitionColumns: Seq[Attribute], - query: LogicalPlan): LogicalPlan -} - -/** - * For datasource table, there two commands can write data to table - * 1. InsertIntoHadoopFsRelationCommand - * 2. CreateDataSourceTableAsSelectCommand - * This rule add a repartition node between write and query - */ -abstract class RepartitionBeforeWritingDatasourceBase extends RepartitionBuilder { - - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE)) { - addRepartition(plan) - } else { - plan - } - } - - private def addRepartition(plan: LogicalPlan): LogicalPlan = plan match { - case i @ InsertIntoHadoopFsRelationCommand(_, sp, _, pc, bucket, _, _, query, _, _, _, _) - if query.resolved && bucket.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = pc.filterNot(attr => sp.contains(attr.name)) - i.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case i @ InsertIntoDataSourceDirCommand(_, _, query, _) - if query.resolved && canInsertRepartitionByExpression(query) => - i.copy(query = buildRepartition(Seq.empty, query)) - - case u @ Union(children, _, _) => - u.copy(children = children.map(addRepartition)) - - case _ => plan - } -} - -/** - * For Hive table, there two commands can write data to table - * 1. InsertIntoHiveTable - * 2. CreateHiveTableAsSelectCommand - * This rule add a repartition node between write and query - */ -abstract class RepartitionBeforeWritingHiveBase extends RepartitionBuilder { - override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(StaticSQLConf.CATALOG_IMPLEMENTATION) == "hive" && - conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE)) { - addRepartition(plan) - } else { - plan - } - } - - def addRepartition(plan: LogicalPlan): LogicalPlan = plan match { - case i @ InsertIntoHiveTable(table, partition, query, _, _, _, _, _, _, _, _) - if query.resolved && table.bucketSpec.isEmpty && canInsertRepartitionByExpression(query) => - val dynamicPartitionColumns = partition.filter(_._2.isEmpty).keys - .flatMap(name => query.output.find(_.name == name)).toSeq - i.copy(query = buildRepartition(dynamicPartitionColumns, query)) - - case i @ InsertIntoHiveDirCommand(_, _, query, _, _) - if query.resolved && canInsertRepartitionByExpression(query) => - i.copy(query = buildRepartition(Seq.empty, query)) - - case u @ Union(children, _, _) => - u.copy(children = children.map(addRepartition)) - - case _ => plan - } -} - -trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] { - private def hasBenefit(plan: LogicalPlan): Boolean = { - def probablyHasShuffle: Boolean = plan.find { - case _: Join => true - case _: Aggregate => true - case _: Distinct => true - case _: Deduplicate => true - case _: Window => true - case s: Sort if s.global => true - case _: RepartitionOperation => true - case _: RebalancePartitions => true - case _: GlobalLimit => true - case _ => false - }.isDefined - - conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE) || probablyHasShuffle - } - - def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = { - def canInsert(p: LogicalPlan): Boolean = p match { - case Project(_, child) => canInsert(child) - case SubqueryAlias(_, child) => canInsert(child) - case Limit(_, _) => false - case _: Sort => false - case _: RepartitionOperation => false - case _: RebalancePartitions => false - case _ => true - } - - // 1. make sure AQE is enabled, otherwise it is no meaning to add a shuffle - // 2. make sure it does not break the semantics of original plan - // 3. try to avoid adding a shuffle if it has potential performance regression - conf.adaptiveExecutionEnabled && canInsert(plan) && hasBenefit(plan) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/WriteUtils.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/WriteUtils.scala deleted file mode 100644 index 89dd8319480..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/WriteUtils.scala +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.sql - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.execution.{SparkPlan, UnionExec} -import org.apache.spark.sql.execution.command.DataWritingCommandExec -import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec - -object WriteUtils { - def isWrite(session: SparkSession, plan: SparkPlan): Boolean = { - plan match { - case _: DataWritingCommandExec => true - case _: V2TableWriteExec => true - case u: UnionExec if u.children.nonEmpty => u.children.forall(isWrite(session, _)) - case _ => false - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala deleted file mode 100644 index 2b4d3940ada..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiUnsupportedOperationsCheck.scala +++ /dev/null @@ -1,35 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ScriptTransformation} - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} - -object KyuubiUnsupportedOperationsCheck extends (LogicalPlan => Unit) with SQLConfHelper { - override def apply(plan: LogicalPlan): Unit = - conf.getConf(KyuubiSQLConf.SCRIPT_TRANSFORMATION_ENABLED) match { - case false => plan foreach { - case _: ScriptTransformation => - throw new KyuubiSQLExtensionException("Script transformation is not allowed") - case _ => - } - case true => - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala deleted file mode 100644 index e44309192a9..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/KyuubiWatchDogException.scala +++ /dev/null @@ -1,30 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -final class MaxPartitionExceedException( - private val reason: String = "", - private val cause: Throwable = None.orNull) - extends KyuubiSQLExtensionException(reason, cause) - -final class MaxFileSizeExceedException( - private val reason: String = "", - private val cause: Throwable = None.orNull) - extends KyuubiSQLExtensionException(reason, cause) diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala deleted file mode 100644 index e647ad3250e..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxScanStrategy.scala +++ /dev/null @@ -1,340 +0,0 @@ -/* - * 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.sql.watchdog - -import org.apache.hadoop.fs.Path -import org.apache.spark.sql.{PruneFileSourcePartitionHelper, SparkSession, Strategy} -import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} -import org.apache.spark.sql.catalyst.planning.ScanOperation -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.datasources.{CatalogFileIndex, HadoopFsRelation, InMemoryFileIndex, LogicalRelation} -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation -import org.apache.spark.sql.types.StructType - -import org.apache.kyuubi.sql.KyuubiSQLConf - -/** - * Add MaxScanStrategy to avoid scan excessive partitions or files - * 1. Check if scan exceed maxPartition of partitioned table - * 2. Check if scan exceed maxFileSize (calculated by hive table and partition statistics) - * This Strategy Add Planner Strategy after LogicalOptimizer - * @param session - */ -case class MaxScanStrategy(session: SparkSession) - extends Strategy - with SQLConfHelper - with PruneFileSourcePartitionHelper { - override def apply(plan: LogicalPlan): Seq[SparkPlan] = { - val maxScanPartitionsOpt = conf.getConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS) - val maxFileSizeOpt = conf.getConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE) - if (maxScanPartitionsOpt.isDefined || maxFileSizeOpt.isDefined) { - checkScan(plan, maxScanPartitionsOpt, maxFileSizeOpt) - } - Nil - } - - private def checkScan( - plan: LogicalPlan, - maxScanPartitionsOpt: Option[Int], - maxFileSizeOpt: Option[Long]): Unit = { - plan match { - case ScanOperation(_, _, _, relation: HiveTableRelation) => - if (relation.isPartitioned) { - relation.prunedPartitions match { - case Some(prunedPartitions) => - if (maxScanPartitionsOpt.exists(_ < prunedPartitions.size)) { - throw new MaxPartitionExceedException( - s""" - |SQL job scan hive partition: ${prunedPartitions.size} - |exceed restrict of hive scan maxPartition ${maxScanPartitionsOpt.get} - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - lazy val scanFileSize = prunedPartitions.flatMap(_.stats).map(_.sizeInBytes).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - Some(relation.tableMeta), - prunedPartitions.flatMap(_.storage.locationUri).map(_.toString), - relation.partitionCols.map(_.name)) - } - case _ => - lazy val scanPartitions: Int = session - .sessionState.catalog.externalCatalog.listPartitionNames( - relation.tableMeta.database, - relation.tableMeta.identifier.table).size - if (maxScanPartitionsOpt.exists(_ < scanPartitions)) { - throw new MaxPartitionExceedException( - s""" - |Your SQL job scan a whole huge table without any partition filter, - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - - lazy val scanFileSize: BigInt = - relation.tableMeta.stats.map(_.sizeInBytes).getOrElse { - session - .sessionState.catalog.externalCatalog.listPartitions( - relation.tableMeta.database, - relation.tableMeta.identifier.table).flatMap(_.stats).map(_.sizeInBytes).sum - } - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |Your SQL job scan a whole huge table without any partition filter, - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${relation.tableMeta.qualifiedName} - |Owner: ${relation.tableMeta.owner} - |Partition Structure: ${relation.partitionCols.map(_.name).mkString(", ")} - |""".stripMargin) - } - } - } else { - lazy val scanFileSize = relation.tableMeta.stats.map(_.sizeInBytes).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - Some(relation.tableMeta)) - } - } - case ScanOperation( - _, - _, - filters, - relation @ LogicalRelation( - fsRelation @ HadoopFsRelation( - fileIndex: InMemoryFileIndex, - partitionSchema, - _, - _, - _, - _), - _, - _, - _)) => - if (fsRelation.partitionSchema.nonEmpty) { - val (partitionKeyFilters, dataFilter) = - getPartitionKeyFiltersAndDataFilters( - SparkSession.active, - relation, - partitionSchema, - filters, - relation.output) - val prunedPartitions = fileIndex.listFiles( - partitionKeyFilters.toSeq, - dataFilter) - if (maxScanPartitionsOpt.exists(_ < prunedPartitions.size)) { - throw maxPartitionExceedError( - prunedPartitions.size, - maxScanPartitionsOpt.get, - relation.catalogTable, - fileIndex.rootPaths, - fsRelation.partitionSchema) - } - lazy val scanFileSize = prunedPartitions.flatMap(_.files).map(_.getLen).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - relation.catalogTable, - fileIndex.rootPaths.map(_.toString), - fsRelation.partitionSchema.map(_.name)) - } - } else { - lazy val scanFileSize = fileIndex.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - relation.catalogTable) - } - } - case ScanOperation( - _, - _, - filters, - logicalRelation @ LogicalRelation( - fsRelation @ HadoopFsRelation( - catalogFileIndex: CatalogFileIndex, - partitionSchema, - _, - _, - _, - _), - _, - _, - _)) => - if (fsRelation.partitionSchema.nonEmpty) { - val (partitionKeyFilters, _) = - getPartitionKeyFiltersAndDataFilters( - SparkSession.active, - logicalRelation, - partitionSchema, - filters, - logicalRelation.output) - - val fileIndex = catalogFileIndex.filterPartitions( - partitionKeyFilters.toSeq) - - lazy val prunedPartitionSize = fileIndex.partitionSpec().partitions.size - if (maxScanPartitionsOpt.exists(_ < prunedPartitionSize)) { - throw maxPartitionExceedError( - prunedPartitionSize, - maxScanPartitionsOpt.get, - logicalRelation.catalogTable, - catalogFileIndex.rootPaths, - fsRelation.partitionSchema) - } - - lazy val scanFileSize = fileIndex - .listFiles(Nil, Nil).flatMap(_.files).map(_.getLen).sum - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw partTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - logicalRelation.catalogTable, - catalogFileIndex.rootPaths.map(_.toString), - fsRelation.partitionSchema.map(_.name)) - } - } else { - lazy val scanFileSize = catalogFileIndex.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw nonPartTableMaxFileExceedError( - scanFileSize, - maxFileSizeOpt.get, - logicalRelation.catalogTable) - } - } - case ScanOperation( - _, - _, - _, - relation @ DataSourceV2ScanRelation(_, _, _, _, _)) => - val table = relation.relation.table - if (table.partitioning().nonEmpty) { - val partitionColumnNames = table.partitioning().map(_.describe()) - val stats = relation.computeStats() - lazy val scanFileSize = stats.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize ${maxFileSizeOpt.get} - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${table.name()} - |Partition Structure: ${partitionColumnNames.mkString(",")} - |""".stripMargin) - } - } else { - val stats = relation.computeStats() - lazy val scanFileSize = stats.sizeInBytes - if (maxFileSizeOpt.exists(_ < scanFileSize)) { - throw new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize ${maxFileSizeOpt.get} - |detail as below: - |Table: ${table.name()} - |""".stripMargin) - } - } - case _ => - } - } - - def maxPartitionExceedError( - prunedPartitionSize: Int, - maxPartitionSize: Int, - tableMeta: Option[CatalogTable], - rootPaths: Seq[Path], - partitionSchema: StructType): Throwable = { - val truncatedPaths = - if (rootPaths.length > 5) { - rootPaths.slice(0, 5).mkString(",") + """... """ + (rootPaths.length - 5) + " more paths" - } else { - rootPaths.mkString(",") - } - - new MaxPartitionExceedException( - s""" - |SQL job scan data source partition: $prunedPartitionSize - |exceed restrict of data source scan maxPartition $maxPartitionSize - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |RootPaths: $truncatedPaths - |Partition Structure: ${partitionSchema.map(_.name).mkString(", ")} - |""".stripMargin) - } - - private def partTableMaxFileExceedError( - scanFileSize: Number, - maxFileSize: Long, - tableMeta: Option[CatalogTable], - rootPaths: Seq[String], - partitions: Seq[String]): Throwable = { - val truncatedPaths = - if (rootPaths.length > 5) { - rootPaths.slice(0, 5).mkString(",") + """... """ + (rootPaths.length - 5) + " more paths" - } else { - rootPaths.mkString(",") - } - - new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize $maxFileSize - |You should optimize your SQL logical according partition structure - |or shorten query scope such as p_date, detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |RootPaths: $truncatedPaths - |Partition Structure: ${partitions.mkString(", ")} - |""".stripMargin) - } - - private def nonPartTableMaxFileExceedError( - scanFileSize: Number, - maxFileSize: Long, - tableMeta: Option[CatalogTable]): Throwable = { - new MaxFileSizeExceedException( - s""" - |SQL job scan file size in bytes: $scanFileSize - |exceed restrict of table scan maxFileSize $maxFileSize - |detail as below: - |Table: ${tableMeta.map(_.qualifiedName).getOrElse("")} - |Owner: ${tableMeta.map(_.owner).getOrElse("")} - |Location: ${tableMeta.map(_.location).getOrElse("")} - |""".stripMargin) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala deleted file mode 100644 index d96b546b94d..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala +++ /dev/null @@ -1,183 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.catalog.CatalogTable -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Expression, NullsLast, SortOrder} -import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.hive.execution.InsertIntoHiveTable - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} - -trait ZorderBuilder { - def buildZorder(children: Seq[Expression]): ZorderBase -} - -trait InsertZorderHelper extends Rule[LogicalPlan] with ZorderBuilder { - private val KYUUBI_ZORDER_ENABLED = "kyuubi.zorder.enabled" - private val KYUUBI_ZORDER_COLS = "kyuubi.zorder.cols" - - def isZorderEnabled(props: Map[String, String]): Boolean = { - props.contains(KYUUBI_ZORDER_ENABLED) && - "true".equalsIgnoreCase(props(KYUUBI_ZORDER_ENABLED)) && - props.contains(KYUUBI_ZORDER_COLS) - } - - def getZorderColumns(props: Map[String, String]): Seq[String] = { - val cols = props.get(KYUUBI_ZORDER_COLS) - assert(cols.isDefined) - cols.get.split(",").map(_.trim) - } - - def canInsertZorder(query: LogicalPlan): Boolean = query match { - case Project(_, child) => canInsertZorder(child) - case _: RepartitionOperation | _: RebalancePartitions - if !conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) => true - // TODO: actually, we can force zorder even if existed some shuffle - case _: Sort => false - case _: RepartitionOperation => false - case _: RebalancePartitions => false - case _ => true - } - - def insertZorder( - catalogTable: CatalogTable, - plan: LogicalPlan, - dynamicPartitionColumns: Seq[Attribute]): LogicalPlan = { - if (!canInsertZorder(plan)) { - return plan - } - val cols = getZorderColumns(catalogTable.properties) - val resolver = session.sessionState.conf.resolver - val output = plan.output - val bound = cols.flatMap(col => output.find(attr => resolver(attr.name, col))) - if (bound.size < cols.size) { - logWarning(s"target table does not contain all zorder cols: ${cols.mkString(",")}, " + - s"please check your table properties ${KYUUBI_ZORDER_COLS}.") - plan - } else { - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) && - conf.getConf(KyuubiSQLConf.REBALANCE_BEFORE_ZORDER)) { - throw new KyuubiSQLExtensionException(s"${KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key} " + - s"and ${KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key} can not be enabled together.") - } - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) && - dynamicPartitionColumns.nonEmpty) { - logWarning(s"Dynamic partition insertion with global sort may produce small files.") - } - - val zorderExpr = - if (bound.length == 1) { - bound - } else if (conf.getConf(KyuubiSQLConf.ZORDER_USING_ORIGINAL_ORDERING_ENABLED)) { - bound.asInstanceOf[Seq[Expression]] - } else { - buildZorder(bound) :: Nil - } - val (global, orderExprs, child) = - if (conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED)) { - (true, zorderExpr, plan) - } else if (conf.getConf(KyuubiSQLConf.REBALANCE_BEFORE_ZORDER)) { - val rebalanceExpr = - if (dynamicPartitionColumns.isEmpty) { - // static partition insert - bound - } else if (conf.getConf(KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED)) { - // improve data compression ratio - dynamicPartitionColumns.asInstanceOf[Seq[Expression]] ++ bound - } else { - dynamicPartitionColumns.asInstanceOf[Seq[Expression]] - } - // for dynamic partition insert, Spark always sort the partition columns, - // so here we sort partition columns + zorder. - val rebalance = - if (dynamicPartitionColumns.nonEmpty && - conf.getConf(KyuubiSQLConf.TWO_PHASE_REBALANCE_BEFORE_ZORDER)) { - // improve compression ratio - RebalancePartitions( - rebalanceExpr, - RebalancePartitions(dynamicPartitionColumns, plan)) - } else { - RebalancePartitions(rebalanceExpr, plan) - } - (false, dynamicPartitionColumns.asInstanceOf[Seq[Expression]] ++ zorderExpr, rebalance) - } else { - (false, zorderExpr, plan) - } - val order = orderExprs.map { expr => - SortOrder(expr, Ascending, NullsLast, Seq.empty) - } - Sort(order, global, child) - } - } - - override def buildZorder(children: Seq[Expression]): ZorderBase = Zorder(children) - - def session: SparkSession - def applyInternal(plan: LogicalPlan): LogicalPlan - - final override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.getConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING)) { - applyInternal(plan) - } else { - plan - } - } -} - -case class InsertZorderBeforeWritingDatasource(session: SparkSession) - extends InsertZorderHelper { - override def applyInternal(plan: LogicalPlan): LogicalPlan = plan match { - case insert: InsertIntoHadoopFsRelationCommand - if insert.query.resolved && - insert.bucketSpec.isEmpty && insert.catalogTable.isDefined && - isZorderEnabled(insert.catalogTable.get.properties) => - val dynamicPartition = - insert.partitionColumns.filterNot(attr => insert.staticPartitions.contains(attr.name)) - val newQuery = insertZorder(insert.catalogTable.get, insert.query, dynamicPartition) - if (newQuery.eq(insert.query)) { - insert - } else { - insert.copy(query = newQuery) - } - - case _ => plan - } -} - -case class InsertZorderBeforeWritingHive(session: SparkSession) - extends InsertZorderHelper { - override def applyInternal(plan: LogicalPlan): LogicalPlan = plan match { - case insert: InsertIntoHiveTable - if insert.query.resolved && - insert.table.bucketSpec.isEmpty && isZorderEnabled(insert.table.properties) => - val dynamicPartition = insert.partition.filter(_._2.isEmpty).keys - .flatMap(name => insert.query.output.find(_.name == name)).toSeq - val newQuery = insertZorder(insert.table, insert.query, dynamicPartition) - if (newQuery.eq(insert.query)) { - insert - } else { - insert.copy(query = newQuery) - } - - case _ => plan - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala deleted file mode 100644 index 21d1cf2a25b..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderCommandBase.scala +++ /dev/null @@ -1,78 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.{Row, SparkSession} -import org.apache.spark.sql.catalyst.catalog.CatalogTable -import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.command.DataWritingCommand -import org.apache.spark.sql.hive.execution.InsertIntoHiveTable - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -/** - * A runnable command for zorder, we delegate to real command to execute - */ -abstract class OptimizeZorderCommandBase extends DataWritingCommand { - def catalogTable: CatalogTable - - override def outputColumnNames: Seq[String] = query.output.map(_.name) - - private def isHiveTable: Boolean = { - catalogTable.provider.isEmpty || - (catalogTable.provider.isDefined && "hive".equalsIgnoreCase(catalogTable.provider.get)) - } - - private def getWritingCommand(session: SparkSession): DataWritingCommand = { - // TODO: Support convert hive relation to datasource relation, can see - // [[org.apache.spark.sql.hive.RelationConversions]] - InsertIntoHiveTable( - catalogTable, - catalogTable.partitionColumnNames.map(p => (p, None)).toMap, - query, - overwrite = true, - ifPartitionNotExists = false, - outputColumnNames) - } - - override def run(session: SparkSession, child: SparkPlan): Seq[Row] = { - // TODO: Support datasource relation - // TODO: Support read and insert overwrite the same table for some table format - if (!isHiveTable) { - throw new KyuubiSQLExtensionException("only support hive table") - } - - val command = getWritingCommand(session) - command.run(session, child) - DataWritingCommand.propogateMetrics(session.sparkContext, command, metrics) - Seq.empty - } -} - -/** - * A runnable command for zorder, we delegate to real command to execute - */ -case class OptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan) - extends OptimizeZorderCommandBase { - protected def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = { - copy(query = newChild) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala deleted file mode 100644 index 895f9e24be3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/OptimizeZorderStatementBase.scala +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} - -/** - * A zorder statement that contains we parsed from SQL. - * We should convert this plan to certain command at Analyzer. - */ -case class OptimizeZorderStatement( - tableIdentifier: Seq[String], - query: LogicalPlan) extends UnaryNode { - override def child: LogicalPlan = query - override def output: Seq[Attribute] = child.output - protected def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = - copy(query = newChild) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala deleted file mode 100644 index 9f735caa7a7..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ResolveZorderBase.scala +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.TableIdentifier -import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} -import org.apache.spark.sql.catalyst.expressions.AttributeSet -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, SubqueryAlias} -import org.apache.spark.sql.catalyst.rules.Rule - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -/** - * Resolve `OptimizeZorderStatement` to `OptimizeZorderCommand` - */ -abstract class ResolveZorderBase extends Rule[LogicalPlan] { - def session: SparkSession - def buildOptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan): OptimizeZorderCommandBase - - protected def checkQueryAllowed(query: LogicalPlan): Unit = query foreach { - case Filter(condition, SubqueryAlias(_, tableRelation: HiveTableRelation)) => - if (tableRelation.partitionCols.isEmpty) { - throw new KyuubiSQLExtensionException("Filters are only supported for partitioned table") - } - - val partitionKeyIds = AttributeSet(tableRelation.partitionCols) - if (condition.references.isEmpty || !condition.references.subsetOf(partitionKeyIds)) { - throw new KyuubiSQLExtensionException("Only partition column filters are allowed") - } - - case _ => - } - - protected def getTableIdentifier(tableIdent: Seq[String]): TableIdentifier = tableIdent match { - case Seq(tbl) => TableIdentifier.apply(tbl) - case Seq(db, tbl) => TableIdentifier.apply(tbl, Some(db)) - case _ => throw new KyuubiSQLExtensionException( - "only support session catalog table, please use db.table instead") - } - - override def apply(plan: LogicalPlan): LogicalPlan = plan match { - case statement: OptimizeZorderStatement if statement.query.resolved => - checkQueryAllowed(statement.query) - val tableIdentifier = getTableIdentifier(statement.tableIdentifier) - val catalogTable = session.sessionState.catalog.getTableMetadata(tableIdentifier) - buildOptimizeZorderCommand(catalogTable, statement.query) - - case _ => plan - } -} - -/** - * Resolve `OptimizeZorderStatement` to `OptimizeZorderCommand` - */ -case class ResolveZorder(session: SparkSession) extends ResolveZorderBase { - override def buildOptimizeZorderCommand( - catalogTable: CatalogTable, - query: LogicalPlan): OptimizeZorderCommandBase = { - OptimizeZorderCommand(catalogTable, query) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala deleted file mode 100644 index e4d98ccbe84..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBase.scala +++ /dev/null @@ -1,95 +0,0 @@ -/* - * 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.sql.zorder - -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.analysis.TypeCheckResult -import org.apache.spark.sql.catalyst.expressions.Expression -import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, FalseLiteral} -import org.apache.spark.sql.catalyst.expressions.codegen.Block._ -import org.apache.spark.sql.types.{BinaryType, DataType} - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -abstract class ZorderBase extends Expression { - override def foldable: Boolean = children.forall(_.foldable) - override def nullable: Boolean = false - override def dataType: DataType = BinaryType - override def prettyName: String = "zorder" - - override def checkInputDataTypes(): TypeCheckResult = { - try { - defaultNullValues - TypeCheckResult.TypeCheckSuccess - } catch { - case e: KyuubiSQLExtensionException => - TypeCheckResult.TypeCheckFailure(e.getMessage) - } - } - - @transient - private[this] lazy val defaultNullValues: Array[Any] = - children.map(_.dataType) - .map(ZorderBytesUtils.defaultValue) - .toArray - - override def eval(input: InternalRow): Any = { - val childrenValues = children.zipWithIndex.map { - case (child: Expression, index) => - val v = child.eval(input) - if (v == null) { - defaultNullValues(index) - } else { - v - } - } - ZorderBytesUtils.interleaveBits(childrenValues.toArray) - } - - override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { - val evals = children.map(_.genCode(ctx)) - val defaultValues = ctx.addReferenceObj("defaultValues", defaultNullValues) - val values = ctx.freshName("values") - val util = ZorderBytesUtils.getClass.getName.stripSuffix("$") - val inputs = evals.zipWithIndex.map { - case (eval, index) => - s""" - |${eval.code} - |if (${eval.isNull}) { - | $values[$index] = $defaultValues[$index]; - |} else { - | $values[$index] = ${eval.value}; - |} - |""".stripMargin - } - ev.copy( - code = - code""" - |byte[] ${ev.value} = null; - |Object[] $values = new Object[${evals.length}]; - |${inputs.mkString("\n")} - |${ev.value} = $util.interleaveBits($values); - |""".stripMargin, - isNull = FalseLiteral) - } -} - -case class Zorder(children: Seq[Expression]) extends ZorderBase { - protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = - copy(children = newChildren) -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala deleted file mode 100644 index d249f1dc32f..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/ZorderBytesUtils.scala +++ /dev/null @@ -1,517 +0,0 @@ -/* - * 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.sql.zorder - -import java.lang.{Double => jDouble, Float => jFloat} - -import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.UTF8String - -import org.apache.kyuubi.sql.KyuubiSQLExtensionException - -object ZorderBytesUtils { - final private val BIT_8_MASK = 1 << 7 - final private val BIT_16_MASK = 1 << 15 - final private val BIT_32_MASK = 1 << 31 - final private val BIT_64_MASK = 1L << 63 - - def interleaveBits(inputs: Array[Any]): Array[Byte] = { - inputs.length match { - // it's a more fast approach, use O(8 * 8) - // can see http://graphics.stanford.edu/~seander/bithacks.html#InterleaveTableObvious - case 1 => longToByte(toLong(inputs(0))) - case 2 => interleave2Longs(toLong(inputs(0)), toLong(inputs(1))) - case 3 => interleave3Longs(toLong(inputs(0)), toLong(inputs(1)), toLong(inputs(2))) - case 4 => - interleave4Longs(toLong(inputs(0)), toLong(inputs(1)), toLong(inputs(2)), toLong(inputs(3))) - case 5 => interleave5Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4))) - case 6 => interleave6Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5))) - case 7 => interleave7Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5)), - toLong(inputs(6))) - case 8 => interleave8Longs( - toLong(inputs(0)), - toLong(inputs(1)), - toLong(inputs(2)), - toLong(inputs(3)), - toLong(inputs(4)), - toLong(inputs(5)), - toLong(inputs(6)), - toLong(inputs(7))) - - case _ => - // it's the default approach, use O(64 * n), n is the length of inputs - interleaveBitsDefault(inputs.map(toByteArray)) - } - } - - private def interleave2Longs(l1: Long, l2: Long): Array[Byte] = { - // output 8 * 16 bits - val result = new Array[Byte](16) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toShort - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toShort - - var z = 0 - var j = 0 - while (j < 8) { - val x_masked = tmp1 & (1 << j) - val y_masked = tmp2 & (1 << j) - z |= (x_masked << j) - z |= (y_masked << (j + 1)) - j = j + 1 - } - result((7 - i) * 2 + 1) = (z & 0xFF).toByte - result((7 - i) * 2) = ((z >> 8) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave3Longs(l1: Long, l2: Long, l3: Long): Array[Byte] = { - // output 8 * 24 bits - val result = new Array[Byte](24) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toInt - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toInt - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toInt - - var z = 0 - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - z |= (r1_mask << (2 * j)) | (r2_mask << (2 * j + 1)) | (r3_mask << (2 * j + 2)) - j = j + 1 - } - result((7 - i) * 3 + 2) = (z & 0xFF).toByte - result((7 - i) * 3 + 1) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 3) = ((z >> 16) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave4Longs(l1: Long, l2: Long, l3: Long, l4: Long): Array[Byte] = { - // output 8 * 32 bits - val result = new Array[Byte](32) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toInt - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toInt - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toInt - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toInt - - var z = 0 - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - z |= (r1_mask << (3 * j)) | (r2_mask << (3 * j + 1)) | (r3_mask << (3 * j + 2)) | - (r4_mask << (3 * j + 3)) - j = j + 1 - } - result((7 - i) * 4 + 3) = (z & 0xFF).toByte - result((7 - i) * 4 + 2) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 4 + 1) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 4) = ((z >> 24) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave5Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long): Array[Byte] = { - // output 8 * 40 bits - val result = new Array[Byte](40) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - z |= (r1_mask << (4 * j)) | (r2_mask << (4 * j + 1)) | (r3_mask << (4 * j + 2)) | - (r4_mask << (4 * j + 3)) | (r5_mask << (4 * j + 4)) - j = j + 1 - } - result((7 - i) * 5 + 4) = (z & 0xFF).toByte - result((7 - i) * 5 + 3) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 5 + 2) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 5 + 1) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 5) = ((z >> 32) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave6Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long): Array[Byte] = { - // output 8 * 48 bits - val result = new Array[Byte](48) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - z |= (r1_mask << (5 * j)) | (r2_mask << (5 * j + 1)) | (r3_mask << (5 * j + 2)) | - (r4_mask << (5 * j + 3)) | (r5_mask << (5 * j + 4)) | (r6_mask << (5 * j + 5)) - j = j + 1 - } - result((7 - i) * 6 + 5) = (z & 0xFF).toByte - result((7 - i) * 6 + 4) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 6 + 3) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 6 + 2) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 6 + 1) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 6) = ((z >> 40) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave7Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long, - l7: Long): Array[Byte] = { - // output 8 * 56 bits - val result = new Array[Byte](56) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - val tmp7 = ((l7 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - val r7_mask = tmp7 & (1 << j) - z |= (r1_mask << (6 * j)) | (r2_mask << (6 * j + 1)) | (r3_mask << (6 * j + 2)) | - (r4_mask << (6 * j + 3)) | (r5_mask << (6 * j + 4)) | (r6_mask << (6 * j + 5)) | - (r7_mask << (6 * j + 6)) - j = j + 1 - } - result((7 - i) * 7 + 6) = (z & 0xFF).toByte - result((7 - i) * 7 + 5) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 7 + 4) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 7 + 3) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 7 + 2) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 7 + 1) = ((z >> 40) & 0xFF).toByte - result((7 - i) * 7) = ((z >> 48) & 0xFF).toByte - i = i + 1 - } - result - } - - private def interleave8Longs( - l1: Long, - l2: Long, - l3: Long, - l4: Long, - l5: Long, - l6: Long, - l7: Long, - l8: Long): Array[Byte] = { - // output 8 * 64 bits - val result = new Array[Byte](64) - var i = 0 - while (i < 8) { - val tmp1 = ((l1 >> (i * 8)) & 0xFF).toLong - val tmp2 = ((l2 >> (i * 8)) & 0xFF).toLong - val tmp3 = ((l3 >> (i * 8)) & 0xFF).toLong - val tmp4 = ((l4 >> (i * 8)) & 0xFF).toLong - val tmp5 = ((l5 >> (i * 8)) & 0xFF).toLong - val tmp6 = ((l6 >> (i * 8)) & 0xFF).toLong - val tmp7 = ((l7 >> (i * 8)) & 0xFF).toLong - val tmp8 = ((l8 >> (i * 8)) & 0xFF).toLong - - var z = 0L - var j = 0 - while (j < 8) { - val r1_mask = tmp1 & (1 << j) - val r2_mask = tmp2 & (1 << j) - val r3_mask = tmp3 & (1 << j) - val r4_mask = tmp4 & (1 << j) - val r5_mask = tmp5 & (1 << j) - val r6_mask = tmp6 & (1 << j) - val r7_mask = tmp7 & (1 << j) - val r8_mask = tmp8 & (1 << j) - z |= (r1_mask << (7 * j)) | (r2_mask << (7 * j + 1)) | (r3_mask << (7 * j + 2)) | - (r4_mask << (7 * j + 3)) | (r5_mask << (7 * j + 4)) | (r6_mask << (7 * j + 5)) | - (r7_mask << (7 * j + 6)) | (r8_mask << (7 * j + 7)) - j = j + 1 - } - result((7 - i) * 8 + 7) = (z & 0xFF).toByte - result((7 - i) * 8 + 6) = ((z >> 8) & 0xFF).toByte - result((7 - i) * 8 + 5) = ((z >> 16) & 0xFF).toByte - result((7 - i) * 8 + 4) = ((z >> 24) & 0xFF).toByte - result((7 - i) * 8 + 3) = ((z >> 32) & 0xFF).toByte - result((7 - i) * 8 + 2) = ((z >> 40) & 0xFF).toByte - result((7 - i) * 8 + 1) = ((z >> 48) & 0xFF).toByte - result((7 - i) * 8) = ((z >> 56) & 0xFF).toByte - i = i + 1 - } - result - } - - def interleaveBitsDefault(arrays: Array[Array[Byte]]): Array[Byte] = { - var totalLength = 0 - var maxLength = 0 - arrays.foreach { array => - totalLength += array.length - maxLength = maxLength.max(array.length * 8) - } - val result = new Array[Byte](totalLength) - var resultBit = 0 - - var bit = 0 - while (bit < maxLength) { - val bytePos = bit / 8 - val bitPos = bit % 8 - - for (arr <- arrays) { - val len = arr.length - if (bytePos < len) { - val resultBytePos = totalLength - 1 - resultBit / 8 - val resultBitPos = resultBit % 8 - result(resultBytePos) = - updatePos(result(resultBytePos), resultBitPos, arr(len - 1 - bytePos), bitPos) - resultBit += 1 - } - } - bit += 1 - } - result - } - - def updatePos(a: Byte, apos: Int, b: Byte, bpos: Int): Byte = { - var temp = (b & (1 << bpos)).toByte - if (apos > bpos) { - temp = (temp << (apos - bpos)).toByte - } else if (apos < bpos) { - temp = (temp >> (bpos - apos)).toByte - } - val atemp = (a & (1 << apos)).toByte - if (atemp == temp) { - return a - } - (a ^ (1 << apos)).toByte - } - - def toLong(a: Any): Long = { - a match { - case b: Boolean => (if (b) 1 else 0).toLong ^ BIT_64_MASK - case b: Byte => b.toLong ^ BIT_64_MASK - case s: Short => s.toLong ^ BIT_64_MASK - case i: Int => i.toLong ^ BIT_64_MASK - case l: Long => l ^ BIT_64_MASK - case f: Float => java.lang.Float.floatToRawIntBits(f).toLong ^ BIT_64_MASK - case d: Double => java.lang.Double.doubleToRawLongBits(d) ^ BIT_64_MASK - case str: UTF8String => str.getPrefix - case dec: Decimal => dec.toLong ^ BIT_64_MASK - case other: Any => - throw new KyuubiSQLExtensionException("Unsupported z-order type: " + other.getClass) - } - } - - def toByteArray(a: Any): Array[Byte] = { - a match { - case bo: Boolean => - booleanToByte(bo) - case b: Byte => - byteToByte(b) - case s: Short => - shortToByte(s) - case i: Int => - intToByte(i) - case l: Long => - longToByte(l) - case f: Float => - floatToByte(f) - case d: Double => - doubleToByte(d) - case str: UTF8String => - // truncate or padding str to 8 byte - paddingTo8Byte(str.getBytes) - case dec: Decimal => - longToByte(dec.toLong) - case other: Any => - throw new KyuubiSQLExtensionException("Unsupported z-order type: " + other.getClass) - } - } - - def booleanToByte(a: Boolean): Array[Byte] = { - if (a) { - byteToByte(1.toByte) - } else { - byteToByte(0.toByte) - } - } - - def byteToByte(a: Byte): Array[Byte] = { - val tmp = (a ^ BIT_8_MASK).toByte - Array(tmp) - } - - def shortToByte(a: Short): Array[Byte] = { - val tmp = a ^ BIT_16_MASK - Array(((tmp >> 8) & 0xFF).toByte, (tmp & 0xFF).toByte) - } - - def intToByte(a: Int): Array[Byte] = { - val result = new Array[Byte](4) - var i = 0 - val tmp = a ^ BIT_32_MASK - while (i <= 3) { - val offset = i * 8 - result(3 - i) = ((tmp >> offset) & 0xFF).toByte - i += 1 - } - result - } - - def longToByte(a: Long): Array[Byte] = { - val result = new Array[Byte](8) - var i = 0 - val tmp = a ^ BIT_64_MASK - while (i <= 7) { - val offset = i * 8 - result(7 - i) = ((tmp >> offset) & 0xFF).toByte - i += 1 - } - result - } - - def floatToByte(a: Float): Array[Byte] = { - val fi = jFloat.floatToRawIntBits(a) - intToByte(fi) - } - - def doubleToByte(a: Double): Array[Byte] = { - val dl = jDouble.doubleToRawLongBits(a) - longToByte(dl) - } - - def paddingTo8Byte(a: Array[Byte]): Array[Byte] = { - val len = a.length - if (len == 8) { - a - } else if (len > 8) { - val result = new Array[Byte](8) - System.arraycopy(a, 0, result, 0, 8) - result - } else { - val result = new Array[Byte](8) - System.arraycopy(a, 0, result, 8 - len, len) - result - } - } - - def defaultByteArrayValue(dataType: DataType): Array[Byte] = toByteArray { - defaultValue(dataType) - } - - def defaultValue(dataType: DataType): Any = { - dataType match { - case BooleanType => - true - case ByteType => - Byte.MaxValue - case ShortType => - Short.MaxValue - case IntegerType | DateType => - Int.MaxValue - case LongType | TimestampType | _: DecimalType => - Long.MaxValue - case FloatType => - Float.MaxValue - case DoubleType => - Double.MaxValue - case StringType => - // we pad string to 8 bytes so it's equal to long - UTF8String.fromBytes(longToByte(Long.MaxValue)) - case other: Any => - throw new KyuubiSQLExtensionException(s"Unsupported z-order type: ${other.catalogString}") - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala deleted file mode 100644 index 81873476cc4..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala +++ /dev/null @@ -1,289 +0,0 @@ -/* - * 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.spark.sql - -import scala.annotation.tailrec -import scala.collection.mutable -import scala.collection.mutable.ArrayBuffer - -import org.apache.spark.{ExecutorAllocationClient, MapOutputTrackerMaster, SparkContext, SparkEnv} -import org.apache.spark.internal.Logging -import org.apache.spark.resource.ResourceProfile -import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SortExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive._ -import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec -import org.apache.spark.sql.execution.command.DataWritingCommandExec -import org.apache.spark.sql.execution.datasources.WriteFilesExec -import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec -import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec} - -import org.apache.kyuubi.sql.{KyuubiSQLConf, WriteUtils} - -/** - * This rule assumes the final write stage has less cores requirement than previous, otherwise - * this rule would take no effect. - * - * It provide a feature: - * 1. Kill redundant executors before running final write stage - */ -case class FinalStageResourceManager(session: SparkSession) - extends Rule[SparkPlan] with FinalRebalanceStageHelper { - override def apply(plan: SparkPlan): SparkPlan = { - if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED)) { - return plan - } - - if (!WriteUtils.isWrite(session, plan)) { - return plan - } - - val sc = session.sparkContext - val dra = sc.getConf.getBoolean("spark.dynamicAllocation.enabled", false) - val coresPerExecutor = sc.getConf.getInt("spark.executor.cores", 1) - val minExecutors = sc.getConf.getInt("spark.dynamicAllocation.minExecutors", 0) - val maxExecutors = sc.getConf.getInt("spark.dynamicAllocation.maxExecutors", Int.MaxValue) - val factor = conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_PARTITION_FACTOR) - val hasImprovementRoom = maxExecutors - 1 > minExecutors * factor - // Fast fail if: - // 1. DRA off - // 2. only work with yarn and k8s - // 3. maxExecutors is not bigger than minExecutors * factor - if (!dra || !sc.schedulerBackend.isInstanceOf[CoarseGrainedSchedulerBackend] || - !hasImprovementRoom) { - return plan - } - - val stageOpt = findFinalRebalanceStage(plan) - if (stageOpt.isEmpty) { - return plan - } - - // It's not safe to kill executors if this plan contains table cache. - // If the executor loses then the rdd would re-compute those partition. - if (hasTableCache(plan) && - conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_SKIP_KILLING_EXECUTORS_FOR_TABLE_CACHE)) { - return plan - } - - // TODO: move this to query stage optimizer when updating Spark to 3.5.x - // Since we are in `prepareQueryStage`, the AQE shuffle read has not been applied. - // So we need to apply it by self. - val shuffleRead = queryStageOptimizerRules.foldLeft(stageOpt.get.asInstanceOf[SparkPlan]) { - case (latest, rule) => rule.apply(latest) - } - val (targetCores, stage) = shuffleRead match { - case AQEShuffleReadExec(stage: ShuffleQueryStageExec, partitionSpecs) => - (partitionSpecs.length, stage) - case stage: ShuffleQueryStageExec => - // we can still kill executors if no AQE shuffle read, e.g., `.repartition(2)` - (stage.shuffle.numPartitions, stage) - case _ => - // it should never happen in current Spark, but to be safe do nothing if happens - logWarning("BUG, Please report to Apache Kyuubi community") - return plan - } - // The condition whether inject custom resource profile: - // - target executors < active executors - // - active executors - target executors > min executors - val numActiveExecutors = sc.getExecutorIds().length - val targetExecutors = (math.ceil(targetCores.toFloat / coresPerExecutor) * factor).toInt - .max(1) - val hasBenefits = targetExecutors < numActiveExecutors && - (numActiveExecutors - targetExecutors) > minExecutors - logInfo(s"The snapshot of current executors view, " + - s"active executors: $numActiveExecutors, min executor: $minExecutors, " + - s"target executors: $targetExecutors, has benefits: $hasBenefits") - if (hasBenefits) { - val shuffleId = stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleDependency.shuffleId - val numReduce = stage.plan.asInstanceOf[ShuffleExchangeExec].numPartitions - // Now, there is only a final rebalance stage waiting to execute and all tasks of previous - // stage are finished. Kill redundant existed executors eagerly so the tasks of final - // stage can be centralized scheduled. - killExecutors(sc, targetExecutors, shuffleId, numReduce) - } - - plan - } - - /** - * The priority of kill executors follow: - * 1. kill executor who is younger than other (The older the JIT works better) - * 2. kill executor who produces less shuffle data first - */ - private def findExecutorToKill( - sc: SparkContext, - targetExecutors: Int, - shuffleId: Int, - numReduce: Int): Seq[String] = { - val tracker = SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] - val shuffleStatusOpt = tracker.shuffleStatuses.get(shuffleId) - if (shuffleStatusOpt.isEmpty) { - return Seq.empty - } - val shuffleStatus = shuffleStatusOpt.get - val executorToBlockSize = new mutable.HashMap[String, Long] - shuffleStatus.withMapStatuses { mapStatus => - mapStatus.foreach { status => - var i = 0 - var sum = 0L - while (i < numReduce) { - sum += status.getSizeForBlock(i) - i += 1 - } - executorToBlockSize.getOrElseUpdate(status.location.executorId, sum) - } - } - - val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] - val executorsWithRegistrationTs = backend.getExecutorsWithRegistrationTs() - val existedExecutors = executorsWithRegistrationTs.keys.toSet - val expectedNumExecutorToKill = existedExecutors.size - targetExecutors - if (expectedNumExecutorToKill < 1) { - return Seq.empty - } - - val executorIdsToKill = new ArrayBuffer[String]() - // We first kill executor who does not hold shuffle block. It would happen because - // the last stage is running fast and finished in a short time. The existed executors are - // from previous stages that have not been killed by DRA, so we can not find it by tracking - // shuffle status. - // We should evict executors by their alive time first and retain all of executors which - // have better locality for shuffle block. - executorsWithRegistrationTs.toSeq.sortBy(_._2).foreach { case (id, _) => - if (executorIdsToKill.length < expectedNumExecutorToKill && - !executorToBlockSize.contains(id)) { - executorIdsToKill.append(id) - } - } - - // Evict the rest executors according to the shuffle block size - executorToBlockSize.toSeq.sortBy(_._2).foreach { case (id, _) => - if (executorIdsToKill.length < expectedNumExecutorToKill && existedExecutors.contains(id)) { - executorIdsToKill.append(id) - } - } - - executorIdsToKill.toSeq - } - - private def killExecutors( - sc: SparkContext, - targetExecutors: Int, - shuffleId: Int, - numReduce: Int): Unit = { - val executorAllocationClient = sc.schedulerBackend.asInstanceOf[ExecutorAllocationClient] - - val executorsToKill = - if (conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL)) { - executorAllocationClient.getExecutorIds() - } else { - findExecutorToKill(sc, targetExecutors, shuffleId, numReduce) - } - logInfo(s"Request to kill executors, total count ${executorsToKill.size}, " + - s"[${executorsToKill.mkString(", ")}].") - if (executorsToKill.isEmpty) { - return - } - - // Note, `SparkContext#killExecutors` does not allow with DRA enabled, - // see `https://github.com/apache/spark/pull/20604`. - // It may cause the status in `ExecutorAllocationManager` inconsistent with - // `CoarseGrainedSchedulerBackend` for a while. But it should be synchronous finally. - // - // We should adjust target num executors, otherwise `YarnAllocator` might re-request original - // target executors if DRA has not updated target executors yet. - // Note, DRA would re-adjust executors if there are more tasks to be executed, so we are safe. - // - // * We kill executor - // * YarnAllocator re-request target executors - // * DRA can not release executors since they are new added - // ----------------------------------------------------------------> timeline - executorAllocationClient.killExecutors( - executorIds = executorsToKill, - adjustTargetNumExecutors = true, - countFailures = false, - force = false) - - FinalStageResourceManager.getAdjustedTargetExecutors(sc) - .filter(_ < targetExecutors).foreach { adjustedExecutors => - val delta = targetExecutors - adjustedExecutors - logInfo(s"Target executors after kill ($adjustedExecutors) is lower than required " + - s"($targetExecutors). Requesting $delta additional executor(s).") - executorAllocationClient.requestExecutors(delta) - } - } - - @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq( - OptimizeSkewInRebalancePartitions, - CoalesceShufflePartitions(session), - OptimizeShuffleWithLocalRead) -} - -object FinalStageResourceManager extends Logging { - - private[sql] def getAdjustedTargetExecutors(sc: SparkContext): Option[Int] = { - sc.schedulerBackend match { - case schedulerBackend: CoarseGrainedSchedulerBackend => - try { - val field = classOf[CoarseGrainedSchedulerBackend] - .getDeclaredField("requestedTotalExecutorsPerResourceProfile") - field.setAccessible(true) - schedulerBackend.synchronized { - val requestedTotalExecutorsPerResourceProfile = - field.get(schedulerBackend).asInstanceOf[mutable.HashMap[ResourceProfile, Int]] - val defaultRp = sc.resourceProfileManager.defaultResourceProfile - requestedTotalExecutorsPerResourceProfile.get(defaultRp) - } - } catch { - case e: Exception => - logWarning("Failed to get requestedTotalExecutors of Default ResourceProfile", e) - None - } - case _ => None - } - } -} - -trait FinalRebalanceStageHelper extends AdaptiveSparkPlanHelper { - @tailrec - final protected def findFinalRebalanceStage(plan: SparkPlan): Option[ShuffleQueryStageExec] = { - plan match { - case write: DataWritingCommandExec => findFinalRebalanceStage(write.child) - case write: V2TableWriteExec => findFinalRebalanceStage(write.child) - case write: WriteFilesExec => findFinalRebalanceStage(write.child) - case p: ProjectExec => findFinalRebalanceStage(p.child) - case f: FilterExec => findFinalRebalanceStage(f.child) - case s: SortExec if !s.global => findFinalRebalanceStage(s.child) - case stage: ShuffleQueryStageExec - if stage.isMaterialized && stage.mapStats.isDefined && - stage.plan.isInstanceOf[ShuffleExchangeExec] && - stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleOrigin != ENSURE_REQUIREMENTS => - Some(stage) - case _ => None - } - } - - final protected def hasTableCache(plan: SparkPlan): Boolean = { - find(plan) { - case _: InMemoryTableScanExec => true - case _ => false - }.isDefined - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala deleted file mode 100644 index 64421d6bfab..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.{CustomResourceProfileExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive._ - -import org.apache.kyuubi.sql.{KyuubiSQLConf, WriteUtils} - -/** - * Inject custom resource profile for final write stage, so we can specify custom - * executor resource configs. - */ -case class InjectCustomResourceProfile(session: SparkSession) - extends Rule[SparkPlan] with FinalRebalanceStageHelper { - override def apply(plan: SparkPlan): SparkPlan = { - if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED)) { - return plan - } - - if (!WriteUtils.isWrite(session, plan)) { - return plan - } - - val stage = findFinalRebalanceStage(plan) - if (stage.isEmpty) { - return plan - } - - // TODO: Ideally, We can call `CoarseGrainedSchedulerBackend.requestTotalExecutors` eagerly - // to reduce the task submit pending time, but it may lose task locality. - // - // By default, it would request executors when catch stage submit event. - injectCustomResourceProfile(plan, stage.get.id) - } - - private def injectCustomResourceProfile(plan: SparkPlan, id: Int): SparkPlan = { - plan match { - case stage: ShuffleQueryStageExec if stage.id == id => - CustomResourceProfileExec(stage) - case _ => plan.mapChildren(child => injectCustomResourceProfile(child, id)) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala deleted file mode 100644 index ce496eb474c..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/PruneFileSourcePartitionHelper.scala +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExpressionSet, PredicateHelper, SubqueryExpression} -import org.apache.spark.sql.catalyst.plans.logical.LeafNode -import org.apache.spark.sql.execution.datasources.DataSourceStrategy -import org.apache.spark.sql.types.StructType - -trait PruneFileSourcePartitionHelper extends PredicateHelper { - - def getPartitionKeyFiltersAndDataFilters( - sparkSession: SparkSession, - relation: LeafNode, - partitionSchema: StructType, - filters: Seq[Expression], - output: Seq[AttributeReference]): (ExpressionSet, Seq[Expression]) = { - val normalizedFilters = DataSourceStrategy.normalizeExprs( - filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f)), - output) - val partitionColumns = - relation.resolve(partitionSchema, sparkSession.sessionState.analyzer.resolver) - val partitionSet = AttributeSet(partitionColumns) - val (partitionFilters, dataFilters) = normalizedFilters.partition(f => - f.references.subsetOf(partitionSet)) - val extraPartitionFilter = - dataFilters.flatMap(extractPredicatesWithinOutputSet(_, partitionSet)) - - (ExpressionSet(partitionFilters ++ extraPartitionFilter), dataFilters) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala deleted file mode 100644 index 043d6496b22..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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.spark.sql.execution - -import org.apache.spark.network.util.{ByteUnit, JavaUtils} -import org.apache.spark.rdd.RDD -import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfileBuilder} -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} -import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils - -import org.apache.kyuubi.sql.KyuubiSQLConf._ - -/** - * This node wraps the final executed plan and inject custom resource profile to the RDD. - * It assumes that, the produced RDD would create the `ResultStage` in `DAGScheduler`, - * so it makes resource isolation between previous and final stage. - * - * Note that, Spark does not support config `minExecutors` for each resource profile. - * Which means, it would retain `minExecutors` for each resource profile. - * So, suggest set `spark.dynamicAllocation.minExecutors` to 0 if enable this feature. - */ -case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode { - override def output: Seq[Attribute] = child.output - override def outputPartitioning: Partitioning = child.outputPartitioning - override def outputOrdering: Seq[SortOrder] = child.outputOrdering - override def supportsColumnar: Boolean = child.supportsColumnar - override def supportsRowBased: Boolean = child.supportsRowBased - override protected def doCanonicalize(): SparkPlan = child.canonicalized - - private val executorCores = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_CORES).getOrElse( - sparkContext.getConf.getInt("spark.executor.cores", 1)) - private val executorMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY).getOrElse( - sparkContext.getConf.get("spark.executor.memory", "2G")) - private val executorMemoryOverhead = - conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD) - .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G")) - private val executorOffHeapMemory = - if (sparkContext.getConf.getBoolean("spark.memory.offHeap.enabled", false)) { - conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY) - } else { - None - } - - override lazy val metrics: Map[String, SQLMetric] = { - val base = Map( - "executorCores" -> SQLMetrics.createMetric(sparkContext, "executor cores"), - "executorMemory" -> SQLMetrics.createMetric(sparkContext, "executor memory (MiB)"), - "executorMemoryOverhead" -> SQLMetrics.createMetric( - sparkContext, - "executor memory overhead (MiB)")) - val addition = executorOffHeapMemory.map(_ => - "executorOffHeapMemory" -> - SQLMetrics.createMetric(sparkContext, "executor off heap memory (MiB)")).toMap - base ++ addition - } - - private def wrapResourceProfile[T](rdd: RDD[T]): RDD[T] = { - if (Utils.isTesting) { - // do nothing for local testing - return rdd - } - - metrics("executorCores") += executorCores - metrics("executorMemory") += JavaUtils.byteStringAs(executorMemory, ByteUnit.MiB) - metrics("executorMemoryOverhead") += JavaUtils.byteStringAs( - executorMemoryOverhead, - ByteUnit.MiB) - executorOffHeapMemory.foreach(m => - metrics("executorOffHeapMemory") += JavaUtils.byteStringAs(m, ByteUnit.MiB)) - - val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) - SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) - - val resourceProfileBuilder = new ResourceProfileBuilder() - val executorResourceRequests = new ExecutorResourceRequests() - executorResourceRequests.cores(executorCores) - executorResourceRequests.memory(executorMemory) - executorResourceRequests.memoryOverhead(executorMemoryOverhead) - executorOffHeapMemory.foreach(executorResourceRequests.offHeapMemory) - resourceProfileBuilder.require(executorResourceRequests) - rdd.withResources(resourceProfileBuilder.build()) - rdd - } - - override protected def doExecute(): RDD[InternalRow] = { - val rdd = child.execute() - wrapResourceProfile(rdd) - } - - override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { - val rdd = child.executeColumnar() - wrapResourceProfile(rdd) - } - - override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = { - this.copy(child = newChild) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/resources/log4j2-test.xml b/extensions/spark/kyuubi-extension-spark-3-4/src/test/resources/log4j2-test.xml deleted file mode 100644 index 3110216c17c..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/resources/log4j2-test.xml +++ /dev/null @@ -1,43 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala deleted file mode 100644 index bbc61fb4408..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/DropIgnoreNonexistentSuite.scala +++ /dev/null @@ -1,45 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.plans.logical.{DropNamespace, NoopCommand} -import org.apache.spark.sql.execution.command._ - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class DropIgnoreNonexistentSuite extends KyuubiSparkSQLExtensionTest { - - test("drop ignore nonexistent") { - withSQLConf(KyuubiSQLConf.DROP_IGNORE_NONEXISTENT.key -> "true") { - // drop nonexistent database - val df1 = sql("DROP DATABASE nonexistent_database") - assert(df1.queryExecution.analyzed.asInstanceOf[DropNamespace].ifExists == true) - - // drop nonexistent function - val df4 = sql("DROP FUNCTION nonexistent_function") - assert(df4.queryExecution.analyzed.isInstanceOf[NoopCommand]) - - // drop nonexistent PARTITION - withTable("test") { - sql("CREATE TABLE IF NOT EXISTS test(i int) PARTITIONED BY (p int)") - val df5 = sql("ALTER TABLE test DROP PARTITION (p = 1)") - assert(df5.queryExecution.analyzed - .asInstanceOf[AlterTableDropPartitionCommand].ifExists == true) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala deleted file mode 100644 index 96c8ae6e8b0..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala +++ /dev/null @@ -1,203 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.{FinalStageConfigIsolation, KyuubiSQLConf} - -class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - test("final stage config set reset check") { - withSQLConf( - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "false", - "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum" -> "1", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "100") { - // use loop to double check final stage config doesn't affect the sql query each other - (1 to 3).foreach { _ => - sql("SELECT COUNT(*) FROM VALUES(1) as t(c)").collect() - assert(spark.sessionState.conf.getConfString( - "spark.sql.previousStage.adaptive.coalescePartitions.minPartitionNum") === - FinalStageConfigIsolation.INTERNAL_UNSET_CONFIG_TAG) - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.coalescePartitions.minPartitionNum") === - "1") - assert(spark.sessionState.conf.getConfString( - "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum") === - "1") - - // 64MB - assert(spark.sessionState.conf.getConfString( - "spark.sql.previousStage.adaptive.advisoryPartitionSizeInBytes") === - "67108864b") - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.advisoryPartitionSizeInBytes") === - "100") - assert(spark.sessionState.conf.getConfString( - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes") === - "100") - } - - sql("SET spark.sql.adaptive.advisoryPartitionSizeInBytes=1") - assert(spark.sessionState.conf.getConfString( - "spark.sql.adaptive.advisoryPartitionSizeInBytes") === - "1") - assert(!spark.sessionState.conf.contains( - "spark.sql.previousStage.adaptive.advisoryPartitionSizeInBytes")) - - sql("SET a=1") - assert(spark.sessionState.conf.getConfString("a") === "1") - - sql("RESET spark.sql.adaptive.coalescePartitions.minPartitionNum") - assert(!spark.sessionState.conf.contains( - "spark.sql.adaptive.coalescePartitions.minPartitionNum")) - assert(!spark.sessionState.conf.contains( - "spark.sql.previousStage.adaptive.coalescePartitions.minPartitionNum")) - - sql("RESET a") - assert(!spark.sessionState.conf.contains("a")) - } - } - - test("final stage config isolation") { - def checkPartitionNum( - sqlString: String, - previousPartitionNum: Int, - finalPartitionNum: Int): Unit = { - val df = sql(sqlString) - df.collect() - val shuffleReaders = collect(df.queryExecution.executedPlan) { - case customShuffleReader: AQEShuffleReadExec => customShuffleReader - } - assert(shuffleReaders.nonEmpty) - // reorder stage by stage id to ensure we get the right stage - val sortedShuffleReaders = shuffleReaders.sortWith { - case (s1, s2) => - s1.child.asInstanceOf[QueryStageExec].id < s2.child.asInstanceOf[QueryStageExec].id - } - if (sortedShuffleReaders.length > 1) { - assert(sortedShuffleReaders.head.partitionSpecs.length === previousPartitionNum) - } - assert(sortedShuffleReaders.last.partitionSpecs.length === finalPartitionNum) - assert(df.rdd.partitions.length === finalPartitionNum) - } - - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", - SQLConf.SHUFFLE_PARTITIONS.key -> "3", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "false", - "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "1", - "spark.sql.adaptive.coalescePartitions.minPartitionSize" -> "1", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "10000000") { - - // use loop to double check final stage config doesn't affect the sql query each other - (1 to 3).foreach { _ => - checkPartitionNum( - "SELECT c1, count(*) FROM t1 GROUP BY c1", - 1, - 1) - - checkPartitionNum( - "SELECT c2, count(*) FROM (SELECT c1, count(*) as c2 FROM t1 GROUP BY c1) GROUP BY c2", - 3, - 1) - - checkPartitionNum( - "SELECT t1.c1, count(*) FROM t1 JOIN t2 ON t1.c2 = t2.c2 GROUP BY t1.c1", - 3, - 1) - - checkPartitionNum( - """ - | SELECT /*+ REPARTITION */ - | t1.c1, count(*) FROM t1 - | JOIN t2 ON t1.c2 = t2.c2 - | JOIN t3 ON t1.c1 = t3.c1 - | GROUP BY t1.c1 - |""".stripMargin, - 3, - 1) - - // one shuffle reader - checkPartitionNum( - """ - | SELECT /*+ BROADCAST(t1) */ - | t1.c1, t2.c2 FROM t1 - | JOIN t2 ON t1.c2 = t2.c2 - | DISTRIBUTE BY c1 - |""".stripMargin, - 1, - 1) - - // test ReusedExchange - checkPartitionNum( - """ - |SELECT /*+ REPARTITION */ t0.c2 FROM ( - |SELECT t1.c1, (count(*) + c1) as c2 FROM t1 GROUP BY t1.c1 - |) t0 JOIN ( - |SELECT t1.c1, (count(*) + c1) as c2 FROM t1 GROUP BY t1.c1 - |) t1 ON t0.c2 = t1.c2 - |""".stripMargin, - 3, - 1) - - // one shuffle reader - checkPartitionNum( - """ - |SELECT t0.c1 FROM ( - |SELECT t1.c1 FROM t1 GROUP BY t1.c1 - |) t0 JOIN ( - |SELECT t1.c1 FROM t1 GROUP BY t1.c1 - |) t1 ON t0.c1 = t1.c1 - |""".stripMargin, - 1, - 1) - } - } - } - - test("final stage config isolation write only") { - withSQLConf( - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION_WRITE_ONLY.key -> "true", - "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "7") { - sql("set spark.sql.adaptive.advisoryPartitionSizeInBytes=5") - sql("SELECT * FROM t1").count() - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("5")) - - withTable("tmp") { - sql("CREATE TABLE t1 USING PARQUET SELECT /*+ repartition */ 1 AS c1, 'a' AS c2") - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("7")) - } - - sql("SELECT * FROM t1").count() - assert(spark.conf.getOption("spark.sql.adaptive.advisoryPartitionSizeInBytes") - .contains("5")) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala deleted file mode 100644 index 4b9991ef6f2..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/FinalStageResourceManagerSuite.scala +++ /dev/null @@ -1,62 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.scalatest.time.{Minutes, Span} - -import org.apache.kyuubi.sql.KyuubiSQLConf -import org.apache.kyuubi.tags.SparkLocalClusterTest - -@SparkLocalClusterTest -class FinalStageResourceManagerSuite extends KyuubiSparkSQLExtensionTest { - - override def sparkConf(): SparkConf = { - // It is difficult to run spark in local-cluster mode when spark.testing is set. - sys.props.remove("spark.testing") - - super.sparkConf().set("spark.master", "local-cluster[3, 1, 1024]") - .set("spark.dynamicAllocation.enabled", "true") - .set("spark.dynamicAllocation.initialExecutors", "3") - .set("spark.dynamicAllocation.minExecutors", "1") - .set("spark.dynamicAllocation.shuffleTracking.enabled", "true") - .set(KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key, "true") - .set(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED.key, "true") - } - - test("[KYUUBI #5136][Bug] Final Stage hangs forever") { - // Prerequisite to reproduce the bug: - // 1. Dynamic allocation is enabled. - // 2. Dynamic allocation min executors is 1. - // 3. target executors < active executors. - // 4. No active executor is left after FinalStageResourceManager killed executors. - // This is possible because FinalStageResourceManager retained executors may already be - // requested to be killed but not died yet. - // 5. Final Stage required executors is 1. - withSQLConf( - (KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_KILL_ALL.key, "true")) { - withTable("final_stage") { - eventually(timeout(Span(10, Minutes))) { - sql( - "CREATE TABLE final_stage AS SELECT id, count(*) as num FROM (SELECT 0 id) GROUP BY id") - } - assert(FinalStageResourceManager.getAdjustedTargetExecutors(spark.sparkContext).get == 1) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala deleted file mode 100644 index b0767b18708..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} -import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class InjectResourceProfileSuite extends KyuubiSparkSQLExtensionTest { - private def checkCustomResourceProfile(sqlString: String, exists: Boolean): Unit = { - @volatile var lastEvent: SparkListenerSQLAdaptiveExecutionUpdate = null - val listener = new SparkListener { - override def onOtherEvent(event: SparkListenerEvent): Unit = { - event match { - case e: SparkListenerSQLAdaptiveExecutionUpdate => lastEvent = e - case _ => - } - } - } - - spark.sparkContext.addSparkListener(listener) - try { - sql(sqlString).collect() - spark.sparkContext.listenerBus.waitUntilEmpty() - assert(lastEvent != null) - var current = lastEvent.sparkPlanInfo - var shouldStop = false - while (!shouldStop) { - if (current.nodeName != "CustomResourceProfile") { - if (current.children.isEmpty) { - assert(!exists) - shouldStop = true - } else { - current = current.children.head - } - } else { - assert(exists) - shouldStop = true - } - } - } finally { - spark.sparkContext.removeSparkListener(listener) - } - } - - test("Inject resource profile") { - withTable("t") { - withSQLConf( - "spark.sql.adaptive.forceApply" -> "true", - KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", - KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED.key -> "true") { - - sql("CREATE TABLE t (c1 int, c2 string) USING PARQUET") - - checkCustomResourceProfile("INSERT INTO TABLE t VALUES(1, 'a')", false) - checkCustomResourceProfile("SELECT 1", false) - checkCustomResourceProfile( - "INSERT INTO TABLE t SELECT /*+ rebalance */ * FROM VALUES(1, 'a')", - true) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala deleted file mode 100644 index f0d38465734..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala +++ /dev/null @@ -1,19 +0,0 @@ -/* - * 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.spark.sql - -class InsertShuffleNodeBeforeJoinSuite extends InsertShuffleNodeBeforeJoinSuiteBase diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala deleted file mode 100644 index c657dee49f3..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuiteBase.scala +++ /dev/null @@ -1,98 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeLike} -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} - -import org.apache.kyuubi.sql.KyuubiSQLConf - -trait InsertShuffleNodeBeforeJoinSuiteBase extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - override def sparkConf(): SparkConf = { - super.sparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLCommonExtension") - } - - test("force shuffle before join") { - def checkShuffleNodeNum(sqlString: String, num: Int): Unit = { - var expectedResult: Seq[Row] = Seq.empty - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - expectedResult = sql(sqlString).collect() - } - val df = sql(sqlString) - checkAnswer(df, expectedResult) - assert( - collect(df.queryExecution.executedPlan) { - case shuffle: ShuffleExchangeLike if shuffle.shuffleOrigin == ENSURE_REQUIREMENTS => - shuffle - }.size == num) - } - - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - KyuubiSQLConf.FORCE_SHUFFLE_BEFORE_JOIN.key -> "true") { - Seq("SHUFFLE_HASH", "MERGE").foreach { joinHint => - // positive case - checkShuffleNodeNum( - s""" - |SELECT /*+ $joinHint(t2, t3) */ t1.c1, t1.c2, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN t3 ON t1.c1 = t3.c1 - | """.stripMargin, - 4) - - // negative case - checkShuffleNodeNum( - s""" - |SELECT /*+ $joinHint(t2, t3) */ t1.c1, t1.c2, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN t3 ON t1.c2 = t3.c2 - | """.stripMargin, - 4) - } - - checkShuffleNodeNum( - """ - |SELECT t1.c1, t2.c1, t3.c2 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN ( - | SELECT c2, count(*) FROM t1 GROUP BY c2 - | ) t3 ON t1.c1 = t3.c2 - | """.stripMargin, - 5) - - checkShuffleNodeNum( - """ - |SELECT t1.c1, t2.c1, t3.c1 from t1 - | JOIN t2 ON t1.c1 = t2.c1 - | JOIN ( - | SELECT c1, count(*) FROM t1 GROUP BY c1 - | ) t3 ON t1.c1 = t3.c1 - | """.stripMargin, - 5) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala deleted file mode 100644 index 996bef763a2..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala +++ /dev/null @@ -1,121 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.hadoop.hive.conf.HiveConf.ConfVars -import org.apache.spark.SparkConf -import org.apache.spark.sql.execution.QueryExecution -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.command.{DataWritingCommand, DataWritingCommandExec} -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.test.SQLTestData.TestData -import org.apache.spark.sql.test.SQLTestUtils -import org.apache.spark.sql.util.QueryExecutionListener -import org.apache.spark.util.Utils - -trait KyuubiSparkSQLExtensionTest extends QueryTest - with SQLTestUtils - with AdaptiveSparkPlanHelper { - sys.props.put("spark.testing", "1") - - private var _spark: Option[SparkSession] = None - protected def spark: SparkSession = _spark.getOrElse { - throw new RuntimeException("test spark session don't initial before using it.") - } - - override protected def beforeAll(): Unit = { - if (_spark.isEmpty) { - _spark = Option(SparkSession.builder() - .master("local[1]") - .config(sparkConf) - .enableHiveSupport() - .getOrCreate()) - } - super.beforeAll() - } - - override protected def afterAll(): Unit = { - super.afterAll() - cleanupData() - _spark.foreach(_.stop) - } - - protected def setupData(): Unit = { - val self = spark - import self.implicits._ - spark.sparkContext.parallelize( - (1 to 100).map(i => TestData(i, i.toString)), - 10) - .toDF("c1", "c2").createOrReplaceTempView("t1") - spark.sparkContext.parallelize( - (1 to 10).map(i => TestData(i, i.toString)), - 5) - .toDF("c1", "c2").createOrReplaceTempView("t2") - spark.sparkContext.parallelize( - (1 to 50).map(i => TestData(i, i.toString)), - 2) - .toDF("c1", "c2").createOrReplaceTempView("t3") - } - - private def cleanupData(): Unit = { - spark.sql("DROP VIEW IF EXISTS t1") - spark.sql("DROP VIEW IF EXISTS t2") - spark.sql("DROP VIEW IF EXISTS t3") - } - - def sparkConf(): SparkConf = { - val basePath = Utils.createTempDir() + "/" + getClass.getCanonicalName - val metastorePath = basePath + "/metastore_db" - val warehousePath = basePath + "/warehouse" - new SparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLExtension") - .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true") - .set("spark.hadoop.hive.exec.dynamic.partition.mode", "nonstrict") - .set("spark.hadoop.hive.metastore.client.capability.check", "false") - .set( - ConfVars.METASTORECONNECTURLKEY.varname, - s"jdbc:derby:;databaseName=$metastorePath;create=true") - .set(StaticSQLConf.WAREHOUSE_PATH, warehousePath) - .set("spark.ui.enabled", "false") - } - - def withListener(sqlString: String)(callback: DataWritingCommand => Unit): Unit = { - withListener(sql(sqlString))(callback) - } - - def withListener(df: => DataFrame)(callback: DataWritingCommand => Unit): Unit = { - val listener = new QueryExecutionListener { - override def onFailure(f: String, qe: QueryExecution, e: Exception): Unit = {} - - override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = { - qe.executedPlan match { - case write: DataWritingCommandExec => callback(write.cmd) - case _ => - } - } - } - spark.listenerManager.register(listener) - try { - df.collect() - sparkContext.listenerBus.waitUntilEmpty() - } finally { - spark.listenerManager.unregister(listener) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala deleted file mode 100644 index 9aa75f53dc5..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala +++ /dev/null @@ -1,319 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, RebalancePartitions, Sort} -import org.apache.spark.sql.execution.command.{DataWritingCommand, InsertIntoDataSourceDirCommand} -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.hive.HiveUtils -import org.apache.spark.sql.hive.execution.{InsertIntoHiveDirCommand, InsertIntoHiveTable} - -import org.apache.kyuubi.sql.KyuubiSQLConf - -class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest { - - test("check rebalance exists") { - def check( - df: => DataFrame, - expectedRebalanceNumEnabled: Int = 1, - expectedRebalanceNumDisabled: Int = 0): Unit = { - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withListener(df) { write => - assert(write.collect { - case r: RebalancePartitions => r - }.size == expectedRebalanceNumEnabled) - } - } - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "false") { - withListener(df) { write => - assert(write.collect { - case r: RebalancePartitions => r - }.size == expectedRebalanceNumDisabled) - } - } - } - - // It's better to set config explicitly in case of we change the default value. - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2='a') " + - "SELECT * FROM VALUES(1),(2) AS t(c1)")) - } - - withTable("tmp1", "tmp2") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - sql(s"CREATE TABLE tmp2 (c1 int) $storage PARTITIONED BY (c2 string)") - check( - sql( - """FROM VALUES(1),(2) - |INSERT INTO TABLE tmp1 PARTITION(c2='a') SELECT * - |INSERT INTO TABLE tmp2 PARTITION(c2='a') SELECT * - |""".stripMargin), - 2) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check(sql("INSERT INTO TABLE tmp1 SELECT * FROM VALUES(1),(2),(3) AS t(c1)")) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check( - sql("INSERT INTO TABLE tmp1 SELECT /*+ REBALANCE */ * FROM VALUES(1),(2),(3) AS t(c1)"), - 1, - 1) - } - - withTable("tmp1", "tmp2") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - sql(s"CREATE TABLE tmp2 (c1 int) $storage") - check( - sql( - """FROM VALUES(1),(2),(3) - |INSERT INTO TABLE tmp1 SELECT * - |INSERT INTO TABLE tmp2 SELECT * - |""".stripMargin), - 2) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage AS SELECT * FROM VALUES(1),(2),(3) AS t(c1)") - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 $storage PARTITIONED BY(c2) AS " + - s"SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)") - } - } - } - } - - test("check rebalance does not exists") { - def check(df: DataFrame): Unit = { - withListener(df) { write => - assert(write.collect { - case r: RebalancePartitions => r - }.isEmpty) - } - } - - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - // test no write command - check(sql("SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - check(sql("SELECT count(*) FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - - // test not supported plan - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT /*+ repartition(10) */ * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2) ORDER BY c1")) - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2) LIMIT 10")) - } - } - - withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "false") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - check(sql("INSERT INTO TABLE tmp1 PARTITION(c2) " + - "SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")) - } - - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage") - check(sql("INSERT INTO TABLE tmp1 SELECT * FROM VALUES(1),(2),(3) AS t(c1)")) - } - } - } - } - - test("test dynamic partition write") { - def checkRepartitionExpression(sqlString: String): Unit = { - withListener(sqlString) { write => - assert(write.isInstanceOf[InsertIntoHiveTable]) - assert(write.collect { - case r: RebalancePartitions if r.partitionExpressions.size == 1 => - assert(r.partitionExpressions.head.asInstanceOf[Attribute].name === "c2") - r - }.size == 1) - } - } - - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - Seq("USING PARQUET", "").foreach { storage => - withTable("tmp1") { - sql(s"CREATE TABLE tmp1 (c1 int) $storage PARTITIONED BY (c2 string)") - checkRepartitionExpression("INSERT INTO TABLE tmp1 SELECT 1 as c1, 'a' as c2 ") - } - - withTable("tmp1") { - checkRepartitionExpression( - "CREATE TABLE tmp1 PARTITIONED BY(C2) SELECT 1 as c1, 'a' as c2") - } - } - } - } - - test("OptimizedCreateHiveTableAsSelectCommand") { - withSQLConf( - HiveUtils.CONVERT_METASTORE_PARQUET.key -> "true", - HiveUtils.CONVERT_METASTORE_CTAS.key -> "true", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTable("t") { - withListener("CREATE TABLE t STORED AS parquet AS SELECT 1 as a") { write => - assert(write.isInstanceOf[InsertIntoHadoopFsRelationCommand]) - assert(write.collect { - case _: RebalancePartitions => true - }.size == 1) - } - } - } - } - - test("Infer rebalance and sorder orders") { - def checkShuffleAndSort(dataWritingCommand: LogicalPlan, sSize: Int, rSize: Int): Unit = { - assert(dataWritingCommand.isInstanceOf[DataWritingCommand]) - val plan = dataWritingCommand.asInstanceOf[DataWritingCommand].query - assert(plan.collect { - case s: Sort => s - }.size == sSize) - assert(plan.collect { - case r: RebalancePartitions if r.partitionExpressions.size == rSize => r - }.nonEmpty || rSize == 0) - } - - withView("v") { - withTable("t", "input1", "input2") { - withSQLConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS.key -> "true") { - sql(s"CREATE TABLE t (c1 int, c2 long) USING PARQUET PARTITIONED BY (p string)") - sql(s"CREATE TABLE input1 USING PARQUET AS SELECT * FROM VALUES(1,2),(1,3)") - sql(s"CREATE TABLE input2 USING PARQUET AS SELECT * FROM VALUES(1,3),(1,3)") - sql(s"CREATE VIEW v as SELECT col1, count(*) as col2 FROM input1 GROUP BY col1") - - val df0 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT /*+ broadcast(input2) */ input1.col1, input2.col1 - |FROM input1 - |JOIN input2 - |ON input1.col1 = input2.col1 - |""".stripMargin) - checkShuffleAndSort(df0.queryExecution.analyzed, 1, 1) - - val df1 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT /*+ broadcast(input2) */ input1.col1, input1.col2 - |FROM input1 - |LEFT JOIN input2 - |ON input1.col1 = input2.col1 and input1.col2 = input2.col2 - |""".stripMargin) - checkShuffleAndSort(df1.queryExecution.analyzed, 1, 2) - - val df2 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT col1 as c1, count(*) as c2 - |FROM input1 - |GROUP BY col1 - |HAVING count(*) > 0 - |""".stripMargin) - checkShuffleAndSort(df2.queryExecution.analyzed, 1, 1) - - // dynamic partition - val df3 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p) - |SELECT /*+ broadcast(input2) */ input1.col1, input1.col2, input1.col2 - |FROM input1 - |JOIN input2 - |ON input1.col1 = input2.col1 - |""".stripMargin) - checkShuffleAndSort(df3.queryExecution.analyzed, 0, 1) - - // non-deterministic - val df4 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT col1 + rand(), count(*) as c2 - |FROM input1 - |GROUP BY col1 - |""".stripMargin) - checkShuffleAndSort(df4.queryExecution.analyzed, 0, 0) - - // view - val df5 = sql( - s""" - |INSERT INTO TABLE t PARTITION(p='a') - |SELECT * FROM v - |""".stripMargin) - checkShuffleAndSort(df5.queryExecution.analyzed, 1, 1) - } - } - } - } - - test("Test rebalance in InsertIntoHiveDirCommand") { - withSQLConf( - HiveUtils.CONVERT_METASTORE_PARQUET.key -> "false", - HiveUtils.CONVERT_METASTORE_CTAS.key -> "false", - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTempDir(tmpDir => { - spark.range(0, 1000, 1, 10).createOrReplaceTempView("tmp_table") - val df = sql(s"INSERT OVERWRITE DIRECTORY '${tmpDir.getPath}' " + - s"STORED AS PARQUET SELECT * FROM tmp_table") - val insertHiveDirCommand = df.queryExecution.analyzed.collect { - case _: InsertIntoHiveDirCommand => true - } - assert(insertHiveDirCommand.size == 1) - val repartition = df.queryExecution.analyzed.collect { - case _: RebalancePartitions => true - } - assert(repartition.size == 1) - }) - } - } - - test("Test rebalance in InsertIntoDataSourceDirCommand") { - withSQLConf( - KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") { - withTempDir(tmpDir => { - spark.range(0, 1000, 1, 10).createOrReplaceTempView("tmp_table") - val df = sql(s"INSERT OVERWRITE DIRECTORY '${tmpDir.getPath}' " + - s"USING PARQUET SELECT * FROM tmp_table") - assert(df.queryExecution.analyzed.isInstanceOf[InsertIntoDataSourceDirCommand]) - val repartition = - df.queryExecution.analyzed.asInstanceOf[InsertIntoDataSourceDirCommand].query.collect { - case _: RebalancePartitions => true - } - assert(repartition.size == 1) - }) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala deleted file mode 100644 index 670d9ce7e43..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsAndPartitionAwareDataSource.scala +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.spark.sql - -import java.util.OptionalLong - -import org.apache.spark.sql.connector.{RangeInputPartition, SimpleBatchTable, SimpleScanBuilder, SimpleWritableDataSource} -import org.apache.spark.sql.connector.catalog.Table -import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, Transform} -import org.apache.spark.sql.connector.read._ -import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, Partitioning} -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -class ReportStatisticsAndPartitionAwareDataSource extends SimpleWritableDataSource { - - class MyScanBuilder( - val partitionKeys: Seq[String]) extends SimpleScanBuilder - with SupportsReportStatistics with SupportsReportPartitioning { - - override def estimateStatistics(): Statistics = { - new Statistics { - override def sizeInBytes(): OptionalLong = OptionalLong.of(80) - - override def numRows(): OptionalLong = OptionalLong.of(10) - - } - } - - override def planInputPartitions(): Array[InputPartition] = { - Array(RangeInputPartition(0, 5), RangeInputPartition(5, 10)) - } - - override def outputPartitioning(): Partitioning = { - new KeyGroupedPartitioning(partitionKeys.map(FieldReference(_)).toArray, 10) - } - } - - override def getTable(options: CaseInsensitiveStringMap): Table = { - new SimpleBatchTable { - override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new MyScanBuilder(Seq("i")) - } - - override def partitioning(): Array[Transform] = { - Array(Expressions.identity("i")) - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala deleted file mode 100644 index 2035d352562..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ReportStatisticsDataSource.scala +++ /dev/null @@ -1,53 +0,0 @@ -/* - * 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.spark.sql - -import java.util.OptionalLong - -import org.apache.spark.sql.connector._ -import org.apache.spark.sql.connector.catalog.Table -import org.apache.spark.sql.connector.read._ -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -class ReportStatisticsDataSource extends SimpleWritableDataSource { - - class MyScanBuilder extends SimpleScanBuilder - with SupportsReportStatistics { - - override def estimateStatistics(): Statistics = { - new Statistics { - override def sizeInBytes(): OptionalLong = OptionalLong.of(80) - - override def numRows(): OptionalLong = OptionalLong.of(10) - } - } - - override def planInputPartitions(): Array[InputPartition] = { - Array(RangeInputPartition(0, 5), RangeInputPartition(5, 10)) - } - - } - - override def getTable(options: CaseInsensitiveStringMap): Table = { - new SimpleBatchTable { - override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new MyScanBuilder - } - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala deleted file mode 100644 index 957089340ca..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala +++ /dev/null @@ -1,20 +0,0 @@ -/* - * 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.spark.sql - -class WatchDogSuite extends WatchDogSuiteBase {} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala deleted file mode 100644 index a392c969423..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/WatchDogSuiteBase.scala +++ /dev/null @@ -1,251 +0,0 @@ -/* - * 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.spark.sql - -import java.io.File - -import scala.collection.JavaConverters._ - -import org.apache.commons.io.FileUtils -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} -import org.apache.kyuubi.sql.watchdog.{MaxFileSizeExceedException, MaxPartitionExceedException} - -trait WatchDogSuiteBase extends KyuubiSparkSQLExtensionTest { - override protected def beforeAll(): Unit = { - super.beforeAll() - setupData() - } - - case class LimitAndExpected(limit: Int, expected: Int) - - val limitAndExpecteds = List(LimitAndExpected(1, 1), LimitAndExpected(11, 10)) - - private def checkMaxPartition: Unit = { - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS.key -> "100") { - checkAnswer(sql("SELECT count(distinct(p)) FROM test"), Row(10) :: Nil) - } - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_PARTITIONS.key -> "5") { - sql("SELECT * FROM test where p=1").queryExecution.sparkPlan - - sql(s"SELECT * FROM test WHERE p in (${Range(0, 5).toList.mkString(",")})") - .queryExecution.sparkPlan - - intercept[MaxPartitionExceedException]( - sql("SELECT * FROM test where p != 1").queryExecution.sparkPlan) - - intercept[MaxPartitionExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - - intercept[MaxPartitionExceedException](sql( - s"SELECT * FROM test WHERE p in (${Range(0, 6).toList.mkString(",")})") - .queryExecution.sparkPlan) - } - } - - test("watchdog with scan maxPartitions -- hive") { - Seq("textfile", "parquet").foreach { format => - withTable("test", "temp") { - sql( - s""" - |CREATE TABLE test(i int) - |PARTITIONED BY (p int) - |STORED AS $format""".stripMargin) - spark.range(0, 10, 1).selectExpr("id as col") - .createOrReplaceTempView("temp") - - for (part <- Range(0, 10)) { - sql( - s""" - |INSERT OVERWRITE TABLE test PARTITION (p='$part') - |select col from temp""".stripMargin) - } - checkMaxPartition - } - } - } - - test("watchdog with scan maxPartitions -- data source") { - withTempDir { dir => - withTempView("test") { - spark.range(10).selectExpr("id", "id as p") - .write - .partitionBy("p") - .mode("overwrite") - .save(dir.getCanonicalPath) - spark.read.load(dir.getCanonicalPath).createOrReplaceTempView("test") - checkMaxPartition - } - } - } - - private def checkMaxFileSize(tableSize: Long, nonPartTableSize: Long): Unit = { - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> tableSize.toString) { - checkAnswer(sql("SELECT count(distinct(p)) FROM test"), Row(10) :: Nil) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (tableSize / 2).toString) { - sql("SELECT * FROM test where p=1").queryExecution.sparkPlan - - sql(s"SELECT * FROM test WHERE p in (${Range(0, 3).toList.mkString(",")})") - .queryExecution.sparkPlan - - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test where p != 1").queryExecution.sparkPlan) - - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - - intercept[MaxFileSizeExceedException](sql( - s"SELECT * FROM test WHERE p in (${Range(0, 6).toList.mkString(",")})") - .queryExecution.sparkPlan) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> nonPartTableSize.toString) { - checkAnswer(sql("SELECT count(*) FROM test_non_part"), Row(10000) :: Nil) - } - - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (nonPartTableSize - 1).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan) - } - } - - test("watchdog with scan maxFileSize -- hive") { - Seq(false).foreach { convertMetastoreParquet => - withTable("test", "test_non_part", "temp") { - spark.range(10000).selectExpr("id as col") - .createOrReplaceTempView("temp") - - // partitioned table - sql( - s""" - |CREATE TABLE test(i int) - |PARTITIONED BY (p int) - |STORED AS parquet""".stripMargin) - for (part <- Range(0, 10)) { - sql( - s""" - |INSERT OVERWRITE TABLE test PARTITION (p='$part') - |select col from temp""".stripMargin) - } - - val tablePath = new File(spark.sessionState.catalog.externalCatalog - .getTable("default", "test").location) - val tableSize = FileUtils.listFiles(tablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // non-partitioned table - sql( - s""" - |CREATE TABLE test_non_part(i int) - |STORED AS parquet""".stripMargin) - sql( - s""" - |INSERT OVERWRITE TABLE test_non_part - |select col from temp""".stripMargin) - sql("ANALYZE TABLE test_non_part COMPUTE STATISTICS") - - val nonPartTablePath = new File(spark.sessionState.catalog.externalCatalog - .getTable("default", "test_non_part").location) - val nonPartTableSize = FileUtils.listFiles(nonPartTablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(nonPartTableSize > 0) - - // check - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> convertMetastoreParquet.toString) { - checkMaxFileSize(tableSize, nonPartTableSize) - } - } - } - } - - test("watchdog with scan maxFileSize -- data source") { - withTempDir { dir => - withTempView("test", "test_non_part") { - // partitioned table - val tablePath = new File(dir, "test") - spark.range(10).selectExpr("id", "id as p") - .write - .partitionBy("p") - .mode("overwrite") - .parquet(tablePath.getCanonicalPath) - spark.read.load(tablePath.getCanonicalPath).createOrReplaceTempView("test") - - val tableSize = FileUtils.listFiles(tablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // non-partitioned table - val nonPartTablePath = new File(dir, "test_non_part") - spark.range(10000).selectExpr("id", "id as p") - .write - .mode("overwrite") - .parquet(nonPartTablePath.getCanonicalPath) - spark.read.load(nonPartTablePath.getCanonicalPath).createOrReplaceTempView("test_non_part") - - val nonPartTableSize = FileUtils.listFiles(nonPartTablePath, Array("parquet"), true).asScala - .map(_.length()).sum - assert(tableSize > 0) - - // check - checkMaxFileSize(tableSize, nonPartTableSize) - } - } - } - - test("disable script transformation") { - withSQLConf(KyuubiSQLConf.SCRIPT_TRANSFORMATION_ENABLED.key -> "false") { - val e = intercept[KyuubiSQLExtensionException] { - sql("SELECT TRANSFORM('') USING 'ls /'") - } - assert(e.getMessage == "Script transformation is not allowed") - } - } - - test("watchdog with scan maxFileSize -- data source v2") { - val df = spark.read.format(classOf[ReportStatisticsAndPartitionAwareDataSource].getName).load() - df.createOrReplaceTempView("test") - val logical = df.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - val tableSize = logical.computeStats().sizeInBytes.toLong - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> tableSize.toString) { - sql("SELECT * FROM test").queryExecution.sparkPlan - } - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (tableSize / 2).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test").queryExecution.sparkPlan) - } - val nonPartDf = spark.read.format(classOf[ReportStatisticsDataSource].getName).load() - nonPartDf.createOrReplaceTempView("test_non_part") - val nonPartLogical = nonPartDf.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - val nonPartTableSize = nonPartLogical.computeStats().sizeInBytes.toLong - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> nonPartTableSize.toString) { - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan - } - withSQLConf(KyuubiSQLConf.WATCHDOG_MAX_FILE_SIZE.key -> (nonPartTableSize / 2).toString) { - intercept[MaxFileSizeExceedException]( - sql("SELECT * FROM test_non_part").queryExecution.sparkPlan) - } - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala deleted file mode 100644 index 7af1ca0482b..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderCoreBenchmark.scala +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.benchmark.Benchmark -import org.apache.spark.sql.benchmark.KyuubiBenchmarkBase -import org.apache.spark.sql.internal.StaticSQLConf - -import org.apache.kyuubi.sql.zorder.ZorderBytesUtils - -/** - * Benchmark to measure performance with zorder core. - * - * {{{ - * RUN_BENCHMARK=1 ./build/mvn clean test \ - * -pl extensions/spark/kyuubi-extension-spark-3-5 -am \ - * -Pspark-3.5,kyuubi-extension-spark-3-5 \ - * -Dtest=none -DwildcardSuites=org.apache.spark.sql.ZorderCoreBenchmark - * }}} - */ -class ZorderCoreBenchmark extends KyuubiSparkSQLExtensionTest with KyuubiBenchmarkBase { - private val runBenchmark = sys.env.contains("RUN_BENCHMARK") - private val numRows = 1 * 1000 * 1000 - - private def randomInt(numColumns: Int): Seq[Array[Any]] = { - (1 to numRows).map { l => - val arr = new Array[Any](numColumns) - (0 until numColumns).foreach(col => arr(col) = l) - arr - } - } - - private def randomLong(numColumns: Int): Seq[Array[Any]] = { - (1 to numRows).map { l => - val arr = new Array[Any](numColumns) - (0 until numColumns).foreach(col => arr(col) = l.toLong) - arr - } - } - - private def interleaveMultiByteArrayBenchmark(): Unit = { - val benchmark = - new Benchmark(s"$numRows rows zorder core benchmark", numRows, output = output) - benchmark.addCase("2 int columns benchmark", 3) { _ => - randomInt(2).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("3 int columns benchmark", 3) { _ => - randomInt(3).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("4 int columns benchmark", 3) { _ => - randomInt(4).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("2 long columns benchmark", 3) { _ => - randomLong(2).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("3 long columns benchmark", 3) { _ => - randomLong(3).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.addCase("4 long columns benchmark", 3) { _ => - randomLong(4).foreach(ZorderBytesUtils.interleaveBits) - } - - benchmark.run() - } - - private def paddingTo8ByteBenchmark() { - val iterations = 10 * 1000 * 1000 - - val b2 = Array('a'.toByte, 'b'.toByte) - val benchmark = - new Benchmark(s"$iterations iterations paddingTo8Byte benchmark", iterations, output = output) - benchmark.addCase("2 length benchmark", 3) { _ => - (1 to iterations).foreach(_ => ZorderBytesUtils.paddingTo8Byte(b2)) - } - - val b16 = Array.tabulate(16) { i => i.toByte } - benchmark.addCase("16 length benchmark", 3) { _ => - (1 to iterations).foreach(_ => ZorderBytesUtils.paddingTo8Byte(b16)) - } - - benchmark.run() - } - - test("zorder core benchmark") { - assume(runBenchmark) - - withHeader { - interleaveMultiByteArrayBenchmark() - paddingTo8ByteBenchmark() - } - } - - override def sparkConf(): SparkConf = { - super.sparkConf().remove(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key) - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuite.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuite.scala deleted file mode 100644 index c2fa1619707..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuite.scala +++ /dev/null @@ -1,123 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.sql.catalyst.parser.ParserInterface -import org.apache.spark.sql.catalyst.plans.logical.{RebalancePartitions, Sort} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.kyuubi.sql.{KyuubiSQLConf, SparkKyuubiSparkSQLParser} -import org.apache.kyuubi.sql.zorder.Zorder - -trait ZorderSuiteSpark extends ZorderSuiteBase { - - test("Add rebalance before zorder") { - Seq("true" -> false, "false" -> true).foreach { case (useOriginalOrdering, zorder) => - withSQLConf( - KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false", - KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED.key -> "true", - KyuubiSQLConf.ZORDER_USING_ORIGINAL_ORDERING_ENABLED.key -> useOriginalOrdering) { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) PARTITIONED BY (d string) - | TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,C2') - |""".stripMargin) - val p = sql("INSERT INTO TABLE t PARTITION(d='a') SELECT * FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p.collect { - case sort: Sort - if !sort.global && - ((sort.order.exists(_.child.isInstanceOf[Zorder]) && zorder) || - (!sort.order.exists(_.child.isInstanceOf[Zorder]) && !zorder)) => sort - }.size == 1) - assert(p.collect { - case rebalance: RebalancePartitions - if rebalance.references.map(_.name).exists(_.equals("c1")) => rebalance - }.size == 1) - - val p2 = sql("INSERT INTO TABLE t PARTITION(d) SELECT * FROM VALUES(1,'a','b')") - .queryExecution.analyzed - assert(p2.collect { - case sort: Sort - if (!sort.global && Seq("c1", "c2", "d").forall(x => - sort.references.map(_.name).exists(_.equals(x)))) && - ((sort.order.exists(_.child.isInstanceOf[Zorder]) && zorder) || - (!sort.order.exists(_.child.isInstanceOf[Zorder]) && !zorder)) => sort - }.size == 1) - assert(p2.collect { - case rebalance: RebalancePartitions - if Seq("c1", "c2", "d").forall(x => - rebalance.references.map(_.name).exists(_.equals(x))) => rebalance - }.size == 1) - } - } - } - } - - test("Two phase rebalance before Z-Order") { - withSQLConf( - SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> - "org.apache.spark.sql.catalyst.optimizer.CollapseRepartition", - KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false", - KyuubiSQLConf.REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.TWO_PHASE_REBALANCE_BEFORE_ZORDER.key -> "true", - KyuubiSQLConf.REBALANCE_ZORDER_COLUMNS_ENABLED.key -> "true") { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int) PARTITIONED BY (d string) - | TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1') - |""".stripMargin) - val p = sql("INSERT INTO TABLE t PARTITION(d) SELECT * FROM VALUES(1,'a')") - val rebalance = p.queryExecution.optimizedPlan.innerChildren - .flatMap(_.collect { case r: RebalancePartitions => r }) - assert(rebalance.size == 2) - assert(rebalance.head.partitionExpressions.flatMap(_.references.map(_.name)) - .contains("d")) - assert(rebalance.head.partitionExpressions.flatMap(_.references.map(_.name)) - .contains("c1")) - - assert(rebalance(1).partitionExpressions.flatMap(_.references.map(_.name)) - .contains("d")) - assert(!rebalance(1).partitionExpressions.flatMap(_.references.map(_.name)) - .contains("c1")) - } - } - } -} - -trait ParserSuite { self: ZorderSuiteBase => - override def createParser: ParserInterface = { - new SparkKyuubiSparkSQLParser(spark.sessionState.sqlParser) - } -} - -class ZorderWithCodegenEnabledSuite - extends ZorderWithCodegenEnabledSuiteBase - with ZorderSuiteSpark - with ParserSuite {} -class ZorderWithCodegenDisabledSuite - extends ZorderWithCodegenDisabledSuiteBase - with ZorderSuiteSpark - with ParserSuite {} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala deleted file mode 100644 index d1615699517..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala +++ /dev/null @@ -1,833 +0,0 @@ -/* - * 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.spark.sql - -import org.apache.spark.SparkConf -import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} -import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar} -import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, EqualTo, Expression, ExpressionEvalHelper, Literal, NullsLast, SortOrder} -import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, OneRowRelation, Project, Sort} -import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.hive.execution.InsertIntoHiveTable -import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.types._ - -import org.apache.kyuubi.sql.{KyuubiSQLConf, KyuubiSQLExtensionException} -import org.apache.kyuubi.sql.zorder.{OptimizeZorderCommandBase, OptimizeZorderStatement, Zorder, ZorderBytesUtils} - -trait ZorderSuiteBase extends KyuubiSparkSQLExtensionTest with ExpressionEvalHelper { - override def sparkConf(): SparkConf = { - super.sparkConf() - .set( - StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, - "org.apache.kyuubi.sql.KyuubiSparkSQLCommonExtension") - } - - test("optimize unpartitioned table") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("up") { - sql(s"DROP TABLE IF EXISTS up") - - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - sql(s"CREATE TABLE up (c1 INT, c2 INT, c3 INT)") - sql(s"INSERT INTO TABLE up VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - val e = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE up WHERE c1 > 1 ZORDER BY c1, c2") - } - assert(e.getMessage == "Filters are only supported for partitioned table") - - sql("OPTIMIZE up ZORDER BY c1, c2") - val res = sql("SELECT c1, c2 FROM up").collect() - - assert(res.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r = res(i) - assert(t(0) == r.getInt(0)) - assert(t(1) == r.getInt(1)) - } - } - } - } - - test("optimize partitioned table") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("p") { - sql("DROP TABLE IF EXISTS p") - - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - - sql(s"CREATE TABLE p (c1 INT, c2 INT, c3 INT) PARTITIONED BY (id INT)") - sql(s"ALTER TABLE p ADD PARTITION (id = 1)") - sql(s"ALTER TABLE p ADD PARTITION (id = 2)") - sql(s"INSERT INTO TABLE p PARTITION (id = 1) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - sql(s"INSERT INTO TABLE p PARTITION (id = 2) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - sql(s"OPTIMIZE p ZORDER BY c1, c2") - - val res1 = sql(s"SELECT c1, c2 FROM p WHERE id = 1").collect() - val res2 = sql(s"SELECT c1, c2 FROM p WHERE id = 2").collect() - - assert(res1.length == 16) - assert(res2.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r1 = res1(i) - assert(t(0) == r1.getInt(0)) - assert(t(1) == r1.getInt(1)) - - val r2 = res2(i) - assert(t(0) == r2.getInt(0)) - assert(t(1) == r2.getInt(1)) - } - } - } - } - - test("optimize partitioned table with filters") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("p") { - sql("DROP TABLE IF EXISTS p") - - val target1 = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - val target2 = Seq( - Seq(0, 0), - Seq(0, 1), - Seq(0, 2), - Seq(0, 3), - Seq(1, 0), - Seq(1, 1), - Seq(1, 2), - Seq(1, 3), - Seq(2, 0), - Seq(2, 1), - Seq(2, 2), - Seq(2, 3), - Seq(3, 0), - Seq(3, 1), - Seq(3, 2), - Seq(3, 3)) - sql(s"CREATE TABLE p (c1 INT, c2 INT, c3 INT) PARTITIONED BY (id INT)") - sql(s"ALTER TABLE p ADD PARTITION (id = 1)") - sql(s"ALTER TABLE p ADD PARTITION (id = 2)") - sql(s"INSERT INTO TABLE p PARTITION (id = 1) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - sql(s"INSERT INTO TABLE p PARTITION (id = 2) VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - val e = intercept[KyuubiSQLExtensionException]( - sql(s"OPTIMIZE p WHERE id = 1 AND c1 > 1 ZORDER BY c1, c2")) - assert(e.getMessage == "Only partition column filters are allowed") - - sql(s"OPTIMIZE p WHERE id = 1 ZORDER BY c1, c2") - - val res1 = sql(s"SELECT c1, c2 FROM p WHERE id = 1").collect() - val res2 = sql(s"SELECT c1, c2 FROM p WHERE id = 2").collect() - - assert(res1.length == 16) - assert(res2.length == 16) - - for (i <- target1.indices) { - val t1 = target1(i) - val r1 = res1(i) - assert(t1(0) == r1.getInt(0)) - assert(t1(1) == r1.getInt(1)) - - val t2 = target2(i) - val r2 = res2(i) - assert(t2(0) == r2.getInt(0)) - assert(t2(1) == r2.getInt(1)) - } - } - } - } - - test("optimize zorder with datasource table") { - // TODO remove this if we support datasource table - withTable("t") { - sql("CREATE TABLE t (c1 int, c2 int) USING PARQUET") - val msg = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE t ZORDER BY c1, c2") - }.getMessage - assert(msg.contains("only support hive table")) - } - } - - private def checkZorderTable( - enabled: Boolean, - cols: String, - planHasRepartition: Boolean, - resHasSort: Boolean): Unit = { - def checkSort(plan: LogicalPlan): Unit = { - assert(plan.isInstanceOf[Sort] === resHasSort) - plan match { - case sort: Sort => - val colArr = cols.split(",") - val refs = - if (colArr.length == 1) { - sort.order.head - .child.asInstanceOf[AttributeReference] :: Nil - } else { - sort.order.head - .child.asInstanceOf[Zorder].children.map(_.references.head) - } - assert(refs.size === colArr.size) - refs.zip(colArr).foreach { case (ref, col) => - assert(ref.name === col.trim) - } - case _ => - } - } - - val repartition = - if (planHasRepartition) { - "/*+ repartition */" - } else { - "" - } - withSQLConf("spark.sql.shuffle.partitions" -> "1") { - // hive - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> "false") { - withTable("zorder_t1", "zorder_t2_true", "zorder_t2_false") { - sql( - s""" - |CREATE TABLE zorder_t1 (c1 int, c2 string, c3 long, c4 double) STORED AS PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - |""".stripMargin) - val df1 = sql(s""" - |INSERT INTO TABLE zorder_t1 - |SELECT $repartition * FROM VALUES(1,'a',2,4D),(2,'b',3,6D) - |""".stripMargin) - assert(df1.queryExecution.analyzed.isInstanceOf[InsertIntoHiveTable]) - checkSort(df1.queryExecution.analyzed.children.head) - - Seq("true", "false").foreach { optimized => - withSQLConf( - "spark.sql.hive.convertMetastoreCtas" -> optimized, - "spark.sql.hive.convertMetastoreParquet" -> optimized) { - - withListener( - s""" - |CREATE TABLE zorder_t2_$optimized STORED AS PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - | - |SELECT $repartition * FROM - |VALUES(1,'a',2,4D),(2,'b',3,6D) AS t(c1 ,c2 , c3, c4) - |""".stripMargin) { write => - if (optimized.toBoolean) { - assert(write.isInstanceOf[InsertIntoHadoopFsRelationCommand]) - } else { - assert(write.isInstanceOf[InsertIntoHiveTable]) - } - checkSort(write.query) - } - } - } - } - } - - // datasource - withTable("zorder_t3", "zorder_t4") { - sql( - s""" - |CREATE TABLE zorder_t3 (c1 int, c2 string, c3 long, c4 double) USING PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - |""".stripMargin) - val df1 = sql(s""" - |INSERT INTO TABLE zorder_t3 - |SELECT $repartition * FROM VALUES(1,'a',2,4D),(2,'b',3,6D) - |""".stripMargin) - assert(df1.queryExecution.analyzed.isInstanceOf[InsertIntoHadoopFsRelationCommand]) - checkSort(df1.queryExecution.analyzed.children.head) - - withListener( - s""" - |CREATE TABLE zorder_t4 USING PARQUET - |TBLPROPERTIES ( - | 'kyuubi.zorder.enabled' = '$enabled', - | 'kyuubi.zorder.cols' = '$cols') - | - |SELECT $repartition * FROM - |VALUES(1,'a',2,4D),(2,'b',3,6D) AS t(c1 ,c2 , c3, c4) - |""".stripMargin) { write => - assert(write.isInstanceOf[InsertIntoHadoopFsRelationCommand]) - checkSort(write.query) - } - } - } - } - - test("Support insert zorder by table properties") { - withSQLConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING.key -> "false") { - checkZorderTable(true, "c1", false, false) - checkZorderTable(false, "c1", false, false) - } - withSQLConf(KyuubiSQLConf.INSERT_ZORDER_BEFORE_WRITING.key -> "true") { - checkZorderTable(true, "", false, false) - checkZorderTable(true, "c5", false, false) - checkZorderTable(true, "c1,c5", false, false) - checkZorderTable(false, "c3", false, false) - checkZorderTable(true, "c3", true, false) - checkZorderTable(true, "c3", false, true) - checkZorderTable(true, "c2,c4", false, true) - checkZorderTable(true, "c4, c2, c1, c3", false, true) - } - } - - test("zorder: check unsupported data type") { - def checkZorderPlan(zorder: Expression): Unit = { - val msg = intercept[AnalysisException] { - val plan = Project(Seq(Alias(zorder, "c")()), OneRowRelation()) - spark.sessionState.analyzer.checkAnalysis(plan) - }.getMessage - assert(msg.contains("Unsupported z-order type: void")) - } - - checkZorderPlan(Zorder(Seq(Literal(null, NullType)))) - checkZorderPlan(Zorder(Seq(Literal(1, IntegerType), Literal(null, NullType)))) - } - - test("zorder: check supported data type") { - val children = Seq( - Literal.create(false, BooleanType), - Literal.create(null, BooleanType), - Literal.create(1.toByte, ByteType), - Literal.create(null, ByteType), - Literal.create(1.toShort, ShortType), - Literal.create(null, ShortType), - Literal.create(1, IntegerType), - Literal.create(null, IntegerType), - Literal.create(1L, LongType), - Literal.create(null, LongType), - Literal.create(1f, FloatType), - Literal.create(null, FloatType), - Literal.create(1d, DoubleType), - Literal.create(null, DoubleType), - Literal.create("1", StringType), - Literal.create(null, StringType), - Literal.create(1L, TimestampType), - Literal.create(null, TimestampType), - Literal.create(1, DateType), - Literal.create(null, DateType), - Literal.create(BigDecimal(1, 1), DecimalType(1, 1)), - Literal.create(null, DecimalType(1, 1))) - val zorder = Zorder(children) - val plan = Project(Seq(Alias(zorder, "c")()), OneRowRelation()) - spark.sessionState.analyzer.checkAnalysis(plan) - assert(zorder.foldable) - -// // scalastyle:off -// val resultGen = org.apache.commons.codec.binary.Hex.encodeHex( -// zorder.eval(InternalRow.fromSeq(children)).asInstanceOf[Array[Byte]], false) -// resultGen.grouped(2).zipWithIndex.foreach { case (char, i) => -// print("0x" + char(0) + char(1) + ", ") -// if ((i + 1) % 10 == 0) { -// println() -// } -// } -// // scalastyle:on - - val expected = Array( - 0xFB, 0xEA, 0xAA, 0xBA, 0xAE, 0xAB, 0xAA, 0xEA, 0xBA, 0xAE, 0xAB, 0xAA, 0xEA, 0xBA, 0xA6, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xBA, 0xBB, 0xAA, 0xAA, 0xAA, - 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0xBA, 0xAA, 0x9A, 0xAA, 0xAA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xEA, - 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, - 0xAA, 0xAA, 0xBE, 0xAA, 0xAA, 0x8A, 0xBA, 0xAA, 0x2A, 0xEA, 0xA8, 0xAA, 0xAA, 0xA2, 0xAA, - 0xAA, 0x8A, 0xAA, 0xAA, 0x2F, 0xEB, 0xFE) - .map(_.toByte) - checkEvaluation(zorder, expected, InternalRow.fromSeq(children)) - } - - private def checkSort(input: DataFrame, expected: Seq[Row], dataType: Array[DataType]): Unit = { - withTempDir { dir => - input.repartition(3).write.mode("overwrite").format("parquet").save(dir.getCanonicalPath) - val df = spark.read.format("parquet") - .load(dir.getCanonicalPath) - .repartition(1) - assert(df.schema.fields.map(_.dataType).sameElements(dataType)) - val exprs = Seq("c1", "c2").map(col).map(_.expr) - val sortOrder = SortOrder(Zorder(exprs), Ascending, NullsLast, Seq.empty) - val zorderSort = Sort(Seq(sortOrder), true, df.logicalPlan) - val result = Dataset.ofRows(spark, zorderSort) - checkAnswer(result, expected) - } - } - - test("sort with zorder -- boolean column") { - val schema = StructType(StructField("c1", BooleanType) :: StructField("c2", BooleanType) :: Nil) - val nonNullDF = spark.createDataFrame( - spark.sparkContext.parallelize( - Seq(Row(false, false), Row(false, true), Row(true, false), Row(true, true))), - schema) - val expected = - Row(false, false) :: Row(true, false) :: Row(false, true) :: Row(true, true) :: Nil - checkSort(nonNullDF, expected, Array(BooleanType, BooleanType)) - val df = spark.createDataFrame( - spark.sparkContext.parallelize( - Seq(Row(false, false), Row(false, null), Row(null, false), Row(null, null))), - schema) - val expected2 = - Row(false, false) :: Row(null, false) :: Row(false, null) :: Row(null, null) :: Nil - checkSort(df, expected2, Array(BooleanType, BooleanType)) - } - - test("sort with zorder -- int column") { - // TODO: add more datatype unit test - val session = spark - import session.implicits._ - // generate 4 * 4 matrix - val len = 3 - val input = spark.range(len + 1).selectExpr("cast(id as int) as c1") - .select($"c1", explode(sequence(lit(0), lit(len))) as "c2") - val expected = - Row(0, 0) :: Row(1, 0) :: Row(0, 1) :: Row(1, 1) :: - Row(2, 0) :: Row(3, 0) :: Row(2, 1) :: Row(3, 1) :: - Row(0, 2) :: Row(1, 2) :: Row(0, 3) :: Row(1, 3) :: - Row(2, 2) :: Row(3, 2) :: Row(2, 3) :: Row(3, 3) :: Nil - checkSort(input, expected, Array(IntegerType, IntegerType)) - - // contains null value case. - val nullDF = spark.range(1).selectExpr("cast(null as int) as c1") - val input2 = spark.range(len).selectExpr("cast(id as int) as c1") - .union(nullDF) - .select( - $"c1", - explode(concat(sequence(lit(0), lit(len - 1)), array(lit(null)))) as "c2") - val expected2 = Row(0, 0) :: Row(1, 0) :: Row(0, 1) :: Row(1, 1) :: - Row(2, 0) :: Row(2, 1) :: Row(0, 2) :: Row(1, 2) :: - Row(2, 2) :: Row(null, 0) :: Row(null, 1) :: Row(null, 2) :: - Row(0, null) :: Row(1, null) :: Row(2, null) :: Row(null, null) :: Nil - checkSort(input2, expected2, Array(IntegerType, IntegerType)) - } - - test("sort with zorder -- string column") { - val schema = StructType(StructField("c1", StringType) :: StructField("c2", StringType) :: Nil) - val rdd = spark.sparkContext.parallelize(Seq( - Row("a", "a"), - Row("a", "b"), - Row("a", "c"), - Row("a", "d"), - Row("b", "a"), - Row("b", "b"), - Row("b", "c"), - Row("b", "d"), - Row("c", "a"), - Row("c", "b"), - Row("c", "c"), - Row("c", "d"), - Row("d", "a"), - Row("d", "b"), - Row("d", "c"), - Row("d", "d"))) - val input = spark.createDataFrame(rdd, schema) - val expected = Row("a", "a") :: Row("b", "a") :: Row("c", "a") :: Row("a", "b") :: - Row("a", "c") :: Row("b", "b") :: Row("c", "b") :: Row("b", "c") :: - Row("c", "c") :: Row("d", "a") :: Row("d", "b") :: Row("d", "c") :: - Row("a", "d") :: Row("b", "d") :: Row("c", "d") :: Row("d", "d") :: Nil - checkSort(input, expected, Array(StringType, StringType)) - - val rdd2 = spark.sparkContext.parallelize(Seq( - Row(null, "a"), - Row("a", "b"), - Row("a", "c"), - Row("a", null), - Row("b", "a"), - Row(null, "b"), - Row("b", null), - Row("b", "d"), - Row("c", "a"), - Row("c", null), - Row(null, "c"), - Row("c", "d"), - Row("d", null), - Row("d", "b"), - Row("d", "c"), - Row(null, "d"), - Row(null, null))) - val input2 = spark.createDataFrame(rdd2, schema) - val expected2 = Row("b", "a") :: Row("c", "a") :: Row("a", "b") :: Row("a", "c") :: - Row("d", "b") :: Row("d", "c") :: Row("b", "d") :: Row("c", "d") :: - Row(null, "a") :: Row(null, "b") :: Row(null, "c") :: Row(null, "d") :: - Row("a", null) :: Row("b", null) :: Row("c", null) :: Row("d", null) :: - Row(null, null) :: Nil - checkSort(input2, expected2, Array(StringType, StringType)) - } - - test("test special value of short int long type") { - val df1 = spark.createDataFrame(Seq( - (-1, -1L), - (Int.MinValue, Int.MinValue.toLong), - (1, 1L), - (Int.MaxValue - 1, Int.MaxValue.toLong), - (Int.MaxValue - 1, Int.MaxValue.toLong - 1), - (Int.MaxValue, Int.MaxValue.toLong + 1), - (Int.MaxValue, Int.MaxValue.toLong))).toDF("c1", "c2") - val expected1 = - Row(Int.MinValue, Int.MinValue.toLong) :: - Row(-1, -1L) :: - Row(1, 1L) :: - Row(Int.MaxValue - 1, Int.MaxValue.toLong - 1) :: - Row(Int.MaxValue - 1, Int.MaxValue.toLong) :: - Row(Int.MaxValue, Int.MaxValue.toLong) :: - Row(Int.MaxValue, Int.MaxValue.toLong + 1) :: Nil - checkSort(df1, expected1, Array(IntegerType, LongType)) - - val df2 = spark.createDataFrame(Seq( - (-1, -1.toShort), - (Short.MinValue.toInt, Short.MinValue), - (1, 1.toShort), - (Short.MaxValue.toInt, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toInt + 1, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toInt, Short.MaxValue), - (Short.MaxValue.toInt + 1, Short.MaxValue))).toDF("c1", "c2") - val expected2 = - Row(Short.MinValue.toInt, Short.MinValue) :: - Row(-1, -1.toShort) :: - Row(1, 1.toShort) :: - Row(Short.MaxValue.toInt, Short.MaxValue - 1) :: - Row(Short.MaxValue.toInt, Short.MaxValue) :: - Row(Short.MaxValue.toInt + 1, Short.MaxValue - 1) :: - Row(Short.MaxValue.toInt + 1, Short.MaxValue) :: Nil - checkSort(df2, expected2, Array(IntegerType, ShortType)) - - val df3 = spark.createDataFrame(Seq( - (-1L, -1.toShort), - (Short.MinValue.toLong, Short.MinValue), - (1L, 1.toShort), - (Short.MaxValue.toLong, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toLong + 1, (Short.MaxValue - 1).toShort), - (Short.MaxValue.toLong, Short.MaxValue), - (Short.MaxValue.toLong + 1, Short.MaxValue))).toDF("c1", "c2") - val expected3 = - Row(Short.MinValue.toLong, Short.MinValue) :: - Row(-1L, -1.toShort) :: - Row(1L, 1.toShort) :: - Row(Short.MaxValue.toLong, Short.MaxValue - 1) :: - Row(Short.MaxValue.toLong, Short.MaxValue) :: - Row(Short.MaxValue.toLong + 1, Short.MaxValue - 1) :: - Row(Short.MaxValue.toLong + 1, Short.MaxValue) :: Nil - checkSort(df3, expected3, Array(LongType, ShortType)) - } - - test("skip zorder if only requires one column") { - withTable("t") { - withSQLConf("spark.sql.hive.convertMetastoreParquet" -> "false") { - sql("CREATE TABLE t (c1 int, c2 string) stored as parquet") - val order1 = sql("OPTIMIZE t ZORDER BY c1").queryExecution.analyzed - .asInstanceOf[OptimizeZorderCommandBase].query.asInstanceOf[Sort].order.head.child - assert(!order1.isInstanceOf[Zorder]) - assert(order1.isInstanceOf[AttributeReference]) - } - } - } - - test("Add config to control if zorder using global sort") { - withTable("t") { - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,c2') - |""".stripMargin) - val p1 = sql("OPTIMIZE t ZORDER BY c1, c2").queryExecution.analyzed - assert(p1.collect { - case shuffle: Sort if !shuffle.global => shuffle - }.size == 1) - - val p2 = sql("INSERT INTO TABLE t SELECT * FROM VALUES(1,'a')").queryExecution.analyzed - assert(p2.collect { - case shuffle: Sort if !shuffle.global => shuffle - }.size == 1) - } - } - } - - test("Allow insert zorder after repartition if zorder using local sort") { - withTable("t") { - sql( - """ - |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES ( - |'kyuubi.zorder.enabled'= 'true', - |'kyuubi.zorder.cols'= 'c1,c2') - |""".stripMargin) - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") { - val p1 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p1.collect { - case sort: Sort if !sort.global => sort - }.size == 1) - } - withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "true") { - val p2 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')") - .queryExecution.analyzed - assert(p2.collect { - case sort: Sort if !sort.global => sort - }.size == 0) - } - } - } - - test("fast approach test") { - Seq[Seq[Any]]( - Seq(1L, 2L), - Seq(1L, 2L, 3L), - Seq(1L, 2L, 3L, 4L), - Seq(1L, 2L, 3L, 4L, 5L), - Seq(1L, 2L, 3L, 4L, 5L, 6L), - Seq(1L, 2L, 3L, 4L, 5L, 6L, 7L), - Seq(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L)) - .foreach { inputs => - assert(java.util.Arrays.equals( - ZorderBytesUtils.interleaveBits(inputs.toArray), - ZorderBytesUtils.interleaveBitsDefault(inputs.map(ZorderBytesUtils.toByteArray).toArray))) - } - } - - test("OPTIMIZE command is parsed as expected") { - val parser = createParser - val globalSort = spark.conf.get(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) - - assert(parser.parsePlan("OPTIMIZE p zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project(Seq(UnresolvedStar(None)), UnresolvedRelation(TableIdentifier("p")))))) - - assert(parser.parsePlan("OPTIMIZE p zorder by c1, c2") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder( - Zorder(Seq(UnresolvedAttribute("c1"), UnresolvedAttribute("c2"))), - Ascending, - NullsLast, - Seq.empty) :: Nil, - globalSort, - Project(Seq(UnresolvedStar(None)), UnresolvedRelation(TableIdentifier("p")))))) - - assert(parser.parsePlan("OPTIMIZE p where id = 1 zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo(UnresolvedAttribute("id"), Literal(1)), - UnresolvedRelation(TableIdentifier("p"))))))) - - assert(parser.parsePlan("OPTIMIZE p where id = 1 zorder by c1, c2") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder( - Zorder(Seq(UnresolvedAttribute("c1"), UnresolvedAttribute("c2"))), - Ascending, - NullsLast, - Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo(UnresolvedAttribute("id"), Literal(1)), - UnresolvedRelation(TableIdentifier("p"))))))) - - assert(parser.parsePlan("OPTIMIZE p where id = current_date() zorder by c1") === - OptimizeZorderStatement( - Seq("p"), - Sort( - SortOrder(UnresolvedAttribute("c1"), Ascending, NullsLast, Seq.empty) :: Nil, - globalSort, - Project( - Seq(UnresolvedStar(None)), - Filter( - EqualTo( - UnresolvedAttribute("id"), - UnresolvedFunction("current_date", Seq.empty, false)), - UnresolvedRelation(TableIdentifier("p"))))))) - - // TODO: add following case support - intercept[ParseException] { - parser.parsePlan("OPTIMIZE p zorder by (c1)") - } - - intercept[ParseException] { - parser.parsePlan("OPTIMIZE p zorder by (c1, c2)") - } - } - - test("OPTIMIZE partition predicates constraint") { - withTable("p") { - sql("CREATE TABLE p (c1 INT, c2 INT) PARTITIONED BY (event_date DATE)") - val e1 = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE p WHERE event_date = current_date as c ZORDER BY c1, c2") - } - assert(e1.getMessage.contains("unsupported partition predicates")) - - val e2 = intercept[KyuubiSQLExtensionException] { - sql("OPTIMIZE p WHERE c1 = 1 ZORDER BY c1, c2") - } - assert(e2.getMessage == "Only partition column filters are allowed") - } - } - - test("optimize sort by backquoted column name") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - withTable("up") { - sql(s"DROP TABLE IF EXISTS up") - val target = Seq( - Seq(0, 0), - Seq(1, 0), - Seq(0, 1), - Seq(1, 1), - Seq(2, 0), - Seq(3, 0), - Seq(2, 1), - Seq(3, 1), - Seq(0, 2), - Seq(1, 2), - Seq(0, 3), - Seq(1, 3), - Seq(2, 2), - Seq(3, 2), - Seq(2, 3), - Seq(3, 3)) - sql(s"CREATE TABLE up (c1 INT, `@c2` INT, c3 INT)") - sql(s"INSERT INTO TABLE up VALUES" + - "(0,0,2),(0,1,2),(0,2,1),(0,3,3)," + - "(1,0,4),(1,1,2),(1,2,1),(1,3,3)," + - "(2,0,2),(2,1,1),(2,2,5),(2,3,5)," + - "(3,0,3),(3,1,4),(3,2,9),(3,3,0)") - - sql("OPTIMIZE up ZORDER BY c1, `@c2`") - val res = sql("SELECT c1, `@c2` FROM up").collect() - - assert(res.length == 16) - - for (i <- target.indices) { - val t = target(i) - val r = res(i) - assert(t(0) == r.getInt(0)) - assert(t(1) == r.getInt(1)) - } - } - } - } - - def createParser: ParserInterface -} - -trait ZorderWithCodegenEnabledSuiteBase extends ZorderSuiteBase { - override def sparkConf(): SparkConf = { - val conf = super.sparkConf - conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") - conf - } -} - -trait ZorderWithCodegenDisabledSuiteBase extends ZorderSuiteBase { - override def sparkConf(): SparkConf = { - val conf = super.sparkConf - conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "false") - conf.set(SQLConf.CODEGEN_FACTORY_MODE.key, "NO_CODEGEN") - conf - } -} diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala deleted file mode 100644 index b891a7224a0..00000000000 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/benchmark/KyuubiBenchmarkBase.scala +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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.spark.sql.benchmark - -import java.io.{File, FileOutputStream, OutputStream} - -import scala.collection.JavaConverters._ - -import com.google.common.reflect.ClassPath -import org.scalatest.Assertions._ - -trait KyuubiBenchmarkBase { - var output: Option[OutputStream] = None - - private val prefix = { - val benchmarkClasses = ClassPath.from(Thread.currentThread.getContextClassLoader) - .getTopLevelClassesRecursive("org.apache.spark.sql").asScala.toArray - assert(benchmarkClasses.nonEmpty) - val benchmark = benchmarkClasses.find(_.load().getName.endsWith("Benchmark")) - val targetDirOrProjDir = - new File(benchmark.get.load().getProtectionDomain.getCodeSource.getLocation.toURI) - .getParentFile.getParentFile - if (targetDirOrProjDir.getName == "target") { - targetDirOrProjDir.getParentFile.getCanonicalPath + "/" - } else { - targetDirOrProjDir.getCanonicalPath + "/" - } - } - - def withHeader(func: => Unit): Unit = { - val version = System.getProperty("java.version").split("\\D+")(0).toInt - val jdkString = if (version > 8) s"-jdk$version" else "" - val resultFileName = - s"${this.getClass.getSimpleName.replace("$", "")}$jdkString-results.txt" - val dir = new File(s"${prefix}benchmarks/") - if (!dir.exists()) { - // scalastyle:off println - println(s"Creating ${dir.getAbsolutePath} for benchmark results.") - // scalastyle:on println - dir.mkdirs() - } - val file = new File(dir, resultFileName) - if (!file.exists()) { - file.createNewFile() - } - output = Some(new FileOutputStream(file)) - - func - - output.foreach { o => - if (o != null) { - o.close() - } - } - } -} diff --git a/extensions/spark/kyuubi-spark-authz/README.md b/extensions/spark/kyuubi-spark-authz/README.md index 55291122de3..9755db9d81c 100644 --- a/extensions/spark/kyuubi-spark-authz/README.md +++ b/extensions/spark/kyuubi-spark-authz/README.md @@ -37,12 +37,8 @@ build/mvn clean package -DskipTests -pl :kyuubi-spark-authz_2.12 -am -Dspark.ver - [x] 4.1.x - [x] 4.0.x - [x] 3.5.x (default) -- [x] 3.4.x -- [x] 3.3.x -- [ ] 3.2.x -- [ ] 3.1.x -- [ ] 3.0.x -- [ ] 2.4.x and earlier +- [ ] 3.4.x +- [ ] 3.3.x ### Supported Apache Ranger Versions @@ -55,10 +51,4 @@ build/mvn clean package -DskipTests -pl :kyuubi-spark-authz_2.12 -am -Dspark.ver - [x] 2.3.x - [x] 2.2.x - [x] 2.1.x -- [x] 2.0.x -- [x] 1.2.x -- [x] 1.1.x -- [x] 1.0.x -- [x] 0.7.x -- [ ] 0.6.x - +- [ ] 2.0.x diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala index f9b4af74c5f..1bac3642aeb 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala @@ -92,8 +92,6 @@ private[authz] object AuthZUtils { } lazy val SPARK_RUNTIME_VERSION: SemanticVersion = SemanticVersion(SPARK_VERSION) - lazy val isSparkV34OrGreater: Boolean = SPARK_RUNTIME_VERSION >= "3.4" - lazy val isSparkV35OrGreater: Boolean = SPARK_RUNTIME_VERSION >= "3.5" lazy val isSparkV40OrGreater: Boolean = SPARK_RUNTIME_VERSION >= "4.0" lazy val SCALA_RUNTIME_VERSION: SemanticVersion = diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/IcebergCatalogPrivilegesBuilderSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/IcebergCatalogPrivilegesBuilderSuite.scala index a64431b2269..8b92d902668 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/IcebergCatalogPrivilegesBuilderSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/IcebergCatalogPrivilegesBuilderSuite.scala @@ -20,7 +20,6 @@ package org.apache.kyuubi.plugin.spark.authz import org.apache.kyuubi.Utils import org.apache.kyuubi.plugin.spark.authz.OperationType._ import org.apache.kyuubi.plugin.spark.authz.ranger.AccessType -import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._ import org.apache.kyuubi.tags.IcebergTest import org.apache.kyuubi.util.AssertionUtils._ @@ -52,15 +51,11 @@ class IcebergCatalogPrivilegesBuilderSuite extends V2CommandsPrivilegesSuite { val plan = sql(s"DELETE FROM $catalogTable WHERE key = 1 ").queryExecution.analyzed val (inputs, outputs, operationType) = PrivilegesBuilder.build(plan, spark) assert(operationType === QUERY) - if (isSparkV34OrGreater) { - assert(inputs.size === 1) - val po = inputs.head - assertEqualsIgnoreCase(namespace)(po.dbname) - assertEqualsIgnoreCase(catalogTableShort)(po.objectName) - assertContains(po.columns, "key", "value") - } else { - assert(inputs.size === 0) - } + assert(inputs.size === 1) + val inputPo = inputs.head + assertEqualsIgnoreCase(namespace)(inputPo.dbname) + assertEqualsIgnoreCase(catalogTableShort)(inputPo.objectName) + assertContains(inputPo.columns, "key", "value") assert(outputs.size === 1) val po = outputs.head assert(po.actionType === PrivilegeObjectActionType.UPDATE) @@ -77,15 +72,11 @@ class IcebergCatalogPrivilegesBuilderSuite extends V2CommandsPrivilegesSuite { val plan = sql(s"UPDATE $catalogTable SET value = 'b' WHERE key = 1 ").queryExecution.analyzed val (inputs, outputs, operationType) = PrivilegesBuilder.build(plan, spark) assert(operationType === QUERY) - if (isSparkV35OrGreater) { - assert(inputs.size === 1) - val po = inputs.head - assertEqualsIgnoreCase(namespace)(po.dbname) - assertEqualsIgnoreCase(catalogTableShort)(po.objectName) - assertContains(po.columns, "key", "value") - } else { - assert(inputs.size === 0) - } + assert(inputs.size === 1) + val inputPo = inputs.head + assertEqualsIgnoreCase(namespace)(inputPo.dbname) + assertEqualsIgnoreCase(catalogTableShort)(inputPo.objectName) + assertContains(inputPo.columns, "key", "value") assert(outputs.size === 1) val po = outputs.head assert(po.actionType === PrivilegeObjectActionType.UPDATE) @@ -109,19 +100,15 @@ class IcebergCatalogPrivilegesBuilderSuite extends V2CommandsPrivilegesSuite { s"WHEN NOT MATCHED THEN INSERT *").queryExecution.analyzed val (inputs, outputs, operationType) = PrivilegesBuilder.build(plan, spark) assert(operationType === QUERY) - if (isSparkV35OrGreater) { - assert(inputs.size === 2) - val po = inputs.head - assert(po.actionType === PrivilegeObjectActionType.OTHER) - assert(po.privilegeObjectType === PrivilegeObjectType.TABLE_OR_VIEW) - assertEqualsIgnoreCase(namespace)(po.dbname) - assertEqualsIgnoreCase(table)(po.objectName) - assertContains(po.columns, "key", "value") - // The properties of RowLevelOperationTable are empty, so owner is none - assert(po.owner.isEmpty) - } else { - assert(inputs.size === 1) - } + assert(inputs.size === 2) + val inputPo = inputs.head + assert(inputPo.actionType === PrivilegeObjectActionType.OTHER) + assert(inputPo.privilegeObjectType === PrivilegeObjectType.TABLE_OR_VIEW) + assertEqualsIgnoreCase(namespace)(inputPo.dbname) + assertEqualsIgnoreCase(table)(inputPo.objectName) + assertContains(inputPo.columns, "key", "value") + // The properties of RowLevelOperationTable are empty, so owner is none + assert(inputPo.owner.isEmpty) val po0 = inputs.last assert(po0.actionType === PrivilegeObjectActionType.OTHER) assert(po0.privilegeObjectType === PrivilegeObjectType.TABLE_OR_VIEW) diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilderSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilderSuite.scala index 1e90e9f4e0b..864e00e7393 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilderSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilderSuite.scala @@ -31,7 +31,6 @@ import org.apache.kyuubi.plugin.spark.authz.OperationType._ import org.apache.kyuubi.plugin.spark.authz.RangerTestNamespace._ import org.apache.kyuubi.plugin.spark.authz.RangerTestUsers._ import org.apache.kyuubi.plugin.spark.authz.ranger.AccessType -import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._ import org.apache.kyuubi.util.AssertionUtils._ abstract class PrivilegesBuilderSuite extends KyuubiFunSuite with SparkSessionProvider { @@ -111,24 +110,6 @@ abstract class PrivilegesBuilderSuite extends KyuubiFunSuite with SparkSessionPr super.beforeEach() } - test("AlterDatabasePropertiesCommand") { - assume(SPARK_RUNTIME_VERSION <= "3.2") - val plan = sql("ALTER DATABASE default SET DBPROPERTIES (abc = '123')").queryExecution.analyzed - val (in, out, operationType) = PrivilegesBuilder.build(plan, spark) - assertResult(plan.getClass.getName)( - "org.apache.spark.sql.execution.command.AlterDatabasePropertiesCommand") - assert(operationType === ALTERDATABASE) - assert(in.isEmpty) - assert(out.size === 1) - val po = out.head - assert(po.actionType === PrivilegeObjectActionType.OTHER) - assert(po.privilegeObjectType === PrivilegeObjectType.DATABASE) - assert(po.catalog.isEmpty) - assertEqualsIgnoreCase(defaultDb)(po.dbname) - assertEqualsIgnoreCase(defaultDb)(po.objectName) - assert(po.columns.isEmpty) - } - test("AlterTableRenameCommand") { withTable(s"$reusedDb.efg") { t => withTable(s"${reusedTable}_old") { oldTable => @@ -159,51 +140,6 @@ abstract class PrivilegesBuilderSuite extends KyuubiFunSuite with SparkSessionPr } } - test("CreateDatabaseCommand") { - assume(SPARK_RUNTIME_VERSION <= "3.2") - withDatabase("CreateDatabaseCommand") { db => - val plan = sql(s"CREATE DATABASE $db").queryExecution.analyzed - val (in, out, operationType) = PrivilegesBuilder.build(plan, spark) - assertResult(plan.getClass.getName)( - "org.apache.spark.sql.execution.command.CreateDatabaseCommand") - assert(operationType === CREATEDATABASE) - assert(in.isEmpty) - assert(out.size === 1) - val po = out.head - assert(po.actionType === PrivilegeObjectActionType.OTHER) - assert(po.privilegeObjectType === PrivilegeObjectType.DATABASE) - assert(po.catalog.isEmpty) - assertEqualsIgnoreCase(db)(po.dbname) - assertEqualsIgnoreCase(db)(po.objectName) - assert(po.columns.isEmpty) - val accessType = ranger.AccessType(po, operationType, isInput = false) - assert(accessType === AccessType.CREATE) - } - } - - test("DropDatabaseCommand") { - assume(SPARK_RUNTIME_VERSION <= "3.2") - withDatabase("DropDatabaseCommand") { db => - sql(s"CREATE DATABASE $db") - val plan = sql(s"DROP DATABASE DropDatabaseCommand").queryExecution.analyzed - val (in, out, operationType) = PrivilegesBuilder.build(plan, spark) - assertResult(plan.getClass.getName)( - "org.apache.spark.sql.execution.command.DropDatabaseCommand") - assert(operationType === DROPDATABASE) - assert(in.isEmpty) - assert(out.size === 1) - val po = out.head - assert(po.actionType === PrivilegeObjectActionType.OTHER) - assert(po.privilegeObjectType === PrivilegeObjectType.DATABASE) - assert(po.catalog.isEmpty) - assertEqualsIgnoreCase(db)(po.dbname) - assertEqualsIgnoreCase(db)(po.objectName) - assert(po.columns.isEmpty) - val accessType = ranger.AccessType(po, operationType, isInput = false) - assert(accessType === AccessType.DROP) - } - } - test("AlterTableAddPartitionCommand") { val plan = sql(s"ALTER TABLE $reusedPartTable ADD IF NOT EXISTS PARTITION (pid=1)") .queryExecution.analyzed @@ -750,25 +686,6 @@ abstract class PrivilegesBuilderSuite extends KyuubiFunSuite with SparkSessionPr assert(out.size === 0) } - test("DescribeDatabaseCommand") { - assume(SPARK_RUNTIME_VERSION <= "3.2") - val plan = sql(s"DESC DATABASE $reusedDb").queryExecution.analyzed - val (in, out, operationType) = PrivilegesBuilder.build(plan, spark) - assert(operationType === DESCDATABASE) - assert(in.size === 1) - val po = in.head - assert(po.actionType === PrivilegeObjectActionType.OTHER) - assert(po.privilegeObjectType === PrivilegeObjectType.DATABASE) - assert(po.catalog.isEmpty) - assertEqualsIgnoreCase(reusedDb)(po.dbname) - assertEqualsIgnoreCase(reusedDb)(po.objectName) - assert(po.columns.isEmpty) - val accessType = ranger.AccessType(po, operationType, isInput = false) - assert(accessType === AccessType.USE) - - assert(out.size === 0) - } - test("SetDatabaseCommand") { try { val plan = sql(s"USE $reusedDb").queryExecution.analyzed @@ -1242,37 +1159,6 @@ class InMemoryPrivilegeBuilderSuite extends PrivilegesBuilderSuite { override protected val catalogImpl: String = "in-memory" // some hive version does not support set database location - test("AlterDatabaseSetLocationCommand") { - assume(SPARK_RUNTIME_VERSION <= "3.2") - val newLoc = spark.conf.get("spark.sql.warehouse.dir") + "/new_db_location" - val plan = sql(s"ALTER DATABASE default SET LOCATION '$newLoc'") - .queryExecution.analyzed - val (in, out, operationType) = PrivilegesBuilder.build(plan, spark) - assertResult(plan.getClass.getName)( - "org.apache.spark.sql.execution.command.AlterDatabaseSetLocationCommand") - assert(operationType === ALTERDATABASE_LOCATION) - assert(in.isEmpty) - assert(out.size === 2) - val po0 = out.head - assert(po0.actionType === PrivilegeObjectActionType.OTHER) - assert(po0.privilegeObjectType === PrivilegeObjectType.DATABASE) - assert(po0.catalog.isEmpty) - assertEqualsIgnoreCase(defaultDb)(po0.dbname) - assertEqualsIgnoreCase(defaultDb)(po0.objectName) - assert(po0.columns.isEmpty) - val accessType0 = ranger.AccessType(po0, operationType, isInput = false) - assert(accessType0 === AccessType.ALTER) - - val po1 = out.last - assert(po1.actionType === PrivilegeObjectActionType.OTHER) - assert(po1.catalog.isEmpty) - assertEqualsIgnoreCase(defaultDb)(po0.dbname) - assertEqualsIgnoreCase(defaultDb)(po0.objectName) - assert(po1.columns.isEmpty) - val accessType1 = ranger.AccessType(po1, operationType, isInput = false) - assert(accessType1 === AccessType.WRITE) - } - test("CreateDataSourceTableAsSelectCommand") { val plan = sql(s"CREATE TABLE CreateDataSourceTableAsSelectCommand USING parquet" + s" AS SELECT key, value FROM $reusedTable") diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/V2CommandsPrivilegesSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/V2CommandsPrivilegesSuite.scala index 40cfc38774e..6647780e72a 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/V2CommandsPrivilegesSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/V2CommandsPrivilegesSuite.scala @@ -26,7 +26,6 @@ import org.apache.kyuubi.plugin.spark.authz.OperationType._ import org.apache.kyuubi.plugin.spark.authz.RangerTestNamespace._ import org.apache.kyuubi.plugin.spark.authz.ranger.AccessType import org.apache.kyuubi.plugin.spark.authz.serde.{Database, DB_COMMAND_SPECS} -import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._ import org.apache.kyuubi.util.AssertionUtils._ abstract class V2CommandsPrivilegesSuite extends PrivilegesBuilderSuite { @@ -161,11 +160,7 @@ abstract class V2CommandsPrivilegesSuite extends PrivilegesBuilderSuite { assertEqualsIgnoreCase(namespace)(po.dbname) assertEqualsIgnoreCase(table)(po.objectName) assert(po.columns.isEmpty) - if (isSparkV34OrGreater) { - checkV2TableOwner(po) - } else { - assert(po.owner.isEmpty) - } + checkV2TableOwner(po) val accessType = AccessType(po, operationType, isInput = false) assert(accessType === AccessType.CREATE) } @@ -197,11 +192,7 @@ abstract class V2CommandsPrivilegesSuite extends PrivilegesBuilderSuite { assertEqualsIgnoreCase(namespace)(po.dbname) assertEqualsIgnoreCase(table)(po.objectName) assert(po.columns.isEmpty) - if (isSparkV34OrGreater) { - checkV2TableOwner(po) - } else { - assert(po.owner.isEmpty) - } + checkV2TableOwner(po) val accessType = AccessType(po, operationType, isInput = false) assert(accessType === AccessType.CREATE) } diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/DeltaCatalogRangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/DeltaCatalogRangerSparkExtensionSuite.scala index db5ad5d50ad..edade2c1792 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/DeltaCatalogRangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/DeltaCatalogRangerSparkExtensionSuite.scala @@ -23,7 +23,6 @@ import org.apache.kyuubi.plugin.spark.authz.AccessControlException import org.apache.kyuubi.plugin.spark.authz.RangerTestNamespace._ import org.apache.kyuubi.plugin.spark.authz.RangerTestUsers._ import org.apache.kyuubi.plugin.spark.authz.ranger.DeltaCatalogRangerSparkExtensionSuite._ -import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils.isSparkV35OrGreater import org.apache.kyuubi.tags.DeltaTest import org.apache.kyuubi.util.AssertionUtils._ @@ -543,9 +542,7 @@ class DeltaCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { // "Cannot drop column from a struct type with a single field: // StructType(StructField(birthDate,TimestampType,true))". // For details, see https://github.com/delta-io/delta/pull/1822 - if (isSparkV35OrGreater) { - doAs(admin, sql(replaceColumnsSql)) - } + doAs(admin, sql(replaceColumnsSql)) } } } diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/IcebergCatalogRangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/IcebergCatalogRangerSparkExtensionSuite.scala index c4f99623a08..740e17a0722 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/IcebergCatalogRangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/IcebergCatalogRangerSparkExtensionSuite.scala @@ -28,7 +28,6 @@ import org.apache.kyuubi.Utils import org.apache.kyuubi.plugin.spark.authz.AccessControlException import org.apache.kyuubi.plugin.spark.authz.RangerTestNamespace._ import org.apache.kyuubi.plugin.spark.authz.RangerTestUsers._ -import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._ import org.apache.kyuubi.tags.IcebergTest import org.apache.kyuubi.util.AssertionUtils._ @@ -113,14 +112,8 @@ class IcebergCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite withSingleCallEnabled { interceptEndsWith[AccessControlException](doAs(someone, sql(mergeIntoSql)))( - if (isSparkV35OrGreater) { - s"does not have [select] privilege on [$namespace1/table1/city" + - s",$namespace1/$table1/id,$namespace1/$table1/name]" - } else { - "does not have " + - s"[select] privilege on [$namespace1/$table1/city,$namespace1/$table1/id,$namespace1/$table1/name]," + - s" [update] privilege on [$bobNamespace/$bobSelectTable]" - }) + s"does not have [select] privilege on [$namespace1/table1/city" + + s",$namespace1/$table1/id,$namespace1/$table1/name]") interceptEndsWith[AccessControlException] { doAs(bob, sql(mergeIntoSql)) @@ -138,14 +131,10 @@ class IcebergCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite doAs( someone, sql(s"UPDATE $catalogV2.$namespace1.$table1 SET city='Guangzhou' WHERE id=1")) - }(if (isSparkV35OrGreater) { - s"does not have [select] privilege on " + - s"[$namespace1/$table1/_file,$namespace1/$table1/_pos," + - s"$namespace1/$table1/id,$namespace1/$table1/name,$namespace1/$table1/city], " + - s"[update] privilege on [$namespace1/$table1]" - } else { - s"does not have [update] privilege on [$namespace1/$table1]" - }) + }(s"does not have [select] privilege on " + + s"[$namespace1/$table1/_file,$namespace1/$table1/_pos," + + s"$namespace1/$table1/id,$namespace1/$table1/name,$namespace1/$table1/city], " + + s"[update] privilege on [$namespace1/$table1]") doAs( admin, @@ -159,14 +148,10 @@ class IcebergCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite // DeleteFromTable interceptEndsWith[AccessControlException] { doAs(someone, sql(s"DELETE FROM $catalogV2.$namespace1.$table1 WHERE id=2")) - }(if (isSparkV34OrGreater) { - s"does not have [select] privilege on " + - s"[$namespace1/$table1/_file,$namespace1/$table1/_pos," + - s"$namespace1/$table1/city,$namespace1/$table1/id,$namespace1/$table1/name], " + - s"[update] privilege on [$namespace1/$table1]" - } else { - s"does not have [update] privilege on [$namespace1/$table1]" - }) + }(s"does not have [select] privilege on " + + s"[$namespace1/$table1/_file,$namespace1/$table1/_pos," + + s"$namespace1/$table1/city,$namespace1/$table1/id,$namespace1/$table1/name], " + + s"[update] privilege on [$namespace1/$table1]") interceptEndsWith[AccessControlException] { doAs(bob, sql(s"DELETE FROM $catalogV2.$bobNamespace.$bobSelectTable WHERE id=2")) diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/PaimonCatalogRangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/PaimonCatalogRangerSparkExtensionSuite.scala index 1cdc4b57a9a..c5419435cfa 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/PaimonCatalogRangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/PaimonCatalogRangerSparkExtensionSuite.scala @@ -578,58 +578,56 @@ class PaimonCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { } test("Producers") { - if (isSparkV34OrGreater) { - withCleanTmpResources(Seq( - (s"$catalogV2.$namespace1.$table1", "table"))) { - try { - doAs(admin, sql(createTableSql(namespace1, table1))) - doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (1, 'a'), (2, 'b')")) - - var currentCatalogName = - doAs(admin, spark.sessionState.catalogManager.currentCatalog.name()) - - doAs(admin, sql(s"use $catalogV2")) - currentCatalogName = doAs(admin, spark.sessionState.catalogManager.currentCatalog.name()) - assert(currentCatalogName.equals(catalogV2)) - - // Create Tag - val createTagSql = s"Call sys.create_tag(table =>" + - s"'$catalogV2.$namespace1.$table1', tag => 'test_tag', snapshot => 1)" - interceptEndsWith[AccessControlException] { - doAs(table1OnlyUserForNs, sql(createTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - interceptEndsWith[AccessControlException] { - doAs(someone, sql(createTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - doAs(admin, sql(createTagSql)) - - // Delete Tag - val deleteTagSql = s"Call sys.delete_tag(table =>" + - s"'$catalogV2.$namespace1.$table1', tag => 'test_tag')" - interceptEndsWith[AccessControlException] { - doAs(table1OnlyUserForNs, sql(deleteTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - interceptEndsWith[AccessControlException] { - doAs(someone, sql(deleteTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - doAs(admin, sql(deleteTagSql)) - - // Rollback - doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (3, 'a'), (4, 'b')")) - doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (5, 'a'), (6, 'b')")) - val rollbackTagSql = s"Call sys.rollback(table =>" + - s"'$catalogV2.$namespace1.$table1', version => '2')" - interceptEndsWith[AccessControlException] { - doAs(table1OnlyUserForNs, sql(rollbackTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - interceptEndsWith[AccessControlException] { - doAs(someone, sql(rollbackTagSql)) - }(s"does not have [alter] privilege on [$namespace1/$table1]") - doAs(admin, sql(rollbackTagSql)) - - } finally { - doAs(admin, sql(s"use spark_catalog")) - } + withCleanTmpResources(Seq( + (s"$catalogV2.$namespace1.$table1", "table"))) { + try { + doAs(admin, sql(createTableSql(namespace1, table1))) + doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (1, 'a'), (2, 'b')")) + + var currentCatalogName = + doAs(admin, spark.sessionState.catalogManager.currentCatalog.name()) + + doAs(admin, sql(s"use $catalogV2")) + currentCatalogName = doAs(admin, spark.sessionState.catalogManager.currentCatalog.name()) + assert(currentCatalogName.equals(catalogV2)) + + // Create Tag + val createTagSql = s"Call sys.create_tag(table =>" + + s"'$catalogV2.$namespace1.$table1', tag => 'test_tag', snapshot => 1)" + interceptEndsWith[AccessControlException] { + doAs(table1OnlyUserForNs, sql(createTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + interceptEndsWith[AccessControlException] { + doAs(someone, sql(createTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + doAs(admin, sql(createTagSql)) + + // Delete Tag + val deleteTagSql = s"Call sys.delete_tag(table =>" + + s"'$catalogV2.$namespace1.$table1', tag => 'test_tag')" + interceptEndsWith[AccessControlException] { + doAs(table1OnlyUserForNs, sql(deleteTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + interceptEndsWith[AccessControlException] { + doAs(someone, sql(deleteTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + doAs(admin, sql(deleteTagSql)) + + // Rollback + doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (3, 'a'), (4, 'b')")) + doAs(admin, sql(s"INSERT INTO $catalogV2.$namespace1.$table1 VALUES (5, 'a'), (6, 'b')")) + val rollbackTagSql = s"Call sys.rollback(table =>" + + s"'$catalogV2.$namespace1.$table1', version => '2')" + interceptEndsWith[AccessControlException] { + doAs(table1OnlyUserForNs, sql(rollbackTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + interceptEndsWith[AccessControlException] { + doAs(someone, sql(rollbackTagSql)) + }(s"does not have [alter] privilege on [$namespace1/$table1]") + doAs(admin, sql(rollbackTagSql)) + + } finally { + doAs(admin, sql(s"use spark_catalog")) } } } diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala index c97d2237543..8805f0cd527 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala @@ -25,15 +25,12 @@ import scala.util.Try import org.apache.hadoop.security.UserGroupInformation import org.apache.logging.log4j.Level import org.apache.spark.SparkConf -import org.apache.spark.sql.{DataFrame, Row, SparkSessionExtensions} +import org.apache.spark.sql.{Row, SparkSessionExtensions} import org.apache.spark.sql.catalyst.analysis.NoSuchTableException import org.apache.spark.sql.catalyst.catalog.HiveTableRelation -import org.apache.spark.sql.catalyst.expressions.PythonUDF import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.datasources.LogicalRelation -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.kyuubi.{KyuubiFunSuite, Utils} import org.apache.kyuubi.plugin.lineage.Lineage @@ -946,15 +943,9 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { doAs( someone, sql(s"SELECT id as new_id, name, max_scope FROM $db1.$view1".stripMargin).show())) - if (isSparkV35OrGreater) { - assert(e2.getMessage.contains( - s"does not have [select] privilege on " + - s"[$db1/$view1/id,$db1/$view1/max_scope,$db1/$view1/name]")) - } else { - assert(e2.getMessage.contains( - s"does not have [select] privilege on " + - s"[$db1/$view1/name,$db1/$view1/id,$db1/$view1/max_scope]")) - } + assert(e2.getMessage.contains( + s"does not have [select] privilege on " + + s"[$db1/$view1/id,$db1/$view1/max_scope,$db1/$view1/name]")) } } } @@ -1311,13 +1302,8 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { s""" |CREATE TABLE IF NOT EXISTS $db1.$table1(id int, scope int) |LOCATION '$path'""".stripMargin)))( - if (!isSparkV35OrGreater) { - s"does not have [create] privilege on [$db1/$table1], " + - s"[write] privilege on [[$path, $path/]]" - } else { - s"does not have [create] privilege on [$db1/$table1], " + - s"[write] privilege on [[file://$path, file://$path/]]" - }) + s"does not have [create] privilege on [$db1/$table1], " + + s"[write] privilege on [[file://$path, file://$path/]]") doAs( admin, sql( @@ -1346,11 +1332,7 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { |AS |SELECT * FROM $db1.$table1 |""".stripMargin)))( - if (!isSparkV35OrGreater) { - s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope], " + - s"[create] privilege on [$db1/$table2/id,$db1/$table2/scope], " + - s"[write] privilege on [[$path, $path/]]" - } else if (isSparkV40OrGreater) { + if (isSparkV40OrGreater) { // Spark 4.0 no longer propagates CTAS output columns into the create privilege s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope], " + s"[create] privilege on [$db1/$table2], " + @@ -1529,55 +1511,6 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { } } - test("[KYUUBI #5594][AUTHZ] BuildQuery should respect normal node's input ") { - assume(!isSparkV35OrGreater, "mapInPandas not supported after spark 3.5") - val db1 = defaultDb - val table1 = "table1" - val view1 = "view1" - withSingleCallEnabled { - withCleanTmpResources(Seq((s"$db1.$table1", "table"), (s"$db1.$view1", "view"))) { - doAs(admin, sql(s"CREATE TABLE IF NOT EXISTS $db1.$table1 (id int, scope int)")) - doAs(admin, sql(s"CREATE VIEW $db1.$view1 AS SELECT * FROM $db1.$table1")) - - val table = spark.read.table(s"$db1.$table1") - val mapTableInPandasUDF = PythonUDF( - "mapInPandasUDF", - null, - StructType(Seq(StructField("id", IntegerType), StructField("scope", IntegerType))), - table.queryExecution.analyzed.output, - 205, - true) - interceptContains[AccessControlException]( - doAs( - someone, - invokeAs( - table, - "mapInPandas", - (classOf[PythonUDF], mapTableInPandasUDF)) - .asInstanceOf[DataFrame].select(col("id"), col("scope")).limit(1).show(true)))( - s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]") - - val view = spark.read.table(s"$db1.$view1") - val mapViewInPandasUDF = PythonUDF( - "mapInPandasUDF", - null, - StructType(Seq(StructField("id", IntegerType), StructField("scope", IntegerType))), - view.queryExecution.analyzed.output, - 205, - true) - interceptContains[AccessControlException]( - doAs( - someone, - invokeAs( - view, - "mapInPandas", - (classOf[PythonUDF], mapViewInPandasUDF)) - .asInstanceOf[DataFrame].select(col("id"), col("scope")).limit(1).show(true)))( - s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]") - } - } - } - test("[KYUUBI #5594][AUTHZ] BuildQuery should respect sort agg input") { val db1 = defaultDb val table1 = "table1" diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveConnectorUtils.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveConnectorUtils.scala index 0fb3bf0dda3..29753abac1e 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveConnectorUtils.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveConnectorUtils.scala @@ -22,8 +22,7 @@ import java.net.URI import scala.util.Try -import org.apache.hadoop.fs.{FileStatus, Path} -import org.apache.hadoop.hive.ql.plan.{FileSinkDesc, TableDesc} +import org.apache.hadoop.fs.Path import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} @@ -31,43 +30,14 @@ import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, Cat import org.apache.spark.sql.connector.catalog.TableChange import org.apache.spark.sql.connector.catalog.TableChange._ import org.apache.spark.sql.execution.command.CommandUtils -import org.apache.spark.sql.execution.datasources.{PartitionDirectory, PartitionedFile} -import org.apache.spark.sql.hive.execution.HiveFileFormat +import org.apache.spark.sql.execution.datasources.PartitionedFile import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, MapType, StructField, StructType} -import org.apache.kyuubi.util.reflect.{DynClasses, DynConstructors, DynMethods} -import org.apache.kyuubi.util.reflect.ReflectUtils.invokeAs +import org.apache.kyuubi.util.reflect.{DynClasses, DynMethods} object HiveConnectorUtils extends Logging { - def getHiveFileFormat(fileSinkConf: FileSinkDesc): HiveFileFormat = - Try { // SPARK-43186: 3.5.0 - DynConstructors.builder() - .impl(classOf[HiveFileFormat], classOf[FileSinkDesc]) - .build[HiveFileFormat]() - .newInstance(fileSinkConf) - }.recover { case _: Exception => - val shimFileSinkDescClz = DynClasses.builder() - .impl("org.apache.spark.sql.hive.HiveShim$ShimFileSinkDesc") - .build() - val shimFileSinkDesc = DynConstructors.builder() - .impl( - "org.apache.spark.sql.hive.HiveShim$ShimFileSinkDesc", - classOf[String], - classOf[TableDesc], - classOf[Boolean]) - .build[AnyRef]() - .newInstance( - fileSinkConf.getDirName.toString, - fileSinkConf.getTableInfo, - fileSinkConf.getCompressed.asInstanceOf[JBoolean]) - DynConstructors.builder() - .impl(classOf[HiveFileFormat], shimFileSinkDescClz) - .build[HiveFileFormat]() - .newInstance(shimFileSinkDesc) - }.get - // `serdeName` widened the case-class `apply` from 6 to 7 args. `DynMethods.invoke` // truncates trailing args to the matched arity, so the trailing `serdeName` is // silently dropped on the 6-arg impl. @@ -258,13 +228,6 @@ object HiveConnectorUtils extends Logging { } // scalastyle:on parameter.number - def partitionedFilePath(file: PartitionedFile): String = - Try { // SPARK-41970: 3.4.0 - invokeAs[String](file, "urlEncodedPath") - }.recover { case _: Exception => - invokeAs[String](file, "filePath") - }.get - def splitFiles( sparkSession: SparkSession, file: AnyRef, @@ -356,51 +319,6 @@ object HiveConnectorUtils extends Logging { isSplitable, maxSplitBytes, partitionValues) - }.recover { case _: Exception => - DynMethods - .builder("splitFiles") - .impl( - "org.apache.spark.sql.execution.PartitionedFileUtil", - classOf[SparkSession], - classOf[FileStatus], - classOf[Path], - classOf[Boolean], - classOf[Long], - classOf[InternalRow]) - .buildChecked() - .invokeChecked[Seq[PartitionedFile]]( - null, - sparkSession, - file, - filePath, - isSplitable, - maxSplitBytes, - partitionValues) - }.get - - def createPartitionDirectory(values: InternalRow, files: Seq[FileStatus]): PartitionDirectory = - Try { // SPARK-43039: 3.5.0 - new DynMethods.Builder("apply") - .impl(classOf[PartitionDirectory], classOf[InternalRow], classOf[Array[FileStatus]]) - .buildChecked() - .asStatic() - .invoke[PartitionDirectory](values, files.toArray) - }.recover { case _: Exception => - new DynMethods.Builder("apply") - .impl(classOf[PartitionDirectory], classOf[InternalRow], classOf[Seq[FileStatus]]) - .buildChecked() - .asStatic() - .invoke[PartitionDirectory](values, files) - }.get - - def getPartitionFilePath(file: AnyRef): Path = - Try { // SPARK-43039: 3.5.0 - new DynMethods.Builder("getPath") - .impl("org.apache.spark.sql.execution.datasources.FileStatusWithMetadata") - .build() - .invoke[Path](file) - }.recover { case _: Exception => - file.asInstanceOf[FileStatus].getPath }.get private def calculateMultipleLocationSizes( diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala index 35f5189b055..3d5c27ebb88 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala @@ -658,36 +658,6 @@ private object HiveTableCatalog extends Logging { HIVE_OUTPUT_FORMAT, HIVE_INPUT_FORMAT) - /** - * Attaches the KSHC catalog name to a V1 [[TableIdentifier]]. - * - * Since Spark 3.4 (SPARK-46283), [[TableIdentifier]] carries a `catalog` field. Spark 4.1's - * `SessionCatalog.requireTableExists` reads it via `name.catalog.get`, which throws - * `NoSuchElementException` when it is `None`. KSHC deliberately avoids the wrong default - * catalog name (`spark_catalog`) by attaching its *own* catalog name here instead of relying - * on the `spark.sql.legacy.v1IdentifierNoCatalog` workaround, so `TableIdentifier.catalog` is - * never `None` on Spark 3.4+. On Spark 3.3 (no `catalog` field) the 3-arg constructor is - * absent and the identifier is returned unchanged. - */ - private def attachCatalogName( - identifier: TableIdentifier, - catalogName: String): TableIdentifier = { - Try { // Spark 3.4+ (SPARK-46283): TableIdentifier(table, database, catalog) - DynConstructors.builder() - .impl( - classOf[TableIdentifier], - classOf[String], - classOf[Option[String]], - classOf[Option[String]]) - .buildChecked() - .invokeChecked[TableIdentifier]( - null, - identifier.table, - identifier.database, - Some(catalogName)) - }.recover { case _: Exception => identifier }.get - } - private def toCatalogDatabase( db: String, metadata: util.Map[String, String], @@ -785,15 +755,12 @@ private object HiveTableCatalog extends Logging { def asMultipartIdentifier: Seq[String] = ident.namespace :+ ident.name - def asTableIdentifier(catalogName: String): TableIdentifier = { - val base = ident.namespace match { - case ns if ns.isEmpty => TableIdentifier(ident.name) - case Array(dbName) => TableIdentifier(ident.name, Some(dbName)) - case _ => - throw KyuubiHiveConnectorException( - s"$quoted is not a valid TableIdentifier as it has more than 2 name parts.") - } - attachCatalogName(base, catalogName) + def asTableIdentifier(catalogName: String): TableIdentifier = ident.namespace match { + case ns if ns.isEmpty => TableIdentifier(ident.name, None, Some(catalogName)) + case Array(dbName) => TableIdentifier(ident.name, Some(dbName), Some(catalogName)) + case _ => + throw KyuubiHiveConnectorException( + s"$quoted is not a valid TableIdentifier as it has more than 2 name parts.") } } diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveFileIndex.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveFileIndex.scala index ec9fd3aeeb1..899e32fb563 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveFileIndex.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveFileIndex.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.connector.catalog.CatalogPlugin import org.apache.spark.sql.execution.datasources._ import org.apache.spark.sql.types.StructType -import org.apache.kyuubi.spark.connector.hive.{HiveConnectorUtils, HiveTableCatalog, KyuubiHiveConnectorException} +import org.apache.kyuubi.spark.connector.hive.{HiveTableCatalog, KyuubiHiveConnectorException} class HiveCatalogFileIndex( sparkSession: SparkSession, @@ -183,9 +183,9 @@ class HiveInMemoryFileIndex( } val selectedPartitions = if (partitionSpec().partitionColumns.isEmpty) { - HiveConnectorUtils.createPartitionDirectory( + PartitionDirectory( InternalRow.empty, - allFiles().filter(isNonEmptyFile)) :: Nil + allFiles().filter(isNonEmptyFile).toArray) :: Nil } else { if (recursiveFileLookup) { throw new IllegalArgumentException( @@ -202,7 +202,7 @@ class HiveInMemoryFileIndex( // Directory does not exist, or has no children files Nil } - val partDir = HiveConnectorUtils.createPartitionDirectory(values, files) + val partDir = PartitionDirectory(values, files.toArray) // Update Partition Directory -> binding Hive part map updatePartDirHivePartitionMapping(partDir, partPath) diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HivePartitionReaderFactory.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HivePartitionReaderFactory.scala index 6a2a7f1d6ed..4ff15b8879f 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HivePartitionReaderFactory.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HivePartitionReaderFactory.scala @@ -41,8 +41,6 @@ import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types._ import org.apache.spark.util.SerializableConfiguration -import org.apache.kyuubi.spark.connector.hive.HiveConnectorUtils - case class HivePartitionReaderFactory( sqlConf: SQLConf, broadcastHiveConf: Broadcast[SerializableConfiguration], @@ -112,7 +110,7 @@ case class HivePartitionReaderFactory( val jobConf = new JobConf(broadcastHiveConf.value.value) - val filePath = new Path(new URI(HiveConnectorUtils.partitionedFilePath(file))) + val filePath = new Path(new URI(file.urlEncodedPath)) if (tableDesc != null) { configureJobPropertiesForStorageHandler(tableDesc, jobConf, true) diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveRuntimeFilterSupport.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveRuntimeFilterSupport.scala index 764b4697940..f4b0c994184 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveRuntimeFilterSupport.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveRuntimeFilterSupport.scala @@ -35,10 +35,6 @@ import org.apache.spark.sql.types.StructType * runtime filter, so translation here handles `In` only. Any filter whose * attribute does not match a known partition column is dropped; drops are * logged at DEBUG. - * - * We deliberately use the V1 `SupportsRuntimeFiltering` instead of the newer - * `SupportsRuntimeV2Filtering` to keep this connector compilable against - * Spark 3.3, where `SupportsRuntimeV2Filtering` was introduced in Spark 3.4. */ object HiveRuntimeFilterSupport extends Logging { diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala index 56725d4f177..088fa6167e2 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala @@ -37,6 +37,7 @@ import org.apache.spark.sql.types.StructType import org.apache.spark.util.SerializableConfiguration import org.apache.kyuubi.spark.connector.hive.{HiveConnectorUtils, KyuubiHiveConnectorException} +import org.apache.kyuubi.util.reflect.ReflectUtils.invokeAs case class HiveScan( sparkSession: SparkSession, @@ -122,7 +123,7 @@ case class HiveScan( partition.values } partition.files.asInstanceOf[Seq[AnyRef]].flatMap { file => - val filePath = HiveConnectorUtils.getPartitionFilePath(file) + val filePath = invokeAs[Path](file, "getPath") val partFiles = HiveConnectorUtils.splitFiles( sparkSession = sparkSession, file = file, @@ -141,7 +142,7 @@ case class HiveScan( } if (splitFiles.length == 1) { - val path = new Path(HiveConnectorUtils.partitionedFilePath(splitFiles(0))) + val path = new Path(splitFiles(0).urlEncodedPath) if (!isSplitable(path) && splitFiles(0).length > sparkSession.sparkContext.getConf.getOption("spark.io.warning.largeFileThreshold") .getOrElse("1024000000").toLong) { diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/SparkFilePartitionReader.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/SparkFilePartitionReader.scala index f785694d125..2dade15db59 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/SparkFilePartitionReader.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/SparkFilePartitionReader.scala @@ -26,8 +26,6 @@ import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupporte import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper.InputFileBlockHolder import org.apache.spark.sql.internal.SQLConf -import org.apache.kyuubi.spark.connector.hive.HiveConnectorUtils - // scalastyle:off line.size.limit // copy from https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FilePartitionReader.scala // scalastyle:on line.size.limit @@ -101,7 +99,7 @@ class SparkFilePartitionReader[T](readers: Iterator[HivePartitionedFileReader[T] // Sets InputFileBlockHolder for the file block's information val file = reader.file InputFileBlockHolder.set( - HiveConnectorUtils.partitionedFilePath(file), + file.urlEncodedPath, file.start, file.length) reader diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/write/HiveWrite.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/write/HiveWrite.scala index 842f6bac8bc..a7f139595a8 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/write/HiveWrite.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/write/HiveWrite.scala @@ -39,12 +39,11 @@ import org.apache.spark.sql.connector.write.{BatchWrite, LogicalWriteInfo, Requi import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, WriteJobDescription} import org.apache.spark.sql.execution.datasources.v2.FileBatchWrite import org.apache.spark.sql.execution.metric.SQLMetric -import org.apache.spark.sql.hive.execution.HiveOptions +import org.apache.spark.sql.hive.execution.{HiveFileFormat, HiveOptions} import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper.{HiveClientImpl, StructTypeHelper} import org.apache.spark.sql.types.StructType import org.apache.spark.util.SerializableConfiguration -import org.apache.kyuubi.spark.connector.hive.HiveConnectorUtils.getHiveFileFormat import org.apache.kyuubi.spark.connector.hive.HiveTableCatalog case class HiveWrite( @@ -150,7 +149,7 @@ case class HiveWrite( pathName: String, customPartitionLocations: Map[TablePartitionSpec, String], options: Map[String, String]): WriteJobDescription = { - val hiveFileFormat = getHiveFileFormat(fileSinkConf) + val hiveFileFormat = new HiveFileFormat(fileSinkConf) val dataSchema = StructType( info.schema().filterNot(field => isPartitionColumn(field.name))) val outputWriterFactory = hiveFileFormat.prepareWrite(sparkSession, job, options, dataSchema) diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/command/DropNamespaceSuite.scala b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/command/DropNamespaceSuite.scala index 9f4a64efc2c..830c137bdef 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/command/DropNamespaceSuite.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/command/DropNamespaceSuite.scala @@ -20,7 +20,6 @@ package org.apache.kyuubi.spark.connector.hive.command import org.apache.spark.sql.{AnalysisException, Row} import org.apache.spark.sql.types.{StringType, StructType} -import org.apache.kyuubi.spark.connector.common.SparkUtils.SPARK_RUNTIME_VERSION import org.apache.kyuubi.spark.connector.hive.command.DDLCommandTestUtils.{V1_COMMAND_VERSION, V2_COMMAND_VERSION} import org.apache.kyuubi.util.AssertionUtils.interceptContains @@ -75,11 +74,7 @@ trait DropNamespaceSuiteBase extends DDLCommandTestUtils { // $catalog.ns.table is present, thus $catalog.ns cannot be dropped. interceptContains[AnalysisException] { sql(s"DROP NAMESPACE $catalogName.$namespace") - }(if (SPARK_RUNTIME_VERSION >= "3.4") { - s"[SCHEMA_NOT_EMPTY] Cannot drop a schema `$namespace` because it contains objects" - } else { - "Use CASCADE option to drop a non-empty database" - }) + }(s"[SCHEMA_NOT_EMPTY] Cannot drop a schema `$namespace` because it contains objects") sql(s"DROP TABLE $catalogName.$namespace.table") diff --git a/extensions/spark/kyuubi-spark-lineage/README.md b/extensions/spark/kyuubi-spark-lineage/README.md index 955ad06c352..61026a52354 100644 --- a/extensions/spark/kyuubi-spark-lineage/README.md +++ b/extensions/spark/kyuubi-spark-lineage/README.md @@ -37,6 +37,6 @@ build/mvn clean package -DskipTests -pl :kyuubi-spark-lineage_2.12 -am -Dspark.v - [x] 4.1.x - [x] 4.0.x - [x] 3.5.x (default) -- [x] 3.4.x -- [x] 3.3.x +- [ ] 3.4.x +- [ ] 3.3.x diff --git a/extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala b/extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala index 27d8ae24677..665664ae9a1 100644 --- a/extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala +++ b/extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala @@ -38,7 +38,6 @@ import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.kyuubi.plugin.lineage.Lineage -import org.apache.kyuubi.plugin.lineage.helper.SparkListenerHelper.SPARK_RUNTIME_VERSION import org.apache.kyuubi.util.reflect.ReflectUtils._ trait LineageParser { @@ -211,12 +210,7 @@ trait LineageParser { val commandPlan = getField[LogicalPlan](plan, "commandLogicalPlan") extractColumnsLineage(commandPlan, parentColumnsLineage, inputTablesByPlan) case p if p.nodeName == "AlterViewAsCommand" => - val query = - if (SPARK_RUNTIME_VERSION <= "3.1") { - sparkSession.sessionState.analyzer.execute(getQuery(plan)) - } else { - getQuery(plan) - } + val query = getQuery(plan) val view = getV1TableName(getField[TableIdentifier](plan, "name").unquotedString) extractColumnsLineage(query, parentColumnsLineage, inputTablesByPlan).map { case (k, v) => k.withName(s"$view.${k.name}") -> v @@ -228,12 +222,7 @@ trait LineageParser { val view = getV1TableName(getField[TableIdentifier](plan, "name").unquotedString) val outputCols = getField[Seq[(String, Option[String])]](plan, "userSpecifiedColumns").map(_._1) - val query = - if (SPARK_RUNTIME_VERSION <= "3.1") { - sparkSession.sessionState.analyzer.execute(getField[LogicalPlan](plan, "child")) - } else { - getField[LogicalPlan](plan, "plan") - } + val query = getField[LogicalPlan](plan, "plan") val lineages = extractColumnsLineage( query, diff --git a/extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/RowLevelCatalogLineageParserSuite.scala b/extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/RowLevelCatalogLineageParserSuite.scala index 966fb708882..0390254d45d 100644 --- a/extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/RowLevelCatalogLineageParserSuite.scala +++ b/extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/RowLevelCatalogLineageParserSuite.scala @@ -18,7 +18,6 @@ package org.apache.kyuubi.plugin.lineage.helper import org.apache.kyuubi.plugin.lineage.Lineage -import org.apache.kyuubi.plugin.lineage.helper.SparkListenerHelper.SPARK_RUNTIME_VERSION class RowLevelCatalogLineageParserSuite extends SparkSQLLineageParserHelperSuite { @@ -26,9 +25,6 @@ class RowLevelCatalogLineageParserSuite extends SparkSQLLineageParserHelperSuite "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog" test("columns lineage extract - WriteDelta") { - assume( - SPARK_RUNTIME_VERSION >= "3.5", - "WriteDelta is only supported in SPARK_RUNTIME_VERSION >= 3.5") val ddls = """ |create table v2_catalog.db.target_t(pk int not null, name string, price float) @@ -110,9 +106,6 @@ class RowLevelCatalogLineageParserSuite extends SparkSQLLineageParserHelperSuite } test("columns lineage extract - ReplaceData") { - assume( - SPARK_RUNTIME_VERSION >= "3.5", - "ReplaceData[SPARK-43963] for merge into is supported in SPARK_RUNTIME_VERSION >= 3.5") val ddls = """ |create table v2_catalog.db.target_t(id int, name string, price float) diff --git a/externals/kyuubi-spark-sql-engine/src/main/resources/python/execute_python.py b/externals/kyuubi-spark-sql-engine/src/main/resources/python/execute_python.py index 1fff10a676e..950d56a105a 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/resources/python/execute_python.py +++ b/externals/kyuubi-spark-sql-engine/src/main/resources/python/execute_python.py @@ -28,9 +28,6 @@ import base64 from glob import glob -if sys.version_info[0] < 3: - sys.exit("Python < 3 is unsupported.") - os.environ["PYSPARK_PYTHON"] = os.environ.get("PYSPARK_PYTHON", sys.executable) # add pyspark to sys.path @@ -59,15 +56,7 @@ # import kyuubi_util after preparing sys.path import kyuubi_util -# ast api is changed after python 3.8, see https://github.com/ipython/ipython/pull/11593 -if sys.version_info >= (3, 8): - from ast import Module -else: - # mock the new API, ignore second argument - # see https://github.com/ipython/ipython/issues/11590 - from ast import Module as OriginalModule - - Module = lambda nodelist, type_ignores: OriginalModule(nodelist) +from ast import Module TOP_FRAME_REGEX = re.compile(r'\s*File "".*in ') @@ -136,9 +125,7 @@ def __init__(self, exc_info): class UnicodeDecodingStringIO(io.StringIO): def write(self, s): - if isinstance(s, bytes): - s = s.decode("utf-8") - super(UnicodeDecodingStringIO, self).write(s) + super().write(s) def clearOutputs(): @@ -428,9 +415,7 @@ def magic_matplot(name): imgdata = io.BytesIO() fig.savefig(imgdata, format="png") imgdata.seek(0) - encode = base64.b64encode(imgdata.getvalue()) - if sys.version >= "3": - encode = encode.decode() + encode = base64.b64encode(imgdata.getvalue()).decode() except: exc_type, exc_value, tb = sys.exc_info() diff --git a/externals/kyuubi-spark-sql-engine/src/main/resources/python/kyuubi_util.py b/externals/kyuubi-spark-sql-engine/src/main/resources/python/kyuubi_util.py index 35ab885111e..8b34ebb0375 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/resources/python/kyuubi_util.py +++ b/externals/kyuubi-spark-sql-engine/src/main/resources/python/kyuubi_util.py @@ -25,7 +25,7 @@ from pyspark.sql import SparkSession -def connect_to_exist_gateway() -> "JavaGateway": +def connect_to_exist_gateway() -> JavaGateway: conn_info_file = os.environ.get("PYTHON_GATEWAY_CONNECTION_INFO") if conn_info_file is None: raise SystemExit("the python gateway connection information file not found!") @@ -45,8 +45,6 @@ def connect_to_exist_gateway() -> "JavaGateway": port=gateway_port, auth_token=gateway_secret, auto_convert=True ) ) - # gateway.proc = proc - # Import the classes used by PySpark java_import(gateway.jvm, "org.apache.spark.SparkConf") java_import(gateway.jvm, "org.apache.spark.api.java.*") @@ -62,7 +60,7 @@ def connect_to_exist_gateway() -> "JavaGateway": return gateway -def get_spark_session(uuid=None) -> "SparkSession": +def get_spark_session(uuid=None) -> SparkSession: gateway = connect_to_exist_gateway() jjsc = gateway.jvm.JavaSparkContext( gateway.jvm.org.apache.spark.SparkContext.getOrCreate() diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala index 8f157698c91..c793d846ee6 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala @@ -390,10 +390,6 @@ object SparkSQLEngine extends Logging { } def main(args: Array[String]): Unit = { - if (KyuubiSparkUtil.SPARK_ENGINE_RUNTIME_VERSION <= "3.4") { - warn("The support for Spark 3.3 and 3.4 is deprecated, and will be removed " + - "in a future version.") - } val startedTime = System.currentTimeMillis() val submitTime = kyuubiConf.getOption(KYUUBI_ENGINE_SUBMIT_TIME_KEY) match { case Some(t) => t.toLong diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/schema/SchemaHelper.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/schema/SchemaHelper.scala index 21702c581e4..6d51745265a 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/schema/SchemaHelper.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/schema/SchemaHelper.scala @@ -27,11 +27,6 @@ import org.apache.kyuubi.shaded.hive.service.rpc.thrift._ object SchemaHelper { - /** - * Spark 3.4.0 DataType TimestampNTZType's class name. - */ - final val TIMESTAMP_NTZ = "TimestampNTZType$" - /** * Spark 4.0.0 DataType VariantType's class name. */ @@ -50,7 +45,7 @@ object SchemaHelper { case _: DecimalType => TTypeId.DECIMAL_TYPE case DateType => TTypeId.DATE_TYPE case TimestampType => TTypeId.TIMESTAMP_TYPE - case ntz if ntz.getClass.getSimpleName.equals(TIMESTAMP_NTZ) => TTypeId.TIMESTAMP_TYPE + case _: TimestampNTZType => TTypeId.TIMESTAMP_TYPE case variant if variant.getClass.getSimpleName.equals(VARIANT) => TTypeId.STRING_TYPE case BinaryType => TTypeId.BINARY_TYPE case CalendarIntervalType => TTypeId.STRING_TYPE @@ -117,7 +112,7 @@ object SchemaHelper { case _: DecimalType => java.sql.Types.DECIMAL case DateType => java.sql.Types.DATE case TimestampType => java.sql.Types.TIMESTAMP - case ntz if ntz.getClass.getSimpleName.equals(TIMESTAMP_NTZ) => java.sql.Types.TIMESTAMP + case _: TimestampNTZType => java.sql.Types.TIMESTAMP case variant if variant.getClass.getSimpleName.equals(VARIANT) => java.sql.Types.OTHER case BinaryType => java.sql.Types.BINARY case _: ArrayType => java.sql.Types.ARRAY @@ -133,7 +128,7 @@ object SchemaHelper { * For array, map, string, and binaries, the column size is variable, return null as unknown. */ def getColumnSize(sparkType: DataType): Option[Int] = sparkType match { - case dt if dt.getClass.getSimpleName == TIMESTAMP_NTZ => + case dt: TimestampNTZType => Some(dt.defaultSize) case dt: DecimalType => Some(dt.precision) @@ -164,7 +159,7 @@ object SchemaHelper { case DoubleType => Some(15) case d: DecimalType => Some(d.scale) case TimestampType => Some(6) - case ntz if ntz.getClass.getSimpleName.equals(TIMESTAMP_NTZ) => Some(6) + case _: TimestampNTZType => Some(6) case _ => None } diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala index 75618e3124b..b8f590ddfbd 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/execution/arrow/KyuubiArrowConverters.scala @@ -18,7 +18,6 @@ package org.apache.spark.sql.execution.arrow import java.io.{ByteArrayInputStream, ByteArrayOutputStream} -import java.lang.{Boolean => JBoolean} import java.nio.channels.Channels import scala.collection.JavaConverters._ @@ -27,7 +26,6 @@ import scala.collection.mutable.ArrayBuffer import org.apache.arrow.vector._ import org.apache.arrow.vector.ipc.{ArrowStreamWriter, ReadChannel, WriteChannel} import org.apache.arrow.vector.ipc.message.{IpcOption, MessageSerializer} -import org.apache.arrow.vector.types.pojo.{Schema => ArrowSchema} import org.apache.spark.TaskContext import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.{InternalRow, SQLConfHelper} @@ -37,8 +35,6 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.util.ArrowUtils import org.apache.spark.util.Utils -import org.apache.kyuubi.util.reflect.DynMethods - object KyuubiArrowConverters extends SQLConfHelper with Logging { type Batch = (Array[Byte], Long) @@ -63,7 +59,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { "slice", 0, Long.MaxValue) - val arrowSchema = toArrowSchema(schema, timeZoneId, true, false) + val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true, false) vectorSchemaRoot = VectorSchemaRoot.create(arrowSchema, sliceAllocator) try { val recordBatch = MessageSerializer.deserializeRecordBatch( @@ -139,7 +135,9 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { while (bufferedRowSize < n && partsScanned < totalParts) { // The number of partitions to try in this iteration. It is ok for this number to be // greater than totalParts because we actually cap it at totalParts in runJob. - var numPartsToTry = limitInitialNumPartitions + // SPARK-40211 (3.4.0) introduced spark.sql.limit.initialNumPartitions + var numPartsToTry = + conf.getConfString("spark.sql.limit.initialNumPartitions", "1").toInt if (partsScanned > 0) { // If we didn't find any rows after the previous iteration, multiply by // limitScaleUpFactor and retry. Otherwise, interpolate the number of partitions we need @@ -192,14 +190,6 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { } } - /** - * Spark introduced the config `spark.sql.limit.initialNumPartitions` since 3.4.0. see SPARK-40211 - */ - private def limitInitialNumPartitions: Int = { - conf.getConfString("spark.sql.limit.initialNumPartitions", "1") - .toInt - } - /** * Different from [[org.apache.spark.sql.execution.arrow.ArrowConverters.toBatchIterator]], * each output arrow batch contains this batch row count. @@ -239,7 +229,7 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { context: TaskContext) extends Iterator[Array[Byte]] { - protected val arrowSchema = toArrowSchema(schema, timeZoneId, true, false) + protected val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true, false) private val allocator = ArrowUtils.rootAllocator.newChildAllocator( s"to${this.getClass.getSimpleName}", @@ -322,34 +312,4 @@ object KyuubiArrowConverters extends SQLConfHelper with Logging { } } - // the signature of function [[ArrowUtils.toArrowSchema]] is changed in SPARK-41971 (since Spark - // 3.5) - private lazy val toArrowSchemaMethod = DynMethods.builder("toArrowSchema") - .impl( // for Spark 3.4 or previous - "org.apache.spark.sql.util.ArrowUtils", - classOf[StructType], - classOf[String]) - .impl( // for Spark 3.5 or later - "org.apache.spark.sql.util.ArrowUtils", - classOf[StructType], - classOf[String], - classOf[Boolean], - classOf[Boolean]) - .build() - - /** - * this function uses reflective calls to the [[ArrowUtils.toArrowSchema]]. - */ - private def toArrowSchema( - schema: StructType, - timeZone: String, - errorOnDuplicatedFieldNames: JBoolean, - largeVarTypes: JBoolean): ArrowSchema = { - toArrowSchemaMethod.invoke[ArrowSchema]( - ArrowUtils, - schema, - timeZone, - errorOnDuplicatedFieldNames, - largeVarTypes) - } } diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala index 7d6e5f8fdb7..c503d780341 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala @@ -48,7 +48,7 @@ object SparkDatasetHelper extends Logging { case adaptiveSparkPlan: AdaptiveSparkPlanExec => executeArrowBatchCollect(adaptiveSparkPlan.finalPhysicalPlan) // TODO: avoid extra shuffle if `offset` > 0 - case collectLimit: CollectLimitExec if offset(collectLimit) > 0 => + case collectLimit: CollectLimitExec if collectLimit.offset > 0 => logWarning("unsupported offset > 0, an extra shuffle will be introduced.") toArrowBatchRdd(collectLimit).collect() case collectLimit: CollectLimitExec if collectLimit.limit >= 0 => @@ -227,20 +227,6 @@ object SparkDatasetHelper extends Logging { spark.sessionState.conf.sessionLocalTimeZone).toArray } - /** - * offset support was add in SPARK-28330(3.4.0), to ensure backward compatibility with - * earlier versions of Spark, this function uses reflective calls to the "offset". - */ - private def offset(collectLimitExec: CollectLimitExec): Int = { - Option( - DynMethods.builder("offset") - .impl(collectLimitExec.getClass) - .orNoop() - .build() - .invoke[Int](collectLimitExec)) - .getOrElse(0) - } - /** * refer to org.apache.spark.sql.Dataset#withAction(), assign a new execution id for arrow-based * operation, so that we can track the arrow-based queries on the UI tab. diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala index 14f900d90f9..8e246b3c6e6 100644 --- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkArrowbasedOperationSuite.scala @@ -261,7 +261,6 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp } test("result offset support") { - assume(SPARK_ENGINE_RUNTIME_VERSION >= "3.4") var numStages = 0 val listener = new SparkListener { override def onJobStart(jobStart: SparkListenerJobStart): Unit = { @@ -528,12 +527,6 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp classOf[String], classOf[Boolean], classOf[TaskContext]) - .hiddenImpl( // for Spark 3.4 or previous - "org.apache.spark.sql.execution.arrow.ArrowConverters$", - classOf[Iterator[Array[Byte]]], - classOf[StructType], - classOf[String], - classOf[TaskContext]) .buildChecked() private lazy val arrowConvertersObject = DynFields.builder() @@ -565,13 +558,6 @@ class SparkArrowbasedOperationSuite extends WithSparkSQLEngine with SparkDataTyp timeZoneId, errorOnDuplicatedFieldNames, context) - }.recover { case _: Exception => // for Spark 3.4 or previous - fromBatchIteratorMethod.invokeChecked[Iterator[InternalRow]]( - arrowConvertersObject, - arrowBatchIter, - schema, - timeZoneId, - context) }.get class JobCountListener extends SparkListener { diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkOperationSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkOperationSuite.scala index feff2d6130e..1f4aa85318b 100644 --- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkOperationSuite.scala +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkOperationSuite.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.types._ import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.engine.spark.WithSparkSQLEngine -import org.apache.kyuubi.engine.spark.schema.SchemaHelper.{TIMESTAMP_NTZ, VARIANT} +import org.apache.kyuubi.engine.spark.schema.SchemaHelper.VARIANT import org.apache.kyuubi.engine.spark.util.SparkCatalogUtils import org.apache.kyuubi.jdbc.hive.KyuubiStatement import org.apache.kyuubi.operation.{HiveMetadataTests, SparkQueryTests} @@ -94,9 +94,7 @@ class SparkOperationSuite extends WithSparkSQLEngine with HiveMetadataTests with .add("c18", "interval day", nullable = true, "18") .add("c19", "interval year", nullable = true, "19") // since spark3.4.0 - if (SPARK_ENGINE_RUNTIME_VERSION >= "3.4") { - schema = schema.add("c20", "timestamp_ntz", nullable = true, "20") - } + schema = schema.add("c20", "timestamp_ntz", nullable = true, "20") if (SPARK_ENGINE_RUNTIME_VERSION >= "4.0") { schema = schema.add("c21", "variant", nullable = true, "21") } @@ -169,7 +167,7 @@ class SparkOperationSuite extends WithSparkSQLEngine with HiveMetadataTests with case FloatType => assert(decimalDigits === 7) case DoubleType => assert(decimalDigits === 15) case TimestampType => assert(decimalDigits === 6) - case ntz if ntz.getClass.getSimpleName.equals(TIMESTAMP_NTZ) => + case _: TimestampNTZType => assert(decimalDigits === 6) case _ => assert(decimalDigits === 0) // nulls } @@ -569,11 +567,7 @@ class SparkOperationSuite extends WithSparkSQLEngine with HiveMetadataTests with val status = tOpenSessionResp.getStatus val errorMessage = status.getErrorMessage assert(status.getStatusCode === TStatusCode.ERROR_STATUS) - if (SPARK_ENGINE_RUNTIME_VERSION >= "3.4") { - assert(errorMessage.contains("[SCHEMA_NOT_FOUND]")) - } else { - assert(errorMessage.contains(s"Database '$dbName' not found")) - } + assert(errorMessage.contains("[SCHEMA_NOT_FOUND]")) } } diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkDataTypeTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkDataTypeTests.scala index 963b8e8953a..1bcbb15bfdf 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkDataTypeTests.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkDataTypeTests.scala @@ -215,7 +215,6 @@ trait SparkDataTypeTests extends HiveJDBCTestHelper with SparkVersionUtil { } test("execute statement - select timestamp_ntz") { - assume(SPARK_ENGINE_RUNTIME_VERSION >= "3.4") withJdbcStatement() { statement => val resultSet = statement.executeQuery( "SELECT make_timestamp_ntz(2022, 03, 24, 18, 08, 31.8888) AS col") diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkQueryTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkQueryTests.scala index f39e55f5b73..cb3a09b3db7 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkQueryTests.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/SparkQueryTests.scala @@ -186,12 +186,8 @@ trait SparkQueryTests extends SparkDataTypeTests with HiveJDBCTestHelper { withJdbcStatement("t") { statement => try { val assertTableOrViewNotfound: (Exception, String) => Unit = (e, tableName) => { - if (SPARK_ENGINE_RUNTIME_VERSION >= "3.4") { - assert(e.getMessage.contains("[TABLE_OR_VIEW_NOT_FOUND]")) - assert(e.getMessage.contains(s"The table or view `$tableName` cannot be found.")) - } else { - assert(e.getMessage.contains("Table or view not found")) - } + assert(e.getMessage.contains("[TABLE_OR_VIEW_NOT_FOUND]")) + assert(e.getMessage.contains(s"The table or view `$tableName` cannot be found.")) } statement.execute("SET kyuubi.operation.plan.only.mode=optimize") val set = statement.executeQuery(ddl) diff --git a/pom.xml b/pom.xml index 5361ddea6ff..a7d02e9f759 100644 --- a/pom.xml +++ b/pom.xml @@ -2016,41 +2016,6 @@ - - spark-3.3 - - extensions/spark/kyuubi-extension-spark-3-3 - extensions/spark/kyuubi-spark-connector-hive - - - 3.3.4 - 3.3 - 2.3.0 - delta-core_${scala.binary.version} - - 1.6.1 - org.scalatest.tags.Slow - - - - - spark-3.4 - - extensions/spark/kyuubi-extension-spark-3-4 - extensions/spark/kyuubi-spark-connector-hive - - - 3.4.4 - 3.4 - 2.4.0 - delta-core_${scala.binary.version} - org.scalatest.tags.Slow - - - spark-3.5