Fix null literal handling for null intolerant functions in multi-stage query engine

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArithmeticFunctions.java

@@ -20,12 +20,14 @@
 
 import java.math.BigDecimal;
 import java.math.RoundingMode;
+import org.apache.calcite.linq4j.function.Strict;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 
 
 /**
  * Arithmetic scalar functions.
  */
+@Strict
 public class ArithmeticFunctions {
   private ArithmeticFunctions() {
   }

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java

@@ -25,6 +25,7 @@
 import it.unimi.dsi.fastutil.objects.ObjectSet;
 import java.math.BigDecimal;
 import java.util.Arrays;
+import org.apache.calcite.linq4j.function.SemiStrict;
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
@@ -33,6 +34,7 @@
 /**
  * Inbuilt array scalar functions. See {@link ArrayUtils} for details.
  */
+@SemiStrict
 public class ArrayFunctions {
   private ArrayFunctions() {
   }

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ComparisonFunctions.java

@@ -18,9 +18,10 @@
  */
 package org.apache.pinot.common.function.scalar;
 
+import org.apache.calcite.linq4j.function.Strict;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 
-
+@Strict
 public class ComparisonFunctions {
 
   private static final double DOUBLE_COMPARISON_TOLERANCE = 1e-7d;

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DateTimeFunctions.java

@@ -21,6 +21,7 @@
 import java.sql.Timestamp;
 import java.time.Duration;
 import java.util.concurrent.TimeUnit;
+import org.apache.calcite.linq4j.function.Strict;
 import org.apache.pinot.common.function.DateTimePatternHandler;
 import org.apache.pinot.common.function.DateTimeUtils;
 import org.apache.pinot.common.function.TimeZoneKey;
@@ -69,6 +70,7 @@
  *     }]
  *   </code>
  */
+@Strict
 public class DateTimeFunctions {
   private DateTimeFunctions() {
   }

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/LogicalFunctions.java

@@ -18,11 +18,13 @@
  */
 package org.apache.pinot.common.function.scalar;
 
+import org.apache.calcite.linq4j.function.Strict;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 
-
-// Logical transformation on boolean values.
-// Currently, only not is supported.
+/**
+ * Logical transformation on boolean values. Currently, only not is supported.
+ */
+@Strict
 public class LogicalFunctions {
   private LogicalFunctions() {
   }

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/TrigonometricFunctions.java

@@ -18,9 +18,10 @@
  */
 package org.apache.pinot.common.function.scalar;
 
+import org.apache.calcite.linq4j.function.Strict;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 
-
+@Strict
 public class TrigonometricFunctions {
   private TrigonometricFunctions() {
   }

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/NullHandlingIntegrationTest.java

@@ -27,6 +27,7 @@
 import org.apache.pinot.util.TestUtils;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static org.testng.AssertJUnit.assertEquals;
@@ -129,11 +130,6 @@ protected List<String> getBloomFilterColumns() {
     return null;
   }
 
-  @Override
-  protected boolean getNullHandlingEnabled() {
-    return true;
-  }
-
   @Override
   protected long getCountStarResult() {
     return 100;
@@ -151,7 +147,6 @@ public void testTotalCount(boolean useMultiStageQueryEngine)
   public void testCountWithNullDescription(boolean useMultiStageQueryEngine)
       throws Exception {
     setUseMultiStageQueryEngine(useMultiStageQueryEngine);
-    notSupportedInV2();
     String query = "SELECT COUNT(*) FROM " + getTableName() + " WHERE description IS NOT NULL";
     testQuery(query);
   }
@@ -160,7 +155,6 @@ public void testCountWithNullDescription(boolean useMultiStageQueryEngine)
   public void testCountWithNullDescriptionAndSalary(boolean useMultiStageQueryEngine)
       throws Exception {
     setUseMultiStageQueryEngine(useMultiStageQueryEngine);
-    notSupportedInV2();
     String query = "SELECT COUNT(*) FROM " + getTableName() + " WHERE description IS NOT NULL AND salary IS NOT NULL";
     testQuery(query);
   }
@@ -215,101 +209,49 @@ public void testTotalCountWithNullHandlingQueryOptionEnabled(boolean useMultiSta
     DataTableBuilderFactory.setDataTableVersion(DataTableBuilderFactory.DEFAULT_VERSION);
   }
 
-  @Test(dataProvider = "useBothQueryEngines")
-  public void testNullLiteralSelectionOnlyBroker(boolean useMultiStageQueryEngine)
+  @Test(dataProvider = "nullLiteralQueries")
+  public void testNullLiteralSelectionOnlyBroker(String sqlQuery, Object expectedResult)
       throws Exception {
-    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
-    notSupportedInV2();
-    // Null literal only
-    String sqlQuery = "SELECT null FROM mytable OPTION(enableNullHandling=true)";
+    // V2 handles such queries in the servers where all rows in the table are matched and hence the number of rows
+    // returned in the result set is equal to the total number of rows in the table (instead of a single row like in
+    // V1 where it is handled in the broker). The V2 way is more standard though and matches behavior in other
+    // databases like Postgres.
+    setUseMultiStageQueryEngine(false);
+
     JsonNode response = postQuery(sqlQuery);
     JsonNode rows = response.get("resultTable").get("rows");
     assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asText(), "null");
-
-    // Null related functions
-    sqlQuery = "SELECT isNull(null) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), true);
-
-    sqlQuery = "SELECT isNotNull(null) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), false);
-
-
-    sqlQuery = "SELECT coalesce(null, 1) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asInt(), 1);
-
-    sqlQuery = "SELECT coalesce(null, null) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asText(), "null");
-
-    sqlQuery = "SELECT isDistinctFrom(null, null) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), false);
-
-    sqlQuery = "SELECT isNotDistinctFrom(null, null) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), true);
-
-
-    sqlQuery = "SELECT isDistinctFrom(null, 1) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), true);
-
-    sqlQuery = "SELECT isNotDistinctFrom(null, 1) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asBoolean(), false);
-
-    sqlQuery = "SELECT case when true then null end FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asText(), "null");
-
-
-    sqlQuery = "SELECT case when false then 1 end FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
-    assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asText(), "null");
+    assertEquals(1, rows.size());
+
+    if (expectedResult instanceof String) {
+      assertEquals(expectedResult, rows.get(0).get(0).asText());
+    } else if (expectedResult instanceof Integer) {
+      assertEquals(expectedResult, rows.get(0).get(0).asInt());
+    } else if (expectedResult instanceof Boolean) {
+      assertEquals(expectedResult, rows.get(0).get(0).asBoolean());
+    } else {
+      throw new IllegalArgumentException("Unexpected type for expectedResult: " + expectedResult.getClass());
+    }
+  }
 
+  @Test(dataProvider = "nullLiteralQueries")
+  public void testNullLiteralSelectionInV2(String sqlQuery, Object expectedResult) throws Exception {
+    setUseMultiStageQueryEngine(true);
 
-    // Null intolerant functions
-    sqlQuery = "SELECT add(null, 1) FROM " + getTableName() + "  OPTION (enableNullHandling=true);";
-    response = postQuery(sqlQuery);
-    rows = response.get("resultTable").get("rows");
+    JsonNode response = postQuery(sqlQuery);
+    JsonNode rows = response.get("resultTable").get("rows");
     assertTrue(response.get("exceptions").isEmpty());
-    assertEquals(rows.size(), 1);
-    assertEquals(rows.get(0).get(0).asText(), "null");
+    assertEquals(getCountStarResult(), rows.size());
+
+    if (expectedResult instanceof String) {
+      assertEquals(expectedResult, rows.get(0).get(0).asText());
+    } else if (expectedResult instanceof Integer) {
+      assertEquals(expectedResult, rows.get(0).get(0).asInt());
+    } else if (expectedResult instanceof Boolean) {
+      assertEquals(expectedResult, rows.get(0).get(0).asBoolean());
+    } else {
+      throw new IllegalArgumentException("Unexpected type for expectedResult: " + expectedResult.getClass());
+    }
   }
 
   @Test
@@ -378,4 +320,38 @@ public void testCaseWhenAllLiteral(boolean useMultiStageQueryEngine)
 
     assertEquals(response.get("resultTable").get("rows").get(0).get(0).asInt(), 1);
   }
+
+  @DataProvider(name = "nullLiteralQueries")
+  public Object[][] nullLiteralQueries() {
+    // Query string, expected value
+    return new Object[][]{
+        // Null literal only
+        {String.format("SELECT null FROM %s OPTION(enableNullHandling=true)", getTableName()), "null"},
+        // Null related functions
+        {String.format("SELECT isNull(null) FROM %s OPTION (enableNullHandling=true)", getTableName()), true},
+        {String.format("SELECT isNotNull(null) FROM %s OPTION (enableNullHandling=true)", getTableName()), false},
+        {String.format("SELECT coalesce(null, 1) FROM %s OPTION (enableNullHandling=true)", getTableName()), 1},
+        {String.format("SELECT coalesce(null, null) FROM %s OPTION (enableNullHandling=true)", getTableName()), "null"},
+        {String.format("SELECT isDistinctFrom(null, null) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            false},
+        {String.format("SELECT isNotDistinctFrom(null, null) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            true},
+        {String.format("SELECT isDistinctFrom(null, 1) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            true},
+        {String.format("SELECT isNotDistinctFrom(null, 1) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            false},
+        {String.format("SELECT case when true then null end FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            "null"},
+        {String.format("SELECT case when false then 1 end FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            "null"},
+        // Null intolerant functions
+        {String.format("SELECT add(null, 1) FROM %s OPTION (enableNullHandling=true)", getTableName()), "null"},
+        {String.format("SELECT greater_than(null, 1) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            "null"},
+        {String.format("SELECT to_epoch_seconds(null) FROM %s OPTION (enableNullHandling=true)", getTableName()),
+            "null"},
+        {String.format("SELECT not(null) FROM %s OPTION (enableNullHandling=true)", getTableName()), "null"},
+        {String.format("SELECT tan(null) FROM %s OPTION (enableNullHandling=true)", getTableName()), "null"}
+    };
+  }
 }

pinot-integration-tests/src/test/resources/test_null_handling.schema

@@ -3,22 +3,26 @@
     {
       "dataType": "INT",
       "singleValueField": true,
-      "name": "clientId"
+      "name": "clientId",
+      "notNull": "true"
     },
     {
       "dataType": "STRING",
       "singleValueField": true,
-      "name": "city"
+      "name": "city",
+      "notNull": "true"
     },
     {
       "dataType": "STRING",
       "singleValueField": true,
-      "name": "description"
+      "name": "description",
+      "notNull": "false"
     },
     {
       "dataType": "INT",
       "singleValueField": true,
-      "name": "salary"
+      "name": "salary",
+      "notNull": "false"
     }
   ],
   "timeFieldSpec": {
@@ -28,5 +32,6 @@
       "name": "DaysSinceEpoch"
     }
   },
-  "schemaName": "mytable"
+  "schemaName": "mytable",
+  "enableColumnBasedNullHandling": true
 }