Examples with PartitioningScheme io.trino.sql.planner.PartitioningScheme used on opensource projects

Example 1 with PartitioningScheme

use of io.trino.sql.planner.PartitioningScheme in project trino by trinodb.

the class AddIntermediateAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Lookup lookup = context.getLookup();
    PlanNodeIdAllocator idAllocator = context.getIdAllocator();
    Session session = context.getSession();
    Optional<PlanNode> rewrittenSource = recurseToPartial(lookup.resolve(aggregation.getSource()), lookup, idAllocator);
    if (rewrittenSource.isEmpty()) {
        return Result.empty();
    }
    PlanNode source = rewrittenSource.get();
    if (getTaskConcurrency(session) > 1) {
        source = ExchangeNode.partitionedExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source, new PartitioningScheme(Partitioning.create(FIXED_ARBITRARY_DISTRIBUTION, ImmutableList.of()), source.getOutputSymbols()));
        source = new AggregationNode(idAllocator.getNextId(), source, inputsAsOutputs(aggregation.getAggregations()), aggregation.getGroupingSets(), aggregation.getPreGroupedSymbols(), AggregationNode.Step.INTERMEDIATE, aggregation.getHashSymbol(), aggregation.getGroupIdSymbol());
        source = ExchangeNode.gatheringExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source);
    }
    return Result.ofPlanNode(aggregation.replaceChildren(ImmutableList.of(source)));
}

Example 2 with PartitioningScheme

use of io.trino.sql.planner.PartitioningScheme in project trino by trinodb.

the class PruneTableWriterSourceColumns method apply.

@Override
public Result apply(TableWriterNode tableWriterNode, Captures captures, Context context) {
    ImmutableSet.Builder<Symbol> requiredInputs = ImmutableSet.<Symbol>builder().addAll(tableWriterNode.getColumns());
    if (tableWriterNode.getPartitioningScheme().isPresent()) {
        PartitioningScheme partitioningScheme = tableWriterNode.getPartitioningScheme().get();
        partitioningScheme.getPartitioning().getColumns().forEach(requiredInputs::add);
        partitioningScheme.getHashColumn().ifPresent(requiredInputs::add);
    }
    if (tableWriterNode.getStatisticsAggregation().isPresent()) {
        StatisticAggregations aggregations = tableWriterNode.getStatisticsAggregation().get();
        requiredInputs.addAll(aggregations.getGroupingSymbols());
        aggregations.getAggregations().values().stream().map(SymbolsExtractor::extractUnique).forEach(requiredInputs::addAll);
    }
    return restrictChildOutputs(context.getIdAllocator(), tableWriterNode, requiredInputs.build()).map(Result::ofPlanNode).orElse(Result.empty());
}

Example 3 with PartitioningScheme

use of io.trino.sql.planner.PartitioningScheme in project trino by trinodb.

the class TestSqlStage method createExchangePlanFragment.

private static PlanFragment createExchangePlanFragment() {
    PlanNode planNode = new RemoteSourceNode(new PlanNodeId("exchange"), ImmutableList.of(new PlanFragmentId("source")), ImmutableList.of(new Symbol("column")), Optional.empty(), REPARTITION, RetryPolicy.NONE);
    ImmutableMap.Builder<Symbol, Type> types = ImmutableMap.builder();
    for (Symbol symbol : planNode.getOutputSymbols()) {
        types.put(symbol, VARCHAR);
    }
    return new PlanFragment(new PlanFragmentId("exchange_fragment_id"), planNode, types.buildOrThrow(), SOURCE_DISTRIBUTION, ImmutableList.of(planNode.getId()), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), planNode.getOutputSymbols()), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
}

Example 4 with PartitioningScheme

use of io.trino.sql.planner.PartitioningScheme in project trino by trinodb.

the class MockRemoteTaskFactory method createTableScanTask.

public MockRemoteTask createTableScanTask(TaskId taskId, InternalNode newNode, List<Split> splits, PartitionedSplitCountTracker partitionedSplitCountTracker) {
    Symbol symbol = new Symbol("column");
    PlanNodeId sourceId = new PlanNodeId("sourceId");
    PlanFragment testFragment = new PlanFragment(new PlanFragmentId("test"), TableScanNode.newInstance(sourceId, TEST_TABLE_HANDLE, ImmutableList.of(symbol), ImmutableMap.of(symbol, new TestingColumnHandle("column")), false, Optional.empty()), ImmutableMap.of(symbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(sourceId), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
    ImmutableMultimap.Builder<PlanNodeId, Split> initialSplits = ImmutableMultimap.builder();
    for (Split sourceSplit : splits) {
        initialSplits.put(sourceId, sourceSplit);
    }
    return createRemoteTask(TEST_SESSION, taskId, newNode, testFragment, initialSplits.build(), createInitialEmptyOutputBuffers(BROADCAST), partitionedSplitCountTracker, ImmutableSet.of(), true);
}

Example 5 with PartitioningScheme

use of io.trino.sql.planner.PartitioningScheme in project trino by trinodb.

the class TestApplyPreferredTableWriterPartitioning method testThresholdWithNullFraction.

@Test
public void testThresholdWithNullFraction() {
    // Null value in partition column should increase the number of partitions by 1
    PlanNodeStatsEstimate stats = PlanNodeStatsEstimate.builder().addSymbolStatistics(ImmutableMap.of(new Symbol("col_one"), new SymbolStatsEstimate(0, 0, .5, 0, 49))).build();
    assertPreferredPartitioning(new PartitioningScheme(Partitioning.create(FIXED_HASH_DISTRIBUTION, ImmutableList.of(new Symbol("col_one"))), ImmutableList.of(new Symbol("col_one")))).withSession(SESSION_WITH_PREFERRED_PARTITIONING_DEFAULT_THRESHOLD).overrideStats(NODE_ID, stats).matches(SUCCESSFUL_MATCH);
}