Implement hasSideEffects for GpuGetArrayItem, GpuElementAt, GpuGetMap…

…Value, GpuUnaryMinus, and GpuAbs (#5097)
NVIDIA · Mar 31, 2022 · cca2b37 · cca2b37
1 parent 4b542ce
commit cca2b37
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 0 deletions.
diff --git a/integration_tests/src/main/python/conditionals_test.py b/integration_tests/src/main/python/conditionals_test.py
@@ -226,3 +226,47 @@ def test_conditional_with_side_effects_sequence(data_gen):
             WHEN length(a) > 0 THEN sequence(1, length(a), 1) \
             ELSE null END'),
         conf = ansi_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [ArrayGen(mk_str_gen('[a-z]{0,3}'))], ids=idfn)
+@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
+def test_conditional_with_side_effects_element_at(data_gen, ansi_enabled):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : unary_op_df(spark, data_gen).selectExpr(
+            'CASE WHEN size(a) > 1 THEN element_at(a, 2) ELSE null END'),
+        conf = {'spark.sql.ansi.enabled': ansi_enabled})
+
+@pytest.mark.parametrize('data_gen', [ArrayGen(mk_str_gen('[a-z]{0,3}'))], ids=idfn)
+@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
+def test_conditional_with_side_effects_array_index(data_gen, ansi_enabled):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : unary_op_df(spark, data_gen).selectExpr(
+            'CASE WHEN size(a) > 1 THEN a[1] ELSE null END'),
+        conf = {'spark.sql.ansi.enabled': ansi_enabled})
+
+@pytest.mark.parametrize('map_gen',
+                         [MapGen(StringGen(pattern='key_[0-9]', nullable=False),
+                                 mk_str_gen('[a-z]{0,3}'), max_length=6)])
+@pytest.mark.parametrize('data_gen', [StringGen(pattern='neverempty_[0-9]', nullable=False)])
+@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
+def test_conditional_with_side_effects_map_key_not_found(map_gen, data_gen, ansi_enabled):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: two_col_df(spark, map_gen, data_gen).selectExpr(
+            'CASE WHEN length(b) = 0 THEN a["not_found"] ELSE null END'),
+        conf = {'spark.sql.ansi.enabled': ansi_enabled})
+
+@pytest.mark.parametrize('data_gen', [ShortGen().with_special_case(SHORT_MIN)], ids=idfn)
+@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
+def test_conditional_with_side_effects_abs(data_gen, ansi_enabled):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : unary_op_df(spark, data_gen).selectExpr(
+            'CASE WHEN a > -32768 THEN abs(a) ELSE null END'),
+        conf = {'spark.sql.ansi.enabled': ansi_enabled})
+
+@pytest.mark.parametrize('data_gen', [ShortGen().with_special_case(SHORT_MIN)], ids=idfn)
+@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
+def test_conditional_with_side_effects_unary_minus(data_gen, ansi_enabled):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : unary_op_df(spark, data_gen).selectExpr(
+            'CASE WHEN a > -32768 THEN -a ELSE null END'),
+        conf = {'spark.sql.ansi.enabled': ansi_enabled})
+
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala
@@ -67,6 +67,8 @@ case class GpuUnaryMinus(child: Expression, failOnError: Boolean) extends GpuUna
 
   override def sql: String = s"(- ${child.sql})"
 
+  override def hasSideEffects: Boolean = failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)
+
   override def doColumnar(input: GpuColumnVector) : ColumnVector = {
     if (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) {
       // Because of 2s compliment we need to only worry about the min value for integer types.
@@ -122,6 +124,8 @@ case class GpuAbs(child: Expression, failOnError: Boolean) extends CudfUnaryExpr
 
   override def unaryOp: UnaryOp = UnaryOp.ABS
 
+  override def hasSideEffects: Boolean = failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)
+
   override def doColumnar(input: GpuColumnVector) : ColumnVector = {
     if (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) {
       // Because of 2s compliment we need to only worry about the min value for integer types.

diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala
@@ -84,6 +84,8 @@ case class GpuConcat(children: Seq[Expression]) extends GpuComplexTypeMergingExp
 case class GpuElementAt(left: Expression, right: Expression, failOnError: Boolean)
   extends GpuBinaryExpression with ExpectsInputTypes {
 
+  override def hasSideEffects: Boolean = failOnError
+
   override lazy val dataType: DataType = left.dataType match {
     case ArrayType(elementType, _) => elementType
     case MapType(_, valueType, _) => valueType

diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeExtractors.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeExtractors.scala
@@ -107,6 +107,8 @@ case class GpuGetArrayItem(child: Expression, ordinal: Expression, failOnError:
   override def nullable: Boolean = true
   override def dataType: DataType = child.dataType.asInstanceOf[ArrayType].elementType
 
+  override def hasSideEffects: Boolean = failOnError
+
   override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
     val (array, indices) = (lhs.getBase, rhs.getBase)
     val indicesCol = withResource(Scalar.fromInt(0)) { zeroS =>
@@ -230,6 +232,8 @@ case class GpuGetMapValue(child: Expression, key: Expression, failOnError: Boole
 
   override def prettyName: String = "getMapValue"
 
+  override def hasSideEffects: Boolean = failOnError
+
   override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
     if (failOnError){
       withResource(lhs.getBase.getMapKeyExistence(rhs.getBase)) { keyExistenceColumn =>