Switch and/or to use new cudf binops to improve performance

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -271,27 +271,11 @@ trait GpuBinaryExpression extends ShimBinaryExpression with GpuExpression {
 
 trait GpuBinaryOperator extends BinaryOperator with GpuBinaryExpression
 
-object CudfBinaryExpression {
-  lazy val opToAstMap: Map[BinaryOp, ast.BinaryOperator] = Map(
-    BinaryOp.ADD -> ast.BinaryOperator.ADD,
-    BinaryOp.BITWISE_AND -> ast.BinaryOperator.BITWISE_AND,
-    BinaryOp.BITWISE_OR -> ast.BinaryOperator.BITWISE_OR,
-    BinaryOp.BITWISE_XOR -> ast.BinaryOperator.BITWISE_XOR,
-    BinaryOp.GREATER -> ast.BinaryOperator.GREATER,
-    BinaryOp.GREATER_EQUAL -> ast.BinaryOperator.GREATER_EQUAL,
-    BinaryOp.LESS -> ast.BinaryOperator.LESS,
-    BinaryOp.LESS_EQUAL -> ast.BinaryOperator.LESS_EQUAL,
-    BinaryOp.LOGICAL_AND -> ast.BinaryOperator.NULL_LOGICAL_AND,
-    BinaryOp.LOGICAL_OR -> ast.BinaryOperator.NULL_LOGICAL_OR,
-    BinaryOp.MUL -> ast.BinaryOperator.MUL,
-    BinaryOp.POW -> ast.BinaryOperator.POW,
-    BinaryOp.SUB -> ast.BinaryOperator.SUB)
-}
-
 trait CudfBinaryExpression extends GpuBinaryExpression {
   def binaryOp: BinaryOp
   def outputTypeOverride: DType = null
   def castOutputAtEnd: Boolean = false
+  def astOperator: Option[ast.BinaryOperator] = None
 
   def outputType(l: BinaryOperable, r: BinaryOperable): DType = {
     val over = outputTypeOverride
@@ -339,7 +323,7 @@ trait CudfBinaryExpression extends GpuBinaryExpression {
   }
 
   override def convertToAst(numFirstTableColumns: Int): ast.AstExpression = {
-    val astOp = CudfBinaryExpression.opToAstMap.getOrElse(binaryOp,
+    val astOp = astOperator.getOrElse(
       throw new IllegalStateException(s"$this is not supported by AST"))
     assert(left.dataType == right.dataType)
     new ast.BinaryOperation(astOp,

sql-plugin/src/main/scala/com/nvidia/spark/rapids/conditionalExpressions.scala

@@ -16,8 +16,7 @@
 
 package com.nvidia.spark.rapids
 
-import ai.rapids.cudf.{ColumnVector, NullPolicy, Scalar, ScanAggregation, ScanType, Table, UnaryOp}
-import ai.rapids.cudf.ast
+import ai.rapids.cudf.{BinaryOp, ColumnVector, DType, NullPolicy, Scalar, ScanAggregation, ScanType, Table, UnaryOp}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.shims.v2.ShimExpression
 
@@ -441,15 +440,8 @@ case class GpuCaseWhen(
     cumulativePred match {
       case Some(prev) =>
         withResource(prev) { _ =>
-          val result = withResource(new Table(prev.getBase, whenBool.getBase)) { t =>
-            val or = new ast.BinaryOperation(ast.BinaryOperator.NULL_LOGICAL_OR,
-              new ast.ColumnReference(0),
-              new ast.ColumnReference(1)
-            )
-            withResource(or.compile()) {
-              _.computeColumn(t)
-            }
-          }
+          val result = prev.getBase.binaryOp(BinaryOp.NULL_LOGICAL_OR,
+            whenBool.getBase, DType.BOOL8)
           GpuColumnVector.from(result, DataTypes.BooleanType)
         }
       case _ =>

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.rapids
 import java.math.BigInteger
 
 import ai.rapids.cudf._
+import ai.rapids.cudf.ast.BinaryOperator
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.shims.v2.ShimExpression
@@ -219,6 +220,7 @@ case class GpuAdd(
   override def symbol: String = "+"
 
   override def binaryOp: BinaryOp = BinaryOp.ADD
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.ADD)
 
   override def doColumnar(lhs: BinaryOperable, rhs: BinaryOperable): ColumnVector = {
     val ret = super.doColumnar(lhs, rhs)
@@ -246,6 +248,7 @@ case class GpuSubtract(
   override def symbol: String = "-"
 
   override def binaryOp: BinaryOp = BinaryOp.SUB
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.SUB)
 
   private[this] def basicOpOverflowCheck(
       lhs: BinaryOperable,
@@ -534,6 +537,7 @@ case class GpuMultiply(
   override def symbol: String = "*"
 
   override def binaryOp: BinaryOp = BinaryOp.MUL
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.MUL)
 }
 
 object GpuDivModLike extends Arm {

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/bitwise.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-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.
@@ -16,7 +16,8 @@
 
 package org.apache.spark.sql.rapids
 
-import ai.rapids.cudf.{BinaryOp, ColumnVector, DType, Scalar, UnaryOp}
+import ai.rapids.cudf.{ast, BinaryOp, ColumnVector, DType, Scalar, UnaryOp}
+import ai.rapids.cudf.ast.BinaryOperator
 import com.nvidia.spark.rapids._
 
 import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes, Expression, ImplicitCastInputTypes}
@@ -120,6 +121,7 @@ case class GpuBitwiseAnd(left: Expression, right: Expression) extends CudfBinary
   override def symbol: String = "&"
 
   override def binaryOp: BinaryOp = BinaryOp.BITWISE_AND
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.BITWISE_AND)
 }
 
 case class GpuBitwiseOr(left: Expression, right: Expression) extends CudfBinaryArithmetic {
@@ -128,6 +130,7 @@ case class GpuBitwiseOr(left: Expression, right: Expression) extends CudfBinaryA
   override def symbol: String = "|"
 
   override def binaryOp: BinaryOp = BinaryOp.BITWISE_OR
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.BITWISE_OR)
 }
 
 case class GpuBitwiseXor(left: Expression, right: Expression) extends CudfBinaryArithmetic {
@@ -136,6 +139,7 @@ case class GpuBitwiseXor(left: Expression, right: Expression) extends CudfBinary
   override def symbol: String = "^"
 
   override def binaryOp: BinaryOp = BinaryOp.BITWISE_XOR
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.BITWISE_XOR)
 }
 
 case class GpuBitwiseNot(child: Expression) extends CudfUnaryExpression with ExpectsInputTypes {

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/mathExpressions.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -19,6 +19,7 @@ package org.apache.spark.sql.rapids
 import java.io.Serializable
 
 import ai.rapids.cudf._
+import ai.rapids.cudf.ast.BinaryOperator
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.RapidsPluginImplicits.ReallyAGpuExpression
 
@@ -650,6 +651,7 @@ case class GpuRound(child: Expression, scale: Expression) extends
 case class GpuPow(left: Expression, right: Expression)
     extends CudfBinaryMathExpression("POWER") {
   override def binaryOp: BinaryOp = BinaryOp.POW
+  override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.POW)
   override def outputTypeOverride: DType = DType.FLOAT64
 }