Search in sources :

Example 1 with MockConnectorColumnHandle

use of io.trino.connector.MockConnectorColumnHandle 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 2 with MockConnectorColumnHandle

use of io.trino.connector.MockConnectorColumnHandle 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 3 with MockConnectorColumnHandle

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

the class TestPushDistinctLimitIntoTableScan method testNoEffect.

@Test
public void testNoEffect() {
    AtomicInteger applyCallCounter = new AtomicInteger();
    testApplyAggregation = (session, handle, aggregates, assignments, groupingSets) -> {
        applyCallCounter.incrementAndGet();
        return Optional.empty();
    };
    tester().assertThat(rule).on(p -> {
        Symbol regionkey = p.symbol("regionkey");
        return p.distinctLimit(10, List.of(regionkey), p.tableScan(tableHandle, ImmutableList.of(regionkey), ImmutableMap.of(regionkey, new MockConnectorColumnHandle("regionkey", BIGINT))));
    }).doesNotFire();
    assertThat(applyCallCounter).as("applyCallCounter").hasValue(1);
}
Also used : AggregateFunction(io.trino.spi.connector.AggregateFunction) Assertions.assertThat(org.assertj.core.api.Assertions.assertThat) Test(org.testng.annotations.Test) AtomicReference(java.util.concurrent.atomic.AtomicReference) Collectors.toUnmodifiableList(java.util.stream.Collectors.toUnmodifiableList) PlanMatchPattern.limit(io.trino.sql.planner.assertions.PlanMatchPattern.limit) CatalogName(io.trino.connector.CatalogName) ImmutableList(com.google.common.collect.ImmutableList) MockConnectorFactory(io.trino.connector.MockConnectorFactory) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) LocalQueryRunner(io.trino.testing.LocalQueryRunner) Map(java.util.Map) ColumnHandle(io.trino.spi.connector.ColumnHandle) Symbol(io.trino.sql.planner.Symbol) ApplyAggregation(io.trino.connector.MockConnectorFactory.ApplyAggregation) ImmutableMap(com.google.common.collect.ImmutableMap) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) BeforeClass(org.testng.annotations.BeforeClass) BeforeMethod(org.testng.annotations.BeforeMethod) SchemaTableName(io.trino.spi.connector.SchemaTableName) AggregationApplicationResult(io.trino.spi.connector.AggregationApplicationResult) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorTransactionHandle(io.trino.connector.MockConnectorTransactionHandle) List(java.util.List) TableHandle(io.trino.metadata.TableHandle) BIGINT(io.trino.spi.type.BigintType.BIGINT) PlanMatchPattern.project(io.trino.sql.planner.assertions.PlanMatchPattern.project) TestingSession(io.trino.testing.TestingSession) Optional(java.util.Optional) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) Session(io.trino.Session) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Symbol(io.trino.sql.planner.Symbol) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Test(org.testng.annotations.Test) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest)

Example 4 with MockConnectorColumnHandle

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

the class TestPruneTableScanColumns method testPushColumnPruningProjection.

@Test
public void testPushColumnPruningProjection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        String mockCatalog = "mock_catalog";
        String testSchema = "test_schema";
        String testTable = "test_table";
        SchemaTableName testSchemaTable = new SchemaTableName(testSchema, testTable);
        ColumnHandle columnHandleA = new MockConnectorColumnHandle("cola", DATE);
        ColumnHandle columnHandleB = new MockConnectorColumnHandle("colb", DOUBLE);
        Map<String, ColumnHandle> assignments = ImmutableMap.of("cola", columnHandleA, "colb", columnHandleB);
        // Create catalog with applyProjection
        MockConnectorFactory factory = MockConnectorFactory.builder().withListSchemaNames(connectorSession -> ImmutableList.of(testSchema)).withListTables((connectorSession, schema) -> testSchema.equals(schema) ? ImmutableList.of(testSchemaTable) : ImmutableList.of()).withGetColumns(schemaTableName -> assignments.entrySet().stream().map(entry -> new ColumnMetadata(entry.getKey(), ((MockConnectorColumnHandle) entry.getValue()).getType())).collect(toImmutableList())).withApplyProjection(this::mockApplyProjection).build();
        ruleTester.getQueryRunner().createCatalog(mockCatalog, factory, ImmutableMap.of());
        ruleTester.assertThat(new PruneTableScanColumns(ruleTester.getMetadata())).on(p -> {
            Symbol symbolA = p.symbol("cola", DATE);
            Symbol symbolB = p.symbol("colb", DOUBLE);
            return p.project(Assignments.of(p.symbol("x"), symbolB.toSymbolReference()), p.tableScan(new TableHandle(new CatalogName(mockCatalog), new MockConnectorTableHandle(testSchemaTable), MockConnectorTransactionHandle.INSTANCE), ImmutableList.of(symbolA, symbolB), ImmutableMap.of(symbolA, columnHandleA, symbolB, columnHandleB)));
        }).withSession(testSessionBuilder().setCatalog(mockCatalog).setSchema(testSchema).build()).matches(strictProject(ImmutableMap.of("expr", PlanMatchPattern.expression("COLB")), tableScan(new MockConnectorTableHandle(testSchemaTable, TupleDomain.all(), Optional.of(ImmutableList.of(columnHandleB)))::equals, TupleDomain.all(), ImmutableMap.of("COLB", columnHandleB::equals))));
    }
}
Also used : ColumnMetadata(io.trino.spi.connector.ColumnMetadata) PlanMatchPattern(io.trino.sql.planner.assertions.PlanMatchPattern) Test(org.testng.annotations.Test) Variable(io.trino.spi.expression.Variable) CatalogName(io.trino.connector.CatalogName) ImmutableList(com.google.common.collect.ImmutableList) MockConnectorFactory(io.trino.connector.MockConnectorFactory) TpchTableHandle(io.trino.plugin.tpch.TpchTableHandle) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) Symbol(io.trino.sql.planner.Symbol) TINY_SCHEMA_NAME(io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) ImmutableMap(com.google.common.collect.ImmutableMap) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest) TINY_SCALE_FACTOR(io.trino.plugin.tpch.TpchMetadata.TINY_SCALE_FACTOR) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Domain(io.trino.spi.predicate.Domain) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Assignments(io.trino.sql.planner.plan.Assignments) ConnectorSession(io.trino.spi.connector.ConnectorSession) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) TupleDomain(io.trino.spi.predicate.TupleDomain) SchemaTableName(io.trino.spi.connector.SchemaTableName) TpchTransactionHandle(io.trino.plugin.tpch.TpchTransactionHandle) DOUBLE(io.trino.spi.type.DoubleType.DOUBLE) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorTransactionHandle(io.trino.connector.MockConnectorTransactionHandle) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) TableHandle(io.trino.metadata.TableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) PlanMatchPattern.strictProject(io.trino.sql.planner.assertions.PlanMatchPattern.strictProject) Assignment(io.trino.spi.connector.Assignment) PlanMatchPattern.strictConstrainedTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictConstrainedTableScan) Optional(java.util.Optional) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) PlanBuilder.expression(io.trino.sql.planner.iterative.rule.test.PlanBuilder.expression) DATE(io.trino.spi.type.DateType.DATE) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) PlanMatchPattern.strictTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictTableScan) ColumnHandle(io.trino.spi.connector.ColumnHandle) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) MockConnectorFactory(io.trino.connector.MockConnectorFactory) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) Symbol(io.trino.sql.planner.Symbol) SchemaTableName(io.trino.spi.connector.SchemaTableName) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) TpchTableHandle(io.trino.plugin.tpch.TpchTableHandle) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) TableHandle(io.trino.metadata.TableHandle) CatalogName(io.trino.connector.CatalogName) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Test(org.testng.annotations.Test) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest)

Example 5 with MockConnectorColumnHandle

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

the class TestPruneTableScanColumns method mockApplyProjection.

private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
    List<Variable> variables = projections.stream().map(Variable.class::cast).collect(toImmutableList());
    List<ColumnHandle> newColumns = variables.stream().map(variable -> assignments.get(variable.getName())).collect(toImmutableList());
    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)), projections, variables.stream().map(variable -> new Assignment(variable.getName(), assignments.get(variable.getName()), ((MockConnectorColumnHandle) assignments.get(variable.getName())).getType())).collect(toImmutableList()), false));
}
Also used : ColumnMetadata(io.trino.spi.connector.ColumnMetadata) PlanMatchPattern(io.trino.sql.planner.assertions.PlanMatchPattern) Test(org.testng.annotations.Test) Variable(io.trino.spi.expression.Variable) CatalogName(io.trino.connector.CatalogName) ImmutableList(com.google.common.collect.ImmutableList) MockConnectorFactory(io.trino.connector.MockConnectorFactory) TpchTableHandle(io.trino.plugin.tpch.TpchTableHandle) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) Symbol(io.trino.sql.planner.Symbol) TINY_SCHEMA_NAME(io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) ImmutableMap(com.google.common.collect.ImmutableMap) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest) TINY_SCALE_FACTOR(io.trino.plugin.tpch.TpchMetadata.TINY_SCALE_FACTOR) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Domain(io.trino.spi.predicate.Domain) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Assignments(io.trino.sql.planner.plan.Assignments) ConnectorSession(io.trino.spi.connector.ConnectorSession) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) TupleDomain(io.trino.spi.predicate.TupleDomain) SchemaTableName(io.trino.spi.connector.SchemaTableName) TpchTransactionHandle(io.trino.plugin.tpch.TpchTransactionHandle) DOUBLE(io.trino.spi.type.DoubleType.DOUBLE) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorTransactionHandle(io.trino.connector.MockConnectorTransactionHandle) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) TableHandle(io.trino.metadata.TableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) PlanMatchPattern.strictProject(io.trino.sql.planner.assertions.PlanMatchPattern.strictProject) Assignment(io.trino.spi.connector.Assignment) PlanMatchPattern.strictConstrainedTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictConstrainedTableScan) Optional(java.util.Optional) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) PlanBuilder.expression(io.trino.sql.planner.iterative.rule.test.PlanBuilder.expression) DATE(io.trino.spi.type.DateType.DATE) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) PlanMatchPattern.strictTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictTableScan) Assignment(io.trino.spi.connector.Assignment) ColumnHandle(io.trino.spi.connector.ColumnHandle) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) Variable(io.trino.spi.expression.Variable) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle)

Aggregations

MockConnectorColumnHandle (io.trino.connector.MockConnectorColumnHandle)6 MockConnectorTableHandle (io.trino.connector.MockConnectorTableHandle)6 ColumnHandle (io.trino.spi.connector.ColumnHandle)6 ImmutableList (com.google.common.collect.ImmutableList)5 Test (org.testng.annotations.Test)5 ImmutableMap (com.google.common.collect.ImmutableMap)4 CatalogName (io.trino.connector.CatalogName)4 MockConnectorFactory (io.trino.connector.MockConnectorFactory)4 MockConnectorTransactionHandle (io.trino.connector.MockConnectorTransactionHandle)4 TableHandle (io.trino.metadata.TableHandle)4 SchemaTableName (io.trino.spi.connector.SchemaTableName)4 Symbol (io.trino.sql.planner.Symbol)4 PlanMatchPattern.tableScan (io.trino.sql.planner.assertions.PlanMatchPattern.tableScan)4 BaseRuleTest (io.trino.sql.planner.iterative.rule.test.BaseRuleTest)4 List (java.util.List)4 Map (java.util.Map)4 Optional (java.util.Optional)4 Assignment (io.trino.spi.connector.Assignment)3 ConnectorExpression (io.trino.spi.expression.ConnectorExpression)3 Variable (io.trino.spi.expression.Variable)3