Implemented bloat configuration from QueryContext

sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java

@@ -83,10 +83,8 @@ public class CalciteRulesManager
   private static final int HEP_DEFAULT_MATCH_LIMIT = Integer.parseInt(
       System.getProperty(HEP_DEFAULT_MATCH_LIMIT_CONFIG_STRING, "1200")
   );
-  private static final String BLOAT_PROPERTY = "druid.sql.planner.bloat";
-  private static final int BLOAT = Integer.parseInt(
-      System.getProperty(BLOAT_PROPERTY, "1000")
-  );
+  public static final String BLOAT_PROPERTY = "sql.planner.bloat";
+  public static final int DEFAULT_BLOAT = 1000;
 
   /**
    * Rules from {@link org.apache.calcite.plan.RelOptRules#BASE_RULES}, minus:
@@ -100,16 +98,14 @@ public class CalciteRulesManager
    * and {@link CoreRules#FILTER_INTO_JOIN}, which are part of {@link #FANCY_JOIN_RULES}.
    * 4) {@link CoreRules#PROJECT_FILTER_TRANSPOSE} because PartialDruidQuery would like to have the Project on top of the Filter -
    * this rule could create a lot of non-useful plans.
-   *
-   * {@link CoreRules#PROJECT_MERGE} includes configurable bloat parameter, as a workaround for Calcite exception
+   * 5) {@link CoreRules#PROJECT_MERGE} added later with bloat parameter configured from query context as a workaround for Calcite exception
    * (there are not enough rules to produce a node with desired properties) thrown while running complex sql-queries with
-   * big amount of subqueries. `druid.sql.planner.bloat` should be set in broker's `jvm.config` file.
+   * big amount of subqueries.
    */
   private static final List<RelOptRule> BASE_RULES =
       ImmutableList.of(
           CoreRules.AGGREGATE_STAR_TABLE,
           CoreRules.AGGREGATE_PROJECT_STAR_TABLE,
-          ProjectMergeRule.Config.DEFAULT.withBloat(BLOAT).toRule(),
           CoreRules.FILTER_SCAN,
           CoreRules.FILTER_PROJECT_TRANSPOSE,
           CoreRules.JOIN_PUSH_EXPRESSIONS,
@@ -444,6 +440,15 @@ public List<RelOptRule> bindableConventionRuleSet(final PlannerContext plannerCo
                         .build();
   }
 
+  public List<RelOptRule> configurableRuleSet(PlannerContext plannerContext){
+    return ImmutableList.of(ProjectMergeRule.Config.DEFAULT.withBloat(getBloatProperty(plannerContext)).toRule());
+  }
+
+  private int getBloatProperty(PlannerContext plannerContext) {
+    final Integer bloat = plannerContext.queryContext().getInt(BLOAT_PROPERTY);
+    return (bloat != null) ? bloat : DEFAULT_BLOAT;
+  }
+
   public List<RelOptRule> baseRuleSet(final PlannerContext plannerContext)
   {
     final PlannerConfig plannerConfig = plannerContext.getPlannerConfig();
@@ -453,6 +458,7 @@ public List<RelOptRule> baseRuleSet(final PlannerContext plannerContext)
     rules.addAll(BASE_RULES);
     rules.addAll(ABSTRACT_RULES);
     rules.addAll(ABSTRACT_RELATIONAL_RULES);
+    rules.addAll(configurableRuleSet(plannerContext));
 
     if (plannerContext.getJoinAlgorithm().requiresSubquery()) {
       rules.addAll(FANCY_JOIN_RULES);

sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java

@@ -30,6 +30,7 @@
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.rel.logical.LogicalTableScan;
+import org.apache.calcite.rel.rules.ProjectMergeRule;
 import org.apache.calcite.schema.Schema;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.jackson.DefaultObjectMapper;
@@ -61,17 +62,21 @@
 import javax.validation.Validator;
 
 import java.util.Collections;
+import java.util.Optional;
 import java.util.Set;
 
 import static org.apache.calcite.plan.RelOptRule.any;
 import static org.apache.calcite.plan.RelOptRule.operand;
+import static org.apache.druid.sql.calcite.planner.CalciteRulesManager.BLOAT_PROPERTY;
+import static org.apache.druid.sql.calcite.planner.CalciteRulesManager.DEFAULT_BLOAT;
 
 @ExtendWith(EasyMockExtension.class)
 public class CalcitePlannerModuleTest extends CalciteTestBase
 {
   private static final String SCHEMA_1 = "SCHEMA_1";
   private static final String SCHEMA_2 = "SCHEMA_2";
   private static final String DRUID_SCHEMA_NAME = "DRUID_SCHEMA_NAME";
+  private static final int BLOAT = 1200;
 
   @Mock
   private NamedSchema druidSchema1;
@@ -204,4 +209,50 @@ public void testExtensionCalciteRule()
                                          .contains(customRule);
     Assert.assertTrue(containsCustomRule);
   }
+
+  @Test
+  public void testConfigurableBloat()
+  {
+    ObjectMapper mapper = new DefaultObjectMapper();
+    PlannerToolbox toolbox = new PlannerToolbox(
+            injector.getInstance(DruidOperatorTable.class),
+            macroTable,
+            mapper,
+            injector.getInstance(PlannerConfig.class),
+            rootSchema,
+            joinableFactoryWrapper,
+            CatalogResolver.NULL_RESOLVER,
+            "druid",
+            new CalciteRulesManager(ImmutableSet.of()),
+            CalciteTests.TEST_AUTHORIZER_MAPPER,
+            AuthConfig.newBuilder().build()
+    );
+
+    PlannerContext contextWithBloat = PlannerContext.create(
+            toolbox,
+            "SELECT 1",
+            new NativeSqlEngine(queryLifecycleFactory, mapper),
+            Collections.singletonMap(BLOAT_PROPERTY, BLOAT),
+            null
+    );
+
+    PlannerContext contextWithoutBloat = PlannerContext.create(
+            toolbox,
+            "SELECT 1",
+            new NativeSqlEngine(queryLifecycleFactory, mapper),
+            Collections.emptyMap(),
+            null
+    );
+
+    assertBloat(contextWithBloat, BLOAT);
+    assertBloat(contextWithoutBloat, DEFAULT_BLOAT);
+  }
+
+  private void assertBloat(PlannerContext context, int expectedBloat) {
+    Optional<ProjectMergeRule> firstProjectMergeRule = injector.getInstance(CalciteRulesManager.class).baseRuleSet(context).stream()
+            .filter(rule -> rule instanceof ProjectMergeRule)
+            .map(rule -> (ProjectMergeRule) rule)
+            .findAny();
+    Assert.assertTrue(firstProjectMergeRule.isPresent() && firstProjectMergeRule.get().config.bloat() == expectedBloat);
+  }
 }