Search in sources :

Example 1 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestTableScanRedirectionWithPushdown method testRedirectionAfterPredicatePushdownIntoTableScan.

@Test
public void testRedirectionAfterPredicatePushdownIntoTableScan() {
    // the connector can detect a filter
    try (LocalQueryRunner queryRunner = createLocalQueryRunner(getMockApplyRedirectAfterPredicatePushdown(REDIRECTION_MAPPING_A, Optional.empty()), Optional.empty(), Optional.of(getMockApplyFilter(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_HANDLE_A))))) {
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table WHERE source_col_a = 1", output(ImmutableList.of("DEST_COL"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A, singleValue(INTEGER, 1L))), Optional.empty())::equals, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A::equals, singleValue(INTEGER, 1L))), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals))));
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table", output(ImmutableList.of("SOURCE_COL"), tableScan(TEST_TABLE, ImmutableMap.of("SOURCE_COL", SOURCE_COLUMN_NAME_A))));
    }
}
Also used : MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) LocalQueryRunner(io.trino.testing.LocalQueryRunner) Test(org.testng.annotations.Test)

Example 2 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestTableScanRedirectionWithPushdown method createLocalQueryRunner.

private LocalQueryRunner createLocalQueryRunner(ApplyTableScanRedirect applyTableScanRedirect, Optional<ApplyProjection> applyProjection, Optional<ApplyFilter> applyFilter) {
    LocalQueryRunner queryRunner = LocalQueryRunner.create(MOCK_SESSION);
    MockConnectorFactory.Builder builder = MockConnectorFactory.builder().withGetTableHandle((session, schemaTableName) -> new MockConnectorTableHandle(schemaTableName)).withGetColumns(name -> {
        if (name.equals(SOURCE_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(SOURCE_COLUMN_NAME_A, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_B, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_C, VARCHAR), new ColumnMetadata(SOURCE_COLUMN_NAME_D, ROW_TYPE));
        } else if (name.equals(DESTINATION_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(DESTINATION_COLUMN_NAME_A, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_B, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_C, ROW_TYPE), new ColumnMetadata(DESTINATION_COLUMN_NAME_D, BOGUS));
        }
        throw new IllegalArgumentException();
    }).withApplyTableScanRedirect(applyTableScanRedirect);
    applyProjection.ifPresent(builder::withApplyProjection);
    applyFilter.ifPresent(builder::withApplyFilter);
    queryRunner.createCatalog(MOCK_CATALOG, builder.build(), ImmutableMap.of());
    return queryRunner;
}
Also used : ApplyTableScanRedirect(io.trino.connector.MockConnectorFactory.ApplyTableScanRedirect) TransactionBuilder.transaction(io.trino.transaction.TransactionBuilder.transaction) OPTIMIZED_AND_VALIDATED(io.trino.sql.planner.LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED) PlanMatchPattern(io.trino.sql.planner.assertions.PlanMatchPattern) Test(org.testng.annotations.Test) ApplyFilter(io.trino.connector.MockConnectorFactory.ApplyFilter) MockConnectorFactory(io.trino.connector.MockConnectorFactory) Arrays.asList(java.util.Arrays.asList) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) INTEGER(io.trino.spi.type.IntegerType.INTEGER) PlanMatchPattern.expression(io.trino.sql.planner.assertions.PlanMatchPattern.expression) RowType(io.trino.spi.type.RowType) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) Set(java.util.Set) TrinoException(io.trino.spi.TrinoException) SchemaTableName(io.trino.spi.connector.SchemaTableName) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) BIGINT(io.trino.spi.type.BigintType.BIGINT) CatalogSchemaTableName(io.trino.spi.connector.CatalogSchemaTableName) Domain.singleValue(io.trino.spi.predicate.Domain.singleValue) Assignment(io.trino.spi.connector.Assignment) Optional(java.util.Optional) RowType.field(io.trino.spi.type.RowType.field) TableScanRedirectApplicationResult(io.trino.spi.connector.TableScanRedirectApplicationResult) PlanMatchPattern.output(io.trino.sql.planner.assertions.PlanMatchPattern.output) Session(io.trino.Session) ApplyProjection(io.trino.connector.MockConnectorFactory.ApplyProjection) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) Type(io.trino.spi.type.Type) PlanMatchPattern.filter(io.trino.sql.planner.assertions.PlanMatchPattern.filter) BOGUS(io.trino.tests.BogusType.BOGUS) Variable(io.trino.spi.expression.Variable) PlanAssert(io.trino.sql.planner.assertions.PlanAssert) VARCHAR(io.trino.spi.type.VarcharType.VARCHAR) ImmutableList(com.google.common.collect.ImmutableList) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) PlanOptimizer(io.trino.sql.planner.optimizations.PlanOptimizer) LocalQueryRunner(io.trino.testing.LocalQueryRunner) ColumnHandle(io.trino.spi.connector.ColumnHandle) ConstraintApplicationResult(io.trino.spi.connector.ConstraintApplicationResult) Language(org.intellij.lang.annotations.Language) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) FieldDereference(io.trino.spi.expression.FieldDereference) ConnectorSession(io.trino.spi.connector.ConnectorSession) TupleDomain(io.trino.spi.predicate.TupleDomain) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) PlanMatchPattern.project(io.trino.sql.planner.assertions.PlanMatchPattern.project) WarningCollector(io.trino.execution.warnings.WarningCollector) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) MockConnectorFactory(io.trino.connector.MockConnectorFactory) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) LocalQueryRunner(io.trino.testing.LocalQueryRunner)

Example 3 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestTableScanRedirectionWithPushdown method testRedirectionBeforeDeferencePushdown.

@Test
public void testRedirectionBeforeDeferencePushdown() {
    // the connector can detect that source_col_a and source_col_d is projected
    try (LocalQueryRunner queryRunner = createLocalQueryRunner(mockApplyRedirectAfterProjectionPushdown(ROW_TYPE_REDIRECTION_MAPPING_AD, Optional.of(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, SOURCE_COLUMN_HANDLE_D))), Optional.of(this::mockApplyProjection), Optional.empty())) {
        // Pushdown of dereference for source_col_d.a into table scan results in a new column handle
        // Table scan redirection would not take place if dereference pushdown has already taken place before redirection
        ColumnHandle destinationColumnHandleC0 = new MockConnectorColumnHandle(DESTINATION_COLUMN_NAME_C + "#0", BIGINT);
        assertPlan(queryRunner, "SELECT source_col_a, source_col_d.a FROM test_table", output(ImmutableList.of("DEST_COL_A", "DEST_COL_C#0"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.all(), Optional.of(ImmutableList.of(DESTINATION_COLUMN_HANDLE_A, destinationColumnHandleC0)))::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL_A", DESTINATION_COLUMN_HANDLE_A::equals, "DEST_COL_C#0", destinationColumnHandleC0::equals))));
    }
}
Also used : ColumnHandle(io.trino.spi.connector.ColumnHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) LocalQueryRunner(io.trino.testing.LocalQueryRunner) Test(org.testng.annotations.Test)

Example 4 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestTableScanRedirectionWithPushdown method mockApplyProjection.

private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
    ImmutableList.Builder<ColumnHandle> newColumnsBuilder = ImmutableList.builder();
    ImmutableList.Builder<ConnectorExpression> outputExpressions = ImmutableList.builder();
    ImmutableList.Builder<Assignment> outputAssignments = ImmutableList.builder();
    for (ConnectorExpression projection : projections) {
        String newVariableName;
        ColumnHandle newColumnHandle;
        if (projection instanceof Variable) {
            newVariableName = ((Variable) projection).getName();
            newColumnHandle = assignments.get(newVariableName);
        } else if (projection instanceof FieldDereference) {
            FieldDereference dereference = (FieldDereference) projection;
            if (!(dereference.getTarget() instanceof Variable)) {
                throw new UnsupportedOperationException();
            }
            String dereferenceTargetName = ((Variable) dereference.getTarget()).getName();
            newVariableName = ((MockConnectorColumnHandle) assignments.get(dereferenceTargetName)).getName() + "#" + dereference.getField();
            newColumnHandle = new MockConnectorColumnHandle(newVariableName, projection.getType());
        } else {
            throw new UnsupportedOperationException();
        }
        Variable newVariable = new Variable(newVariableName, projection.getType());
        newColumnsBuilder.add(newColumnHandle);
        outputExpressions.add(newVariable);
        outputAssignments.add(new Assignment(newVariableName, newColumnHandle, projection.getType()));
    }
    List<ColumnHandle> newColumns = newColumnsBuilder.build();
    if (handle.getColumns().isPresent() && newColumns.equals(handle.getColumns().get())) {
        return Optional.empty();
    }
    return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(handle.getTableName(), handle.getConstraint(), Optional.of(newColumns)), outputExpressions.build(), outputAssignments.build(), false));
}
Also used : ColumnHandle(io.trino.spi.connector.ColumnHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Variable(io.trino.spi.expression.Variable) ImmutableList(com.google.common.collect.ImmutableList) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Assignment(io.trino.spi.connector.Assignment) FieldDereference(io.trino.spi.expression.FieldDereference) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle)

Example 5 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestPushProjectionIntoTableScan method mockApplyProjection.

private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    // Prepare new table handle
    SchemaTableName inputSchemaTableName = ((MockConnectorTableHandle) tableHandle).getTableName();
    SchemaTableName projectedTableName = new SchemaTableName(inputSchemaTableName.getSchemaName(), "projected_" + inputSchemaTableName.getTableName());
    // Prepare new column handles
    ImmutableList.Builder<ConnectorExpression> outputExpressions = ImmutableList.builder();
    ImmutableList.Builder<Assignment> outputAssignments = ImmutableList.builder();
    ImmutableList.Builder<ColumnHandle> projectedColumnsBuilder = ImmutableList.builder();
    for (ConnectorExpression projection : projections) {
        String variablePrefix;
        if (projection instanceof Variable) {
            variablePrefix = "projected_variable_";
        } else if (projection instanceof FieldDereference) {
            variablePrefix = "projected_dereference_";
        } else if (projection instanceof Constant) {
            variablePrefix = "projected_constant_";
        } else if (projection instanceof Call) {
            variablePrefix = "projected_call_";
        } else {
            throw new UnsupportedOperationException();
        }
        String newVariableName = variablePrefix + projection.toString();
        Variable newVariable = new Variable(newVariableName, projection.getType());
        ColumnHandle newColumnHandle = new TpchColumnHandle(newVariableName, projection.getType());
        outputExpressions.add(newVariable);
        outputAssignments.add(new Assignment(newVariableName, newColumnHandle, projection.getType()));
        projectedColumnsBuilder.add(newColumnHandle);
    }
    return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(projectedTableName, TupleDomain.all(), Optional.of(projectedColumnsBuilder.build())), outputExpressions.build(), outputAssignments.build(), false));
}
Also used : TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) ColumnHandle(io.trino.spi.connector.ColumnHandle) FunctionCall(io.trino.sql.tree.FunctionCall) Call(io.trino.spi.expression.Call) Variable(io.trino.spi.expression.Variable) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) ImmutableList(com.google.common.collect.ImmutableList) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Constant(io.trino.spi.expression.Constant) SchemaTableName(io.trino.spi.connector.SchemaTableName) Assignment(io.trino.spi.connector.Assignment) FieldDereference(io.trino.spi.expression.FieldDereference) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle)

Aggregations

MockConnectorTableHandle (io.trino.connector.MockConnectorTableHandle)35 Test (org.testng.annotations.Test)32 MockConnectorFactory (io.trino.connector.MockConnectorFactory)29 ImmutableList (com.google.common.collect.ImmutableList)25 SchemaTableName (io.trino.spi.connector.SchemaTableName)25 TestingSession.testSessionBuilder (io.trino.testing.TestingSession.testSessionBuilder)21 Session (io.trino.Session)20 ImmutableMap (com.google.common.collect.ImmutableMap)19 Optional (java.util.Optional)18 ColumnHandle (io.trino.spi.connector.ColumnHandle)17 ColumnMetadata (io.trino.spi.connector.ColumnMetadata)17 Symbol (io.trino.sql.planner.Symbol)16 BIGINT (io.trino.spi.type.BigintType.BIGINT)15 RuleTester (io.trino.sql.planner.iterative.rule.test.RuleTester)15 RuleTester.defaultRuleTester (io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester)15 List (java.util.List)14 TableHandle (io.trino.metadata.TableHandle)13 TupleDomain (io.trino.spi.predicate.TupleDomain)12 PlanMatchPattern.tableScan (io.trino.sql.planner.assertions.PlanMatchPattern.tableScan)12 Assertions.assertThatThrownBy (org.assertj.core.api.Assertions.assertThatThrownBy)12