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

@Override
public SelectionVector4 getFinalSort(BufferAllocator allocator, int targetBatchSize){
  Stopwatch watch = Stopwatch.createStarted();

  intVector.setValueCount(totalCount);
  QuickSort qs = new QuickSort();
  if (totalCount > 0) {
    qs.sort(this, 0, totalCount);
  }

  SelectionVector4 finalSortedSV4 = new SelectionVector4(allocator.buffer(totalCount * 4), totalCount, targetBatchSize);
  for (int i = 0; i < totalCount; i++) {
    finalSortedSV4.set(i, intVector.get(i));
  }

  logger.debug("Took {} us to final sort {} records in {} batches",
    watch.elapsed(TimeUnit.MICROSECONDS), totalCount, hyperBatch.size());

  return finalSortedSV4;
}
 

@Override
public void setup(
    FunctionContext context,
    BufferAllocator allocator,
    DiskRunIterator[] iterators,
    VectorAccessible incoming,
    VectorContainer outgoing) throws SchemaChangeException, IOException {
  this.incoming = new Sv4HyperContainer(allocator, incoming.getSchema());
  this.size = iterators.length;
  final ArrowBuf arrowBuf = allocator.buffer(4 * size);
  this.vector4 = new SelectionVector4(arrowBuf, size, Character.MAX_VALUE);
  this.iterators = iterators;
  this.outgoing = outgoing;

  doSetup(context, incoming, outgoing);

  queueSize = 0;
  for (int i = 0; i < size; i++) {
    vector4.set(i, i, iterators[i].getNextId());
    siftUp();
    queueSize++;
  }
}
 

@Override
  public SelectionVector4 getFinalSort(BufferAllocator allocator, int targetBatchSize){
    Stopwatch watch = Stopwatch.createStarted();

    final int totalCount = tree.getTotalCount();
    final SelectionVector4 sortVector = new SelectionVector4(allocator.buffer(totalCount * 4), totalCount, targetBatchSize);
    final SplayIterator iterator = tree.iterator();
    int i = 0;
    while(iterator.hasNext()){
      final int index = iterator.next();
//      logger.debug("Setting {}={}", i, index);
      sortVector.set(i, index);
      i++;
    }
    logger.debug("Took {} us to generate final order", watch.elapsed(TimeUnit.MICROSECONDS));

    // release the memory associated with the tree.
    return sortVector;
  }
 

private static void reallocOffsetBuffer(final BaseRepeatedValueVector vector,
                                        final BufferAllocator allocator) {
  long currentBufferCapacity = vector.offsetBuffer.capacity();
  long baseSize = vector.offsetAllocationSizeInBytes;
  if(baseSize < currentBufferCapacity) {
    baseSize = currentBufferCapacity;
  }
  long newAllocationSize = baseSize * 2L;
  newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
  int size = LargeMemoryUtil.checkedCastToInt(newAllocationSize);
  ArrowBuf newBuf = allocator.buffer(size);
  newBuf.setZero(0, newBuf.capacity());
  newBuf.setBytes(0, vector.offsetBuffer, 0, currentBufferCapacity);
  vector.offsetBuffer.release(1);
  vector.offsetBuffer = newBuf;
  vector.offsetAllocationSizeInBytes = size;
}
 

/**
 * Creates MatchBitSet, allocates the buffer and initialized values to 0
 * @param numBits       the number of bits
 * @param allocator     the allocator used to allocate direct memory
 */
public MatchBitSet(final int numBits, final BufferAllocator allocator) {
  Preconditions.checkArgument(numBits >= 0, "bits count in constructor of BitSet is invalid");

  this.allocator = allocator;
  this.numBits = numBits;
  numWords = bitsToWords(this.numBits);
  buffer = allocator.buffer(numWords * BYTES_PER_WORD);
  bufferAddr = buffer.memoryAddress();
  long maxBufferAddr = bufferAddr + numWords * BYTES_PER_WORD;
  for (long wordAddr = bufferAddr; wordAddr < maxBufferAddr; wordAddr += BYTES_PER_WORD) {
    PlatformDependent.putLong(wordAddr, 0);
  }
}
 

private ArrowBuf allocateHelper(BufferAllocator alloc, final int requestSize) throws Exception{
  try {
    return alloc.buffer(requestSize);
  } catch (OutOfMemoryException e) {
    throw UserException
      .memoryError(e)
      .addContext(MemoryDebugInfo.getDetailsOnAllocationFailure(e, alloc))
      .build(logger);
  }
}
 

/**
 * Ensures there is headroom in the allocator. Throws an exception if not sufficient.
 * @param allocator
 * @param headRoom
 */
public static void ensureHeadroom(BufferAllocator allocator, long headRoom) throws OutOfMemoryException {
  if (allocator.getHeadroom() >= headRoom) {
    // nothing to do in this case.
    return;
  }

  // do a dummy allocation so that we get the allocation states as part of the exception.
  try (ArrowBuf buffer = allocator.buffer(headRoom)) {}
}
 

public FixedBlockVector(BufferAllocator allocator, int blockWidth, int initialCapacity, boolean allowExpansion) {
  this.allocator = allocator;
  this.blockWidth = blockWidth;
  this.allowExpansion = allowExpansion;
  this.buf = allocator.buffer(0);
  this.capacity = 0;
  resizeBuffer(initialCapacity);
  resetPositions();
}
 

@Test
public void ensureZeroBufferIsValid() throws Exception {
  try(RollbackCloseable closeables = new RollbackCloseable(true)) {
    BufferAllocator alloc = closeables.add(this.rootAllocator.newChildAllocator("child", 0, Long.MAX_VALUE));
    ArrowBuf buffer = alloc.buffer(0);
    assertTrue(buffer.memoryAddress() != 0);
    closeables.add(buffer);
  }
}
 

@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"));
    }
  }
}
 

@Override
public void init(Sv4HyperContainer hyperBatch, int limit, FunctionContext context, BufferAllocator allocator,  boolean hasSv2, int maxSize) {
  this.limit = limit;
  this.context = context;
  this.allocator = allocator;
  final ArrowBuf ArrowBuf = allocator.buffer(4 * (limit + 1));
  this.heapSv4 = new SelectionVector4(ArrowBuf, limit, Character.MAX_VALUE);
  this.hasSv2 = hasSv2;
  this.hyperBatch = hyperBatch;
  this.maxSize = maxSize;
  doSetup(context, hyperBatch, null);
}
 

public SplaySorter(ExternalSort sortConfig, ClassProducer classProducer, Schema schema, BufferAllocator allocator) {
  this.sortConfig = sortConfig;
  this.classProducer = classProducer;
  this.schema = schema;
  this.allocator = allocator;
  this.splayTreeBuffer = allocator.buffer(4096 * SplayTree.NODE_SIZE);
  splayTreeBuffer.setZero(0, splayTreeBuffer.capacity());
}
 

public PartitionToLoadSpilledData(final BufferAllocator allocator,
                                  final int fixedDataLength,
                                  final int variableDataLength,
                                  final List<Field> postSpillAccumulatorVectorTypes,
                                  final int batchSize) throws Exception {
  Preconditions.checkArgument(allocator != null, "Error: need a valid allocator to pre-allocate memory");
  this.allocator = allocator;
  /* we use Numbers.nextPowerOfTwo because that is how memory allocation happens
   * inside FixedBlockVector and VariableBlockVector when inserting into hashtable.
   * if we don't use nextPowerOfTwo for actual allocation size, we might run into
   * situation where a spilled batch has more data than the capacity we have pre-allocated
   * here to load the spilled batch into memory.
   */
  try(AutoCloseables.RollbackCloseable rollbackable = new AutoCloseables.RollbackCloseable()) {
    fixedKeyColPivotedData = allocator.buffer(Numbers.nextPowerOfTwo(fixedDataLength));
    rollbackable.add(fixedKeyColPivotedData);
    variableKeyColPivotedData = allocator.buffer(Numbers.nextPowerOfTwo(variableDataLength));
    rollbackable.add(variableKeyColPivotedData);
    this.postSpillAccumulatorVectors = new FieldVector[postSpillAccumulatorVectorTypes.size()];
    this.fixedDataLength = fixedDataLength;
    this.variableDataLength = variableDataLength;
    this.batchSize = batchSize;
    this.recordsInBatch = 0;
    this.accumulatorTypes = new byte[postSpillAccumulatorVectorTypes.size()];
    initBuffers();
    initPostSpillAccumulatorVectors(postSpillAccumulatorVectorTypes, batchSize, rollbackable);
    rollbackable.commit();
    logger.debug("Extra Partition Pre-allocation, fixed-data length: {}, variable-data length: {}, actual fixed-data capacity: {}, actual variable-data capacty: {}, batchSize: {}",
                 fixedDataLength, variableDataLength, fixedKeyColPivotedData.capacity(), variableKeyColPivotedData.capacity(), batchSize);
  }
}
 

public TwosComplementValuePair(BufferAllocator allocator, Field field, byte[] value) {
  super(field.getName(), value != null ? DecimalTools.signExtend16(value) : null);
  CompleteType type = CompleteType.fromField(field);
  scale = type.getScale();
  precision = type.getPrecision();
  if (value != null) {
    buf = allocator.buffer(16);
    /* set the bytes in LE format in the buffer of decimal vector. since we
     * are populating the decimal vector multiple times with the same buffer, it
     * is fine to swap the bytes once here as opposed to swapping while writing.
     */
    byte [] decimalBytesInLE = DecimalUtils.convertDecimalBytesToArrowByteArray(value);
    buf.setBytes(0, decimalBytesInLE);
  }
}
 

public static VarCharHolder getVarCharHolder(BufferAllocator a, String s){
  VarCharHolder vch = new VarCharHolder();

  byte[] b = s.getBytes(Charsets.UTF_8);
  vch.start = 0;
  vch.end = b.length;
  vch.buffer = a.buffer(b.length); //
  vch.buffer.setBytes(0, b);
  return vch;
}
 

public void setup(
    BufferAllocator allocator,
    final ExpandableHyperContainer buildBatch,
    final VectorAccessible probeBatch,
    // Contains all vectors in probe side output
    final List<FieldVector> probeOutputs,
    /* Contains only carry over vectors in build side output for VECTORIZED_GENERIC
     * Contains all field vectors in build side output for VECTORIZED_BIGINT
     */
    final List<FieldVector> buildOutputs,
    /* Contains the key field vectors in incoming probe side batch for VECTORIZED_GENERIC
     * Only for VECTORIZED_GENERIC
     */
    final List<FieldVector> probeIncomingKeys,
    /* Contains the key field vectors in build side output for VECTORIZED_GENERIC
     * Only for VECTORIZED_GENERIC
     */
    final List<FieldVector> buildOutputKeys,
    VectorizedHashJoinOperator.Mode mode,
    JoinRelType joinRelType,
    List<BuildInfo> buildInfos,
    List<ArrowBuf> startIndices,
    List<MatchBitSet> keyMatchBitVectors,
    int maxHashTableIndex,
    JoinTable table,
    // Used to pivot the keys in incoming build batch into hash table
    PivotDef pivot,
    // Used to unpivot the keys in hash table to build side output
    PivotDef buildUnpivot,
    int targetRecordsPerBatch,
    final NullComparator nullMask){

  this.nullMask = nullMask;
  this.pivot = pivot;
  this.buildUnpivot = buildUnpivot;
  this.allocator = allocator;
  this.table = table;
  this.links = new ArrowBuf[buildInfos.size()];

  for (int i =0; i < links.length; i++) {
    links[i] = buildInfos.get(i).getLinks();
  }

  this.starts = new ArrowBuf[startIndices.size()];
  this.keyMatches = new MatchBitSet[keyMatchBitVectors.size()];

  if (startIndices.size() > 0) {
    this.maxOffsetForLastBatch = maxHashTableIndex - (startIndices.size() - 1) * HashTable.BATCH_SIZE;
  } else {
    this.maxOffsetForLastBatch = -1;
  }
  this.maxHashTableIndex = maxHashTableIndex;
  for (int i = 0; i < starts.length; i++) {
    starts[i] = startIndices.get(i);
    keyMatches[i] = keyMatchBitVectors.get(i);
  }

  this.projectUnmatchedBuild = joinRelType == JoinRelType.RIGHT || joinRelType == JoinRelType.FULL;
  this.projectUnmatchedProbe = joinRelType == JoinRelType.LEFT || joinRelType == JoinRelType.FULL;
  this.targetRecordsPerBatch = targetRecordsPerBatch;
  this.projectProbeSv2 = allocator.buffer(targetRecordsPerBatch * BATCH_OFFSET_SIZE);
  this.probeSv2Addr = projectProbeSv2.memoryAddress();
  // first 4 bytes (int) are for batch index and rest 2 bytes are offset within the batch
  this.projectBuildOffsetBuf = allocator.buffer(targetRecordsPerBatch * BUILD_RECORD_LINK_SIZE);
  this.projectBuildOffsetAddr = projectBuildOffsetBuf.memoryAddress();
  this.projectBuildKeyOffsetBuf = allocator.buffer(targetRecordsPerBatch * ORDINAL_SIZE);
  this.projectBuildKeyOffsetAddr = projectBuildKeyOffsetBuf.memoryAddress();

  this.mode = mode;

  this.buildOutputs = buildOutputs;
  if (table.size() > 0) {
    this.buildCopiers = projectUnmatchedProbe  ?
      ConditionalFieldBufferCopier6.getFourByteCopiers(VectorContainer.getHyperFieldVectors(buildBatch), buildOutputs) :
      FieldBufferCopier6.getFourByteCopiers(VectorContainer.getHyperFieldVectors(buildBatch), buildOutputs);
  } else {
    this.buildCopiers = Collections.emptyList();
  }

  /* For VECTORIZED_GENERIC, we don't keep the key vectors in hyper container,
   * and then we need to copy keys from probe batch to build side in output for matched and non matched records,
   * otherwise eight byte hash table is used, we keep the key vector in hyper container,
   * and then we don't need to copy keys from probe batch to build side in output for matched and non matched records.
   */
  if (this.mode == VectorizedHashJoinOperator.Mode.VECTORIZED_GENERIC) {
    // create copier for copying keys from probe batch to build side output
    if (probeIncomingKeys.size() > 0) {
      this.keysCopiers = FieldBufferCopier.getCopiers(probeIncomingKeys, buildOutputKeys);
    } else {
      this.keysCopiers = Collections.emptyList();
    }

    this.projectNullKeyOffset = allocator.buffer(targetRecordsPerBatch * BATCH_OFFSET_SIZE);
    this.projectNullKeyOffsetAddr = projectNullKeyOffset.memoryAddress();
  } else {
    this.projectNullKeyOffsetAddr = 0;
    this.keysCopiers = null;
  }

  this.probeCopiers = FieldBufferCopier.getCopiers(VectorContainer.getFieldVectors(probeBatch), probeOutputs);
}
 

public MatchedVector(BufferAllocator allocator) {
  this.buffer = allocator.buffer(4096);
}
 

private void insertAndAccumulateForAllPartitions(final BufferAllocator allocator,
                                                 final int records,
                                                 final PivotDef pivot,
                                                 final LBlockHashTable hashTable,
                                                 final AccumulatorSet accumulator,
                                                 final int[] expectedOrdinals) {
  try (final FixedBlockVector fbv = new FixedBlockVector(allocator, pivot.getBlockWidth());
       final VariableBlockVector var = new VariableBlockVector(allocator, pivot.getVariableCount());
       final ArrowBuf offsets = allocator.buffer(records * VectorizedHashAggOperator.PARTITIONINDEX_HTORDINAL_WIDTH);
       final SimpleBigIntVector hashValues = new SimpleBigIntVector("hashvalues", allocator)) {

    /* pivot the data into temporary space */
    Pivots.pivot(pivot, records, fbv, var);

    final long keyFixedVectorAddr = fbv.getMemoryAddress();
    final long keyVarVectorAddr = var.getMemoryAddress();
    int[] actualOrdinals = new int[expectedOrdinals.length];
    final boolean fixedOnly = pivot.getVariableCount() == 0;

    /* compute hash on the pivoted data */
    hashValues.allocateNew(records);
    final BlockChunk blockChunk = new BlockChunk(keyFixedVectorAddr, keyVarVectorAddr, fixedOnly,
      pivot.getBlockWidth(), records, hashValues.getBufferAddress(), 0);
    HashComputation.computeHash(blockChunk);

    /* insert */
    long offsetAddr = offsets.memoryAddress();
    for (int keyIndex = 0; keyIndex < records; keyIndex++, offsetAddr += VectorizedHashAggOperator.PARTITIONINDEX_HTORDINAL_WIDTH) {
      final int keyHash = (int)hashValues.get(keyIndex);
      actualOrdinals[keyIndex] = hashTable.add(keyFixedVectorAddr, keyVarVectorAddr, keyIndex, keyHash);
      PlatformDependent.putByte(offsetAddr, (byte)0);
      PlatformDependent.putInt(offsetAddr + VectorizedHashAggOperator.HTORDINAL_OFFSET, actualOrdinals[keyIndex]);
      PlatformDependent.putInt(offsetAddr + VectorizedHashAggOperator.KEYINDEX_OFFSET, keyIndex);
    }

    assertArrayEquals(expectedOrdinals, actualOrdinals);

    /* accumulate */
    accumulator.accumulate(offsets.memoryAddress(), records, hashTable.getBitsInChunk(), hashTable.getChunkOffsetMask());
  }
}
 

private void insertAndAccumulateForAllPartitions(final BufferAllocator allocator,
                                                 final int records,
                                                 final PivotDef pivot,
                                                 final VectorizedHashAggPartition[] partitions,
                                                 final int[] expectedOrdinals) {
  try (final FixedBlockVector fbv = new FixedBlockVector(allocator, pivot.getBlockWidth());
       final VariableBlockVector var = new VariableBlockVector(allocator, pivot.getVariableCount());
       final ArrowBuf offsets = allocator.buffer(records * VectorizedHashAggOperator.PARTITIONINDEX_HTORDINAL_WIDTH);
       final SimpleBigIntVector hashValues = new SimpleBigIntVector("hashvalues", allocator)) {

    /* pivot the data into temporary space */
    Pivots.pivot(pivot, records, fbv, var);

    final long keyFixedVectorAddr = fbv.getMemoryAddress();
    final long keyVarVectorAddr = var.getMemoryAddress();
    int[] actualOrdinals = new int[expectedOrdinals.length];
    final boolean fixedOnly = pivot.getVariableCount() == 0;

    /* compute hash on the pivoted data */
    hashValues.allocateNew(records);
    final BlockChunk blockChunk = new BlockChunk(keyFixedVectorAddr, keyVarVectorAddr, fixedOnly,
      pivot.getBlockWidth(), records, hashValues.getBufferAddress(), 0);
    HashComputation.computeHash(blockChunk);

    /* since we are mocking partitions to test spill handler,
     * just insert same data into all partitions
     */
    int hashPartitionIndex = 0;
    for (int i = 0; i < partitions.length; i++) {

      /* insert */
      long offsetAddr = offsets.memoryAddress();
      LBlockHashTable hashTable = null;
      for (int keyIndex = 0; keyIndex < records; keyIndex++, offsetAddr += VectorizedHashAggOperator.PARTITIONINDEX_HTORDINAL_WIDTH) {
        final int keyHash = (int)hashValues.get(keyIndex);
        hashTable = partitions[i].getHashTable();
        actualOrdinals[keyIndex] = hashTable.add(keyFixedVectorAddr, keyVarVectorAddr, keyIndex, keyHash);
        PlatformDependent.putByte(offsetAddr, (byte)hashPartitionIndex);
        PlatformDependent.putInt(offsetAddr + VectorizedHashAggOperator.HTORDINAL_OFFSET, actualOrdinals[keyIndex]);
        PlatformDependent.putInt(offsetAddr + VectorizedHashAggOperator.KEYINDEX_OFFSET, keyIndex);
      }

      /* accumulate */
      partitions[i].getAccumulator().accumulate(offsets.memoryAddress(), records,
                                                partitions[0].getHashTable().getBitsInChunk(),
                                                partitions[0].getHashTable().getChunkOffsetMask());

      /* check hashtable */
      assertArrayEquals(expectedOrdinals, actualOrdinals);
      assertEquals(1, hashTable.getFixedBlockBuffers().size());
      assertEquals(1, hashTable.getFixedBlockBuffers().size());
    }
  }
}
 

public ComplexDataGenerator(int numRows, BufferAllocator allocator) {
  Preconditions.checkState(numRows > 0);
  intsVals = ints(numRows);
  stringsVals = strings(numRows);
  l1IntsVals = l1Ints(numRows);
  l1StringsVals = l1Strings(numRows);
  l2IntsVals = l2Ints(numRows);
  l2StringsVals = l2Strings(numRows);
  l3IntsVals = l3Ints(numRows);
  l3StringsVals = l3Strings(numRows);
  mapVals = mapVals(numRows);

  BatchSchema schema = BatchSchema.newBuilder()
      .addField(ints)
      .addField(strings)
      .addField(l1Ints)
      .addField(l1Strings)
      .addField(l2Ints)
      .addField(l2Strings)
      .addField(l3Ints)
      .addField(l3Strings)
      .addField(map)
      .build();

  container = VectorContainer.create(allocator, schema);

  intsVector = container.addOrGet(ints);
  stringsVector = container.addOrGet(strings);
  l1IntsVector = container.addOrGet(l1Ints);
  addChild(l1IntsVector, INT);
  l1StringsVector = container.addOrGet(l1Strings);
  addChild(l1StringsVector, VARCHAR);
  l2IntsVector = container.addOrGet(l2Ints);
  addChild(addChild(l2IntsVector, LIST), INT);
  l2StringsVector = container.addOrGet(l2Strings);
  addChild(addChild(l2StringsVector, LIST), VARCHAR);
  l3IntsVector = container.addOrGet(l3Ints);
  addChild(addChild(addChild(l3IntsVector, LIST), LIST), INT);
  l3StringsVector = container.addOrGet(l3Strings);
  addChild(addChild(addChild(l3StringsVector, LIST), LIST), VARCHAR);
  structVector = container.addOrGet(map);
  structVector.addOrGet("varchar", nullable(VARCHAR.getType()), VarCharVector.class);
  structVector.addOrGet("int", nullable(INT.getType()), IntVector.class);
  ListVector listVector = structVector.addOrGet("bits", nullable(LIST.getType()), ListVector.class);
  addChild(listVector, BIT);

  tempBuf = allocator.buffer(2048);
}