Verify types in IR

core/trino-main/src/main/java/io/trino/sql/ir/Arithmetic.java

@@ -20,6 +20,7 @@
 
 import java.util.List;
 
+import static io.trino.sql.ir.IrUtils.validateType;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize
@@ -48,10 +49,9 @@ public String getValue()
 
     public Arithmetic
     {
-        requireNonNull(function, "function is null");
         requireNonNull(operator, "operator is null");
-        requireNonNull(left, "left is null");
-        requireNonNull(right, "right is null");
+        validateType(function.getSignature().getArgumentType(0), left);
+        validateType(function.getSignature().getArgumentType(1), right);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Between.java

@@ -21,7 +21,7 @@
 import java.util.List;
 
 import static io.trino.spi.type.BooleanType.BOOLEAN;
-import static java.util.Objects.requireNonNull;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
 public record Between(Expression value, Expression min, Expression max)
@@ -30,9 +30,8 @@ public record Between(Expression value, Expression min, Expression max)
     @JsonCreator
     public Between
     {
-        requireNonNull(value, "value is null");
-        requireNonNull(min, "min is null");
-        requireNonNull(max, "max is null");
+        validateType(value.type(), min);
+        validateType(value.type(), max);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Bind.java

@@ -24,7 +24,7 @@
 import java.util.stream.Collectors;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static java.util.Objects.requireNonNull;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 /**
  * Bind(value, targetFunction)
@@ -54,8 +54,10 @@ public record Bind(List<Expression> values, Lambda function)
 {
     public Bind
     {
-        requireNonNull(function, "function is null");
         values = ImmutableList.copyOf(values);
+        for (int i = 0; i < values.size(); i++) {
+            validateType(function.arguments().get(i).getType(), values.get(i));
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Call.java

@@ -21,16 +21,21 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
-import static java.util.Objects.requireNonNull;
+import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
 public record Call(ResolvedFunction function, List<Expression> arguments)
         implements Expression
 {
     public Call
     {
-        requireNonNull(function, "function is null");
         arguments = ImmutableList.copyOf(arguments);
+
+        checkArgument(function.getSignature().getArgumentTypes().size() == arguments.size(), "Expected %s arguments, found: %s", function.getSignature().getArgumentTypes().size(), arguments.size());
+        for (int i = 0; i < arguments.size(); i++) {
+            validateType(function.getSignature().getArgumentType(i), arguments.get(i));
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Case.java

@@ -21,6 +21,8 @@
 import java.util.Optional;
 import java.util.stream.Collectors;
 
+import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.sql.ir.IrUtils.validateType;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize
@@ -31,6 +33,18 @@ public record Case(List<WhenClause> whenClauses, Optional<Expression> defaultVal
     {
         whenClauses = ImmutableList.copyOf(whenClauses);
         requireNonNull(defaultValue, "defaultValue is null");
+
+        for (WhenClause clause : whenClauses) {
+            validateType(BOOLEAN, clause.getOperand());
+        }
+
+        for (int i = 1; i < whenClauses.size(); i++) {
+            validateType(whenClauses.getFirst().getResult().type(), whenClauses.get(i).getResult());
+        }
+
+        if (defaultValue.isPresent()) {
+            validateType(whenClauses.getFirst().getResult().type(), defaultValue.get());
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Coalesce.java

@@ -20,6 +20,7 @@
 import java.util.List;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
 public record Coalesce(List<Expression> operands)
@@ -43,6 +44,10 @@ public Type type()
     {
         checkArgument(operands.size() >= 2, "must have at least two operands");
         operands = ImmutableList.copyOf(operands);
+
+        for (int i = 1; i < operands.size(); i++) {
+            validateType(operands.getFirst().type(), operands.get(i));
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Comparison.java

@@ -20,6 +20,7 @@
 import java.util.List;
 
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.sql.ir.IrUtils.validateType;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize
@@ -87,8 +88,7 @@ public Operator negate()
     public Comparison
     {
         requireNonNull(operator, "operator is null");
-        requireNonNull(left, "left is null");
-        requireNonNull(right, "right is null");
+        validateType(left.type(), right);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/ExpressionTreeRewriter.java

@@ -149,7 +149,7 @@ protected Expression visitSubscript(Subscript node, Context<C> context)
             Expression index = rewrite(node.index(), context.get());
 
             if (base != node.base() || index != node.index()) {
-                return new Subscript(node.type(), base, index);
+                return new Subscript(base, index);
             }
 
             return node;

core/trino-main/src/main/java/io/trino/sql/ir/In.java

@@ -20,6 +20,7 @@
 import java.util.List;
 
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
 public record In(Expression value, List<Expression> valueList)
@@ -28,6 +29,10 @@ public record In(Expression value, List<Expression> valueList)
     public In
     {
         valueList = ImmutableList.copyOf(valueList);
+
+        for (Expression item : valueList) {
+            validateType(value.type(), item);
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/IrUtils.java

@@ -18,6 +18,7 @@
 import com.google.common.graph.SuccessorsFunction;
 import com.google.common.graph.Traverser;
 import io.trino.metadata.Metadata;
+import io.trino.spi.type.Type;
 import io.trino.sql.planner.DeterminismEvaluator;
 import io.trino.sql.planner.Symbol;
 import io.trino.sql.planner.SymbolsExtractor;
@@ -31,6 +32,7 @@
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Predicates.not;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Streams.stream;
@@ -43,6 +45,11 @@ public final class IrUtils
 {
     private IrUtils() {}
 
+    static void validateType(Type expected, Expression expression)
+    {
+        checkArgument(expected.equals(expression.type()), "Expected '%s' type but found '%s' for expression: %s", expected, expression.type(), expression);
+    }
+
     public static List<Expression> extractConjuncts(Expression expression)
     {
         return extractPredicates(Logical.Operator.AND, expression);

core/trino-main/src/main/java/io/trino/sql/ir/Logical.java

@@ -22,6 +22,7 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.sql.ir.IrUtils.validateType;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize
@@ -48,6 +49,11 @@ public Operator flip()
     {
         requireNonNull(operator, "operator is null");
         checkArgument(terms.size() >= 2, "Expected at least 2 terms");
+
+        for (Expression term : terms) {
+            validateType(BOOLEAN, term);
+        }
+
         terms = ImmutableList.copyOf(terms);
     }
 

core/trino-main/src/main/java/io/trino/sql/ir/Not.java

@@ -20,15 +20,15 @@
 import java.util.List;
 
 import static io.trino.spi.type.BooleanType.BOOLEAN;
-import static java.util.Objects.requireNonNull;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
 public record Not(Expression value)
         implements Expression
 {
     public Not
     {
-        requireNonNull(value, "value is null");
+        validateType(BOOLEAN, value);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/NullIf.java

@@ -32,6 +32,8 @@ public record NullIf(Expression first, Expression second)
     {
         requireNonNull(first, "first is null");
         requireNonNull(second, "second is null");
+
+        // TODO: verify that first and second can be coerced to the same type
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Subscript.java

@@ -15,26 +15,35 @@
 
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import com.google.common.collect.ImmutableList;
+import io.trino.spi.type.RowType;
 import io.trino.spi.type.Type;
 
 import java.util.List;
 
-import static java.util.Objects.requireNonNull;
+import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.spi.type.IntegerType.INTEGER;
+import static io.trino.sql.ir.IrUtils.validateType;
 
 @JsonSerialize
-public record Subscript(Type type, Expression base, Expression index)
+public record Subscript(Expression base, Expression index)
         implements Expression
 {
     public Subscript
     {
-        requireNonNull(base, "base is null");
-        requireNonNull(index, "index is null");
+        if (!(base.type() instanceof RowType rowType)) {
+            throw new IllegalArgumentException("Expected 'row' type but found '%s' for expression: %s".formatted(base.type(), base));
+        }
+
+        validateType(INTEGER, index);
+        int field = (int) (long) ((Constant) index).value() - 1;
+        checkArgument(field < rowType.getFields().size(), "Expected 'row' type to have at least %s fields, but has: %s", field + 1, rowType.getFields().size());
     }
 
     @Override
     public Type type()
     {
-        return type;
+        int field = (int) (long) ((Constant) index).value() - 1;
+        return ((RowType) base.type()).getFields().get(field).getType();
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/ir/Switch.java

@@ -20,6 +20,7 @@
 import java.util.List;
 import java.util.Optional;
 
+import static io.trino.sql.ir.IrUtils.validateType;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize
@@ -30,7 +31,18 @@ public record Switch(Expression operand, List<WhenClause> whenClauses, Optional<
     {
         requireNonNull(operand, "operand is null");
         whenClauses = ImmutableList.copyOf(whenClauses);
-        requireNonNull(defaultValue, "defaultValue is null");
+
+        for (WhenClause clause : whenClauses) {
+            validateType(operand.type(), clause.getOperand());
+        }
+
+        for (int i = 1; i < whenClauses.size(); i++) {
+            validateType(whenClauses.getFirst().getResult().type(), whenClauses.get(i).getResult());
+        }
+
+        if (defaultValue.isPresent()) {
+            validateType(whenClauses.getFirst().getResult().type(), defaultValue.get());
+        }
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/planner/ConnectorExpressionTranslator.java

@@ -208,7 +208,7 @@ public Optional<Expression> translate(ConnectorExpression expression)
 
             if (expression instanceof FieldDereference dereference) {
                 return translate(dereference.getTarget())
-                        .map(base -> new Subscript(dereference.getType(), base, new Constant(INTEGER, (long) (dereference.getField() + 1))));
+                        .map(base -> new Subscript(base, new Constant(INTEGER, (long) (dereference.getField() + 1))));
             }
 
             if (expression instanceof io.trino.spi.expression.Call call) {

core/trino-main/src/main/java/io/trino/sql/planner/IrExpressionInterpreter.java

@@ -898,7 +898,7 @@ protected Object visitSubscript(Subscript node, Object context)
             }
 
             if (hasUnresolvedValue(base, index)) {
-                return new Subscript(node.type(), toExpression(base, node.base().type()), toExpression(index, node.index().type()));
+                return new Subscript(toExpression(base, node.base().type()), toExpression(index, node.index().type()));
             }
 
             // Subscript on Row hasn't got a dedicated operator. It is interpreted by hand.

core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java

@@ -889,7 +889,7 @@ public MergeWriterNode plan(Merge merge)
                 subPlanProject,
                 Assignments.builder()
                         .putIdentities(subPlanProject.getOutputSymbols())
-                        .put(caseNumberSymbol, new Subscript(INTEGER, mergeRowSymbol.toSymbolReference(), new Constant(INTEGER, (long) mergeAnalysis.getMergeRowType().getFields().size())))
+                        .put(caseNumberSymbol, new Subscript(mergeRowSymbol.toSymbolReference(), new Constant(INTEGER, (long) mergeAnalysis.getMergeRowType().getFields().size())))
                         .build());
 
         // Mark distinct combinations of the unique_id value and the case_number

core/trino-main/src/main/java/io/trino/sql/planner/TranslationMap.java

@@ -671,7 +671,6 @@ private io.trino.sql.ir.Expression translate(DereferenceExpression expression)
         checkState(index >= 0, "could not find field name: %s", fieldName);
 
         return new Subscript(
-                rowType.getFields().get(index).getType(),
                 translateExpression(expression.getBase()),
                 new Constant(INTEGER, (long) (index + 1)));
     }
@@ -959,7 +958,6 @@ private io.trino.sql.ir.Expression translate(SubscriptExpression node)
             io.trino.sql.ir.Expression rewrittenBase = translateExpression(node.getBase());
             LongLiteral index = (LongLiteral) node.getIndex();
             return new Subscript(
-                    analysis.getType(node),
                     rewrittenBase, new Constant(INTEGER, index.getParsedValue()));
         }
 

core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/UnwrapSingleColumnRowInApply.java

@@ -148,8 +148,8 @@ private Optional<Unwrapping> unwrapSingleColumnRow(Context context, Expression v
                 Symbol valueSymbol = context.getSymbolAllocator().newSymbol("input", elementType);
                 Symbol listSymbol = context.getSymbolAllocator().newSymbol("subquery", elementType);
 
-                Assignment inputAssignment = new Assignment(valueSymbol, new Subscript(elementType, value, new Constant(INTEGER, 1L)));
-                Assignment nestedPlanAssignment = new Assignment(listSymbol, new Subscript(elementType, list, new Constant(INTEGER, 1L)));
+                Assignment inputAssignment = new Assignment(valueSymbol, new Subscript(value, new Constant(INTEGER, 1L)));
+                Assignment nestedPlanAssignment = new Assignment(listSymbol, new Subscript(list, new Constant(INTEGER, 1L)));
                 ApplyNode.SetExpression comparison = function.apply(valueSymbol, listSymbol);
 
                 return Optional.of(new Unwrapping(comparison, inputAssignment, nestedPlanAssignment));