org.apache.arrow.vector.types.pojo.ArrowType Java Examples

@Test
public void toArrowType()
{
    Assert.assertEquals(Types.MinorType.BIT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BIT, 0, 0));
    Assert.assertEquals(Types.MinorType.BIT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BOOLEAN, 0, 0));
    Assert.assertEquals(Types.MinorType.TINYINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TINYINT, 0, 0));
    Assert.assertEquals(Types.MinorType.SMALLINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.SMALLINT, 0, 0));
    Assert.assertEquals(Types.MinorType.INT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.INTEGER, 0, 0));
    Assert.assertEquals(Types.MinorType.BIGINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BIGINT, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT4.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.REAL, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT4.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.FLOAT, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT8.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DOUBLE, 0, 0));
    Assert.assertEquals(new ArrowType.Decimal(5, 3), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DECIMAL, 5, 3));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.CHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.NCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.VARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.NVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGNVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.VARBINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGVARBINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.DATEMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DATE, 0, 0));
    Assert.assertEquals(Types.MinorType.TIMEMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TIME, 0, 0));
    Assert.assertEquals(Types.MinorType.TIMESTAMPMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TIMESTAMP, 0, 0));
}
 

@Test
public void testSupportedTypesHaveSerializers()
{
    ArrowTypeSerDe.Serializer arrowTypeSerializer = new ArrowTypeSerDe.Serializer();
    Map<String, TypedSerializer<ArrowType>> delegateSerDeMap = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
    delegateSerDeMap.putAll(arrowTypeSerializer.getDelegateSerDeMap());
    for (SupportedTypes supportedType : SupportedTypes.values()) {
        String arrowTypeClass;
        try {
            ArrowType arrowType = supportedType.getArrowMinorType().getType();
            arrowTypeClass = arrowType.getClass().getSimpleName();
        }
        catch (UnsupportedOperationException e) {
            // fall back to enum name
            arrowTypeClass = FALL_BACK_ARROW_TYPE_CLASS.getOrDefault(supportedType, supportedType.name());
        }
        assertTrue("No serializer for supported type " + supportedType + " with ArrowType " + arrowTypeClass, delegateSerDeMap.containsKey(arrowTypeClass));
    }
}
 

@Test
public void getUnSupportedFunctions() throws GandivaException {
  FunctionImplementationRegistry fnRegistry = FUNCTIONS();
  ArrayListMultimap<String, AbstractFunctionHolder> functions = fnRegistry.getRegisteredFunctions();
  Set<String> fns = Sets.newHashSet();
  Set<FunctionSignature> supportedFunctions = ExpressionRegistry.getInstance()
    .getSupportedFunctions();
  for (FunctionSignature signature : supportedFunctions ) {
    String fnName = (signature.getName().toLowerCase() +"##");
    for (ArrowType param : signature.getParamTypes()) {
      fnName = fnName + "##" + param.toString();
    }
    fns.add(fnName);
  }
  for (Map.Entry<String, AbstractFunctionHolder> holders : functions.entries()) {
    String name = holders.getKey();
    AbstractFunctionHolder holder = holders.getValue();
    totalFuncs++;
    isFunctionSupported(name, (BaseFunctionHolder)holder, fns);

  }
  System.out.println("Total : " + totalFuncs + " unSupported : " + unSupportedFn);
}
 

private static VectorContainer buildBinaryGlobalDictionary(List<Dictionary> dictionaries, VectorContainer existingDict, ColumnDescriptor columnDescriptor, BufferAllocator bufferAllocator) {
  final Field field = new Field(SchemaPath.getCompoundPath(columnDescriptor.getPath()).getAsUnescapedPath(), true, new ArrowType.Binary(), null);
  final VectorContainer input = new VectorContainer(bufferAllocator);
  final VarBinaryVector binaryVector = input.addOrGet(field);
  binaryVector.allocateNew();
  final SortedSet<Binary> values = new TreeSet<>();
  for (Dictionary dictionary : dictionaries) {
    for (int i = 0; i <= dictionary.getMaxId(); ++i) {
      values.add(dictionary.decodeToBinary(i));
    }
  }
  if (existingDict != null) {
    final VarBinaryVector existingDictValues = existingDict.getValueAccessorById(VarBinaryVector.class, 0).getValueVector();
    for (int i = 0; i < existingDict.getRecordCount(); ++i) {
      values.add(Binary.fromConstantByteArray(existingDictValues.get(i)));
    }
  }
  final Iterator<Binary> iter = values.iterator();
  int recordCount = 0;
  while (iter.hasNext()) {
    final byte[] data = iter.next().getBytes();
    binaryVector.setSafe(recordCount++, data, 0, data.length);
  }
  binaryVector.setValueCount(recordCount);
  input.setRecordCount(recordCount);
  input.buildSchema(BatchSchema.SelectionVectorMode.NONE);
  return input;
}
 

private Schema fromMetadata(List<ResultColumnMetadata> rcmd) {

      Schema schema = new Schema(rcmd.stream().map(md -> {
        ArrowType arrowType = SqlTypeNameToArrowType.toArrowType(md);
        FieldType fieldType = new FieldType(md.getIsNullable(), arrowType, null, null);
        return new Field(md.getColumnName(), fieldType, null);
      }).collect(Collectors.toList()));
      return schema;
    }
 

@Test
public void falseWhenChildrenAreSelected() {
  BatchSchema schema = mock(BatchSchema.class);
  when(schema.iterator())
    .thenReturn(Lists.newArrayList(
        new Field("a1", new FieldType(true, new ArrowType.Struct(), null),
            Lists.newArrayList(Field.nullable("a2", new ArrowType.Bool()))),
        Field.nullable("a3", new ArrowType.Bool())).iterator());
  assertFalse(EasyScanOperatorCreator.selectsAllColumns(schema,
    Lists.newArrayList(SchemaPath.getSimplePath("a1"), SchemaPath.getCompoundPath("a1", "a2"),
        SchemaPath.getSimplePath("a3"))));
}
 

@Before
public void setup()
{
    allocator = new BlockAllocatorImpl();

    Schema schema = SchemaBuilder.newBuilder()
            .addField("col1", new ArrowType.Int(32, true))
            .addField("col2", new ArrowType.Utf8())
            .build();

    spillConfig = SpillConfig.newBuilder().withEncryptionKey(keyFactory.create())
            .withRequestId(requestId)
            .withSpillLocation(S3SpillLocation.newBuilder()
                    .withBucket(bucket)
                    .withPrefix(prefix)
                    .withQueryId(requestId)
                    .withSplitId(splitId)
                    .withIsDirectory(true)
                    .build())
            .withRequestId(requestId)
            .build();

    blockWriter = new S3BlockSpiller(mockS3, spillConfig, allocator, schema, ConstraintEvaluator.emptyEvaluator());

    expected = allocator.createBlock(schema);
    BlockUtils.setValue(expected.getFieldVector("col1"), 1, 100);
    BlockUtils.setValue(expected.getFieldVector("col2"), 1, "VarChar");
    BlockUtils.setValue(expected.getFieldVector("col1"), 1, 101);
    BlockUtils.setValue(expected.getFieldVector("col2"), 1, "VarChar1");
    expected.setRowCount(2);
}
 

/**
 * Adds a new LIST child field with the given name to the builder.
 *
 * @param fieldName The name to use for the newly added child field.
 * @param type The concrete type for values in the List
 * @return This FieldBuilder itself.
 */
public FieldBuilder addListField(String fieldName, ArrowType type)
{
    Field baseField = new Field("", FieldType.nullable(type), null);
    Field field = new Field(fieldName,
            FieldType.nullable(Types.MinorType.LIST.getType()),
            Collections.singletonList(baseField));
    this.children.put(fieldName, field);
    return this;
}
 

@Test
public void testParquetFixedBinaryToArrowDecimal() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(FIXED_LEN_BYTE_ARRAY).length(5).as(DECIMAL).precision(8).scale(2).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Decimal(8, 2))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

@Test(expected = JsonMappingException.class)
public void testStrictSerializer()
        throws JsonProcessingException
{
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        ObjectMapper mapper = VersionedObjectMapperFactory.create(allocator);
        mapper.writeValueAsString(new ArrowType.Null());
    }
}
 

@Override
protected void doTypedSerialize(ArrowType arrowType, JsonGenerator jgen, SerializerProvider provider)
        throws IOException
{
    ArrowType.Decimal decimal = (ArrowType.Decimal) arrowType;
    jgen.writeNumberField(PRECISION_FIELD, decimal.getPrecision());
    jgen.writeNumberField(SCALE_FIELD, decimal.getScale());
}
 

public static IMemoryAllocator getFromListVector( final IField schema , final String columnName , final BufferAllocator allocator , final ListVector vector , final int rowCount ) throws IOException{
  switch( schema.getFieldType() ){
    case UNION:
      return NullMemoryAllocator.INSTANCE;
    case ARRAY:
      AddOrGetResult<ListVector> listVector =  vector.addOrGetVector( new FieldType( true , ArrowType.List.INSTANCE , null , null ) );
      return new ArrowFixedSchemaArrayMemoryAllocator( (ArrayContainerField)schema , allocator , listVector.getVector() , rowCount );
    case MAP:
      AddOrGetResult<StructVector> mapVector =  vector.addOrGetVector( new FieldType( true , ArrowType.Struct.INSTANCE , null , null ) );
      return new ArrowFixedSchemaMapMemoryAllocator( (MapContainerField)schema , allocator , mapVector.getVector() , rowCount );
    case STRUCT:
      AddOrGetResult<StructVector> structVector =  vector.addOrGetVector( new FieldType( true , ArrowType.Struct.INSTANCE , null , null ) );
      return new ArrowFixedSchemaStructMemoryAllocator( (StructContainerField)schema , allocator , structVector.getVector() , rowCount );

    case BOOLEAN:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.BOOLEAN , columnName ,  allocator , vector , rowCount );
    case BYTE:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.BYTE , columnName ,  allocator , vector , rowCount );
    case SHORT:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.SHORT , columnName ,  allocator , vector , rowCount );
    case INTEGER:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.INTEGER , columnName ,  allocator , vector , rowCount );
    case LONG:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.LONG , columnName ,  allocator , vector , rowCount );
    case FLOAT:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.FLOAT , columnName ,  allocator , vector , rowCount );
    case DOUBLE:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.DOUBLE , columnName ,  allocator , vector , rowCount );
    case STRING:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.STRING , columnName ,  allocator , vector , rowCount );
    case BYTES:
      return ArrowMemoryAllocatorFactory.getFromListVector( ColumnType.BYTES , columnName ,  allocator , vector , rowCount );

    default:
      return NullMemoryAllocator.INSTANCE;
  }
}
 

/**
 * Constructs a new AllOrNoneValueSet.
 *
 * @param type The Apache Arrow type of the field that this ValueSet applies to.
 * @param all True if all non-null values are in this ValueSet.
 * @param nullAllowed True is all null values are in this ValueSet.
 */
@JsonCreator
public AllOrNoneValueSet(@JsonProperty("type") ArrowType type,
        @JsonProperty("all") boolean all,
        @JsonProperty("nullAllowed") boolean nullAllowed)
{
    this.type = requireNonNull(type, "type is null");
    this.all = all;
    this.nullAllowed = nullAllowed;
}
 

@Override
public CompleteType getOutputType(CompleteType baseReturn, List<LogicalExpression> args) {
  return new CompleteType(
      ArrowType.List.INSTANCE,
      new CompleteType(
          ArrowType.Struct.INSTANCE,
          CompleteType.INT.toField(OFFSET_FIELD),
          CompleteType.INT.toField(LENGTH_FIELD)
      ).toField(ListVector.DATA_VECTOR_NAME)
  );
}
 

/**
 * Converts from DynamoDB Attribute Type to Arrow type.
 *
 * @param attributeName the DDB Attribute name
 * @param attributeType the DDB Attribute type
 * @return the converted-to Arrow Field
 */
public static Field getArrowFieldFromDDBType(String attributeName, String attributeType)
{
    switch (attributeType) {
        case STRING:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null);
        case NUMBER:
            return new Field(attributeName, FieldType.nullable(new ArrowType.Decimal(38, 9)), null);
        case BOOLEAN:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.BIT.getType()), null);
        case BINARY:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.VARBINARY.getType()), null);
        case STRING_SET:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.LIST.getType()),
                    Collections.singletonList(new Field("", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null)));
        case NUMBER_SET:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.LIST.getType()),
                    Collections.singletonList(new Field("", FieldType.nullable(new ArrowType.Decimal(38, 9)), null)));
        case BINARY_SET:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.LIST.getType()),
                    Collections.singletonList(new Field("", FieldType.nullable(Types.MinorType.VARBINARY.getType()), null)));
        case LIST:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.LIST.getType()), null);
        case MAP:
            return new Field(attributeName, FieldType.nullable(Types.MinorType.STRUCT.getType()), null);
        default:
            throw new RuntimeException("Unknown type[" + attributeType + "] for field[" + attributeName + "]");
    }
}
 

@Test
public void testParquetFixedBinaryToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(FIXED_LEN_BYTE_ARRAY).length(12).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Binary())
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

private FunctionCall getDecimalV2NamesIfEnabled(FunctionCall functionCall) {
  // If decimal v2 is on, switch to function definitions that return decimal vectors
  // for decimal accumulation.
  FunctionCall functionCallToResolve = functionCall;
  if (isDecimalV2Enabled) {
    if (aggrFunctionNames.contains(functionCall.getName().toLowerCase())) {
      LogicalExpression aggrColumn = functionCall.args.get(0);
      if (aggrColumn.getCompleteType().getType().getTypeID() == ArrowType.ArrowTypeID.Decimal) {
        functionCallToResolve = new FunctionCall(functionCall.getName() +
          "_v2", functionCall.args);
      }
    }
  }
  return functionCallToResolve;
}
 

@Override
protected void doTypedSerialize(ArrowType arrowType, JsonGenerator jgen, SerializerProvider provider)
        throws IOException
{
    ArrowType.FixedSizeBinary fixedSizeBinary = (ArrowType.FixedSizeBinary) arrowType;
    jgen.writeNumberField(BYTE_WIDTH_FIELD, fixedSizeBinary.getByteWidth());
}
 

@Before
public void beforeTest()
        throws IOException
{
    String yearCol = "year";
    String monthCol = "month";
    String dayCol = "day";

    Schema schema = SchemaBuilder.newBuilder()
            .addField(yearCol, new ArrowType.Int(32, true))
            .addField(monthCol, new ArrowType.Int(32, true))
            .addField(dayCol, new ArrowType.Int(32, true))
            .addField("col3", new ArrowType.Utf8())
            .build();

    Block partitions = allocator.createBlock(schema);
    int num_partitions = 10;
    for (int i = 0; i < num_partitions; i++) {
        BlockUtils.setValue(partitions.getFieldVector(yearCol), i, 2016 + i);
        BlockUtils.setValue(partitions.getFieldVector(monthCol), i, (i % 12) + 1);
        BlockUtils.setValue(partitions.getFieldVector(dayCol), i, (i % 28) + 1);
    }
    partitions.setRowCount(num_partitions);

    expected = new GetTableLayoutResponse(
            "test-catalog",
            new TableName("test-schema", "test-table"),
            partitions);

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

/** Helper method which creates a empty list vector */
private static ListVector testEmptyListVector(BufferAllocator allocator) {
  final ListVector vector =
      new ListVector("emptyListVector", allocator, FieldType.nullable(ArrowType.Null.INSTANCE), null);
  vector.allocateNew();
  return vector;
}
 

private SortedRangeSet(ArrowType type, NavigableMap<ValueMarker, Range> lowIndexedRanges, boolean nullAllowed)
{
    requireNonNull(type, "type is null");
    requireNonNull(lowIndexedRanges, "lowIndexedRanges is null");

    this.type = type;
    this.lowIndexedRanges = lowIndexedRanges;
    this.nullAllowed = nullAllowed;
}
 

@Test
public void testParquetInt64TimeMicrosToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(INT64).as(TIME_MICROS).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.MICROSECOND, 64))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

/**
 * This method takes in a table schema and a String representing the set of
 * simple contraints to be ANDed together and applied to that table.
 *
 * @param schema The schema of the table in question
 * @param input A comma-separated constraint String in the form of {field_name}{operator}{value}.
 *              The operators must be one of those available in {@link LogicalOperator}.
 *              Currently, we only support Boolean, Integer, Floating Point, Decimal, and String operands
 *              for this validator's constraints.
 * @return a {@link Constraints} object populated from the input string, un-constrained if input is not present
 */
public static Constraints parseConstraints(Schema schema, Optional<String> input)
{
  if (!input.isPresent() || input.get().trim().isEmpty()) {
    return new Constraints(Collections.EMPTY_MAP);
  }

  Map<String, ArrowType> fieldTypes = schema.getFields().stream()
                                              .collect(Collectors.toMap(Field::getName, Field::getType));

  Map<String, ValueSet> constraints = new HashMap<>();
  Iterable<String> constraintStrings = CONSTRAINT_SPLITTER.split(input.get());
  constraintStrings.forEach(str -> parseAndAddConstraint(fieldTypes, constraints, str));
  return new Constraints(constraints);
}
 

/**
 *
 * @param allocator
 * @param name
 * @param data
 * @return
 */
public static BigIntVector vectorFor(BufferAllocator allocator,String name,long[] data) {
    BigIntVector float8Vector = new BigIntVector(name,FieldType.nullable(new ArrowType.Int(64,true)),allocator);
    float8Vector.allocateNew(data.length);
    for(int i = 0; i < data.length; i++) {
        float8Vector.setSafe(i,data[i]);
    }

    float8Vector.setValueCount(data.length);

    return float8Vector;
}
 

@Override
protected ArrowType doTypedDeserialize(JsonParser jparser, DeserializationContext ctxt)
        throws IOException
{
    return new ArrowType.Struct();
}
 

/**
 * Provided Ranges are unioned together to form the SortedRangeSet
 */
static SortedRangeSet copyOf(ArrowType type, Iterable<Range> ranges, boolean nullAllowed)
{
    return new Builder(type, nullAllowed).addAll(ranges).build();
}
 

@Override
public Void visit(ArrowType.Null aNull) {
  return writeString("OTHER");
}
 

public static CompleteType getOutputType(OperationType operationType, List<LogicalExpression> args) {
  return new CompleteType(DecimalTypeUtil.getResultTypeForOperation(operationType,
    args.get(0).getCompleteType().getType(ArrowType.Decimal.class),
    args.get(1).getCompleteType().getType(ArrowType.Decimal.class)));
}
 

private Serializer()
{
    super(ArrowType.class, ArrowType.Interval.class);
}
 

private Deserializer()
{
    super(ArrowType.class, ArrowType.Time.class);
}