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

Example 26 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle 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 27 with MockConnectorTableHandle

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

the class TestPushTopNIntoTableScan method testPushPartialTopNIntoTableScanNotGuaranteed.

@Test
public void testPushPartialTopNIntoTableScanNotGuaranteed() {
    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), TopNNode.Step.PARTIAL, 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.PARTIAL, tableScan(connectorHandle::equals, TupleDomain.all(), ImmutableMap.of(dimensionName, dimensionColumn::equals, metricName, metricColumn::equals))));
    }
}

Example 28 with MockConnectorTableHandle

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

the class TestPushTopNIntoTableScan method testPushSingleTopNIntoTableScan.

@Test
public void testPushSingleTopNIntoTableScan() {
    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), 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 29 with MockConnectorTableHandle

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

the class TestPushJoinIntoTableScan method testPushJoinIntoTableScan.

@Test(dataProvider = "testPushJoinIntoTableScanParams")
public void testPushJoinIntoTableScan(JoinNode.Type joinType, Optional<ComparisonExpression.Operator> filterComparisonOperator) {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory connectorFactory = createMockConnectorFactory((session, applyJoinType, left, right, joinConditions, leftAssignments, rightAssignments) -> {
            assertThat(((MockConnectorTableHandle) left).getTableName()).isEqualTo(TABLE_A_SCHEMA_TABLE_NAME);
            assertThat(((MockConnectorTableHandle) right).getTableName()).isEqualTo(TABLE_B_SCHEMA_TABLE_NAME);
            Assertions.assertThat(applyJoinType).isEqualTo(toSpiJoinType(joinType));
            JoinCondition.Operator expectedOperator = filterComparisonOperator.map(this::getConditionOperator).orElse(JoinCondition.Operator.EQUAL);
            Assertions.assertThat(joinConditions).containsExactly(new JoinCondition(expectedOperator, COLUMN_A1_VARIABLE, COLUMN_B1_VARIABLE));
            return 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));
            TableScanNode right = p.tableScan(TABLE_B_HANDLE, ImmutableList.of(columnB1Symbol), ImmutableMap.of(columnB1Symbol, COLUMN_B1_HANDLE));
            if (filterComparisonOperator.isEmpty()) {
                return p.join(joinType, left, right, new JoinNode.EquiJoinClause(columnA1Symbol, columnB1Symbol));
            }
            return p.join(joinType, left, right, new ComparisonExpression(filterComparisonOperator.get(), columnA1Symbol.toSymbolReference(), columnB1Symbol.toSymbolReference()));
        }).withSession(MOCK_SESSION).matches(project(tableScan(JOIN_PUSHDOWN_SCHEMA_TABLE_NAME.getTableName())));
    }
}

Example 30 with MockConnectorTableHandle

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

the class TestPushDistinctLimitIntoTableScan method testPushDistinct.

@Test
public void testPushDistinct() {
    AtomicInteger applyCallCounter = new AtomicInteger();
    AtomicReference<List<AggregateFunction>> applyAggregates = new AtomicReference<>();
    AtomicReference<Map<String, ColumnHandle>> applyAssignments = new AtomicReference<>();
    AtomicReference<List<List<ColumnHandle>>> applyGroupingSets = new AtomicReference<>();
    testApplyAggregation = (session, handle, aggregates, assignments, groupingSets) -> {
        applyCallCounter.incrementAndGet();
        applyAggregates.set(List.copyOf(aggregates));
        applyAssignments.set(Map.copyOf(assignments));
        applyGroupingSets.set(groupingSets.stream().map(List::copyOf).collect(toUnmodifiableList()));
        return Optional.of(new AggregationApplicationResult<>(new MockConnectorTableHandle(new SchemaTableName("mock_schema", "mock_nation_aggregated")), List.of(), List.of(), Map.of(), false));
    };
    MockConnectorColumnHandle regionkeyHandle = new MockConnectorColumnHandle("regionkey", BIGINT);
    tester().assertThat(rule).on(p -> {
        Symbol regionkeySymbol = p.symbol("regionkey_symbol");
        return p.distinctLimit(43, List.of(regionkeySymbol), p.tableScan(tableHandle, ImmutableList.of(regionkeySymbol), ImmutableMap.of(regionkeySymbol, regionkeyHandle)));
    }).matches(limit(43, project(tableScan("mock_nation_aggregated"))));
    assertThat(applyCallCounter).as("applyCallCounter").hasValue(1);
    assertThat(applyAggregates).as("applyAggregates").hasValue(List.of());
    assertThat(applyAssignments).as("applyAssignments").hasValue(Map.of("regionkey_symbol", regionkeyHandle));
    assertThat(applyGroupingSets).as("applyGroupingSets").hasValue(List.of(List.of(regionkeyHandle)));
}