Restrict size of automatically broadcasted tables

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -28,6 +28,7 @@
 import javax.inject.Inject;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.stream.Stream;
 
@@ -53,6 +54,7 @@ public final class SystemSessionProperties
 {
     public static final String OPTIMIZE_HASH_GENERATION = "optimize_hash_generation";
     public static final String JOIN_DISTRIBUTION_TYPE = "join_distribution_type";
+    public static final String JOIN_MAX_BROADCAST_TABLE_SIZE = "join_max_broadcast_table_size";
     public static final String DISTRIBUTED_JOIN = "distributed_join";
     public static final String DISTRIBUTED_INDEX_JOIN = "distributed_index_join";
     public static final String HASH_PARTITION_COUNT = "hash_partition_count";
@@ -152,6 +154,15 @@ public SystemSessionProperties(
                         false,
                         value -> JoinDistributionType.valueOf(((String) value).toUpperCase()),
                         JoinDistributionType::name),
+                new PropertyMetadata<>(
+                        JOIN_MAX_BROADCAST_TABLE_SIZE,
+                        "Maximum estimated size of a table that can be broadcast for JOIN.",
+                        VARCHAR,
+                        DataSize.class,
+                        featuresConfig.getJoinMaxBroadcastTableSize(),
+                        true,
+                        value -> DataSize.valueOf((String) value),
+                        DataSize::toString),
                 booleanProperty(
                         DISTRIBUTED_INDEX_JOIN,
                         "Distribute index joins on join keys instead of executing inline",
@@ -534,6 +545,11 @@ public static JoinDistributionType getJoinDistributionType(Session session)
         return session.getSystemProperty(JOIN_DISTRIBUTION_TYPE, JoinDistributionType.class);
     }
 
+    public static Optional<DataSize> getJoinMaxBroadcastTableSize(Session session)
+    {
+        return Optional.ofNullable(session.getSystemProperty(JOIN_MAX_BROADCAST_TABLE_SIZE, DataSize.class));
+    }
+
     public static boolean isDistributedIndexJoinEnabled(Session session)
     {
         return session.getSystemProperty(DISTRIBUTED_INDEX_JOIN, Boolean.class);
@@ -811,6 +827,7 @@ public static OptionalInt getMaxDriversPerTask(Session session)
         }
         return OptionalInt.of(value);
     }
+
     private static int validateValueIsPowerOfTwo(Object value, String property)
     {
         int intValue = ((Number) requireNonNull(value, "value is null")).intValue();

presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java

@@ -65,6 +65,7 @@ public class FeaturesConfig
     private double networkCostWeight = 15;
     private boolean distributedIndexJoinsEnabled;
     private JoinDistributionType joinDistributionType = PARTITIONED;
+    private DataSize joinMaxBroadcastTableSize;
     private boolean colocatedJoinsEnabled;
     private boolean groupedExecutionForAggregationEnabled;
     private int concurrentLifespansPerTask;
@@ -290,6 +291,18 @@ public FeaturesConfig setJoinDistributionType(JoinDistributionType joinDistribut
         return this;
     }
 
+    public DataSize getJoinMaxBroadcastTableSize()
+    {
+        return joinMaxBroadcastTableSize;
+    }
+
+    @Config("join-max-broadcast-table-size")
+    public FeaturesConfig setJoinMaxBroadcastTableSize(DataSize joinMaxBroadcastTableSize)
+    {
+        this.joinMaxBroadcastTableSize = joinMaxBroadcastTableSize;
+        return this;
+    }
+
     public boolean isGroupedExecutionForAggregationEnabled()
     {
         return groupedExecutionForAggregationEnabled;

presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/DetermineJoinDistributionType.java

@@ -17,18 +17,22 @@
 import com.facebook.presto.cost.CostComparator;
 import com.facebook.presto.cost.CostProvider;
 import com.facebook.presto.cost.PlanNodeCostEstimate;
+import com.facebook.presto.cost.PlanNodeStatsEstimate;
 import com.facebook.presto.matching.Captures;
 import com.facebook.presto.matching.Pattern;
 import com.facebook.presto.sql.analyzer.FeaturesConfig.JoinDistributionType;
 import com.facebook.presto.sql.planner.iterative.Rule;
 import com.facebook.presto.sql.planner.plan.JoinNode;
 import com.facebook.presto.sql.planner.plan.PlanNode;
 import com.google.common.collect.Ordering;
+import io.airlift.units.DataSize;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 import static com.facebook.presto.SystemSessionProperties.getJoinDistributionType;
+import static com.facebook.presto.SystemSessionProperties.getJoinMaxBroadcastTableSize;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.JoinDistributionType.AUTOMATIC;
 import static com.facebook.presto.sql.planner.optimizations.QueryCardinalityUtil.isAtMostScalar;
 import static com.facebook.presto.sql.planner.plan.JoinNode.DistributionType.PARTITIONED;
@@ -68,26 +72,31 @@ public Result apply(JoinNode joinNode, Captures captures, Context context)
         return Result.ofPlanNode(getSyntacticOrderJoin(joinNode, context, joinDistributionType));
     }
 
-    private PlanNode getCostBasedJoin(JoinNode joinNode, Context context)
+    public static boolean canReplicate(JoinNode joinNode, Context context)
     {
-        CostProvider costProvider = context.getCostProvider();
-        List<PlanNodeWithCost> possibleJoinNodes = new ArrayList<>();
-
-        if (!mustPartition(joinNode)) {
-            possibleJoinNodes.add(getJoinNodeWithCost(costProvider, joinNode.withDistributionType(REPLICATED)));
-        }
-        if (!mustReplicate(joinNode, context)) {
-            possibleJoinNodes.add(getJoinNodeWithCost(costProvider, joinNode.withDistributionType(PARTITIONED)));
+        JoinDistributionType joinDistributionType = getJoinDistributionType(context.getSession());
+        if (!joinDistributionType.canReplicate()) {
+            return false;
         }
 
-        JoinNode flipped = joinNode.flipChildren();
-        if (!mustPartition(flipped)) {
-            possibleJoinNodes.add(getJoinNodeWithCost(costProvider, flipped.withDistributionType(REPLICATED)));
-        }
-        if (!mustReplicate(flipped, context)) {
-            possibleJoinNodes.add(getJoinNodeWithCost(costProvider, flipped.withDistributionType(PARTITIONED)));
+        Optional<DataSize> joinMaxBroadcastTableSize = getJoinMaxBroadcastTableSize(context.getSession());
+        if (!joinMaxBroadcastTableSize.isPresent()) {
+            return true;
         }
 
+        PlanNode buildSide = joinNode.getRight();
+        PlanNodeStatsEstimate buildSideStatsEstimate = context.getStatsProvider().getStats(buildSide);
+        double buildSideSizeInBytes = buildSideStatsEstimate.getOutputSizeInBytes(buildSide.getOutputSymbols(), context.getSymbolAllocator().getTypes());
+        return buildSideSizeInBytes <= joinMaxBroadcastTableSize.get().toBytes();
+    }
+
+    private PlanNode getCostBasedJoin(JoinNode joinNode, Context context)
+    {
+        List<PlanNodeWithCost> possibleJoinNodes = new ArrayList<>();
+
+        addJoinsWithDifferentDistributions(joinNode, possibleJoinNodes, context);
+        addJoinsWithDifferentDistributions(joinNode.flipChildren(), possibleJoinNodes, context);
+
         if (possibleJoinNodes.stream().anyMatch(result -> result.getCost().hasUnknownComponents()) || possibleJoinNodes.isEmpty()) {
             return getSyntacticOrderJoin(joinNode, context, AUTOMATIC);
         }
@@ -97,6 +106,16 @@ private PlanNode getCostBasedJoin(JoinNode joinNode, Context context)
         return planNodeOrderings.min(possibleJoinNodes).getPlanNode();
     }
 
+    private void addJoinsWithDifferentDistributions(JoinNode joinNode, List<PlanNodeWithCost> possibleJoinNodes, Context context)
+    {
+        if (!mustPartition(joinNode) && canReplicate(joinNode, context)) {
+            possibleJoinNodes.add(getJoinNodeWithCost(context.getCostProvider(), joinNode.withDistributionType(REPLICATED)));
+        }
+        if (!mustReplicate(joinNode, context)) {
+            possibleJoinNodes.add(getJoinNodeWithCost(context.getCostProvider(), joinNode.withDistributionType(PARTITIONED)));
+        }
+    }
+
     private PlanNode getSyntacticOrderJoin(JoinNode joinNode, Context context, JoinDistributionType joinDistributionType)
     {
         if (mustPartition(joinNode)) {

presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ReorderJoins.java

@@ -62,6 +62,7 @@
 import static com.facebook.presto.sql.planner.DeterminismEvaluator.isDeterministic;
 import static com.facebook.presto.sql.planner.EqualityInference.createEqualityInference;
 import static com.facebook.presto.sql.planner.EqualityInference.nonInferrableConjuncts;
+import static com.facebook.presto.sql.planner.iterative.rule.DetermineJoinDistributionType.canReplicate;
 import static com.facebook.presto.sql.planner.iterative.rule.ReorderJoins.JoinEnumerationResult.INFINITE_COST_RESULT;
 import static com.facebook.presto.sql.planner.iterative.rule.ReorderJoins.JoinEnumerationResult.UNKNOWN_COST_RESULT;
 import static com.facebook.presto.sql.planner.iterative.rule.ReorderJoins.MultiJoinNode.toMultiJoinNode;
@@ -374,7 +375,7 @@ private JoinEnumerationResult setJoinNodeProperties(JoinNode joinNode)
                 possibleJoinNodes.add(createJoinEnumerationResult(joinNode.withDistributionType(PARTITIONED)));
                 possibleJoinNodes.add(createJoinEnumerationResult(joinNode.flipChildren().withDistributionType(PARTITIONED)));
             }
-            if (joinDistributionType.canReplicate()) {
+            if (canReplicate(joinNode, context)) {
                 possibleJoinNodes.add(createJoinEnumerationResult(joinNode.withDistributionType(REPLICATED)));
                 possibleJoinNodes.add(createJoinEnumerationResult(joinNode.flipChildren().withDistributionType(REPLICATED)));
             }

presto-main/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java

@@ -53,6 +53,7 @@ public void testDefaults()
                 .setNetworkCostWeight(15)
                 .setDistributedIndexJoinsEnabled(false)
                 .setJoinDistributionType(PARTITIONED)
+                .setJoinMaxBroadcastTableSize(null)
                 .setGroupedExecutionForAggregationEnabled(false)
                 .setConcurrentLifespansPerTask(0)
                 .setFastInequalityJoins(true)
@@ -125,6 +126,7 @@ public void testExplicitPropertyMappings()
                 .put("deprecated.legacy-char-to-varchar-coercion", "true")
                 .put("distributed-index-joins-enabled", "true")
                 .put("join-distribution-type", "BROADCAST")
+                .put("join-max-broadcast-table-size", "42GB")
                 .put("grouped-execution-for-aggregation-enabled", "true")
                 .put("concurrent-lifespans-per-task", "1")
                 .put("fast-inequality-joins", "false")
@@ -179,6 +181,7 @@ public void testExplicitPropertyMappings()
                 .setEnableStatsCalculator(false)
                 .setDistributedIndexJoinsEnabled(true)
                 .setJoinDistributionType(BROADCAST)
+                .setJoinMaxBroadcastTableSize(new DataSize(42, GIGABYTE))
                 .setGroupedExecutionForAggregationEnabled(true)
                 .setConcurrentLifespansPerTask(1)
                 .setFastInequalityJoins(false)