Java Code Examples for org.apache.arrow.vector.IntVector#allocateNew()

private Pair<IntVector, ResultVerifier> testIntVector() {
  final String colName = "colInt";
  final List<Integer> values =  asList(20, 50, -2000, 327345, null);

  IntVector valuesVector = new IntVector(colName, allocator);
  valuesVector.allocateNew(values.size());
  for (int i = 0; i < values.size(); i++) {
    if (values.get(i) == null) {
      valuesVector.setNull(i);
    } else {
      valuesVector.set(i, values.get(i));
    }
  }

  ResultVerifier verifier = new ResultVerifier() {
    @Override
    public void verify(DataPOJO output) {
      verifyIntValues(values, output, colName);
    }
  };

  return Pair.of(valuesVector, verifier);
}
 

private void populateInt(IntVector vector, Integer[] data) {
  vector.allocateNew();
  Random r = new Random();
  for(int i =0; i < data.length; i++){
    Integer val = data[i];
    if(val != null){
      vector.setSafe(i, val);
    } else {
      vector.setSafe(i, 0, r.nextInt());
    }
  }
  vector.setValueCount(data.length);
}
 

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

private void populateInt(IntVector vector, Integer[] data) {
  vector.allocateNew();
  Random r = new Random();
  for(int i =0; i < data.length; i++){
    Integer val = data[i];
    if(val != null){
      vector.setSafe(i, val);
    } else {
      vector.setSafe(i, 0, r.nextInt());
    }
  }
  vector.setValueCount(data.length);
}
 

static Integer[] populate4ByteValuesWithoutNull(IntVector vector, int size){
  vector.allocateNew();
  Integer values[] = new Integer[size];
  for(int i =0; i < size; i++){
    values[i] = RAND.nextInt();
    vector.setSafe(i, values[i]);
  }
  vector.setValueCount(size);
  return values;
}
 

static Integer[] populate4ByteValues(IntVector vector, int size){
  vector.allocateNew();
  Integer values[] = new Integer[size];
  for(int i =0; i < size; i++){
    if(RAND.nextBoolean()){
      values[i] = RAND.nextInt();
      vector.setSafe(i, values[i]);
    }
  }
  vector.setValueCount(size);
  return values;
}
 

private void populateInt(IntVector vector, Integer[] data) {
  vector.allocateNew();
  Random r = new Random();
  for(int i =0; i < data.length; i++){
    Integer val = data[i];
    if(val != null){
      vector.setSafe(i, val);
    } else {
      vector.setSafe(i, 0, r.nextInt());
    }
  }
  vector.setValueCount(data.length);
}
 

static void populate(IntVector vector, Integer[] values){
  vector.allocateNew();
  Random r = new Random();
  for(int i =0; i < values.length; i++){
    Integer val = values[i];
    if(val != null){
      vector.setSafe(i, val);
    } else {
      // add noise so we make sure not to read/see noise.
      vector.setSafe(i, 0, r.nextInt());
    }
  }
  vector.setValueCount(values.length);
}
 

private IntVector intVector(String name) {
  IntVector vec = new IntVector(name, allocator);
  vec.allocateNew(5);
  vec.set(0, 20);
  vec.set(1, 50);
  vec.set(2, -2000);
  vec.set(3, 327345);
  vec.setNull(4);

  vec.setValueCount(COUNT);
  return vec;
}
 

@Override
public void setup(VectorAccessible incoming) throws Exception {
  state.is(State.NEEDS_SETUP);
  Preconditions.checkState(incoming.getSchema().getSelectionVectorMode() == BatchSchema.SelectionVectorMode.NONE,
    "Vectorized Partition Sender doesn't support SV " + incoming.getSchema().getSelectionVectorMode());

  checkSchema(incoming.getSchema());

  // how many records we can keep in memory before we are forced to flush the outgoing batch
  numRecordsBeforeFlush = config.getProps().getTargetBatchSize();
  stats.setLongStat(Metric.BUCKET_SIZE, numRecordsBeforeFlush);

  //
  final BufferAllocator allocator = context.getAllocator();
  // we need to synchronize this to ensure no receiver termination message is lost
  synchronized (batchCreationLock) {
    initBatchesAndLookup(incoming);

    // some receivers may have finished already, make sure to terminate corresponding outgoing batches
    for (int index = 0; index < terminations.size(); index++) {
      final int p = terminations.buffer[index];
      batches[p].terminate();
      batches[p + numReceivers].terminate();
    }
    terminations.clear();
  }

  copiers = MultiDestCopier.getCopiers(VectorContainer.getFieldVectors(incoming), batches, copyWatches);

  copyIndices = new IntVector("copy-compound-indices", allocator);
  copyIndices.allocateNew(numRecordsBeforeFlush);

  initHashVector(incoming);

  state = State.CAN_CONSUME;
}
 

private static VectorContainer buildIntegerGlobalDictionary(List<Dictionary> dictionaries, VectorContainer existingDict, ColumnDescriptor columnDescriptor, BufferAllocator bufferAllocator) {
  final Field field = new Field(SchemaPath.getCompoundPath(columnDescriptor.getPath()).getAsUnescapedPath(), true, new ArrowType.Int(32, true), null);
  final VectorContainer input = new VectorContainer(bufferAllocator);
  final IntVector intVector = input.addOrGet(field);
  intVector.allocateNew();
  final SortedSet<Integer> values = Sets.newTreeSet();
  for (Dictionary dictionary : dictionaries) {
    for (int i = 0; i <= dictionary.getMaxId(); ++i) {
      values.add(dictionary.decodeToInt(i));
    }
  }
  if (existingDict != null) {
    final IntVector existingDictValues = existingDict.getValueAccessorById(IntVector.class, 0).getValueVector();
    for (int i = 0; i < existingDict.getRecordCount(); ++i) {
      values.add(existingDictValues.get(i));
    }
  }
  final Iterator<Integer> iter = values.iterator();
  int recordCount = 0;
  while (iter.hasNext()) {
    intVector.setSafe(recordCount++, iter.next());
  }
  intVector.setValueCount(recordCount);
  input.setRecordCount(recordCount);
  input.buildSchema(BatchSchema.SelectionVectorMode.NONE);
  return input;
}
 

@Override
public <T extends FieldVector> void generate(T intVector1, int from, int items, int isSet) {
  IntVector intVector = (IntVector) intVector1;
  intVector.setInitialCapacity(items);
  intVector.allocateNew();
  for (int i = 0; i < items; i++) {
    intVector.setSafe(i, isSet, (int) dataList.get(from + i));
  }
  intVector.setValueCount(items);
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  IntVector vector = new IntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , 0 );  
  vector.setSafe( 1 , 1 );  
  vector.setSafe( 2 , 0 );  
  vector.setNull( 3 );  
  vector.setSafe( 4 , 1 );  
  vector.setSafe( 5 , 1 );  
  vector.setSafe( 6 , 1 );  
  vector.setNull( 7 );  
  vector.setValueCount( 8 );

  IColumn column = ArrowColumnFactory.convert( "test" , vector );
  assertEquals( column.getColumnName() , "test" );
  assertEquals( column.size() , 8 );
  assertTrue( ( column.getColumnType() == ColumnType.INTEGER ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getInt() , 0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getInt() , 1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getInt() , 0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getInt() , 1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  IntVector vector = new IntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , 0 );
  vector.setSafe( 1 , 1 );
  vector.setSafe( 2 , 0 );
  vector.setNull( 3 );
  vector.setSafe( 4 , 1 );
  vector.setSafe( 5 , 1 );
  vector.setSafe( 6 , 1 );
  vector.setNull( 7 );
  vector.setValueCount( 8 );

  IColumn column = ArrowColumnFactory.convert( "test" , vector );
  assertEquals( column.getColumnName() , "test" );
  assertEquals( column.size() , 8 );
  assertTrue( ( column.getColumnType() == ColumnType.INTEGER ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getInt() , 0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getInt() , 1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getInt() , 0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getInt() , 1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

public ArrowIntegerMemoryAllocator( final IntVector vector , final int rowCount ) {
  vector.allocateNew( rowCount );
  this.vector = vector;
}
 

@Test
public void intManyRoundtrip() throws Exception {
  final int count = 1024;
  final int mult = 80;

  ValueVector[] in = new IntVector[mult];
  ValueVector[] out = new IntVector[mult];
  List<FieldVectorPair> pairs = new ArrayList<>();
  try {

    for (int x = 0; x < mult; x++) {
      IntVector inv = new IntVector("in", allocator);
      in[x] = inv;
      inv.allocateNew(count);
      IntVector outv = new IntVector("out", allocator);
      out[x] = outv;
      for (int i = 0; i < count; i++) {
        if (i % 5 == 0) {
          inv.setSafe(i, Integer.MAX_VALUE - i);
        }
      }
      inv.setValueCount(count);
      pairs.add(new FieldVectorPair(inv, outv));
    }

    final PivotDef pivot = PivotBuilder.getBlockDefinition(pairs);
    try (
      final FixedBlockVector fbv = new FixedBlockVector(allocator, pivot.getBlockWidth());
      final VariableBlockVector vbv = new VariableBlockVector(allocator, pivot.getVariableCount());
    ) {
      fbv.ensureAvailableBlocks(count);

      for (int x = 0; x < mult; x++) {
        Pivots.pivot4Bytes(pivot.getFixedPivots().get(x), fbv, count);
      }

      unpivotHelper(pivot, fbv, vbv, in, out, 0, count);
      unpivotHelper(pivot, fbv, vbv, in, out, 0, 100);
      unpivotHelper(pivot, fbv, vbv, in, out, 100, 924);
    }
  } finally {
    AutoCloseables.close(ImmutableList.copyOf(in));
    AutoCloseables.close(ImmutableList.copyOf(out));
  }
}
 

public ArrowIntegerMemoryAllocator( final IntVector vector , final int rowCount ){
  vector.allocateNew( rowCount );
  this.vector = vector;
}