use of org.apache.arrow.vector.holders.NullableBigIntHolder in project aws-athena-query-federation by awslabs.
the class VertexRowWriter method writeRowTemplate.
public static void writeRowTemplate(RowWriterBuilder rowWriterBuilder, Field field) {
ArrowType arrowType = field.getType();
Types.MinorType minorType = Types.getMinorTypeForArrowType(arrowType);
switch(minorType) {
case BIT:
rowWriterBuilder.withExtractor(field.getName(), (BitExtractor) (Object context, NullableBitHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
value.isSet = 0;
if (objValues != null && objValues.get(0) != null) {
Boolean booleanValue = Boolean.parseBoolean(objValues.get(0).toString());
value.value = booleanValue ? 1 : 0;
value.isSet = 1;
}
});
break;
case VARCHAR:
rowWriterBuilder.withExtractor(field.getName(), (VarCharExtractor) (Object context, NullableVarCharHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
String fieldName = field.getName().toLowerCase().trim();
value.isSet = 0;
// check for special keys and parse them separately
if (fieldName.equals(SpecialKeys.ID.toString().toLowerCase())) {
Object fieldValue = obj.get(T.id);
if (fieldValue != null) {
value.value = fieldValue.toString();
value.isSet = 1;
}
} else {
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
if (objValues != null && objValues.get(0) != null) {
value.value = objValues.get(0).toString();
value.isSet = 1;
}
}
});
break;
case INT:
rowWriterBuilder.withExtractor(field.getName(), (IntExtractor) (Object context, NullableIntHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
value.isSet = 0;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
if (objValues != null && objValues.get(0) != null) {
value.value = Integer.parseInt(objValues.get(0).toString());
value.isSet = 1;
}
});
break;
case BIGINT:
rowWriterBuilder.withExtractor(field.getName(), (BigIntExtractor) (Object context, NullableBigIntHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
value.isSet = 0;
if (objValues != null && objValues.get(0) != null) {
value.value = Long.parseLong(objValues.get(0).toString());
value.isSet = 1;
}
});
break;
case FLOAT4:
rowWriterBuilder.withExtractor(field.getName(), (Float4Extractor) (Object context, NullableFloat4Holder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
value.isSet = 0;
if (objValues != null && objValues.get(0) != null) {
value.value = Float.parseFloat(objValues.get(0).toString());
value.isSet = 1;
}
});
break;
case FLOAT8:
rowWriterBuilder.withExtractor(field.getName(), (Float8Extractor) (Object context, NullableFloat8Holder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
value.isSet = 0;
if (objValues != null && objValues.get(0) != null) {
value.value = Double.parseDouble(objValues.get(0).toString());
value.isSet = 1;
}
});
break;
}
}
use of org.apache.arrow.vector.holders.NullableBigIntHolder in project aws-athena-query-federation by awslabs.
the class TimestreamRecordHandler method buildRowWriter.
private GeneratedRowWriter buildRowWriter(ReadRecordsRequest request) {
GeneratedRowWriter.RowWriterBuilder builder = GeneratedRowWriter.newBuilder(request.getConstraints());
int fieldNum = 0;
for (Field nextField : request.getSchema().getFields()) {
int curFieldNum = fieldNum++;
switch(Types.getMinorTypeForArrowType(nextField.getType())) {
case VARCHAR:
builder.withExtractor(nextField.getName(), (VarCharExtractor) (Object context, NullableVarCharHolder value) -> {
value.isSet = 1;
value.value = ((Row) context).getData().get(curFieldNum).getScalarValue();
});
break;
case FLOAT8:
builder.withExtractor(nextField.getName(), (Float8Extractor) (Object context, NullableFloat8Holder value) -> {
value.isSet = 1;
value.value = Double.valueOf(((Row) context).getData().get(curFieldNum).getScalarValue());
});
break;
case BIT:
builder.withExtractor(nextField.getName(), (BitExtractor) (Object context, NullableBitHolder value) -> {
value.isSet = 1;
value.value = Boolean.valueOf(((Row) context).getData().get(curFieldNum).getScalarValue()) == false ? 0 : 1;
});
break;
case BIGINT:
builder.withExtractor(nextField.getName(), (BigIntExtractor) (Object context, NullableBigIntHolder value) -> {
value.isSet = 1;
value.value = Long.valueOf(((Row) context).getData().get(curFieldNum).getScalarValue());
});
break;
case DATEMILLI:
builder.withExtractor(nextField.getName(), (DateMilliExtractor) (Object context, NullableDateMilliHolder value) -> {
value.isSet = 1;
value.value = TIMESTAMP_FORMATTER.parse(((Row) context).getData().get(curFieldNum).getScalarValue()).getTime();
});
break;
case LIST:
// TODO: This presently only supports TimeSeries results but it is possible that customers may
// generate LIST type results for other reasons when using VIEWs. For now this seems like an OK
// compromise since it enables an important capability of TimeStream even if it doesn't enable arbitrary
// complex types.
buildTimeSeriesExtractor(builder, nextField, curFieldNum);
break;
default:
throw new RuntimeException("Unsupported field type[" + nextField.getType() + "] for field[" + nextField.getName() + "]");
}
}
return builder.build();
}
use of org.apache.arrow.vector.holders.NullableBigIntHolder 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;
}
use of org.apache.arrow.vector.holders.NullableBigIntHolder 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;
}
}
use of org.apache.arrow.vector.holders.NullableBigIntHolder in project aws-athena-query-federation by awslabs.
the class EdgeRowWriter method writeRowTemplate.
public static void writeRowTemplate(RowWriterBuilder rowWriterBuilder, Field field) {
ArrowType arrowType = field.getType();
Types.MinorType minorType = Types.getMinorTypeForArrowType(arrowType);
switch(minorType) {
case BIT:
rowWriterBuilder.withExtractor(field.getName(), (BitExtractor) (Object context, NullableBitHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
Object fieldValue = obj.get(field.getName());
value.isSet = 0;
if (fieldValue != null) {
Boolean booleanValue = Boolean.parseBoolean(fieldValue.toString());
value.value = booleanValue ? 1 : 0;
value.isSet = 1;
}
});
break;
case VARCHAR:
rowWriterBuilder.withExtractor(field.getName(), (VarCharExtractor) (Object context, NullableVarCharHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
value.isSet = 0;
String fieldName = field.getName().toLowerCase().trim();
// check for special keys and parse them separately
if (fieldName.equals(SpecialKeys.ID.toString().toLowerCase())) {
Object fieldValue = obj.get(T.id);
if (fieldValue != null) {
value.value = fieldValue.toString();
value.isSet = 1;
}
} else if (fieldName.equals(SpecialKeys.IN.toString().toLowerCase())) {
Object fieldValue = ((LinkedHashMap) obj.get(Direction.IN)).get(T.id);
if (fieldValue != null) {
value.value = fieldValue.toString();
value.isSet = 1;
}
} else if (fieldName.equals(SpecialKeys.OUT.toString().toLowerCase())) {
Object fieldValue = ((LinkedHashMap) obj.get(Direction.OUT)).get(T.id);
if (fieldValue != null) {
value.value = fieldValue.toString();
value.isSet = 1;
}
} else {
Object fieldValue = obj.get(field.getName());
if (fieldValue != null) {
value.value = fieldValue.toString();
value.isSet = 1;
}
}
});
break;
case INT:
rowWriterBuilder.withExtractor(field.getName(), (IntExtractor) (Object context, NullableIntHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
value.isSet = 0;
Object fieldValue = obj.get(field.getName());
if (fieldValue != null) {
value.value = Integer.parseInt(fieldValue.toString());
value.isSet = 1;
}
});
break;
case BIGINT:
rowWriterBuilder.withExtractor(field.getName(), (BigIntExtractor) (Object context, NullableBigIntHolder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
value.isSet = 0;
Object fieldValue = obj.get(field.getName());
if (fieldValue != null) {
value.value = Long.parseLong(fieldValue.toString());
value.isSet = 1;
}
});
break;
case FLOAT4:
rowWriterBuilder.withExtractor(field.getName(), (Float4Extractor) (Object context, NullableFloat4Holder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
ArrayList<Object> objValues = (ArrayList) obj.get(field.getName());
value.isSet = 0;
Object fieldValue = obj.get(field.getName());
if (fieldValue != null) {
value.value = Float.parseFloat(fieldValue.toString());
value.isSet = 1;
}
});
break;
case FLOAT8:
rowWriterBuilder.withExtractor(field.getName(), (Float8Extractor) (Object context, NullableFloat8Holder value) -> {
Map<Object, Object> obj = (Map<Object, Object>) context;
value.isSet = 0;
Object fieldValue = obj.get(field.getName());
if (fieldValue != null) {
value.value = Double.parseDouble(fieldValue.toString());
value.isSet = 1;
}
});
break;
}
}
Aggregations