Java Code Examples for org.apache.arrow.vector.types.pojo.ArrowType#Int

public static ColumnMetaData metaDataFromField(Field field) {
    ArrowType arrowType = field.getFieldType().getType();
    if (arrowType instanceof ArrowType.Int) {
        ArrowType.Int intType = (ArrowType.Int) arrowType;
        return intType.getBitWidth() == 32 ? new IntegerMetaData(field.getName()) : new LongMetaData(field.getName());
    } else if (arrowType instanceof ArrowType.Bool) {
        return new BooleanMetaData(field.getName());
    } else if (arrowType instanceof ArrowType.FloatingPoint) {
        ArrowType.FloatingPoint floatingPointType = (ArrowType.FloatingPoint) arrowType;
        return floatingPointType.getPrecision() == FloatingPointPrecision.DOUBLE ? new DoubleMetaData(field.getName()) : new FloatMetaData(field.getName());
    } else if (arrowType instanceof ArrowType.Binary) {
        return new BinaryMetaData(field.getName());
    } else if (arrowType instanceof ArrowType.Utf8) {
        return new StringMetaData(field.getName());
    } else if (arrowType instanceof ArrowType.Date) {
        return new TimeMetaData(field.getName());
    } else {
        throw new IllegalStateException("Illegal type " + field.getFieldType().getType());
    }
}
 

/**
 * Convert arrow type to the arrow type supported by dremio
 * @param arrowType original arrow type
 * @return the arrow type supported by dremio
 */
private static ArrowType convertToSupportedArrowType(ArrowType arrowType) {
  switch (arrowType.getTypeID()) {
    case Int:
      ArrowType.Int arrowInt = (ArrowType.Int)arrowType;
      return (arrowInt.getBitWidth() < 32) ? CompleteType.INT.getType() : arrowType;
    case Date:
      // We don't support DateDay, so we should convert it to DataMilli.
      return CompleteType.DATE.getType();
    case Timestamp:
      // We always treat timezone as null.
      return CompleteType.TIMESTAMP.getType();
    default:
      return arrowType;
  }
}
 

@Override
protected void doTypedSerialize(ArrowType arrowType, JsonGenerator jgen, SerializerProvider provider)
        throws IOException
{
    ArrowType.Int arrowInt = (ArrowType.Int) arrowType;
    jgen.writeNumberField(BIT_WIDTH_FIELD, arrowInt.getBitWidth());
    jgen.writeBooleanField(IS_SIGNED_FIELD, arrowInt.getIsSigned());
}
 

private static ColumnMetaData metaDataFromField(Field field) {
    ArrowType arrowType = field.getFieldType().getType();
    if(arrowType instanceof ArrowType.Int) {
        val intType = (ArrowType.Int) arrowType;
        if(intType.getBitWidth() == 32)
            return new IntegerMetaData(field.getName());
        else {
            return new LongMetaData(field.getName());
        }
    }
    else if(arrowType instanceof ArrowType.Bool) {
        return new BooleanMetaData(field.getName());
    }
    else if(arrowType  instanceof ArrowType.FloatingPoint) {
        val floatingPointType = (ArrowType.FloatingPoint) arrowType;
        if(floatingPointType.getPrecision() == FloatingPointPrecision.DOUBLE)
            return new DoubleMetaData(field.getName());
        else {
            return new FloatMetaData(field.getName());
        }
    }
    else if(arrowType instanceof  ArrowType.Binary) {
        return new BinaryMetaData(field.getName());
    }
    else if(arrowType instanceof ArrowType.Utf8) {
        return new StringMetaData(field.getName());

    }
    else if(arrowType instanceof ArrowType.Date) {
        return new TimeMetaData(field.getName());
    }
    else {
        throw new IllegalStateException("Illegal type " + field.getFieldType().getType());
    }

}
 

private static ColumnMetaData metaDataFromField(Field field) {
    ArrowType arrowType = field.getFieldType().getType();
    if(arrowType instanceof ArrowType.Int) {
        val intType = (ArrowType.Int) arrowType;
        if(intType.getBitWidth() == 32)
            return new IntegerMetaData(field.getName());
        else {
            return new LongMetaData(field.getName());
        }
    }
    else if(arrowType instanceof ArrowType.Bool) {
        return new BooleanMetaData(field.getName());
    }
    else if(arrowType  instanceof ArrowType.FloatingPoint) {
        val floatingPointType = (ArrowType.FloatingPoint) arrowType;
        if(floatingPointType.getPrecision() == FloatingPointPrecision.DOUBLE)
            return new DoubleMetaData(field.getName());
        else {
            return new FloatMetaData(field.getName());
        }
    }
    else if(arrowType instanceof  ArrowType.Binary) {
        return new BinaryMetaData(field.getName());
    }
    else if(arrowType instanceof ArrowType.Utf8) {
        return new StringMetaData(field.getName());

    }
    else if(arrowType instanceof ArrowType.Date) {
        return new TimeMetaData(field.getName());
    }
    else {
        throw new IllegalStateException("Illegal type " + field.getFieldType().getType());
    }

}
 

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 static VectorContainer buildIntegerGlobalDictionary(List<Dictionary> dictionaries, VectorContainer existingDict, ColumnDescriptor columnDescriptor, BufferAllocator bufferAllocator) {
  final Field field = new Field(SchemaPath.getCompoundPath(columnDescriptor.getPath()).getAsUnescapedPath(), true, new ArrowType.Int(32, true), null);
  final VectorContainer input = new VectorContainer(bufferAllocator);
  final IntVector intVector = input.addOrGet(field);
  intVector.allocateNew();
  final SortedSet<Integer> values = Sets.newTreeSet();
  for (Dictionary dictionary : dictionaries) {
    for (int i = 0; i <= dictionary.getMaxId(); ++i) {
      values.add(dictionary.decodeToInt(i));
    }
  }
  if (existingDict != null) {
    final IntVector existingDictValues = existingDict.getValueAccessorById(IntVector.class, 0).getValueVector();
    for (int i = 0; i < existingDict.getRecordCount(); ++i) {
      values.add(existingDictValues.get(i));
    }
  }
  final Iterator<Integer> iter = values.iterator();
  int recordCount = 0;
  while (iter.hasNext()) {
    intVector.setSafe(recordCount++, iter.next());
  }
  intVector.setValueCount(recordCount);
  input.setRecordCount(recordCount);
  input.buildSchema(BatchSchema.SelectionVectorMode.NONE);
  return input;
}
 

@Override
protected ArrowType doTypedDeserialize(JsonParser jparser, DeserializationContext ctxt)
        throws IOException
{
    int bitWidth = getNextIntField(jparser, BIT_WIDTH_FIELD);
    boolean isSigned = getNextBoolField(jparser, IS_SIGNED_FIELD);
    return new ArrowType.Int(bitWidth, isSigned);
}
 

private DataType sparkFromArrow(FieldType fieldType) {
  switch (fieldType.getType().getTypeID()) {
    case Null:
      return DataTypes.NullType;
    case Struct:
      throw new UnsupportedOperationException("have not implemented Struct type yet");
    case List:
      throw new UnsupportedOperationException("have not implemented List type yet");
    case FixedSizeList:
      throw new UnsupportedOperationException("have not implemented FixedSizeList type yet");
    case Union:
      throw new UnsupportedOperationException("have not implemented Union type yet");
    case Int:
      ArrowType.Int intType = (ArrowType.Int) fieldType.getType();
      int bitWidth = intType.getBitWidth();
      if (bitWidth == 8) {
        return DataTypes.ByteType;
      } else if (bitWidth == 16) {
        return DataTypes.ShortType;
      } else if (bitWidth == 32) {
        return DataTypes.IntegerType;
      } else if (bitWidth == 64) {
        return DataTypes.LongType;
      }
      throw new UnsupportedOperationException("unknown int type with bitwidth " + bitWidth);
    case FloatingPoint:
      ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) fieldType.getType();
      FloatingPointPrecision precision = floatType.getPrecision();
      switch (precision) {
        case HALF:
        case SINGLE:
          return DataTypes.FloatType;
        case DOUBLE:
          return DataTypes.DoubleType;
      }
    case Utf8:
      return DataTypes.StringType;
    case Binary:
    case FixedSizeBinary:
      return DataTypes.BinaryType;
    case Bool:
      return DataTypes.BooleanType;
    case Decimal:
      throw new UnsupportedOperationException("have not implemented Decimal type yet");
    case Date:
      return DataTypes.DateType;
    case Time:
      return DataTypes.TimestampType; //note i don't know what this will do!
    case Timestamp:
      return DataTypes.TimestampType;
    case Interval:
      return DataTypes.CalendarIntervalType;
    case NONE:
      return DataTypes.NullType;
  }
  throw new IllegalStateException("Unexpected value: " + fieldType);
}
 

@Test
public void testFlippedCodeGenerator() throws Exception {
  ArrowType strType = new ArrowType.Utf8();
  ArrowType bigIntType = new ArrowType.Int(64, true);
  FunctionSignature substrFn = new FunctionSignature("substr", strType, Lists.newArrayList(strType, bigIntType, bigIntType));
  GandivaPushdownSieveHelper helper = new GandivaPushdownSieveHelper();
  try {
    // hide substr function if implemented in Gandiva
    helper.addFunctionToHide(substrFn);
    // enable decimal
    testContext.getOptions().setOption(OptionValue.createBoolean(
      OptionValue.OptionType.SYSTEM,
      PlannerSettings.ENABLE_DECIMAL_DATA_TYPE_KEY,
      true));
    // enable decimal v2
    testContext.getOptions().setOption(OptionValue.createBoolean(
      OptionValue.OptionType.SYSTEM,
      PlannerSettings.ENABLE_DECIMAL_V2_KEY,
      true));
    // enabled mixed mode execution
    testContext.getOptions().setOption(OptionValue.createString(
    OptionValue.OptionType.SYSTEM,
      ExecConstants.QUERY_EXEC_OPTION_KEY,
      execPreferenceMixed
    ));
    // increase the threshold for flipping the code generator
    testContext.getOptions().setOption(OptionValue.createDouble(
      OptionValue.OptionType.SYSTEM,
    ExecConstants.WORK_THRESHOLD_FOR_SPLIT_KEY,
      10.0));
    testFunctionsCompiledOnly(new Object[][]{
      {"c1 = 'CA' or c1 = 'WA' or c1 = 'GA' or c2 > 500 or substr(c0,1,4) = '8566' or substr(c0, 1, 4) = '8619' or substr(c0, 1, 4) = '8827' or substr(c0, 1, 4) = '8340' or substr(c0, 1, 4) = '1111' or substr(c0, 1, 4) = '1234' or substr(c0, 1, 4) = 2222",
        "3031", "TN", BigDecimal.valueOf(100, 2), false}
    });
  } finally {
    testContext.getOptions().setOption(OptionValue.createBoolean(
      OptionValue.OptionType.SYSTEM,
      PlannerSettings.ENABLE_DECIMAL_DATA_TYPE_KEY,
      PlannerSettings.ENABLE_DECIMAL_DATA_TYPE.getDefault().getBoolVal()));
    testContext.getOptions().setOption(OptionValue.createBoolean(
      OptionValue.OptionType.SYSTEM,
      PlannerSettings.ENABLE_DECIMAL_V2_KEY,
      PlannerSettings.ENABLE_DECIMAL_V2.getDefault().getBoolVal()));
    testContext.getOptions().setOption(OptionValue.createDouble(
      OptionValue.OptionType.SYSTEM,
      ExecConstants.WORK_THRESHOLD_FOR_SPLIT_KEY,
      ExecConstants.WORK_THRESHOLD_FOR_SPLIT.getDefault().getFloatVal()));
    testContext.getOptions().setOption(OptionValue.createString(
      OptionValue.OptionType.SYSTEM,
      ExecConstants.QUERY_EXEC_OPTION_KEY,
      execPreferenceGandivaOnly
    ));
    helper.removeFunctionToHide(substrFn);
  }
}
 

@Override
ArrowType getType() {
  return new ArrowType.Int(64, true);
}
 

@Override
ArrowType getType() {
  return new ArrowType.Int(32, true);
}
 

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

@Override
public Void visit(ArrowType.Int anInt) {
  return writeString(sqlTypeNameVisitor.visit(anInt));
}
 

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

@Override
public ArrowType visit(TinyIntType tinyIntType) {
	return new ArrowType.Int(8, true);
}
 

@Override
public ArrowType visit(SmallIntType smallIntType) {
	return new ArrowType.Int(2 * 8, true);
}
 

@Override
public ArrowType visit(IntType intType) {
	return new ArrowType.Int(4 * 8, true);
}
 

@Override
public ArrowType visit(BigIntType bigIntType) {
	return new ArrowType.Int(8 * 8, true);
}
 

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