From 90e4b25620b85e3c948832bff939d5ab3a58ef4f Mon Sep 17 00:00:00 2001
From: Zoltan Haindrich <zoltan.haindrich@imply.io>
Date: Wed, 4 Oct 2023 14:08:46 +0200
Subject: [PATCH] Fix lead/lag to be usable without offset (#15057)

---
 .../druid/sql/calcite/rel/Windowing.java      |  4 +-
 .../calcite/tests/window/lead_lag.sqlTest     | 50 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 sql/src/test/resources/calcite/tests/window/lead_lag.sqlTest

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/Windowing.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/Windowing.java
index e486a958812..07c5544441d 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/Windowing.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/Windowing.java
@@ -88,9 +88,9 @@ public class Windowing
   private static final ImmutableMap<String, ProcessorMaker> KNOWN_WINDOW_FNS = ImmutableMap
       .<String, ProcessorMaker>builder()
       .put("LAG", (agg) ->
-          new WindowOffsetProcessor(agg.getColumn(0), agg.getOutputName(), -agg.getConstantInt(1)))
+          new WindowOffsetProcessor(agg.getColumn(0), agg.getOutputName(), -agg.getConstantInt(1, 1)))
       .put("LEAD", (agg) ->
-          new WindowOffsetProcessor(agg.getColumn(0), agg.getOutputName(), agg.getConstantInt(1)))
+          new WindowOffsetProcessor(agg.getColumn(0), agg.getOutputName(), agg.getConstantInt(1, 1)))
       .put("FIRST_VALUE", (agg) ->
           new WindowFirstProcessor(agg.getColumn(0), agg.getOutputName()))
       .put("LAST_VALUE", (agg) ->
diff --git a/sql/src/test/resources/calcite/tests/window/lead_lag.sqlTest b/sql/src/test/resources/calcite/tests/window/lead_lag.sqlTest
new file mode 100644
index 00000000000..1ce080b08b6
--- /dev/null
+++ b/sql/src/test/resources/calcite/tests/window/lead_lag.sqlTest
@@ -0,0 +1,50 @@
+type: "operatorValidation"
+
+sql: |
+  SELECT
+    dim1,
+    LAG(dim1,2) OVER (),
+    LAG(dim1) OVER (),
+    LAG(dim1,0) OVER (),
+    LEAD(dim1,0) OVER (),
+    LEAD(dim1) OVER (),
+    LEAD(dim1,2) OVER ()
+  FROM foo
+  WHERE length(dim1) > 1
+  GROUP BY dim1
+
+expectedOperators:
+  - type: "naivePartition"
+    partitionColumns: []
+  - type: "window"
+    processor:
+      type: "composing"
+      processors:
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w0"
+          offset: -2
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w1"
+          offset: -1
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w2"
+          offset: 0
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w3"
+          offset: 0
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w4"
+          offset: 1
+        - type: "offset"
+          inputColumn: "d0"
+          outputColumn: "w5"
+          offset: 2
+expectedResults:
+  - ["10.1",null,null,"10.1","10.1","abc","def"]
+  - ["abc",null,"10.1","abc","abc","def",null]
+  - ["def","10.1","abc","def","def",null,null]