This is an automated email from the ASF dual-hosted git repository.
abhishek pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new d732de9948 Allow adding calcite rules from extensions (#12715)
d732de9948 is described below
commit d732de99486988a4311e1bf5621d4a6d5511b950
Author: Rohan Garg <[email protected]>
AuthorDate: Wed Jul 6 19:32:35 2022 +0530
Allow adding calcite rules from extensions (#12715)
* Allow adding calcite rules from extensions
* fixup! Allow adding calcite rules from extensions
* Move Rules to CalciteRulesManager
* fixup! Move Rules to CalciteRulesManager
---
.../apache/druid/benchmark/query/SqlBenchmark.java | 5 ++-
.../benchmark/query/SqlExpressionBenchmark.java | 5 ++-
.../benchmark/query/SqlVsNativeBenchmark.java | 5 ++-
.../sql/calcite/planner/CalcitePlannerModule.java | 3 ++
.../{Rules.java => CalciteRulesManager.java} | 52 ++++++++++++++--------
.../druid/sql/calcite/planner/DruidPlanner.java | 4 +-
.../druid/sql/calcite/planner/PlannerFactory.java | 7 ++-
.../calcite/rule/ExtensionCalciteRuleProvider.java | 35 +++++++++++++++
.../druid/sql/avatica/DruidAvaticaHandlerTest.java | 8 +++-
.../druid/sql/avatica/DruidStatementTest.java | 5 ++-
.../druid/sql/calcite/BaseCalciteQueryTest.java | 5 ++-
.../calcite/SqlVectorizedExpressionSanityTest.java | 5 ++-
.../calcite/planner/CalcitePlannerModuleTest.java | 44 +++++++++++++++++-
.../org/apache/druid/sql/http/SqlResourceTest.java | 5 ++-
14 files changed, 155 insertions(+), 33 deletions(-)
diff --git
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java
index a64e4710df..916403a1bf 100644
---
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java
+++
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java
@@ -21,6 +21,7 @@ package org.apache.druid.benchmark.query;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.guava.Sequence;
@@ -39,6 +40,7 @@ import org.apache.druid.server.security.AuthTestUtils;
import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
import
org.apache.druid.sql.calcite.expression.builtin.QueryLookupOperatorConversion;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
import org.apache.druid.sql.calcite.planner.DruidPlanner;
@@ -431,7 +433,8 @@ public class SqlBenchmark
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
}
diff --git
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
index b991b8922a..6f3c8bedb6 100644
---
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
+++
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
@@ -21,6 +21,7 @@ package org.apache.druid.benchmark.query;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.guava.Sequence;
@@ -37,6 +38,7 @@ import org.apache.druid.segment.generator.SegmentGenerator;
import org.apache.druid.server.QueryStackTests;
import org.apache.druid.server.security.AuthTestUtils;
import org.apache.druid.sql.calcite.SqlVectorizedExpressionSanityTest;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.DruidPlanner;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
@@ -317,7 +319,8 @@ public class SqlExpressionBenchmark
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
try {
diff --git
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java
index 8b64ea4c64..20db919354 100644
---
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java
+++
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java
@@ -19,6 +19,7 @@
package org.apache.druid.benchmark.query;
+import com.google.common.collect.ImmutableSet;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.granularity.Granularities;
@@ -38,6 +39,7 @@ import org.apache.druid.segment.generator.GeneratorSchemaInfo;
import org.apache.druid.segment.generator.SegmentGenerator;
import org.apache.druid.server.QueryStackTests;
import org.apache.druid.server.security.AuthTestUtils;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.DruidPlanner;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
import org.apache.druid.sql.calcite.planner.PlannerFactory;
@@ -121,7 +123,8 @@ public class SqlVsNativeBenchmark
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
groupByQuery = GroupByQuery
.builder()
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModule.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModule.java
index 16aaa3a1b7..7e67ec5020 100644
---
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModule.java
+++
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModule.java
@@ -22,9 +22,11 @@ package org.apache.druid.sql.calcite.planner;
import com.google.inject.Binder;
import com.google.inject.Key;
import com.google.inject.Module;
+import com.google.inject.multibindings.Multibinder;
import org.apache.druid.guice.JsonConfigProvider;
import org.apache.druid.guice.LazySingleton;
import org.apache.druid.guice.PolyBind;
+import org.apache.druid.sql.calcite.rule.ExtensionCalciteRuleProvider;
import org.apache.druid.sql.calcite.run.NativeQueryMakerFactory;
import org.apache.druid.sql.calcite.run.QueryMakerFactory;
@@ -54,5 +56,6 @@ public class CalcitePlannerModule implements Module
binder.bind(PlannerFactory.class).in(LazySingleton.class);
binder.bind(DruidOperatorTable.class).in(LazySingleton.class);
+ Multibinder.newSetBinder(binder, ExtensionCalciteRuleProvider.class);
}
}
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/Rules.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java
similarity index 88%
rename from sql/src/main/java/org/apache/druid/sql/calcite/planner/Rules.java
rename to
sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java
index 8169a244c0..088a9cb861 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/Rules.java
+++
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java
@@ -20,6 +20,7 @@
package org.apache.druid.sql.calcite.planner;
import com.google.common.collect.ImmutableList;
+import com.google.inject.Inject;
import org.apache.calcite.interpreter.Bindables;
import org.apache.calcite.plan.RelOptLattice;
import org.apache.calcite.plan.RelOptMaterialization;
@@ -78,13 +79,15 @@ import
org.apache.druid.sql.calcite.rule.DruidLogicalValuesRule;
import org.apache.druid.sql.calcite.rule.DruidRelToDruidRule;
import org.apache.druid.sql.calcite.rule.DruidRules;
import org.apache.druid.sql.calcite.rule.DruidTableScanRule;
+import org.apache.druid.sql.calcite.rule.ExtensionCalciteRuleProvider;
import org.apache.druid.sql.calcite.rule.FilterJoinExcludePushToChildRule;
import org.apache.druid.sql.calcite.rule.ProjectAggregatePruneUnusedCallRule;
import org.apache.druid.sql.calcite.rule.SortCollapseRule;
import java.util.List;
+import java.util.Set;
-public class Rules
+public class CalciteRulesManager
{
public static final int DRUID_CONVENTION_RULES = 0;
public static final int BINDABLE_CONVENTION_RULES = 1;
@@ -95,7 +98,7 @@ public class Rules
// Calcite 1.23.0 fixes this issue by not consider expression as reduced if
this case happens. However, while
// we are still using Calcite 1.21.0, a workaround is to limit the number of
pattern matches to avoid infinite loop.
private static final String HEP_DEFAULT_MATCH_LIMIT_CONFIG_STRING =
"druid.sql.planner.hepMatchLimit";
- private static final int HEP_DEFAULT_MATCH_LIMIT = Integer.valueOf(
+ private final int HEP_DEFAULT_MATCH_LIMIT = Integer.valueOf(
System.getProperty(HEP_DEFAULT_MATCH_LIMIT_CONFIG_STRING, "1200")
);
@@ -107,7 +110,7 @@ public class Rules
// functions).
// 3) JoinCommuteRule (we don't support reordering joins yet).
// 4) JoinPushThroughJoinRule (we don't support reordering joins yet).
- private static final List<RelOptRule> BASE_RULES =
+ private final List<RelOptRule> BASE_RULES =
ImmutableList.of(
AggregateStarTableRule.INSTANCE,
AggregateStarTableRule.INSTANCE2,
@@ -130,7 +133,7 @@ public class Rules
);
// Rules for scanning via Bindable, embedded directly in RelOptUtil's
registerDefaultRules.
- private static final List<RelOptRule> DEFAULT_BINDABLE_RULES =
+ private final List<RelOptRule> DEFAULT_BINDABLE_RULES =
ImmutableList.of(
Bindables.BINDABLE_TABLE_SCAN_RULE,
ProjectTableScanRule.INSTANCE,
@@ -142,7 +145,7 @@ public class Rules
// 1) ReduceExpressionsRule.JOIN_INSTANCE
// Removed by https://github.com/apache/druid/pull/9941 due to issue in
https://github.com/apache/druid/issues/9942
// TODO: Re-enable when https://github.com/apache/druid/issues/9942 is
fixed
- private static final List<RelOptRule> REDUCTION_RULES =
+ private final List<RelOptRule> REDUCTION_RULES =
ImmutableList.of(
ReduceExpressionsRule.PROJECT_INSTANCE,
ReduceExpressionsRule.FILTER_INSTANCE,
@@ -158,7 +161,7 @@ public class Rules
// Omit DateRangeRules due to
https://issues.apache.org/jira/browse/CALCITE-1601
// Omit UnionMergeRule since it isn't very effective given how Druid unions
currently operate and is potentially
// expensive in terms of planning time.
- private static final List<RelOptRule> ABSTRACT_RULES =
+ private final List<RelOptRule> ABSTRACT_RULES =
ImmutableList.of(
AggregateProjectPullUpConstantsRule.INSTANCE2,
UnionPullUpConstantsRule.INSTANCE,
@@ -186,7 +189,7 @@ public class Rules
// 4) FilterJoinRule.FILTER_ON_JOIN and FilterJoinRule.JOIN
// Removed by https://github.com/apache/druid/pull/9773 due to issue in
https://github.com/apache/druid/issues/9843
// TODO: Re-enable when https://github.com/apache/druid/issues/9843 is
fixed
- private static final List<RelOptRule> ABSTRACT_RELATIONAL_RULES =
+ private final List<RelOptRule> ABSTRACT_RELATIONAL_RULES =
ImmutableList.of(
AbstractConverter.ExpandConversionRule.INSTANCE,
AggregateRemoveRule.INSTANCE,
@@ -198,15 +201,21 @@ public class Rules
SortRemoveRule.INSTANCE
);
- private Rules()
+ private final Set<ExtensionCalciteRuleProvider>
extensionCalciteRuleProviderSet;
+
+ /**
+ * Manages the rules for planning of SQL queries via Calcite. Also provides
methods for extensions to provide custom
+ * rules for planning.
+ * @param extensionCalciteRuleProviderSet the set of custom rules coming
from extensions
+ */
+ @Inject
+ public CalciteRulesManager(final Set<ExtensionCalciteRuleProvider>
extensionCalciteRuleProviderSet)
{
- // No instantiation.
+ this.extensionCalciteRuleProviderSet = extensionCalciteRuleProviderSet;
}
- public static List<Program> programs(final PlannerContext plannerContext)
+ public List<Program> programs(final PlannerContext plannerContext)
{
-
-
// Program that pre-processes the tree before letting the full-on
VolcanoPlanner loose.
final Program preProgram =
Programs.sequence(
@@ -221,10 +230,12 @@ public class Rules
);
}
- private static Program buildHepProgram(Iterable<? extends RelOptRule> rules,
- boolean noDag,
- RelMetadataProvider metadataProvider,
- int matchLimit)
+ public Program buildHepProgram(
+ final Iterable<? extends RelOptRule> rules,
+ final boolean noDag,
+ final RelMetadataProvider metadataProvider,
+ final int matchLimit
+ )
{
final HepProgramBuilder builder = HepProgram.builder();
builder.addMatchLimit(matchLimit);
@@ -234,7 +245,7 @@ public class Rules
return Programs.of(builder.build(), noDag, metadataProvider);
}
- private static List<RelOptRule> druidConventionRuleSet(final PlannerContext
plannerContext)
+ public List<RelOptRule> druidConventionRuleSet(final PlannerContext
plannerContext)
{
final ImmutableList.Builder<RelOptRule> retVal = ImmutableList
.<RelOptRule>builder()
@@ -245,10 +256,13 @@ public class Rules
.add(new ExternalTableScanRule(plannerContext))
.addAll(DruidRules.rules(plannerContext));
+ for (ExtensionCalciteRuleProvider extensionCalciteRuleProvider :
extensionCalciteRuleProviderSet) {
+ retVal.add(extensionCalciteRuleProvider.getRule(plannerContext));
+ }
return retVal.build();
}
- private static List<RelOptRule> bindableConventionRuleSet(final
PlannerContext plannerContext)
+ public List<RelOptRule> bindableConventionRuleSet(final PlannerContext
plannerContext)
{
return ImmutableList.<RelOptRule>builder()
.addAll(baseRuleSet(plannerContext))
@@ -258,7 +272,7 @@ public class Rules
.build();
}
- private static List<RelOptRule> baseRuleSet(final PlannerContext
plannerContext)
+ public List<RelOptRule> baseRuleSet(final PlannerContext plannerContext)
{
final PlannerConfig plannerConfig = plannerContext.getPlannerConfig();
final ImmutableList.Builder<RelOptRule> rules = ImmutableList.builder();
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java
index ff139c7153..6167e82ad5 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java
@@ -315,7 +315,7 @@ public class DruidPlanner implements Closeable
RelNode parameterized =
rewriteRelDynamicParameters(possiblyLimitedRoot.rel);
final DruidRel<?> druidRel = (DruidRel<?>) planner.transform(
- Rules.DRUID_CONVENTION_RULES,
+ CalciteRulesManager.DRUID_CONVENTION_RULES,
planner.getEmptyTraitSet()
.replace(DruidConvention.instance())
.plus(root.collation),
@@ -362,7 +362,7 @@ public class DruidPlanner implements Closeable
) throws RelConversionException
{
BindableRel bindableRel = (BindableRel) planner.transform(
- Rules.BINDABLE_CONVENTION_RULES,
+ CalciteRulesManager.BINDABLE_CONVENTION_RULES,
planner.getEmptyTraitSet().replace(BindableConvention.INSTANCE).plus(root.collation),
root.rel
);
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
index f1d09abb36..12ed756190 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
@@ -72,6 +72,7 @@ public class PlannerFactory
private final ObjectMapper jsonMapper;
private final AuthorizerMapper authorizerMapper;
private final String druidSchemaName;
+ private final CalciteRulesManager calciteRuleManager;
@Inject
public PlannerFactory(
@@ -82,7 +83,8 @@ public class PlannerFactory
final PlannerConfig plannerConfig,
final AuthorizerMapper authorizerMapper,
final @Json ObjectMapper jsonMapper,
- final @DruidSchemaName String druidSchemaName
+ final @DruidSchemaName String druidSchemaName,
+ final CalciteRulesManager calciteRuleManager
)
{
this.rootSchema = rootSchema;
@@ -93,6 +95,7 @@ public class PlannerFactory
this.authorizerMapper = authorizerMapper;
this.jsonMapper = jsonMapper;
this.druidSchemaName = druidSchemaName;
+ this.calciteRuleManager = calciteRuleManager;
}
/**
@@ -163,7 +166,7 @@ public class PlannerFactory
.traitDefs(ConventionTraitDef.INSTANCE, RelCollationTraitDef.INSTANCE)
.convertletTable(new DruidConvertletTable(plannerContext))
.operatorTable(operatorTable)
- .programs(Rules.programs(plannerContext))
+ .programs(calciteRuleManager.programs(plannerContext))
.executor(new DruidRexExecutor(plannerContext))
.typeSystem(DruidTypeSystem.INSTANCE)
.defaultSchema(rootSchema.getSubSchema(druidSchemaName))
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/rule/ExtensionCalciteRuleProvider.java
b/sql/src/main/java/org/apache/druid/sql/calcite/rule/ExtensionCalciteRuleProvider.java
new file mode 100644
index 0000000000..2e2ffb2017
--- /dev/null
+++
b/sql/src/main/java/org/apache/druid/sql/calcite/rule/ExtensionCalciteRuleProvider.java
@@ -0,0 +1,35 @@
+/*
+ * 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.druid.sql.calcite.rule;
+
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.druid.guice.annotations.UnstableApi;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+
+/**
+ * This interface provides a way to supply custom calcite planning rules from
extensions. All the custom rules are
+ * collected and supplied to the planner which invokes {@link
ExtensionCalciteRuleProvider#getRule(PlannerContext)}
+ * for each of the rule provider per query.
+ */
+@UnstableApi
+public interface ExtensionCalciteRuleProvider
+{
+ RelOptRule getRule(PlannerContext plannerContext);
+}
diff --git
a/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java
b/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java
index f20fca2af3..5baa3de83a 100644
---
a/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java
+++
b/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java
@@ -67,6 +67,7 @@ import org.apache.druid.server.security.AuthTestUtils;
import org.apache.druid.server.security.AuthenticatorMapper;
import org.apache.druid.server.security.AuthorizerMapper;
import org.apache.druid.server.security.Escalator;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
@@ -219,6 +220,7 @@ public abstract class DruidAvaticaHandlerTest extends
CalciteTestBase
.in(LazySingleton.class);
binder.bind(QueryMakerFactory.class).to(NativeQueryMakerFactory.class);
binder.bind(new
TypeLiteral<Supplier<DefaultQueryConfig>>(){}).toInstance(Suppliers.ofInstance(new
DefaultQueryConfig(ImmutableMap.of())));
+ binder.bind(CalciteRulesManager.class).toInstance(new
CalciteRulesManager(ImmutableSet.of()));
}
}
)
@@ -906,7 +908,8 @@ public abstract class DruidAvaticaHandlerTest extends
CalciteTestBase
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
)
),
smallFrameConfig,
@@ -996,7 +999,8 @@ public abstract class DruidAvaticaHandlerTest extends
CalciteTestBase
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
)
),
smallFrameConfig,
diff --git
a/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java
b/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java
index dd643aaf14..c04f429ff1 100644
--- a/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java
@@ -20,6 +20,7 @@
package org.apache.druid.sql.avatica;
import com.google.common.base.Function;
+import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.Meta;
@@ -33,6 +34,7 @@ import org.apache.druid.server.QueryStackTests;
import org.apache.druid.server.security.AllowAllAuthenticator;
import org.apache.druid.server.security.AuthTestUtils;
import org.apache.druid.sql.SqlLifecycleFactory;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
import org.apache.druid.sql.calcite.planner.PlannerFactory;
@@ -98,7 +100,8 @@ public class DruidStatementTest extends CalciteTestBase
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
this.sqlLifecycleFactory =
CalciteTests.createSqlLifecycleFactory(plannerFactory);
}
diff --git
a/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
b/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
index e5b190d900..cf542d81db 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
@@ -26,6 +26,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import org.apache.calcite.plan.RelOptPlanner;
import org.apache.druid.annotations.UsedByJUnitParamsRunner;
import org.apache.druid.common.config.NullHandling;
@@ -84,6 +85,7 @@ import org.apache.druid.sql.SqlLifecycle;
import org.apache.druid.sql.SqlLifecycleFactory;
import org.apache.druid.sql.calcite.expression.DruidExpression;
import org.apache.druid.sql.calcite.external.ExternalDataSource;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
@@ -1020,7 +1022,8 @@ public class BaseCalciteQueryTest extends CalciteTestBase
plannerConfig,
authorizerMapper,
objectMapper,
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
final SqlLifecycleFactory sqlLifecycleFactory =
CalciteTests.createSqlLifecycleFactory(plannerFactory, authConfig);
diff --git
a/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java
b/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java
index 79ccebf9ab..1aca4f1adc 100644
---
a/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java
+++
b/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java
@@ -21,6 +21,7 @@ package org.apache.druid.sql.calcite;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.calcite.tools.RelConversionException;
import org.apache.calcite.tools.ValidationException;
@@ -40,6 +41,7 @@ import org.apache.druid.segment.generator.GeneratorSchemaInfo;
import org.apache.druid.segment.generator.SegmentGenerator;
import org.apache.druid.server.QueryStackTests;
import org.apache.druid.server.security.AuthTestUtils;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.DruidPlanner;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
@@ -147,7 +149,8 @@ public class SqlVectorizedExpressionSanityTest extends
InitializedNullHandlingTe
plannerConfig,
AuthTestUtils.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
}
diff --git
a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java
b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java
index b84ec33455..72278c19a8 100644
---
a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java
+++
b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java
@@ -25,15 +25,22 @@ import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.Scopes;
import com.google.inject.TypeLiteral;
+import com.google.inject.multibindings.Multibinder;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.logical.LogicalTableScan;
import org.apache.calcite.schema.Schema;
import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.jackson.DefaultObjectMapper;
import org.apache.druid.jackson.JacksonModule;
import org.apache.druid.math.expr.ExprMacroTable;
+import org.apache.druid.query.QueryContext;
import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.server.security.AuthorizerMapper;
import org.apache.druid.server.security.ResourceType;
import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
+import org.apache.druid.sql.calcite.rule.ExtensionCalciteRuleProvider;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.schema.DruidSchemaName;
import org.apache.druid.sql.calcite.schema.NamedSchema;
@@ -50,6 +57,9 @@ import javax.validation.Validation;
import javax.validation.Validator;
import java.util.Set;
+import static org.apache.calcite.plan.RelOptRule.any;
+import static org.apache.calcite.plan.RelOptRule.operand;
+
@RunWith(EasyMockRunner.class)
public class CalcitePlannerModuleTest extends CalciteTestBase
{
@@ -79,6 +89,7 @@ public class CalcitePlannerModuleTest extends CalciteTestBase
private CalcitePlannerModule target;
private Injector injector;
+ private RelOptRule customRule;
@Before
public void setUp()
@@ -93,6 +104,14 @@ public class CalcitePlannerModuleTest extends
CalciteTestBase
aggregators = ImmutableSet.of();
operatorConversions = ImmutableSet.of();
target = new CalcitePlannerModule();
+ customRule = new RelOptRule(operand(LogicalTableScan.class, any()),
"customRule")
+ {
+ @Override
+ public void onMatch(RelOptRuleCall call)
+ {
+
+ }
+ };
injector = Guice.createInjector(
new JacksonModule(),
binder -> {
@@ -106,7 +125,12 @@ public class CalcitePlannerModuleTest extends
CalciteTestBase
binder.bind(Key.get(new
TypeLiteral<Set<SqlOperatorConversion>>(){})).toInstance(operatorConversions);
binder.bind(DruidSchemaCatalog.class).toInstance(rootSchema);
},
- target
+ target,
+ binder -> {
+ Multibinder.newSetBinder(binder, ExtensionCalciteRuleProvider.class)
+ .addBinding()
+ .toInstance(plannerContext -> customRule);
+ }
);
}
@@ -138,4 +162,22 @@ public class CalcitePlannerModuleTest extends
CalciteTestBase
PlannerConfig plannerConfig = injector.getInstance(PlannerConfig.class);
Assert.assertNotNull(plannerConfig);
}
+
+ @Test
+ public void testExtensionCalciteRule()
+ {
+ PlannerContext context = PlannerContext.create(
+ "SELECT 1",
+ injector.getInstance(DruidOperatorTable.class),
+ macroTable,
+ new DefaultObjectMapper(),
+ injector.getInstance(PlannerConfig.class),
+ rootSchema,
+ new QueryContext()
+ );
+ boolean containsCustomRule =
injector.getInstance(CalciteRulesManager.class)
+ .druidConventionRuleSet(context)
+ .contains(customRule);
+ Assert.assertTrue(containsCustomRule);
+ }
}
diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java
b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java
index 77a7c7720b..59a8d0b3e1 100644
--- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java
@@ -26,6 +26,7 @@ import com.google.common.base.Strings;
import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
@@ -73,6 +74,7 @@ import org.apache.druid.sql.SqlLifecycle;
import org.apache.druid.sql.SqlLifecycleFactory;
import org.apache.druid.sql.SqlLifecycleManager;
import org.apache.druid.sql.SqlPlanningException.PlanningError;
+import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
import org.apache.druid.sql.calcite.planner.PlannerContext;
@@ -242,7 +244,8 @@ public class SqlResourceTest extends CalciteTestBase
plannerConfig,
CalciteTests.TEST_AUTHORIZER_MAPPER,
CalciteTests.getJsonMapper(),
- CalciteTests.DRUID_SCHEMA_NAME
+ CalciteTests.DRUID_SCHEMA_NAME,
+ new CalciteRulesManager(ImmutableSet.of())
);
lifecycleManager = new SqlLifecycleManager()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
