use of io.trino.operator.exchange.LocalPartitionGenerator 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();
}
Aggregations