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

Example 6 with MockConnectorFactory

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

the class TestPushJoinIntoTableScan method testPushJoinIntoTableScanPreservesEnforcedConstraint.

@Test(dataProvider = "testPushJoinIntoTableScanPreservesEnforcedConstraintParams")
public void testPushJoinIntoTableScanPreservesEnforcedConstraint(JoinNode.Type joinType, TupleDomain<ColumnHandle> leftConstraint, TupleDomain<ColumnHandle> rightConstraint, TupleDomain<Predicate<ColumnHandle>> expectedConstraint) {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory connectorFactory = createMockConnectorFactory((session, applyJoinType, left, right, joinConditions, leftAssignments, rightAssignments) -> Optional.of(new JoinApplicationResult<>(JOIN_CONNECTOR_TABLE_HANDLE, JOIN_TABLE_A_COLUMN_MAPPING, JOIN_TABLE_B_COLUMN_MAPPING, false)));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, connectorFactory, ImmutableMap.of());
        ruleTester.assertThat(new PushJoinIntoTableScan(ruleTester.getMetadata())).on(p -> {
            Symbol columnA1Symbol = p.symbol(COLUMN_A1);
            Symbol columnA2Symbol = p.symbol(COLUMN_A2);
            Symbol columnB1Symbol = p.symbol(COLUMN_B1);
            TableScanNode left = p.tableScan(TABLE_A_HANDLE, ImmutableList.of(columnA1Symbol, columnA2Symbol), ImmutableMap.of(columnA1Symbol, COLUMN_A1_HANDLE, columnA2Symbol, COLUMN_A2_HANDLE), leftConstraint);
            TableScanNode right = p.tableScan(TABLE_B_HANDLE, ImmutableList.of(columnB1Symbol), ImmutableMap.of(columnB1Symbol, COLUMN_B1_HANDLE), rightConstraint);
            return p.join(joinType, left, right, new JoinNode.EquiJoinClause(columnA1Symbol, columnB1Symbol));
        }).withSession(MOCK_SESSION).matches(project(tableScan(tableHandle -> JOIN_PUSHDOWN_SCHEMA_TABLE_NAME.equals(((MockConnectorTableHandle) tableHandle).getTableName()), expectedConstraint, ImmutableMap.of())));
    }
}

Example 7 with MockConnectorFactory

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

the class TestPushJoinIntoTableScan method testPushJoinIntoTableScanDoesNotFireForDifferentCatalogs.

@Test
public void testPushJoinIntoTableScanDoesNotFireForDifferentCatalogs() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory connectorFactory = createMockConnectorFactory((session, applyJoinType, left, right, joinConditions, leftAssignments, rightAssignments) -> {
            throw new IllegalStateException("applyJoin should not be called!");
        });
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, connectorFactory, ImmutableMap.of());
        ruleTester.getQueryRunner().createCatalog("another_catalog", "mock", ImmutableMap.of());
        TableHandle tableBHandleAnotherCatalog = createTableHandle(new MockConnectorTableHandle(new SchemaTableName(SCHEMA, TABLE_B)), "another_catalog");
        ruleTester.assertThat(new PushJoinIntoTableScan(ruleTester.getMetadata())).on(p -> {
            Symbol columnA1Symbol = p.symbol(COLUMN_A1);
            Symbol columnA2Symbol = p.symbol(COLUMN_A2);
            Symbol columnB1Symbol = p.symbol(COLUMN_B1);
            TableScanNode left = p.tableScan(TABLE_A_HANDLE, ImmutableList.of(columnA1Symbol, columnA2Symbol), ImmutableMap.of(columnA1Symbol, COLUMN_A1_HANDLE, columnA2Symbol, COLUMN_A2_HANDLE));
            TableScanNode right = p.tableScan(tableBHandleAnotherCatalog, ImmutableList.of(columnB1Symbol), ImmutableMap.of(columnB1Symbol, COLUMN_B1_HANDLE));
            return p.join(INNER, left, right, new JoinNode.EquiJoinClause(columnA1Symbol, columnB1Symbol));
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 8 with MockConnectorFactory

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

the class TestPushTopNIntoTableScan method testPushPartialTopNIntoTableScan.

@Test
public void testPushPartialTopNIntoTableScan() {
    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.PARTIAL, 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 9 with MockConnectorFactory

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

the class TestPushTopNIntoTableScan method testPushSingleTopNIntoTableScanNotGuaranteed.

@Test
public void testPushSingleTopNIntoTableScanNotGuaranteed() {
    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, false, 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), p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(dimension, metric), ImmutableMap.of(dimension, dimensionColumn, metric, metricColumn)));
        }).withSession(MOCK_SESSION).matches(topN(1, ImmutableList.of(sort(dimensionName, ASCENDING, FIRST)), TopNNode.Step.SINGLE, tableScan(connectorHandle::equals, TupleDomain.all(), ImmutableMap.of(dimensionName, dimensionColumn::equals, metricName, metricColumn::equals))));
    }
}

Example 10 with MockConnectorFactory

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