Implement AGGREATE pushdown in Exasol Connector

skyglass · skyglass · commit 5d42f00fcb80 · 2025-09-20T18:03:06.000+02:00
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/StandardColumnMappings.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/StandardColumnMappings.java
@@ -177,12 +177,12 @@ public static ColumnMapping decimalColumnMapping(DecimalType decimalType, Roundi
             checkArgument(roundingMode == UNNECESSARY, "Round mode is not supported for short decimal, map the type to long decimal instead");
             return ColumnMapping.longMapping(
                     decimalType,
-                    shortDecimalReadFunction(decimalType),
+                    shortDecimalReadFunction(decimalType, roundingMode),
                     shortDecimalWriteFunction(decimalType));
         }
         return ColumnMapping.objectMapping(
                 decimalType,
-                longDecimalReadFunction(decimalType, roundingMode),
+                longDecimalReadFunction(decimalType),
                 longDecimalWriteFunction(decimalType));
     }
 
diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java
@@ -150,12 +150,12 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
         return switch (connectorBehavior) {
             case SUPPORTS_UPDATE -> true;
             case SUPPORTS_ADD_COLUMN_WITH_POSITION,
-                 SUPPORTS_CREATE_MATERIALIZED_VIEW,
-                 SUPPORTS_CREATE_VIEW,
-                 SUPPORTS_DEFAULT_COLUMN_VALUE,
-                 SUPPORTS_MERGE,
-                 SUPPORTS_PREDICATE_EXPRESSION_PUSHDOWN,
-                 SUPPORTS_ROW_LEVEL_UPDATE -> false;
+                     SUPPORTS_CREATE_MATERIALIZED_VIEW,
+                     SUPPORTS_CREATE_VIEW,
+                     SUPPORTS_DEFAULT_COLUMN_VALUE,
+                     SUPPORTS_MERGE,
+                     SUPPORTS_PREDICATE_EXPRESSION_PUSHDOWN,
+                     SUPPORTS_ROW_LEVEL_UPDATE -> false;
             // Dynamic filters can be pushed down only if predicate push down is supported.
             // It is possible for a connector to have predicate push down support but not push down dynamic filters.
             // TODO default SUPPORTS_DYNAMIC_FILTER_PUSHDOWN to SUPPORTS_PREDICATE_PUSHDOWN
@@ -621,15 +621,15 @@ public void testNumericAggregationPushdown()
             assertThat(query("SELECT min(short_decimal), min(long_decimal), min(a_bigint), min(t_double) FROM " + emptyTable.getName())).isFullyPushedDown();
             assertThat(query("SELECT max(short_decimal), max(long_decimal), max(a_bigint), max(t_double) FROM " + emptyTable.getName())).isFullyPushedDown();
             assertThat(query("SELECT sum(short_decimal), sum(long_decimal), sum(a_bigint), sum(t_double) FROM " + emptyTable.getName())).isFullyPushedDown();
-            assertThat(query("SELECT avg(short_decimal), avg(long_decimal), avg(a_bigint), avg(t_double) FROM " + emptyTable.getName())).isFullyPushedDown();
+            assertNumericAveragePushdown(emptyTable);
         }
 
         try (TestTable testTable = createAggregationTestTable(schemaName + ".test_num_agg_pd",
                 ImmutableList.of("100.000, 100000000.000000000, 100.000, 100000000", "123.321, 123456789.987654321, 123.321, 123456789"))) {
             assertThat(query("SELECT min(short_decimal), min(long_decimal), min(a_bigint), min(t_double) FROM " + testTable.getName())).isFullyPushedDown();
             assertThat(query("SELECT max(short_decimal), max(long_decimal), max(a_bigint), max(t_double) FROM " + testTable.getName())).isFullyPushedDown();
             assertThat(query("SELECT sum(short_decimal), sum(long_decimal), sum(a_bigint), sum(t_double) FROM " + testTable.getName())).isFullyPushedDown();
-            assertThat(query("SELECT avg(short_decimal), avg(long_decimal), avg(a_bigint), avg(t_double) FROM " + testTable.getName())).isFullyPushedDown();
+            assertNumericAveragePushdown(testTable);
 
             // smoke testing of more complex cases
             // WHERE on aggregation column
@@ -647,6 +647,11 @@ public void testNumericAggregationPushdown()
         }
     }
 
+    protected void assertNumericAveragePushdown(TestTable testTable)
+    {
+        assertThat(query("SELECT avg(short_decimal), avg(long_decimal), avg(a_bigint), avg(t_double) FROM " + testTable.getName())).isFullyPushedDown();
+    }
+
     @Test
     public void testCountDistinctWithStringTypes()
     {
@@ -1153,12 +1158,12 @@ public void testArithmeticPredicatePushdown()
 
         assertThat(query("SELECT nationkey, name, regionkey FROM nation WHERE nationkey > 0 AND (nationkey - regionkey) % nationkey = 2"))
                 .isFullyPushedDown()
-                .matches("VALUES (BIGINT '3', CAST('CANADA' AS varchar(25)), BIGINT '1')");
+                .matches(getArithmeticPredicatePushdownExpectedValues());
 
         // some databases calculate remainder instead of modulus when one of the values is negative
         assertThat(query("SELECT nationkey, name, regionkey FROM nation WHERE nationkey > 0 AND (nationkey - regionkey) % -nationkey = 2"))
                 .isFullyPushedDown()
-                .matches("VALUES (BIGINT '3', CAST('CANADA' AS varchar(25)), BIGINT '1')");
+                .matches(getArithmeticPredicatePushdownExpectedValues());
 
         assertThat(query("SELECT nationkey, name, regionkey FROM nation WHERE nationkey > 0 AND (nationkey - regionkey) % 0 = 2"))
                 .failure().hasMessageContaining("by zero");
@@ -1170,6 +1175,11 @@ public void testArithmeticPredicatePushdown()
         // TODO add coverage for other arithmetic pushdowns https://github.com/trinodb/trino/issues/14808
     }
 
+    protected String getArithmeticPredicatePushdownExpectedValues()
+    {
+        return "VALUES (BIGINT '3', CAST('CANADA' AS varchar(25)), BIGINT '1')";
+    }
+
     @Test
     public void testCaseSensitiveTopNPushdown()
     {
@@ -1307,7 +1317,8 @@ public void testJoinPushdown()
                 assertThat(query(session, format("SELECT n.name FROM nation n %s orders o ON DATE '2025-03-19' = o.orderdate", joinOperator))).joinIsNotFullyPushedDown();
 
                 // no projection on the probe side, only filter
-                assertJoinConditionallyPushedDown(session, format("SELECT n.name FROM nation n %s orders o ON n.regionkey = 1", joinOperator),
+                // reduced the size of the join table to make the test faster: instead of joining on large orders table join only on one record
+                assertJoinConditionallyPushedDown(session, format("SELECT n.name FROM nation n %s (SELECT * FROM orders WHERE orderkey = 1) o ON n.regionkey = 1", joinOperator),
                         expectJoinPushdownOnEmptyProjection(joinOperator));
 
                 // pushdown when using USING
diff --git a/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java b/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java
@@ -16,6 +16,8 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
 import io.airlift.slice.Slices;
+import io.trino.plugin.base.aggregation.AggregateFunctionRewriter;
+import io.trino.plugin.base.aggregation.AggregateFunctionRule;
 import io.trino.plugin.base.expression.ConnectorExpressionRewriter;
 import io.trino.plugin.base.mapping.IdentifierMapping;
 import io.trino.plugin.jdbc.BaseJdbcClient;
@@ -36,6 +38,22 @@
 import io.trino.plugin.jdbc.SliceWriteFunction;
 import io.trino.plugin.jdbc.WriteFunction;
 import io.trino.plugin.jdbc.WriteMapping;
+import io.trino.plugin.jdbc.aggregation.ImplementAvgDecimal;
+import io.trino.plugin.jdbc.aggregation.ImplementAvgFloatingPoint;
+import io.trino.plugin.jdbc.aggregation.ImplementCorr;
+import io.trino.plugin.jdbc.aggregation.ImplementCount;
+import io.trino.plugin.jdbc.aggregation.ImplementCountAll;
+import io.trino.plugin.jdbc.aggregation.ImplementCountDistinct;
+import io.trino.plugin.jdbc.aggregation.ImplementCovariancePop;
+import io.trino.plugin.jdbc.aggregation.ImplementCovarianceSamp;
+import io.trino.plugin.jdbc.aggregation.ImplementMinMax;
+import io.trino.plugin.jdbc.aggregation.ImplementRegrIntercept;
+import io.trino.plugin.jdbc.aggregation.ImplementRegrSlope;
+import io.trino.plugin.jdbc.aggregation.ImplementStddevPop;
+import io.trino.plugin.jdbc.aggregation.ImplementStddevSamp;
+import io.trino.plugin.jdbc.aggregation.ImplementSum;
+import io.trino.plugin.jdbc.aggregation.ImplementVariancePop;
+import io.trino.plugin.jdbc.aggregation.ImplementVarianceSamp;
 import io.trino.plugin.jdbc.expression.JdbcConnectorExpressionRewriterBuilder;
 import io.trino.plugin.jdbc.expression.ParameterizedExpression;
 import io.trino.plugin.jdbc.expression.RewriteIn;
@@ -50,6 +68,7 @@
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.type.DecimalType;
 import io.trino.spi.type.Type;
+import io.trino.spi.type.VarcharType;
 
 import java.sql.Connection;
 import java.sql.Date;
@@ -73,6 +92,7 @@
 import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalWriteFunction;
 import static io.trino.plugin.jdbc.StandardColumnMappings.smallintColumnMapping;
+import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction;
 import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling;
 import static io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
@@ -91,7 +111,9 @@ public class ExasolClient
             .add("EXA_STATISTICS")
             .add("SYS")
             .build();
+    public static final int MAX_EXASOL_DECIMAL_PRECISION = 36;
     private final ConnectorExpressionRewriter<ParameterizedExpression> connectorExpressionRewriter;
+    private final AggregateFunctionRewriter<JdbcExpression, ?> aggregateFunctionRewriter;
 
     @Inject
     public ExasolClient(
@@ -102,12 +124,13 @@ public ExasolClient(
             RemoteQueryModifier queryModifier)
     {
         super("\"", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, false);
-        // Basic implementation required to enable JOIN pushdown support
-        // It is covered by "testJoinpushdown" integration tests.
-        // More detailed test case scenarios are covered by Unit tests in "TestConvertPredicate"
+        // Basic implementation required to enable JOIN and AGGREGATION pushdown support
+        // It is covered by "testJoinPushdown" and "testAggregationPushdown" integration tests.
+        // More detailed test case scenarios are covered by Unit tests in "TestExasolClient"
         this.connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder()
                 .addStandardRules(this::quoted)
                 .add(new RewriteIn())
+                .withTypeClass("numeric_type", ImmutableSet.of("decimal", "double"))
                 .map("$equal(left, right)").to("left = right")
                 .map("$not_equal(left, right)").to("left <> right")
                 // Exasol doesn't support "IS NOT DISTINCT FROM" expression,
@@ -120,7 +143,38 @@ public ExasolClient(
                 .map("$not($is_null(value))").to("value IS NOT NULL")
                 .map("$not(value: boolean)").to("NOT value")
                 .map("$is_null(value)").to("value IS NULL")
+                .map("$add(left: numeric_type, right: numeric_type)").to("left + right")
+                .map("$subtract(left: numeric_type, right: numeric_type)").to("left - right")
+                .map("$multiply(left: numeric_type, right: numeric_type)").to("left * right")
+                .map("$divide(left: numeric_type, right: numeric_type)").to("left / right")
+                .map("$modulus(left: numeric_type, right: numeric_type)").to("mod(left, right)")
+                .map("$negate(value: numeric_type)").to("-value")
+                .map("$like(value: varchar, pattern: varchar): boolean").to("value LIKE pattern")
+                .map("$like(value: varchar, pattern: varchar, escape: varchar(1)): boolean").to("value LIKE pattern ESCAPE escape")
+                .map("$nullif(first, second)").to("NULLIF(first, second)")
                 .build();
+        JdbcTypeHandle bigintTypeHandle = new JdbcTypeHandle(Types.BIGINT, Optional.of("bigint"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
+        this.aggregateFunctionRewriter = new AggregateFunctionRewriter<>(
+                this.connectorExpressionRewriter,
+                ImmutableSet.<AggregateFunctionRule<JdbcExpression, ParameterizedExpression>>builder()
+                        .add(new ImplementCountAll(bigintTypeHandle))
+                        .add(new ImplementMinMax(true))
+                        .add(new ImplementCount(bigintTypeHandle))
+                        .add(new ImplementCountDistinct(bigintTypeHandle, true))
+                        .add(new ImplementSum(ExasolClient::toSumTypeHandle))
+                        .add(new ImplementAvgFloatingPoint())
+                        .add(new ImplementAvgDecimal())
+                        .add(new ImplementExasolAvgBigInt())
+                        .add(new ImplementStddevSamp())
+                        .add(new ImplementStddevPop())
+                        .add(new ImplementVarianceSamp())
+                        .add(new ImplementVariancePop())
+                        .add(new ImplementCovarianceSamp())
+                        .add(new ImplementCovariancePop())
+                        .add(new ImplementCorr())
+                        .add(new ImplementRegrIntercept())
+                        .add(new ImplementRegrSlope())
+                        .build());
     }
 
     @Override
@@ -234,11 +288,22 @@ protected boolean isSupportedJoinCondition(ConnectorSession session, JdbcJoinCon
         return true;
     }
 
+    @Override
+    public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, List<JdbcSortItem> sortOrder)
+    {
+        return true;
+    }
+
     @Override
     public Optional<JdbcExpression> implementAggregation(ConnectorSession session, AggregateFunction aggregate, Map<String, ColumnHandle> assignments)
     {
-        // Deactivated because test 'testCaseSensitiveAggregationPushdown()' requires write access which is not implemented for Exasol
-        return Optional.empty();
+        return aggregateFunctionRewriter.rewrite(session, aggregate, assignments);
+    }
+
+    @Override
+    public boolean supportsAggregationPushdown(ConnectorSession session, JdbcTableHandle table, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets)
+    {
+        return true;
     }
 
     @Override
@@ -263,9 +328,9 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
             case Types.DOUBLE:
                 return Optional.of(doubleColumnMapping());
             case Types.DECIMAL:
-                int decimalDigits = typeHandle.requiredDecimalDigits();
-                int columnSize = typeHandle.requiredColumnSize();
-                return Optional.of(decimalColumnMapping(createDecimalType(columnSize, decimalDigits)));
+                int precision = typeHandle.requiredColumnSize();
+                int scale = typeHandle.requiredDecimalDigits();
+                return Optional.of(decimalColumnMapping(createDecimalType(precision, scale)));
             case Types.CHAR:
                 return Optional.of(defaultCharColumnMapping(typeHandle.requiredColumnSize(), true));
             case Types.VARCHAR:
@@ -289,6 +354,12 @@ private boolean isHashType(JdbcTypeHandle typeHandle)
                 && typeHandle.jdbcTypeName().get().equalsIgnoreCase("HASHTYPE");
     }
 
+    private static Optional<JdbcTypeHandle> toSumTypeHandle(DecimalType decimalType)
+    {
+        return Optional.of(new JdbcTypeHandle(Types.DECIMAL, Optional.of("decimal"),
+                Optional.of(decimalType.getPrecision()), Optional.of(decimalType.getScale()), Optional.empty(), Optional.empty()));
+    }
+
     private static ColumnMapping dateColumnMapping()
     {
         // Exasol driver does not support LocalDate
@@ -350,6 +421,16 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
             }
             return WriteMapping.objectMapping(dataType, longDecimalWriteFunction(decimalType));
         }
+        if (type instanceof VarcharType varcharType) {
+            String dataType;
+            if (varcharType.isUnbounded()) {
+                dataType = "varchar";
+            }
+            else {
+                dataType = "varchar(" + varcharType.getBoundedLength() + ")";
+            }
+            return WriteMapping.sliceMapping(dataType, varcharWriteFunction());
+        }
 
         throw new TrinoException(NOT_SUPPORTED, "Unsupported column type: " + type.getDisplayName());
     }
@@ -398,10 +479,4 @@ public boolean isLimitGuaranteed(ConnectorSession session)
     {
         return true;
     }
-
-    @Override
-    public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, List<JdbcSortItem> sortOrder)
-    {
-        return true;
-    }
 }
diff --git a/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ImplementExasolAvgBigInt.java b/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ImplementExasolAvgBigInt.java
@@ -0,0 +1,26 @@
+/*
+ * 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.exasol;
+
+import io.trino.plugin.jdbc.aggregation.BaseImplementAvgBigint;
+
+public class ImplementExasolAvgBigInt
+        extends BaseImplementAvgBigint
+{
+    @Override
+    protected String getRewriteFormatExpression()
+    {
+        return "avg(CAST(%s AS double))";
+    }
+}
diff --git a/plugin/trino-exasol/src/test/java/io/trino/plugin/exasol/TestExasolConnectorTest.java b/plugin/trino-exasol/src/test/java/io/trino/plugin/exasol/TestExasolConnectorTest.java

Original file line number	Diff line number	Diff line change
`@@ -177,12 +177,12 @@ public static ColumnMapping decimalColumnMapping(DecimalType decimalType, Roundi`
`177`	`177`	`checkArgument(roundingMode == UNNECESSARY, "Round mode is not supported for short decimal, map the type to long decimal instead");`
`178`	`178`	`return ColumnMapping.longMapping(`
`179`	`179`	`decimalType,`
`180`		`- shortDecimalReadFunction(decimalType),`
	`180`	`+ shortDecimalReadFunction(decimalType, roundingMode),`
`181`	`181`	`shortDecimalWriteFunction(decimalType));`
`182`	`182`	`}`
`183`	`183`	`return ColumnMapping.objectMapping(`
`184`	`184`	`decimalType,`
`185`		`- longDecimalReadFunction(decimalType, roundingMode),`
	`185`	`+ longDecimalReadFunction(decimalType),`
`186`	`186`	`longDecimalWriteFunction(decimalType));`
`187`	`187`	`}`
`188`	`188`