Fixed issue for contains when searching for an empty string (NVIDIA#202)

nartal1 · Jun 17, 2020 · faf34bd · faf34bd
1 parent a6ecf8f
commit faf34bd
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 18 deletions.
diff --git a/integration_tests/src/main/python/string_test.py b/integration_tests/src/main/python/string_test.py
@@ -55,23 +55,13 @@ def test_locate():
                 'locate("A", a, 500)',
                 'locate("_", a, NULL)'))
 
-# Once https://github.com/NVIDIA/spark-rapids/issues/121
-# is fixed this should go away and test_contains should be updated
-@pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/121')
-def test_contains_empty():
-    gen = mk_str_gen('.{0,3}Z?_Z?.{0,3}A?.{0,3}')
-    assert_gpu_and_cpu_are_equal_collect(
-            lambda spark: unary_op_df(spark, gen).select(
-                f.col('a').contains('')))
-
 def test_contains():
     gen = mk_str_gen('.{0,3}Z?_Z?.{0,3}A?.{0,3}')
     assert_gpu_and_cpu_are_equal_collect(
             lambda spark: unary_op_df(spark, gen).select(
                 f.col('a').contains('Z'),
                 f.col('a').contains('Z_'),
-                # https://github.com/NVIDIA/spark-rapids/issues/121
-                #f.col('a').contains(''),
+                f.col('a').contains(''),
                 f.col('a').contains(None)
                 ))
 

diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/stringFunctions.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/stringFunctions.scala
@@ -341,16 +341,22 @@ case class GpuContains(left: GpuExpression, right: GpuExpression) extends GpuBin
   override def toString: String = s"gpucontains($left, $right)"
 
   def doColumnar(lhs: GpuColumnVector, rhs: Scalar): GpuColumnVector = {
-    if (rhs.getJavaString.isEmpty) {
-      val boolScalar = Scalar.fromBool(true)
-      try {
-        GpuColumnVector.from(ColumnVector.fromScalar(boolScalar, lhs.getRowCount.toInt))
-      } finally {
-        boolScalar.close()
+    val ret = if (rhs.getJavaString.isEmpty) {
+      withResource(Scalar.fromBool(true)) { trueScalar =>
+        if (left.nullable) {
+          withResource(Scalar.fromBool(null)) { nullBool =>
+            withResource(lhs.getBase.isNull) { isNull =>
+              isNull.ifElse(nullBool, trueScalar)
+            }
+          }
+        } else {
+          ColumnVector.fromScalar(trueScalar, lhs.getRowCount.toInt)
+        }
       }
     } else {
-      GpuColumnVector.from(lhs.getBase.stringContains(rhs))
+      lhs.getBase.stringContains(rhs)
     }
+    GpuColumnVector.from(ret)
   }
 
   override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): GpuColumnVector =