Fix lead/lag to be usable without offset (apache#15057)

CaseyPan · Oct 4, 2023 · 90e4b25 · 90e4b25
1 parent c888ac5
commit 90e4b25
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
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
@@ -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
@@ -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]