private void createCopier(
VectorAccessible batch,
List<BatchGroup> batchGroupList,
VectorContainer outputContainer,
boolean spilling)
throws SchemaChangeException {
try {
if (copier == null) {
CodeGenerator<PriorityQueueCopier> cg =
CodeGenerator.get(
PriorityQueueCopier.TEMPLATE_DEFINITION,
context.getFunctionRegistry(),
context.getOptions());
ClassGenerator<PriorityQueueCopier> g = cg.getRoot();
generateComparisons(g, batch);
g.setMappingSet(COPIER_MAPPING_SET);
CopyUtil.generateCopies(g, batch, true);
g.setMappingSet(MAIN_MAPPING);
copier = context.getImplementationClass(cg);
} else {
copier.close();
}
BufferAllocator allocator = spilling ? copierAllocator : oAllocator;
for (VectorWrapper<?> i : batch) {
ValueVector v = TypeHelper.getNewVector(i.getField(), allocator);
outputContainer.add(v);
}
copier.setup(context, allocator, batch, batchGroupList, outputContainer);
} catch (ClassTransformationException | IOException e) {
throw new RuntimeException(e);
}
}
private void addRecordValues(
ClassGenerator<StreamingAggregator> cg, LogicalExpression[] valueExprs) {
cg.setMappingSet(EVAL);
for (final LogicalExpression ex : valueExprs) {
final HoldingContainer hc = cg.addExpr(ex);
}
}
/**
* Sets up projection that will transfer all of the columns in batch, and also populate the
* partition column based on which partition a record falls into in the partition table
*
* @param batch
* @throws SchemaChangeException
*/
protected void setupNewSchema(VectorAccessible batch) throws SchemaChangeException {
container.clear();
final ErrorCollector collector = new ErrorCollectorImpl();
final List<TransferPair> transfers = Lists.newArrayList();
final ClassGenerator<OrderedPartitionProjector> cg =
CodeGenerator.getRoot(
OrderedPartitionProjector.TEMPLATE_DEFINITION, context.getFunctionRegistry());
for (VectorWrapper<?> vw : batch) {
TransferPair tp = vw.getValueVector().getTransferPair();
transfers.add(tp);
container.add(tp.getTo());
}
cg.setMappingSet(mainMapping);
int count = 0;
for (Ordering od : popConfig.getOrderings()) {
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
od.getExpr(), batch, collector, context.getFunctionRegistry());
if (collector.hasErrors())
throw new SchemaChangeException(
"Failure while materializing expression. " + collector.toErrorString());
cg.setMappingSet(incomingMapping);
ClassGenerator.HoldingContainer left = cg.addExpr(expr, false);
cg.setMappingSet(partitionMapping);
ClassGenerator.HoldingContainer right =
cg.addExpr(
new ValueVectorReadExpression(new TypedFieldId(expr.getMajorType(), count++)), false);
cg.setMappingSet(mainMapping);
LogicalExpression fh =
FunctionGenerationHelper.getComparator(left, right, context.getFunctionRegistry());
ClassGenerator.HoldingContainer out = cg.addExpr(fh, false);
JConditional jc = cg.getEvalBlock()._if(out.getValue().ne(JExpr.lit(0)));
if (od.getDirection() == Direction.ASCENDING) {
jc._then()._return(out.getValue());
} else {
jc._then()._return(out.getValue().minus());
}
}
cg.getEvalBlock()._return(JExpr.lit(0));
container.add(this.partitionKeyVector);
container.buildSchema(batch.getSchema().getSelectionVectorMode());
try {
this.projector = context.getImplementationClass(cg);
projector.setup(
context, batch, this, transfers, partitionVectors, partitions, popConfig.getRef());
} catch (ClassTransformationException | IOException e) {
throw new SchemaChangeException("Failure while attempting to load generated class", e);
}
}
private void outputRecordKeys(
ClassGenerator<StreamingAggregator> cg,
TypedFieldId[] keyOutputIds,
LogicalExpression[] keyExprs) {
cg.setMappingSet(RECORD_KEYS);
for (int i = 0; i < keyExprs.length; i++) {
final HoldingContainer hc =
cg.addExpr(new ValueVectorWriteExpression(keyOutputIds[i], keyExprs[i], true));
}
}
private void getIndex(ClassGenerator<StreamingAggregator> g) {
switch (incoming.getSchema().getSelectionVectorMode()) {
case FOUR_BYTE:
{
JVar var = g.declareClassField("sv4_", g.getModel()._ref(SelectionVector4.class));
g.getBlock("setupInterior")
.assign(var, JExpr.direct("incoming").invoke("getSelectionVector4"));
g.getBlock("getVectorIndex")._return(var.invoke("get").arg(JExpr.direct("recordIndex")));
;
return;
}
case NONE:
{
g.getBlock("getVectorIndex")._return(JExpr.direct("recordIndex"));
;
return;
}
case TWO_BYTE:
{
JVar var = g.declareClassField("sv2_", g.getModel()._ref(SelectionVector2.class));
g.getBlock("setupInterior")
.assign(var, JExpr.direct("incoming").invoke("getSelectionVector2"));
g.getBlock("getVectorIndex")
._return(var.invoke("getIndex").arg(JExpr.direct("recordIndex")));
;
return;
}
default:
throw new IllegalStateException();
}
}
protected HoldingContainer generateEvalBody(
ClassGenerator<?> g, HoldingContainer[] inputVariables, String body, JVar[] workspaceJVars) {
// g.getBlock().directStatement(String.format("//---- start of eval portion of %s function.
// ----//", functionName));
JBlock sub = new JBlock(true, true);
JBlock topSub = sub;
HoldingContainer out = null;
MajorType returnValueType = returnValue.type;
// add outside null handling if it is defined.
if (nullHandling == NullHandling.NULL_IF_NULL) {
JExpression e = null;
for (HoldingContainer v : inputVariables) {
if (v.isOptional()) {
if (e == null) {
e = v.getIsSet();
} else {
e = e.mul(v.getIsSet());
}
}
}
if (e != null) {
// if at least one expression must be checked, set up the conditional.
returnValueType = returnValue.type.toBuilder().setMode(DataMode.OPTIONAL).build();
out = g.declare(returnValueType);
e = e.eq(JExpr.lit(0));
JConditional jc = sub._if(e);
jc._then().assign(out.getIsSet(), JExpr.lit(0));
sub = jc._else();
}
}
if (out == null) out = g.declare(returnValueType);
// add the subblock after the out declaration.
g.getEvalBlock().add(topSub);
JVar internalOutput =
sub.decl(
JMod.FINAL,
g.getHolderType(returnValueType),
returnValue.name,
JExpr._new(g.getHolderType(returnValueType)));
addProtectedBlock(g, sub, body, inputVariables, workspaceJVars, false);
if (sub != topSub)
sub.assign(internalOutput.ref("isSet"), JExpr.lit(1)); // Assign null if NULL_IF_NULL mode
sub.assign(out.getHolder(), internalOutput);
if (sub != topSub)
sub.assign(internalOutput.ref("isSet"), JExpr.lit(1)); // Assign null if NULL_IF_NULL mode
return out;
}
/**
* Creates a copier that does a project for every Nth record from a VectorContainer incoming into
* VectorContainer outgoing. Each Ordering in orderings generates a column, and evaluation of the
* expression associated with each Ordering determines the value of each column. These records
* will later be sorted based on the values in each column, in the same order as the orderings.
*
* @param sv4
* @param incoming
* @param outgoing
* @param orderings
* @return
* @throws SchemaChangeException
*/
private SampleCopier getCopier(
SelectionVector4 sv4,
VectorContainer incoming,
VectorContainer outgoing,
List<Ordering> orderings,
List<ValueVector> localAllocationVectors)
throws SchemaChangeException {
final ErrorCollector collector = new ErrorCollectorImpl();
final ClassGenerator<SampleCopier> cg =
CodeGenerator.getRoot(SampleCopier.TEMPLATE_DEFINITION, context.getFunctionRegistry());
int i = 0;
for (Ordering od : orderings) {
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
od.getExpr(), incoming, collector, context.getFunctionRegistry());
SchemaPath schemaPath = SchemaPath.getSimplePath("f" + i++);
TypeProtos.MajorType.Builder builder =
TypeProtos.MajorType.newBuilder()
.mergeFrom(expr.getMajorType())
.clearMode()
.setMode(TypeProtos.DataMode.REQUIRED);
TypeProtos.MajorType newType = builder.build();
MaterializedField outputField = MaterializedField.create(schemaPath, newType);
if (collector.hasErrors()) {
throw new SchemaChangeException(
String.format(
"Failure while trying to materialize incoming schema. Errors:\n %s.",
collector.toErrorString()));
}
ValueVector vector = TypeHelper.getNewVector(outputField, oContext.getAllocator());
localAllocationVectors.add(vector);
TypedFieldId fid = outgoing.add(vector);
ValueVectorWriteExpression write = new ValueVectorWriteExpression(fid, expr, true);
HoldingContainer hc = cg.addExpr(write);
cg.getEvalBlock()._if(hc.getValue().eq(JExpr.lit(0)))._then()._return(JExpr.FALSE);
}
cg.rotateBlock();
cg.getEvalBlock()._return(JExpr.TRUE);
outgoing.buildSchema(BatchSchema.SelectionVectorMode.NONE);
try {
SampleCopier sampleCopier = context.getImplementationClass(cg);
sampleCopier.setupCopier(context, sv4, incoming, outgoing);
return sampleCopier;
} catch (ClassTransformationException | IOException e) {
throw new SchemaChangeException(e);
}
}
private void setupIsSame(ClassGenerator<StreamingAggregator> cg, LogicalExpression[] keyExprs) {
cg.setMappingSet(IS_SAME_I1);
for (final LogicalExpression expr : keyExprs) {
// first, we rewrite the evaluation stack for each side of the comparison.
cg.setMappingSet(IS_SAME_I1);
final HoldingContainer first = cg.addExpr(expr, false);
cg.setMappingSet(IS_SAME_I2);
final HoldingContainer second = cg.addExpr(expr, false);
final LogicalExpression fh =
FunctionGenerationHelper.getOrderingComparatorNullsHigh(
first, second, context.getFunctionRegistry());
final HoldingContainer out = cg.addExpr(fh, false);
cg.getEvalBlock()._if(out.getValue().ne(JExpr.lit(0)))._then()._return(JExpr.FALSE);
}
cg.getEvalBlock()._return(JExpr.TRUE);
}
private void outputRecordKeysPrev(
ClassGenerator<StreamingAggregator> cg,
TypedFieldId[] keyOutputIds,
LogicalExpression[] keyExprs) {
cg.setMappingSet(RECORD_KEYS_PREV);
for (int i = 0; i < keyExprs.length; i++) {
// IMPORTANT: there is an implicit assertion here that the TypedFieldIds for the previous
// batch and the current batch are the same. This is possible because InternalBatch
// guarantees this.
logger.debug("Writing out expr {}", keyExprs[i]);
cg.rotateBlock();
cg.setMappingSet(RECORD_KEYS_PREV);
final HoldingContainer innerExpression = cg.addExpr(keyExprs[i], false);
cg.setMappingSet(RECORD_KEYS_PREV_OUT);
final HoldingContainer outerExpression =
cg.addExpr(
new ValueVectorWriteExpression(
keyOutputIds[i], new HoldingContainerExpression(innerExpression), true),
false);
}
}
private StreamingAggregator createAggregatorInternal()
throws SchemaChangeException, ClassTransformationException, IOException {
ClassGenerator<StreamingAggregator> cg =
CodeGenerator.getRoot(
StreamingAggTemplate.TEMPLATE_DEFINITION, context.getFunctionRegistry());
container.clear();
LogicalExpression[] keyExprs = new LogicalExpression[popConfig.getKeys().length];
LogicalExpression[] valueExprs = new LogicalExpression[popConfig.getExprs().length];
TypedFieldId[] keyOutputIds = new TypedFieldId[popConfig.getKeys().length];
ErrorCollector collector = new ErrorCollectorImpl();
for (int i = 0; i < keyExprs.length; i++) {
final NamedExpression ne = popConfig.getKeys()[i];
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
ne.getExpr(), incoming, collector, context.getFunctionRegistry());
if (expr == null) {
continue;
}
keyExprs[i] = expr;
final MaterializedField outputField =
MaterializedField.create(ne.getRef(), expr.getMajorType());
final ValueVector vector = TypeHelper.getNewVector(outputField, oContext.getAllocator());
keyOutputIds[i] = container.add(vector);
}
for (int i = 0; i < valueExprs.length; i++) {
final NamedExpression ne = popConfig.getExprs()[i];
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
ne.getExpr(), incoming, collector, context.getFunctionRegistry());
if (expr instanceof IfExpression) {
throw UserException.unsupportedError(
new UnsupportedOperationException(
"Union type not supported in aggregate functions"))
.build(logger);
}
if (expr == null) {
continue;
}
final MaterializedField outputField =
MaterializedField.create(ne.getRef(), expr.getMajorType());
ValueVector vector = TypeHelper.getNewVector(outputField, oContext.getAllocator());
TypedFieldId id = container.add(vector);
valueExprs[i] = new ValueVectorWriteExpression(id, expr, true);
}
if (collector.hasErrors()) {
throw new SchemaChangeException(
"Failure while materializing expression. " + collector.toErrorString());
}
setupIsSame(cg, keyExprs);
setupIsSameApart(cg, keyExprs);
addRecordValues(cg, valueExprs);
outputRecordKeys(cg, keyOutputIds, keyExprs);
outputRecordKeysPrev(cg, keyOutputIds, keyExprs);
cg.getBlock("resetValues")._return(JExpr.TRUE);
getIndex(cg);
container.buildSchema(SelectionVectorMode.NONE);
StreamingAggregator agg = context.getImplementationClass(cg);
agg.setup(oContext, incoming, this);
return agg;
}
private void generateComparisons(ClassGenerator<?> g, VectorAccessible batch)
throws SchemaChangeException {
g.setMappingSet(MAIN_MAPPING);
for (Ordering od : popConfig.getOrderings()) {
// first, we rewrite the evaluation stack for each side of the comparison.
ErrorCollector collector = new ErrorCollectorImpl();
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
od.getExpr(), batch, collector, context.getFunctionRegistry());
if (collector.hasErrors()) {
throw new SchemaChangeException(
"Failure while materializing expression. " + collector.toErrorString());
}
g.setMappingSet(LEFT_MAPPING);
HoldingContainer left = g.addExpr(expr, ClassGenerator.BlkCreateMode.FALSE);
g.setMappingSet(RIGHT_MAPPING);
HoldingContainer right = g.addExpr(expr, ClassGenerator.BlkCreateMode.FALSE);
g.setMappingSet(MAIN_MAPPING);
// next we wrap the two comparison sides and add the expression block for the comparison.
LogicalExpression fh =
FunctionGenerationHelper.getOrderingComparator(
od.nullsSortHigh(), left, right, context.getFunctionRegistry());
HoldingContainer out = g.addExpr(fh, ClassGenerator.BlkCreateMode.FALSE);
JConditional jc = g.getEvalBlock()._if(out.getValue().ne(JExpr.lit(0)));
if (od.getDirection() == Direction.ASCENDING) {
jc._then()._return(out.getValue());
} else {
jc._then()._return(out.getValue().minus());
}
g.rotateBlock();
}
g.rotateBlock();
g.getEvalBlock()._return(JExpr.lit(0));
}
private MSorter createNewMSorter(
FragmentContext context,
List<Ordering> orderings,
VectorAccessible batch,
MappingSet mainMapping,
MappingSet leftMapping,
MappingSet rightMapping)
throws ClassTransformationException, IOException, SchemaChangeException {
CodeGenerator<MSorter> cg =
CodeGenerator.get(
MSorter.TEMPLATE_DEFINITION, context.getFunctionRegistry(), context.getOptions());
ClassGenerator<MSorter> g = cg.getRoot();
g.setMappingSet(mainMapping);
for (Ordering od : orderings) {
// first, we rewrite the evaluation stack for each side of the comparison.
ErrorCollector collector = new ErrorCollectorImpl();
final LogicalExpression expr =
ExpressionTreeMaterializer.materialize(
od.getExpr(), batch, collector, context.getFunctionRegistry());
if (collector.hasErrors()) {
throw new SchemaChangeException(
"Failure while materializing expression. " + collector.toErrorString());
}
g.setMappingSet(leftMapping);
HoldingContainer left = g.addExpr(expr, ClassGenerator.BlkCreateMode.FALSE);
g.setMappingSet(rightMapping);
HoldingContainer right = g.addExpr(expr, ClassGenerator.BlkCreateMode.FALSE);
g.setMappingSet(mainMapping);
// next we wrap the two comparison sides and add the expression block for the comparison.
LogicalExpression fh =
FunctionGenerationHelper.getOrderingComparator(
od.nullsSortHigh(), left, right, context.getFunctionRegistry());
HoldingContainer out = g.addExpr(fh, ClassGenerator.BlkCreateMode.FALSE);
JConditional jc = g.getEvalBlock()._if(out.getValue().ne(JExpr.lit(0)));
if (od.getDirection() == Direction.ASCENDING) {
jc._then()._return(out.getValue());
} else {
jc._then()._return(out.getValue().minus());
}
g.rotateBlock();
}
g.rotateBlock();
g.getEvalBlock()._return(JExpr.lit(0));
return context.getImplementationClass(cg);
}
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;
}
