Fix correctness issue with CASE WHEN with expressions that have side-effects

integration_tests/src/main/python/conditionals_test.py

@@ -186,28 +186,36 @@ def test_ifnull(data_gen):
                 'ifnull({}, b)'.format(null_lit),
                 'ifnull(a, {})'.format(null_lit)))
 
-@pytest.mark.parametrize('data_gen', int_n_long_gens, ids=idfn)
+@pytest.mark.parametrize('data_gen', [IntegerGen().with_special_case(2147483647)], 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)'),
+            lambda spark : unary_op_df(spark, data_gen).selectExpr(
+                'IF(a < 2147483647, a + 1, a)'),
             conf = {'spark.sql.ansi.enabled':True})
 
-@pytest.mark.parametrize('data_gen', int_n_long_gens, ids=idfn)
+@pytest.mark.parametrize('data_gen', [IntegerGen().with_special_case(2147483647)], 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)'),
+            lambda spark : unary_op_df(spark, data_gen).selectExpr(
+                'IF(a < 2147483647, a + 1, 2147483647)',
+                'IF(a >= 2147483646, 2147483647, a + 1)'),
             conf = {'spark.sql.ansi.enabled':True})
 
-@pytest.mark.parametrize('data_gen', int_n_long_gens, ids=idfn)
+@pytest.mark.parametrize('data_gen', [mk_str_gen('[0-9]{1,20}')], ids=idfn)
 def test_conditional_with_side_effects_cast(data_gen):
-    gen = mk_str_gen('[0-9]{1,20}')
     assert_gpu_and_cpu_are_equal_collect(
-            lambda spark : two_col_df(spark, data_gen, gen).selectExpr(
+            lambda spark : unary_op_df(spark, data_gen).selectExpr(
                 'IF(a RLIKE "^[0-9]{1,5}$", CAST(a AS INT), 0)'),
             conf = {'spark.sql.ansi.enabled':True,
-                    'spark.rapids.sql.expression.RLike': True})
+                    'spark.rapids.sql.expression.RLike': True})
+
+@pytest.mark.parametrize('data_gen', [mk_str_gen('[0-9]{1,9}')], ids=idfn)
+def test_conditional_with_side_effects_case_when(data_gen):
+    assert_gpu_and_cpu_are_equal_collect(
+            lambda spark : unary_op_df(spark, data_gen).selectExpr(
+                'CASE \
+                WHEN a RLIKE "^[0-9]{1,3}$" THEN CAST(a AS INT) \
+                WHEN a RLIKE "^[0-9]{4,6}$" THEN CAST(a AS INT) + 123 \
+                ELSE -1 END'),
+                conf = {'spark.sql.ansi.enabled':True,
+                        'spark.rapids.sql.expression.RLike': True})

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

@@ -80,6 +80,61 @@ trait GpuConditionalExpression extends ComplexTypeMergingExpression with GpuExpr
       !anyTrue.getBoolean
     }
   }
+
+  protected 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)
+      }
+    }
+  }
+
+  protected def gather(predicate: ColumnVector, t: GpuColumnVector): ColumnVector = {
+    // 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(Scalar.fromInt(Int.MinValue)) {
+          outOfBoundsFlag => predicate.ifElse(prefixSumExclusive, outOfBoundsFlag)
+        }
+
+        withResource(new Table(t.getBase)) { tbl =>
+          withResource(gatherMap) { _ =>
+            withResource(tbl.gather(gatherMap)) { gatherTbl =>
+              gatherTbl.getColumn(0).incRefCount()
+            }
+          }
+        }
+      }
+    }
+  }
+
+
 }
 
 case class GpuIf(
@@ -203,59 +258,6 @@ case class GpuIf(
     }
   }
 
-  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, t: GpuColumnVector): ColumnVector = {
-    // 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(Scalar.fromInt(Int.MinValue)) {
-          outOfBoundsFlag => predicate.ifElse(prefixSumExclusive, outOfBoundsFlag)
-        }
-
-        withResource(new Table(t.getBase)) { tbl =>
-          withResource(gatherMap) { _ =>
-            withResource(tbl.gather(gatherMap)) { gatherTbl =>
-              gatherTbl.getColumn(0).incRefCount()
-            }
-          }
-        }
-      }
-    }
-  }
-
   override def toString: String = s"if ($predicateExpr) $trueExpr else $falseExpr"
 
   override def sql: String = s"(IF(${predicateExpr.sql}, ${trueExpr.sql}, ${falseExpr.sql}))"
@@ -274,6 +276,9 @@ case class GpuCaseWhen(
     branches.map(_._2.dataType) ++ elseValue.map(_.dataType)
   }
 
+  private lazy val branchesWithSideEffects =
+    branches.exists(_._2.asInstanceOf[GpuExpression].hasSideEffects)
+
   override def nullable: Boolean = {
     // Result is nullable if any of the branch is nullable, or if the else value is nullable
     branches.exists(_._2.nullable) || elseValue.forall(_.nullable)
@@ -301,12 +306,139 @@ case class GpuCaseWhen(
   }
 
   override def columnarEval(batch: ColumnarBatch): Any = {
-    // `elseRet` will be closed in `computeIfElse`.
-    val elseRet = elseValue
-      .map(_.columnarEval(batch))
-      .getOrElse(GpuScalar(null, branches.last._2.dataType))
-    branches.foldRight[Any](elseRet) { case ((predicateExpr, trueExpr), falseRet) =>
-      computeIfElse(batch, predicateExpr, trueExpr, falseRet)
+    if (branchesWithSideEffects) {
+      columnarEvalWithSideEffects(batch)
+    } else {
+      // `elseRet` will be closed in `computeIfElse`.
+      val elseRet = elseValue
+        .map(_.columnarEval(batch))
+        .getOrElse(GpuScalar(null, branches.last._2.dataType))
+      branches.foldRight[Any](elseRet) {
+        case ((predicateExpr, trueExpr), falseRet) =>
+          computeIfElse(batch, predicateExpr, trueExpr, falseRet)
+        }
+      }
+    }
+
+  /**
+   * Perform lazy evaluation of each branch sa that we only evaluate the THEN expressions
+   * against rows where the WHEN expression is true.
+   */
+  private def columnarEvalWithSideEffects(batch: ColumnarBatch): Any = {
+    val colTypes = GpuColumnVector.extractTypes(batch)
+
+    // track cumulative state of predicate evaluation per row so that we never evaluate expressions
+    // for a row if an earlier expression has already been evaluated to true for that row
+    var cumulativePred: Option[GpuColumnVector] = None
+
+    // this variable contains the currently evaluated value for each row and gets updated
+    // as each branch is evaluated
+    var currentValue: Option[GpuColumnVector] = None
+
+    try {
+      withResource(GpuColumnVector.from(batch)) { tbl =>
+
+        // iterate over the WHEN THEN branches first
+        branches.foreach {
+          case (whenExpr, thenExpr) =>
+            // evaluate the WHEN predicate
+            withResource(GpuExpressionsUtils.columnarEvalToColumn(whenExpr, batch)) { whenBool =>
+              // we only want to evaluate where this WHEN is true and no previous WHEN has been true
+              val whenBoolAdjusted: GpuColumnVector = cumulativePred match {
+                case Some(prev) =>
+                  withResource(prev.getBase.not()) { notPreviouslyTrue =>
+                    GpuColumnVector.from(
+                      whenBool.getBase.and(notPreviouslyTrue), DataTypes.BooleanType)
+                  }
+                case None =>
+                  whenBool.incRefCount()
+              }
+
+              withResource(whenBoolAdjusted) { _ =>
+                if (isAllTrue(whenBoolAdjusted)) {
+                  // if this WHEN predicate is true for all rows and no previous predicate has
+                  // been true then we can return immediately
+                  return GpuExpressionsUtils.columnarEvalToColumn(thenExpr, batch)
+                }
+                val thenValues = withResource(filterBatch(tbl,
+                    whenBoolAdjusted.getBase, colTypes)) {
+                  trueBatch => GpuExpressionsUtils.columnarEvalToColumn(thenExpr, trueBatch)
+                }
+                withResource(thenValues) { _ =>
+                  withResource(gather(whenBoolAdjusted.getBase, thenValues)) { gather =>
+                    currentValue = Some(currentValue match {
+                      case Some(v) =>
+                        withResource(v) { _ =>
+                          GpuColumnVector.from(whenBoolAdjusted.getBase.ifElse(gather, v.getBase),
+                            dataType)
+                        }
+                      case _ =>
+                        GpuColumnVector.from(gather.incRefCount(), dataType)
+                    })
+                  }
+
+                  cumulativePred = Some(cumulativePred match {
+                    case Some(prev) =>
+                      withResource(prev) { _ =>
+                        GpuColumnVector.from(whenBool.getBase.or(prev.getBase),
+                          DataTypes.BooleanType)
+                      }
+                    case _ =>
+                      whenBoolAdjusted.incRefCount()
+                  })
+
+                  if (isAllTrue(cumulativePred.get)) {
+                    // no need to process any more branches or the else condition
+                    return currentValue.get.incRefCount()
+                  }
+                }
+              }
+            }
+        }
+
+        // invert the cumulative predicate to get the ELSE predicate
+        withResource(cumulativePred.get.getBase.not()) { elsePredicate =>
+          // replace null predicates with true (because this is an inverted predicate)
+          val elsePredNoNulls = withResource(Scalar.fromBool(true)) { t =>
+            GpuColumnVector.from(elsePredicate.replaceNulls(t), DataTypes.BooleanType)
+          }
+          withResource(elsePredNoNulls) { _ =>
+            elseValue match {
+              case Some(expr) =>
+                if (isAllTrue(elsePredNoNulls)) {
+                  GpuExpressionsUtils.columnarEvalToColumn(expr, batch)
+                } else {
+                  val elseValues = withResource(
+                      filterBatch(tbl, elsePredNoNulls.getBase, colTypes)) {
+                    elseBatch => GpuExpressionsUtils.columnarEvalToColumn(expr, elseBatch)
+                  }
+                  withResource(elseValues) { _ =>
+                    withResource(gather(elsePredNoNulls.getBase, elseValues)) { gather =>
+                      GpuColumnVector.from(elsePredNoNulls.getBase.ifElse(
+                        gather, currentValue.get.getBase), dataType)
+                    }
+                  }
+                }
+
+              case None =>
+                // if there is no ELSE condition then we return NULL for any rows not matched by
+                // previous branches
+                withResource(GpuScalar.from(null, dataType)) { nullScalar =>
+                  if (isAllTrue(elsePredNoNulls) && elsePredNoNulls.getRowCount <= Int.MaxValue) {
+                    GpuColumnVector.from(nullScalar, elsePredNoNulls.getRowCount.toInt, dataType)
+                  } else {
+                    GpuColumnVector.from(
+                      elsePredNoNulls.getBase.ifElse(nullScalar, currentValue.get.getBase),
+                        dataType)
+                  }
+                }
+            }
+          }
+        }
+      }
+    } finally {
+      currentValue.foreach(_.safeClose())
+      cumulativePred.foreach(_.safeClose())
     }
   }