case when improvement: avoid copy_if_else

src/main/cpp/src/case_when.cu

@@ -54,7 +54,9 @@ struct select_first_true_fn {
     bool found_true                  = false;
     cudf::size_type first_true_index = col_num;
     for (auto col_idx = 0; !found_true && col_idx < col_num; col_idx++) {
-      if (d_table.column(col_idx).element<bool>(row_idx)) {
+      auto const& col = d_table.column(col_idx);
+      if (!col.is_null(row_idx) && col.element<bool>(row_idx)) {
+        // Predicate is true and not null
         found_true       = true;
         first_true_index = col_idx;
       }
@@ -114,7 +116,7 @@ std::unique_ptr<cudf::column> select_from_index(
                       cudf::size_type scalar_idx = d_select_index.element<cudf::size_type>(row_idx);
 
                       // return <str_ptr, str_size> pair
-                      if (scalar_idx < num_of_scalar) {
+                      if (scalar_idx < num_of_scalar && !d_scalars.is_null(scalar_idx)) {
                         auto const d_str = d_scalars.element<cudf::string_view>(scalar_idx);
                         return str_view{d_str.data(), d_str.size_bytes()};
                       } else {

src/test/java/com/nvidia/spark/rapids/jni/CaseWhenTest.java

@@ -17,6 +17,8 @@
 package com.nvidia.spark.rapids.jni;
 
 import ai.rapids.cudf.*;
+import ai.rapids.cudf.HostColumnVector.Builder;
+
 import org.junit.jupiter.api.Test;
 
 import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual;
@@ -42,6 +44,47 @@ void selectIndexTest() {
     }
   }
 
+  public static ColumnVector fromBooleansWithNulls(Boolean... values) {
+    Byte[] bytes = new Byte[values.length];
+    for (int i = 0; i < values.length; i++) {
+      if (values[i] == null) {
+        bytes[i] = null;
+      } else {
+        bytes[i] = values[i] ? (byte) 1 : (byte) 0;
+      }
+    }
+
+    try (Builder builder = HostColumnVector.builder(DType.BOOL8, bytes.length)) {
+      for (Byte bool : bytes) {
+        if (bool == null) {
+          builder.appendNull();
+        } else {
+          builder.append(bool.byteValue());
+        }
+      }
+      return builder.buildAndPutOnDevice();
+    }
+  }
+
+  @Test
+  void selectIndexTestWithNull() {
+    try (
+        ColumnVector b0 = fromBooleansWithNulls(
+            null, false, false, null, false);
+        ColumnVector b1 = fromBooleansWithNulls(
+            null, null, false, true, true);
+        ColumnVector b2 = fromBooleansWithNulls(
+            null, null, false, true, false);
+        ColumnVector b3 = fromBooleansWithNulls(
+            null, null, null, true, null);
+        ColumnVector expected = ColumnVector.fromInts(4, 4, 4, 1, 1)) {
+      ColumnVector[] boolColumns = new ColumnVector[] { b0, b1, b2, b3 };
+      try (ColumnVector actual = CaseWhen.selectFirstTrueIndex(boolColumns)) {
+        assertColumnsAreEqual(expected, actual);
+      }
+    }
+  }
+
   @Test
   void selectTest() {
     try (ColumnVector values = ColumnVector.fromStrings(
@@ -53,4 +96,16 @@ void selectTest() {
       assertColumnsAreEqual(expected, actual);
     }
   }
+
+  @Test
+  void selectTestWillNull() {
+    try (ColumnVector values = ColumnVector.fromStrings(
+        "s0", null, "s2", "s3");
+        ColumnVector selects = ColumnVector.fromInts(0, 1, 2, 3, 3, 2, 1, 0, 4, 5, 6);
+        ColumnVector expected = ColumnVector.fromStrings("s0", null, "s2", "s3", "s3", "s2", null, "s0", null, null,
+            null);
+        ColumnVector actual = CaseWhen.selectFromIndex(values, selects)) {
+      assertColumnsAreEqual(expected, actual);
+    }
+  }
 }