Added UTs for null handling in CaseTransform function. (apache#11721)

vvivekiyer · Oct 2, 2023 · c3cb5c0 · c3cb5c0
1 parent cb29d3c
commit c3cb5c0
Showing 1 changed file with 66 additions and 0 deletions.
diff --git a/...est/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunctionTest.java b/...est/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunctionTest.java
@@ -26,6 +26,7 @@
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.request.context.RequestContextUtils;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.roaringbitmap.RoaringBitmap;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -196,6 +197,71 @@ public void testCaseTransformFunctionWithStringResults() {
     }
   }
 
+  @Test
+  public void testCaseTransformationWithNullColumn() {
+    ExpressionContext expression =
+        RequestContextUtils.getExpression(
+            String.format("CASE WHEN %s IS NULL THEN 'aaa' ELSE 'bbb' END", STRING_ALPHANUM_NULL_SV_COLUMN));
+    TransformFunction transformFunction = TransformFunctionFactory.getNullHandlingEnabled(expression, _dataSourceMap);
+    Assert.assertTrue(transformFunction instanceof CaseTransformFunction);
+    Assert.assertEquals(transformFunction.getName(), "case");
+    Assert.assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.STRING);
+
+    String[] expectedValues = new String[NUM_ROWS];
+    for (int i = 0; i < NUM_ROWS; i++) {
+      if (isNullRow(i)) {
+        expectedValues[i] = "aaa";
+      } else {
+        expectedValues[i] = "bbb";
+      }
+    }
+    testTransformFunctionWithNull(transformFunction, expectedValues, new RoaringBitmap());
+  }
+
+  @Test
+  public void testCaseTransformationWithNullThenClause() {
+    ExpressionContext expression =
+        RequestContextUtils.getExpression(
+            String.format("CASE WHEN %s IS NULL THEN NULL ELSE 'bbb' END", STRING_ALPHANUM_NULL_SV_COLUMN));
+    TransformFunction transformFunction = TransformFunctionFactory.getNullHandlingEnabled(expression, _dataSourceMap);
+    Assert.assertTrue(transformFunction instanceof CaseTransformFunction);
+    Assert.assertEquals(transformFunction.getName(), "case");
+    Assert.assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.STRING);
+    String[] expectedValues = new String[NUM_ROWS];
+    RoaringBitmap bitmap = new RoaringBitmap();
+    for (int i = 0; i < NUM_ROWS; i++) {
+      if (isNullRow(i)) {
+        bitmap.add(i);
+      } else {
+        expectedValues[i] = "bbb";
+      }
+    }
+    testTransformFunctionWithNull(transformFunction, expectedValues, bitmap);
+  }
+
+  @Test
+  public void testCaseTransformationWithNullElseClause() {
+    ExpressionContext expression =
+        RequestContextUtils.getExpression(
+            String.format("CASE WHEN %s IS NULL THEN 'aaa' END", STRING_ALPHANUM_NULL_SV_COLUMN));
+    TransformFunction transformFunction = TransformFunctionFactory.getNullHandlingEnabled(expression, _dataSourceMap);
+    Assert.assertTrue(transformFunction instanceof CaseTransformFunction);
+    Assert.assertEquals(transformFunction.getName(), "case");
+    Assert.assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.STRING);
+
+    String[] expectedValues = new String[NUM_ROWS];
+    RoaringBitmap bitmap = new RoaringBitmap();
+    for (int i = 0; i < NUM_ROWS; i++) {
+      if (isNullRow(i)) {
+        expectedValues[i] = "aaa";
+      } else {
+        bitmap.add(i);
+      }
+    }
+    testTransformFunctionWithNull(transformFunction, expectedValues, bitmap);
+  }
+
+
   private void testCaseQueryWithIntResults(String predicate, int[] expectedValues) {
     ExpressionContext expression =
         RequestContextUtils.getExpression(String.format("CASE WHEN %s THEN 100 ELSE 10 END", predicate));