use of org.apache.drill.exec.record.TransferPair in project drill by apache.
the class FlattenRecordBatch method getFlattenFieldTransferPair.
/**
* The data layout is the same for the actual data within a repeated field, as it is in a scalar vector for
* the same sql type. For example, a repeated int vector has a vector of offsets into a regular int vector to
* represent the lists. As the data layout for the actual values in the same in the repeated vector as in the
* scalar vector of the same type, we can avoid making individual copies for the column being flattened, and just
* use vector copies between the inner vector of the repeated field to the resulting scalar vector from the flatten
* operation. This is completed after we determine how many records will fit (as we will hit either a batch end, or
* the end of one of the other vectors while we are copying the data of the other vectors alongside each new flattened
* value coming out of the repeated field.)
*/
@SuppressWarnings("resource")
private TransferPair getFlattenFieldTransferPair(FieldReference reference) {
final TypedFieldId fieldId = incoming.getValueVectorId(popConfig.getColumn());
final Class<?> vectorClass = incoming.getSchema().getColumn(fieldId.getFieldIds()[0]).getValueClass();
final ValueVector flattenField = incoming.getValueAccessorById(vectorClass, fieldId.getFieldIds()).getValueVector();
TransferPair tp = null;
if (flattenField instanceof RepeatedMapVector) {
tp = ((RepeatedMapVector) flattenField).getTransferPairToSingleMap(reference.getAsNamePart().getName(), oContext.getAllocator());
} else if (!(flattenField instanceof RepeatedValueVector)) {
if (incoming.getRecordCount() != 0) {
throw UserException.unsupportedError().message("Flatten does not support inputs of non-list values.").build(logger);
}
logger.error("Cannot cast {} to RepeatedValueVector", flattenField);
//when incoming recordCount is 0, don't throw exception since the type being seen here is not solid
final ValueVector vv = new RepeatedMapVector(flattenField.getField(), oContext.getAllocator(), null);
tp = RepeatedValueVector.class.cast(vv).getTransferPair(reference.getAsNamePart().getName(), oContext.getAllocator());
} else {
final ValueVector vvIn = RepeatedValueVector.class.cast(flattenField).getDataVector();
// vvIn may be null because of fast schema return for repeated list vectors
if (vvIn != null) {
tp = vvIn.getTransferPair(reference.getAsNamePart().getName(), oContext.getAllocator());
}
}
return tp;
}
use of org.apache.drill.exec.record.TransferPair in project drill by apache.
the class FlattenRecordBatch method setupNewSchema.
@Override
protected boolean setupNewSchema() throws SchemaChangeException {
this.allocationVectors = Lists.newArrayList();
container.clear();
final List<NamedExpression> exprs = getExpressionList();
final ErrorCollector collector = new ErrorCollectorImpl();
final List<TransferPair> transfers = Lists.newArrayList();
final ClassGenerator<Flattener> cg = CodeGenerator.getRoot(Flattener.TEMPLATE_DEFINITION, context.getFunctionRegistry(), context.getOptions());
cg.getCodeGenerator().plainJavaCapable(true);
// Uncomment out this line to debug the generated code.
// cg.getCodeGenerator().saveCodeForDebugging(true);
final IntHashSet transferFieldIds = new IntHashSet();
final NamedExpression flattenExpr = new NamedExpression(popConfig.getColumn(), new FieldReference(popConfig.getColumn()));
final ValueVectorReadExpression vectorRead = (ValueVectorReadExpression) ExpressionTreeMaterializer.materialize(flattenExpr.getExpr(), incoming, collector, context.getFunctionRegistry(), true);
final FieldReference fieldReference = flattenExpr.getRef();
final TransferPair transferPair = getFlattenFieldTransferPair(fieldReference);
if (transferPair != null) {
final ValueVector flattenVector = transferPair.getTo();
// checks that list has only default ValueVector and replaces resulting ValueVector to INT typed ValueVector
if (exprs.size() == 0 && flattenVector.getField().getType().equals(Types.LATE_BIND_TYPE)) {
final MaterializedField outputField = MaterializedField.create(fieldReference.getAsNamePart().getName(), Types.OPTIONAL_INT);
final ValueVector vector = TypeHelper.getNewVector(outputField, oContext.getAllocator());
container.add(vector);
} else {
transfers.add(transferPair);
container.add(flattenVector);
transferFieldIds.add(vectorRead.getFieldId().getFieldIds()[0]);
}
}
logger.debug("Added transfer for project expression.");
ClassifierResult result = new ClassifierResult();
for (int i = 0; i < exprs.size(); i++) {
final NamedExpression namedExpression = exprs.get(i);
result.clear();
String outputName = getRef(namedExpression).getRootSegment().getPath();
if (result != null && result.outputNames != null && result.outputNames.size() > 0) {
for (int j = 0; j < result.outputNames.size(); j++) {
if (!result.outputNames.get(j).equals(EMPTY_STRING)) {
outputName = result.outputNames.get(j);
break;
}
}
}
final LogicalExpression expr = ExpressionTreeMaterializer.materialize(namedExpression.getExpr(), incoming, collector, context.getFunctionRegistry(), true);
if (collector.hasErrors()) {
throw new SchemaChangeException(String.format("Failure while trying to materialize incoming schema. Errors:\n %s.", collector.toErrorString()));
}
if (expr instanceof DrillFuncHolderExpr && ((DrillFuncHolderExpr) expr).getHolder().isComplexWriterFuncHolder()) {
// Lazy initialization of the list of complex writers, if not done yet.
if (complexWriters == null) {
complexWriters = Lists.newArrayList();
}
// The reference name will be passed to ComplexWriter, used as the name of the output vector from the writer.
((DrillFuncHolderExpr) expr).getFieldReference(namedExpression.getRef());
cg.addExpr(expr);
} else {
// need to do evaluation.
final MaterializedField outputField;
if (expr instanceof ValueVectorReadExpression) {
final TypedFieldId id = ValueVectorReadExpression.class.cast(expr).getFieldId();
@SuppressWarnings("resource") final ValueVector incomingVector = incoming.getValueAccessorById(id.getIntermediateClass(), id.getFieldIds()).getValueVector();
// when the first batch will be empty.
if (incomingVector != null) {
outputField = incomingVector.getField().clone();
} else {
outputField = MaterializedField.create(outputName, expr.getMajorType());
}
} else {
outputField = MaterializedField.create(outputName, expr.getMajorType());
}
@SuppressWarnings("resource") final ValueVector vector = TypeHelper.getNewVector(outputField, oContext.getAllocator());
allocationVectors.add(vector);
TypedFieldId fid = container.add(vector);
ValueVectorWriteExpression write = new ValueVectorWriteExpression(fid, expr, true);
cg.addExpr(write);
logger.debug("Added eval for project expression.");
}
}
cg.rotateBlock();
cg.getEvalBlock()._return(JExpr.TRUE);
container.buildSchema(SelectionVectorMode.NONE);
try {
this.flattener = context.getImplementationClass(cg.getCodeGenerator());
flattener.setup(context, incoming, this, transfers);
} catch (ClassTransformationException | IOException e) {
throw new SchemaChangeException("Failure while attempting to load generated class", e);
}
return true;
}
use of org.apache.drill.exec.record.TransferPair in project drill by apache.
the class OrderedPartitionRecordBatch method setupNewSchema.
/**
* 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(), context.getOptions());
for (VectorWrapper<?> vw : batch) {
TransferPair tp = vw.getValueVector().getTransferPair(oContext.getAllocator());
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, ClassGenerator.BlkCreateMode.FALSE);
cg.setMappingSet(partitionMapping);
ClassGenerator.HoldingContainer right = cg.addExpr(new ValueVectorReadExpression(new TypedFieldId(expr.getMajorType(), count++)), ClassGenerator.BlkCreateMode.FALSE);
cg.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());
ClassGenerator.HoldingContainer out = cg.addExpr(fh, ClassGenerator.BlkCreateMode.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);
}
}
use of org.apache.drill.exec.record.TransferPair in project drill by apache.
the class ProducerConsumerBatch method load.
private boolean load(final RecordBatchData batch) {
final VectorContainer newContainer = batch.getContainer();
if (schema != null && newContainer.getSchema().equals(schema)) {
container.zeroVectors();
final BatchSchema schema = container.getSchema();
for (int i = 0; i < container.getNumberOfColumns(); i++) {
final MaterializedField field = schema.getColumn(i);
final MajorType type = field.getType();
final ValueVector vOut = container.getValueAccessorById(TypeHelper.getValueVectorClass(type.getMinorType(), type.getMode()), container.getValueVectorId(SchemaPath.getSimplePath(field.getPath())).getFieldIds()).getValueVector();
final ValueVector vIn = newContainer.getValueAccessorById(TypeHelper.getValueVectorClass(type.getMinorType(), type.getMode()), newContainer.getValueVectorId(SchemaPath.getSimplePath(field.getPath())).getFieldIds()).getValueVector();
final TransferPair tp = vIn.makeTransferPair(vOut);
tp.transfer();
}
return false;
} else {
container.clear();
for (final VectorWrapper<?> w : newContainer) {
container.add(w.getValueVector());
}
container.buildSchema(SelectionVectorMode.NONE);
schema = container.getSchema();
return true;
}
}
use of org.apache.drill.exec.record.TransferPair in project drill by apache.
the class LimitRecordBatch method doWork.
@Override
protected IterOutcome doWork() {
if (first) {
first = false;
}
skipBatch = false;
final int recordCount = incoming.getRecordCount();
if (recordCount == 0) {
skipBatch = true;
return IterOutcome.OK;
}
for (final TransferPair tp : transfers) {
tp.transfer();
}
if (recordCount <= recordsToSkip) {
recordsToSkip -= recordCount;
skipBatch = true;
} else {
outgoingSv.allocateNew(recordCount);
limit(recordCount);
}
return IterOutcome.OK;
}
Aggregations