Skip to content

Commit

Permalink
Convert range predicates to discrete set in Redshift
Browse files Browse the repository at this point in the history
  • Loading branch information
@mayankvadariya @raunaqmorarka
mayankvadariya authored and raunaqmorarka committed Oct 7, 2024
1 parent 9c09b6a commit 3999f62
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 6 deletions.
17 changes: 17 additions & 0 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/PredicatePushdownController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,14 @@
import io.trino.spi.predicate.DiscreteValues;
import io.trino.spi.predicate.Domain;
import io.trino.spi.predicate.Ranges;
import io.trino.spi.predicate.ValueSet;
import io.trino.spi.type.CharType;
import io.trino.spi.type.Type;
import io.trino.spi.type.VarcharType;

import java.util.Collection;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkArgument;
import static io.trino.plugin.jdbc.JdbcMetadataSessionProperties.getDomainCompactionThreshold;
import static java.util.Objects.requireNonNull;
Expand Down Expand Up @@ -62,6 +67,18 @@ public interface PredicatePushdownController
return new DomainPushdownResult(simplifiedDomain, domain);
};

static PredicatePushdownController pushdownDiscreteValues(Type type)
{
return (session, domain) -> {
Optional<Collection<Object>> expandedRange = domain.getValues().tryExpandRanges(getDomainCompactionThreshold(session));
if (expandedRange.isPresent()) {
Domain convertedDiscreteDomain = Domain.create(ValueSet.copyOf(type, expandedRange.get()), domain.isNullAllowed());
return new DomainPushdownResult(convertedDiscreteDomain, Domain.all(domain.getType()));
}
return FULL_PUSHDOWN.apply(session, domain);
};
}

DomainPushdownResult apply(ConnectorSession session, Domain domain);

final class DomainPushdownResult
Expand Down
13 changes: 7 additions & 6 deletions plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,22 +113,20 @@
import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_NON_TRANSIENT_ERROR;
import static io.trino.plugin.jdbc.JdbcJoinPushdownUtil.implementJoinCostAware;
import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping;
import static io.trino.plugin.jdbc.PredicatePushdownController.pushdownDiscreteValues;
import static io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.booleanColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.booleanWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.charReadFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.decimalColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.doubleWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.integerColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.integerWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalReadFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.realColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.realWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.smallintColumnMapping;
import static io.trino.plugin.jdbc.StandardColumnMappings.smallintWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.tinyintWriteFunction;
import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryReadFunction;
Expand Down Expand Up @@ -618,11 +616,14 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect

// case Types.TINYINT: -- Redshift doesn't support tinyint
case Types.SMALLINT:
return Optional.of(smallintColumnMapping());
// IN clause query in Redshift performs better compared to range queries, hence convert range queries to discrete set where possible.
return Optional.of(ColumnMapping.longMapping(SMALLINT, ResultSet::getShort, smallintWriteFunction(), pushdownDiscreteValues(SMALLINT)));
case Types.INTEGER:
return Optional.of(integerColumnMapping());
// IN clause query in Redshift performs better compared to range queries, hence convert range queries to discrete set where possible.
return Optional.of(ColumnMapping.longMapping(INTEGER, ResultSet::getInt, integerWriteFunction(), pushdownDiscreteValues(INTEGER)));
case Types.BIGINT:
return Optional.of(bigintColumnMapping());
// IN clause query in Redshift performs better compared to range queries, hence convert range queries to discrete set where possible.
return Optional.of(ColumnMapping.longMapping(BIGINT, ResultSet::getLong, bigintWriteFunction(), pushdownDiscreteValues(BIGINT)));

case Types.REAL:
return Optional.of(realColumnMapping());
Expand Down
42 changes: 42 additions & 0 deletions plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,15 @@
import io.airlift.units.Duration;
import io.trino.Session;
import io.trino.plugin.jdbc.BaseJdbcConnectorTest;
import io.trino.plugin.jdbc.JdbcColumnHandle;
import io.trino.plugin.jdbc.JdbcTableHandle;
import io.trino.plugin.jdbc.JdbcTypeHandle;
import io.trino.plugin.jdbc.RemoteDatabaseEvent;
import io.trino.plugin.jdbc.RemoteDatabaseEvent.Status;
import io.trino.plugin.jdbc.RemoteLogTracingEvent;
import io.trino.spi.predicate.Domain;
import io.trino.spi.predicate.TupleDomain;
import io.trino.sql.planner.plan.TableScanNode;
import io.trino.testing.QueryRunner;
import io.trino.testing.TestingConnectorBehavior;
import io.trino.testing.sql.SqlExecutor;
Expand All @@ -30,7 +36,9 @@
import org.jdbi.v3.core.Jdbi;
import org.junit.jupiter.api.Test;

import java.sql.Types;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.Set;
Expand All @@ -51,6 +59,8 @@
import static io.trino.plugin.redshift.TestingRedshiftServer.TEST_SCHEMA;
import static io.trino.plugin.redshift.TestingRedshiftServer.executeInRedshift;
import static io.trino.plugin.redshift.TestingRedshiftServer.executeWithRedshift;
import static io.trino.spi.type.BigintType.BIGINT;
import static io.trino.sql.planner.assertions.PlanMatchPattern.node;
import static io.trino.testing.TestingNames.randomNameSuffix;
import static java.lang.Math.round;
import static java.lang.String.format;
Expand Down Expand Up @@ -267,6 +277,38 @@ public void testRedshiftAddNotNullColumn()
}
}

@Test
public void testRangeQueryConvertedToInClauseQuery()
{
assertThat(query("SELECT regionkey FROM region WHERE regionkey >= 1 AND regionkey <= 4"))
.isFullyPushedDown();
assertThat(query("SELECT regionkey FROM region WHERE regionkey >= 1 AND regionkey <= 4"))
.isNotFullyPushedDown(node(TableScanNode.class)
.with(TableScanNode.class, tableScanNode -> {
TupleDomain<?> effectivePredicate = ((JdbcTableHandle) tableScanNode.getTable().connectorHandle()).getConstraint();
TupleDomain<?> expectedPredicate =
TupleDomain.withColumnDomains(
Map.of(
new JdbcColumnHandle.Builder()
.setColumnName("regionkey")
.setJdbcTypeHandle(
new JdbcTypeHandle(
Types.BIGINT,
Optional.of("int8"),
Optional.of(19),
Optional.of(0),
Optional.empty(),
Optional.empty()))
.setComment(Optional.of("Dynamic Column."))
.setColumnType(BIGINT)
.setNullable(true)
.build(),
Domain.multipleValues(BIGINT, List.of(1L, 2L, 3L, 4L), false)));
assertThat(effectivePredicate).isEqualTo(expectedPredicate);
return true;
}));
}

@Test
@Override
public void testDelete()
Expand Down

0 comments on commit 3999f62

Please sign in to comment.