Allow complex expression rewrites with double and real constant

Author: Praveen2112

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteExactNumericConstant.java

@@ -27,15 +27,18 @@
 import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.constant;
 import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
 import static io.trino.spi.type.BigintType.BIGINT;
+import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.IntegerType.INTEGER;
+import static io.trino.spi.type.RealType.REAL;
 import static io.trino.spi.type.SmallintType.SMALLINT;
 import static io.trino.spi.type.TinyintType.TINYINT;
 
 public class RewriteExactNumericConstant
         implements ConnectorExpressionRule<Constant, ParameterizedExpression>
 {
     private static final Pattern<Constant> PATTERN = constant().with(type().matching(type ->
-            type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT || type instanceof DecimalType));
+            type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT
+                    || type == REAL || type == DOUBLE || type instanceof DecimalType));
 
     @Override
     public Pattern<Constant> getPattern()
@@ -52,7 +55,7 @@ public Optional<ParameterizedExpression> rewrite(Constant constant, Captures cap
             // TODO we could handle NULL values too
             return Optional.empty();
         }
-        if (type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT || type instanceof DecimalType) {
+        if (type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT || type == REAL || type == DOUBLE || type instanceof DecimalType) {
             return Optional.of(new ParameterizedExpression("?", ImmutableList.of(new QueryParameter(type, Optional.of(value)))));
         }
 

plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlClient.java

@@ -68,9 +68,11 @@
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
 import static io.trino.spi.type.DoubleType.DOUBLE;
+import static io.trino.spi.type.RealType.REAL;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.spi.type.VarcharType.createVarcharType;
 import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER;
+import static java.lang.Float.floatToIntBits;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -93,6 +95,13 @@ public class TestPostgreSqlClient
                     .setJdbcTypeHandle(new JdbcTypeHandle(Types.DOUBLE, Optional.of("double"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()))
                     .build();
 
+    private static final JdbcColumnHandle REAL_COLUMN =
+            JdbcColumnHandle.builder()
+                    .setColumnName("c_real")
+                    .setColumnType(REAL)
+                    .setJdbcTypeHandle(new JdbcTypeHandle(Types.REAL, Optional.of("real"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()))
+                    .build();
+
     private static final JdbcColumnHandle VARCHAR_COLUMN =
             JdbcColumnHandle.builder()
                     .setColumnName("c_varchar")
@@ -235,15 +244,18 @@ public void testConvertOr()
                                         Logical.Operator.OR,
                                         List.of(
                                                 new Comparison(Comparison.Operator.EQUAL, new Reference(BIGINT, "c_bigint_symbol"), new Constant(BIGINT, 42L)),
-                                                new Comparison(Comparison.Operator.EQUAL, new Reference(BIGINT, "c_bigint_symbol_2"), new Constant(BIGINT, 415L))))),
+                                                new Comparison(Comparison.Operator.EQUAL, new Reference(REAL, "c_real_symbol"), new Constant(REAL, (long) floatToIntBits(3.14f))),
+                                                new Comparison(Comparison.Operator.EQUAL, new Reference(DOUBLE, "c_double_symbol"), new Constant(DOUBLE, 3.14))))),
                         Map.of(
                                 "c_bigint_symbol", BIGINT_COLUMN,
-                                "c_bigint_symbol_2", BIGINT_COLUMN))
+                                "c_real_symbol", REAL_COLUMN,
+                                "c_double_symbol", DOUBLE_COLUMN))
                 .orElseThrow();
-        assertThat(converted.expression()).isEqualTo("((\"c_bigint\") = (?)) OR ((\"c_bigint\") = (?))");
+        assertThat(converted.expression()).isEqualTo("((\"c_bigint\") = (?)) OR ((\"c_real\") = (?)) OR ((\"c_double\") = (?))");
         assertThat(converted.parameters()).isEqualTo(List.of(
                 new QueryParameter(BIGINT, Optional.of(42L)),
-                new QueryParameter(BIGINT, Optional.of(415L))));
+                new QueryParameter(REAL, Optional.of((long) floatToIntBits(3.14f))),
+                new QueryParameter(DOUBLE, Optional.of(3.14))));
     }
 
     @Test
@@ -279,8 +291,16 @@ public void testConvertComparison()
             Optional<ParameterizedExpression> converted = JDBC_CLIENT.convertPredicate(
                     SESSION,
                     translateToConnectorExpression(
-                            new Comparison(operator, new Reference(BIGINT, "c_bigint_symbol"), new Constant(BIGINT, 42L))),
-                    Map.of("c_bigint_symbol", BIGINT_COLUMN));
+                            new Logical(
+                                    Logical.Operator.OR,
+                            List.of(
+                                    new Comparison(operator, new Reference(BIGINT, "c_bigint_symbol"), new Constant(BIGINT, 42L)),
+                                    new Comparison(operator, new Reference(REAL, "c_real_symbol"), new Constant(REAL, (long) floatToIntBits(3.14f))),
+                                    new Comparison(operator, new Reference(DOUBLE, "c_double_symbol"), new Constant(DOUBLE, 3.14))))),
+                    Map.of(
+                            "c_bigint_symbol", BIGINT_COLUMN,
+                            "c_real_symbol", REAL_COLUMN,
+                            "c_double_symbol", DOUBLE_COLUMN));
 
             switch (operator) {
                 case EQUAL:
@@ -290,8 +310,11 @@ public void testConvertComparison()
                 case GREATER_THAN:
                 case GREATER_THAN_OR_EQUAL:
                     assertThat(converted).isPresent();
-                    assertThat(converted.get().expression()).isEqualTo(format("(\"c_bigint\") %s (?)", operator.getValue()));
-                    assertThat(converted.get().parameters()).isEqualTo(List.of(new QueryParameter(BIGINT, Optional.of(42L))));
+                    assertThat(converted.get().expression()).isEqualTo(format("((\"c_bigint\") %1$s (?)) OR ((\"c_real\") %1$s (?)) OR ((\"c_double\") %1$s (?))", operator.getValue()));
+                    assertThat(converted.get().parameters()).isEqualTo(List.of(
+                            new QueryParameter(BIGINT, Optional.of(42L)),
+                            new QueryParameter(REAL, Optional.of((long) floatToIntBits(3.14f))),
+                            new QueryParameter(DOUBLE, Optional.of(3.14))));
                     break;
                 case IDENTICAL:
                     assertThat(converted).isPresent();