public static List<HiveColumnHandle> getPartitionKeyColumnHandles(
String connectorId, Table table, boolean forceIntegralToBigint) {
ImmutableList.Builder<HiveColumnHandle> columns = ImmutableList.builder();
List<Column> partitionKeys = table.getPartitionColumns();
for (Column field : partitionKeys) {
HiveType hiveType = field.getType();
if (!hiveType.isSupportedType()) {
throw new PrestoException(
NOT_SUPPORTED,
format(
"Unsupported Hive type %s found in partition keys of table %s.%s",
hiveType, table.getDatabaseName(), table.getTableName()));
}
columns.add(
new HiveColumnHandle(
connectorId,
field.getName(),
hiveType,
hiveType.getTypeSignature(forceIntegralToBigint),
-1,
PARTITION_KEY));
}
return columns.build();
}
@Override
protected String extractField(Object target) {
if (target instanceof HiveType) {
HiveType type = (HiveType) target;
ObjectInspector inspector = type.getObjectInspector();
if (inspector instanceof StructObjectInspector) {
StructObjectInspector soi = (StructObjectInspector) inspector;
StructField field = soi.getStructFieldRef(fieldName);
ObjectInspector foi = field.getFieldObjectInspector();
Assert.isTrue(
foi.getCategory() == ObjectInspector.Category.PRIMITIVE,
String.format(
"Field [%s] needs to be a primitive; found [%s]", fieldName, foi.getTypeName()));
// expecting a writeable - simply do a toString
Object data = soi.getStructFieldData(type.getObject(), field);
if (data == null || data instanceof NullWritable) {
return StringUtils.EMPTY;
}
return data.toString();
}
}
return null;
}
protected List<HiveColumnHandle> getColumnHandles(List<TestColumn> testColumns) {
List<HiveColumnHandle> columns = new ArrayList<>();
int nextHiveColumnIndex = 0;
for (int i = 0; i < testColumns.size(); i++) {
TestColumn testColumn = testColumns.get(i);
int columnIndex = testColumn.isPartitionKey() ? -1 : nextHiveColumnIndex++;
HiveType hiveType = HiveType.valueOf(testColumn.getObjectInspector().getTypeName());
columns.add(
new HiveColumnHandle(
"client_id",
testColumn.getName(),
hiveType,
hiveType.getTypeSignature(),
columnIndex,
testColumn.isPartitionKey()));
}
return columns;
}
protected void checkCursor(RecordCursor cursor, List<TestColumn> testColumns, int numRows)
throws IOException {
for (int row = 0; row < numRows; row++) {
assertTrue(cursor.advanceNextPosition());
for (int i = 0, testColumnsSize = testColumns.size(); i < testColumnsSize; i++) {
TestColumn testColumn = testColumns.get(i);
Object fieldFromCursor;
Type type =
HiveType.valueOf(testColumn.getObjectInspector().getTypeName()).getType(TYPE_MANAGER);
if (cursor.isNull(i)) {
fieldFromCursor = null;
} else if (BOOLEAN.equals(type)) {
fieldFromCursor = cursor.getBoolean(i);
} else if (BIGINT.equals(type)) {
fieldFromCursor = cursor.getLong(i);
} else if (DOUBLE.equals(type)) {
fieldFromCursor = cursor.getDouble(i);
} else if (VARCHAR.equals(type)) {
fieldFromCursor = cursor.getSlice(i);
} else if (VARBINARY.equals(type)) {
fieldFromCursor = cursor.getSlice(i);
} else if (DateType.DATE.equals(type)) {
fieldFromCursor = cursor.getLong(i);
} else if (TimestampType.TIMESTAMP.equals(type)) {
fieldFromCursor = cursor.getLong(i);
} else if (isStructuralType(type)) {
fieldFromCursor = cursor.getObject(i);
} else {
throw new RuntimeException("unknown type");
}
if (fieldFromCursor == null) {
assertEquals(
null,
testColumn.getExpectedValue(),
String.format("Expected null for column %s", testColumn.getName()));
} else if (testColumn.getObjectInspector().getTypeName().equals("float")
|| testColumn.getObjectInspector().getTypeName().equals("double")) {
assertEquals((double) fieldFromCursor, (double) testColumn.getExpectedValue(), EPSILON);
} else if (testColumn.getObjectInspector().getCategory() == Category.PRIMITIVE) {
assertEquals(
fieldFromCursor,
testColumn.getExpectedValue(),
String.format("Wrong value for column %s", testColumn.getName()));
} else {
Block expected = (Block) testColumn.getExpectedValue();
Block actual = (Block) fieldFromCursor;
assertBlockEquals(
actual, expected, String.format("Wrong value for column %s", testColumn.getName()));
}
}
}
}
public static List<HiveColumnHandle> getRegularColumnHandles(
String connectorId, Table table, boolean forceIntegralToBigint) {
ImmutableList.Builder<HiveColumnHandle> columns = ImmutableList.builder();
int hiveColumnIndex = 0;
for (Column field : table.getDataColumns()) {
// ignore unsupported types rather than failing
HiveType hiveType = field.getType();
if (hiveType.isSupportedType()) {
columns.add(
new HiveColumnHandle(
connectorId,
field.getName(),
hiveType,
hiveType.getTypeSignature(forceIntegralToBigint),
hiveColumnIndex,
REGULAR));
}
hiveColumnIndex++;
}
return columns.build();
}
public static boolean isStructuralType(HiveType hiveType) {
return hiveType.getCategory() == Category.LIST
|| hiveType.getCategory() == Category.MAP
|| hiveType.getCategory() == Category.STRUCT;
}
public static Optional<DecimalType> getDecimalType(HiveType hiveType) {
return getDecimalType(hiveType.getHiveTypeName());
}