Java Code Examples for org.apache.arrow.memory.BufferAllocator#newChildAllocator()

/**
 * Create a new exec service. Note that at start time, the provider to one of
 * the handlers may be a noop implementation. This is allowed if this is a
 * single role node.
 *  @param config
 * @param allocator
 * @param fabricService
 * @param executorService
 * @param execResults
 * @param execStatus
 */
public CoordExecService(
  SabotConfig config,
  BufferAllocator allocator,
  Provider<FabricService> fabricService,
  Provider<ExecutorService> executorService,
  Provider<ExecToCoordResultsHandler> execResults,
  Provider<ExecToCoordStatusHandler> execStatus,
  Provider<CoordinationProtos.NodeEndpoint> selfEndpoint,
  Provider<JobTelemetryClient> jobTelemetryClient) {

  super();

  this.fabricService = fabricService;
  this.allocator =  allocator.newChildAllocator("coord-exec-rpc",
          config.getLong("dremio.exec.rpc.bit.server.memory.data.reservation"),
          config.getLong("dremio.exec.rpc.bit.server.memory.data.maximum"));
  this.executorService = executorService;
  this.execResults = execResults;
  this.execStatus = execStatus;
  this.selfEndpoint = selfEndpoint;
  this.jobTelemetryClient = jobTelemetryClient;
  this.config = getMapping(config);
}
 

public SpoolingRawBatchBuffer(SharedResource resource, final SabotConfig config, FragmentWorkQueue workQueue,
                              FragmentHandle handle, SpillService spillService, BufferAllocator allocator,
                              int fragmentCount, int oppositeId, int bufferIndex) {
  super(resource, config, handle, allocator, fragmentCount);
  final String name = String.format("%s:spoolingBatchBuffer", QueryIdHelper.getFragmentId(handle));
  this.allocator = allocator.newChildAllocator(name, ALLOCATOR_INITIAL_RESERVATION, ALLOCATOR_MAX_RESERVATION);
  this.threshold = config.getLong(ExecConstants.SPOOLING_BUFFER_SIZE);
  this.oppositeId = oppositeId;
  this.bufferIndex = bufferIndex;
  this.bufferQueue = new SpoolingBufferQueue();
  this.workQueue = workQueue;
  this.spillService = spillService;
  this.inputStream = null;
  this.inputStreamLastKnownLen = 0;

  workQueue.put(new Runnable() {
    @Override
    public void run() {
      setupOutputStream();
    }
  });

}
 

public PDFSService(
    Provider<FabricService> fabricService,
    Provider<NodeEndpoint> identityProvider,
    Provider<Iterable<NodeEndpoint>> nodeProvider,
    Tracer tracer,
    SabotConfig config,
    BufferAllocator allocator,
    PDFSMode mode) {
  this.fabricService = fabricService;
  this.identityProvider = identityProvider;
  this.nodeProvider = nodeProvider;
  this.tracer = tracer;
  this.config = config;
  this.allocator = allocator.newChildAllocator("pdfs-allocator", 0, Long.MAX_VALUE);
  this.allowLocalAccess = mode == PDFSMode.DATA;
}
 

public WorkloadTicketDepot(BufferAllocator parentAllocator, SabotConfig config, GroupManager<AsyncTaskWrapper> manager) {
  this.manager = Preconditions.checkNotNull(manager, "Task manager required");

  nrtWorkloadTicket = new WorkloadTicket(parentAllocator.newChildAllocator("nrt-workload-allocator", 0,
    getLongConfig(config, INSTANT_MAX_ALLOCATION_CONFIG, Long.MAX_VALUE)), manager.newGroup(NRT_WEIGHT));
  nrtWorkloadTicket.reserve();

  generalWorkloadTicket = new WorkloadTicket(parentAllocator.newChildAllocator("general-workload-allocator", 0,
    getLongConfig(config, GENERAL_MAX_ALLOCATION_CONFIG, Long.MAX_VALUE)), manager.newGroup(GENERAL_WEIGHT));
  generalWorkloadTicket.reserve();

  backgroundWorkloadTicket = new WorkloadTicket(parentAllocator.newChildAllocator("background-workload-allocator", 0,
    getLongConfig(config, BACKGROUND_MAX_ALLOCATION_CONFIG, Long.MAX_VALUE)), manager.newGroup(BACKGROUND_WEIGHT));
  backgroundWorkloadTicket.reserve();
}
 

@Test
public void testPrune() throws Exception {
  try (RollbackCloseable closeable = new RollbackCloseable(true)) {
    BufferAllocator root = allocatorRule.newAllocator("test-memory-debug-info", 0, 1024 * 1024);
    closeable.add(root);

    BufferAllocator twig = root.newChildAllocator("twig",0, 1024 * 1024);
    closeable.add(twig);

    for (int i = 0; i < 20; ++i) {
      boolean isBig = (i % 2 == 0);
      BufferAllocator child =
          twig.newChildAllocator((isBig ? "big" : "small") + i, 0, Integer.MAX_VALUE);
      closeable.add(child);

      if (isBig) {
        closeable.add(child.buffer(8192));
      } else {
        closeable.add(child.buffer(4096));
      }
    }

    // allocate to hit the root limit.
    try (ArrowBuf buf = twig.buffer(1024 * 1024)) {
      assertTrue("expected allocation above limit to fail", false); // should not reach here
    } catch (OutOfMemoryException e) {
      String info = MemoryDebugInfo.getDetailsOnAllocationFailure(e, twig);

      assertTrue(!info.contains("test-memory-debug-info"));
      assertTrue(info.contains("twig"));
      assertTrue(info.contains("big"));
      assertTrue(!info.contains("small"));
    }
  }
}
 

public Producer(BufferAllocator allocator, FlightDescriptor descriptor, BlockingQueue<ArrowRecordBatch> exchanger) {
  this.allocator = allocator.newChildAllocator("producer", 0, Long.MAX_VALUE);
  this.descriptor = descriptor;
  this.exchanger = exchanger;
}
 

public Consumer(BufferAllocator allocator, FlightDescriptor descriptor, BlockingQueue<ArrowRecordBatch> exchanger) {
  this.allocator = allocator.newChildAllocator("consumer", 0, Long.MAX_VALUE);
  this.descriptor = descriptor;
  this.exchanger = exchanger;
}
 

/**
 * Create a query allocator for usage with the given query
 */
protected BufferAllocator makeQueryAllocator(BufferAllocator parent, QueryId queryId, long maxAllocation) {
  return parent.newChildAllocator("query-" + QueryIdHelper.getQueryId(queryId), 0, maxAllocation);
}
 

public IncomingBuffers(
    DeferredException exception,
    SharedResourceGroup resourceGroup,
    FragmentWorkQueue workQueue,
    TunnelProvider tunnelProvider,
    PlanFragmentFull fragment,
    BufferAllocator incomingAllocator,
    SabotConfig config,
    ExecutionControls executionControls,
    SpillService spillService,
    PlanFragmentsIndex planFragmentsIndex
    ) {
  this.deferredException = exception;
  this.resourceGroup = resourceGroup;

  final FragmentHandle handle = fragment.getMajor().getHandle();
  final String allocatorName = String.format("op:%s:incoming",
    QueryIdHelper.getFragmentId(fragment.getHandle()));
  this.allocator = incomingAllocator.newChildAllocator(allocatorName, 0, Long.MAX_VALUE);

  final Map<Integer, DataCollector> collectors = Maps.newHashMap();
  EndpointsIndex endpointsIndex = planFragmentsIndex.getEndpointsIndex();
  try (AutoCloseables.RollbackCloseable rollbackCloseable = new AutoCloseables.RollbackCloseable(true, allocator)) {
    for (int i = 0; i < fragment.getMinor().getCollectorCount(); i++) {
      Collector collector = fragment.getMinor().getCollector(i);

      DataCollector newCollector = collector.getSupportsOutOfOrder() ?
        new MergingCollector(resourceGroup, collector, allocator, config, fragment.getHandle(), workQueue, tunnelProvider, spillService, endpointsIndex) :
        new PartitionedCollector(resourceGroup, collector, allocator, config, fragment.getHandle(), workQueue, tunnelProvider, spillService, endpointsIndex);
      rollbackCloseable.add(newCollector);
      collectors.put(collector.getOppositeMajorFragmentId(), newCollector);
    }

    injector.injectChecked(executionControls, INJECTOR_DO_WORK, OutOfMemoryException.class);
    rollbackCloseable.commit();
  } catch (Exception e) {
    throw Throwables.propagate(e);
  }

  collectorMap = ImmutableMap.copyOf(collectors);
}
 

public ReflectionServiceImpl(
  SabotConfig config,
  Provider<LegacyKVStoreProvider> storeProvider,
  Provider<SchedulerService> schedulerService,
  Provider<JobsService> jobsService,
  Provider<CatalogService> catalogService,
  final Provider<SabotContext> sabotContext,
  Provider<ReflectionStatusService> reflectionStatusService,
  ExecutorService executorService,
  boolean isMaster,
  BufferAllocator allocator) {
  this.schedulerService = Preconditions.checkNotNull(schedulerService, "scheduler service required");
  this.jobsService = Preconditions.checkNotNull(jobsService, "jobs service required");
  this.catalogService = Preconditions.checkNotNull(catalogService, "catalog service required");
  this.sabotContext = Preconditions.checkNotNull(sabotContext, "acceleration plugin required");
  this.reflectionStatusService = Preconditions.checkNotNull(reflectionStatusService, "reflection status service required");
  this.executorService = Preconditions.checkNotNull(executorService, "executor service required");
  this.namespaceService = new Provider<NamespaceService>() {
    @Override
    public NamespaceService get() {
      return sabotContext.get().getNamespaceService(SYSTEM_USERNAME);
    }
  };
  this.reflectionSettings = new ReflectionSettings(namespaceService, storeProvider);
  this.isMaster = isMaster;
  this.allocator = allocator.newChildAllocator(getClass().getName(), 0, Long.MAX_VALUE);

  userStore = new ReflectionGoalsStore(storeProvider);
  internalStore = new ReflectionEntriesStore(storeProvider);
  materializationStore = new MaterializationStore(storeProvider);
  externalReflectionStore = new ExternalReflectionStore(storeProvider);
  dependenciesStore = new DependenciesStore(storeProvider);
  requestsStore = new RefreshRequestsStore(storeProvider);

  this.queryContext = new Supplier<QueryContext>() {
    @Override
    public QueryContext get() {
      final UserSession session = systemSession(getOptionManager());
      return new QueryContext(session, sabotContext.get(), new AttemptId().toQueryId(),
          java.util.Optional.of(false));
    }
  };

  this.expansionHelper = new Supplier<ExpansionHelper>() {
    @Override
    public ExpansionHelper get() {
      return new ExpansionHelper(queryContext.get());
    }
  };

  this.validator = new ReflectionValidator(catalogService);
  this.materializationDescriptorFactory = config.getInstance(
      "dremio.reflection.materialization.descriptor.factory",
      MaterializationDescriptorFactory.class,
      DEFAULT_MATERIALIZATION_DESCRIPTOR_FACTORY);
}
 

public BufferAllocatorFactory(BufferAllocator bootStrapAllocator, String name) {
  allocator = bootStrapAllocator.newChildAllocator(name, 0, Long.MAX_VALUE);
}
 

/**
 * Opens an iterator to read read batches from spill run. It is responsibility of the caller to release iterator
 * when done.
 *
 * @param batchId
 * @param container
 * @return
 * @throws IOException
 */
private DiskRunIterator openRun(BufferAllocator copierAllocator, int batchId, ExpandableHyperContainer container) throws IOException{
  Preconditions.checkState(iterator == null);
  final long memCapacity = nextPowerOfTwo(largestBatch);
  final BufferAllocator allocator = copierAllocator.newChildAllocator("diskrun", 0, memCapacity);
  iterator = new DiskRunIterator(batchCount, spillFile, container, allocator);

  return iterator;
}