Add IGNORE NULLS support to FIRST_VALUE and LAST_VALUE window functions

pinot-common/src/main/proto/expressions.proto

@@ -92,6 +92,7 @@ message FunctionCall {
   string functionName = 2;
   repeated Expression functionOperands = 3;
   bool isDistinct = 4;
+  bool ignoreNulls = 5;
 }
 
 message Expression {

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

@@ -326,6 +326,31 @@ public void testAggregateServerReturnFinalResult(boolean useMultiStageQueryEngin
     assertTrue(response.get("resultTable").get("rows").get(0).get(0).isNull());
   }
 
+  @Test
+  public void testWindowFunctionIgnoreNulls()
+      throws Exception {
+    // Window functions are only supported in the multi-stage query engine
+    setUseMultiStageQueryEngine(true);
+    String sqlQuery =
+        "SELECT salary, LAST_VALUE(salary) IGNORE NULLS OVER (ORDER BY DaysSinceEpoch) AS gapfilledSalary from "
+            + "mytable";
+    JsonNode response = postQuery(sqlQuery);
+    assertNoError(response);
+
+    // Check if the LAST_VALUE window function with IGNORE NULLS has effectively gap-filled the salary values
+    Integer lastSalary = null;
+    JsonNode rows = response.get("resultTable").get("rows");
+    for (int i = 0; i < rows.size(); i++) {
+      JsonNode row = rows.get(i);
+      if (!row.get(0).isNull()) {
+        assertEquals(row.get(0).asInt(), row.get(1).asInt());
+        lastSalary = row.get(0).asInt();
+      } else {
+        assertEquals(lastSalary, row.get(1).numberValue());
+      }
+    }
+  }
+
   @Override
   protected void overrideBrokerConf(PinotConfiguration brokerConf) {
     brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_QUERY_ENABLE_NULL_HANDLING, "true");

pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java

@@ -32,6 +32,7 @@
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlOperatorTable;
 import org.apache.calcite.sql.SqlSyntax;
+import org.apache.calcite.sql.fun.SqlLeadLagAggFunction;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.ReturnTypes;
@@ -134,10 +135,6 @@ public static PinotOperatorTable instance() {
       SqlStdOperatorTable.MODE,
       SqlStdOperatorTable.MIN,
       SqlStdOperatorTable.MAX,
-      SqlStdOperatorTable.LAST_VALUE,
-      SqlStdOperatorTable.FIRST_VALUE,
-      SqlStdOperatorTable.LEAD,
-      SqlStdOperatorTable.LAG,
       SqlStdOperatorTable.AVG,
       SqlStdOperatorTable.STDDEV_POP,
       SqlStdOperatorTable.COVAR_POP,
@@ -152,7 +149,17 @@ public static PinotOperatorTable instance() {
       SqlStdOperatorTable.RANK,
       SqlStdOperatorTable.ROW_NUMBER,
 
+      // WINDOW Functions (non-aggregate)
+      SqlStdOperatorTable.LAST_VALUE,
+      SqlStdOperatorTable.FIRST_VALUE,
+      // TODO: Replace these with SqlStdOperatorTable.LEAD and SqlStdOperatorTable.LAG when the function implementations
+      // are updated to support the IGNORE NULLS option.
+      PinotLeadWindowFunction.INSTANCE,
+      PinotLagWindowFunction.INSTANCE,
+
       // SPECIAL OPERATORS
+      SqlStdOperatorTable.IGNORE_NULLS,
+      SqlStdOperatorTable.RESPECT_NULLS,
       SqlStdOperatorTable.BETWEEN,
       SqlStdOperatorTable.SYMMETRIC_BETWEEN,
       SqlStdOperatorTable.NOT_BETWEEN,
@@ -372,4 +379,30 @@ public void lookupOperatorOverloads(SqlIdentifier opName, @Nullable SqlFunctionC
   public List<SqlOperator> getOperatorList() {
     return _operatorList;
   }
+
+  private static class PinotLeadWindowFunction extends SqlLeadLagAggFunction {
+    static final SqlOperator INSTANCE = new PinotLeadWindowFunction();
+
+    public PinotLeadWindowFunction() {
+      super(SqlKind.LEAD);
+    }
+
+    @Override
+    public boolean allowsNullTreatment() {
+      return false;
+    }
+  }
+
+  private static class PinotLagWindowFunction extends SqlLeadLagAggFunction {
+    static final SqlOperator INSTANCE = new PinotLagWindowFunction();
+
+    public PinotLagWindowFunction() {
+      super(SqlKind.LAG);
+    }
+
+    @Override
+    public boolean allowsNullTreatment() {
+      return false;
+    }
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java

@@ -324,7 +324,7 @@ public Void visitWindow(WindowNode node, Void context) {
           RelDataType relDataType = funCall.getDataType().toType(_builder.getTypeFactory());
           Window.RexWinAggCall winCall = new Window.RexWinAggCall(aggFunction, relDataType, operands, aggCalls.size(),
               // same as the one used in LogicalWindow.create
-              funCall.isDistinct(), false);
+              funCall.isDistinct(), funCall.isIgnoreNulls());
           aggCalls.add(winCall);
         }
 

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java

@@ -111,17 +111,20 @@ class FunctionCall implements RexExpression {
     private final List<RexExpression> _functionOperands;
     // whether the function is a distinct function.
     private final boolean _isDistinct;
+    // whether the function should ignore nulls (relevant to certain window functions like LAST_VALUE).
+    private final boolean _ignoreNulls;
 
     public FunctionCall(ColumnDataType dataType, String functionName, List<RexExpression> functionOperands) {
-      this(dataType, functionName, functionOperands, false);
+      this(dataType, functionName, functionOperands, false, false);
     }
 
     public FunctionCall(ColumnDataType dataType, String functionName, List<RexExpression> functionOperands,
-        boolean isDistinct) {
+        boolean isDistinct, boolean ignoreNulls) {
       _dataType = dataType;
       _functionName = functionName;
       _functionOperands = functionOperands;
       _isDistinct = isDistinct;
+      _ignoreNulls = ignoreNulls;
     }
 
     public ColumnDataType getDataType() {
@@ -140,6 +143,10 @@ public boolean isDistinct() {
       return _isDistinct;
     }
 
+    public boolean isIgnoreNulls() {
+      return _ignoreNulls;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) {

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java

@@ -405,12 +405,12 @@ private static List<RexExpression> toFunctionOperands(RexInputRef rexInputRef, S
   public static RexExpression.FunctionCall fromAggregateCall(AggregateCall aggregateCall) {
     return new RexExpression.FunctionCall(RelToPlanNodeConverter.convertToColumnDataType(aggregateCall.type),
         getFunctionName(aggregateCall.getAggregation()), fromRexNodes(aggregateCall.rexList),
-        aggregateCall.isDistinct());
+        aggregateCall.isDistinct(), false);
   }
 
   public static RexExpression.FunctionCall fromWindowAggregateCall(Window.RexWinAggCall winAggCall) {
     return new RexExpression.FunctionCall(RelToPlanNodeConverter.convertToColumnDataType(winAggCall.type),
-        getFunctionName(winAggCall.op), fromRexNodes(winAggCall.operands), winAggCall.distinct);
+        getFunctionName(winAggCall.op), fromRexNodes(winAggCall.operands), winAggCall.distinct, winAggCall.ignoreNulls);
   }
 
   public static Integer getValueAsInt(@Nullable RexNode in) {

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java

@@ -58,7 +58,7 @@ public static RexExpression.FunctionCall convertFunctionCall(Expressions.Functio
       operands.add(convertExpression(protoOperand));
     }
     return new RexExpression.FunctionCall(convertColumnDataType(functionCall.getDataType()),
-        functionCall.getFunctionName(), operands, functionCall.getIsDistinct());
+        functionCall.getFunctionName(), operands, functionCall.getIsDistinct(), functionCall.getIgnoreNulls());
   }
 
   public static RexExpression.Literal convertLiteral(Expressions.Literal literal) {

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java

@@ -64,9 +64,15 @@ public static Expressions.FunctionCall convertFunctionCall(RexExpression.Functio
     for (RexExpression operand : operands) {
       protoOperands.add(convertExpression(operand));
     }
-    return Expressions.FunctionCall.newBuilder().setDataType(convertColumnDataType(functionCall.getDataType()))
-        .setFunctionName(functionCall.getFunctionName()).addAllFunctionOperands(protoOperands)
-        .setIsDistinct(functionCall.isDistinct()).build();
+    //@formatter:off
+    return Expressions.FunctionCall.newBuilder()
+        .setDataType(convertColumnDataType(functionCall.getDataType()))
+        .setFunctionName(functionCall.getFunctionName())
+        .addAllFunctionOperands(protoOperands)
+        .setIsDistinct(functionCall.isDistinct())
+        .setIgnoreNulls(functionCall.isIgnoreNulls())
+        .build();
+    //@formatter:on
   }
 
   public static Expressions.Literal convertLiteral(RexExpression.Literal literal) {