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

Example 6 with MockConnectorColumnHandle

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