Examples with PartitionFunction io.trino.operator.PartitionFunction used on opensource projects

Example 1 with PartitionFunction

use of io.trino.operator.PartitionFunction in project trino by trinodb.

the class BenchmarkHashBuildAndJoinOperators method buildHash.

private static void buildHash(BuildContext buildContext, JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager, List<Integer> outputChannels, int partitionCount) {
    HashBuilderOperatorFactory hashBuilderOperatorFactory = new HashBuilderOperatorFactory(HASH_BUILD_OPERATOR_ID, TEST_PLAN_NODE_ID, lookupSourceFactoryManager, outputChannels, buildContext.getHashChannels(), buildContext.getHashChannel(), Optional.empty(), Optional.empty(), ImmutableList.of(), 10_000, new PagesIndex.TestingFactory(false), false, SingleStreamSpillerFactory.unsupportedSingleStreamSpillerFactory(), incrementalLoadFactorHashArraySizeSupplier(buildContext.getSession()));
    Operator[] operators = IntStream.range(0, partitionCount).mapToObj(i -> buildContext.createTaskContext().addPipelineContext(0, true, true, partitionCount > 1).addDriverContext()).map(hashBuilderOperatorFactory::createOperator).toArray(Operator[]::new);
    if (partitionCount == 1) {
        for (Page page : buildContext.getBuildPages()) {
            operators[0].addInput(page);
        }
    } else {
        PartitionFunction partitionGenerator = new LocalPartitionGenerator(new InterpretedHashGenerator(buildContext.getHashChannels().stream().map(channel -> buildContext.getTypes().get(channel)).collect(toImmutableList()), buildContext.getHashChannels(), TYPE_OPERATOR_FACTORY), partitionCount);
        for (Page page : buildContext.getBuildPages()) {
            Page[] partitionedPages = partitionPages(page, buildContext.getTypes(), partitionCount, partitionGenerator);
            for (int i = 0; i < partitionCount; i++) {
                operators[i].addInput(partitionedPages[i]);
            }
        }
    }
    LookupSourceFactory lookupSourceFactory = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide());
    ListenableFuture<LookupSourceProvider> lookupSourceProvider = lookupSourceFactory.createLookupSourceProvider();
    for (Operator operator : operators) {
        operator.finish();
    }
    if (!lookupSourceProvider.isDone()) {
        throw new AssertionError("Expected lookup source provider to be ready");
    }
    getFutureValue(lookupSourceProvider).close();
}

Example 2 with PartitionFunction

use of io.trino.operator.PartitionFunction in project trino by trinodb.

the class LocalExecutionPlanner method plan.

public LocalExecutionPlan plan(TaskContext taskContext, PlanNode plan, TypeProvider types, PartitioningScheme partitioningScheme, StageExecutionDescriptor stageExecutionDescriptor, List<PlanNodeId> partitionedSourceOrder, OutputBuffer outputBuffer) {
    List<Symbol> outputLayout = partitioningScheme.getOutputLayout();
    if (partitioningScheme.getPartitioning().getHandle().equals(FIXED_BROADCAST_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(SCALED_WRITER_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(SINGLE_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(COORDINATOR_DISTRIBUTION)) {
        return plan(taskContext, stageExecutionDescriptor, plan, outputLayout, types, partitionedSourceOrder, new TaskOutputFactory(outputBuffer));
    }
    // We can convert the symbols directly into channels, because the root must be a sink and therefore the layout is fixed
    List<Integer> partitionChannels;
    List<Optional<NullableValue>> partitionConstants;
    List<Type> partitionChannelTypes;
    if (partitioningScheme.getHashColumn().isPresent()) {
        partitionChannels = ImmutableList.of(outputLayout.indexOf(partitioningScheme.getHashColumn().get()));
        partitionConstants = ImmutableList.of(Optional.empty());
        partitionChannelTypes = ImmutableList.of(BIGINT);
    } else {
        partitionChannels = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument.isConstant()) {
                return -1;
            }
            return outputLayout.indexOf(argument.getColumn());
        }).collect(toImmutableList());
        partitionConstants = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument.isConstant()) {
                return Optional.of(argument.getConstant());
            }
            return Optional.<NullableValue>empty();
        }).collect(toImmutableList());
        partitionChannelTypes = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument.isConstant()) {
                return argument.getConstant().getType();
            }
            return types.get(argument.getColumn());
        }).collect(toImmutableList());
    }
    PartitionFunction partitionFunction = nodePartitioningManager.getPartitionFunction(taskContext.getSession(), partitioningScheme, partitionChannelTypes);
    OptionalInt nullChannel = OptionalInt.empty();
    Set<Symbol> partitioningColumns = partitioningScheme.getPartitioning().getColumns();
    // partitioningColumns expected to have one column in the normal case, and zero columns when partitioning on a constant
    checkArgument(!partitioningScheme.isReplicateNullsAndAny() || partitioningColumns.size() <= 1);
    if (partitioningScheme.isReplicateNullsAndAny() && partitioningColumns.size() == 1) {
        nullChannel = OptionalInt.of(outputLayout.indexOf(getOnlyElement(partitioningColumns)));
    }
    return plan(taskContext, stageExecutionDescriptor, plan, outputLayout, types, partitionedSourceOrder, new PartitionedOutputFactory(partitionFunction, partitionChannels, partitionConstants, partitioningScheme.isReplicateNullsAndAny(), nullChannel, outputBuffer, maxPagePartitioningBufferSize, positionsAppenderFactory));
}