/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.drill.exec.planner.logical.partition;
import java.util.ArrayDeque;
import java.util.BitSet;
import java.util.Deque;
import java.util.List;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexCorrelVariable;
import org.apache.calcite.rex.RexDynamicParam;
import org.apache.calcite.rex.RexFieldAccess;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexOver;
import org.apache.calcite.rex.RexRangeRef;
import org.apache.calcite.rex.RexVisitorImpl;
import org.apache.calcite.sql.SqlBinaryOperator;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlOperator;
import org.apache.calcite.sql.SqlSyntax;
import org.apache.calcite.sql.fun.SqlRowOperator;
import org.apache.calcite.util.Util;
import com.google.common.collect.Lists;
public class FindPartitionConditions extends RexVisitorImpl<Void> {
/** Whether an expression is a directory filter, and if so, whether
* it can be pushed into the scan.
*/
enum PushDirFilter {
NO_PUSH, PUSH, PARTIAL_PUSH
}
/**
* During top-down traversal of the expression tree, keep track of the
* current operators that are directory filters. Children that are
* directory filters add themselves to their parent operators.
*
* NOTE: this auxiliary class is necessary because RexNodes are immutable.
* If they were mutable, we could have easily added/dropped inputs as we
* encountered directory filters.
*/
public class OpState {
private SqlOperator sqlOperator;
private List<RexNode> children = Lists.newArrayList();
public OpState(SqlOperator op) {
sqlOperator = op;
}
public SqlOperator getOp() {
return sqlOperator;
}
public void addChild(RexNode n) {
if (!children.contains(n)) {
children.add(n);
}
}
public List<RexNode> getChildren() {
return children;
}
public void clear() {
children.clear();
}
}
private final BitSet dirs;
// The Scan could be projecting several dirN columns but we are only interested in the
// ones that are referenced by the Filter, so keep track of such referenced dirN columns.
private final BitSet referencedDirs;
private final List<PushDirFilter> pushStatusStack = Lists.newArrayList();
private final Deque<OpState> opStack = new ArrayDeque<OpState>();
/* While traversing the filter tree, some RexCalls need special handling where
* a partial expression cannot be built bottom up and the whole expression needs
* to be built. Eg: cast(VALUE/COLUMN as Type). We need to treat this as holistic
* expression and cannot treat the VALUE/COLUMN to be casted as a separate expression
* Below count keeps track of expressions that need such handling. Its type is an integer
* and not a boolean because we can have such nested expressions.
*/
private int holisticExpression = 0;
private RexBuilder builder = null;
private RexNode resultCondition = null;
public FindPartitionConditions(BitSet dirs) {
// go deep
super(true);
this.dirs = dirs;
this.referencedDirs = new BitSet(dirs.size());
}
public FindPartitionConditions(BitSet dirs, RexBuilder builder) {
// go deep
super(true);
this.dirs = dirs;
this.builder = builder;
this.referencedDirs = new BitSet(dirs.size());
}
public void analyze(RexNode exp) {
assert pushStatusStack.isEmpty();
exp.accept(this);
// Deal with top of stack
assert pushStatusStack.size() == 1;
PushDirFilter rootPushDirFilter = pushStatusStack.get(0);
if (rootPushDirFilter == PushDirFilter.PUSH) {
// The entire subtree was directory filter, so add it to the result.
addResult(exp);
}
pushStatusStack.clear();
}
public RexNode getFinalCondition() {
return resultCondition;
}
public BitSet getReferencedDirs() {
return referencedDirs;
}
private Void pushVariable() {
pushStatusStack.add(PushDirFilter.NO_PUSH);
return null;
}
private void addResult(RexNode exp) {
// processing a special expression, will be handled holistically by the parent
if (holisticExpression > 0) {
return;
}
// when we find a directory filter, add it to the current operator's
// children (if one exists)
if (!opStack.isEmpty()) {
OpState op = opStack.peek();
op.addChild(exp);
} else {
resultCondition = exp;
}
}
/*
* One of the children for the current OP is a NO PUSH. Clear all the children
* added for the current OP that were a PUSH.
*/
private void clearChildren() {
if (!opStack.isEmpty()) {
OpState op = opStack.peek();
if (op.getChildren().size() >= 1) {
op.clear();
}
}
}
/*
* Pop the current operation from the stack (opStack) as we are done with it.
* If it has children, then it means the current OP itself is a PUSH
* and hence add itself as a child to the parent OP. If no more parents
* remains, it means that the current OP is the root condition and
* set the root condition in that case.
*/
private void popOpStackAndBuildFilter() {
// Parsing a special expression; handled holistically by the parent
if (holisticExpression > 0) {
return;
}
OpState currentOp = opStack.pop();
int size = currentOp.getChildren().size();
RexNode newFilter = null;
if (size >= 1) {
if (size == 1 && currentOp.getOp() instanceof SqlBinaryOperator) {
/* The only operator for which we allow partial pushes is AND.
* For all other operators we clear the children if one of the
* children is a no push.
*/
assert currentOp.getOp().getKind() == SqlKind.AND;
newFilter = currentOp.getChildren().get(0);
} else {
newFilter = builder.makeCall(currentOp.getOp(), currentOp.getChildren());
}
}
if (newFilter != null) {
// add this new filter to my parent boolean operator's children
if (!opStack.isEmpty()) {
OpState parentOp = opStack.peek();
parentOp.addChild(newFilter);
} else {
resultCondition = newFilter;
}
}
}
private boolean isHolisticExpression(RexCall call) {
/* If we are already processing a holistic expression then all
* nested expressions should be treated holistically as well
*/
if (holisticExpression > 0) {
return true;
}
if (call.getOperator().getSyntax() == SqlSyntax.SPECIAL ||
call.getOperator().getSyntax() == SqlSyntax.FUNCTION) {
return true;
}
return false;
}
public Void visitInputRef(RexInputRef inputRef) {
if(dirs.get(inputRef.getIndex())){
pushStatusStack.add(PushDirFilter.PUSH);
addResult(inputRef);
referencedDirs.set(inputRef.getIndex());
}else{
pushStatusStack.add(PushDirFilter.NO_PUSH);
}
return null;
}
public Void visitLiteral(RexLiteral literal) {
pushStatusStack.add(PushDirFilter.PUSH);
addResult(literal);
return null;
}
public Void visitOver(RexOver over) {
// assume NO_PUSH until proven otherwise
analyzeCall(over, PushDirFilter.NO_PUSH);
return null;
}
public Void visitCorrelVariable(RexCorrelVariable correlVariable) {
return pushVariable();
}
public Void visitCall(RexCall call) {
analyzeCall(call, PushDirFilter.PUSH);
return null;
}
/*
* Traverse over the RexCall recursively. All children that are potential
* PUSH are added as children to the current operator. Once all children
* of the current operator are processed and if it satisfies to be a
* PUSH then the current operator is added as a child to its parent and so on.
*
* This bottom up building does not work for some expressions like 'cast'.
* In that case, we don't add the current operator or its children in the
* opStack and based on whether the special expression is a PUSH we add the
* whole expression to its parent op. Fpr such special expressions we maintain
* specialExpression count.
*/
private void analyzeCall(RexCall call, PushDirFilter callPushDirFilter) {
if (isHolisticExpression(call)) {
holisticExpression++;
} else {
opStack.push(new OpState(call.getOperator()));
}
// visit operands, pushing their states onto stack
super.visitCall(call);
// look for NO_PUSH operands
int operandCount = call.getOperands().size();
List<PushDirFilter> operandStack = Util.last(pushStatusStack, operandCount);
for (PushDirFilter operandPushDirFilter : operandStack) {
if (operandPushDirFilter == PushDirFilter.NO_PUSH) {
callPushDirFilter = PushDirFilter.NO_PUSH;
} else if (operandPushDirFilter == PushDirFilter.PARTIAL_PUSH) {
callPushDirFilter = PushDirFilter.PARTIAL_PUSH;
}
}
// Even if all operands are PUSH, the call itself may
// be non-deterministic.
if (!call.getOperator().isDeterministic()) {
callPushDirFilter = PushDirFilter.NO_PUSH;
} else if (call.getOperator().isDynamicFunction()) {
// For now, treat it same as non-deterministic.
callPushDirFilter = PushDirFilter.NO_PUSH;
}
// Row operator itself can't be reduced to a PUSH
if ((callPushDirFilter == PushDirFilter.PUSH)
&& (call.getOperator() instanceof SqlRowOperator)) {
callPushDirFilter = PushDirFilter.NO_PUSH;
}
if (callPushDirFilter == PushDirFilter.NO_PUSH) {
OpState currentOp = opStack.peek();
if (currentOp != null) {
if (currentOp.sqlOperator.getKind() != SqlKind.AND) {
clearChildren();
} else {
// AND op, check if we pushed some children
if (currentOp.children.size() > 0) {
callPushDirFilter = PushDirFilter.PARTIAL_PUSH;
}
}
}
}
// pop operands off of the stack
operandStack.clear();
if (isHolisticExpression(call)) {
assert holisticExpression > 0;
holisticExpression--;
if (callPushDirFilter == PushDirFilter.PUSH) {
addResult(call);
}
} else {
popOpStackAndBuildFilter();
}
// push PushDirFilter result for this call onto stack
pushStatusStack.add(callPushDirFilter);
}
public Void visitDynamicParam(RexDynamicParam dynamicParam) {
return pushVariable();
}
public Void visitRangeRef(RexRangeRef rangeRef) {
return pushVariable();
}
public Void visitFieldAccess(RexFieldAccess fieldAccess) {
return pushVariable();
}
}