Examples with MockConnectorTableHandle io.trino.connector.MockConnectorTableHandle used on opensource projects

Example 11 with MockConnectorTableHandle

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

the class TestPushTopNIntoTableScan method testPushFinalTopNIntoTableScan.

/**
 * Ensure FINAL TopN can be pushed into table scan.
 * <p>
 * In case of TopN over outer join, TopN may become eligible for push down
 * only after PARTIAL TopN was pushed down and only then the join was
 * pushed down as well -- the connector may decide to accept Join pushdown
 * only after it learns there is TopN in play which limits results size.
 * <p>
 * Thus the optimization sequence can be:
 * <ol>
 * <li>Try to push Join into Table Scan -- connector rejects that (e.g. too big data set size)
 * <li>Create FINAL/PARTIAL TopN
 * <li>Push PARTIAL TopN through Outer Join
 * <li>Push PARTIAL TopN into Table Scan -- connector accepts that.
 * <li>Push Join into Table Scan -- connector now accepts join pushdown.
 * <li>Push FINAL TopN into Table Scan
 * </ol>
 */
@Test
public void testPushFinalTopNIntoTableScan() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorTableHandle connectorHandle = new MockConnectorTableHandle(TEST_SCHEMA_TABLE);
        // make the mock connector return a new connectorHandle
        MockConnectorFactory.ApplyTopN applyTopN = (session, handle, topNCount, sortItems, tableAssignments) -> Optional.of(new TopNApplicationResult<>(connectorHandle, true, false));
        MockConnectorFactory mockFactory = createMockFactory(assignments, Optional.of(applyTopN));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new PushTopNIntoTableScan(ruleTester.getMetadata())).on(p -> {
            Symbol dimension = p.symbol(dimensionName, VARCHAR);
            Symbol metric = p.symbol(metricName, BIGINT);
            return p.topN(1, ImmutableList.of(dimension), TopNNode.Step.FINAL, p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(dimension, metric), ImmutableMap.of(dimension, dimensionColumn, metric, metricColumn)));
        }).withSession(MOCK_SESSION).matches(tableScan(connectorHandle::equals, TupleDomain.all(), new HashMap<>()));
    }
}

Example 12 with MockConnectorTableHandle

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

the class TestApplyTableScanRedirection method testApplyTableScanRedirectionWithFilter.

@Test
public void testApplyTableScanRedirectionWithFilter() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // make the mock connector return a table scan on different table
        // source table handle has a pushed down predicate
        ApplyTableScanRedirect applyTableScanRedirect = getMockApplyRedirect(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_NAME_A, SOURCE_COLUMN_HANDLE_B, DESTINATION_COLUMN_NAME_B));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ApplyTableScanRedirection applyTableScanRedirection = new ApplyTableScanRedirection(ruleTester.getPlannerContext());
        TupleDomain<ColumnHandle> constraint = TupleDomain.withColumnDomains(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, singleValue(VARCHAR, utf8Slice("foo"))));
        ruleTester.assertThat(applyTableScanRedirection).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(createTableHandle(new MockConnectorTableHandle(SOURCE_TABLE, constraint, Optional.empty())), ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A), constraint);
        }).withSession(MOCK_SESSION).matches(filter("DEST_COL = VARCHAR 'foo'", tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals))));
        ruleTester.assertThat(applyTableScanRedirection).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_B, VARCHAR);
            return p.tableScan(createTableHandle(new MockConnectorTableHandle(SOURCE_TABLE, constraint, Optional.empty())), ImmutableList.of(column), // predicate on non-projected column
            ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_B), TupleDomain.all());
        }).withSession(MOCK_SESSION).matches(project(ImmutableMap.of("expr", expression("DEST_COL_B")), filter("DEST_COL_A = VARCHAR 'foo'", tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL_A", DESTINATION_COLUMN_HANDLE_A::equals, "DEST_COL_B", DESTINATION_COLUMN_HANDLE_B::equals)))));
    }
}

Example 13 with MockConnectorTableHandle

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

the class TestPushDistinctLimitIntoTableScan method init.

@BeforeClass
public void init() {
    rule = new PushDistinctLimitIntoTableScan(tester().getPlannerContext());
    tableHandle = new TableHandle(TEST_CATALOG, new MockConnectorTableHandle(new SchemaTableName("mock_schema", "mock_nation")), MockConnectorTransactionHandle.INSTANCE);
}

Example 14 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle 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))));
    }
}

Example 15 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle 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));
}