Update GpuIf to support expressions with side effects

integration_tests/src/main/python/conditionals_test.py

@@ -182,3 +182,20 @@ def test_ifnull(data_gen):
                 'ifnull({}, b)'.format(s1),
                 'ifnull({}, b)'.format(null_lit),
                 'ifnull(a, {})'.format(null_lit)))
+
+@pytest.mark.parametrize('data_gen', int_n_long_gens, ids=idfn)
+def test_conditional_with_side_effects_col_col(data_gen):
+    gen = IntegerGen().with_special_case(2147483647)
+    assert_gpu_and_cpu_are_equal_collect(
+            lambda spark : two_col_df(spark, data_gen, gen).selectExpr(
+                'IF(b < 2147483647, b + 1, b)'),
+            conf = {'spark.sql.ansi.enabled':True})
+
+@pytest.mark.parametrize('data_gen', int_n_long_gens, ids=idfn)
+def test_conditional_with_side_effects_col_scalar(data_gen):
+    gen = IntegerGen().with_special_case(2147483647)
+    assert_gpu_and_cpu_are_equal_collect(
+            lambda spark : two_col_df(spark, data_gen, gen).selectExpr(
+                'IF(b < 2147483647, b + 1, 2147483647)',
+                'IF(b >= 2147483646, 2147483647, b + 1)'),
+            conf = {'spark.sql.ansi.enabled':True})

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

@@ -1451,6 +1451,10 @@ case class GpuCast(
 
   import GpuCast._
 
+  // when ansi mode is enabled, some cast expressions can throw exceptions on invalid inputs
+  //TODO not all casts can throw exceptions so this needs to be more specific
+  override def hasSideEffects: Boolean = ansiMode
+
   override def toString: String = if (ansiMode) {
     s"ansi_cast($child as ${dataType.simpleString})"
   } else {

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

@@ -160,6 +160,13 @@ trait GpuExpression extends Expression with Arm {
    */
   def convertToAst(numFirstTableColumns: Int): ast.AstExpression =
     throw new IllegalStateException(s"Cannot convert ${this.getClass.getSimpleName} to AST")
+
+  /** Could evaluating this expression cause side-effects, such as throwing an exception? */
+  def hasSideEffects: Boolean =
+    children.exists {
+      case c: GpuExpression => c.hasSideEffects
+      case _ => false // This path should never really happen
+    }
 }
 
 abstract class GpuLeafExpression extends GpuExpression with ShimExpression {

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

@@ -41,6 +41,8 @@ trait GpuUserDefinedFunction extends GpuExpression
   /** True if the UDF is deterministic */
   val udfDeterministic: Boolean
 
+  override def hasSideEffects: Boolean = true
+
   override lazy val deterministic: Boolean = udfDeterministic && children.forall(_.deterministic)
 
   private[this] val nvtxRangeName = s"UDF: $name"

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

@@ -16,18 +16,19 @@
 
 package com.nvidia.spark.rapids
 
+import ai.rapids.cudf.{ColumnVector, NullPolicy, ScanAggregation, ScanType, Table, UnaryOp}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.shims.v2.ShimExpression
 
 import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
 import org.apache.spark.sql.catalyst.expressions.{ComplexTypeMergingExpression, Expression}
-import org.apache.spark.sql.types.{BooleanType, DataType}
+import org.apache.spark.sql.types.{BooleanType, DataType, DataTypes}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 trait GpuConditionalExpression extends ComplexTypeMergingExpression with GpuExpression
-    with ShimExpression {
+  with ShimExpression {
 
-  protected def computeIfElse(
+  def computeIfElse(
       batch: ColumnarBatch,
       predExpr: Expression,
       trueExpr: Expression,
@@ -54,6 +55,31 @@ trait GpuConditionalExpression extends ComplexTypeMergingExpression with GpuExpr
     }
   }
 
+  protected def isAllTrue(col: GpuColumnVector): Boolean = {
+    assert(BooleanType == col.dataType())
+    if (col.getRowCount == 0) {
+      return true
+    }
+    if (col.hasNull) {
+      return false
+    }
+    withResource(col.getBase.all()) { allTrue =>
+      // Guaranteed there is at least one row and no nulls so result must be valid
+      allTrue.getBoolean
+    }
+  }
+
+  protected def isAllFalse(col: GpuColumnVector): Boolean = {
+    assert(BooleanType == col.dataType())
+    if (col.getRowCount == col.numNulls()) {
+      // all nulls, and null values are false values here
+      return true
+    }
+    withResource(col.getBase.any()) { anyTrue =>
+      // null values are considered false values in this context
+      !anyTrue.isValid || !anyTrue.getBoolean
+    }
+  }
 }
 
 case class GpuIf(
@@ -67,6 +93,7 @@ case class GpuIf(
   }
 
   override def children: Seq[Expression] = predicateExpr :: trueExpr :: falseExpr :: Nil
+
   override def nullable: Boolean = trueExpr.nullable || falseExpr.nullable
 
   override def checkInputDataTypes(): TypeCheckResult = {
@@ -82,8 +109,142 @@ case class GpuIf(
     }
   }
 
-  override def columnarEval(batch: ColumnarBatch): Any = computeIfElse(batch, predicateExpr,
-    trueExpr, falseExpr.columnarEval(batch))
+  override def columnarEval(batch: ColumnarBatch): Any = {
+
+    val gpuTrueExpr = trueExpr.asInstanceOf[GpuExpression]
+    val gpuFalseExpr = falseExpr.asInstanceOf[GpuExpression]
+    val colTypes = GpuColumnVector.extractTypes(batch)
+
+    withResource(GpuExpressionsUtils.columnarEvalToColumn(predicateExpr, batch)) { pred =>
+      if (isAllTrue(pred)) {
+        trueExpr.columnarEval(batch)
+      } else if (isAllFalse(pred)) {
+        falseExpr.columnarEval(batch)
+      } else if (gpuTrueExpr.hasSideEffects || gpuFalseExpr.hasSideEffects) {
+        withResource(GpuColumnVector.from(batch)) { tbl =>
+          withResource(pred.getBase.unaryOp(UnaryOp.NOT)) { inverted =>
+            withResource(filterBatch(tbl, pred.getBase, colTypes)) { trueBatch =>
+              withResource(filterBatch(tbl, inverted, colTypes)) { falseBatch =>
+                withResourceIfAllowed(gpuTrueExpr.columnarEval(trueBatch)) { tt =>
+                  withResourceIfAllowed(gpuFalseExpr.columnarEval(falseBatch)) { ff =>
+                    val finalRet = (tt, ff) match {
+                      case (t: GpuColumnVector, f: GpuColumnVector) =>
+                        withResource(GpuColumnVector.from(new Table(t.getBase),
+                            Array(trueExpr.dataType))) { trueTable =>
+                          withResource(GpuColumnVector.from(new Table(f.getBase),
+                              Array(falseExpr.dataType))) { falseTable =>
+                            withResource(gather(pred.getBase, trueTable)) { trueValues =>
+                              withResource(gather(inverted, falseTable)) { falseValues =>
+                                pred.getBase.ifElse(
+                                  trueValues.getColumn(0),
+                                  falseValues.getColumn(0))
+                              }
+                            }
+                          }
+                        }
+                      case (t: GpuScalar, f: GpuColumnVector) =>
+                        withResource(GpuColumnVector.from(new Table(f.getBase),
+                            Array(falseExpr.dataType))) { falseTable =>
+                          withResource(gather(inverted, falseTable)) { falseValues =>
+                            pred.getBase.ifElse(
+                              t.getBase,
+                              falseValues.getColumn(0))
+                          }
+                        }
+                      case (t: GpuColumnVector, f: GpuScalar) =>
+                        withResource(GpuColumnVector.from(new Table(t.getBase),
+                            Array(trueExpr.dataType))) { trueTable =>
+                          withResource(gather(pred.getBase, trueTable)) { trueValues =>
+                            pred.getBase.ifElse(
+                              trueValues.getColumn(0),
+                              f.getBase)
+                          }
+                        }
+                      case (_: GpuScalar, _: GpuScalar) =>
+                        throw new IllegalStateException(
+                          "scalar expressions can never have side effects")
+                    }
+                    GpuColumnVector.from(finalRet, dataType)
+                  }
+                }
+              }
+            }
+          }
+        }
+      } else {
+        // simple approach (original GpuIf code)
+        withResourceIfAllowed(trueExpr.columnarEval(batch)) { trueRet =>
+          withResourceIfAllowed(falseExpr.columnarEval(batch)) { falseRet =>
+            val finalRet = (trueRet, falseRet) match {
+              case (t: GpuColumnVector, f: GpuColumnVector) =>
+                pred.getBase.ifElse(t.getBase, f.getBase)
+              case (t: GpuScalar, f: GpuColumnVector) =>
+                pred.getBase.ifElse(t.getBase, f.getBase)
+              case (t: GpuColumnVector, f: GpuScalar) =>
+                pred.getBase.ifElse(t.getBase, f.getBase)
+              case (t: GpuScalar, f: GpuScalar) =>
+                pred.getBase.ifElse(t.getBase, f.getBase)
+              case (t, f) =>
+                throw new IllegalStateException(s"Unexpected inputs" +
+                  s" ($t: ${t.getClass}, $f: ${f.getClass})")
+            }
+            GpuColumnVector.from(finalRet, dataType)
+          }
+        }
+      }
+    }
+  }
+
+  private def filterBatch(
+      tbl: Table,
+      pred: ColumnVector,
+      colTypes: Array[DataType]): ColumnarBatch = {
+    withResource(tbl.filter(pred)) { filteredData =>
+      GpuColumnVector.from(filteredData, colTypes)
+    }
+  }
+
+  private def boolToInt(cv: ColumnVector): ColumnVector = {
+    withResource(GpuScalar.from(1, DataTypes.IntegerType)) { one =>
+      withResource(GpuScalar.from(0, DataTypes.IntegerType)) { zero =>
+        cv.ifElse(one, zero)
+      }
+    }
+  }
+
+  private def gather(predicate: ColumnVector, batch: ColumnarBatch): Table = {
+    // convert the predicate boolean column to numeric where 1 = true
+    // amd 0 = false and then use `scan` with `sum` to convert to
+    // indices.
+    //
+    // For example, if the predicate evaluates to [F, F, T, F, T] then this
+    // gets translated first to [0, 0, 1, 0, 1] and then the scan operation
+    // will perform an exclusive sum on these values and
+    // produce [0, 0, 0, 1, 1]. Combining this with the original
+    // predicate boolean array results in the two T values mapping to
+    // indices 0 and 1, respectively.
+
+    withResource(boolToInt(predicate)) { boolsAsInts =>
+      withResource(boolsAsInts.scan(
+        ScanAggregation.sum(),
+        ScanType.EXCLUSIVE,
+        NullPolicy.INCLUDE)) { prefixSumExclusive =>
+
+        // for the entries in the gather map that do not represent valid
+        // values to be gathered, we change the value to -MAX_INT which
+        // will be treated as null values in the gather algorithm
+        val gatherMap = withResource(GpuScalar.from(Int.MinValue, DataTypes.IntegerType)) {
+          outOfBoundsFlag => predicate.ifElse(prefixSumExclusive, outOfBoundsFlag)
+        }
+
+        withResource(gatherMap) { _ =>
+          withResource(GpuColumnVector.from(batch)) { table =>
+            table.gather(gatherMap)
+          }
+        }
+      }
+    }
+  }
 
   override def toString: String = s"if ($predicateExpr) $trueExpr else $falseExpr"
 

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

@@ -142,6 +142,8 @@ case class GpuAbs(child: Expression, failOnError: Boolean) extends CudfUnaryExpr
 
 abstract class CudfBinaryArithmetic extends CudfBinaryOperator with NullIntolerant {
   override def dataType: DataType = left.dataType
+  // arithmetic operations can overflow and throw exceptions in ANSI mode
+  override def hasSideEffects: Boolean = SQLConf.get.ansiEnabled
 }
 
 object GpuAdd extends Arm {

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

@@ -171,6 +171,8 @@ case class GpuCeil(child: Expression) extends CudfUnaryMathExpression("CEIL") {
     case _ => LongType
   }
 
+  override def hasSideEffects: Boolean = true
+
   override def inputTypes: Seq[AbstractDataType] =
     Seq(TypeCollection(DoubleType, DecimalType, LongType))
 
@@ -245,6 +247,8 @@ case class GpuFloor(child: Expression) extends CudfUnaryMathExpression("FLOOR")
     case _ => LongType
   }
 
+  override def hasSideEffects: Boolean = true
+
   override def inputTypes: Seq[AbstractDataType] =
     Seq(TypeCollection(DoubleType, DecimalType, LongType))