Search in sources :

Example 1 with ConstraintProjector

use of com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector in project aws-athena-query-federation by awslabs.

the class JdbcRecordHandler method makeFactory.

/**
 * Create a field extractor for complex List type.
 * @param field Field's metadata information.
 * @return Extractor for the List type.
 */
protected FieldWriterFactory makeFactory(Field field) {
    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
        Array arrayField = ((ResultSet) context).getArray(field.getName());
        if (!((ResultSet) context).wasNull()) {
            List<Object> fieldValue = new ArrayList<>(Arrays.asList((Object[]) arrayField.getArray()));
            BlockUtils.setComplexValue(vector, rowNum, FieldResolver.DEFAULT, fieldValue);
        }
        return true;
    };
}
Also used : Array(java.sql.Array) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) ResultSet(java.sql.ResultSet) ArrayList(java.util.ArrayList) FieldVector(org.apache.arrow.vector.FieldVector) 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) FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter)

Example 2 with ConstraintProjector

use of com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector 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 ConstraintProjector

use of com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector in project aws-athena-query-federation by awslabs.

the class ExampleRecordHandler method makeFactory.

/**
 * Since GeneratedRowWriter doesn't yet support complex types (STRUCT, LIST) we use this to
 * create our own FieldWriters via customer FieldWriterFactory. In this case we are producing
 * FieldWriters that only work for our exact example schema. This will be enhanced with a more
 * generic solution in a future release.
 */
private FieldWriterFactory makeFactory(Field field, RowContext rowContext) {
    Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());
    switch(fieldType) {
        case LIST:
            Field child = field.getChildren().get(0);
            Types.MinorType childType = Types.getMinorTypeForArrowType(child.getType());
            switch(childType) {
                case LIST:
                    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                        UnionListWriter writer = ((ListVector) vector).getWriter();
                        writer.setPosition(rowNum);
                        writer.startList();
                        BaseWriter.ListWriter innerWriter = writer.list();
                        innerWriter.startList();
                        for (int i = 0; i < 3; i++) {
                            byte[] bytes = String.valueOf(1000 + i).getBytes(Charsets.UTF_8);
                            try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                                buf.writeBytes(bytes);
                                innerWriter.varChar().writeVarChar(0, (int) (buf.readableBytes()), buf);
                            }
                        }
                        innerWriter.endList();
                        writer.endList();
                        ((ListVector) vector).setNotNull(rowNum);
                        return true;
                    };
                case STRUCT:
                    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                        UnionListWriter writer = ((ListVector) vector).getWriter();
                        writer.setPosition(rowNum);
                        writer.startList();
                        BaseWriter.StructWriter structWriter = writer.struct();
                        structWriter.start();
                        byte[] bytes = "chars".getBytes(Charsets.UTF_8);
                        try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                            buf.writeBytes(bytes);
                            structWriter.varChar("varchar").writeVarChar(0, (int) (buf.readableBytes()), buf);
                        }
                        structWriter.bigInt("bigint").writeBigInt(100L);
                        structWriter.end();
                        writer.endList();
                        ((ListVector) vector).setNotNull(rowNum);
                        return true;
                    };
                default:
                    throw new IllegalArgumentException("Unsupported type " + childType);
            }
        case MAP:
            return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                UnionMapWriter writer = ((MapVector) vector).getWriter();
                writer.setPosition(rowNum);
                writer.startMap();
                writer.startEntry();
                byte[] bytes = "chars".getBytes(Charsets.UTF_8);
                try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                    buf.writeBytes(bytes);
                    writer.key().varChar("key").writeVarChar(0, (int) (buf.readableBytes()), buf);
                }
                writer.value().integer("value").writeInt(1001);
                writer.endEntry();
                writer.endMap();
                ((MapVector) vector).setNotNull(rowNum);
                return true;
            };
        default:
            throw new IllegalArgumentException("Unsupported type " + fieldType);
    }
}
Also used : Types(org.apache.arrow.vector.types.Types) BaseWriter(org.apache.arrow.vector.complex.writer.BaseWriter) ArrowBuf(org.apache.arrow.memory.ArrowBuf) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) FieldVector(org.apache.arrow.vector.FieldVector) FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter) Field(org.apache.arrow.vector.types.pojo.Field) UnionMapWriter(org.apache.arrow.vector.complex.impl.UnionMapWriter) UnionListWriter(org.apache.arrow.vector.complex.impl.UnionListWriter) ListVector(org.apache.arrow.vector.complex.ListVector) 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) MapVector(org.apache.arrow.vector.complex.MapVector)

Example 4 with ConstraintProjector

use of com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector in project aws-athena-query-federation by awslabs.

the class TimestreamRecordHandler method buildTimeSeriesExtractor.

private void buildTimeSeriesExtractor(GeneratedRowWriter.RowWriterBuilder builder, Field field, int curFieldNum) {
    builder.withFieldWriterFactory(field.getName(), (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (Object context, int rowNum) -> {
        Row row = (Row) context;
        Datum datum = row.getData().get(curFieldNum);
        Field timeField = field.getChildren().get(0).getChildren().get(0);
        Field valueField = field.getChildren().get(0).getChildren().get(1);
        if (datum.getTimeSeriesValue() != null) {
            List<Map<String, Object>> values = new ArrayList<>();
            for (TimeSeriesDataPoint nextDatum : datum.getTimeSeriesValue()) {
                Map<String, Object> eventMap = new HashMap<>();
                eventMap.put(timeField.getName(), TIMESTAMP_FORMATTER.parse(nextDatum.getTime()).getTime());
                switch(Types.getMinorTypeForArrowType(valueField.getType())) {
                    case FLOAT8:
                        eventMap.put(valueField.getName(), Double.valueOf(nextDatum.getValue().getScalarValue()));
                        break;
                    case BIGINT:
                        eventMap.put(valueField.getName(), Long.valueOf(nextDatum.getValue().getScalarValue()));
                        break;
                    case INT:
                        eventMap.put(valueField.getName(), Integer.valueOf(nextDatum.getValue().getScalarValue()));
                        break;
                    case BIT:
                        eventMap.put(valueField.getName(), Boolean.valueOf(((Row) context).getData().get(curFieldNum).getScalarValue()) == false ? 0 : 1);
                        break;
                }
                values.add(eventMap);
            }
            BlockUtils.setComplexValue(vector, rowNum, FieldResolver.DEFAULT, values);
        } else {
            throw new RuntimeException("Only LISTs of type TimeSeries are presently supported.");
        }
        // we don't yet support predicate pushdown on complex types
        return true;
    });
}
Also used : TimeSeriesDataPoint(com.amazonaws.services.timestreamquery.model.TimeSeriesDataPoint) Datum(com.amazonaws.services.timestreamquery.model.Datum) HashMap(java.util.HashMap) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) ArrayList(java.util.ArrayList) FieldVector(org.apache.arrow.vector.FieldVector) Field(org.apache.arrow.vector.types.pojo.Field) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor) Row(com.amazonaws.services.timestreamquery.model.Row) HashMap(java.util.HashMap) Map(java.util.Map)

Aggregations

BigIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor)4 BitExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor)4 DateMilliExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor)4 Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor)4 Float8Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor)4 VarCharExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor)4 ConstraintProjector (com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector)4 DateDayExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor)3 DecimalExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor)3 Float4Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor)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 VarBinaryExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor)3 FieldVector (org.apache.arrow.vector.FieldVector)3 Field (org.apache.arrow.vector.types.pojo.Field)3 FieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter)2 ArrayList (java.util.ArrayList)2 Types (org.apache.arrow.vector.types.Types)2 BigIntFieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BigIntFieldWriter)1