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

Example 26 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method createMockFactory.

private MockConnectorFactory createMockFactory(Optional<MockConnectorFactory.ApplyTableScanRedirect> applyTableScanRedirect) {
    MockConnectorFactory.Builder builder = MockConnectorFactory.builder().withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(SOURCE_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(SOURCE_COLUMN_NAME_A, VARCHAR), new ColumnMetadata(SOURCE_COLUMN_NAME_B, VARCHAR));
        } else if (schemaTableName.equals(DESTINATION_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(DESTINATION_COLUMN_NAME_A, VARCHAR), new ColumnMetadata(DESTINATION_COLUMN_NAME_B, VARCHAR), new ColumnMetadata(DESTINATION_COLUMN_NAME_C, BIGINT), new ColumnMetadata(DESTINATION_COLUMN_NAME_D, BOGUS));
        }
        throw new IllegalArgumentException();
    });
    applyTableScanRedirect.ifPresent(builder::withApplyTableScanRedirect);
    return builder.build();
}

Example 27 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method testDoesNotFire.

@Test
public void testDoesNotFire() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory mockFactory = createMockFactory(Optional.empty());
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new ApplyTableScanRedirection(ruleTester.getPlannerContext())).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A));
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 28 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method testMismatchedTypesWithCoercion.

@Test
public void testMismatchedTypesWithCoercion() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // make the mock connector return a table scan on different table
        ApplyTableScanRedirect applyTableScanRedirect = getMockApplyRedirect(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_NAME_C));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        LocalQueryRunner runner = ruleTester.getQueryRunner();
        runner.createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new ApplyTableScanRedirection(ruleTester.getPlannerContext())).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A));
        }).withSession(MOCK_SESSION).matches(project(ImmutableMap.of("COL", expression("CAST(DEST_COL AS VARCHAR)")), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_C::equals))));
    }
}

Example 29 with MockConnectorFactory

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

the class TestPushProjectionIntoTableScan method testPushProjection.

@Test
public void testPushProjection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // Building context for input
        String columnName = "col0";
        Type columnType = ROW_TYPE;
        Symbol baseColumn = new Symbol(columnName);
        ColumnHandle columnHandle = new TpchColumnHandle(columnName, columnType);
        // Create catalog with applyProjection enabled
        MockConnectorFactory factory = createMockFactory(ImmutableMap.of(columnName, columnHandle), Optional.of(this::mockApplyProjection));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, factory, ImmutableMap.of());
        TypeAnalyzer typeAnalyzer = createTestingTypeAnalyzer(ruleTester.getPlannerContext());
        // Prepare project node symbols and types
        Symbol identity = new Symbol("symbol_identity");
        Symbol dereference = new Symbol("symbol_dereference");
        Symbol constant = new Symbol("symbol_constant");
        Symbol call = new Symbol("symbol_call");
        ImmutableMap<Symbol, Type> types = ImmutableMap.of(baseColumn, ROW_TYPE, identity, ROW_TYPE, dereference, BIGINT, constant, BIGINT, call, VARCHAR);
        // Prepare project node assignments
        ImmutableMap<Symbol, Expression> inputProjections = ImmutableMap.of(identity, baseColumn.toSymbolReference(), dereference, new SubscriptExpression(baseColumn.toSymbolReference(), new LongLiteral("1")), constant, new LongLiteral("5"), call, new FunctionCall(QualifiedName.of("STARTS_WITH"), ImmutableList.of(new StringLiteral("abc"), new StringLiteral("ab"))));
        // Compute expected symbols after applyProjection
        TransactionId transactionId = ruleTester.getQueryRunner().getTransactionManager().beginTransaction(false);
        Session session = MOCK_SESSION.beginTransactionId(transactionId, ruleTester.getQueryRunner().getTransactionManager(), ruleTester.getQueryRunner().getAccessControl());
        ImmutableMap<Symbol, String> connectorNames = inputProjections.entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, e -> translate(session, e.getValue(), typeAnalyzer, viewOf(types), ruleTester.getPlannerContext()).get().toString()));
        ImmutableMap<Symbol, String> newNames = ImmutableMap.of(identity, "projected_variable_" + connectorNames.get(identity), dereference, "projected_dereference_" + connectorNames.get(dereference), constant, "projected_constant_" + connectorNames.get(constant), call, "projected_call_" + connectorNames.get(call));
        Map<String, ColumnHandle> expectedColumns = newNames.entrySet().stream().collect(toImmutableMap(Map.Entry::getValue, e -> column(e.getValue(), types.get(e.getKey()))));
        ruleTester.assertThat(createRule(ruleTester)).on(p -> {
            // Register symbols
            types.forEach((symbol, type) -> p.symbol(symbol.getName(), type));
            return p.project(new Assignments(inputProjections), p.tableScan(tableScan -> tableScan.setTableHandle(TEST_TABLE_HANDLE).setSymbols(ImmutableList.copyOf(types.keySet())).setAssignments(types.keySet().stream().collect(Collectors.toMap(Function.identity(), v -> columnHandle))).setStatistics(Optional.of(PlanNodeStatsEstimate.builder().setOutputRowCount(42).addSymbolStatistics(baseColumn, SymbolStatsEstimate.builder().setNullsFraction(0).setDistinctValuesCount(33).build()).build()))));
        }).withSession(MOCK_SESSION).matches(project(newNames.entrySet().stream().collect(toImmutableMap(e -> e.getKey().getName(), e -> expression(symbolReference(e.getValue())))), tableScan(new MockConnectorTableHandle(new SchemaTableName(TEST_SCHEMA, "projected_" + TEST_TABLE), TupleDomain.all(), Optional.of(ImmutableList.copyOf(expectedColumns.values())))::equals, TupleDomain.all(), expectedColumns.entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, e -> e.getValue()::equals)), Optional.of(PlanNodeStatsEstimate.builder().setOutputRowCount(42).addSymbolStatistics(new Symbol(newNames.get(constant)), SymbolStatsEstimate.builder().setDistinctValuesCount(1).setNullsFraction(0).setLowValue(5).setHighValue(5).build()).addSymbolStatistics(new Symbol(newNames.get(call).toLowerCase(ENGLISH)), SymbolStatsEstimate.builder().setDistinctValuesCount(1).setNullsFraction(0).build()).addSymbolStatistics(new Symbol(newNames.get(identity)), SymbolStatsEstimate.builder().setDistinctValuesCount(33).setNullsFraction(0).build()).addSymbolStatistics(new Symbol(newNames.get(dereference)), SymbolStatsEstimate.unknown()).build())::equals)));
    }
}

Example 30 with MockConnectorFactory

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

the class TestPushProjectionIntoTableScan method testPartitioningChanged.

@Test
public void testPartitioningChanged() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        String columnName = "col0";
        ColumnHandle columnHandle = new TpchColumnHandle(columnName, VARCHAR);
        // Create catalog with applyProjection enabled
        MockConnectorFactory factory = createMockFactory(ImmutableMap.of(columnName, columnHandle), Optional.of(this::mockApplyProjection));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, factory, ImmutableMap.of());
        assertThatThrownBy(() -> ruleTester.assertThat(createRule(ruleTester)).on(p -> p.project(Assignments.of(), p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(p.symbol("col", VARCHAR)), ImmutableMap.of(p.symbol("col", VARCHAR), columnHandle), Optional.of(true)))).withSession(MOCK_SESSION).matches(anyTree())).hasMessage("Partitioning must not change after projection is pushed down");
    }
}