From 7e90b64e5dd66b837457fb76f6df3850f82f2101 Mon Sep 17 00:00:00 2001 From: Vladyslav Lyutenko Date: Mon, 15 Jul 2024 19:58:44 +0200 Subject: [PATCH 1/3] Allow customising fetch size for Oracle connector --- .../java/io/trino/plugin/oracle/OracleClient.java | 10 ++++++++-- .../java/io/trino/plugin/oracle/OracleConfig.java | 14 ++++++++++++++ .../io/trino/plugin/oracle/TestOracleConfig.java | 7 +++++-- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java index 60cefaf68c41..ae2b4320714b 100644 --- a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java +++ b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java @@ -216,6 +216,7 @@ public class OracleClient private final boolean synonymsEnabled; private final ConnectorExpressionRewriter connectorExpressionRewriter; private final AggregateFunctionRewriter aggregateFunctionRewriter; + private final Optional fetchSize; @Inject public OracleClient( @@ -260,6 +261,8 @@ public OracleClient( .add(new ImplementCovarianceSamp()) .add(new ImplementCovariancePop()) .build()); + + this.fetchSize = oracleConfig.getFetchSize(); } @Override @@ -287,8 +290,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql, PreparedStatement statement = connection.prepareStatement(sql); // This is a heuristic, not exact science. A better formula can perhaps be found with measurements. // Column count is not known for non-SELECT queries. Not setting fetch size for these. - if (columnCount.isPresent()) { - statement.setFetchSize(max(100_000 / columnCount.get(), 1_000)); + Optional fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() -> + columnCount.map(count -> max(100_000 / count, 1_000)) + .orElse(null))); + if (fetchSize.isPresent()) { + statement.setFetchSize(fetchSize.get()); } return statement; } diff --git a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java index 1c214315368e..09e249c68036 100644 --- a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java +++ b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java @@ -38,6 +38,7 @@ public class OracleConfig private int connectionPoolMinSize = 1; private int connectionPoolMaxSize = 30; private Duration inactiveConnectionTimeout = new Duration(20, MINUTES); + private Integer fetchSize; public boolean isSynonymsEnabled() { @@ -141,6 +142,19 @@ public OracleConfig setInactiveConnectionTimeout(Duration inactiveConnectionTime return this; } + public Optional<@Min(0) Integer> getFetchSize() + { + return Optional.ofNullable(fetchSize); + } + + @Config("oracle.fetch-size") + @ConfigDescription("Oracle fetch size, trino specific heuristic is applied if empty") + public OracleConfig setFetchSize(Integer fetchSize) + { + this.fetchSize = fetchSize; + return this; + } + @AssertTrue(message = "Pool min size cannot be larger than max size") public boolean isPoolSizedProperly() { diff --git a/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java b/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java index 659dec813f51..8d4822c61665 100644 --- a/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java +++ b/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java @@ -42,7 +42,8 @@ public void testDefaults() .setConnectionPoolEnabled(true) .setConnectionPoolMinSize(1) .setConnectionPoolMaxSize(30) - .setInactiveConnectionTimeout(new Duration(20, MINUTES))); + .setInactiveConnectionTimeout(new Duration(20, MINUTES)) + .setFetchSize(null)); } @Test @@ -57,6 +58,7 @@ public void testExplicitPropertyMappings() .put("oracle.connection-pool.min-size", "10") .put("oracle.connection-pool.max-size", "20") .put("oracle.connection-pool.inactive-timeout", "30s") + .put("oracle.fetch-size", "2000") .buildOrThrow(); OracleConfig expected = new OracleConfig() @@ -67,7 +69,8 @@ public void testExplicitPropertyMappings() .setConnectionPoolEnabled(false) .setConnectionPoolMinSize(10) .setConnectionPoolMaxSize(20) - .setInactiveConnectionTimeout(new Duration(30, SECONDS)); + .setInactiveConnectionTimeout(new Duration(30, SECONDS)) + .setFetchSize(2000); assertFullMapping(properties, expected); } From 6625c2ac7e4e7c0fa944185716664ce0ae01f2c2 Mon Sep 17 00:00:00 2001 From: Vladyslav Lyutenko Date: Mon, 15 Jul 2024 19:59:02 +0200 Subject: [PATCH 2/3] Allow customising fetch size for PostgreSql connector --- .../plugin/postgresql/PostgreSqlClient.java | 9 +++++++-- .../plugin/postgresql/PostgreSqlConfig.java | 18 ++++++++++++++++++ .../postgresql/TestPostgreSqlConfig.java | 7 +++++-- 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java index 24a1a6ea6446..127a4bb4b435 100644 --- a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java +++ b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java @@ -282,6 +282,7 @@ public class PostgreSqlClient private final ConnectorExpressionRewriter connectorExpressionRewriter; private final ProjectFunctionRewriter projectFunctionRewriter; private final AggregateFunctionRewriter aggregateFunctionRewriter; + private final Optional fetchSize; @Inject public PostgreSqlClient( @@ -368,6 +369,7 @@ public PostgreSqlClient( .add(new ImplementRegrIntercept()) .add(new ImplementRegrSlope()) .build()); + this.fetchSize = postgreSqlConfig.getFetchSize(); } @Override @@ -432,8 +434,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql, PreparedStatement statement = connection.prepareStatement(sql); // This is a heuristic, not exact science. A better formula can perhaps be found with measurements. // Column count is not known for non-SELECT queries. Not setting fetch size for these. - if (columnCount.isPresent()) { - statement.setFetchSize(max(100_000 / columnCount.get(), 1_000)); + Optional fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() -> + columnCount.map(count -> max(100_000 / count, 1_000)) + .orElse(null))); + if (fetchSize.isPresent()) { + statement.setFetchSize(fetchSize.get()); } return statement; } diff --git a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java index f7b0ca32f962..07aa0df1db82 100644 --- a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java +++ b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java @@ -14,16 +14,21 @@ package io.trino.plugin.postgresql; import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.DefunctConfig; import io.airlift.configuration.LegacyConfig; +import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; +import java.util.Optional; + @DefunctConfig("postgresql.disable-automatic-fetch-size") public class PostgreSqlConfig { private ArrayMapping arrayMapping = ArrayMapping.DISABLED; private boolean includeSystemTables; private boolean enableStringPushdownWithCollate; + private Integer fetchSize; public enum ArrayMapping { @@ -69,4 +74,17 @@ public PostgreSqlConfig setEnableStringPushdownWithCollate(boolean enableStringP this.enableStringPushdownWithCollate = enableStringPushdownWithCollate; return this; } + + public Optional<@Min(0) Integer> getFetchSize() + { + return Optional.ofNullable(fetchSize); + } + + @Config("postgresql.fetch-size") + @ConfigDescription("Postgresql fetch size, trino specific heuristic is applied if empty") + public PostgreSqlConfig setFetchSize(Integer fetchSize) + { + this.fetchSize = fetchSize; + return this; + } } diff --git a/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java b/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java index aade2bde29e7..5251be22235d 100644 --- a/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java +++ b/plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java @@ -30,7 +30,8 @@ public void testDefaults() assertRecordedDefaults(recordDefaults(PostgreSqlConfig.class) .setArrayMapping(PostgreSqlConfig.ArrayMapping.DISABLED) .setIncludeSystemTables(false) - .setEnableStringPushdownWithCollate(false)); + .setEnableStringPushdownWithCollate(false) + .setFetchSize(null)); } @Test @@ -40,12 +41,14 @@ public void testExplicitPropertyMappings() .put("postgresql.array-mapping", "AS_ARRAY") .put("postgresql.include-system-tables", "true") .put("postgresql.experimental.enable-string-pushdown-with-collate", "true") + .put("postgresql.fetch-size", "2000") .buildOrThrow(); PostgreSqlConfig expected = new PostgreSqlConfig() .setArrayMapping(PostgreSqlConfig.ArrayMapping.AS_ARRAY) .setIncludeSystemTables(true) - .setEnableStringPushdownWithCollate(true); + .setEnableStringPushdownWithCollate(true) + .setFetchSize(2000); assertFullMapping(properties, expected); } From d880bd47b8a25a8bded3134a7898dfd9dcabc244 Mon Sep 17 00:00:00 2001 From: Vladyslav Lyutenko Date: Mon, 15 Jul 2024 19:59:13 +0200 Subject: [PATCH 3/3] Allow customising fetch size for Redshift connector --- plugin/trino-redshift/pom.xml | 5 +++++ .../trino/plugin/redshift/RedshiftClient.java | 10 ++++++++-- .../trino/plugin/redshift/RedshiftConfig.java | 19 +++++++++++++++++++ .../plugin/redshift/TestRedshiftConfig.java | 7 +++++-- 4 files changed, 37 insertions(+), 4 deletions(-) diff --git a/plugin/trino-redshift/pom.xml b/plugin/trino-redshift/pom.xml index 3f2c8f78a5a6..f17509a4ccb7 100644 --- a/plugin/trino-redshift/pom.xml +++ b/plugin/trino-redshift/pom.xml @@ -50,6 +50,11 @@ trino-plugin-toolkit + + jakarta.validation + jakarta.validation-api + + org.jdbi jdbi3-core diff --git a/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java b/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java index fc41a8359099..c9d8d5a0f0cc 100644 --- a/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java +++ b/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java @@ -228,10 +228,12 @@ public class RedshiftClient private final boolean statisticsEnabled; private final RedshiftTableStatisticsReader statisticsReader; private final ConnectorExpressionRewriter connectorExpressionRewriter; + private final Optional fetchSize; @Inject public RedshiftClient( BaseJdbcConfig config, + RedshiftConfig redshiftConfig, ConnectionFactory connectionFactory, JdbcStatisticsConfig statisticsConfig, QueryBuilder queryBuilder, @@ -269,6 +271,7 @@ public RedshiftClient( this.statisticsEnabled = requireNonNull(statisticsConfig, "statisticsConfig is null").isEnabled(); this.statisticsReader = new RedshiftTableStatisticsReader(connectionFactory); + this.fetchSize = redshiftConfig.getFetchSize(); } private static Optional toTypeHandle(DecimalType decimalType) @@ -487,8 +490,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql, PreparedStatement statement = connection.prepareStatement(sql); // This is a heuristic, not exact science. A better formula can perhaps be found with measurements. // Column count is not known for non-SELECT queries. Not setting fetch size for these. - if (columnCount.isPresent()) { - statement.setFetchSize(max(100_000 / columnCount.get(), 1_000)); + Optional fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() -> + columnCount.map(count -> max(100_000 / count, 1_000)) + .orElse(null))); + if (fetchSize.isPresent()) { + statement.setFetchSize(fetchSize.get()); } return statement; } diff --git a/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java b/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java index 88924f8c6fee..370f31e7dc8b 100644 --- a/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java +++ b/plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java @@ -13,7 +13,12 @@ */ package io.trino.plugin.redshift; +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.DefunctConfig; +import jakarta.validation.constraints.Min; + +import java.util.Optional; @DefunctConfig({ "redshift.disable-automatic-fetch-size", @@ -21,4 +26,18 @@ }) public class RedshiftConfig { + private Integer fetchSize; + + public Optional<@Min(0) Integer> getFetchSize() + { + return Optional.ofNullable(fetchSize); + } + + @Config("redshift.fetch-size") + @ConfigDescription("Redshift fetch size, trino specific heuristic is applied if empty") + public RedshiftConfig setFetchSize(Integer fetchSize) + { + this.fetchSize = fetchSize; + return this; + } } diff --git a/plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java b/plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java index f9385e04e1d6..6a507d29d87b 100644 --- a/plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java +++ b/plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java @@ -27,16 +27,19 @@ public class TestRedshiftConfig @Test public void testDefaults() { - assertRecordedDefaults(recordDefaults(RedshiftConfig.class)); + assertRecordedDefaults(recordDefaults(RedshiftConfig.class) + .setFetchSize(null)); } @Test public void testExplicitPropertyMappings() { Map properties = ImmutableMap.builder() + .put("redshift.fetch-size", "2000") .buildOrThrow(); - RedshiftConfig expected = new RedshiftConfig(); + RedshiftConfig expected = new RedshiftConfig() + .setFetchSize(2000); assertFullMapping(properties, expected); }