[CALCITE-6382] Type inference for SqlLeadLagAggFunction is incorrect

Signed-off-by: Mihai Budiu <[email protected]>

core/src/main/java/org/apache/calcite/sql/fun/SqlLeadLagAggFunction.java

@@ -22,13 +22,17 @@
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.type.OperandTypes;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeTransform;
 import org.apache.calcite.sql.type.SqlTypeTransforms;
 import org.apache.calcite.util.Optionality;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static com.google.common.base.Preconditions.checkArgument;
 
 /**
@@ -40,10 +44,29 @@ public class SqlLeadLagAggFunction extends SqlAggFunction {
       OperandTypes.ANY
           .or(OperandTypes.ANY_NUMERIC)
           .or(OperandTypes.ANY_NUMERIC_ANY
+              // "same" only checks if two types are comparable
               .and(OperandTypes.same(3, 0, 2)));
 
+  // A version of least restrictive which only looks at operands 0 and 2,
+  // if the latter exists.
+  private static final SqlReturnTypeInference ARG03 = opBinding -> {
+    List<RelDataType> toCheck = new ArrayList<>();
+    toCheck.add(opBinding.getOperandType(0));
+    if (opBinding.getOperandCount() >= 3) {
+      RelDataType op2 = opBinding.getOperandType(2);
+      toCheck.add(op2);
+    }
+    // If any operand is in the CHAR type family, return VARCHAR.
+    for (RelDataType type : toCheck) {
+      if (type.getFamily() == SqlTypeFamily.CHARACTER) {
+        return opBinding.getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
+      }
+    }
+    return opBinding.getTypeFactory().leastRestrictive(toCheck);
+  };
+
   private static final SqlReturnTypeInference RETURN_TYPE =
-      ReturnTypes.ARG0.andThen(SqlLeadLagAggFunction::transformType);
+      ARG03.andThen(SqlLeadLagAggFunction::transformType);
 
   public SqlLeadLagAggFunction(SqlKind kind) {
     super(kind.name(),

core/src/test/java/org/apache/calcite/test/JdbcTest.java

@@ -4631,9 +4631,6 @@ private void startOfGroupStep3(String startOfGroup) {
             "RN=8; L=2");
   }
 
-  /**
-   * Tests expression in offset value of LAG function.
-   */
   @Disabled("Have no idea how to validate that expression is constant")
   @Test void testNtileConstantArgs() {
     CalciteAssert.that()
@@ -4652,9 +4649,6 @@ private void startOfGroupStep3(String startOfGroup) {
             "RN=8; L=2");
   }
 
-  /**
-   * Tests expression in offset value of LAG function.
-   */
   @Test void testNtileNegativeArg() {
     CalciteAssert.that()
         .query("select rn, ntile(-1) over (order by rn) l\n"
@@ -4663,9 +4657,6 @@ private void startOfGroupStep3(String startOfGroup) {
             "Argument to function 'NTILE' must be a positive integer literal");
   }
 
-  /**
-   * Tests expression in offset value of LAG function.
-   */
   @Test void testNtileDecimalArg() {
     CalciteAssert.that()
         .query("select rn, ntile(3.141592653) over (order by rn) l\n"

core/src/test/java/org/apache/calcite/test/SqlValidatorTest.java

@@ -3109,6 +3109,15 @@ void testWinPartClause() {
         .fails(RANK_REQUIRES_ORDER_BY);
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6382">[CALCITE-6382]
+   * Type inference for SqlLeadLagAggFunction is incorrect</a>. */
+  @Test void testWindowLagInference() {
+    sql("select lead(sal, 4, 0.5) over (w)\n"
+        + " from emp window w as (order by empno)")
+        .type("RecordType(DECIMAL(11, 1) NOT NULL EXPR$0) NOT NULL");
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-883">[CALCITE-883]
    * Give error if the aggregate function don't support null treatment</a>. */