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

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

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 testGetBooleanNoScale() throws SFException
{
  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "0");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  IntVector vector = new IntVector("col_one", fieldType, allocator);
  vector.setSafe(0, 0);
  vector.setSafe(1, 1);
  vector.setNull(2);
  vector.setSafe(3, 5);

  ArrowVectorConverter converter = new IntToFixedConverter(vector, 0, this);

  assertThat(false, is(converter.toBoolean(0)));
  assertThat(true, is(converter.toBoolean(1)));
  assertThat(false, is(converter.toBoolean(2)));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toBoolean(3));

  vector.close();
}
 

@Test
public void testInvalidConversion()
{
  // try convert to int/long/byte/short with scale > 0
  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "3");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  IntVector vector = new IntVector("col_one", fieldType, allocator);
  vector.setSafe(0, 33000);

  final ArrowVectorConverter converter = new IntToScaledFixedConverter(vector, 0, this, 3);
  final int invalidConversionErrorCode =
      ErrorCode.INVALID_VALUE_CONVERT.getMessageCode();

  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toBoolean(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toLong(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toInt(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toShort(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toByte(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toDate(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toTime(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toTimestamp(0, TimeZone.getDefault()));
  vector.clear();
}
 

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

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

@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 testGetBooleanWithScale() throws SFException
{
  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "3");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  IntVector vector = new IntVector("col_one", fieldType, allocator);
  vector.setSafe(0, 0);
  vector.setSafe(1, 1);
  vector.setNull(2);
  vector.setSafe(3, 5);

  final ArrowVectorConverter converter = new IntToScaledFixedConverter(vector, 0, this, 3);

  assertThat(false, is(converter.toBoolean(0)));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toBoolean(3));
  assertThat(false, is(converter.toBoolean(2)));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converter.toBoolean(3));

  vector.close();
}
 

@Before
public void beforeTest()
        throws IOException
{
    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);

    expected = new UserDefinedFunctionRequest(federatedIdentity,
            inputRecords,
            outputSchema,
            "test-method",
            UserDefinedFunctionType.SCALAR);


    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "UserDefinedFunctionRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
 

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

private static void insertIntoIntVector(int index, int value, BaseFixedWidthVector vector) {
  IntVector vec = (IntVector)vector;
  vec.setSafe(index, value);
}
 

@Test
public void testFixedNoScale() throws SFException
{
  final int rowCount = 1000;
  List<Integer> expectedValues = new ArrayList<>();
  Set<Integer> nullValIndex = new HashSet<>();
  for (int i = 0; i < rowCount; i++)
  {
    expectedValues.add(random.nextInt());
  }

  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "0");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  IntVector vector = new IntVector("col_one", fieldType, allocator);
  for (int i = 0; i < rowCount; i++)
  {
    boolean isNull = random.nextBoolean();
    if (isNull)
    {
      vector.setNull(i);
      nullValIndex.add(i);
    }
    else
    {
      vector.setSafe(i, expectedValues.get(i));
    }
  }

  ArrowVectorConverter converter = new IntToFixedConverter(vector, 0, this);

  for (int i = 0; i < rowCount; i++)
  {
    int intVal = converter.toInt(i);
    Object longObj = converter.toObject(i);
    String intString = converter.toString(i);

    if (nullValIndex.contains(i))
    {
      assertThat(intVal, is(0));
      assertThat(longObj, is(nullValue()));
      assertThat(intString, is(nullValue()));
      assertThat(converter.toBytes(i), is(nullValue()));
    }
    else
    {
      assertThat(intVal, is(expectedValues.get(i)));
      assertEquals(longObj, (long) expectedValues.get(i));
      assertThat(intString, is(expectedValues.get(i).toString()));
      bb = ByteBuffer.wrap(converter.toBytes(i));
      assertThat(intVal, is(bb.getInt()));
    }
  }
  vector.clear();
}
 

@Test
public void testFixedWithScale() throws SFException
{
  final int rowCount = 1000;
  List<Integer> expectedValues = new ArrayList<>();
  Set<Integer> nullValIndex = new HashSet<>();
  for (int i = 0; i < rowCount; i++)
  {
    expectedValues.add(random.nextInt());
  }

  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "3");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  IntVector vector = new IntVector("col_one", fieldType, allocator);
  for (int i = 0; i < rowCount; i++)
  {
    boolean isNull = random.nextBoolean();
    if (isNull)
    {
      vector.setNull(i);
      nullValIndex.add(i);
    }
    else
    {
      vector.setSafe(i, expectedValues.get(i));
    }
  }

  ArrowVectorConverter converter = new IntToScaledFixedConverter(vector, 0, this, 3);

  for (int i = 0; i < rowCount; i++)
  {
    BigDecimal bigDecimalVal = converter.toBigDecimal(i);
    Object objectVal = converter.toObject(i);
    String stringVal = converter.toString(i);

    if (nullValIndex.contains(i))
    {
      assertThat(bigDecimalVal, nullValue());
      assertThat(objectVal, nullValue());
      assertThat(stringVal, nullValue());
      assertThat(converter.toBytes(i), is(nullValue()));
    }
    else
    {
      BigDecimal expectedVal = BigDecimal.valueOf(expectedValues.get(i), 3);
      assertThat(bigDecimalVal, is(expectedVal));
      assertThat(objectVal, is(expectedVal));
      assertThat(stringVal, is(expectedVal.toString()));
      assertThat(converter.toBytes(i), is(notNullValue()));
    }
  }

  vector.clear();
}
 

@Test
public void testGetSmallerIntegralType() throws SFException
{
  // try convert to int/long/byte/short with scale > 0
  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "FIXED");
  customFieldMeta.put("precision", "10");
  customFieldMeta.put("scale", "0");

  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.INT.getType(),
                                      null, customFieldMeta);

  // test value which is out of range of int, but falls in long
  IntVector vectorFoo = new IntVector("col_one", fieldType, allocator);
  vectorFoo.setSafe(0, 33000);
  vectorFoo.setSafe(1, -33000);

  final ArrowVectorConverter converterFoo =
      new IntToFixedConverter(vectorFoo, 0, this);
  final int invalidConversionErrorCode =
      ErrorCode.INVALID_VALUE_CONVERT.getMessageCode();

  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converterFoo.toShort(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converterFoo.toByte(0));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converterFoo.toShort(1));
  TestUtil.assertSFException(invalidConversionErrorCode,
                             () -> converterFoo.toByte(1));

  assertThat(converterFoo.toLong(0), is(33000L));
  assertThat(converterFoo.toLong(1), is(-33000L));
  vectorFoo.clear();

  // test value which is in range of byte, all get method should return
  IntVector vectorBar = new IntVector("col_one", fieldType, allocator);
  // set value which is out of range of int, but falls in long
  vectorBar.setSafe(0, 10);
  vectorBar.setSafe(1, -10);

  final ArrowVectorConverter converterBar =
      new IntToFixedConverter(vectorBar, 0, this);

  assertThat(converterBar.toByte(0), is((byte) 10));
  assertThat(converterBar.toByte(1), is((byte) -10));
  assertThat(converterBar.toShort(0), is((short) 10));
  assertThat(converterBar.toShort(1), is((short) -10));
  assertThat(converterBar.toLong(0), is(10L));
  assertThat(converterBar.toLong(1), is(-10L));
  vectorBar.clear();
}
 

private void writeColumn(FieldVector fieldVector, int idx)
{
    if (fieldVector instanceof IntVector) {
        IntVector intVector = (IntVector) fieldVector;
        intVector.setSafe(idx, idx + 100);
        return;
    }

    if (fieldVector instanceof Float4Vector) {
        Float4Vector float4Vector = (Float4Vector) fieldVector;
        float4Vector.setSafe(idx, idx + 100.1f);
        return;
    }

    if (fieldVector instanceof Float8Vector) {
        Float8Vector float8Vector = (Float8Vector) fieldVector;
        float8Vector.setSafe(idx, idx + 100.2);
        return;
    }

    if (fieldVector instanceof VarCharVector) {
        VarCharVector varCharVector = (VarCharVector) fieldVector;
        varCharVector.setSafe(idx, new Text(idx + "-my-varchar"));
        return;
    }

    if (fieldVector instanceof ListVector) {
        BlockUtils.setComplexValue(fieldVector,
                idx,
                FieldResolver.DEFAULT,
                ImmutableList.of(idx + 100, idx + 200, idx + 300));
        return;
    }

    if (fieldVector instanceof StructVector) {
        Map<String, Object> input = ImmutableMap.of("intVal", idx + 100, "doubleVal", idx + 200.2);
        BlockUtils.setComplexValue(fieldVector,
                idx,
                FieldResolver.DEFAULT,
                input);
        return;
    }

    throw new IllegalArgumentException("Unsupported fieldVector " + fieldVector.getClass().getCanonicalName());
}