Examples with TestingColumnHandle io.trino.testing.TestingMetadata.TestingColumnHandle used on opensource projects

Example 6 with TestingColumnHandle

use of io.trino.testing.TestingMetadata.TestingColumnHandle in project trino by trinodb.

the class TestSourcePartitionedScheduler method testDynamicFiltersUnblockedOnBlockedBuildSource.

@Test
public void testDynamicFiltersUnblockedOnBlockedBuildSource() {
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(nodeManager, new NodeSchedulerConfig().setIncludeCoordinator(false), nodeTaskMap));
    DynamicFilterService dynamicFilterService = new DynamicFilterService(metadata, functionManager, typeOperators, new DynamicFilterConfig());
    dynamicFilterService.registerQuery(QUERY_ID, TEST_SESSION, ImmutableSet.of(DYNAMIC_FILTER_ID), ImmutableSet.of(DYNAMIC_FILTER_ID), ImmutableSet.of(DYNAMIC_FILTER_ID));
    StageScheduler scheduler = newSourcePartitionedSchedulerAsStageScheduler(stage, TABLE_SCAN_NODE_ID, new ConnectorAwareSplitSource(CONNECTOR_ID, createBlockedSplitSource()), new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(session, Optional.of(CONNECTOR_ID)), stage::getAllTasks), 2, dynamicFilterService, new TableExecuteContextManager(), () -> true);
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(QUERY_ID, ImmutableList.of(new DynamicFilters.Descriptor(DYNAMIC_FILTER_ID, symbol.toSymbolReference())), ImmutableMap.of(symbol, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    // make sure dynamic filtering collecting task was created immediately
    assertEquals(stage.getState(), PLANNED);
    scheduler.start();
    assertEquals(stage.getAllTasks().size(), 1);
    assertEquals(stage.getState(), SCHEDULING);
    // make sure dynamic filter is initially blocked
    assertFalse(dynamicFilter.isBlocked().isDone());
    // make sure dynamic filter is unblocked due to build side source tasks being blocked
    ScheduleResult scheduleResult = scheduler.schedule();
    assertTrue(dynamicFilter.isBlocked().isDone());
    // no new probe splits should be scheduled
    assertEquals(scheduleResult.getSplitsScheduled(), 0);
}

Example 7 with TestingColumnHandle

use of io.trino.testing.TestingMetadata.TestingColumnHandle in project trino by trinodb.

the class TestFaultTolerantStageScheduler method createPlanFragment.

private PlanFragment createPlanFragment() {
    Symbol probeColumnSymbol = new Symbol("probe_column");
    Symbol buildColumnSymbol = new Symbol("build_column");
    TableScanNode tableScan = new TableScanNode(TABLE_SCAN_NODE_ID, TEST_TABLE_HANDLE, ImmutableList.of(probeColumnSymbol), ImmutableMap.of(probeColumnSymbol, new TestingColumnHandle("column")), TupleDomain.none(), Optional.empty(), false, Optional.empty());
    RemoteSourceNode remoteSource = new RemoteSourceNode(new PlanNodeId("remote_source_id"), ImmutableList.of(SOURCE_FRAGMENT_ID_1, SOURCE_FRAGMENT_ID_2), ImmutableList.of(buildColumnSymbol), Optional.empty(), REPLICATE, TASK);
    return new PlanFragment(FRAGMENT_ID, new JoinNode(new PlanNodeId("join_id"), INNER, tableScan, remoteSource, ImmutableList.of(), tableScan.getOutputSymbols(), remoteSource.getOutputSymbols(), false, Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(REPLICATED), Optional.empty(), ImmutableMap.of(), Optional.empty()), ImmutableMap.of(probeColumnSymbol, VARCHAR, buildColumnSymbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(TABLE_SCAN_NODE_ID), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(probeColumnSymbol, buildColumnSymbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
}

Example 8 with TestingColumnHandle

use of io.trino.testing.TestingMetadata.TestingColumnHandle in project trino by trinodb.

the class TestEffectivePredicateExtractor method setUp.

@BeforeMethod
public void setUp() {
    scanAssignments = ImmutableMap.<Symbol, ColumnHandle>builder().put(A, new TestingColumnHandle("a")).put(B, new TestingColumnHandle("b")).put(C, new TestingColumnHandle("c")).put(D, new TestingColumnHandle("d")).put(E, new TestingColumnHandle("e")).put(F, new TestingColumnHandle("f")).put(R, new TestingColumnHandle("r")).buildOrThrow();
    Map<Symbol, ColumnHandle> assignments = Maps.filterKeys(scanAssignments, Predicates.in(ImmutableList.of(A, B, C, D, E, F)));
    baseTableScan = TableScanNode.newInstance(newId(), makeTableHandle(TupleDomain.all()), ImmutableList.copyOf(assignments.keySet()), assignments, false, Optional.empty());
    expressionNormalizer = new ExpressionIdentityNormalizer();
}