diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunctionTest.java index 329d6625e25..12a35e9ed1a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunctionTest.java +++ b/pinot-core/src/test/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));