Skip to content

Commit

Permalink
Add TINYINT and SMALLINT support to Hive connector
Browse files Browse the repository at this point in the history
  • Loading branch information
Raghav Sethi committed May 20, 2016
1 parent a232ecd commit a601fe4
Show file tree
Hide file tree
Showing 19 changed files with 428 additions and 90 deletions.
27 changes: 25 additions & 2 deletions presto-hive/src/main/java/com/facebook/presto/hive/ColumnarBinaryHiveRecordCursor.java
View file
Expand Up @@ -62,7 +62,9 @@
import static com.facebook.presto.hive.HiveUtil.isStructuralType;
import static com.facebook.presto.hive.HiveUtil.longDecimalPartitionKey;
import static com.facebook.presto.hive.HiveUtil.shortDecimalPartitionKey;
import static com.facebook.presto.hive.HiveUtil.smallintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.timestampPartitionKey;
import static com.facebook.presto.hive.HiveUtil.tinyintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.varcharPartitionKey;
import static com.facebook.presto.hive.util.DecimalUtils.getLongDecimalValue;
import static com.facebook.presto.hive.util.DecimalUtils.getShortDecimalValue;
Expand All @@ -75,8 +77,10 @@
import static com.facebook.presto.spi.type.Decimals.isShortDecimal;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.StandardTypes.DECIMAL;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.facebook.presto.spi.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.spi.type.Varchars.isVarcharType;
import static com.facebook.presto.spi.type.Varchars.truncateToLength;
Expand Down Expand Up @@ -206,6 +210,12 @@ public ColumnarBinaryHiveRecordCursor(RecordReader<K, BytesRefArrayWritable> rec
else if (BOOLEAN.equals(type)) {
booleans[columnIndex] = booleanPartitionKey(partitionKey.getValue(), name);
}
else if (TINYINT.equals(type)) {
longs[columnIndex] = tinyintPartitionKey(partitionKey.getValue(), name);
}
else if (SMALLINT.equals(type)) {
longs[columnIndex] = smallintPartitionKey(partitionKey.getValue(), name);
}
else if (INTEGER.equals(type)) {
longs[columnIndex] = integerPartitionKey(partitionKey.getValue(), name);
}
Expand Down Expand Up @@ -347,9 +357,16 @@ public long getLong(int fieldId)
{
checkState(!closed, "Cursor is closed");

if (!types[fieldId].equals(BIGINT) && !types[fieldId].equals(INTEGER) && !types[fieldId].equals(DATE) && !types[fieldId].equals(TIMESTAMP) && !isShortDecimal(types[fieldId])) {
if (!types[fieldId].equals(BIGINT) &&
!types[fieldId].equals(INTEGER) &&
!types[fieldId].equals(SMALLINT) &&
!types[fieldId].equals(TINYINT) &&
!types[fieldId].equals(DATE) &&
!types[fieldId].equals(TIMESTAMP) &&
!isShortDecimal(types[fieldId])) {
// we don't use Preconditions.checkArgument because it requires boxing fieldId, which affects inner loop performance
throw new IllegalArgumentException(format("Expected field to be %s, %s, %s, %s or %s , actual %s (field %s)", INTEGER, BIGINT, DATE, TIMESTAMP, DECIMAL, types[fieldId], fieldId));
throw new IllegalArgumentException(
format("Expected field to be %s, %s, %s, %s, %s, %s or %s , actual %s (field %s)", TINYINT, SMALLINT, INTEGER, BIGINT, DATE, TIMESTAMP, DECIMAL, types[fieldId], fieldId));
}
if (!loaded[fieldId]) {
parseLongColumn(fieldId);
Expand Down Expand Up @@ -726,6 +743,12 @@ else if (BIGINT.equals(type)) {
else if (INTEGER.equals(type)) {
parseLongColumn(column);
}
else if (SMALLINT.equals(type)) {
parseLongColumn(column);
}
else if (TINYINT.equals(type)) {
parseLongColumn(column);
}
else if (DOUBLE.equals(type)) {
parseDoubleColumn(column);
}
Expand Down
27 changes: 25 additions & 2 deletions presto-hive/src/main/java/com/facebook/presto/hive/ColumnarTextHiveRecordCursor.java
View file
Expand Up @@ -54,7 +54,9 @@
import static com.facebook.presto.hive.HiveUtil.parseHiveDate;
import static com.facebook.presto.hive.HiveUtil.parseHiveTimestamp;
import static com.facebook.presto.hive.HiveUtil.shortDecimalPartitionKey;
import static com.facebook.presto.hive.HiveUtil.smallintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.timestampPartitionKey;
import static com.facebook.presto.hive.HiveUtil.tinyintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.varcharPartitionKey;
import static com.facebook.presto.hive.NumberParser.parseDouble;
import static com.facebook.presto.hive.NumberParser.parseLong;
Expand All @@ -67,8 +69,10 @@
import static com.facebook.presto.spi.type.Decimals.isShortDecimal;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.StandardTypes.DECIMAL;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.facebook.presto.spi.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.spi.type.Varchars.isVarcharType;
import static com.facebook.presto.spi.type.Varchars.truncateToLength;
Expand Down Expand Up @@ -192,6 +196,12 @@ public ColumnarTextHiveRecordCursor(
else if (BOOLEAN.equals(type)) {
booleans[columnIndex] = booleanPartitionKey(partitionKey.getValue(), name);
}
else if (TINYINT.equals(type)) {
longs[columnIndex] = tinyintPartitionKey(partitionKey.getValue(), name);
}
else if (SMALLINT.equals(type)) {
longs[columnIndex] = smallintPartitionKey(partitionKey.getValue(), name);
}
else if (INTEGER.equals(type)) {
longs[columnIndex] = integerPartitionKey(partitionKey.getValue(), name);
}
Expand Down Expand Up @@ -339,9 +349,16 @@ public long getLong(int fieldId)
{
checkState(!closed, "Cursor is closed");

if (!types[fieldId].equals(BIGINT) && !types[fieldId].equals(INTEGER) && !types[fieldId].equals(DATE) && !types[fieldId].equals(TIMESTAMP) && !isShortDecimal(types[fieldId])) {
if (!types[fieldId].equals(BIGINT) &&
!types[fieldId].equals(INTEGER) &&
!types[fieldId].equals(SMALLINT) &&
!types[fieldId].equals(TINYINT) &&
!types[fieldId].equals(DATE) &&
!types[fieldId].equals(TIMESTAMP) &&
!isShortDecimal(types[fieldId])) {
// we don't use Preconditions.checkArgument because it requires boxing fieldId, which affects inner loop performance
throw new IllegalArgumentException(String.format("Expected field to be %s, %s, %s, %s or %s , actual %s (field %s)", INTEGER, BIGINT, DATE, TIMESTAMP, DECIMAL, types[fieldId], fieldId));
throw new IllegalArgumentException(
format("Expected field to be %s, %s, %s, %s, %s or %s , actual %s (field %s)", TINYINT, SMALLINT, INTEGER, BIGINT, DECIMAL, DATE, TIMESTAMP, types[fieldId], fieldId));
}

if (!loaded[fieldId]) {
Expand Down Expand Up @@ -662,6 +679,12 @@ else if (type.equals(BIGINT)) {
else if (type.equals(INTEGER)) {
parseLongColumn(column);
}
else if (type.equals(SMALLINT)) {
parseLongColumn(column);
}
else if (type.equals(TINYINT)) {
parseLongColumn(column);
}
else if (type.equals(DOUBLE)) {
parseDoubleColumn(column);
}
Expand Down
16 changes: 16 additions & 0 deletions presto-hive/src/main/java/com/facebook/presto/hive/GenericHiveRecordCursor.java
View file
Expand Up @@ -54,7 +54,9 @@
import static com.facebook.presto.hive.HiveUtil.isStructuralType;
import static com.facebook.presto.hive.HiveUtil.longDecimalPartitionKey;
import static com.facebook.presto.hive.HiveUtil.shortDecimalPartitionKey;
import static com.facebook.presto.hive.HiveUtil.smallintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.timestampPartitionKey;
import static com.facebook.presto.hive.HiveUtil.tinyintPartitionKey;
import static com.facebook.presto.hive.HiveUtil.varcharPartitionKey;
import static com.facebook.presto.hive.util.SerDeUtils.getBlockObject;
import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
Expand All @@ -66,7 +68,9 @@
import static com.facebook.presto.spi.type.Decimals.rescale;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.facebook.presto.spi.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.spi.type.Varchars.isVarcharType;
import static com.facebook.presto.spi.type.Varchars.truncateToLength;
Expand Down Expand Up @@ -198,6 +202,12 @@ else if (BIGINT.equals(type)) {
else if (INTEGER.equals(type)) {
longs[columnIndex] = integerPartitionKey(partitionKey.getValue(), name);
}
else if (SMALLINT.equals(type)) {
longs[columnIndex] = smallintPartitionKey(partitionKey.getValue(), name);
}
else if (TINYINT.equals(type)) {
longs[columnIndex] = tinyintPartitionKey(partitionKey.getValue(), name);
}
else if (DOUBLE.equals(type)) {
doubles[columnIndex] = doublePartitionKey(partitionKey.getValue(), name);
}
Expand Down Expand Up @@ -534,6 +544,12 @@ else if (BIGINT.equals(type)) {
else if (INTEGER.equals(type)) {
parseLongColumn(column);
}
else if (SMALLINT.equals(type)) {
parseLongColumn(column);
}
else if (TINYINT.equals(type)) {
parseLongColumn(column);
}
else if (DOUBLE.equals(type)) {
parseDoubleColumn(column);
}
Expand Down
12 changes: 10 additions & 2 deletions presto-hive/src/main/java/com/facebook/presto/hive/HiveType.java
View file
Expand Up @@ -49,7 +49,9 @@
import static com.facebook.presto.spi.type.DecimalType.createDecimalType;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.facebook.presto.spi.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.spi.type.VarcharType.createUnboundedVarcharType;
import static com.facebook.presto.spi.type.VarcharType.createVarcharType;
Expand Down Expand Up @@ -223,6 +225,12 @@ private static TypeInfo toTypeInfo(Type type)
if (INTEGER.equals(type)) {
return HIVE_INT.typeInfo;
}
if (SMALLINT.equals(type)) {
return HIVE_SHORT.typeInfo;
}
if (TINYINT.equals(type)) {
return HIVE_BYTE.typeInfo;
}
if (DOUBLE.equals(type)) {
return HIVE_DOUBLE.typeInfo;
}
Expand Down Expand Up @@ -320,9 +328,9 @@ public static Type getPrimitiveType(PrimitiveTypeInfo typeInfo)
case BOOLEAN:
return BOOLEAN;
case BYTE:
return INTEGER;
return TINYINT;
case SHORT:
return INTEGER;
return SMALLINT;
case INT:
return INTEGER;
case LONG:
Expand Down
44 changes: 44 additions & 0 deletions presto-hive/src/main/java/com/facebook/presto/hive/HiveUtil.java
View file
Expand Up @@ -87,16 +87,20 @@
import static com.facebook.presto.spi.type.DecimalType.createDecimalType;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Predicates.not;
import static com.google.common.base.Strings.emptyToNull;
import static com.google.common.base.Strings.nullToEmpty;
import static com.google.common.collect.Iterables.filter;
import static com.google.common.collect.Lists.transform;
import static java.lang.Byte.parseByte;
import static java.lang.Double.parseDouble;
import static java.lang.Integer.parseInt;
import static java.lang.Long.parseLong;
import static java.lang.Short.parseShort;
import static java.lang.String.format;
import static java.math.BigDecimal.ROUND_UNNECESSARY;
import static java.nio.charset.StandardCharsets.UTF_8;
Expand Down Expand Up @@ -387,6 +391,26 @@ public static NullableValue parsePartitionValue(String partitionName, String val
return NullableValue.of(BOOLEAN, booleanPartitionKey(value, partitionName));
}

if (TINYINT.equals(type)) {
if (isNull) {
return NullableValue.asNull(TINYINT);
}
if (value.isEmpty()) {
return NullableValue.of(TINYINT, 0L);
}
return NullableValue.of(TINYINT, tinyintPartitionKey(value, partitionName));
}

if (SMALLINT.equals(type)) {
if (isNull) {
return NullableValue.asNull(SMALLINT);
}
if (value.isEmpty()) {
return NullableValue.of(SMALLINT, 0L);
}
return NullableValue.of(SMALLINT, smallintPartitionKey(value, partitionName));
}

if (INTEGER.equals(type)) {
if (isNull) {
return NullableValue.asNull(INTEGER);
Expand Down Expand Up @@ -535,6 +559,26 @@ public static long integerPartitionKey(String value, String name)
}
}

public static long smallintPartitionKey(String value, String name)
{
try {
return parseShort(value);
}
catch (NumberFormatException e) {
throw new PrestoException(HIVE_INVALID_PARTITION_VALUE, format("Invalid partition value '%s' for SMALLINT partition key: %s", value, name));
}
}

public static long tinyintPartitionKey(String value, String name)
{
try {
return parseByte(value);
}
catch (NumberFormatException e) {
throw new PrestoException(HIVE_INVALID_PARTITION_VALUE, format("Invalid partition value '%s' for TINYINT partition key: %s", value, name));
}
}

public static double doublePartitionKey(String value, String name)
{
try {
Expand Down

0 comments on commit a601fe4

Please sign in to comment.