How to use

PartitionedLookupSourceFactory

in

io.prestosql.operator

private void freePartitions()
{
  // Let the HashBuilderOperators reduce their accounted memory
  partitionsNoLongerNeeded.set(null);
  lock.writeLock().lock();
  try {
    // Remove out references to partitions to actually free memory
    Arrays.fill(partitions, null);
    lookupSourceSupplier = null;
    closeCachedLookupSources();
  }
  finally {
    lock.writeLock().unlock();
  }
}

@Override
public ListenableFuture<?> whenBuildFinishes()
{
  return transform(
      this.createLookupSourceProvider(),
      lookupSourceProvider -> {
        // Close the lookupSourceProvider we just created.
        // The only reason we created it is to wait until lookup source is ready.
        lookupSourceProvider.close();
        return null;
      },
      directExecutor());
}

@Override
public void destroy()
{
  lock.writeLock().lock();
  try {
    freePartitions();
    spilledPartitions.values().forEach(SpilledLookupSourceHandle::dispose);
    // Setting destroyed must be last because it's a part of the state exposed by isDestroyed() without synchronization.
    destroyed.set(null);
  }
  finally {
    lock.writeLock().unlock();
  }
}

this.localRevocableMemoryContext = operatorContext.localRevocableMemoryContext();
this.index = pagesIndexFactory.newPagesIndex(lookupSourceFactory.getTypes(), expectedPositions);
this.lookupSourceFactory = lookupSourceFactory;
lookupSourceFactoryDestroyed = lookupSourceFactory.isDestroyed();

  verify(partitions.length > 1, "Spill occurred when only one partition");
  lookupSourceSupplier = createPartitionedLookupSourceSupplier(ImmutableList.copyOf(partitions), hashChannelTypes, outer);
  closeCachedLookupSources();
supplyLookupSources();

private ListenableFuture<Supplier<LookupSource>> loadSpilledLookupSource(int partitionNumber)
{
  return getSpilledLookupSourceHandle(partitionNumber).getLookupSource();
}

@VisibleForTesting
public static JoinBridgeManager<PartitionedLookupSourceFactory> lookupAllAtOnce(PartitionedLookupSourceFactory factory)
{
  return new JoinBridgeManager<>(
      false,
      UNGROUPED_EXECUTION,
      UNGROUPED_EXECUTION,
      ignored -> factory,
      factory.getOutputTypes());
}

private void finishInput()
{
  checkState(state == State.CONSUMING_INPUT);
  if (lookupSourceFactoryDestroyed.isDone()) {
    close();
    return;
  }
  LookupSourceSupplier partition = buildLookupSource();
  if (spillEnabled) {
    localRevocableMemoryContext.setBytes(partition.get().getInMemorySizeInBytes());
  }
  else {
    localUserMemoryContext.setBytes(partition.get().getInMemorySizeInBytes());
  }
  lookupSourceNotNeeded = Optional.of(lookupSourceFactory.lendPartitionLookupSource(partitionIndex, partition));
  state = State.LOOKUP_SOURCE_BUILT;
}

@Override
public HashBuilderOperator createOperator(DriverContext driverContext)
{
  checkState(!closed, "Factory is already closed");
  OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, HashBuilderOperator.class.getSimpleName());
  PartitionedLookupSourceFactory lookupSourceFactory = this.lookupSourceFactoryManager.getJoinBridge(driverContext.getLifespan());
  int partitionIndex = getAndIncrementPartitionIndex(driverContext.getLifespan());
  verify(partitionIndex < lookupSourceFactory.partitions());
  return new HashBuilderOperator(
      operatorContext,
      lookupSourceFactory,
      partitionIndex,
      outputChannels,
      hashChannels,
      preComputedHashChannel,
      filterFunctionFactory,
      sortChannel,
      searchFunctionFactories,
      expectedPositions,
      pagesIndexFactory,
      spillEnabled,
      singleStreamSpillerFactory);
}

JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = JoinBridgeManager.lookupAllAtOnce(new PartitionedLookupSourceFactory(
    buildPages.getTypes(),
    rangeList(buildPages.getTypes().size()).stream()

this.localRevocableMemoryContext = operatorContext.localRevocableMemoryContext();
this.index = pagesIndexFactory.newPagesIndex(lookupSourceFactory.getTypes(), expectedPositions);
this.lookupSourceFactory = lookupSourceFactory;
lookupSourceFactoryDestroyed = lookupSourceFactory.isDestroyed();

  verify(partitions.length > 1, "Spill occurred when only one partition");
  lookupSourceSupplier = createPartitionedLookupSourceSupplier(ImmutableList.copyOf(partitions), hashChannelTypes, outer);
  closeCachedLookupSources();
supplyLookupSources();

private ListenableFuture<Supplier<LookupSource>> loadSpilledLookupSource(int partitionNumber)
{
  return getSpilledLookupSourceHandle(partitionNumber).getLookupSource();
}

@VisibleForTesting
public static JoinBridgeManager<PartitionedLookupSourceFactory> lookupAllAtOnce(PartitionedLookupSourceFactory factory)
{
  return new JoinBridgeManager<>(
      false,
      UNGROUPED_EXECUTION,
      UNGROUPED_EXECUTION,
      ignored -> factory,
      factory.getOutputTypes());
}

private void finishInput()
{
  checkState(state == State.CONSUMING_INPUT);
  if (lookupSourceFactoryDestroyed.isDone()) {
    close();
    return;
  }
  LookupSourceSupplier partition = buildLookupSource();
  if (spillEnabled) {
    localRevocableMemoryContext.setBytes(partition.get().getInMemorySizeInBytes());
  }
  else {
    localUserMemoryContext.setBytes(partition.get().getInMemorySizeInBytes());
  }
  lookupSourceNotNeeded = Optional.of(lookupSourceFactory.lendPartitionLookupSource(partitionIndex, partition));
  state = State.LOOKUP_SOURCE_BUILT;
}

@Override
public HashBuilderOperator createOperator(DriverContext driverContext)
{
  checkState(!closed, "Factory is already closed");
  OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, HashBuilderOperator.class.getSimpleName());
  PartitionedLookupSourceFactory lookupSourceFactory = this.lookupSourceFactoryManager.getJoinBridge(driverContext.getLifespan());
  int partitionIndex = getAndIncrementPartitionIndex(driverContext.getLifespan());
  verify(partitionIndex < lookupSourceFactory.partitions());
  return new HashBuilderOperator(
      operatorContext,
      lookupSourceFactory,
      partitionIndex,
      outputChannels,
      hashChannels,
      preComputedHashChannel,
      filterFunctionFactory,
      sortChannel,
      searchFunctionFactories,
      expectedPositions,
      pagesIndexFactory,
      spillEnabled,
      singleStreamSpillerFactory);
}

JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = JoinBridgeManager.lookupAllAtOnce(new PartitionedLookupSourceFactory(
    buildPages.getTypes(),
    rangeList(buildPages.getTypes().size()).stream()

@Override
public void destroy()
{
  lock.writeLock().lock();
  try {
    freePartitions();
    spilledPartitions.values().forEach(SpilledLookupSourceHandle::dispose);
    // Setting destroyed must be last because it's a part of the state exposed by isDestroyed() without synchronization.
    destroyed.set(null);
  }
  finally {
    lock.writeLock().unlock();
  }
}

private void disposeSpilledLookupSource(int partitionNumber)
{
  getSpilledLookupSourceHandle(partitionNumber).dispose();
}

@Override
public ListenableFuture<?> whenBuildFinishes()
{
  return transform(
      this.createLookupSourceProvider(),
      lookupSourceProvider -> {
        // Close the lookupSourceProvider we just created.
        // The only reason we created it is to wait until lookup source is ready.
        lookupSourceProvider.close();
        return null;
      },
      directExecutor());
}