Class MemoryManagerImpl

java/core/src/java/org/apache/orc/impl/MemoryManagerImpl.java:42–150

Implements a memory manager that keeps a global context of how many ORC writers there are and manages the memory between them. For use cases with dynamic partitions, it is easy to end up with many writers in the same task. By managing the size of each allocation, we try to cut down the size of each

 * invocations are triggered from the same thread.
 */
public class MemoryManagerImpl implements MemoryManager {

  private final long totalMemoryPool;
  private final Map<Path, WriterInfo> writerList = new HashMap<>();
  private final AtomicLong totalAllocation = new AtomicLong(0);

  private static class WriterInfo {
    long allocation;
    WriterInfo(long allocation) {
      this.allocation = allocation;
    }
  }

  /**
   * Create the memory manager.
   * @param conf use the configuration to find the maximum size of the memory
   *             pool.
   */
  public MemoryManagerImpl(Configuration conf) {
    this(Math.round(ManagementFactory.getMemoryMXBean().
        getHeapMemoryUsage().getMax() * OrcConf.MEMORY_POOL.getDouble(conf)));
  }

  /**
   * Create the memory manager
   * @param poolSize the size of memory to use
   */
  public MemoryManagerImpl(long poolSize) {
    totalMemoryPool = poolSize;
  }

  /**
   * Add a new writer's memory allocation to the pool. We use the path
   * as a unique key to ensure that we don't get duplicates.
   * @param path the file that is being written
   * @param requestedAllocation the requested buffer size
   */
  @Override
  public synchronized void addWriter(Path path, long requestedAllocation,
                              Callback callback) throws IOException {
    WriterInfo oldVal = writerList.get(path);
    // this should always be null, but we handle the case where the memory
    // manager wasn't told that a writer wasn't still in use and the task
    // starts writing to the same path.
    if (oldVal == null) {
      oldVal = new WriterInfo(requestedAllocation);
      writerList.put(path, oldVal);
      totalAllocation.addAndGet(requestedAllocation);
    } else {
      // handle a new writer that is writing to the same path
      totalAllocation.addAndGet(requestedAllocation - oldVal.allocation);
      oldVal.allocation = requestedAllocation;
    }
  }

  /**
   * Remove the given writer from the pool.
   * @param path the file that has been closed