[cherry-pick](branch-21) add more signatures for lag/lead fucntion (#…

…47940) (#48076)
apache · Feb 19, 2025 · bb6641b · bb6641b
1 parent caad4a7
commit bb6641b
Show file tree

Hide file tree

Showing 8 changed files with 254 additions and 64 deletions.
diff --git a/...e-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/window/Lag.java b/...e-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/window/Lag.java
@@ -21,7 +21,9 @@
 import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral;
 import org.apache.doris.nereids.trees.expressions.literal.Literal;
+import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
 import org.apache.doris.nereids.trees.expressions.shape.TernaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.BigIntType;
@@ -39,9 +41,11 @@ public class Lag extends WindowFunction implements TernaryExpression, Explicitly
 
     static {
         List<FunctionSignature> signatures = Lists.newArrayList();
-        trivialTypes.forEach(t ->
-                signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE, t))
-        );
+        trivialTypes.forEach(t -> {
+            signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE, t));
+            signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE));
+            signatures.add(FunctionSignature.ret(t).args(t));
+        });
         SIGNATURES = ImmutableList.copyOf(signatures);
     }
 
@@ -51,21 +55,21 @@ public Lag(Expression child, Expression offset, Expression defaultValue) {
         super("lag", child, offset, defaultValue);
     }
 
-    private Lag(List<Expression> children) {
-        super("lag", children);
+    public Lag(Expression child, Expression offset) {
+        this(child, offset, new NullLiteral(child.getDataType()));
+    }
+
+    public Lag(Expression child) {
+        this(child, new BigIntLiteral(1L), new NullLiteral(child.getDataType()));
     }
 
     public Expression getOffset() {
-        if (children().size() <= 1) {
-            throw new AnalysisException("Not set offset of Lead(): " + this.toSql());
-        }
+        Preconditions.checkArgument(children.size() == 3);
         return child(1);
     }
 
     public Expression getDefaultValue() {
-        if (children.size() <= 2) {
-            throw new AnalysisException("Not set default value of Lead(): " + this.toSql());
-        }
+        Preconditions.checkArgument(children.size() == 3);
         return child(2);
     }
 
@@ -80,7 +84,13 @@ public boolean nullable() {
     @Override
     public Lag withChildren(List<Expression> children) {
         Preconditions.checkArgument(children.size() >= 1 && children.size() <= 3);
-        return new Lag(children);
+        if (children.size() == 1) {
+            return new Lag(children.get(0));
+        } else if (children.size() == 2) {
+            return new Lag(children.get(0), children.get(1));
+        } else {
+            return new Lag(children.get(0), children.get(1), children.get(2));
+        }
     }
 
     @Override

diff --git a/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/window/Lead.java b/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/window/Lead.java
@@ -21,7 +21,9 @@
 import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral;
 import org.apache.doris.nereids.trees.expressions.literal.Literal;
+import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
 import org.apache.doris.nereids.trees.expressions.shape.TernaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.BigIntType;
@@ -41,9 +43,11 @@ public class Lead extends WindowFunction implements TernaryExpression, Explicitl
 
     static {
         List<FunctionSignature> signatures = Lists.newArrayList();
-        trivialTypes.forEach(t ->
-                signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE, t))
-        );
+        trivialTypes.forEach(t -> {
+            signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE, t));
+            signatures.add(FunctionSignature.ret(t).args(t, BigIntType.INSTANCE));
+            signatures.add(FunctionSignature.ret(t).args(t));
+        });
         SIGNATURES = ImmutableList.copyOf(signatures);
     }
 
@@ -53,21 +57,21 @@ public Lead(Expression child, Expression offset, Expression defaultValue) {
         super("lead", child, offset, defaultValue);
     }
 
-    private Lead(List<Expression> children) {
-        super("lead", children);
+    public Lead(Expression child, Expression offset) {
+        this(child, offset, new NullLiteral(child.getDataType()));
+    }
+
+    public Lead(Expression child) {
+        this(child, new BigIntLiteral(1L), new NullLiteral(child.getDataType()));
     }
 
     public Expression getOffset() {
-        if (children().size() <= 1) {
-            throw new AnalysisException("Not set offset of Lead(): " + this.toSql());
-        }
+        Preconditions.checkArgument(children.size() == 3);
         return child(1);
     }
 
     public Expression getDefaultValue() {
-        if (children.size() <= 2) {
-            throw new AnalysisException("Not set default value of Lead(): " + this.toSql());
-        }
+        Preconditions.checkArgument(children.size() == 3);
         return child(2);
     }
 
@@ -114,7 +118,13 @@ public List<FunctionSignature> getSignatures() {
     @Override
     public Lead withChildren(List<Expression> children) {
         Preconditions.checkArgument(children.size() >= 1 && children.size() <= 3);
-        return new Lead(children);
+        if (children.size() == 1) {
+            return new Lead(children.get(0));
+        } else if (children.size() == 2) {
+            return new Lead(children.get(0), children.get(1));
+        } else {
+            return new Lead(children.get(0), children.get(1), children.get(2));
+        }
     }
 
     @Override

diff --git a/regression-test/data/correctness_p0/test_lag_lead_window.out b/regression-test/data/correctness_p0/test_lag_lead_window.out
@@ -19,3 +19,33 @@ c	2022-09-06T00:00:02	2022-09-06T00:00:01
 b	2022-09-06T00:00:01	2022-09-06T00:00
 a	2022-09-06T00:00	2022-08-30T00:00
 
+-- !select_lag_1 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
+-- !select_lag_2 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
+-- !select_lag_3 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
+-- !select_lag_4 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
+-- !select_lag_5 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
+-- !select_lag_6 --
+a	2022-09-06T00:00	\N
+b	2022-09-06T00:00:01	\N
+c	2022-09-06T00:00:02	\N
+
diff --git a/regression-test/data/nereids_p0/sql_functions/window_functions/test_window_function.out b/regression-test/data/nereids_p0/sql_functions/window_functions/test_window_function.out
@@ -432,6 +432,78 @@ USA	Pete	Hello
 9223372036854775807	1
 9223372036854775807	2
 
+-- !lag_1 --
+\N	\N	\N
+-9223372036854775807	false	\N
+-9223372036854775807	true	false
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	\N
+11011902	true	false
+11011902	true	true
+11011903	false	\N
+11011905	false	\N
+11011920	true	\N
+11011920	true	true
+9223372036854775807	false	\N
+9223372036854775807	false	false
+
+-- !lag_2 --
+\N	\N	\N
+-9223372036854775807	false	\N
+-9223372036854775807	true	false
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	\N
+11011902	true	false
+11011902	true	true
+11011903	false	\N
+11011905	false	\N
+11011920	true	\N
+11011920	true	true
+9223372036854775807	false	\N
+9223372036854775807	false	false
+
+-- !lead_1 --
+\N	\N	\N
+-9223372036854775807	false	true
+-9223372036854775807	true	\N
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	true
+11011902	true	true
+11011902	true	\N
+11011903	false	\N
+11011905	false	\N
+11011920	true	true
+11011920	true	\N
+9223372036854775807	false	false
+9223372036854775807	false	\N
+
+-- !lead_2 --
+\N	\N	\N
+-9223372036854775807	false	true
+-9223372036854775807	true	\N
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	true
+11011902	true	true
+11011902	true	\N
+11011903	false	\N
+11011905	false	\N
+11011920	true	true
+11011920	true	\N
+9223372036854775807	false	false
+9223372036854775807	false	\N
+
 -- !window_error1 --
 \N	\N
 -9223372036854775807	true

diff --git a/regression-test/data/query_p0/sql_functions/window_functions/test_window_function.out b/regression-test/data/query_p0/sql_functions/window_functions/test_window_function.out
@@ -432,6 +432,78 @@ USA	Pete	Hello
 9223372036854775807	1
 9223372036854775807	2
 
+-- !lag_1 --
+\N	\N	\N
+-9223372036854775807	false	\N
+-9223372036854775807	true	false
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	\N
+11011902	true	false
+11011902	true	true
+11011903	false	\N
+11011905	false	\N
+11011920	true	\N
+11011920	true	true
+9223372036854775807	false	\N
+9223372036854775807	false	false
+
+-- !lag_2 --
+\N	\N	\N
+-9223372036854775807	false	\N
+-9223372036854775807	true	false
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	\N
+11011902	true	false
+11011902	true	true
+11011903	false	\N
+11011905	false	\N
+11011920	true	\N
+11011920	true	true
+9223372036854775807	false	\N
+9223372036854775807	false	false
+
+-- !lead_1 --
+\N	\N	\N
+-9223372036854775807	false	true
+-9223372036854775807	true	\N
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	true
+11011902	true	true
+11011902	true	\N
+11011903	false	\N
+11011905	false	\N
+11011920	true	true
+11011920	true	\N
+9223372036854775807	false	false
+9223372036854775807	false	\N
+
+-- !lead_2 --
+\N	\N	\N
+-9223372036854775807	false	true
+-9223372036854775807	true	\N
+-11011907	false	\N
+-11011903	true	\N
+123456	true	\N
+7210457	false	\N
+11011902	false	true
+11011902	true	true
+11011902	true	\N
+11011903	false	\N
+11011905	false	\N
+11011920	true	true
+11011920	true	\N
+9223372036854775807	false	false
+9223372036854775807	false	\N
+
 -- !window_error1 --
 \N	\N
 -9223372036854775807	true

diff --git a/regression-test/suites/correctness_p0/test_lag_lead_window.groovy b/regression-test/suites/correctness_p0/test_lag_lead_window.groovy
@@ -51,5 +51,14 @@ suite("test_lag_lead_window") {
     qt_select_default """ select id, create_time, lead(create_time, 1, '2022-09-06 00:00:00') over
                           (order by create_time desc) as "prev_time" from test1; """
     qt_select_default """ select id, create_time, lead(create_time, 1, date_sub('2022-09-06 00:00:00', interval 7 day)) over (order by create_time desc) as "prev_time" from test1; """
+
+    qt_select_lag_1 """ select id, create_time, lag(create_time) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+    qt_select_lag_2 """ select id, create_time, lag(create_time,1) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+    qt_select_lag_3 """ select id, create_time, lag(create_time,1,NULL) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+    qt_select_lag_4 """ select id, create_time, lead(create_time) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+    qt_select_lag_5 """ select id, create_time, lead(create_time,1) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+    qt_select_lag_6 """ select id, create_time, lead(create_time,1,NULL) over(partition by id) as "prev_time" from test1 order by 1 ,2 ; """
+
     sql """ DROP TABLE IF EXISTS test1 """
+
 }
diff --git a/regression-test/suites/nereids_p0/sql_functions/window_functions/test_window_function.groovy b/regression-test/suites/nereids_p0/sql_functions/window_functions/test_window_function.groovy
@@ -366,34 +366,27 @@ suite("test_window_function") {
                     rows between unbounded preceding and current row)
                     as wj from baseall order by ${k1}, wj"""
 
-    // test error
-    test {
-        sql("select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}) over (partition by ${k1} order by ${k3}) from baseall")
-        exception ""
-    }
+    qt_lag_1 "select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1},${k2}, lag(${k2}) over (partition by ${k1} order by ${k3},${k2}) from baseall order by 1,2;"
+
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
         exception ""
     }
-    test {
-        sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        exception ""
-    }
-    test {
-        sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}) over (partition by ${k1} order by ${k3}) from baseall"
-        exception ""
-    }
+
+    qt_lag_2 "select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, ${k2}, lag(${k2}, 1) over (partition by ${k1} order by ${k3},${k2}) from baseall order by 1,2;"
+
+    qt_lead_1 "select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, ${k2}, lead(${k2}) over (partition by ${k1} order by ${k3},${k2}) from baseall order by 1,2;"
+
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
         exception ""
     }
-    test {
-        sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        exception ""
-    }
-    qt_window_error1"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (partition by ${k1}) from baseall order by ${k1}"""
-    qt_window_error2"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (order by ${k3}) from baseall"""
-    qt_window_error3"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, max(${k2}) over (order by ${k3}) from baseall"""
+
+    qt_lead_2 "select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, ${k2}, lead(${k2}, 1) over (partition by ${k1} order by ${k3},${k2}) from baseall order by 1,2;"
+
+    qt_window_error1"""select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, first_value(${k2}) over (partition by ${k1}) from baseall order by ${k1}"""
+    qt_window_error2"""select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, first_value(${k2}) over (order by ${k3}) from baseall"""
+    qt_window_error3"""select /*+SET_VAR(parallel_pipeline_task_num=1) */ ${k1}, max(${k2}) over (order by ${k3}) from baseall"""
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows
             between current row and unbounded preceding) as wj