Search in sources :

Example 16 with JExpression

use of com.sun.codemodel.JExpression in project drill by apache.

the class PartitionSenderRootExec method createClassInstances.

private List<Partitioner> createClassInstances(int actualPartitions) throws SchemaChangeException {
    // set up partitioning function
    final LogicalExpression expr = operator.getExpr();
    final ErrorCollector collector = new ErrorCollectorImpl();
    final ClassGenerator<Partitioner> cg;
    cg = CodeGenerator.getRoot(Partitioner.TEMPLATE_DEFINITION, context.getFunctionRegistry(), context.getOptions());
    cg.getCodeGenerator().plainJavaCapable(true);
    // Uncomment out this line to debug the generated code.
    //    cg.getCodeGenerator().saveCodeForDebugging(true);
    ClassGenerator<Partitioner> cgInner = cg.getInnerGenerator("OutgoingRecordBatch");
    final LogicalExpression materializedExpr = ExpressionTreeMaterializer.materialize(expr, incoming, collector, context.getFunctionRegistry());
    if (collector.hasErrors()) {
        throw new SchemaChangeException(String.format("Failure while trying to materialize incoming schema.  Errors:\n %s.", collector.toErrorString()));
    }
    // generate code to copy from an incoming value vector to the destination partition's outgoing value vector
    JExpression bucket = JExpr.direct("bucket");
    // generate evaluate expression to determine the hash
    ClassGenerator.HoldingContainer exprHolder = cg.addExpr(materializedExpr);
    cg.getEvalBlock().decl(JType.parse(cg.getModel(), "int"), "bucket", exprHolder.getValue().mod(JExpr.lit(outGoingBatchCount)));
    cg.getEvalBlock()._return(cg.getModel().ref(Math.class).staticInvoke("abs").arg(bucket));
    CopyUtil.generateCopies(cgInner, incoming, incoming.getSchema().getSelectionVectorMode() == SelectionVectorMode.FOUR_BYTE);
    try {
        // compile and setup generated code
        List<Partitioner> subPartitioners = context.getImplementationClass(cg, actualPartitions);
        return subPartitioners;
    } catch (ClassTransformationException | IOException e) {
        throw new SchemaChangeException("Failure while attempting to load generated class", e);
    }
}
Also used : ClassTransformationException(org.apache.drill.exec.exception.ClassTransformationException) ErrorCollector(org.apache.drill.common.expression.ErrorCollector) JExpression(com.sun.codemodel.JExpression) IOException(java.io.IOException) ErrorCollectorImpl(org.apache.drill.common.expression.ErrorCollectorImpl) LogicalExpression(org.apache.drill.common.expression.LogicalExpression) SchemaChangeException(org.apache.drill.exec.exception.SchemaChangeException) ClassGenerator(org.apache.drill.exec.expr.ClassGenerator)

Example 17 with JExpression

use of com.sun.codemodel.JExpression in project drill by apache.

the class NestedLoopJoinBatch method setupWorker.

/**
   * Method generates the runtime code needed for NLJ. Other than the setup method to set the input and output value
   * vector references we implement three more methods
   * 1. doEval() -> Evaluates if record from left side matches record from the right side
   * 2. emitLeft() -> Project record from the left side
   * 3. emitRight() -> Project record from the right side (which is a hyper container)
   * @return the runtime generated class that implements the NestedLoopJoin interface
   */
private NestedLoopJoin setupWorker() throws IOException, ClassTransformationException, SchemaChangeException {
    final CodeGenerator<NestedLoopJoin> nLJCodeGenerator = CodeGenerator.get(NestedLoopJoin.TEMPLATE_DEFINITION, context.getFunctionRegistry(), context.getOptions());
    nLJCodeGenerator.plainJavaCapable(true);
    // Uncomment out this line to debug the generated code.
    //    nLJCodeGenerator.saveCodeForDebugging(true);
    final ClassGenerator<NestedLoopJoin> nLJClassGenerator = nLJCodeGenerator.getRoot();
    // generate doEval
    final ErrorCollector collector = new ErrorCollectorImpl();
    /*
        Logical expression may contain fields from left and right batches. During code generation (materialization)
        we need to indicate from which input field should be taken.

        Non-equality joins can belong to one of below categories. For example:
        1. Join on non-equality join predicates:
        select * from t1 inner join t2 on (t1.c1 between t2.c1 AND t2.c2) AND (...)
        2. Join with an OR predicate:
        select * from t1 inner join t2 on on t1.c1 = t2.c1 OR t1.c2 = t2.c2
     */
    Map<VectorAccessible, BatchReference> batches = ImmutableMap.<VectorAccessible, BatchReference>builder().put(left, new BatchReference("leftBatch", "leftIndex")).put(rightContainer, new BatchReference("rightContainer", "rightBatchIndex", "rightRecordIndexWithinBatch")).build();
    LogicalExpression materialize = ExpressionTreeMaterializer.materialize(popConfig.getCondition(), batches, collector, context.getFunctionRegistry(), false, false);
    if (collector.hasErrors()) {
        throw new SchemaChangeException(String.format("Failure while trying to materialize join condition. Errors:\n %s.", collector.toErrorString()));
    }
    nLJClassGenerator.addExpr(new ReturnValueExpression(materialize), ClassGenerator.BlkCreateMode.FALSE);
    // generate emitLeft
    nLJClassGenerator.setMappingSet(emitLeftMapping);
    JExpression outIndex = JExpr.direct("outIndex");
    JExpression leftIndex = JExpr.direct("leftIndex");
    int fieldId = 0;
    int outputFieldId = 0;
    // Set the input and output value vector references corresponding to the left batch
    for (MaterializedField field : leftSchema) {
        final TypeProtos.MajorType fieldType = field.getType();
        // Add the vector to the output container
        container.addOrGet(field);
        JVar inVV = nLJClassGenerator.declareVectorValueSetupAndMember("leftBatch", new TypedFieldId(fieldType, false, fieldId));
        JVar outVV = nLJClassGenerator.declareVectorValueSetupAndMember("outgoing", new TypedFieldId(fieldType, false, outputFieldId));
        nLJClassGenerator.getEvalBlock().add(outVV.invoke("copyFromSafe").arg(leftIndex).arg(outIndex).arg(inVV));
        nLJClassGenerator.rotateBlock();
        fieldId++;
        outputFieldId++;
    }
    // generate emitRight
    fieldId = 0;
    nLJClassGenerator.setMappingSet(emitRightMapping);
    JExpression batchIndex = JExpr.direct("batchIndex");
    JExpression recordIndexWithinBatch = JExpr.direct("recordIndexWithinBatch");
    // Set the input and output value vector references corresponding to the right batch
    for (MaterializedField field : rightSchema) {
        final TypeProtos.MajorType inputType = field.getType();
        TypeProtos.MajorType outputType;
        // if join type is LEFT, make sure right batch output fields data mode is optional
        if (popConfig.getJoinType() == JoinRelType.LEFT && inputType.getMode() == TypeProtos.DataMode.REQUIRED) {
            outputType = Types.overrideMode(inputType, TypeProtos.DataMode.OPTIONAL);
        } else {
            outputType = inputType;
        }
        MaterializedField newField = MaterializedField.create(field.getPath(), outputType);
        container.addOrGet(newField);
        JVar inVV = nLJClassGenerator.declareVectorValueSetupAndMember("rightContainer", new TypedFieldId(inputType, true, fieldId));
        JVar outVV = nLJClassGenerator.declareVectorValueSetupAndMember("outgoing", new TypedFieldId(outputType, false, outputFieldId));
        nLJClassGenerator.getEvalBlock().add(outVV.invoke("copyFromSafe").arg(recordIndexWithinBatch).arg(outIndex).arg(inVV.component(batchIndex)));
        nLJClassGenerator.rotateBlock();
        fieldId++;
        outputFieldId++;
    }
    return context.getImplementationClass(nLJCodeGenerator);
}
Also used : VectorAccessible(org.apache.drill.exec.record.VectorAccessible) ErrorCollector(org.apache.drill.common.expression.ErrorCollector) MaterializedField(org.apache.drill.exec.record.MaterializedField) JExpression(com.sun.codemodel.JExpression) TypeProtos(org.apache.drill.common.types.TypeProtos) ErrorCollectorImpl(org.apache.drill.common.expression.ErrorCollectorImpl) ReturnValueExpression(org.apache.drill.exec.physical.impl.filter.ReturnValueExpression) LogicalExpression(org.apache.drill.common.expression.LogicalExpression) SchemaChangeException(org.apache.drill.exec.exception.SchemaChangeException) BatchReference(org.apache.drill.exec.expr.BatchReference) TypedFieldId(org.apache.drill.exec.record.TypedFieldId) JVar(com.sun.codemodel.JVar)

Aggregations

JExpression (com.sun.codemodel.JExpression)17 JVar (com.sun.codemodel.JVar)13 JBlock (com.sun.codemodel.JBlock)9 JMethod (com.sun.codemodel.JMethod)9 JClass (com.sun.codemodel.JClass)7 JInvocation (com.sun.codemodel.JInvocation)7 JConditional (com.sun.codemodel.JConditional)5 JDefinedClass (com.sun.codemodel.JDefinedClass)5 JFieldVar (com.sun.codemodel.JFieldVar)4 Map (java.util.Map)4 JType (com.sun.codemodel.JType)3 JsonNode (com.fasterxml.jackson.databind.JsonNode)2 ByteString (com.linkedin.data.ByteString)2 ArrayDataSchema (com.linkedin.data.schema.ArrayDataSchema)2 DataSchema (com.linkedin.data.schema.DataSchema)2 EnumDataSchema (com.linkedin.data.schema.EnumDataSchema)2 MapDataSchema (com.linkedin.data.schema.MapDataSchema)2 RecordDataSchema (com.linkedin.data.schema.RecordDataSchema)2 ErrorCollector (org.apache.drill.common.expression.ErrorCollector)2 ErrorCollectorImpl (org.apache.drill.common.expression.ErrorCollectorImpl)2