Fix tests failing on Spark 3.1 (#557)

Signed-off-by: Jason Lowe <[email protected]>
NVIDIA · Aug 14, 2020 · 48e7ed7 · 48e7ed7
1 parent 20afca1
commit 48e7ed7
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuCoalesceBatchesSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuCoalesceBatchesSuite.scala
@@ -222,9 +222,10 @@ class GpuCoalesceBatchesSuite extends SparkQueryCompareTestSuite {
       val df = longsCsvDf(spark)
 
       // A coalesce step is added after the filter to help with the case where much of the
-      // data is filtered out
+      // data is filtered out.  The select is there to prevent the coalesce from being
+      // the last thing in the plan which will cause the coalesce to be optimized out.
       val df2 = df
-        .filter(df.col("six").gt(5))
+        .filter(df.col("six").gt(5)).select(df.col("six") * 2)
 
       val coalesce = df2.queryExecution.executedPlan
         .find(_.isInstanceOf[GpuCoalesceBatches]).get

diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala
@@ -18,6 +18,8 @@ package com.nvidia.spark.rapids
 
 import java.sql.Timestamp
 
+import org.apache.spark
+
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.{AnalysisException, DataFrame, SparkSession}
 import org.apache.spark.sql.execution.{SparkPlan, WholeStageCodegenExec}
@@ -1566,14 +1568,19 @@ class HashAggregatesSuite extends SparkQueryCompareTestSuite {
       res
   }
 
-  // A test that verifies that Distinct with Filter is not supported on the CPU or the GPU.
-  testExpectedExceptionStartsWith(
-      "Avg Distinct with filter - unsupported on CPU and GPU",
-      classOf[AnalysisException],
-      "DISTINCT and FILTER cannot be used in aggregate functions at the same time",
-      longsFromCSVDf, conf = floatAggConf) {
-    frame => val res = frame.selectExpr("avg(distinct longs) filter (where longs < 5)")
-      res
+  if (spark.SPARK_VERSION_SHORT < "3.1.0") {
+    // A test that verifies that Distinct with Filter is not supported on the CPU or the GPU.
+    testExpectedExceptionStartsWith(
+        "Avg Distinct with filter - unsupported on CPU and GPU",
+        classOf[AnalysisException],
+        "DISTINCT and FILTER cannot be used in aggregate functions at the same time",
+        longsFromCSVDf, conf = floatAggConf) {
+      frame => frame.selectExpr("avg(distinct longs) filter (where longs < 5)")
+    }
+  } else {
+    testSparkResultsAreEqual("Avg Distinct with filter", longsFromCSVDf, conf = floatAggConf) {
+      frame => frame.selectExpr("avg(distinct longs) filter (where longs < 5)")
+    }
   }
 
   testSparkResultsAreEqualWithCapture("PartMerge:avg_overflow_cast_dbl", veryLargeLongsFromCSVDf,