Support BETWEEN in Delta check constraints

plugin/trino-delta-lake/src/main/antlr4/io/trino/plugin/deltalake/expression/SparkExpressionBase.g4

@@ -36,6 +36,7 @@ booleanExpression
 // workaround for https://github.com/antlr/antlr4/issues/780
 predicate[ParserRuleContext value]
     : comparisonOperator right=valueExpression          #comparison
+    | NOT? BETWEEN lower=valueExpression AND upper=valueExpression        #between
     ;
 
 valueExpression
@@ -48,6 +49,7 @@ valueExpression
 primaryExpression
     : number                                            #numericLiteral
     | booleanValue                                      #booleanLiteral
+    | NULL                                              #nullLiteral
     | string                                            #stringLiteral
     | identifier                                        #columnReference
     ;
@@ -77,16 +79,19 @@ number
     ;
 
 AND: 'AND';
+BETWEEN: 'BETWEEN';
 OR: 'OR';
 FALSE: 'FALSE';
 TRUE: 'TRUE';
+NULL: 'NULL';
 
 EQ: '=';
 NEQ: '<>' | '!=';
 LT: '<';
 LTE: '<=';
 GT: '>';
 GTE: '>=';
+NOT: 'NOT';
 
 PLUS: '+';
 MINUS: '-';

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/expression/BetweenPredicate.java

@@ -0,0 +1,113 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.deltalake.expression;
+
+import java.util.Objects;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+public class BetweenPredicate
+        extends SparkExpression
+{
+    public enum Operator
+    {
+        BETWEEN("BETWEEN"),
+        NOT_BETWEEN("NOT BETWEEN");
+
+        private final String value;
+
+        Operator(String value)
+        {
+            this.value = value;
+        }
+
+        public String getValue()
+        {
+            return value;
+        }
+    }
+
+    private final Operator operator;
+    private final SparkExpression value;
+    private final SparkExpression min;
+    private final SparkExpression max;
+
+    public BetweenPredicate(Operator operator, SparkExpression value, SparkExpression min, SparkExpression max)
+    {
+        this.operator = requireNonNull(operator, "operator is null");
+        this.value = requireNonNull(value, "value is null");
+        this.min = requireNonNull(min, "min is null");
+        this.max = requireNonNull(max, "max is null");
+    }
+
+    public Operator getOperator()
+    {
+        return operator;
+    }
+
+    public SparkExpression getValue()
+    {
+        return value;
+    }
+
+    public SparkExpression getMin()
+    {
+        return min;
+    }
+
+    public SparkExpression getMax()
+    {
+        return max;
+    }
+
+    @Override
+    <R, C> R accept(SparkExpressionTreeVisitor<R, C> visitor, C context)
+    {
+        return visitor.visitBetweenExpression(this, context);
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        BetweenPredicate that = (BetweenPredicate) o;
+        return operator == that.operator &&
+                Objects.equals(value, that.value) &&
+                Objects.equals(min, that.min) &&
+                Objects.equals(max, that.max);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(operator, value, min, max);
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("operator", operator)
+                .add("value", value)
+                .add("min", min)
+                .add("max", max)
+                .toString();
+    }
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/expression/NullLiteral.java

@@ -0,0 +1,33 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.deltalake.expression;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+
+public class NullLiteral
+        extends Literal
+{
+    @Override
+    public <R, C> R accept(SparkExpressionTreeVisitor<R, C> visitor, C context)
+    {
+        return visitor.visitNullLiteral(this, context);
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .toString();
+    }
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/expression/SparkExpressionBuilder.java

@@ -106,6 +106,20 @@ public Object visitOr(SparkExpressionBaseParser.OrContext context)
                 visit(context.right, SparkExpression.class));
     }
 
+    @Override
+    public SparkExpression visitBetween(SparkExpressionBaseParser.BetweenContext context)
+    {
+        BetweenPredicate.Operator operator = BetweenPredicate.Operator.BETWEEN;
+        if (context.NOT() != null) {
+            operator = BetweenPredicate.Operator.NOT_BETWEEN;
+        }
+        return new BetweenPredicate(
+                operator,
+                visit(context.value, SparkExpression.class),
+                visit(context.lower, SparkExpression.class),
+                visit(context.upper, SparkExpression.class));
+    }
+
     @Override
     public Object visitColumnReference(SparkExpressionBaseParser.ColumnReferenceContext context)
     {
@@ -150,6 +164,12 @@ public Object visitUnicodeStringLiteral(SparkExpressionBaseParser.UnicodeStringL
         return new StringLiteral(decodeUnicodeLiteral(context));
     }
 
+    @Override
+    public SparkExpression visitNullLiteral(SparkExpressionBaseParser.NullLiteralContext context)
+    {
+        return new NullLiteral();
+    }
+
     private static String decodeUnicodeLiteral(SparkExpressionBaseParser.UnicodeStringLiteralContext context)
     {
         String rawContent = unquote(context.getText());

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/expression/SparkExpressionConverter.java

@@ -58,6 +58,16 @@ protected String visitLogicalExpression(LogicalExpression node, Void context)
             return "(%s %s %s)".formatted(process(node.getLeft(), context), node.getOperator().toString(), process(node.getRight(), context));
         }
 
+        @Override
+        protected String visitBetweenExpression(BetweenPredicate node, Void context)
+        {
+            return "(%s %s %s AND %s)".formatted(
+                    process(node.getValue(), context),
+                    node.getOperator().getValue(),
+                    process(node.getMin(), context),
+                    process(node.getMax(), context));
+        }
+
         @Override
         protected String visitIdentifier(Identifier node, Void context)
         {
@@ -81,5 +91,11 @@ protected String visitStringLiteral(StringLiteral node, Void context)
         {
             return "'" + node.getValue().replace("'", "''") + "'";
         }
+
+        @Override
+        protected String visitNullLiteral(NullLiteral node, Void context)
+        {
+            return "NULL";
+        }
     }
 }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/expression/SparkExpressionTreeVisitor.java

@@ -29,6 +29,11 @@ protected R visitComparisonExpression(ComparisonExpression node, C context)
         return visitExpression(node, context);
     }
 
+    protected R visitBetweenExpression(BetweenPredicate node, C context)
+    {
+        return visitExpression(node, context);
+    }
+
     protected R visitLogicalExpression(LogicalExpression node, C context)
     {
         return visitExpression(node, context);
@@ -63,4 +68,9 @@ protected R visitStringLiteral(StringLiteral node, C context)
     {
         return visitLiteral(node, context);
     }
+
+    protected R visitNullLiteral(NullLiteral node, C context)
+    {
+        return visitLiteral(node, context);
+    }
 }

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/expression/TestSparkExpressions.java

@@ -138,6 +138,18 @@ public void testArithmeticBinary()
         assertExpressionTranslates("a = b ^ 1", "(\"a\" = (bitwise_xor(\"b\", 1)))");
     }
 
+    @Test
+    public void testBetween()
+    {
+        assertExpressionTranslates("a BETWEEN 1 AND 10", "(\"a\" BETWEEN 1 AND 10)");
+        assertExpressionTranslates("a NOT BETWEEN 1 AND 10", "(\"a\" NOT BETWEEN 1 AND 10)");
+        assertExpressionTranslates("a BETWEEN NULL AND 10", "(\"a\" BETWEEN NULL AND 10)");
+        assertExpressionTranslates("a BETWEEN 1 AND NULL", "(\"a\" BETWEEN 1 AND NULL)");
+        assertExpressionTranslates("a NOT BETWEEN NULL AND NULL", "(\"a\" NOT BETWEEN NULL AND NULL)");
+        assertExpressionTranslates("a not between null and null", "(\"a\" NOT BETWEEN NULL AND NULL)");
+        assertExpressionTranslates("a BETWEEN b AND c", "(\"a\" BETWEEN \"b\" AND \"c\")");
+    }
+
     @Test
     public void testInvalidNotBoolean()
     {
@@ -165,7 +177,6 @@ public void testUnsupportedOperator()
         assertParseFailure("a == 1");
         assertParseFailure("a = b::INTEGER");
         assertParseFailure("a = json_column:root");
-        assertParseFailure("a BETWEEN 1 AND 10");
         assertParseFailure("a IS NULL");
         assertParseFailure("a IS DISTINCT FROM b");
         assertParseFailure("a IS true");

testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeCheckConstraintCompatibility.java

@@ -89,6 +89,12 @@ public static Object[][] checkConstraints()
                 {"a INT, b INT", "a = b % 2", "1, 5", row(1, 5), "1, 6"},
                 {"a INT, b INT", "a = b & 5", "1, 3", row(1, 3), "1, 4"},
                 {"a INT, b INT", "a = b ^ 5", "6, 3", row(6, 3), "6, 4"},
+                // Between
+                {"a INT", "a BETWEEN 1 AND 10", "1", row(1), "0"},
+                {"a INT", "a BETWEEN 1 AND 10", "10", row(10), "11"},
+                {"a INT", "a NOT BETWEEN 1 AND 10", "0", row(0), "1"},
+                {"a INT", "a NOT BETWEEN 1 AND 10", "11", row(11), "10"},
+                {"a INT, b INT, c INT", "a BETWEEN b AND c", "5, 1, 10", row(5, 1, 10), "11, 1, 10"},
                 // Supported types
                 {"a INT", "a < 100", "1", row(1), "100"},
                 {"a STRING", "a = 'valid'", "'valid'", row("valid"), "'invalid'"},