Java Code Examples for org.apache.arrow.vector.FieldVector#setValueCount()

/**
 * This is done after reading every spilled batch so that
 * we can correctly identify length of data read in every batch
 * and the length doesn't end up being cumulative which will then
 * lead to incorrect results, segfaults etc.
 *
 * It is probably not necessary to zero out the buffers again
 * Overwriting the contents should be fine as long as we
 * correctly identify the length read in each iteration.
 *
 * Zeroing out is likely an expensive operation so unless
 * we figure out it is absolutely necessary for correctness,
 * we will just reset the reader/writer indexes.
 */
public void reset() {
  recordsInBatch = 0;
  fixedKeyColPivotedData.readerIndex(0);
  fixedKeyColPivotedData.writerIndex(0);
  variableKeyColPivotedData.readerIndex(0);
  variableKeyColPivotedData.writerIndex(0);
  for (FieldVector vector : postSpillAccumulatorVectors) {
    final ArrowBuf validityBuffer = vector.getValidityBuffer();
    final ArrowBuf dataBuffer = vector.getDataBuffer();
    validityBuffer.readerIndex(0);
    validityBuffer.writerIndex(0);
    dataBuffer.readerIndex(0);
    dataBuffer.writerIndex(0);
    validityBuffer.setZero(0, validityBuffer.capacity());
    dataBuffer.setZero(0, dataBuffer.capacity());
    vector.setValueCount(0);
  }
}
 

private void writeData(Block block, int numOfRows)
{
    for (FieldVector fieldVector : block.getFieldVectors()) {
        fieldVector.setInitialCapacity(numOfRows);
        fieldVector.allocateNew();
        fieldVector.setValueCount(numOfRows);

        for (int idx = 0; idx < numOfRows; ++idx) {
            writeColumn(fieldVector, idx);
        }
    }
}
 

public void nextBatchIntegers(FieldVector fieldVector, int typeWidth, NullabilityHolder holder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchIntegers(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, typeWidth, holder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchLongs(FieldVector fieldVector, int typeWidth, NullabilityHolder holder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchLongs(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, typeWidth, holder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchTimestampMillis(FieldVector fieldVector, int typeWidth, NullabilityHolder holder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchTimestampMillis(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, typeWidth, holder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchFloats(FieldVector fieldVector, int typeWidth, NullabilityHolder holder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchFloats(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, typeWidth, holder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchDoubles(FieldVector fieldVector, int typeWidth, NullabilityHolder holder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchDoubles(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, typeWidth, holder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchIntLongBackedDecimal(
    FieldVector fieldVector,
    int typeWidth,
    NullabilityHolder nullabilityHolder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch =
        vectorizedPageIterator.nextBatchIntLongBackedDecimal(fieldVector, batchSize - rowsReadSoFar,
            rowsReadSoFar, typeWidth, nullabilityHolder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchFixedLengthDecimal(
    FieldVector fieldVector,
    int typeWidth,
    NullabilityHolder nullabilityHolder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch =
        vectorizedPageIterator.nextBatchFixedLengthDecimal(fieldVector, batchSize - rowsReadSoFar,
            rowsReadSoFar, typeWidth, nullabilityHolder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchVarWidthType(FieldVector fieldVector, NullabilityHolder nullabilityHolder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchVarWidthType(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, nullabilityHolder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchFixedWidthBinary(FieldVector fieldVector, int typeWidth, NullabilityHolder nullabilityHolder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch =
        vectorizedPageIterator.nextBatchFixedWidthBinary(fieldVector, batchSize - rowsReadSoFar,
            rowsReadSoFar, typeWidth, nullabilityHolder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

public void nextBatchBoolean(FieldVector fieldVector, NullabilityHolder nullabilityHolder) {
  int rowsReadSoFar = 0;
  while (rowsReadSoFar < batchSize && hasNext()) {
    advance();
    int rowsInThisBatch = vectorizedPageIterator.nextBatchBoolean(fieldVector, batchSize - rowsReadSoFar,
        rowsReadSoFar, nullabilityHolder);
    rowsReadSoFar += rowsInThisBatch;
    this.triplesRead += rowsInThisBatch;
    fieldVector.setValueCount(rowsReadSoFar);
  }
}
 

private void resetFirstAccumulatorVector() {
  Preconditions.checkArgument(accumulators.length == 1, "Error: incorrect number of batches in accumulator");
  final FieldVector vector = accumulators[0];
  Preconditions.checkArgument(vector != null, "Error: expecting a valid accumulator");
  final ArrowBuf validityBuffer = vector.getValidityBuffer();
  final ArrowBuf dataBuffer = vector.getDataBuffer();
  validityBuffer.readerIndex(0);
  validityBuffer.writerIndex(0);
  dataBuffer.readerIndex(0);
  dataBuffer.writerIndex(0);
  initialize(vector);
  vector.setValueCount(0);
}
 

private void writeIntToField(FieldVector fieldVector, Object[] data,
                             int startIndex, int rowsToAppend)
{
  IntVector intVector = (IntVector) fieldVector;
  intVector.setInitialCapacity(rowsToAppend);
  intVector.allocateNew();
  for (int i = 0; i < rowsToAppend; i++)
  {
    intVector.setSafe(i, 1, (int) data[startIndex + i]);
  }
  // how many are set
  fieldVector.setValueCount(rowsToAppend);
}
 

/**
 * Take the accumulator vector (the vector that stores computed values)
 * for a particular batch (identified by batchIndex) and output its contents.
 * Output is done by transferring the contents from accumulator vector
 * to its counterpart in outgoing container. As part of transfer,
 * the memory ownership (along with data) is transferred from source
 * vector's allocator to target vector's allocator and source vector's
 * memory is released.
 *
 * While the transfer is good as it essentially
 * avoids copy, we still want the memory associated with
 * allocator of source vector because of post-spill processing
 * where this accumulator vector will still continue to store
 * the computed values as we start treating spilled batches
 * as new input into the operator.
 *
 * This is why we need to immediately allocate
 * the accumulator vector after transfer is done. However we do this
 * for a singe batch only as once we are done outputting a partition,
 * we anyway get rid of all but 1 batch.
 *
 * @param batchIndex batch to output
 */
@Override
public void output(final int batchIndex) {
  final FieldVector accumulationVector = accumulators[batchIndex];
  final TransferPair transferPair = accumulationVector.makeTransferPair(transferVector);
  transferPair.transfer();
  if (batchIndex == 0) {
    ((FixedWidthVector) accumulationVector).allocateNew(maxValuesPerBatch);
    accumulationVector.setValueCount(0);
    initialize(accumulationVector);
    bitAddresses[batchIndex] = accumulationVector.getValidityBufferAddress();
    valueAddresses[batchIndex] = accumulationVector.getDataBufferAddress();
  }
}