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;
};
}
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");
}
}
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);
}
}
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;
});
}
Aggregations