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

@Test
public void testGetDefaultValueIfNullMethod() throws Exception
{
    Schema inputSchema = SchemaBuilder.newBuilder()
            .addField("input", Types.MinorType.BIGINT.getType())
            .build();
    Schema outputSchema = SchemaBuilder.newBuilder()
            .addField("output", Types.MinorType.BIGINT.getType())
            .build();

    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(2);
    BigIntVector fieldVector = (BigIntVector) inputRecords.getFieldVector("input");
    fieldVector.setSafe(0, 123l);
    fieldVector.setNull(1);

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

    Block outputRecords = response.getRecords();
    assertEquals(2, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("output");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.get_default_value_if_null(123l), arrowValueProjector.project(0));
    assertEquals(exampleUserDefinedFunctionHandler.get_default_value_if_null(null), arrowValueProjector.project(1));
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  BigIntVector vector = new BigIntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , (long)0 );
  vector.setSafe( 1 , (long)1 );
  vector.setSafe( 2 , (long)0 );
  vector.setNull( 3 );
  vector.setSafe( 4 , (long)1 );
  vector.setSafe( 5 , (long)1 );
  vector.setSafe( 6 , (long)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.LONG ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getLong() , (long)0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getLong() , (long)1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getLong() , (long)0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getLong() , (long)1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getLong() , (long)1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getLong() , (long)1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  BigIntVector vector = new BigIntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , (long)0 );  
  vector.setSafe( 1 , (long)1 );  
  vector.setSafe( 2 , (long)0 );  
  vector.setNull( 3 );  
  vector.setSafe( 4 , (long)1 );  
  vector.setSafe( 5 , (long)1 );  
  vector.setSafe( 6 , (long)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.LONG ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getLong() , (long)0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getLong() , (long)1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getLong() , (long)0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getLong() , (long)1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getLong() , (long)1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getLong() , (long)1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

@Override
public <T extends FieldVector> void generate(T bigIntVector1, int from, int items, int isSet) {
  BigIntVector bigIntVector = (BigIntVector) bigIntVector1;
  bigIntVector.setInitialCapacity(items);
  bigIntVector.allocateNew();
  for (int i = 0; i < items; i++) {
    Long l = new Long(dataList.get(from + i).toString());
    bigIntVector.setSafe(i, isSet, l);
  }
  bigIntVector.setValueCount(items);
}
 

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

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

static Long[] populate8ByteValues(BigIntVector vector, int size){
  vector.allocateNew();
  Long values[] = new Long[size];
  for(int i = 0; i < values.length; i++){
    if (RAND.nextBoolean()) {
      values[i] = RAND.nextLong();
      vector.setSafe(i, values[i]);
    }
  }
  vector.setValueCount(values.length);
  return values;
}
 

static Long[] populate8ByteValuesWithoutNull(BigIntVector vector, int size){
  vector.allocateNew();
  Long values[] = new Long[size];
  for(int i = 0; i < values.length; i++){
    values[i] = RAND.nextLong();
    vector.setSafe(i, values[i]);
  }
  vector.setValueCount(values.length);
  return values;
}
 

@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.BIGINT.getType(),
                                      null, customFieldMeta);

  BigIntVector vector = new BigIntVector("col_one", fieldType, allocator);
  vector.setSafe(0, 123456789L);

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

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

@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.BIGINT.getType(),
                                      null, customFieldMeta);

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

  ArrowVectorConverter converter = new BigIntToFixedConverter(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 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.BIGINT.getType(),
                                      null, customFieldMeta);

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

  final ArrowVectorConverter converter = new BigIntToScaledFixedConverter(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();
}
 

@Test
public void testDictionaryLookup() throws Throwable {


  try (final VectorContainer dict1 = new VectorContainer(getTestAllocator());
       final VectorContainer dict2 = new VectorContainer(getTestAllocator());
       final VectorContainer dict3 = new VectorContainer(getTestAllocator())) {

    final Map<String, GlobalDictionaryFieldInfo> dictionaryFieldInfoMap = Maps.newHashMap();
    final Field field1 = new Field(SchemaPath.getSimplePath("c0").getAsUnescapedPath(), true, new ArrowType.Int(64, true), null);
    final BigIntVector longVector = dict1.addOrGet(field1);
    longVector.allocateNew();
    longVector.setSafe(0, 10L);
    longVector.setSafe(1, 20L);
    longVector.setSafe(2, 30L);
    longVector.setSafe(3, 40L);
    longVector.setSafe(4, 50L);
    longVector.setValueCount(5);
    dict1.setRecordCount(5);
    dict1.buildSchema(BatchSchema.SelectionVectorMode.NONE);


    dictionaryFieldInfoMap.put("c0", new GlobalDictionaryFieldInfo(0, "c0", null, field1.getType(), "local"));

    final Field field2 = new Field(SchemaPath.getSimplePath("c1").getAsUnescapedPath(), true, new ArrowType.Binary(), null);
    final VarBinaryVector binaryVector = dict2.addOrGet(field2);
    binaryVector.allocateNew();
    binaryVector.setSafe(0, "abc".getBytes(UTF8), 0, 3);
    binaryVector.setSafe(1, "bcd".getBytes(UTF8), 0, 3);
    binaryVector.setSafe(2, "cde".getBytes(UTF8), 0, 3);
    binaryVector.setSafe(3, "def".getBytes(UTF8), 0, 3);
    binaryVector.setSafe(4, "efg".getBytes(UTF8), 0, 3);
    binaryVector.setValueCount(5);
    dict2.setRecordCount(5);
    dict2.buildSchema(BatchSchema.SelectionVectorMode.NONE);
    dictionaryFieldInfoMap.put("c1", new GlobalDictionaryFieldInfo(0, "c1", null, field2.getType(), "local"));

    final Field field3 = new Field(SchemaPath.getSimplePath("c2").getAsUnescapedPath(), true, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null);
    final Float8Vector doubleVector = dict3.addOrGet(field3);
    doubleVector.allocateNew();
    doubleVector.setSafe(0, 100.1);
    doubleVector.setSafe(1, 200.2);
    doubleVector.setSafe(2, 300.3);
    doubleVector.setSafe(3, 400.4);
    doubleVector.setSafe(4, 500.5);
    doubleVector.setValueCount(5);
    dict3.setRecordCount(5);
    dict3.buildSchema(BatchSchema.SelectionVectorMode.NONE);
    dictionaryFieldInfoMap.put("c2", new GlobalDictionaryFieldInfo(0, "c2", null, field3.getType(), "local"));

    OperatorCreatorRegistry registry = Mockito.mock(OperatorCreatorRegistry.class);
    Mockito.when(registry.getSingleInputOperator(Matchers.any(OperatorContext.class), Matchers.any(PhysicalOperator.class)))
      .thenAnswer(new Answer<SingleInputOperator>() {
        public SingleInputOperator answer(InvocationOnMock invocation) throws Exception {
          Object[] args = invocation.getArguments();
          DictionaryLookupOperator dictionaryLookupOperator = Mockito.spy(new DictionaryLookupOperator(
            (OperatorContext)args[0], (DictionaryLookupPOP)args[1]));

          Mockito.doReturn(dict1).when(dictionaryLookupOperator).loadDictionary(Matchers.eq("c0"));
          Mockito.doReturn(dict2).when(dictionaryLookupOperator).loadDictionary(Matchers.eq("c1"));
          Mockito.doReturn(dict3).when(dictionaryLookupOperator).loadDictionary(Matchers.eq("c2"));
          return dictionaryLookupOperator;
        }
      });

    BaseTestOperator.testContext.setRegistry(registry);

    DictionaryLookupPOP lookup = new DictionaryLookupPOP(null, PROPS, null, dictionaryFieldInfoMap);
    Table input = t(
      th("c0", "c1", "c2"),
      tr(0, 1, 2),
      tr(1, 2, 0),
      tr(2, 0, 1)
    );

    Table output = t(
      th("c0", "c1", "c2"),
      tr(10L, "bcd".getBytes(UTF8), 300.3),
      tr(20L, "cde".getBytes(UTF8), 100.1),
      tr(30L, "abc".getBytes(UTF8), 200.2)
    );

    validateSingle(lookup, DictionaryLookupOperator.class, input, output);
  }
}
 

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

  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.BIGINT.getType(),
                                      null, customFieldMeta);

  BigIntVector vector = new BigIntVector("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 BigIntToFixedConverter(vector, 0, this);

  for (int i = 0; i < rowCount; i++)
  {
    long longVal = converter.toLong(i);
    Object longObject = converter.toObject(i);
    String longString = converter.toString(i);

    if (nullValIndex.contains(i))
    {
      assertThat(longVal, is(0L));
      assertThat(longObject, is(nullValue()));
      assertThat(longString, is(nullValue()));
      assertThat(converter.toBytes(i), is(nullValue()));
    }
    else
    {
      assertThat(longVal, is(expectedValues.get(i)));
      assertThat(longObject, is(expectedValues.get(i)));
      assertThat(longString, is(expectedValues.get(i).toString()));
      bb = ByteBuffer.wrap(converter.toBytes(i));
      assertThat(longVal, is(bb.getLong()));
    }
  }
  vector.clear();
}
 

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

  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.BIGINT.getType(),
                                      null, customFieldMeta);

  BigIntVector vector = new BigIntVector("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 BigIntToScaledFixedConverter(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.BIGINT.getType(),
                                      null, customFieldMeta);

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

  final ArrowVectorConverter converterFoo =
      new BigIntToFixedConverter(vectorFoo, 0, this);

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

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

  final ArrowVectorConverter converterBar =
      new BigIntToFixedConverter(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.toInt(0), is(10));
  assertThat(converterBar.toInt(1), is(-10));
  vectorBar.clear();
}
 

@Test
public void testTime() throws SFException
{
  // test old and new dates
  long[] testTimesInt64 = {
      12345678000000L
  };

  String[] testTimesJson = {
      "12345.678000000"
  };

  Time[] expectedTimes = {
      new Time(12345678)
  };

  Map<String, String> customFieldMeta = new HashMap<>();
  customFieldMeta.put("logicalType", "TIME");
  Set<Integer> nullValIndex = new HashSet<>();
  // test normal date
  FieldType fieldType = new FieldType(true,
                                      Types.MinorType.BIGINT.getType(),
                                      null, customFieldMeta);

  BigIntVector vector = new BigIntVector("date", fieldType, allocator);
  int i = 0, j = 0;
  while (i < testTimesInt64.length)
  {
    boolean isNull = random.nextBoolean();
    if (isNull)
    {
      vector.setNull(j);
      nullValIndex.add(j);
    }
    else
    {
      vector.setSafe(j, testTimesInt64[i++]);
    }
    j++;
  }

  ArrowVectorConverter converter = new BigIntToTimeConverter(vector, 0, this);
  int rowCount = j;
  i = 0;
  j = 0;
  while (j < rowCount)
  {
    String strVal = converter.toString(j);
    Time time = converter.toTime(j);
    Object obj = converter.toObject(j);
    Time oldTime =
        new Time(ResultUtil.getSFTime(
            testTimesJson[i], scale, new SFSession()).getFractionalSeconds(
            ResultUtil.DEFAULT_SCALE_OF_SFTIME_FRACTION_SECONDS));
    if (nullValIndex.contains(j))
    {
      assertThat(obj, is(nullValue()));
      assertThat(strVal, is(nullValue()));
      assertThat(false, is(converter.toBoolean(j)));
      assertThat(converter.toBytes(j), is(nullValue()));
    }
    else
    {
      assertThat(expectedTimes[i], is(time));
      assertThat(expectedTimes[i], is((Time) obj));
      assertThat(oldTime, is(time));
      assertThat(oldTime, is((Time) obj));
      final int x = j;
      TestUtil.assertSFException(invalidConversionErrorCode,
                                 () -> converter.toBoolean(x));
    }
    j++;
  }
  vector.clear();
}