org.apache.arrow.vector.SmallIntVector Java Examples

The following examples show how to use org.apache.arrow.vector.SmallIntVector. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.

Example #1

Source File: TestData.java From dremio-oss with Apache License 2.0

6 votes

private Pair<SmallIntVector, ResultVerifier> testSmallIntVector() {
  final String colName = "colSmallInt";
  final List<Integer> values =  asList(20, null, -2000, 32700, 0);

  SmallIntVector valuesVector = new SmallIntVector(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);
}

Example #2

Source File: TestArrowShortConnector.java From multiple-dimension-spread with Apache License 2.0

5 votes

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

Example #3

Source File: SmallIntToFixedConverterTest.java From snowflake-jdbc with Apache License 2.0

5 votes

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

  SmallIntVector vector = new SmallIntVector("col_one", fieldType, allocator);
  vector.setSafe(0, 200);

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

Example #4

Source File: SmallIntToFixedConverterTest.java From snowflake-jdbc with Apache License 2.0

5 votes

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

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

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

Example #5

Source File: SmallIntToFixedConverterTest.java From snowflake-jdbc with Apache License 2.0

5 votes

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

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

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

Example #6

Source File: SmallIntToFixedConverter.java From snowflake-jdbc with Apache License 2.0

5 votes

public SmallIntToFixedConverter(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.smallIntVector = (SmallIntVector) fieldVector;
}

Example #7

Source File: TestArrowShortConnector.java From yosegi with Apache License 2.0

5 votes

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

Example #8

Source File: RowSmallIntWriter.java From flink with Apache License 2.0

5 votes

@Override
public void doWrite(Row value, int ordinal) {
	if (value.getField(ordinal) == null) {
		((SmallIntVector) getValueVector()).setNull(getCount());
	} else {
		((SmallIntVector) getValueVector()).setSafe(getCount(), (short) value.getField(ordinal));
	}
}

Example #9

Source File: SmallIntWriter.java From flink with Apache License 2.0

5 votes

@Override
public void doWrite(T in, int ordinal) {
	if (isNullAt(in, ordinal)) {
		((SmallIntVector) getValueVector()).setNull(getCount());
	} else {
		((SmallIntVector) getValueVector()).setSafe(getCount(), readShort(in, ordinal));
	}
}

Example #10

Source File: ArrowUtils.java From flink with Apache License 2.0

4 votes

private static ArrowFieldWriter<Row> createRowArrowFieldWriter(ValueVector vector, LogicalType fieldType) {
	if (vector instanceof TinyIntVector) {
		return new RowTinyIntWriter((TinyIntVector) vector);
	} else if (vector instanceof SmallIntVector) {
		return new RowSmallIntWriter((SmallIntVector) vector);
	} else if (vector instanceof IntVector) {
		return new RowIntWriter((IntVector) vector);
	} else if (vector instanceof BigIntVector) {
		return new RowBigIntWriter((BigIntVector) vector);
	} else if (vector instanceof BitVector) {
		return new RowBooleanWriter((BitVector) vector);
	} else if (vector instanceof Float4Vector) {
		return new RowFloatWriter((Float4Vector) vector);
	} else if (vector instanceof Float8Vector) {
		return new RowDoubleWriter((Float8Vector) vector);
	} else if (vector instanceof VarCharVector) {
		return new RowVarCharWriter((VarCharVector) vector);
	} else if (vector instanceof VarBinaryVector) {
		return new RowVarBinaryWriter((VarBinaryVector) vector);
	} else if (vector instanceof DecimalVector) {
		DecimalVector decimalVector = (DecimalVector) vector;
		return new RowDecimalWriter(decimalVector, getPrecision(decimalVector), decimalVector.getScale());
	} else if (vector instanceof DateDayVector) {
		return new RowDateWriter((DateDayVector) vector);
	} else if (vector instanceof TimeSecVector || vector instanceof TimeMilliVector ||
		vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) {
		return new RowTimeWriter(vector);
	} else if (vector instanceof TimeStampVector && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) {
		return new RowTimestampWriter(vector);
	} else if (vector instanceof ListVector) {
		ListVector listVector = (ListVector) vector;
		LogicalType elementType = ((ArrayType) fieldType).getElementType();
		return new RowArrayWriter(listVector, createRowArrowFieldWriter(listVector.getDataVector(), elementType));
	} else if (vector instanceof StructVector) {
		RowType rowType = (RowType) fieldType;
		ArrowFieldWriter<Row>[] fieldsWriters = new ArrowFieldWriter[rowType.getFieldCount()];
		for (int i = 0; i < fieldsWriters.length; i++) {
			fieldsWriters[i] = createRowArrowFieldWriter(
				((StructVector) vector).getVectorById(i),
				rowType.getTypeAt(i));
		}
		return new RowRowWriter((StructVector) vector, fieldsWriters);
	} else {
		throw new UnsupportedOperationException(String.format(
			"Unsupported type %s.", fieldType));
	}
}

Example #11

Source File: RowSmallIntWriter.java From flink with Apache License 2.0

4 votes

public RowSmallIntWriter(SmallIntVector smallIntVector) {
	super(smallIntVector);
}

Example #12

Source File: SmallIntWriter.java From flink with Apache License 2.0

4 votes

private SmallIntWriterForArray(SmallIntVector intVector) {
	super(intVector);
}

Example #13

Source File: ArrowSmallIntColumnVector.java From flink with Apache License 2.0

4 votes

public ArrowSmallIntColumnVector(SmallIntVector smallIntVector) {
	this.smallIntVector = Preconditions.checkNotNull(smallIntVector);
}

Example #14

Source File: SmallIntWriter.java From flink with Apache License 2.0

4 votes

private SmallIntWriterForRow(SmallIntVector intVector) {
	super(intVector);
}

Example #15

Source File: SmallIntFieldReader.java From flink with Apache License 2.0

4 votes

public SmallIntFieldReader(SmallIntVector smallIntVector) {
	super(smallIntVector);
}

Example #16

Source File: SmallIntFieldReader.java From flink with Apache License 2.0

4 votes

@Override
public Short read(int index) {
	return ((SmallIntVector) getValueVector()).getObject(index);
}

Example #17

Source File: SmallIntWriter.java From flink with Apache License 2.0

4 votes

private SmallIntWriter(SmallIntVector intVector) {
	super(intVector);
}

Example #18

Source File: ArrowUtils.java From flink with Apache License 2.0

4 votes

private static ArrowFieldWriter<RowData> createArrowFieldWriterForRow(ValueVector vector, LogicalType fieldType) {
	if (vector instanceof TinyIntVector) {
		return TinyIntWriter.forRow((TinyIntVector) vector);
	} else if (vector instanceof SmallIntVector) {
		return SmallIntWriter.forRow((SmallIntVector) vector);
	} else if (vector instanceof IntVector) {
		return IntWriter.forRow((IntVector) vector);
	} else if (vector instanceof BigIntVector) {
		return BigIntWriter.forRow((BigIntVector) vector);
	} else if (vector instanceof BitVector) {
		return BooleanWriter.forRow((BitVector) vector);
	} else if (vector instanceof Float4Vector) {
		return FloatWriter.forRow((Float4Vector) vector);
	} else if (vector instanceof Float8Vector) {
		return DoubleWriter.forRow((Float8Vector) vector);
	} else if (vector instanceof VarCharVector) {
		return VarCharWriter.forRow((VarCharVector) vector);
	} else if (vector instanceof VarBinaryVector) {
		return VarBinaryWriter.forRow((VarBinaryVector) vector);
	} else if (vector instanceof DecimalVector) {
		DecimalVector decimalVector = (DecimalVector) vector;
		return DecimalWriter.forRow(decimalVector, getPrecision(decimalVector), decimalVector.getScale());
	} else if (vector instanceof DateDayVector) {
		return DateWriter.forRow((DateDayVector) vector);
	} else if (vector instanceof TimeSecVector || vector instanceof TimeMilliVector ||
		vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) {
		return TimeWriter.forRow(vector);
	} else if (vector instanceof TimeStampVector && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) {
		int precision;
		if (fieldType instanceof LocalZonedTimestampType) {
			precision = ((LocalZonedTimestampType) fieldType).getPrecision();
		} else {
			precision = ((TimestampType) fieldType).getPrecision();
		}
		return TimestampWriter.forRow(vector, precision);
	} else if (vector instanceof ListVector) {
		ListVector listVector = (ListVector) vector;
		LogicalType elementType = ((ArrayType) fieldType).getElementType();
		return ArrayWriter.forRow(listVector, createArrowFieldWriterForArray(listVector.getDataVector(), elementType));
	} else if (vector instanceof StructVector) {
		RowType rowType = (RowType) fieldType;
		ArrowFieldWriter<RowData>[] fieldsWriters = new ArrowFieldWriter[rowType.getFieldCount()];
		for (int i = 0; i < fieldsWriters.length; i++) {
			fieldsWriters[i] = createArrowFieldWriterForRow(
				((StructVector) vector).getVectorById(i),
				rowType.getTypeAt(i));
		}
		return RowWriter.forRow((StructVector) vector, fieldsWriters);
	} else {
		throw new UnsupportedOperationException(String.format(
			"Unsupported type %s.", fieldType));
	}
}

Example #19

Source File: ArrowUtils.java From flink with Apache License 2.0

4 votes

private static ArrowFieldWriter<ArrayData> createArrowFieldWriterForArray(ValueVector vector, LogicalType fieldType) {
	if (vector instanceof TinyIntVector) {
		return TinyIntWriter.forArray((TinyIntVector) vector);
	} else if (vector instanceof SmallIntVector) {
		return SmallIntWriter.forArray((SmallIntVector) vector);
	} else if (vector instanceof IntVector) {
		return IntWriter.forArray((IntVector) vector);
	} else if (vector instanceof BigIntVector) {
		return BigIntWriter.forArray((BigIntVector) vector);
	} else if (vector instanceof BitVector) {
		return BooleanWriter.forArray((BitVector) vector);
	} else if (vector instanceof Float4Vector) {
		return FloatWriter.forArray((Float4Vector) vector);
	} else if (vector instanceof Float8Vector) {
		return DoubleWriter.forArray((Float8Vector) vector);
	} else if (vector instanceof VarCharVector) {
		return VarCharWriter.forArray((VarCharVector) vector);
	} else if (vector instanceof VarBinaryVector) {
		return VarBinaryWriter.forArray((VarBinaryVector) vector);
	} else if (vector instanceof DecimalVector) {
		DecimalVector decimalVector = (DecimalVector) vector;
		return DecimalWriter.forArray(decimalVector, getPrecision(decimalVector), decimalVector.getScale());
	} else if (vector instanceof DateDayVector) {
		return DateWriter.forArray((DateDayVector) vector);
	} else if (vector instanceof TimeSecVector || vector instanceof TimeMilliVector ||
		vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) {
		return TimeWriter.forArray(vector);
	} else if (vector instanceof TimeStampVector && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) {
		int precision;
		if (fieldType instanceof LocalZonedTimestampType) {
			precision = ((LocalZonedTimestampType) fieldType).getPrecision();
		} else {
			precision = ((TimestampType) fieldType).getPrecision();
		}
		return TimestampWriter.forArray(vector, precision);
	} else if (vector instanceof ListVector) {
		ListVector listVector = (ListVector) vector;
		LogicalType elementType = ((ArrayType) fieldType).getElementType();
		return ArrayWriter.forArray(listVector, createArrowFieldWriterForArray(listVector.getDataVector(), elementType));
	} else if (vector instanceof StructVector) {
		RowType rowType = (RowType) fieldType;
		ArrowFieldWriter<RowData>[] fieldsWriters = new ArrowFieldWriter[rowType.getFieldCount()];
		for (int i = 0; i < fieldsWriters.length; i++) {
			fieldsWriters[i] = createArrowFieldWriterForRow(
				((StructVector) vector).getVectorById(i),
				rowType.getTypeAt(i));
		}
		return RowWriter.forArray((StructVector) vector, fieldsWriters);
	} else {
		throw new UnsupportedOperationException(String.format(
			"Unsupported type %s.", fieldType));
	}
}

Example #20

Source File: SmallIntToFixedConverterTest.java From snowflake-jdbc with Apache License 2.0

4 votes

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

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

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

  for (int i = 0; i < rowCount; i++)
  {
    short shortVal = converter.toShort(i);
    Object longObject = converter.toObject(i); // the logical type is long
    String shortString = converter.toString(i);

    if (nullValIndex.contains(i))
    {
      assertThat(shortVal, is((short) 0));
      assertThat(longObject, is(nullValue()));
      assertThat(shortString, is(nullValue()));
      assertThat(converter.toBytes(i), is(nullValue()));
    }
    else
    {
      assertThat(shortVal, is(expectedValues.get(i)));
      assertEquals(longObject, (long) expectedValues.get(i));
      assertThat(shortString, is(expectedValues.get(i).toString()));
      bb = ByteBuffer.wrap(converter.toBytes(i));
      assertThat(shortVal, is(bb.getShort()));
    }
  }
  vector.clear();
}

Example #21

Source File: ArrowUtils.java From flink with Apache License 2.0

4 votes

public static ArrowFieldReader createRowArrowFieldReader(ValueVector vector, LogicalType fieldType) {
	if (vector instanceof TinyIntVector) {
		return new TinyIntFieldReader((TinyIntVector) vector);
	} else if (vector instanceof SmallIntVector) {
		return new SmallIntFieldReader((SmallIntVector) vector);
	} else if (vector instanceof IntVector) {
		return new IntFieldReader((IntVector) vector);
	} else if (vector instanceof BigIntVector) {
		return new BigIntFieldReader((BigIntVector) vector);
	} else if (vector instanceof BitVector) {
		return new BooleanFieldReader((BitVector) vector);
	} else if (vector instanceof Float4Vector) {
		return new FloatFieldReader((Float4Vector) vector);
	} else if (vector instanceof Float8Vector) {
		return new DoubleFieldReader((Float8Vector) vector);
	} else if (vector instanceof VarCharVector) {
		return new VarCharFieldReader((VarCharVector) vector);
	} else if (vector instanceof VarBinaryVector) {
		return new VarBinaryFieldReader((VarBinaryVector) vector);
	} else if (vector instanceof DecimalVector) {
		return new DecimalFieldReader((DecimalVector) vector);
	} else if (vector instanceof DateDayVector) {
		return new DateFieldReader((DateDayVector) vector);
	} else if (vector instanceof TimeSecVector || vector instanceof TimeMilliVector ||
		vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) {
		return new TimeFieldReader(vector);
	} else if (vector instanceof TimeStampVector && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) {
		return new TimestampFieldReader(vector);
	} else if (vector instanceof ListVector) {
		ListVector listVector = (ListVector) vector;
		LogicalType elementType = ((ArrayType) fieldType).getElementType();
		return new ArrayFieldReader(listVector,
			createRowArrowFieldReader(listVector.getDataVector(), elementType),
			elementType);
	} else if (vector instanceof StructVector) {
		StructVector structVector = (StructVector) vector;
		ArrowFieldReader[] fieldReaders = new ArrowFieldReader[structVector.size()];
		for (int i = 0; i < fieldReaders.length; i++) {
			fieldReaders[i] = createRowArrowFieldReader(structVector.getVectorById(i), ((RowType) fieldType).getTypeAt(i));
		}
		return new RowFieldReader(structVector, fieldReaders);
	} else {
		throw new UnsupportedOperationException(String.format(
			"Unsupported type %s.", fieldType));
	}
}

Example #22

Source File: ArrowUtils.java From flink with Apache License 2.0

4 votes

public static ColumnVector createColumnVector(ValueVector vector, LogicalType fieldType) {
	if (vector instanceof TinyIntVector) {
		return new ArrowTinyIntColumnVector((TinyIntVector) vector);
	} else if (vector instanceof SmallIntVector) {
		return new ArrowSmallIntColumnVector((SmallIntVector) vector);
	} else if (vector instanceof IntVector) {
		return new ArrowIntColumnVector((IntVector) vector);
	} else if (vector instanceof BigIntVector) {
		return new ArrowBigIntColumnVector((BigIntVector) vector);
	} else if (vector instanceof BitVector) {
		return new ArrowBooleanColumnVector((BitVector) vector);
	} else if (vector instanceof Float4Vector) {
		return new ArrowFloatColumnVector((Float4Vector) vector);
	} else if (vector instanceof Float8Vector) {
		return new ArrowDoubleColumnVector((Float8Vector) vector);
	} else if (vector instanceof VarCharVector) {
		return new ArrowVarCharColumnVector((VarCharVector) vector);
	} else if (vector instanceof VarBinaryVector) {
		return new ArrowVarBinaryColumnVector((VarBinaryVector) vector);
	} else if (vector instanceof DecimalVector) {
		return new ArrowDecimalColumnVector((DecimalVector) vector);
	} else if (vector instanceof DateDayVector) {
		return new ArrowDateColumnVector((DateDayVector) vector);
	} else if (vector instanceof TimeSecVector || vector instanceof TimeMilliVector ||
		vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) {
		return new ArrowTimeColumnVector(vector);
	} else if (vector instanceof TimeStampVector && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) {
		return new ArrowTimestampColumnVector(vector);
	} else if (vector instanceof ListVector) {
		ListVector listVector = (ListVector) vector;
		return new ArrowArrayColumnVector(listVector,
			createColumnVector(listVector.getDataVector(), ((ArrayType) fieldType).getElementType()));
	} else if (vector instanceof StructVector) {
		StructVector structVector = (StructVector) vector;
		ColumnVector[] fieldColumns = new ColumnVector[structVector.size()];
		for (int i = 0; i < fieldColumns.length; ++i) {
			fieldColumns[i] = createColumnVector(structVector.getVectorById(i), ((RowType) fieldType).getTypeAt(i));
		}
		return new ArrowRowColumnVector(structVector, fieldColumns);
	} else {
		throw new UnsupportedOperationException(String.format(
			"Unsupported type %s.", fieldType));
	}
}

Example #23

Source File: SmallIntWriter.java From flink with Apache License 2.0

4 votes

public static SmallIntWriter<RowData> forRow(SmallIntVector intVector) {
	return new SmallIntWriterForRow(intVector);
}

Example #24

Source File: SmallIntWriter.java From flink with Apache License 2.0

4 votes

public static SmallIntWriter<ArrayData> forArray(SmallIntVector intVector) {
	return new SmallIntWriterForArray(intVector);
}

Example #25

Source File: ArrowShortMemoryAllocator.java From multiple-dimension-spread with Apache License 2.0

4 votes

public ArrowShortMemoryAllocator( final SmallIntVector vector , final int rowCount ){
  vector.allocateNew( rowCount );
  this.vector = vector;
}

Example #26

Source File: BlockUtils.java From aws-athena-query-federation with Apache License 2.0

4 votes

/**
 * Used to mark a particular cell as null.
 *
 * @param vector The FieldVector to write the null value to.
 * @param pos The position (row) in the FieldVector to mark as null.
 */
private static void setNullValue(FieldVector vector, int pos)
{
    switch (vector.getMinorType()) {
        case TIMESTAMPMILLITZ:
            ((TimeStampMilliTZVector) vector).setNull(pos);
            break;
        case DATEMILLI:
            ((DateMilliVector) vector).setNull(pos);
            break;
        case DATEDAY:
            ((DateDayVector) vector).setNull(pos);
            break;
        case FLOAT8:
            ((Float8Vector) vector).setNull(pos);
            break;
        case FLOAT4:
            ((Float4Vector) vector).setNull(pos);
            break;
        case INT:
            ((IntVector) vector).setNull(pos);
            break;
        case TINYINT:
            ((TinyIntVector) vector).setNull(pos);
            break;
        case SMALLINT:
            ((SmallIntVector) vector).setNull(pos);
            break;
        case UINT1:
            ((UInt1Vector) vector).setNull(pos);
            break;
        case UINT2:
            ((UInt2Vector) vector).setNull(pos);
            break;
        case UINT4:
            ((UInt4Vector) vector).setNull(pos);
            break;
        case UINT8:
            ((UInt8Vector) vector).setNull(pos);
            break;
        case BIGINT:
            ((BigIntVector) vector).setNull(pos);
            break;
        case VARBINARY:
            ((VarBinaryVector) vector).setNull(pos);
            break;
        case DECIMAL:
            ((DecimalVector) vector).setNull(pos);
            break;
        case VARCHAR:
            ((VarCharVector) vector).setNull(pos);
            break;
        case BIT:
            ((BitVector) vector).setNull(pos);
            break;
        default:
            throw new IllegalArgumentException("Unknown type " + vector.getMinorType());
    }
}

Example #27

Source File: BlockUtils.java From aws-athena-query-federation with Apache License 2.0

4 votes

/**
 * In some filtering situations it can be useful to 'unset' a row as an indication to a later processing stage
 * that the row is irrelevant. The mechanism by which we 'unset' a row is actually field type specific and as such
 * this method is not supported for all field types.
 *
 * @param row The row number to unset in the provided Block.
 * @param block The Block where we'd like to unset the specified row.
 */
public static void unsetRow(int row, Block block)
{
    for (FieldVector vector : block.getFieldVectors()) {
        switch (vector.getMinorType()) {
            case TIMESTAMPMILLITZ:
                ((TimeStampMilliTZVector) vector).setNull(row);
                break;
            case DATEDAY:
                ((DateDayVector) vector).setNull(row);
                break;
            case DATEMILLI:
                ((DateMilliVector) vector).setNull(row);
                break;
            case TINYINT:
                ((TinyIntVector) vector).setNull(row);
                break;
            case UINT1:
                ((UInt1Vector) vector).setNull(row);
                break;
            case SMALLINT:
                ((SmallIntVector) vector).setNull(row);
                break;
            case UINT2:
                ((UInt2Vector) vector).setNull(row);
                break;
            case UINT4:
                ((UInt4Vector) vector).setNull(row);
                break;
            case INT:
                ((IntVector) vector).setNull(row);
                break;
            case UINT8:
                ((UInt8Vector) vector).setNull(row);
                break;
            case BIGINT:
                ((BigIntVector) vector).setNull(row);
                break;
            case FLOAT4:
                ((Float4Vector) vector).setNull(row);
                break;
            case FLOAT8:
                ((Float8Vector) vector).setNull(row);
                break;
            case DECIMAL:
                ((DecimalVector) vector).setNull(row);
                break;
            case VARBINARY:
                ((VarBinaryVector) vector).setNull(row);
                break;
            case VARCHAR:
                ((VarCharVector) vector).setNull(row);
                break;
            case BIT:
                ((BitVector) vector).setNull(row);
                break;
            case STRUCT:
                ((StructVector) vector).setNull(row);
                break;
            case LIST:
                UnionListWriter writer = ((ListVector) vector).getWriter();
                writer.setPosition(row);
                writer.startList();
                writer.endList();
                writer.setValueCount(0);
                break;
            default:
                throw new IllegalArgumentException("Unknown type " + vector.getMinorType());
        }
    }
}

Example #28

Source File: ArrowShortConnector.java From yosegi with Apache License 2.0

4 votes

public ArrowShortConnector( final String columnName , final SmallIntVector vector ) {
  this.columnName = columnName;
  this.vector = vector;
}

Example #29

Source File: ArrowMemoryAllocatorFactory.java From yosegi with Apache License 2.0

4 votes

/**
 * Set the vector of Struct and initialize it.
 */
public static IMemoryAllocator getFromStructVector(
    final ColumnType columnType ,
    final String columnName ,
    final BufferAllocator allocator ,
    final StructVector vector ,
    final int rowCount ) {
  switch ( columnType ) {
    case UNION:
      UnionVector unionVector = vector.addOrGetUnion( columnName );
      return new ArrowUnionMemoryAllocator( allocator , unionVector , rowCount );
    case ARRAY:
      return new ArrowArrayMemoryAllocator(
          allocator , vector.addOrGetList( columnName ) , rowCount );
    case SPREAD:
      StructVector mapVector = vector.addOrGetStruct( columnName );
      return new ArrowMapMemoryAllocator( allocator , mapVector , rowCount );

    case BOOLEAN:
      BitVector bitVector =  vector.addOrGet(
          columnName ,
          new FieldType( true , ArrowType.Bool.INSTANCE , null , null ) ,
          BitVector.class );
      return new ArrowBooleanMemoryAllocator( bitVector , rowCount );
    case BYTE:
      TinyIntVector byteVector =  vector.addOrGet(
          columnName ,
          new FieldType( true , new ArrowType.Int( 8 , true ) , null , null ) ,
          TinyIntVector.class );
      return new ArrowByteMemoryAllocator( byteVector , rowCount );
    case SHORT:
      SmallIntVector shortVector = vector.addOrGet(
          columnName ,
          new FieldType(
            true ,
            new ArrowType.Int( 16 , true ) ,
            null ,
            null ) ,
          SmallIntVector.class );
      return new ArrowShortMemoryAllocator( shortVector , rowCount );
    case INTEGER:
      IntVector integerVector =  vector.addOrGet(
          columnName ,
          new FieldType( true , new ArrowType.Int( 32 , true ) , null , null ) ,
          IntVector.class );
      return new ArrowIntegerMemoryAllocator( integerVector , rowCount );
    case LONG:
      BigIntVector longVector = vector.addOrGet(
          columnName ,
          new FieldType(
            true ,
            new ArrowType.Int( 64 , true ) ,
            null ,
            null ) ,
            BigIntVector.class );
      return new ArrowLongMemoryAllocator( longVector , rowCount );
    case FLOAT:
      Float4Vector floatVector = vector.addOrGet(
          columnName ,
          new FieldType(
            true ,
            new ArrowType.FloatingPoint( FloatingPointPrecision.SINGLE ) ,
            null ,
            null ) ,
          Float4Vector.class );
      return new ArrowFloatMemoryAllocator( floatVector , rowCount );
    case DOUBLE:
      Float8Vector doubleVector = vector.addOrGet(
          columnName ,
          new FieldType(
            true ,
            new ArrowType.FloatingPoint( FloatingPointPrecision.DOUBLE ) ,
            null ,
            null ) ,
          Float8Vector.class );
      return new ArrowDoubleMemoryAllocator( doubleVector , rowCount );
    case STRING:
      VarCharVector charVector = vector.addOrGet(
          columnName ,
          new FieldType( true , ArrowType.Utf8.INSTANCE , null , null ) ,
          VarCharVector.class );
      return new ArrowStringMemoryAllocator( charVector , rowCount );
    case BYTES:
      VarBinaryVector binaryVector = vector.addOrGet(
          columnName ,
          new FieldType( true , ArrowType.Binary.INSTANCE , null , null ) ,
          VarBinaryVector.class );
      return new ArrowBytesMemoryAllocator( binaryVector , rowCount );

    case NULL:
    case EMPTY_ARRAY:
    case EMPTY_SPREAD:
    default:
      return NullMemoryAllocator.INSTANCE;
  }
}

Example #30

Source File: ArrowMemoryAllocatorFactory.java From yosegi with Apache License 2.0

4 votes

/**
 * Set the vector of List and initialize it.
 */
public static IMemoryAllocator getFromListVector(
      final ColumnType columnType ,
      final String columnName ,
      final BufferAllocator allocator ,
      final ListVector vector ,
      final int rowCount ) {
  switch ( columnType ) {
    case UNION:
      AddOrGetResult<UnionVector> unionVector =  vector.addOrGetVector(
          new FieldType( true , MinorType.UNION.getType() , null , null ) );
      return new ArrowUnionMemoryAllocator( allocator , unionVector.getVector() , rowCount );
    case ARRAY:
      AddOrGetResult<ListVector> listVector =  vector.addOrGetVector(
          new FieldType( true , ArrowType.List.INSTANCE , null , null ) );
      return new ArrowArrayMemoryAllocator( allocator , listVector.getVector() , rowCount );
    case SPREAD:
      AddOrGetResult<StructVector> mapVector = vector.addOrGetVector(
          new FieldType( true , ArrowType.Struct.INSTANCE , null , null ) );
      return new ArrowMapMemoryAllocator( allocator , mapVector.getVector() , rowCount );

    case BOOLEAN:
      AddOrGetResult<BitVector> bitVector = vector.addOrGetVector(
          new FieldType( true , ArrowType.Bool.INSTANCE , null , null ) );
      return new ArrowBooleanMemoryAllocator( bitVector.getVector() , rowCount );
    case BYTE:
      AddOrGetResult<TinyIntVector> byteVector = vector.addOrGetVector(
          new FieldType( true , new ArrowType.Int( 8 , true ) , null , null ) );
      return new ArrowByteMemoryAllocator( byteVector.getVector() , rowCount );
    case SHORT:
      AddOrGetResult<SmallIntVector> shortVector = vector.addOrGetVector(
          new FieldType( true , new ArrowType.Int( 16 , true ) , null , null ) );
      return new ArrowShortMemoryAllocator( shortVector.getVector() , rowCount );
    case INTEGER:
      AddOrGetResult<IntVector> integerVector =  vector.addOrGetVector(
          new FieldType( true , new ArrowType.Int( 32 , true ) , null , null ) );
      return new ArrowIntegerMemoryAllocator( integerVector.getVector() , rowCount );
    case LONG:
      AddOrGetResult<BigIntVector> longVector =  vector.addOrGetVector(
          new FieldType( true , new ArrowType.Int( 64 , true ) , null , null ) );
      return new ArrowLongMemoryAllocator( longVector.getVector() , rowCount );
    case FLOAT:
      AddOrGetResult<Float4Vector> floatVector = vector.addOrGetVector(
          new FieldType(
            true ,
            new ArrowType.FloatingPoint( FloatingPointPrecision.HALF ) ,
            null ,
            null ) );
      return new ArrowFloatMemoryAllocator( floatVector.getVector() , rowCount );
    case DOUBLE:
      AddOrGetResult<Float8Vector> doubleVector = vector.addOrGetVector(
          new FieldType(
            true ,
            new ArrowType.FloatingPoint( FloatingPointPrecision.DOUBLE ) ,
            null ,
            null ) );
      return new ArrowDoubleMemoryAllocator( doubleVector.getVector() , rowCount );
    case STRING:
      AddOrGetResult<VarCharVector> charVector = vector.addOrGetVector(
          new FieldType( true , ArrowType.Utf8.INSTANCE , null , null ) );
      return new ArrowStringMemoryAllocator( charVector.getVector() , rowCount );
    case BYTES:
      AddOrGetResult<VarBinaryVector> binaryVector = vector.addOrGetVector(
          new FieldType( true , ArrowType.Binary.INSTANCE , null , null ) );
      return new ArrowBytesMemoryAllocator( binaryVector.getVector() , rowCount );

    case NULL:
    case EMPTY_ARRAY:
    case EMPTY_SPREAD:
    default:
      return NullMemoryAllocator.INSTANCE;
  }
}