Examples with FunctionLookupContext org.apache.drill.exec.expr.fn.FunctionLookupContext used on opensource projects

Example 1 with FunctionLookupContext

use of org.apache.drill.exec.expr.fn.FunctionLookupContext in project drill by apache.

the class TopNBatch method createNewPriorityQueue.

public static PriorityQueue createNewPriorityQueue(MappingSet mainMapping, MappingSet leftMapping, MappingSet rightMapping, List<Ordering> orderings, VectorAccessible batch, boolean unionTypeEnabled, boolean codegenDump, int limit, BufferAllocator allocator, SelectionVectorMode mode, FragmentContext context) {
    OptionSet optionSet = context.getOptions();
    FunctionLookupContext functionLookupContext = context.getFunctionRegistry();
    CodeGenerator<PriorityQueue> cg = CodeGenerator.get(PriorityQueue.TEMPLATE_DEFINITION, optionSet);
    cg.plainJavaCapable(true);
    cg.saveCodeForDebugging(codegenDump);
    // Uncomment out this line to debug the generated code.
    // cg.saveCodeForDebugging(true);
    ClassGenerator<PriorityQueue> 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, functionLookupContext, unionTypeEnabled);
        collector.reportErrors(logger);
        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, functionLookupContext);
        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));
    PriorityQueue q = context.getImplementationClass(cg);
    try {
        q.init(limit, allocator, mode == BatchSchema.SelectionVectorMode.TWO_BYTE);
    } catch (SchemaChangeException e) {
        throw TopNBatch.schemaChangeException(e, "Top N", logger);
    }
    return q;
}