Search in sources :

Example 1 with VarCharExtractor

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

the class ElasticsearchTypeUtilsTest method testField.

/**
 * Uses the correct field extractor to extract values from a document.
 * @param mapping is the metadata definitions of the document being processed.
 * @param document contains the values to be extracted.
 * @return a map of the field names and their associated values extracted from the document.
 * @throws Exception
 */
private Map<String, Object> testField(Schema mapping, Map<String, Object> document) throws Exception {
    Map<String, Object> results = new HashMap<>();
    for (Field field : mapping.getFields()) {
        Extractor extractor = typeUtils.makeExtractor(field);
        if (extractor instanceof VarCharExtractor) {
            NullableVarCharHolder holder = new NullableVarCharHolder();
            ((VarCharExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof BigIntExtractor) {
            NullableBigIntHolder holder = new NullableBigIntHolder();
            ((BigIntExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof IntExtractor) {
            NullableIntHolder holder = new NullableIntHolder();
            ((IntExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof SmallIntExtractor) {
            NullableSmallIntHolder holder = new NullableSmallIntHolder();
            ((SmallIntExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof TinyIntExtractor) {
            NullableTinyIntHolder holder = new NullableTinyIntHolder();
            ((TinyIntExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof Float8Extractor) {
            NullableFloat8Holder holder = new NullableFloat8Holder();
            ((Float8Extractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof Float4Extractor) {
            NullableFloat4Holder holder = new NullableFloat4Holder();
            ((Float4Extractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof DateMilliExtractor) {
            NullableDateMilliHolder holder = new NullableDateMilliHolder();
            ((DateMilliExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        } else if (extractor instanceof BitExtractor) {
            NullableBitHolder holder = new NullableBitHolder();
            ((BitExtractor) extractor).extract(document, holder);
            assertEquals("Could not extract value for: " + field.getName(), 1, holder.isSet);
            results.put(field.getName(), holder.value);
        }
    }
    return results;
}
Also used : BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) NullableDateMilliHolder(org.apache.arrow.vector.holders.NullableDateMilliHolder) HashMap(java.util.HashMap) NullableSmallIntHolder(org.apache.arrow.vector.holders.NullableSmallIntHolder) NullableFloat8Holder(org.apache.arrow.vector.holders.NullableFloat8Holder) NullableTinyIntHolder(org.apache.arrow.vector.holders.NullableTinyIntHolder) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) NullableBigIntHolder(org.apache.arrow.vector.holders.NullableBigIntHolder) NullableBitHolder(org.apache.arrow.vector.holders.NullableBitHolder) Field(org.apache.arrow.vector.types.pojo.Field) NullableVarCharHolder(com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder) NullableFloat4Holder(org.apache.arrow.vector.holders.NullableFloat4Holder) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) NullableIntHolder(org.apache.arrow.vector.holders.NullableIntHolder) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor)

Example 2 with VarCharExtractor

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

the class GeneratedRowWriter method makeFieldWriter.

private FieldWriter makeFieldWriter(FieldVector vector) {
    Field field = vector.getField();
    String fieldName = field.getName();
    Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());
    Extractor extractor = extractors.get(fieldName);
    ConstraintProjector constraint = constraints.get(fieldName);
    FieldWriterFactory factory = fieldWriterFactories.get(fieldName);
    if (factory != null) {
        return factory.create(vector, extractor, constraint);
    }
    if (extractor == null) {
        throw new IllegalStateException("Missing extractor for field[" + fieldName + "]");
    }
    switch(fieldType) {
        case INT:
            return new IntFieldWriter((IntExtractor) extractor, (IntVector) vector, constraint);
        case BIGINT:
            return new BigIntFieldWriter((BigIntExtractor) extractor, (BigIntVector) vector, constraint);
        case DATEMILLI:
            return new DateMilliFieldWriter((DateMilliExtractor) extractor, (DateMilliVector) vector, constraint);
        case DATEDAY:
            return new DateDayFieldWriter((DateDayExtractor) extractor, (DateDayVector) vector, constraint);
        case TINYINT:
            return new TinyIntFieldWriter((TinyIntExtractor) extractor, (TinyIntVector) vector, constraint);
        case SMALLINT:
            return new SmallIntFieldWriter((SmallIntExtractor) extractor, (SmallIntVector) vector, constraint);
        case FLOAT4:
            return new Float4FieldWriter((Float4Extractor) extractor, (Float4Vector) vector, constraint);
        case FLOAT8:
            return new Float8FieldWriter((Float8Extractor) extractor, (Float8Vector) vector, constraint);
        case DECIMAL:
            return new DecimalFieldWriter((DecimalExtractor) extractor, (DecimalVector) vector, constraint);
        case BIT:
            return new BitFieldWriter((BitExtractor) extractor, (BitVector) vector, constraint);
        case VARCHAR:
            return new VarCharFieldWriter((VarCharExtractor) extractor, (VarCharVector) vector, constraint);
        case VARBINARY:
            return new VarBinaryFieldWriter((VarBinaryExtractor) extractor, (VarBinaryVector) vector, constraint);
        default:
            throw new RuntimeException(fieldType + " is not supported");
    }
}
Also used : Types(org.apache.arrow.vector.types.Types) VarBinaryFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.VarBinaryFieldWriter) TinyIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.TinyIntFieldWriter) SmallIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.SmallIntFieldWriter) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) DecimalFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DecimalFieldWriter) BigIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BigIntFieldWriter) FieldWriterFactory(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriterFactory) BitFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BitFieldWriter) DateMilliFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateMilliFieldWriter) DateDayFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateDayFieldWriter) VarCharFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.VarCharFieldWriter) Field(org.apache.arrow.vector.types.pojo.Field) IntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.IntFieldWriter) TinyIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.TinyIntFieldWriter) BigIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BigIntFieldWriter) SmallIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.SmallIntFieldWriter) Float4FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.Float4FieldWriter) Float8FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.Float8FieldWriter) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) DecimalExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor) DateDayExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) VarBinaryExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor)

Example 3 with VarCharExtractor

use of com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor 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)

Aggregations

BigIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor)3 BitExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor)3 DateMilliExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor)3 Float4Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor)3 Float8Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor)3 IntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor)3 SmallIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor)3 TinyIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor)3 VarCharExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor)3 DateDayExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor)2 DecimalExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor)2 Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor)2 VarBinaryExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor)2 NullableVarCharHolder (com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder)2 NullableBigIntHolder (org.apache.arrow.vector.holders.NullableBigIntHolder)2 NullableBitHolder (org.apache.arrow.vector.holders.NullableBitHolder)2 NullableDateMilliHolder (org.apache.arrow.vector.holders.NullableDateMilliHolder)2 NullableFloat4Holder (org.apache.arrow.vector.holders.NullableFloat4Holder)2 NullableFloat8Holder (org.apache.arrow.vector.holders.NullableFloat8Holder)2 NullableIntHolder (org.apache.arrow.vector.holders.NullableIntHolder)2