Support inserting stats for plan unit tests
@@ -18,6 +18,7 @@ import com.facebook.presto.cost.CachingCostProvider;
import com.facebook.presto.cost.CachingStatsProvider;
import com.facebook.presto.cost.CostCalculator;
import com.facebook.presto.cost.CostProvider;
+import com.facebook.presto.cost.PlanNodeStatsEstimate;
import com.facebook.presto.cost.StatsCalculator;
import com.facebook.presto.cost.StatsProvider;
import com.facebook.presto.matching.Match;
@@ -33,9 +34,12 @@ import com.facebook.presto.sql.planner.iterative.Memo;
import com.facebook.presto.sql.planner.iterative.PlanNodeMatcher;
import com.facebook.presto.sql.planner.iterative.Rule;
import com.facebook.presto.sql.planner.plan.PlanNode;
+import com.facebook.presto.sql.planner.plan.PlanNodeId;
import com.facebook.presto.transaction.TransactionManager;
import com.google.common.collect.ImmutableSet;
+import java.util.HashMap;
+import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Stream;
@@ -51,7 +55,7 @@ import static org.testng.Assert.fail;
public class RuleAssert
{
private final Metadata metadata;
- private final StatsCalculator statsCalculator;
+ private TestingStatsCalculator statsCalculator;
private final CostCalculator costCalculator;
private Session session;
private final Rule<?> rule;
@@ -66,7 +70,7 @@ public class RuleAssert
public RuleAssert(Metadata metadata, StatsCalculator statsCalculator, CostCalculator costCalculator, Session session, Rule rule, TransactionManager transactionManager, AccessControl accessControl)
{
this.metadata = metadata;
- this.statsCalculator = statsCalculator;
+ this.statsCalculator = new TestingStatsCalculator(statsCalculator);
this.costCalculator = costCalculator;
this.session = session;
this.rule = rule;
@@ -87,6 +91,12 @@ public class RuleAssert
return this;
}
+ public RuleAssert overrideStats(String nodeId, PlanNodeStatsEstimate nodeStats)
+ {
+ statsCalculator.setNodeStats(new PlanNodeId(nodeId), nodeStats);
+ return this;
+ }
+
public RuleAssert on(Function<PlanBuilder, PlanNode> planProvider)
{
checkArgument(plan == null, "plan has already been set");
@@ -259,4 +269,30 @@ public class RuleAssert
return result.getTransformedPlan().orElseThrow(() -> new IllegalStateException("Rule did not produce transformed plan"));
}
}
+
+ private static class TestingStatsCalculator
+ implements StatsCalculator
+ {
+ private final StatsCalculator delegate;
+ private final Map<PlanNodeId, PlanNodeStatsEstimate> stats = new HashMap<>();
+
+ TestingStatsCalculator(StatsCalculator delegate)
+ {
+ this.delegate = requireNonNull(delegate, "delegate is null");
+ }
+
+ @Override
+ public PlanNodeStatsEstimate calculateStats(PlanNode node, StatsProvider sourceStats, Lookup lookup, Session session, TypeProvider types)
+ {
+ if (stats.containsKey(node.getId())) {
+ return stats.get(node.getId());
+ }
+ return delegate.calculateStats(node, sourceStats, lookup, session, types);
+ }
+
+ public void setNodeStats(PlanNodeId nodeId, PlanNodeStatsEstimate nodeStats)
+ {
+ stats.put(nodeId, nodeStats);
+ }
+ }
}
