Add IGNORE NULLS support for LEAD and LAG window functions

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

@@ -38,7 +38,6 @@
 import org.apache.calcite.sql.SqlPostfixOperator;
 import org.apache.calcite.sql.SqlSplittableAggFunction;
 import org.apache.calcite.sql.SqlSyntax;
-import org.apache.calcite.sql.fun.SqlLeadLagAggFunction;
 import org.apache.calcite.sql.fun.SqlMonotonicBinaryOperator;
 import org.apache.calcite.sql.fun.SqlNtileAggFunction;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -214,10 +213,8 @@ public static PinotOperatorTable instance(boolean nullHandlingEnabled) {
       // 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,
+      SqlStdOperatorTable.LEAD,
+      SqlStdOperatorTable.LAG,
 
       // SPECIAL OPERATORS
       SqlStdOperatorTable.IGNORE_NULLS,
@@ -448,32 +445,6 @@ 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;
-    }
-  }
-
   private static final class PinotNtileWindowFunction extends SqlNtileAggFunction {
     static final SqlOperator INSTANCE = new PinotNtileWindowFunction();
 

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/window/WindowFunction.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.query.runtime.operator.window;
 
 import java.util.List;
+import javax.annotation.Nullable;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.query.planner.logical.RexExpression;
@@ -63,6 +64,7 @@ public WindowFunction(RexExpression.FunctionCall aggCall, DataSchema inputSchema
    */
   public abstract List<Object> processRows(List<Object[]> rows);
 
+  @Nullable
   protected Object extractValueFromRow(Object[] row) {
     return _inputRef == -1 ? _literal : (row == null ? null : row[_inputRef]);
   }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/window/value/LagValueWindowFunction.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.query.runtime.operator.window.value;
 
 import com.google.common.base.Preconditions;
+import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.List;
 import org.apache.calcite.rel.RelFieldCollation;
@@ -28,9 +29,10 @@
 import org.apache.pinot.query.runtime.operator.window.WindowFrame;
 
 
-/**
- * The LAG window function doesn't allow custom window frames (and this is enforced by Calcite).
- */
+/// Window function that returns the value of a column from a preceding row within the partition.
+/// Supports an optional offset (default 1), an optional default value for when no row exists at
+/// that offset, and IGNORE NULLS mode which skips null values when scanning backward.
+/// Custom window frames are not allowed (enforced by Calcite).
 public class LagValueWindowFunction extends ValueWindowFunction {
   private final int _offset;
   private final Object _defaultValue;
@@ -75,6 +77,9 @@ public LagValueWindowFunction(RexExpression.FunctionCall aggCall, DataSchema inp
 
   @Override
   public List<Object> processRows(List<Object[]> rows) {
+    if (_ignoreNulls) {
+      return processRowsIgnoreNulls(rows);
+    }
     int numRows = rows.size();
     Object[] result = new Object[numRows];
     if (_defaultValue != null) {
@@ -88,4 +93,28 @@ public List<Object> processRows(List<Object[]> rows) {
     }
     return Arrays.asList(result);
   }
+
+  /**
+   * LAG with IGNORE NULLS: for each row, find the offset-th non-null value scanning backward.
+   * Uses a bounded deque of size {@code _offset} for O(N) time and O(offset) memory.
+   * Scans left-to-right, maintaining a sliding window of preceding non-null values. The oldest
+   * element in the deque (peekFirst) is always the offset-th non-null value behind the current
+   * row.
+   */
+  private List<Object> processRowsIgnoreNulls(List<Object[]> rows) {
+    int numRows = rows.size();
+    Object[] result = new Object[numRows];
+    ArrayDeque<Object> window = new ArrayDeque<>(_offset);
+    for (int i = 0; i < numRows; i++) {
+      result[i] = (window.size() == _offset) ? window.peekFirst() : _defaultValue;
+      Object val = extractValueFromRow(rows.get(i));
+      if (val != null) {
+        window.addLast(val);
+        if (window.size() > _offset) {
+          window.pollFirst();
+        }
+      }
+    }
+    return Arrays.asList(result);
+  }
 }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/window/value/LeadValueWindowFunction.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.query.runtime.operator.window.value;
 
 import com.google.common.base.Preconditions;
+import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.List;
 import org.apache.calcite.rel.RelFieldCollation;
@@ -28,9 +29,10 @@
 import org.apache.pinot.query.runtime.operator.window.WindowFrame;
 
 
-/**
- * The LAG window function doesn't allow custom window frames (and this is enforced by Calcite).
- */
+/// Window function that returns the value of a column from a subsequent row within the partition.
+/// Supports an optional offset (default 1), an optional default value for when no row exists at
+/// that offset, and IGNORE NULLS mode which skips null values when scanning forward.
+/// Custom window frames are not allowed (enforced by Calcite).
 public class LeadValueWindowFunction extends ValueWindowFunction {
 
   private final int _offset;
@@ -75,6 +77,9 @@ public LeadValueWindowFunction(RexExpression.FunctionCall aggCall, DataSchema in
 
   @Override
   public List<Object> processRows(List<Object[]> rows) {
+    if (_ignoreNulls) {
+      return processRowsIgnoreNulls(rows);
+    }
     int numRows = rows.size();
     Object[] result = new Object[numRows];
     for (int i = 0; i < numRows - _offset; i++) {
@@ -88,4 +93,28 @@ public List<Object> processRows(List<Object[]> rows) {
     }
     return Arrays.asList(result);
   }
+
+  /**
+   * LEAD with IGNORE NULLS: for each row, find the offset-th non-null value scanning forward.
+   * Uses a bounded deque of size {@code _offset} for O(N) time and O(offset) memory.
+   * Scans right-to-left, maintaining a sliding window of upcoming non-null values. The oldest
+   * element in the deque (peekFirst) is always the offset-th non-null value ahead of the current
+   * row.
+   */
+  private List<Object> processRowsIgnoreNulls(List<Object[]> rows) {
+    int numRows = rows.size();
+    Object[] result = new Object[numRows];
+    ArrayDeque<Object> window = new ArrayDeque<>(_offset);
+    for (int i = numRows - 1; i >= 0; i--) {
+      result[i] = (window.size() == _offset) ? window.peekFirst() : _defaultValue;
+      Object val = extractValueFromRow(rows.get(i));
+      if (val != null) {
+        window.addLast(val);
+        if (window.size() > _offset) {
+          window.pollFirst();
+        }
+      }
+    }
+    return Arrays.asList(result);
+  }
 }