Search in sources :

Example 1 with NullableDecimalHolder

use of com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder in project aws-athena-query-federation by awslabs.

the class ExampleRecordHandler method makeExtractor.

/**
 * Creates an Extractor for the given field. In this example the extractor just creates some random data.
 */
private Extractor makeExtractor(Field field, RowContext rowContext) {
    Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());
    // they need to match the split otherwise filtering will brake in unexpected ways.
    if (field.getName().equals("year")) {
        return (IntExtractor) (Object context, NullableIntHolder dst) -> {
            dst.isSet = 1;
            dst.value = rowContext.getYear();
        };
    } else if (field.getName().equals("month")) {
        return (IntExtractor) (Object context, NullableIntHolder dst) -> {
            dst.isSet = 1;
            dst.value = rowContext.getMonth();
        };
    } else if (field.getName().equals("day")) {
        return (IntExtractor) (Object context, NullableIntHolder dst) -> {
            dst.isSet = 1;
            dst.value = rowContext.getDay();
        };
    }
    switch(fieldType) {
        case INT:
            return (IntExtractor) (Object context, NullableIntHolder dst) -> {
                dst.isSet = 1;
                dst.value = ((RowContext) context).seed * (((RowContext) context).negative ? -1 : 1);
            };
        case DATEMILLI:
            return (DateMilliExtractor) (Object context, NullableDateMilliHolder dst) -> {
                dst.isSet = 1;
                dst.value = ((RowContext) context).seed * (((RowContext) context).negative ? -1 : 1);
            };
        case DATEDAY:
            return (DateDayExtractor) (Object context, NullableDateDayHolder dst) -> {
                dst.isSet = 1;
                dst.value = ((RowContext) context).seed * (((RowContext) context).negative ? -1 : 1);
            };
        case TINYINT:
            return (TinyIntExtractor) (Object context, NullableTinyIntHolder dst) -> {
                dst.isSet = 1;
                dst.value = (byte) ((((RowContext) context).seed % 4) * (((RowContext) context).negative ? -1 : 1));
            };
        case SMALLINT:
            return (SmallIntExtractor) (Object context, NullableSmallIntHolder dst) -> {
                dst.isSet = 1;
                dst.value = (short) ((((RowContext) context).seed % 4) * (((RowContext) context).negative ? -1 : 1));
            };
        case FLOAT4:
            return (Float4Extractor) (Object context, NullableFloat4Holder dst) -> {
                dst.isSet = 1;
                dst.value = ((float) ((RowContext) context).seed) * 1.1f * (((RowContext) context).negative ? -1f : 1f);
            };
        case FLOAT8:
            return (Float8Extractor) (Object context, NullableFloat8Holder dst) -> {
                dst.isSet = 1;
                dst.value = ((double) ((RowContext) context).seed) * 1.1D;
            };
        case DECIMAL:
            return (DecimalExtractor) (Object context, NullableDecimalHolder dst) -> {
                dst.isSet = 1;
                double d8Val = ((RowContext) context).seed * 1.1D * (((RowContext) context).negative ? -1d : 1d);
                BigDecimal bdVal = new BigDecimal(d8Val);
                dst.value = bdVal.setScale(((ArrowType.Decimal) field.getType()).getScale(), RoundingMode.HALF_UP);
            };
        case BIT:
            return (BitExtractor) (Object context, NullableBitHolder dst) -> {
                dst.isSet = 1;
                dst.value = ((RowContext) context).seed % 2;
            };
        case BIGINT:
            return (BigIntExtractor) (Object context, NullableBigIntHolder dst) -> {
                dst.isSet = 1;
                dst.value = ((RowContext) context).seed * 1L * (((RowContext) context).negative ? -1 : 1);
            };
        case VARCHAR:
            return (VarCharExtractor) (Object context, NullableVarCharHolder dst) -> {
                dst.isSet = 1;
                dst.value = "VarChar" + ((RowContext) context).seed;
            };
        case VARBINARY:
            return (VarBinaryExtractor) (Object context, NullableVarBinaryHolder dst) -> {
                dst.isSet = 1;
                dst.value = ("VarChar" + ((RowContext) context).seed).getBytes(Charsets.UTF_8);
            };
        default:
            return null;
    }
}
Also used : Types(org.apache.arrow.vector.types.Types) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) NullableDateMilliHolder(org.apache.arrow.vector.holders.NullableDateMilliHolder) NullableFloat8Holder(org.apache.arrow.vector.holders.NullableFloat8Holder) NullableTinyIntHolder(org.apache.arrow.vector.holders.NullableTinyIntHolder) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) DecimalExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor) NullableVarBinaryHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarBinaryHolder) NullableVarCharHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder) BigDecimal(java.math.BigDecimal) NullableIntHolder(org.apache.arrow.vector.holders.NullableIntHolder) DateDayExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) NullableDateDayHolder(org.apache.arrow.vector.holders.NullableDateDayHolder) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) NullableSmallIntHolder(org.apache.arrow.vector.holders.NullableSmallIntHolder) NullableDecimalHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder) BigDecimal(java.math.BigDecimal) NullableBigIntHolder(org.apache.arrow.vector.holders.NullableBigIntHolder) NullableBitHolder(org.apache.arrow.vector.holders.NullableBitHolder) NullableFloat4Holder(org.apache.arrow.vector.holders.NullableFloat4Holder) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) VarBinaryExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor)

Example 2 with NullableDecimalHolder

use of com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder in project aws-athena-query-federation by awslabs.

the class VerticaRecordHandler method makeExtractor.

/**
 * Creates an Extractor for the given field.
 */
private Extractor makeExtractor(Field field, HashMap<String, Types.MinorType> mapOfNamesAndTypes, HashMap<String, Object> mapOfcols) {
    String fieldName = field.getName();
    Types.MinorType fieldType = mapOfNamesAndTypes.get(fieldName);
    switch(fieldType) {
        case BIT:
            return (BitExtractor) (Object context, NullableBitHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = ((boolean) value) ? 1 : 0;
                    dst.isSet = 1;
                }
            };
        case TINYINT:
            return (TinyIntExtractor) (Object context, NullableTinyIntHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = Byte.parseByte(value.toString());
                    dst.isSet = 1;
                }
            };
        case SMALLINT:
            return (SmallIntExtractor) (Object context, NullableSmallIntHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = Short.parseShort(value.toString());
                    dst.isSet = 1;
                }
            };
        case INT:
        case BIGINT:
            return (BigIntExtractor) (Object context, NullableBigIntHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = Long.parseLong(value.toString());
                    dst.isSet = 1;
                }
            };
        case FLOAT4:
            return (Float4Extractor) (Object context, NullableFloat4Holder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = Float.parseFloat(value.toString());
                    dst.isSet = 1;
                }
            };
        case FLOAT8:
            return (Float8Extractor) (Object context, NullableFloat8Holder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = Double.parseDouble(value.toString());
                    dst.isSet = 1;
                }
            };
        case DECIMAL:
            return (DecimalExtractor) (Object context, NullableDecimalHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = new BigDecimal(value.toString());
                    dst.isSet = 1;
                }
            };
        case DATEDAY:
            return (DateDayExtractor) (Object context, NullableDateDayHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.isSet = 1;
                    dst.value = (int) LocalDate.parse(value.toString()).toEpochDay();
                }
            };
        case DATEMILLI:
            return (DateMilliExtractor) (Object context, NullableDateMilliHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName).toString();
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = LocalDateTime.parse(value.toString()).atZone(BlockUtils.UTC_ZONE_ID).toInstant().toEpochMilli();
                    dst.isSet = 1;
                }
            };
        case VARCHAR:
            return (VarCharExtractor) (Object context, NullableVarCharHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName);
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = value.toString();
                    dst.isSet = 1;
                }
            };
        case VARBINARY:
            return (VarBinaryExtractor) (Object context, NullableVarBinaryHolder dst) -> {
                Object value = ((RowContext) context).getNameValue().get(fieldName).toString();
                if (value == null) {
                    dst.isSet = 0;
                } else {
                    dst.value = value.toString().getBytes();
                    dst.isSet = 1;
                }
            };
        default:
            throw new RuntimeException("Unhandled type " + fieldType);
    }
}
Also used : Types(org.apache.arrow.vector.types.Types) NullableVarBinaryHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarBinaryHolder) NullableVarCharHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder) NullableDecimalHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder) BigDecimal(java.math.BigDecimal)

Aggregations

NullableDecimalHolder (com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder)2 NullableVarBinaryHolder (com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarBinaryHolder)2 NullableVarCharHolder (com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder)2 BigDecimal (java.math.BigDecimal)2 Types (org.apache.arrow.vector.types.Types)2 BigIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor)1 BitExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor)1 DateDayExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor)1 DateMilliExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor)1 DecimalExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor)1 Float4Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor)1 Float8Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor)1 IntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor)1 SmallIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor)1 TinyIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor)1 VarBinaryExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor)1 VarCharExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor)1 NullableBigIntHolder (org.apache.arrow.vector.holders.NullableBigIntHolder)1 NullableBitHolder (org.apache.arrow.vector.holders.NullableBitHolder)1 NullableDateDayHolder (org.apache.arrow.vector.holders.NullableDateDayHolder)1