Examples with ColumnHandle io.trino.spi.connector.ColumnHandle used on opensource projects

Example 76 with ColumnHandle

use of io.trino.spi.connector.ColumnHandle in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterCancellation.

@Test
public void testDynamicFilterCancellation() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId = new DynamicFilterId("df");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId = new StageId(queryId, 0);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId), ImmutableSet.of(filterId), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 0, 2);
    ColumnHandle column = new TestingColumnHandle("probeColumnA");
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId, df1)), ImmutableMap.of(symbol1, column), symbolAllocator.getTypes());
    assertFalse(dynamicFilter.isBlocked().isDone());
    assertFalse(dynamicFilter.isComplete());
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.all());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 1L)));
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.all());
    // DynamicFilter future cancellation should not affect DynamicFilterService
    CompletableFuture<?> isBlocked = dynamicFilter.isBlocked();
    assertFalse(isBlocked.isDone());
    assertFalse(isBlocked.cancel(false));
    assertFalse(dynamicFilter.isBlocked().isDone());
    assertFalse(dynamicFilter.isComplete());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 2L)));
    assertTrue(isBlocked.isDone());
    assertTrue(dynamicFilter.isComplete());
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(column, multipleValues(INTEGER, ImmutableList.of(1L, 2L)))));
}

Example 77 with ColumnHandle

use of io.trino.spi.connector.ColumnHandle in project trino by trinodb.

the class TestHttpRemoteTask method testOutboundDynamicFilters.

@Test(timeOut = 30_000)
public void testOutboundDynamicFilters() throws Exception {
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", BIGINT);
    SymbolReference df1 = symbol1.toSymbolReference();
    SymbolReference df2 = symbol2.toSymbolReference();
    ColumnHandle handle1 = new TestingColumnHandle("column1");
    ColumnHandle handle2 = new TestingColumnHandle("column2");
    QueryId queryId = new QueryId("test");
    TestingTaskResource testingTaskResource = new TestingTaskResource(new AtomicLong(System.nanoTime()), FailureScenario.NO_FAILURE);
    DynamicFilterService dynamicFilterService = new DynamicFilterService(PLANNER_CONTEXT.getMetadata(), PLANNER_CONTEXT.getFunctionManager(), new TypeOperators(), newDirectExecutorService());
    dynamicFilterService.registerQuery(queryId, TEST_SESSION, ImmutableSet.of(filterId1, filterId2), ImmutableSet.of(filterId1, filterId2), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(new StageId(queryId, 1), 0, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, handle1, symbol2, handle2), symbolAllocator.getTypes());
    // make sure initial dynamic filter is collected
    CompletableFuture<?> future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L))));
    // Create remote task after dynamic filter is created to simulate new nodes joining
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, dynamicFilterService);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of(filterId1, filterId2));
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 1L);
    // schedule a couple of splits to trigger task updates
    addSplit(remoteTask, testingTaskResource, 1);
    addSplit(remoteTask, testingTaskResource, 2);
    // make sure dynamic filter was sent in task updates only once
    assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L);
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 3L);
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L), handle2, Domain.singleValue(BIGINT, 2L))));
    // dynamic filter should be sent even though there were no further splits scheduled
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 2L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 4L);
    // previously sent dynamic filter should not be repeated
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    httpRemoteTaskFactory.stop();
    dynamicFilterService.stop();
}

Example 78 with ColumnHandle

use of io.trino.spi.connector.ColumnHandle in project trino by trinodb.

the class Util method domainsMatch.

static boolean domainsMatch(TupleDomain<Predicate<ColumnHandle>> expected, TupleDomain<ColumnHandle> actual) {
    Optional<Map<Predicate<ColumnHandle>, Domain>> expectedDomains = expected.getDomains();
    Optional<Map<ColumnHandle, Domain>> actualDomains = actual.getDomains();
    if (expectedDomains.isPresent() != actualDomains.isPresent()) {
        return false;
    }
    if (expectedDomains.isPresent()) {
        if (expectedDomains.get().size() != actualDomains.get().size()) {
            return false;
        }
        for (Map.Entry<Predicate<ColumnHandle>, Domain> entry : expectedDomains.get().entrySet()) {
            // There should be exactly one column matching the expected column matcher
            ColumnHandle actualColumn = Iterables.getOnlyElement(actualDomains.get().keySet().stream().filter(x -> entry.getKey().test(x)).collect(toImmutableList()));
            if (!actualDomains.get().get(actualColumn).contains(entry.getValue())) {
                return false;
            }
        }
    }
    return true;
}

Example 79 with ColumnHandle

use of io.trino.spi.connector.ColumnHandle 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 80 with ColumnHandle

use of io.trino.spi.connector.ColumnHandle in project trino by trinodb.

the class TestDefaultJdbcMetadata method testAggregationPushdownForTableHandle.

@Test
public void testAggregationPushdownForTableHandle() {
    ConnectorSession session = TestingConnectorSession.builder().setPropertyMetadata(new JdbcMetadataSessionProperties(new JdbcMetadataConfig().setAggregationPushdownEnabled(true), Optional.empty()).getSessionProperties()).build();
    ColumnHandle groupByColumn = metadata.getColumnHandles(session, tableHandle).get("text");
    Function<ConnectorTableHandle, Optional<AggregationApplicationResult<ConnectorTableHandle>>> applyAggregation = handle -> metadata.applyAggregation(session, handle, ImmutableList.of(new AggregateFunction("count", BIGINT, List.of(), List.of(), false, Optional.empty())), ImmutableMap.of(), ImmutableList.of(ImmutableList.of(groupByColumn)));
    ConnectorTableHandle baseTableHandle = metadata.getTableHandle(session, new SchemaTableName("example", "numbers"));
    Optional<AggregationApplicationResult<ConnectorTableHandle>> aggregationResult = applyAggregation.apply(baseTableHandle);
    assertThat(aggregationResult).isPresent();
    SchemaTableName noAggregationPushdownTable = new SchemaTableName("example", "no_aggregation_pushdown");
    metadata.createTable(SESSION, new ConnectorTableMetadata(noAggregationPushdownTable, ImmutableList.of(new ColumnMetadata("text", VARCHAR))), false);
    ConnectorTableHandle noAggregationPushdownTableHandle = metadata.getTableHandle(session, noAggregationPushdownTable);
    aggregationResult = applyAggregation.apply(noAggregationPushdownTableHandle);
    assertThat(aggregationResult).isEmpty();
}