org.apache.arrow.vector.IntVector Java Examples

private Timestamp getTimestamp(int index, TimeZone tz) throws SFException
{
  long epoch = epochs.getDataBuffer().getLong(index * BigIntVector.TYPE_WIDTH);
  int timeZoneIndex = timeZoneIndices.getDataBuffer().getInt(index * IntVector.TYPE_WIDTH);

  Timestamp ts = ArrowResultUtil.toJavaTimestamp(epoch, context.getScale(columnIndex));

  if (context.getResultVersion() > 0)
  {
    timeZone = SFTimestamp.convertTimezoneIndexToTimeZone(timeZoneIndex);
  }
  else
  {
    timeZone = TimeZone.getTimeZone("UTC");
  }

  Timestamp adjustedTimestamp = ResultUtil.adjustTimestamp(ts);

  return adjustedTimestamp;
}
 

@Test
public void testMultiplyMethod() throws Exception
{
    Schema inputSchema = SchemaBuilder.newBuilder()
            .addField("factor1", Types.MinorType.INT.getType())
            .addField("factor2", Types.MinorType.INT.getType())
            .build();
    Schema outputSchema = SchemaBuilder.newBuilder()
            .addField("product", Types.MinorType.INT.getType())
            .build();

    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(1);
    IntVector inputVector1 = (IntVector) inputRecords.getFieldVector("factor1");
    IntVector inputVector2 = (IntVector) inputRecords.getFieldVector("factor2");
    inputVector1.setSafe(0, 2);
    inputVector2.setSafe(0, 3);

    UserDefinedFunctionResponse response = runAndAssertSerialization(inputRecords, outputSchema, "multiply");

    Block outputRecords = response.getRecords();
    assertEquals(1, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("product");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.multiply(2, 3), arrowValueProjector.project(0));
}
 

@Test
public void testAllocateFixedReAlloc() {
  boolean allocatePassed = false;

  try (BufferAllocator child = allocator.newChildAllocator("test", 0, 128 * 1024)) {
    try (final ValueVector vector = new IntVector("ints", child)) {
      vector.allocateNew();
      allocatePassed = true;

      // realloc will fail after some iterations.
      while (true) {
        vector.reAlloc();
      }
    }
  } catch (Exception e) {
    assertTrue(e instanceof OutOfMemoryException);
    assertTrue(((OutOfMemoryException)e).getOutcomeDetails().isPresent());
    assertTrue(allocatePassed);
  }
}
 

private void testDualIntVectors(int count, Integer[] vectA, Integer[] vectB){
  try(
      IntVector col1 = new IntVector("col1", allocator);
      IntVector col2 = new IntVector("col2", allocator);){
    PivotDef pivot = PivotBuilder.getBlockDefinition(new FieldVectorPair(col1, col1),
        new FieldVectorPair(col2, col2)
        );
    populate(col1, vectA);
    populate(col2, vectB);

    assertEquals(12, pivot.getBlockWidth());

    try(
        FixedBlockVector fixed = new FixedBlockVector(allocator, pivot.getBlockWidth());
        VariableBlockVector variable = new VariableBlockVector(allocator, pivot.getVariableCount());
        ){
      Pivots.pivot(pivot, count, fixed, variable);
      validateIntValues(pivot.getBlockWidth(), pivot.getFixedPivots().get(0), fixed, vectA);
      validateIntValues(pivot.getBlockWidth(), pivot.getFixedPivots().get(1), fixed, vectB);
    }
  }

}
 

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

@Test
public void intAppend(){
  final int count = 1024;
  try(
    IntVector in = new IntVector("in", allocator);
    IntVector out = new IntVector("out", allocator);
  ){

    in.allocateNew(count);
    for(int i = 0; i < count; i++){
      if(i % 5 == 0){
        in.setSafe(i, i);
      }
    }
    in.setValueCount(count);

    List<FieldBufferCopier> copiers = FieldBufferCopier.getCopiers(ImmutableList.<FieldVector>of(in), ImmutableList.<FieldVector>of(out));
    try(
      final SelectionVector2 sv2 = new SelectionVector2(allocator);
    ){

      sv2.allocateNew(count / 2);
      // set alternate elements.
      int x = 0;
      for(long mem = sv2.memoryAddress(); mem < sv2.memoryAddress() + count; mem+=2){
        PlatformDependent.putShort(mem, (short) (char) x);
        x += 2;
      }
      sv2.setRecordCount(count/2);

      append(copiers, sv2);

      out.setValueCount(count / 2);
      for(int i =0; i < count / 2; i++){
        assertEquals(in.getObject(i * 2), out.getObject(i));
      }
    }
  }
}
 

@Test
public void intRoundtrip() throws Exception {
  final int count = 1024;
  try (
    IntVector in = new IntVector("in", allocator);
    IntVector out = new IntVector("out", allocator);
  ) {

    in.allocateNew(count);

    for (int i = 0; i < count; i++) {
      if (i % 5 == 0) {
        in.setSafe(i, i);
      }
    }
    in.setValueCount(count);

    final PivotDef pivot = PivotBuilder.getBlockDefinition(new FieldVectorPair(in, out));
    try (
      final FixedBlockVector fbv = new FixedBlockVector(allocator, pivot.getBlockWidth());
      final VariableBlockVector vbv = new VariableBlockVector(allocator, pivot.getVariableCount());
    ) {
      fbv.ensureAvailableBlocks(count);
      Pivots.pivot4Bytes(pivot.getFixedPivots().get(0), fbv, count);

      ValueVector[] ins = new ValueVector[]{in};
      ValueVector[] outs = new ValueVector[]{out};
      unpivotHelper(pivot, fbv, vbv, ins, outs, 0, count);
      unpivotHelper(pivot, fbv, vbv, ins, outs, 0, 100);
      unpivotHelper(pivot, fbv, vbv, ins, outs, 100, 924);
    }
  }
}
 

public IntToFixedConverter(ValueVector fieldVector, int columnIndex, DataConversionContext context)
{
  super(String.format("%s(%s,%s)", SnowflakeType.FIXED,
                      fieldVector.getField().getMetadata().get("precision"),
                      fieldVector.getField().getMetadata().get("scale")),
        fieldVector,
        columnIndex,
        context);
  this.intVector = (IntVector) fieldVector;
}
 

void readBatchOfDictionaryIds(IntVector intVector, int startOffset, int numValuesToRead,
                              NullabilityHolder nullabilityHolder) {
  int left = numValuesToRead;
  int idx = startOffset;
  while (left > 0) {
    if (this.currentCount == 0) {
      this.readNextGroup();
    }
    int numValues = Math.min(left, this.currentCount);
    switch (mode) {
      case RLE:
        for (int i = 0; i < numValues; i++) {
          intVector.set(idx, currentValue);
          setNotNull(intVector, nullabilityHolder, idx);
          idx++;
        }
        break;
      case PACKED:
        for (int i = 0; i < numValues; i++) {
          intVector.set(idx, packedValuesBuffer[packedValuesBufferIdx++]);
          setNotNull(intVector, nullabilityHolder, idx);
          idx++;
        }
        break;
    }
    left -= numValues;
    currentCount -= numValues;
  }
}
 

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

@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 testRecordBatchAllocs() {
  OptionManager options = Mockito.mock(OptionManager.class);
  when(options.getOption(ExecConstants.TARGET_BATCH_RECORDS_MIN)).thenReturn(127L);
  when(options.getOption(ExecConstants.TARGET_BATCH_RECORDS_MAX)).thenReturn(4095L);
  when(options.getOption(ExecConstants.TARGET_BATCH_SIZE_BYTES)).thenReturn(1024 * 1024L);

  int batchSize = PhysicalPlanCreator.calculateBatchCountFromRecordSize(options, 0);
  Assert.assertTrue(batchSize >= (int)(0.95 * 4096));

  try (BufferAllocator allocator = allocatorRule.newAllocator("test-batch-size", 0, Long.MAX_VALUE)) {
    final ValueVector bitV = new BitVector("bits", allocator);
    AllocationHelper.allocate(bitV, batchSize, 0);
    Assert.assertEquals(1024, allocator.getAllocatedMemory());
    bitV.close();

    final ValueVector intV = new IntVector("ints", allocator);
    AllocationHelper.allocate(intV, batchSize, 0);
    Assert.assertEquals(4096 * 4, allocator.getAllocatedMemory());
    intV.close();

    final ValueVector longV = new BigIntVector("longs", allocator);
    AllocationHelper.allocate(longV, batchSize, 0);
    Assert.assertEquals(4096 * 8, allocator.getAllocatedMemory());
    longV.close();

    final ValueVector decimalV = new DecimalVector("decimals", allocator, 38, 9);
    AllocationHelper.allocate(decimalV, batchSize, 0);
    Assert.assertEquals(4096 * 16, allocator.getAllocatedMemory());
    decimalV.close();
  }
}
 

@Test
public void fixedOnlyWithoutNull(){
  try(IntVector col1 = new IntVector("col1", allocator);
      BigIntVector col2 = new BigIntVector("col2", allocator);
      DecimalVector col3 = new DecimalVector("col3", allocator, 30, 0);
      BitVector col4 = new BitVector("col4", allocator)){

    PivotDef pivot = PivotBuilder.getBlockDefinition(
      new FieldVectorPair(col1, col1),
      new FieldVectorPair(col2, col2),
      new FieldVectorPair(col3, col3),
      new FieldVectorPair(col4, col4)
    );
    Integer col1Val[] = populate4ByteValuesWithoutNull(col1, 4096);
    Long col2Val[] = populate8ByteValuesWithoutNull(col2, 4096);
    BigDecimal col3Val[] = populate16ByteValuesWithoutNull(col3, 4096);
    Boolean col4Val[] = populateBooleanValuesWithoutNull(col4, 4096);

    assertEquals(32, pivot.getBlockWidth());

    try(FixedBlockVector fixed = new FixedBlockVector(allocator, pivot.getBlockWidth(), 4096, true);
        VariableBlockVector variable = new VariableBlockVector(allocator, pivot.getVariableCount(), 4096 * 10, true)){
      fixedOnlyHelper(pivot, fixed, variable, 0, 4096, false, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 128, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 17, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 76, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 39, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 189, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 123, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 1023, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 2046, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 3069, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 4092, 4, true, col1Val, col2Val, col3Val, col4Val);
    }
  }
}
 

public void assertIsSorted(VectorContainer container, int startIndex) {
  final IntVector idVector = container.addOrGet(ID);
  final VarCharVector valueVector = container.addOrGet(VALUE);

  int recordCount = container.getRecordCount();
  int index = startIndex;
  for (int i = 0; i < recordCount; i++, index++) {
    int rowId = sortedRowIds.get(index);
    String value = values.get(rowId);
    assertEquals("non matching ID at row " + index, rowId, idVector.get(i));
    assertEquals("non matching VALUE at row " + index, value, valueVector.getObject(i).toString());
  }
}
 

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

@Test
public void fixedOnly(){
  try(IntVector col1 = new IntVector("col1", allocator);
      BigIntVector col2 = new BigIntVector("col2", allocator);
      DecimalVector col3 = new DecimalVector("col3", allocator, 30, 0);
      BitVector col4 = new BitVector("col4", allocator)){

    PivotDef pivot = PivotBuilder.getBlockDefinition(
        new FieldVectorPair(col1, col1),
        new FieldVectorPair(col2, col2),
        new FieldVectorPair(col3, col3),
        new FieldVectorPair(col4, col4)
    );
    Integer col1Val[] = populate4ByteValues(col1, 4096);
    Long col2Val[] = populate8ByteValues(col2, 4096);
    BigDecimal col3Val[] = populate16ByteValues(col3, 4096);
    Boolean col4Val[] = populateBooleanValues(col4, 4096);

    assertEquals(32, pivot.getBlockWidth());

    try(FixedBlockVector fixed = new FixedBlockVector(allocator, pivot.getBlockWidth(), 4096, true);
        VariableBlockVector variable = new VariableBlockVector(allocator, pivot.getVariableCount(), 4096 * 10, true)){
      fixedOnlyHelper(pivot, fixed, variable, 0, 4096, false, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 128, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 17, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 76, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 39, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 189, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 5, 123, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 0, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 1023, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 2046, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 3069, 1023, true, col1Val, col2Val, col3Val, col4Val);
      fixedOnlyHelper(pivot, fixed, variable, 4092, 4, true, col1Val, col2Val, col3Val, col4Val);
    }
  }
}
 

@Test
public void fixedVariable(){
  try(IntVector col1 = new IntVector("col1", allocator);
      BigIntVector col2 = new BigIntVector("col2", allocator);
      VarCharVector col3 = new VarCharVector("col3", allocator);){

    PivotDef pivot = PivotBuilder.getBlockDefinition(
        new FieldVectorPair(col1, col1),
        new FieldVectorPair(col2, col2),
        new FieldVectorPair(col3, col3)
    );
    Integer col1Val[] = populate4ByteValues(col1, 4096);
    Long col2Val[] = populate8ByteValues(col2, 4096);
    String col3Val[] = populateVarCharValues(col3, 4096);

    assertEquals(20, pivot.getBlockWidth());

    try(FixedBlockVector fixed = new FixedBlockVector(allocator, pivot.getBlockWidth(), 4096, true);
        VariableBlockVector variable = new VariableBlockVector(allocator, pivot.getVariableCount(), 4096 * 2, true)){
      fixedVariableHelper(pivot, fixed, variable, 0, 4096, false, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 0, 128, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 0, 17, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 0, 76, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 5, 39, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 5, 189, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 5, 123, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 0, 1023, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 1023, 1023, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 2046, 1023, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 3069, 1023, true, col1Val, col2Val, col3Val);
      fixedVariableHelper(pivot, fixed, variable, 4092, 4, true, col1Val, col2Val, col3Val);
    }
  }
}
 

private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) throws SFException
{
  long epoch = epochs.getDataBuffer().getLong(index * BigIntVector.TYPE_WIDTH);
  int fraction = fractions.getDataBuffer().getInt(index * IntVector.TYPE_WIDTH);

  if (ArrowResultUtil.isTimestampOverflow(epoch))
  {
    if (fromToString)
    {
      throw new TimestampOperationNotAvailableException(epoch, fraction);
    }
    else
    {
      return null;
    }
  }
  Timestamp ts = ArrowResultUtil.createTimestamp(epoch, fraction, this.treatNTZasUTC);

  // Note: honorClientTZForTimestampNTZ is not enabled for toString method
  // If JDBC_TREAT_TIMESTAMP_NTZ_AS_UTC=false, default behavior is to honor
  // client timezone for NTZ time. Move NTZ timestamp offset to correspond to
  // client's timezone
  if (!this.treatNTZasUTC && !fromToString && context.getHonorClientTZForTimestampNTZ())
  {
    ts = ArrowResultUtil.moveToTimeZone(ts, NTZ, tz);
  }
  Timestamp adjustedTimestamp = ResultUtil.adjustTimestamp(ts);
  return adjustedTimestamp;
}
 

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

protected void baseManyColumnsPartialProjectLeft() throws Exception {
  int columns = 1000;
  int leftColumns = columns;
  int rightColumns = columns;

  JoinInfo joinInfo = getJoinInfo(Arrays.asList(new JoinCondition("EQUALS", f("left_1"), f("right_1"))), JoinRelType.LEFT, ImmutableSet.of(0), ImmutableSet.of(0));

  Table expected = t(getHeader("right", 1, "left", 1), false, getData(1, 1, 1));
  validateDual(joinInfo.operator, joinInfo.clazz,
    new ManyColumnsGenerator<IntVector>(getTestAllocator(), "left", leftColumns, 1, IntVector
      .class, new ArrowType.Int(32, true),  BaseTestJoin::insertIntoIntVector),
    new ManyColumnsGenerator<IntVector>(getTestAllocator(), "right", rightColumns, 1, IntVector
      .class, new ArrowType.Int(32, true), BaseTestJoin::insertIntoIntVector),
    DEFAULT_BATCH, expected);
}
 

private Date getDate(int index, TimeZone tz) throws SFException
{
  if (isNull(index))
  {
    return null;
  }
  else
  {
    int val = dateVector.getDataBuffer().getInt(index * IntVector.TYPE_WIDTH);
    // Note: use default time zone to match with current getDate() behavior
    return ArrowResultUtil.getDate(val);
  }
}
 

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

public TwoFieldStructToTimestampNTZConverter(ValueVector fieldVector, int columnIndex, DataConversionContext context)
{
  super(SnowflakeType.TIMESTAMP_NTZ.name(), fieldVector, columnIndex, context);
  structVector = (StructVector) fieldVector;
  epochs = structVector.getChild(FIELD_NAME_EPOCH, BigIntVector.class);
  fractions = structVector.getChild(FIELD_NAME_FRACTION, IntVector.class);
}
 

@Override
public int toInt(int index)
{
  if (isNull(index))
  {
    return 0;
  }
  else
  {
    int val = dateVector.getDataBuffer().getInt(index * IntVector.TYPE_WIDTH);
    return val;
  }
}
 

@Test
public void testStructCopier() {
  try (final StructVector s1 = StructVector.empty("s1", allocator);
       final StructVector s2 = StructVector.empty("s2", allocator)) {
    int rowcount = 5000;
    s1.addOrGet("struct_child", FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
    s2.addOrGet("struct_child", FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
    s1.allocateNew();
    BaseWriter.StructWriter structWriter = new NullableStructWriter(s1);
    IntWriter intWriter = structWriter.integer("struct_child");
    for(int i = 0; i<rowcount; i++) {
      if (i % 10 == 0) {
        continue;
      }
      structWriter.setPosition(i);
      structWriter.start();
      intWriter.writeInt(i);
      structWriter.end();
    }
    s1.setValueCount(rowcount);
    HiveParquetCopier.StructCopier structCopier = new HiveParquetCopier.StructCopier(s1, s2);
    structCopier.copyNonDataBufferRefs(rowcount);
    Assert.assertEquals(rowcount, s2.getValueCount());
    Assert.assertEquals(500, s1.getNullCount());
    Assert.assertEquals(500, s2.getNullCount());
    for(int i=0; i<rowcount; ++i) {
      Assert.assertEquals( i % 10 == 0, s2.isNull(i));
    }
  }
}
 

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

@SuppressWarnings({ "unchecked", "rawtypes" })
private VectorAccessible getBatch(MinorType type){
  final IntVector vector = new IntVector("result", RootAllocatorFactory.newRoot(DEFAULT_SABOT_CONFIG));
  VectorWrapper<IntVector> wrapper = Mockito.mock(VectorWrapper.class);
  when(wrapper.getValueVector()).thenReturn(vector);

  final TypedFieldId tfid = new TypedFieldId(CompleteType.fromMinorType(type), false, 0);
  VectorAccessible batch = Mockito.mock(VectorAccessible.class);
  when(batch.getValueVectorId(any(SchemaPath.class))).thenReturn(tfid);
  when(batch.getValueAccessorById(any(Class.class), any(int[].class))).thenReturn((VectorWrapper) wrapper);
  when(batch.getSchema()).thenReturn(BatchSchema.newBuilder().addField(MajorTypeHelper.getFieldForNameAndMajorType("result", Types.optional(type))).setSelectionVectorMode(SelectionVectorMode.NONE).build());
  return batch;
}
 

@Override
public byte[] toBytes(int index) throws SFException
{
  if (isNull(index))
  {
    return null;
  }
  else
  {
    byteBuf.putInt(0, intVector.getDataBuffer().getInt(index * IntVector.TYPE_WIDTH));
    return byteBuf.array();
  }
}
 

public TwoFieldStructToTimestampLTZConverter(ValueVector fieldVector, int columnIndex, DataConversionContext context)
{
  super(SnowflakeType.TIMESTAMP_LTZ.name(), fieldVector, columnIndex, context);
  structVector = (StructVector) fieldVector;
  epochs = structVector.getChild(FIELD_NAME_EPOCH, BigIntVector.class);
  fractions = structVector.getChild(FIELD_NAME_FRACTION, IntVector.class);
}
 

public ThreeFieldStructToTimestampTZConverter(ValueVector fieldVector, int columnIndex, DataConversionContext context)
{
  super(SnowflakeType.TIMESTAMP_LTZ.name(), fieldVector, columnIndex, context);
  structVector = (StructVector) fieldVector;
  epochs = structVector.getChild(FIELD_NAME_EPOCH, BigIntVector.class);
  fractions = structVector.getChild(FIELD_NAME_FRACTION, IntVector.class);
  timeZoneIndices = structVector.getChild(FIELD_NAME_TIME_ZONE_INDEX, IntVector.class);
}