Search in sources :

Example 1 with DrillUnionRelBase

use of org.apache.drill.exec.planner.common.DrillUnionRelBase in project drill by apache.

the class FindLimit0Visitor method addLimitOnTopOfLeafNodes.

public static DrillRel addLimitOnTopOfLeafNodes(final DrillRel rel) {
    final Pointer<Boolean> isUnsupported = new Pointer<>(false);
    // to visit unsupported functions
    final RexShuttle unsupportedFunctionsVisitor = new RexShuttle() {

        @Override
        public RexNode visitCall(RexCall call) {
            final SqlOperator operator = call.getOperator();
            if (isUnsupportedScalarFunction(operator)) {
                isUnsupported.value = true;
                return call;
            }
            return super.visitCall(call);
        }
    };
    // to visit unsupported operators
    final RelShuttle unsupportedOperationsVisitor = new RelShuttleImpl() {

        @Override
        public RelNode visit(RelNode other) {
            if (other instanceof DrillUnionRelBase) {
                isUnsupported.value = true;
                return other;
            } else if (other instanceof DrillProjectRelBase) {
                if (!isUnsupported.value) {
                    other.accept(unsupportedFunctionsVisitor);
                }
                if (isUnsupported.value) {
                    return other;
                }
            }
            return super.visit(other);
        }
    };
    rel.accept(unsupportedOperationsVisitor);
    if (isUnsupported.value) {
        return rel;
    }
    // to add LIMIT (0) on top of leaf nodes
    final RelShuttle addLimitOnScanVisitor = new RelShuttleImpl() {

        private RelNode addLimitAsParent(RelNode node) {
            final RexBuilder builder = node.getCluster().getRexBuilder();
            final RexLiteral offset = builder.makeExactLiteral(BigDecimal.ZERO);
            final RexLiteral fetch = builder.makeExactLiteral(BigDecimal.ZERO);
            return new DrillLimitRel(node.getCluster(), node.getTraitSet(), node, offset, fetch);
        }

        @Override
        public RelNode visit(LogicalValues values) {
            return addLimitAsParent(values);
        }

        @Override
        public RelNode visit(TableScan scan) {
            return addLimitAsParent(scan);
        }

        @Override
        public RelNode visit(RelNode other) {
            if (other.getInputs().isEmpty()) {
                // leaf operator
                return addLimitAsParent(other);
            }
            return super.visit(other);
        }
    };
    return (DrillRel) rel.accept(addLimitOnScanVisitor);
}
Also used : RexLiteral(org.apache.calcite.rex.RexLiteral) TableScan(org.apache.calcite.rel.core.TableScan) RexShuttle(org.apache.calcite.rex.RexShuttle) DrillSqlOperator(org.apache.drill.exec.planner.sql.DrillSqlOperator) SqlOperator(org.apache.calcite.sql.SqlOperator) RelShuttle(org.apache.calcite.rel.RelShuttle) RelShuttleImpl(org.apache.calcite.rel.RelShuttleImpl) Pointer(org.apache.drill.exec.util.Pointer) DrillProjectRelBase(org.apache.drill.exec.planner.common.DrillProjectRelBase) DrillLimitRel(org.apache.drill.exec.planner.logical.DrillLimitRel) LogicalValues(org.apache.calcite.rel.logical.LogicalValues) RexCall(org.apache.calcite.rex.RexCall) RelNode(org.apache.calcite.rel.RelNode) DrillUnionRelBase(org.apache.drill.exec.planner.common.DrillUnionRelBase) RexBuilder(org.apache.calcite.rex.RexBuilder) DrillRel(org.apache.drill.exec.planner.logical.DrillRel)

Aggregations

RelNode (org.apache.calcite.rel.RelNode)1 RelShuttle (org.apache.calcite.rel.RelShuttle)1 RelShuttleImpl (org.apache.calcite.rel.RelShuttleImpl)1 TableScan (org.apache.calcite.rel.core.TableScan)1 LogicalValues (org.apache.calcite.rel.logical.LogicalValues)1 RexBuilder (org.apache.calcite.rex.RexBuilder)1 RexCall (org.apache.calcite.rex.RexCall)1 RexLiteral (org.apache.calcite.rex.RexLiteral)1 RexShuttle (org.apache.calcite.rex.RexShuttle)1 SqlOperator (org.apache.calcite.sql.SqlOperator)1 DrillProjectRelBase (org.apache.drill.exec.planner.common.DrillProjectRelBase)1 DrillUnionRelBase (org.apache.drill.exec.planner.common.DrillUnionRelBase)1 DrillLimitRel (org.apache.drill.exec.planner.logical.DrillLimitRel)1 DrillRel (org.apache.drill.exec.planner.logical.DrillRel)1 DrillSqlOperator (org.apache.drill.exec.planner.sql.DrillSqlOperator)1 Pointer (org.apache.drill.exec.util.Pointer)1