| /* |
| * 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.asterix.optimizer.rules; |
| |
| import java.util.ArrayList; |
| import java.util.List; |
| |
| import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil; |
| import org.apache.asterix.om.functions.BuiltinFunctions; |
| import org.apache.commons.lang3.mutable.Mutable; |
| import org.apache.commons.lang3.mutable.MutableObject; |
| import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException; |
| import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression; |
| import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator; |
| import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext; |
| import org.apache.hyracks.algebricks.core.algebra.base.LogicalExpressionTag; |
| import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag; |
| import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable; |
| import org.apache.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression; |
| import org.apache.hyracks.algebricks.core.algebra.expressions.AggregateFunctionCallExpression; |
| import org.apache.hyracks.algebricks.core.algebra.expressions.ConstantExpression; |
| import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression; |
| import org.apache.hyracks.algebricks.core.algebra.functions.FunctionIdentifier; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractBinaryJoinOperator; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.AggregateOperator; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.GroupByOperator; |
| import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities; |
| import org.apache.hyracks.algebricks.core.algebra.util.OperatorManipulationUtil; |
| import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule; |
| |
| /** |
| * Pushes aggregate functions into a stand alone aggregate operator (no group by). |
| */ |
| public class PushAggFuncIntoStandaloneAggregateRule implements IAlgebraicRewriteRule { |
| |
| @Override |
| public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context) |
| throws AlgebricksException { |
| return false; |
| } |
| |
| @Override |
| public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context) |
| throws AlgebricksException { |
| // Pattern to match: assign <-- aggregate <-- !(group-by) |
| AbstractLogicalOperator op = (AbstractLogicalOperator) opRef.getValue(); |
| if (op.getOperatorTag() != LogicalOperatorTag.ASSIGN) { |
| return false; |
| } |
| AssignOperator assignOp = (AssignOperator) op; |
| |
| Mutable<ILogicalOperator> opRef2 = op.getInputs().get(0); |
| AbstractLogicalOperator op2 = (AbstractLogicalOperator) opRef2.getValue(); |
| if (op2.getOperatorTag() == LogicalOperatorTag.AGGREGATE) { |
| AggregateOperator aggOp = (AggregateOperator) op2; |
| // Make sure the agg expr is a listify. |
| return pushAggregateFunction(assignOp, aggOp, context); |
| } else if (op2.getOperatorTag() == LogicalOperatorTag.INNERJOIN |
| || op2.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) { |
| AbstractBinaryJoinOperator join = (AbstractBinaryJoinOperator) op2; |
| // Tries to push aggregates through the join. |
| if (containsAggregate(assignOp.getExpressions()) && pushableThroughJoin(join)) { |
| return pushAggregateFunctionThroughJoin(assignOp, join, context); |
| } |
| } |
| return false; |
| } |
| |
| /** |
| * Recursively check whether the list of expressions contains an aggregate function. |
| * |
| * @return true if the list contains an aggregate function and false otherwise. |
| */ |
| private boolean containsAggregate(List<Mutable<ILogicalExpression>> exprRefs) { |
| for (Mutable<ILogicalExpression> exprRef : exprRefs) { |
| ILogicalExpression expr = exprRef.getValue(); |
| if (expr.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) { |
| continue; |
| } |
| AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr; |
| FunctionIdentifier funcIdent = BuiltinFunctions.getAggregateFunction(funcExpr.getFunctionIdentifier()); |
| if (funcIdent == null) { |
| // Recursively look in func args. |
| if (containsAggregate(funcExpr.getArguments())) { |
| return true; |
| } |
| } else { |
| // This is an aggregation function. |
| return true; |
| } |
| } |
| return false; |
| } |
| |
| /** |
| * Check whether the join is aggregate-pushable, that is, |
| * 1) the join condition is true; |
| * 2) each join branch produces only one tuple. |
| * |
| * @return true if pushable |
| */ |
| private boolean pushableThroughJoin(AbstractBinaryJoinOperator join) { |
| ILogicalExpression condition = join.getCondition().getValue(); |
| if (condition.equals(ConstantExpression.TRUE)) { |
| // Checks if the aggregation functions are pushable through the join |
| boolean pushable = true; |
| for (Mutable<ILogicalOperator> branchRef : join.getInputs()) { |
| AbstractLogicalOperator branch = (AbstractLogicalOperator) branchRef.getValue(); |
| if (branch.getOperatorTag() == LogicalOperatorTag.AGGREGATE) { |
| pushable &= true; |
| } else if (branch.getOperatorTag() == LogicalOperatorTag.INNERJOIN |
| || branch.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) { |
| AbstractBinaryJoinOperator childJoin = (AbstractBinaryJoinOperator) branch; |
| pushable &= pushableThroughJoin(childJoin); |
| } else { |
| pushable &= false; |
| } |
| } |
| return pushable; |
| } |
| return false; |
| } |
| |
| /** |
| * Does the actual push of aggregates for qualified joins. |
| * |
| * @param assignOp |
| * that contains aggregate function calls. |
| * @param context |
| * @throws AlgebricksException |
| */ |
| private boolean pushAggregateFunctionThroughJoin(AssignOperator assignOp, AbstractBinaryJoinOperator join, |
| IOptimizationContext context) throws AlgebricksException { |
| boolean applied = false; |
| for (Mutable<ILogicalOperator> branchRef : join.getInputs()) { |
| AbstractLogicalOperator branch = (AbstractLogicalOperator) branchRef.getValue(); |
| if (branch.getOperatorTag() == LogicalOperatorTag.AGGREGATE) { |
| AggregateOperator aggOp = (AggregateOperator) branch; |
| applied |= pushAggregateFunction(assignOp, aggOp, context); |
| } else if (branch.getOperatorTag() == LogicalOperatorTag.INNERJOIN |
| || branch.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) { |
| AbstractBinaryJoinOperator childJoin = (AbstractBinaryJoinOperator) branch; |
| applied |= pushAggregateFunctionThroughJoin(assignOp, childJoin, context); |
| } |
| } |
| return applied; |
| } |
| |
| private boolean pushAggregateFunction(AssignOperator assignOp, AggregateOperator aggOp, |
| IOptimizationContext context) throws AlgebricksException { |
| Mutable<ILogicalOperator> aggChilldOpRef = aggOp.getInputs().get(0); |
| AbstractLogicalOperator aggChildOp = (AbstractLogicalOperator) aggChilldOpRef.getValue(); |
| // If there's a group by below the agg, then we want to have the agg pushed into the group by |
| if (aggChildOp.getOperatorTag() == LogicalOperatorTag.GROUP |
| && !((GroupByOperator) aggChildOp).getNestedPlans().isEmpty()) { |
| return false; |
| } |
| |
| List<LogicalVariable> assignUsedVars = new ArrayList<>(); |
| VariableUtilities.getUsedVariables(assignOp, assignUsedVars); |
| |
| List<Mutable<ILogicalExpression>> assignScalarAggExprRefs = new ArrayList<>(); |
| List<LogicalVariable> aggAddVars = null; |
| List<Mutable<ILogicalExpression>> aggAddExprs = null; |
| |
| for (int i = 0, n = aggOp.getVariables().size(); i < n; i++) { |
| LogicalVariable aggVar = aggOp.getVariables().get(i); |
| Mutable<ILogicalExpression> aggExprRef = aggOp.getExpressions().get(i); |
| ILogicalExpression aggExpr = aggExprRef.getValue(); |
| if (aggExpr.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) { |
| continue; |
| } |
| AbstractFunctionCallExpression listifyCandidateExpr = (AbstractFunctionCallExpression) aggExpr; |
| if (listifyCandidateExpr.getFunctionIdentifier() != BuiltinFunctions.LISTIFY) { |
| continue; |
| } |
| if (!assignUsedVars.contains(aggVar)) { |
| continue; |
| } |
| assignScalarAggExprRefs.clear(); |
| findScalarAggFuncExprRef(assignOp.getExpressions(), aggVar, assignScalarAggExprRefs); |
| if (assignScalarAggExprRefs.isEmpty()) { |
| continue; |
| } |
| // perform rewrite |
| if (aggAddVars == null) { |
| aggAddVars = new ArrayList<>(); |
| aggAddExprs = new ArrayList<>(); |
| } |
| for (Mutable<ILogicalExpression> assignScalarAggExprRef : assignScalarAggExprRefs) { |
| AbstractFunctionCallExpression assignScalarAggExpr = |
| (AbstractFunctionCallExpression) assignScalarAggExprRef.getValue(); |
| FunctionIdentifier aggFuncIdent = |
| BuiltinFunctions.getAggregateFunction(assignScalarAggExpr.getFunctionIdentifier()); |
| |
| // Push the scalar aggregate function into the aggregate op. |
| int nArgs = assignScalarAggExpr.getArguments().size(); |
| List<Mutable<ILogicalExpression>> aggArgs = new ArrayList<>(nArgs); |
| aggArgs.add( |
| new MutableObject<>(listifyCandidateExpr.getArguments().get(0).getValue().cloneExpression())); |
| aggArgs.addAll(OperatorManipulationUtil |
| .cloneExpressions(assignScalarAggExpr.getArguments().subList(1, nArgs))); |
| AggregateFunctionCallExpression aggFuncExpr = |
| BuiltinFunctions.makeAggregateFunctionExpression(aggFuncIdent, aggArgs); |
| aggFuncExpr.setSourceLocation(assignScalarAggExpr.getSourceLocation()); |
| |
| LogicalVariable newVar = context.newVar(); |
| aggAddVars.add(newVar); |
| aggAddExprs.add(new MutableObject<>(aggFuncExpr)); |
| // The assign now just "renames" the variable to make sure the upstream plan still works. |
| VariableReferenceExpression newVarRef = new VariableReferenceExpression(newVar); |
| newVarRef.setSourceLocation(assignScalarAggExpr.getSourceLocation()); |
| assignScalarAggExprRef.setValue(newVarRef); |
| } |
| } |
| |
| if (aggAddVars == null) { |
| return false; |
| } |
| |
| // add new variables and expressions to the aggregate operator. |
| aggOp.getVariables().addAll(aggAddVars); |
| aggOp.getExpressions().addAll(aggAddExprs); |
| |
| // Note: we retain the original listify() call in the aggregate operator because |
| // the variable it is assigned to might be used upstream by other operators. |
| // If the variable is not used upstream then it'll later be removed |
| // by {@code RemoveUnusedAssignAndAggregateRule} |
| |
| context.computeAndSetTypeEnvironmentForOperator(aggOp); |
| context.computeAndSetTypeEnvironmentForOperator(assignOp); |
| return true; |
| } |
| |
| private void findScalarAggFuncExprRef(List<Mutable<ILogicalExpression>> exprRefs, LogicalVariable aggVar, |
| List<Mutable<ILogicalExpression>> outScalarAggExprRefs) { |
| for (Mutable<ILogicalExpression> exprRef : exprRefs) { |
| ILogicalExpression expr = exprRef.getValue(); |
| if (expr.getExpressionTag() == LogicalExpressionTag.FUNCTION_CALL) { |
| AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr; |
| FunctionIdentifier funcIdent = BuiltinFunctions.getAggregateFunction(funcExpr.getFunctionIdentifier()); |
| if (funcIdent != null |
| && aggVar.equals(SqlppVariableUtil.getVariable(funcExpr.getArguments().get(0).getValue()))) { |
| outScalarAggExprRefs.add(exprRef); |
| } else { |
| // Recursively look in func args. |
| findScalarAggFuncExprRef(funcExpr.getArguments(), aggVar, outScalarAggExprRefs); |
| } |
| } |
| } |
| } |
| } |
