Support up to 15 boolean expressions for casewhen scalar function (#1…

…1566) Current implementation supports up to 5 boolean expressions. This improves to support up to 15 expressions.
apache · Sep 12, 2023 · b297245 · b297245
1 parent 032a5cc
commit b297245
Show file tree

Hide file tree

Showing 2 changed files with 147 additions and 0 deletions.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ObjectFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ObjectFunctions.java
@@ -128,6 +128,110 @@ public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullab
     return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, oe);
   }
 
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Boolean c11, @Nullable Object o11, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, c11, o11, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Boolean c11, @Nullable Object o11, @Nullable Boolean c12, @Nullable Object o12,
+      @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, c11, o11, c12,
+        o12, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Boolean c11, @Nullable Object o11, @Nullable Boolean c12, @Nullable Object o12,
+      @Nullable Boolean c13, @Nullable Object o13, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, c11, o11, c12,
+        o12, c13, o13, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Boolean c11, @Nullable Object o11, @Nullable Boolean c12, @Nullable Object o12,
+      @Nullable Boolean c13, @Nullable Object o13, @Nullable Boolean c14, @Nullable Object o14, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, c11, o11, c12,
+        o12, c13, o13, c14, o14, oe);
+  }
+
+  @Nullable
+  @ScalarFunction(nullableParameters = true, names = {"case", "caseWhen", "case_when"})
+  public static Object caseWhen(@Nullable Boolean c1, @Nullable Object o1, @Nullable Boolean c2, @Nullable Object o2,
+      @Nullable Boolean c3, @Nullable Object o3, @Nullable Boolean c4, @Nullable Object o4, @Nullable Boolean c5,
+      @Nullable Object o5, @Nullable Boolean c6, @Nullable Object o6, @Nullable Boolean c7, @Nullable Object o7,
+      @Nullable Boolean c8, @Nullable Object o8, @Nullable Boolean c9, @Nullable Object o9, @Nullable Boolean c10,
+      @Nullable Object o10, @Nullable Boolean c11, @Nullable Object o11, @Nullable Boolean c12, @Nullable Object o12,
+      @Nullable Boolean c13, @Nullable Object o13, @Nullable Boolean c14, @Nullable Object o14, @Nullable Boolean c15,
+      @Nullable Object o15, @Nullable Object oe) {
+    return caseWhenVar(c1, o1, c2, o2, c3, o3, c4, o4, c5, o5, c6, o6, c7, o7, c8, o8, c9, o9, c10, o10, c11, o11, c12,
+        o12, c13, o13, c14, o14, c15, o15, oe);
+  }
+
   @Nullable
   private static Object caseWhenVar(Object... objs) {
     for (int i = 0; i < objs.length - 1; i += 2) {

diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/function/ObjectFunctionsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/function/ObjectFunctionsTest.java
@@ -113,6 +113,49 @@ public Object[][] objectFunctionsDataProvider() {
         "coalesce(value1,value2,value3,value4,value5)", Lists.newArrayList("value1", "value2", "value3", "value4",
         "value5"), allValues, "1"
     });
+
+    // Adding a test for case when
+    GenericRow caseWhenCaseValueOne = new GenericRow();
+    caseWhenCaseValueOne.putValue("value1", 1);
+    inputs.add(new Object[]{
+        "CASEWHEN(value1 = 1, 'one', value1 = 2, 'two', 'other')", Lists.newArrayList("value1",
+        "value1"), caseWhenCaseValueOne, "one"
+    });
+
+    GenericRow caseWhenCaseValueTwo = new GenericRow();
+    caseWhenCaseValueTwo.putValue("value1", 2);
+    inputs.add(new Object[]{
+        "CASEWHEN(value1 = 1, 'one', value1 = 2, 'two', 'other')", Lists.newArrayList("value1",
+        "value1"), caseWhenCaseValueTwo, "two"
+    });
+
+    GenericRow caseWhenCaseValueThree = new GenericRow();
+    caseWhenCaseValueThree.putValue("value1", 3);
+    inputs.add(new Object[]{
+        "CASEWHEN(value1 = 1, 'one', value1 = 2, 'two', 'other')", Lists.newArrayList("value1",
+        "value1"), caseWhenCaseValueThree, "other"
+    });
+
+    GenericRow caseWhenCaseMultipleExpression = new GenericRow();
+    caseWhenCaseMultipleExpression.putValue("value1", 10);
+    inputs.add(new Object[]{
+        "CASEWHEN(value1 = 1, 'one', value1 = 2, 'two', value1 = 3, 'three', value1 = 4, 'four', value1 = 5, 'five', "
+            + "value1 = 6, 'six', value1 = 7, 'seven', value1 = 8, 'eight', value1 = 9, 'nine', value1 = 10, 'ten', "
+            + "'other')", Lists.newArrayList("value1", "value1", "value1", "value1", "value1", "value1", "value1",
+        "value1", "value1", "value1"), caseWhenCaseMultipleExpression, "ten"
+    });
+
+    GenericRow caseWhenCaseMultipleExpression2 = new GenericRow();
+    caseWhenCaseMultipleExpression2.putValue("value1", 15);
+    inputs.add(new Object[]{
+        "CASEWHEN(value1 = 1, 'one', value1 = 2, 'two', value1 = 3, 'three', value1 = 4, 'four', value1 = 5, 'five', "
+            + "value1 = 6, 'six', value1 = 7, 'seven', value1 = 8, 'eight', value1 = 9, 'nine', value1 = 10, 'ten', "
+            + "value1 = 11, 'eleven', value1 = 12, 'twelve', value1 = 13, 'thirteen', value1 = 14, 'fourteen', value1"
+            + " = 15, 'fifteen'," + "'other')", Lists.newArrayList("value1", "value1", "value1", "value1", "value1",
+        "value1", "value1", "value1", "value1", "value1", "value1", "value1", "value1", "value1",
+        "value1"), caseWhenCaseMultipleExpression2, "fifteen"
+    });
+
     return inputs.toArray(new Object[0][]);
   }
 }