Merge unwrap cast in comparison optimizer rules

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

@@ -235,7 +235,6 @@
 import io.trino.sql.planner.iterative.rule.UnwrapCastInComparison;
 import io.trino.sql.planner.iterative.rule.UnwrapRowSubscript;
 import io.trino.sql.planner.iterative.rule.UnwrapSingleColumnRowInApply;
-import io.trino.sql.planner.iterative.rule.UnwrapTimestampToDateCastInComparison;
 import io.trino.sql.planner.iterative.rule.UseNonPartitionedJoinLookupSource;
 import io.trino.sql.planner.optimizations.AddExchanges;
 import io.trino.sql.planner.optimizations.AddLocalExchanges;
@@ -365,7 +364,6 @@ public PlanOptimizers(
                 .addAll(new SimplifyExpressions(plannerContext, typeAnalyzer).rules())
                 .addAll(new UnwrapRowSubscript().rules())
                 .addAll(new PushCastIntoRow().rules())
-                .addAll(new UnwrapTimestampToDateCastInComparison(plannerContext, typeAnalyzer).rules())
                 .addAll(new UnwrapCastInComparison(plannerContext, typeAnalyzer).rules())
                 .addAll(new RemoveDuplicateConditions(metadata).rules())
                 .addAll(new CanonicalizeExpressions(plannerContext, typeAnalyzer).rules())

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

@@ -22,6 +22,7 @@
 import io.trino.spi.TrinoException;
 import io.trino.spi.function.InvocationConvention;
 import io.trino.spi.type.CharType;
+import io.trino.spi.type.DateType;
 import io.trino.spi.type.DecimalType;
 import io.trino.spi.type.DoubleType;
 import io.trino.spi.type.LongTimestampWithTimeZone;
@@ -65,6 +66,7 @@
 import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone;
 import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc;
 import static io.trino.spi.type.DateTimeEncoding.unpackZoneKey;
+import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.IntegerType.INTEGER;
 import static io.trino.spi.type.RealType.REAL;
@@ -205,6 +207,10 @@ private Expression unwrapCast(ComparisonExpression expression)
             Type sourceType = typeAnalyzer.getType(session, types, cast.getExpression());
             Type targetType = typeAnalyzer.getType(session, types, expression.getRight());
 
+            if (sourceType instanceof TimestampType && targetType == DATE) {
+                return unwrapTimestampToDateCast(session, (TimestampType) sourceType, (DateType) targetType, operator, cast.getExpression(), (long) right).orElse(expression);
+            }
+
             if (targetType instanceof TimestampWithTimeZoneType) {
                 // Note: two TIMESTAMP WITH TIME ZONE values differing in zone only (same instant) are considered equal.
                 right = withTimeZone(((TimestampWithTimeZoneType) targetType), right, session.getTimeZoneKey());
@@ -404,6 +410,48 @@ private Expression unwrapCast(ComparisonExpression expression)
             return new ComparisonExpression(operator, cast.getExpression(), literalEncoder.toExpression(session, literalInSourceType, sourceType));
         }
 
+        private Optional<Expression> unwrapTimestampToDateCast(Session session, TimestampType sourceType, DateType targetType, ComparisonExpression.Operator operator, Expression timestampExpression, long date)
+        {
+            ResolvedFunction targetToSource;
+            try {
+                targetToSource = plannerContext.getMetadata().getCoercion(session, targetType, sourceType);
+            }
+            catch (OperatorNotFoundException e) {
+                throw new TrinoException(GENERIC_INTERNAL_ERROR, e);
+            }
+
+            Expression dateTimestamp = literalEncoder.toExpression(session, coerce(date, targetToSource), sourceType);
+            Expression nextDateTimestamp = literalEncoder.toExpression(session, coerce(date + 1, targetToSource), sourceType);
+
+            switch (operator) {
+                case EQUAL:
+                    return Optional.of(
+                            and(
+                                    new ComparisonExpression(GREATER_THAN_OR_EQUAL, timestampExpression, dateTimestamp),
+                                    new ComparisonExpression(LESS_THAN, timestampExpression, nextDateTimestamp)));
+                case NOT_EQUAL:
+                    return Optional.of(
+                            or(
+                                    new ComparisonExpression(LESS_THAN, timestampExpression, dateTimestamp),
+                                    new ComparisonExpression(GREATER_THAN_OR_EQUAL, timestampExpression, nextDateTimestamp)));
+                case LESS_THAN:
+                    return Optional.of(new ComparisonExpression(LESS_THAN, timestampExpression, dateTimestamp));
+                case LESS_THAN_OR_EQUAL:
+                    return Optional.of(new ComparisonExpression(LESS_THAN, timestampExpression, nextDateTimestamp));
+                case GREATER_THAN:
+                    return Optional.of(new ComparisonExpression(GREATER_THAN_OR_EQUAL, timestampExpression, nextDateTimestamp));
+                case GREATER_THAN_OR_EQUAL:
+                    return Optional.of(new ComparisonExpression(GREATER_THAN_OR_EQUAL, timestampExpression, dateTimestamp));
+                case IS_DISTINCT_FROM:
+                    return Optional.of(
+                            or(
+                                    new IsNullPredicate(timestampExpression),
+                                    new ComparisonExpression(LESS_THAN, timestampExpression, dateTimestamp),
+                                    new ComparisonExpression(GREATER_THAN_OR_EQUAL, timestampExpression, nextDateTimestamp)));
+            }
+            throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unsupported operator: " + operator);
+        }
+
         private boolean hasInjectiveImplicitCoercion(Type source, Type target, Object value)
         {
             if ((source.equals(BIGINT) && target.equals(DOUBLE)) ||