@Override public Operator createOperator(DriverContext driverContext) { checkState(!closed, "Factory is already closed"); LookupSourceFactory lookupSourceFactory = joinBridgeManager.getJoinBridge(driverContext.getLifespan()); OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, LookupJoinOperator.class.getSimpleName()); lookupSourceFactory.setTaskContext(driverContext.getPipelineContext().getTaskContext()); joinBridgeManager.probeOperatorCreated(driverContext.getLifespan()); return new LookupJoinOperator( operatorContext, probeTypes, buildOutputTypes, joinType, lookupSourceFactory, joinProbeFactory, () -> joinBridgeManager.probeOperatorClosed(driverContext.getLifespan()), totalOperatorsCount, probeHashGenerator, partitioningSpillerFactory); }
NestedLoopJoinPages(List<Page> pages, DataSize estimatedSize, OperatorContext operatorContext) { requireNonNull(pages, "pages is null"); requireNonNull(operatorContext, "operatorContext is null"); this.pages = ImmutableList.copyOf(pages); this.taskContext = operatorContext.getDriverContext().getPipelineContext().getTaskContext(); this.estimatedSize = requireNonNull(estimatedSize, "estimatedSize is null"); operatorContext.transferMemoryToTaskContext(estimatedSize.toBytes()); }
public SharedLookupSource(LookupSource lookupSource, OperatorContext operatorContext) { requireNonNull(operatorContext, "operatorContext is null"); this.lookupSource = requireNonNull(lookupSource, "lookupSource is null"); this.taskContext = operatorContext.getDriverContext().getPipelineContext().getTaskContext(); operatorContext.transferMemoryToTaskContext(lookupSource.getInMemorySizeInBytes()); }
@Override public ListenableFuture<LookupSource> getLookupSource(OperatorContext operatorContext) { IndexLoader indexLoader = indexLoaderSupplier.get(); indexLoader.setContext(operatorContext.getDriverContext().getPipelineContext().getTaskContext()); return Futures.immediateFuture(new IndexLookupSource(indexLoader)); }
public void transferMemoryToTaskContext(long taskBytes) { long bytes = memoryReservation.getAndSet(0); driverContext.transferMemoryToTaskContext(bytes); TaskContext taskContext = driverContext.getPipelineContext().getTaskContext(); if (taskBytes > bytes) { try { taskContext.reserveMemory(taskBytes - bytes); } catch (ExceededMemoryLimitException e) { taskContext.freeMemory(bytes); throw e; } } else { taskContext.freeMemory(bytes - taskBytes); } }