DRILL-6589: Push transitive closure predicate(s) past aggregates

closes #1372

5 files changed, 71 insertions, 74 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRel.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRel.java
index 55cd7bfc2..5a7421b66 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRel.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRel.java
@@ -17,11 +17,15 @@
  */
 package org.apache.drill.exec.planner.logical;
 
-import java.util.List;
-
+import com.google.common.collect.Lists;
 import org.apache.calcite.linq4j.Ord;
+import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelOptCost;
 import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.type.SqlTypeName;
@@ -36,14 +40,8 @@ import org.apache.drill.common.logical.data.GroupingAggregate;
 import org.apache.drill.common.logical.data.LogicalOperator;
 import org.apache.drill.exec.planner.common.DrillAggregateRelBase;
 import org.apache.drill.exec.planner.torel.ConversionContext;
-import org.apache.calcite.rel.core.AggregateCall;
-import org.apache.calcite.rel.core.Aggregate;
-import org.apache.calcite.rel.InvalidRelException;
-import org.apache.calcite.rel.RelNode;
-import org.apache.calcite.plan.RelOptCluster;
-import org.apache.calcite.plan.RelTraitSet;
 
-import com.google.common.collect.Lists;
+import java.util.List;
 
 /**
  * Aggregation implemented in Drill.
@@ -51,17 +49,13 @@ import com.google.common.collect.Lists;
 public class DrillAggregateRel extends DrillAggregateRelBase implements DrillRel {
   /** Creates a DrillAggregateRel. */
   public DrillAggregateRel(RelOptCluster cluster, RelTraitSet traits, RelNode child, boolean indicator, ImmutableBitSet groupSet,
-      List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) throws InvalidRelException {
+      List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls)  {
     super(cluster, traits, child, indicator, groupSet, groupSets, aggCalls);
   }
 
   @Override
   public Aggregate copy(RelTraitSet traitSet, RelNode input, boolean indicator, ImmutableBitSet groupSet, List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
-    try {
-      return new DrillAggregateRel(getCluster(), traitSet, input, indicator, groupSet, groupSets, aggCalls);
-    } catch (InvalidRelException e) {
-      throw new AssertionError(e);
-    }
+    return new DrillAggregateRel(getCluster(), traitSet, input, indicator, groupSet, groupSets, aggCalls);
   }
 
   @Override
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRule.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRule.java
index 2b998b290..75f806f37 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRule.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillAggregateRule.java
@@ -17,15 +17,12 @@
  */
 package org.apache.drill.exec.planner.logical;
 
-import org.apache.calcite.rel.InvalidRelException;
-import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.plan.Convention;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.logical.LogicalAggregate;
-import org.apache.calcite.util.trace.CalciteTrace;
-import org.slf4j.Logger;
 
 /**
  * Rule that converts an {@link LogicalAggregate} to a {@link DrillAggregateRel}, implemented by a Drill "segment" operation
@@ -33,7 +30,6 @@ import org.slf4j.Logger;
  */
 public class DrillAggregateRule extends RelOptRule {
   public static final RelOptRule INSTANCE = new DrillAggregateRule();
-  protected static final Logger tracer = CalciteTrace.getPlannerTracer();
 
   private DrillAggregateRule() {
     super(
@@ -55,11 +51,7 @@ public class DrillAggregateRule extends RelOptRule {
 
     final RelTraitSet traits = aggregate.getTraitSet().plus(DrillRel.DRILL_LOGICAL);
     final RelNode convertedInput = convert(input, input.getTraitSet().plus(DrillRel.DRILL_LOGICAL).simplify());
-    try {
-      call.transformTo(new DrillAggregateRel(aggregate.getCluster(), traits, convertedInput, aggregate.indicator,
-          aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList()));
-    } catch (InvalidRelException e) {
-      tracer.warn(e.toString());
-    }
+    call.transformTo(new DrillAggregateRel(aggregate.getCluster(), traits, convertedInput, aggregate.indicator,
+        aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList()));
   }
 }
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillFilterAggregateTransposeRule.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillFilterAggregateTransposeRule.java
index 9f2088d9b..6bb409ee7 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillFilterAggregateTransposeRule.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillFilterAggregateTransposeRule.java
@@ -24,18 +24,22 @@ import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.core.RelFactories;
 import org.apache.calcite.rel.rules.FilterAggregateTransposeRule;
+import org.apache.calcite.tools.RelBuilderFactory;
 import org.apache.drill.exec.planner.DrillRelBuilder;
 
 public class DrillFilterAggregateTransposeRule extends FilterAggregateTransposeRule{
 
   // Since Calcite's default FilterAggregateTransposeRule would match Filter on top of Aggregate, it potentially will match Rels with mixed CONVENTION trait.
-  // Here override match method, such that the rule matchs with Rel in the same CONVENTION.
+  // Here override match method, such that the rule matches with Rel in the same CONVENTION.
 
-  public static final FilterAggregateTransposeRule INSTANCE = new DrillFilterAggregateTransposeRule();
+  public static final FilterAggregateTransposeRule INSTANCE = new DrillFilterAggregateTransposeRule(
+      DrillRelBuilder.proto(Contexts.of(RelFactories.DEFAULT_FILTER_FACTORY)));
 
-  private DrillFilterAggregateTransposeRule() {
-    super(Filter.class, DrillRelBuilder.proto(Contexts.of(RelFactories.DEFAULT_FILTER_FACTORY)),
-        Aggregate.class);
+  public static final FilterAggregateTransposeRule DRILL_LOGICAL_INSTANCE = new DrillFilterAggregateTransposeRule(
+      DrillRelBuilder.proto(DrillRelFactories.DRILL_LOGICAL_FILTER_FACTORY, DrillRelFactories.DRILL_LOGICAL_AGGREGATE_FACTORY));
+
+  private DrillFilterAggregateTransposeRule(RelBuilderFactory relBuilderFactory) {
+    super(Filter.class, relBuilderFactory, Aggregate.class);
   }
 
   @Override
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillReduceAggregatesRule.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillReduceAggregatesRule.java
index 55b295848..ac814fc4a 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillReduceAggregatesRule.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillReduceAggregatesRule.java
@@ -17,36 +17,19 @@
  */
 package org.apache.drill.exec.planner.logical;
 
-import java.math.BigDecimal;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import org.apache.calcite.plan.Convention;
 import org.apache.calcite.plan.RelOptCluster;
-import org.apache.calcite.rel.InvalidRelException;
-import org.apache.calcite.rel.core.Aggregate;
-import org.apache.calcite.rel.core.Window;
-import org.apache.calcite.rel.logical.LogicalAggregate;
-import org.apache.calcite.sql.SqlKind;
-import org.apache.calcite.sql.SqlOperator;
-import org.apache.calcite.sql.SqlOperatorBinding;
-import org.apache.calcite.sql.fun.SqlCountAggFunction;
-import org.apache.calcite.sql.type.SqlReturnTypeInference;
-import org.apache.calcite.sql.type.SqlTypeName;
-import org.apache.calcite.util.trace.CalciteTrace;
-import org.apache.drill.exec.planner.physical.PlannerSettings;
-import org.apache.drill.exec.planner.sql.DrillCalciteSqlAggFunctionWrapper;
-import org.apache.drill.exec.planner.sql.DrillSqlOperator;
-import org.apache.calcite.rel.core.AggregateCall;
-import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelOptRuleOperand;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.core.Window;
+import org.apache.calcite.rel.logical.LogicalAggregate;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rel.type.RelDataTypeField;
@@ -54,18 +37,31 @@ import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.fun.SqlAvgAggFunction;
+import org.apache.calcite.sql.fun.SqlCountAggFunction;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.fun.SqlSumAggFunction;
 import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.CompositeList;
 import org.apache.calcite.util.ImmutableIntList;
 import org.apache.calcite.util.Util;
-
-import com.google.common.collect.ImmutableList;
+import org.apache.drill.exec.planner.physical.PlannerSettings;
+import org.apache.drill.exec.planner.sql.DrillCalciteSqlAggFunctionWrapper;
+import org.apache.drill.exec.planner.sql.DrillSqlOperator;
 import org.apache.drill.exec.planner.sql.TypeInferenceUtils;
 import org.apache.drill.exec.planner.sql.parser.DrillCalciteWrapperUtility;
-import org.slf4j.Logger;
+
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 /**
  * Rule to reduce aggregates to simpler forms. Currently only AVG(x) to
@@ -693,7 +689,6 @@ public class DrillReduceAggregatesRule extends RelOptRule {
   }
 
   private static class DrillConvertSumToSumZero extends RelOptRule {
-    protected static final Logger tracer = CalciteTrace.getPlannerTracer();
 
     public DrillConvertSumToSumZero(RelOptRuleOperand operand) {
       super(operand, DrillRelFactories.LOGICAL_BUILDER, null);
@@ -744,18 +739,14 @@ public class DrillReduceAggregatesRule extends RelOptRule {
         }
       }
 
-      try {
-        call.transformTo(new DrillAggregateRel(
-            oldAggRel.getCluster(),
-            oldAggRel.getTraitSet(),
-            oldAggRel.getInput(),
-            oldAggRel.indicator,
-            oldAggRel.getGroupSet(),
-            oldAggRel.getGroupSets(),
-            newAggregateCalls));
-      } catch (InvalidRelException e) {
-        tracer.warn(e.toString());
-      }
+      call.transformTo(new DrillAggregateRel(
+          oldAggRel.getCluster(),
+          oldAggRel.getTraitSet(),
+          oldAggRel.getInput(),
+          oldAggRel.indicator,
+          oldAggRel.getGroupSet(),
+          oldAggRel.getGroupSets(),
+          newAggregateCalls));
     }
   }
 
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillRelFactories.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillRelFactories.java
index 59e4d06fc..fcf93fa6c 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillRelFactories.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillRelFactories.java
@@ -17,9 +17,11 @@
  */
 package org.apache.drill.exec.planner.logical;
 
+import com.google.common.collect.ImmutableList;
 import org.apache.calcite.plan.Contexts;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.core.CorrelationId;
 import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.core.RelFactories;
@@ -27,6 +29,7 @@ import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.tools.RelBuilderFactory;
+import org.apache.calcite.util.ImmutableBitSet;
 import org.apache.drill.exec.planner.DrillRelBuilder;
 
 import java.util.List;
@@ -48,7 +51,7 @@ import static org.apache.calcite.rel.core.RelFactories.DEFAULT_VALUES_FACTORY;
  */
 
 public class DrillRelFactories {
-
+  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(DrillRelFactories.class);
   public static final RelFactories.ProjectFactory DRILL_LOGICAL_PROJECT_FACTORY =
       new DrillProjectFactoryImpl();
 
@@ -57,6 +60,7 @@ public class DrillRelFactories {
 
   public static final RelFactories.JoinFactory DRILL_LOGICAL_JOIN_FACTORY = new DrillJoinFactoryImpl();
 
+  public static final RelFactories.AggregateFactory DRILL_LOGICAL_AGGREGATE_FACTORY = new DrillAggregateFactoryImpl();
   /**
    * A {@link RelBuilderFactory} that creates a {@link DrillRelBuilder} that will
    * create logical relational expressions for everything.
@@ -76,7 +80,7 @@ public class DrillRelFactories {
 
   /**
    * Implementation of {@link RelFactories.ProjectFactory} that returns a vanilla
-   * {@link org.apache.calcite.rel.logical.LogicalProject}.
+   * {@link DrillProjectRel}.
    */
   private static class DrillProjectFactoryImpl implements RelFactories.ProjectFactory {
     @Override
@@ -92,7 +96,7 @@ public class DrillRelFactories {
 
   /**
    * Implementation of {@link RelFactories.FilterFactory} that
-   * returns a vanilla {@link LogicalFilter}.
+   * returns a vanilla {@link DrillFilterRel}.
    */
   private static class DrillFilterFactoryImpl implements RelFactories.FilterFactory {
     @Override
@@ -103,7 +107,7 @@ public class DrillRelFactories {
 
   /**
    * Implementation of {@link RelFactories.JoinFactory} that returns a vanilla
-   * {@link org.apache.calcite.rel.logical.LogicalJoin}.
+   * {@link DrillJoinRel}.
    */
   private static class DrillJoinFactoryImpl implements RelFactories.JoinFactory {
 
@@ -122,4 +126,16 @@ public class DrillRelFactories {
     }
   }
 
+  /**
+   * Implementation of {@link RelFactories.AggregateFactory} that returns a vanilla
+   * {@link DrillAggregateRel}.
+   */
+  private static class DrillAggregateFactoryImpl implements RelFactories.AggregateFactory {
+
+    @Override
+    public RelNode createAggregate(RelNode input, boolean indicator, ImmutableBitSet groupSet,
+                                   ImmutableList<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
+      return new DrillAggregateRel(input.getCluster(), input.getTraitSet(), input, indicator, groupSet, groupSets, aggCalls);
+    }
+  }
 }