From 036705c58f160b2a6ad3c1a7497de5f5e8eb0951 Mon Sep 17 00:00:00 2001
From: Raza Jafri <rjafri@nvidia.com>
Date: Fri, 3 Jun 2022 11:34:06 -0700
Subject: [PATCH 1/4] Write Scala test for ORC encrypted write

Signed-off-by: Raza Jafri <rjafri@nvidia.com>
---
 .../src/main/python/orc_write_test.py         | 27 -------
 tests/pom.xml                                 | 76 +++++++++++++++++++
 .../spark/rapids/OrcEncryptionSuite.scala     | 48 ++++++++++++
 3 files changed, 124 insertions(+), 27 deletions(-)
 create mode 100644 tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala

diff --git a/integration_tests/src/main/python/orc_write_test.py b/integration_tests/src/main/python/orc_write_test.py
index 2d58d198903..7afe00076a3 100644
--- a/integration_tests/src/main/python/orc_write_test.py
+++ b/integration_tests/src/main/python/orc_write_test.py
@@ -167,30 +167,3 @@ def create_empty_df(spark, path):
         lambda spark, path: spark.read.orc(path),
         data_path,
         conf={'spark.rapids.sql.format.orc.write.enabled': True})
-
-@allow_non_gpu('DataWritingCommandExec')
-@pytest.mark.parametrize("path", ["", "kms://http@localhost:9600/kms"])
-@pytest.mark.parametrize("provider", ["", "hadoop"])
-@pytest.mark.parametrize("encrypt", ["", "pii:a"])
-@pytest.mark.parametrize("mask", ["", "sha256:a"])
-@pytest.mark.skipif(is_databricks104_or_later(), reason="The test will fail on Databricks10.4 because `HadoopShimsPre2_3$NullKeyProvider` is loaded")
-def test_orc_write_encryption_fallback(spark_tmp_path, spark_tmp_table_factory, path, provider, encrypt, mask):
-    def write_func(spark, write_path):
-        writer = unary_op_df(spark, gen).coalesce(1).write
-        if path != "":
-            writer.option("hadoop.security.key.provider.path", path)
-        if provider != "":
-            writer.option("orc.key.provider", provider)
-        if encrypt != "":
-            writer.option("orc.encrypt", encrypt)
-        if mask != "":
-            writer.option("orc.mask", mask)
-        writer.format("orc").mode('overwrite').option("path", write_path).saveAsTable(spark_tmp_table_factory.get())
-    if path == "" and provider == "" and encrypt == "" and mask == "":
-        pytest.skip("Skip this test when none of the encryption confs are set")
-    gen = IntegerGen()
-    data_path = spark_tmp_path + '/ORC_DATA'
-    assert_gpu_fallback_write(write_func,
-        lambda spark, path: spark.read.orc(path),
-        data_path,
-        'DataWritingCommandExec')
diff --git a/tests/pom.xml b/tests/pom.xml
index 4026d6865d7..e19039dc61e 100644
--- a/tests/pom.xml
+++ b/tests/pom.xml
@@ -148,6 +148,25 @@
                     </exclusions>
                 </dependency>
             </dependencies>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.codehaus.mojo</groupId>
+                        <artifactId>build-helper-maven-plugin</artifactId>
+                        <executions>
+                            <execution>
+                                <id>add-321cdh-test-src</id>
+                                <goals><goal>add-test-source</goal></goals>
+                                <configuration>
+                                    <sources>
+                                        <source>${project.basedir}/src/test/320+/scala</source>
+                                    </sources>
+                                </configuration>
+                            </execution>
+                        </executions>
+                    </plugin>
+                </plugins>
+            </build>
         </profile>
         <profile>
             <!--
@@ -273,6 +292,62 @@
                 </dependency>
             </dependencies>
         </profile>
+        <profile>
+            <id>release321</id>
+            <activation>
+                <property>
+                    <name>buildver</name>
+                    <value>321</value>
+                </property>
+            </activation>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.codehaus.mojo</groupId>
+                        <artifactId>build-helper-maven-plugin</artifactId>
+                        <executions>
+                            <execution>
+                                <id>add-321-test-src</id>
+                                <goals><goal>add-test-source</goal></goals>
+                                <configuration>
+                                    <sources>
+                                        <source>${project.basedir}/src/test/320+/scala</source>
+                                    </sources>
+                                </configuration>
+                            </execution>
+                        </executions>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
+        <profile>
+            <id>release322</id>
+            <activation>
+                <property>
+                    <name>buildver</name>
+                    <value>322</value>
+                </property>
+            </activation>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.codehaus.mojo</groupId>
+                        <artifactId>build-helper-maven-plugin</artifactId>
+                        <executions>
+                            <execution>
+                                <id>add-322-test-src</id>
+                                <goals><goal>add-test-source</goal></goals>
+                                <configuration>
+                                    <sources>
+                                        <source>${project.basedir}/src/test/320+/scala</source>
+                                    </sources>
+                                </configuration>
+                            </execution>
+                        </executions>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
         <profile>
             <id>release330</id>
             <activation>
@@ -293,6 +368,7 @@
                                 <configuration>
                                     <sources>
                                         <source>${project.basedir}/src/test/330/scala</source>
+                                        <source>${project.basedir}/src/test/320+/scala</source>
                                     </sources>
                                 </configuration>
                             </execution>
diff --git a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
new file mode 100644
index 00000000000..f87ff10bf4c
--- /dev/null
+++ b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids
+
+import java.io.File
+import java.util.Random
+
+import org.apache.orc.impl.HadoopShimsFactory
+
+import org.apache.spark.SparkConf
+
+class OrcEncryptionSuite extends SparkQueryCompareTestSuite {
+  testGpuFallback(
+    "Write encrypted ORC fallback",
+    "DataWritingCommandExec",
+    intsDf,
+    conf = new SparkConf().set("hadoop.security.key.provider.path", "test:///")
+        .set("orc.key.provider", "hadoop")
+        .set("orc.encrypt", "pii:ints,more_ints")
+        .set("orc.mask", "sha256:ints,more_ints")) {
+    frame =>
+      withCpuSparkSession(session => {
+        val conf = session.sessionState.newHadoopConf()
+        val provider = HadoopShimsFactory.get.getHadoopKeyProvider(conf, new Random)
+        assume(!provider.getKeyNames.isEmpty,
+          s"$provider doesn't has the test keys. ORC shim is created with old Hadoop libraries")
+      }
+      )
+
+      val tempFile = File.createTempFile("orc-encryption-test", "")
+      frame.write.mode("overwrite").orc(tempFile.getAbsolutePath)
+      frame.selectExpr("*")
+  }
+}

From 8686bc4c16ce761b0d88b04e5d59a7fa91e693e6 Mon Sep 17 00:00:00 2001
From: Raza Jafri <rjafri@nvidia.com>
Date: Fri, 3 Jun 2022 12:27:18 -0700
Subject: [PATCH 2/4] add a defensive piece of code to make sure we don't run
 this code for 3.1.x

Signed-off-by: Raza Jafri <rjafri@nvidia.com>
---
 .../com/nvidia/spark/rapids/OrcEncryptionSuite.scala   | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
index f87ff10bf4c..42f6e442cb2 100644
--- a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
+++ b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
@@ -19,6 +19,7 @@ package com.nvidia.spark.rapids
 import java.io.File
 import java.util.Random
 
+import com.nvidia.spark.rapids.shims.SparkShimImpl
 import org.apache.orc.impl.HadoopShimsFactory
 
 import org.apache.spark.SparkConf
@@ -33,6 +34,15 @@ class OrcEncryptionSuite extends SparkQueryCompareTestSuite {
         .set("orc.encrypt", "pii:ints,more_ints")
         .set("orc.mask", "sha256:ints,more_ints")) {
     frame =>
+      // ORC encryption is only allowed in 3.2+
+      val isValidTestForSparkVersion = SparkShimImpl.getSparkShimVersion match {
+        case SparkShimVersion(major, minor, _) => major == 3 && minor != 1
+        case DatabricksShimVersion(major, minor, _, _) => major == 3 && minor != 1
+        case ClouderaShimVersion(major, minor, _, _) => major == 3 && minor != 1
+        case _ => true
+      }
+      assume(isValidTestForSparkVersion)
+
       withCpuSparkSession(session => {
         val conf = session.sessionState.newHadoopConf()
         val provider = HadoopShimsFactory.get.getHadoopKeyProvider(conf, new Random)

From a899b8bf396499c01815f364c33decb67782b94e Mon Sep 17 00:00:00 2001
From: Raza Jafri <rjafri@nvidia.com>
Date: Fri, 3 Jun 2022 16:52:42 -0700
Subject: [PATCH 3/4] park

Signed-off-by: Raza Jafri <rjafri@nvidia.com>
---
 .../nvidia/spark/rapids/OrcEncryptionSuite.scala | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
index 42f6e442cb2..8475cd2765f 100644
--- a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
+++ b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
@@ -17,10 +17,8 @@
 package com.nvidia.spark.rapids
 
 import java.io.File
-import java.util.Random
 
 import com.nvidia.spark.rapids.shims.SparkShimImpl
-import org.apache.orc.impl.HadoopShimsFactory
 
 import org.apache.spark.SparkConf
 
@@ -43,13 +41,13 @@ class OrcEncryptionSuite extends SparkQueryCompareTestSuite {
       }
       assume(isValidTestForSparkVersion)
 
-      withCpuSparkSession(session => {
-        val conf = session.sessionState.newHadoopConf()
-        val provider = HadoopShimsFactory.get.getHadoopKeyProvider(conf, new Random)
-        assume(!provider.getKeyNames.isEmpty,
-          s"$provider doesn't has the test keys. ORC shim is created with old Hadoop libraries")
-      }
-      )
+//      withCpuSparkSession(session => {
+//        val conf = session.sessionState.newHadoopConf()
+//        val provider = HadoopShimsFactory.get.getHadoopKeyProvider(conf, new Random)
+//        assume(!provider.getKeyNames.isEmpty,
+//          s"$provider doesn't has the test keys. ORC shim is created with old Hadoop libraries")
+//      }
+//      )
 
       val tempFile = File.createTempFile("orc-encryption-test", "")
       frame.write.mode("overwrite").orc(tempFile.getAbsolutePath)

From 8cac906f271a506cea24845cf3f5c4e29b51b194 Mon Sep 17 00:00:00 2001
From: Raza Jafri <rjafri@nvidia.com>
Date: Mon, 6 Jun 2022 13:26:45 -0700
Subject: [PATCH 4/4] remove the redundant test

Signed-off-by: Raza Jafri <rjafri@nvidia.com>
---
 tests/pom.xml                                 | 76 -------------------
 .../spark/rapids/OrcEncryptionSuite.scala     | 56 --------------
 2 files changed, 132 deletions(-)
 delete mode 100644 tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala

diff --git a/tests/pom.xml b/tests/pom.xml
index e19039dc61e..4026d6865d7 100644
--- a/tests/pom.xml
+++ b/tests/pom.xml
@@ -148,25 +148,6 @@
                     </exclusions>
                 </dependency>
             </dependencies>
-            <build>
-                <plugins>
-                    <plugin>
-                        <groupId>org.codehaus.mojo</groupId>
-                        <artifactId>build-helper-maven-plugin</artifactId>
-                        <executions>
-                            <execution>
-                                <id>add-321cdh-test-src</id>
-                                <goals><goal>add-test-source</goal></goals>
-                                <configuration>
-                                    <sources>
-                                        <source>${project.basedir}/src/test/320+/scala</source>
-                                    </sources>
-                                </configuration>
-                            </execution>
-                        </executions>
-                    </plugin>
-                </plugins>
-            </build>
         </profile>
         <profile>
             <!--
@@ -292,62 +273,6 @@
                 </dependency>
             </dependencies>
         </profile>
-        <profile>
-            <id>release321</id>
-            <activation>
-                <property>
-                    <name>buildver</name>
-                    <value>321</value>
-                </property>
-            </activation>
-            <build>
-                <plugins>
-                    <plugin>
-                        <groupId>org.codehaus.mojo</groupId>
-                        <artifactId>build-helper-maven-plugin</artifactId>
-                        <executions>
-                            <execution>
-                                <id>add-321-test-src</id>
-                                <goals><goal>add-test-source</goal></goals>
-                                <configuration>
-                                    <sources>
-                                        <source>${project.basedir}/src/test/320+/scala</source>
-                                    </sources>
-                                </configuration>
-                            </execution>
-                        </executions>
-                    </plugin>
-                </plugins>
-            </build>
-        </profile>
-        <profile>
-            <id>release322</id>
-            <activation>
-                <property>
-                    <name>buildver</name>
-                    <value>322</value>
-                </property>
-            </activation>
-            <build>
-                <plugins>
-                    <plugin>
-                        <groupId>org.codehaus.mojo</groupId>
-                        <artifactId>build-helper-maven-plugin</artifactId>
-                        <executions>
-                            <execution>
-                                <id>add-322-test-src</id>
-                                <goals><goal>add-test-source</goal></goals>
-                                <configuration>
-                                    <sources>
-                                        <source>${project.basedir}/src/test/320+/scala</source>
-                                    </sources>
-                                </configuration>
-                            </execution>
-                        </executions>
-                    </plugin>
-                </plugins>
-            </build>
-        </profile>
         <profile>
             <id>release330</id>
             <activation>
@@ -368,7 +293,6 @@
                                 <configuration>
                                     <sources>
                                         <source>${project.basedir}/src/test/330/scala</source>
-                                        <source>${project.basedir}/src/test/320+/scala</source>
                                     </sources>
                                 </configuration>
                             </execution>
diff --git a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala b/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
deleted file mode 100644
index 8475cd2765f..00000000000
--- a/tests/src/test/320+/scala/com/nvidia/spark/rapids/OrcEncryptionSuite.scala
+++ /dev/null
@@ -1,56 +0,0 @@
-/*
- * Copyright (c) 2022, NVIDIA CORPORATION.
- *
- * Licensed 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 com.nvidia.spark.rapids
-
-import java.io.File
-
-import com.nvidia.spark.rapids.shims.SparkShimImpl
-
-import org.apache.spark.SparkConf
-
-class OrcEncryptionSuite extends SparkQueryCompareTestSuite {
-  testGpuFallback(
-    "Write encrypted ORC fallback",
-    "DataWritingCommandExec",
-    intsDf,
-    conf = new SparkConf().set("hadoop.security.key.provider.path", "test:///")
-        .set("orc.key.provider", "hadoop")
-        .set("orc.encrypt", "pii:ints,more_ints")
-        .set("orc.mask", "sha256:ints,more_ints")) {
-    frame =>
-      // ORC encryption is only allowed in 3.2+
-      val isValidTestForSparkVersion = SparkShimImpl.getSparkShimVersion match {
-        case SparkShimVersion(major, minor, _) => major == 3 && minor != 1
-        case DatabricksShimVersion(major, minor, _, _) => major == 3 && minor != 1
-        case ClouderaShimVersion(major, minor, _, _) => major == 3 && minor != 1
-        case _ => true
-      }
-      assume(isValidTestForSparkVersion)
-
-//      withCpuSparkSession(session => {
-//        val conf = session.sessionState.newHadoopConf()
-//        val provider = HadoopShimsFactory.get.getHadoopKeyProvider(conf, new Random)
-//        assume(!provider.getKeyNames.isEmpty,
-//          s"$provider doesn't has the test keys. ORC shim is created with old Hadoop libraries")
-//      }
-//      )
-
-      val tempFile = File.createTempFile("orc-encryption-test", "")
-      frame.write.mode("overwrite").orc(tempFile.getAbsolutePath)
-      frame.selectExpr("*")
-  }
-}