private IterOutcome doWork()
throws ClassTransformationException, IOException, SchemaChangeException {
if (allocationVectors != null) {
for (ValueVector v : allocationVectors) {
v.clear();
}
}
allocationVectors = Lists.newArrayList();
transfers.clear();
final ClassGenerator<UnionAller> cg =
CodeGenerator.getRoot(UnionAller.TEMPLATE_DEFINITION, context.getFunctionRegistry());
int index = 0;
for (VectorWrapper<?> vw : current) {
ValueVector vvIn = vw.getValueVector();
// get the original input column names
SchemaPath inputPath = vvIn.getField().getPath();
// get the renamed column names
SchemaPath outputPath = outputFields.get(index).getPath();
final ErrorCollector collector = new ErrorCollectorImpl();
// According to input data names, Minortypes, Datamodes, choose to
// transfer directly,
// rename columns or
// cast data types (Minortype or DataMode)
if (hasSameTypeAndMode(outputFields.get(index), vw.getValueVector().getField())) {
// Transfer column
if (outputFields.get(index).getPath().equals(inputPath)) {
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
inputPath, current, collector, context.getFunctionRegistry());
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
ValueVectorReadExpression vectorRead = (ValueVectorReadExpression) expr;
ValueVector vvOut =
container.addOrGet(MaterializedField.create(outputPath, vectorRead.getMajorType()));
TransferPair tp = vvIn.makeTransferPair(vvOut);
transfers.add(tp);
// Copy data in order to rename the column
} else {
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
inputPath, current, collector, context.getFunctionRegistry());
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
MaterializedField outputField = MaterializedField.create(outputPath, expr.getMajorType());
ValueVector vv = container.addOrGet(outputField, callBack);
allocationVectors.add(vv);
TypedFieldId fid = container.getValueVectorId(outputField.getPath());
ValueVectorWriteExpression write = new ValueVectorWriteExpression(fid, expr, true);
cg.addExpr(write);
}
// Cast is necessary
} else {
LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
inputPath, current, collector, context.getFunctionRegistry());
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
// If the inputs' DataMode is required and the outputs' DataMode is not required
// cast to the one with the least restriction
if (vvIn.getField().getType().getMode() == DataMode.REQUIRED
&& outputFields.get(index).getType().getMode() != DataMode.REQUIRED) {
expr =
ExpressionTreeMaterializer.convertToNullableType(
expr,
vvIn.getField().getType().getMinorType(),
context.getFunctionRegistry(),
collector);
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
}
// If two inputs' MinorTypes are different,
// Insert a cast before the Union operation
if (vvIn.getField().getType().getMinorType()
!= outputFields.get(index).getType().getMinorType()) {
expr =
ExpressionTreeMaterializer.addCastExpression(
expr,
outputFields.get(index).getType(),
context.getFunctionRegistry(),
collector);
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
}
final MaterializedField outputField =
MaterializedField.create(outputPath, expr.getMajorType());
ValueVector vector = container.addOrGet(outputField, callBack);
allocationVectors.add(vector);
TypedFieldId fid = container.getValueVectorId(outputField.getPath());
boolean useSetSafe = !(vector instanceof FixedWidthVector);
ValueVectorWriteExpression write = new ValueVectorWriteExpression(fid, expr, useSetSafe);
cg.addExpr(write);
}
++index;
}
unionall = context.getImplementationClass(cg.getCodeGenerator());
unionall.setup(context, current, this, transfers);
if (!schemaAvailable) {
container.buildSchema(BatchSchema.SelectionVectorMode.NONE);
schemaAvailable = true;
}
if (!doAlloc()) {
return IterOutcome.OUT_OF_MEMORY;
}
recordCount = unionall.unionRecords(0, current.getRecordCount(), 0);
setValueCount(recordCount);
return IterOutcome.OK;
}
