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

Example 6 with TestingColumnHandle

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

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

the class TestDynamicFilterService method testShortCircuitOnAllTupleDomain.

@Test
public void testShortCircuitOnAllTupleDomain() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1), ImmutableSet.of(filterId1), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 2);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    // dynamic filter is initially blocked
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertFalse(dynamicFilter.isBlocked().isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 1, 0), ImmutableMap.of(filterId1, Domain.all(INTEGER)));
    // dynamic filter should be unblocked and completed
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertTrue(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isBlocked().isDone());
}

Example 8 with TestingColumnHandle

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

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

the class TestRuleTester method testReportNoFireWithTableScan.

@Test
public void testReportNoFireWithTableScan() {
    try (RuleTester tester = defaultRuleTester()) {
        RuleAssert ruleAssert = tester.assertThat(rule("testReportNoFireWithTableScan rule", Pattern.typeOf(PlanNode.class), (node, captures, context) -> Result.empty())).on(p -> p.tableScan(new TableHandle(tester.getCurrentConnectorId(), new TpchTableHandle("sf1", "nation", 1.0), TestingTransactionHandle.create()), List.of(p.symbol("x")), Map.of(p.symbol("x"), new TestingColumnHandle("column"))));
        PlanMatchPattern expected = values(List.of("whatever"), List.of());
        assertThatThrownBy(() -> ruleAssert.matches(expected)).isInstanceOf(AssertionError.class).hasMessageMatching("testReportNoFireWithTableScan rule did not fire for:\n" + "(?s:.*)" + "\\QEstimates: {rows: 25 (225B), cpu: 225, memory: 0B, network: 0B}\\E\n" + "(?s:.*)");
    }
}

Example 10 with TestingColumnHandle

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

the class TestCardinalityExtractorPlanVisitor method testAggregation.

@Test
public void testAggregation() {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), dummyMetadata(), TEST_SESSION);
    Symbol symbol = planBuilder.symbol("symbol");
    ColumnHandle columnHandle = new TestingColumnHandle("column");
    // single default aggregation
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet().source(planBuilder.values(10)))), Range.singleton(1L));
    // multiple grouping sets with default aggregation with source that produces arbitrary number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.tableScan(ImmutableList.of(symbol), ImmutableMap.of(symbol, columnHandle))))), Range.atLeast(1L));
    // multiple grouping sets with default aggregation with source that produces exact number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.values(10, symbol)))), Range.closed(1L, 10L));
    // multiple grouping sets with default aggregation with source that produces no rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.values(0, symbol)))), Range.singleton(1L));
    // single non-default aggregation with source that produces arbitrary number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.tableScan(ImmutableList.of(symbol), ImmutableMap.of(symbol, columnHandle))))), Range.atLeast(0L));
    // single non-default aggregation with source that produces at least single row
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.values(10, symbol)))), Range.closed(1L, 10L));
    // single non-default aggregation with source that produces no rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.values(0, symbol)))), Range.singleton(0L));
}