This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new bc2f955bd11 Add RuleSetCustomizer SPI for multi-stage planner Calcite
rules (#18387)
bc2f955bd11 is described below
commit bc2f955bd11224ebacc0aab03761ab4e3cc1ebbd
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Fri May 22 13:40:49 2026 +0200
Add RuleSetCustomizer SPI for multi-stage planner Calcite rules (#18387)
---
.../pom.xml | 28 ++---
.../org/apache/pinot/query/planner/spi/Phase.java | 59 ++++++++++
.../pinot/query/planner/spi/RuleSetCustomizer.java | 89 +++++++++++++++
.../pinot/spi/plugin/PluginRealmExportTest.java | 78 +++++++++++++
pinot-query-planner/pom.xml | 8 ++
.../calcite/rel/rules/PinotQueryRuleSets.java | 80 ++++++-------
.../org/apache/pinot/query/QueryEnvironment.java | 56 ++++++---
.../planner/rules/DefaultRuleSetCustomizer.java | 70 ++++++++++++
.../pinot/query/planner/rules/PinotRuleSet.java | 127 +++++++++++++++++++++
.../query/planner/rules/PinotRuleSetTest.java | 97 ++++++++++++++++
.../org/apache/pinot/spi/plugin/PluginManager.java | 45 +++++++-
.../apache/pinot/spi/plugin/PluginManagerTest.java | 41 +++++++
pom.xml | 6 +
13 files changed, 700 insertions(+), 84 deletions(-)
diff --git a/pinot-query-planner/pom.xml b/pinot-query-planner-spi/pom.xml
similarity index 70%
copy from pinot-query-planner/pom.xml
copy to pinot-query-planner-spi/pom.xml
index 484fcd62524..3f88b02dfa5 100644
--- a/pinot-query-planner/pom.xml
+++ b/pinot-query-planner-spi/pom.xml
@@ -27,8 +27,8 @@
<groupId>org.apache.pinot</groupId>
<version>1.6.0-SNAPSHOT</version>
</parent>
- <artifactId>pinot-query-planner</artifactId>
- <name>Pinot Query Planner</name>
+ <artifactId>pinot-query-planner-spi</artifactId>
+ <name>Pinot Query Planner SPI</name>
<url>https://pinot.apache.org/</url>
<properties>
@@ -37,37 +37,23 @@
<dependencies>
<dependency>
- <groupId>org.apache.pinot</groupId>
- <artifactId>pinot-core</artifactId>
- </dependency>
- <dependency>
- <groupId>org.codehaus.janino</groupId>
- <artifactId>janino</artifactId>
- </dependency>
- <dependency>
- <groupId>org.codehaus.janino</groupId>
- <artifactId>commons-compiler</artifactId>
+ <groupId>org.apache.calcite</groupId>
+ <artifactId>calcite-core</artifactId>
</dependency>
<dependency>
- <groupId>org.immutables</groupId>
- <artifactId>value-annotations</artifactId>
+ <groupId>com.google.code.findbugs</groupId>
+ <artifactId>jsr305</artifactId>
</dependency>
<dependency>
<groupId>org.apache.pinot</groupId>
- <artifactId>pinot-core</artifactId>
+ <artifactId>pinot-spi</artifactId>
<scope>test</scope>
- <type>test-jar</type>
</dependency>
<dependency>
<groupId>org.testng</groupId>
<artifactId>testng</artifactId>
<scope>test</scope>
</dependency>
- <dependency>
- <groupId>org.mockito</groupId>
- <artifactId>mockito-core</artifactId>
- <scope>test</scope>
- </dependency>
</dependencies>
</project>
diff --git
a/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/Phase.java
b/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/Phase.java
new file mode 100644
index 00000000000..22923951c61
--- /dev/null
+++
b/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/Phase.java
@@ -0,0 +1,59 @@
+/**
+ * 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.pinot.query.planner.spi;
+
+
+/// HEP phases the multi-stage planner exposes to plugin
+/// [RuleSetCustomizer]s. Every phase corresponds to one slot in the rule
+/// programs built by `QueryEnvironment#getOptProgram` /
+/// `QueryEnvironment#getTraitProgram`.
+///
+/// **Stability**: this enum is append-only — new phases may be added at the
+/// end without breaking plugins compiled against an older version. Existing
+/// phases must never be reordered or removed.
+public enum Phase {
+ /// Basic logical-rewrite phase. HEP, depth-first.
+ /// OSS defaults: `DefaultRuleSetCustomizer.BASIC_RULES`.
+ BASIC,
+
+ /// Filter pushdown rules. The HEP program runs this phase twice (around the
+ /// project pushdown phase). OSS defaults:
+ /// `DefaultRuleSetCustomizer.FILTER_PUSHDOWN_RULES`.
+ FILTER_PUSHDOWN,
+
+ /// Project pushdown rules.
+ /// OSS defaults: `DefaultRuleSetCustomizer.PROJECT_PUSHDOWN_RULES`.
+ PROJECT_PUSHDOWN,
+
+ /// Top-down pruning rules.
+ /// OSS defaults: `DefaultRuleSetCustomizer.PRUNE_RULES`.
+ PRUNE,
+
+ /// Post-logical rules used when the physical optimizer is **not** enabled.
+ /// OSS defaults: `DefaultRuleSetCustomizer.POST_LOGICAL_RULES`. The list
+ /// includes `PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY` configured with
+ /// the rule's hard-coded default `fetchLimitThreshold`. Per-query overrides
+ /// (and broker-config overrides) are applied by `QueryEnvironment` swapping
+ /// the rule on a per-query copy of this list.
+ POST_LOGICAL,
+
+ /// Post-logical rules used when the physical optimizer **is** enabled.
+ /// OSS defaults: `DefaultRuleSetCustomizer.POST_LOGICAL_PHYSICAL_RULES`.
+ POST_LOGICAL_PHYSICAL_OPT
+}
diff --git
a/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/RuleSetCustomizer.java
b/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/RuleSetCustomizer.java
new file mode 100644
index 00000000000..a8ee936a1cb
--- /dev/null
+++
b/pinot-query-planner-spi/src/main/java/org/apache/pinot/query/planner/spi/RuleSetCustomizer.java
@@ -0,0 +1,89 @@
+/**
+ * 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.pinot.query.planner.spi;
+
+import java.util.List;
+import org.apache.calcite.plan.RelOptRule;
+
+
+/// Plugin SPI for customizing the multi-stage planner's Calcite rule sets.
+///
+/// Implementations are discovered via [java.util.ServiceLoader]. Declare your
+/// customizer in
+/// `META-INF/services/org.apache.pinot.query.planner.spi.RuleSetCustomizer`
+/// inside the plugin JAR. Implementations must have a public no-arg
constructor.
+///
+///
[org.apache.pinot.query.planner.rules.PinotRuleSet#loadFromServiceLoader()]
first discovers
+/// application-classpath customizers, then iterates every plugin classloader
registered with
+/// [org.apache.pinot.spi.plugin.PluginManager]. Plugin JARs must be loaded
+/// before
[org.apache.pinot.query.planner.rules.PinotRuleSet#defaultInstance()] is first
called;
+/// the singleton is initialized once and not refreshed.
+///
+///
+/// [org.apache.pinot.query.planner.rules.PinotRuleSet] invokes every
discovered customizer once
+/// per [Phase] at broker startup. The customizer receives the mutable
per-phase rule list
+/// and can append, remove, replace, or reorder rules using standard `List`
+/// operations. After every customizer has run, `PinotRuleSet` defensively
+/// copies each list to an immutable form; mutations to the supplied list
+/// after `customize` returns have no effect.
+///
+/// ### Iteration order
+///
+/// Customizers run in `ServiceLoader` iteration order — typically classpath
+/// order, which is JVM-dependent. The OSS defaults are themselves contributed
+/// by [org.apache.pinot.query.planner.rules.DefaultRuleSetCustomizer]; plugin
authors that depend
+/// on observing the OSS defaults should not assume a specific position. To
guarantee ordering
+/// across customizers, drop and re-add rules from your own `customize`
+/// implementation.
+///
+/// ### Example
+///
+/// ```java
+/// public class MyPluginRules implements RuleSetCustomizer {
+/// @Override public void customize(Phase phase, List<RelOptRule> rules) {
+/// if (phase == Phase.BASIC) {
+/// rules.add(MyOptimizationRule.INSTANCE);
+/// rules.removeIf(r -> "BadOldRule".equals(r.toString()));
+/// }
+/// }
+/// }
+/// ```
+///
+/// And
`META-INF/services/org.apache.pinot.query.planner.spi.RuleSetCustomizer`:
+/// ```
+/// com.example.MyPluginRules
+/// ```
+///
+/// ### Thread safety
+///
+/// Customizers are invoked single-threaded during
+/// [org.apache.pinot.query.planner.rules.PinotRuleSet] construction.
+/// Implementations need not be thread-safe and must not retain references to
+/// the supplied list — the list is defensively copied after every customizer
+/// has run. A customizer that throws aborts construction and propagates the
+/// exception out of
[org.apache.pinot.query.planner.rules.PinotRuleSet#loadFromServiceLoader()].
+public interface RuleSetCustomizer {
+
+ /// Modify the broker's rule list for the given phase. The list is mutable
+ /// during this call only; it is defensively copied after every customizer
+ /// has run. Implementations may simply return without mutating when the
+ /// phase is not one they care about. Implementations must not insert `null`
+ /// rules — `PinotRuleSet` will reject the resulting list.
+ void customize(Phase phase, List<RelOptRule> rules);
+}
diff --git
a/pinot-query-planner-spi/src/test/java/org/apache/pinot/spi/plugin/PluginRealmExportTest.java
b/pinot-query-planner-spi/src/test/java/org/apache/pinot/spi/plugin/PluginRealmExportTest.java
new file mode 100644
index 00000000000..5592437d16b
--- /dev/null
+++
b/pinot-query-planner-spi/src/test/java/org/apache/pinot/spi/plugin/PluginRealmExportTest.java
@@ -0,0 +1,78 @@
+/**
+ * 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.pinot.spi.plugin;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.nio.file.Files;
+import java.util.Set;
+import java.util.jar.JarOutputStream;
+import org.apache.pinot.query.planner.spi.RuleSetCustomizer;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/// Verifies that {@link PluginManager} exports the SPI packages from the pinot
+/// realm into each plugin realm, so plugin JARs can implement
+/// {@link RuleSetCustomizer} without bundling or shading those classes.
+///
+/// This test lives in the {@code org.apache.pinot.spi.plugin} package (in a
+/// different Maven module) to access the package-private {@link PluginManager}
+/// constructor.
+public class PluginRealmExportTest {
+
+ @Test
+ public void pluginRealmCanLoadRuleSetCustomizer()
+ throws Exception {
+ File tempDir = Files.createTempDirectory("pinot-spi-export-test").toFile();
+ try {
+ // Minimal new-style plugin: just pinot-plugin.properties + an empty JAR.
+ Files.createFile(tempDir.toPath().resolve("pinot-plugin.properties"));
+ try (JarOutputStream jos = new JarOutputStream(new FileOutputStream(new
File(tempDir, "dummy.jar")))) {
+ // intentionally empty
+ }
+
+ PluginManager pm = new PluginManager();
+ pm.load("spi-export-test-plugin", tempDir);
+
+ Set<ClassLoader> loaders = pm.getPluginClassLoaders();
+ Assert.assertEquals(loaders.size(), 1, "Exactly one plugin realm should
be registered");
+
+ ClassLoader pluginRealm = loaders.iterator().next();
+
+ // RuleSetCustomizer is in org.apache.pinot.query.planner.spi, which
PluginManager
+ // exports from the pinotRealm. Loading it from the plugin realm must
succeed.
+ Class<?> customizerClass =
pluginRealm.loadClass(RuleSetCustomizer.class.getName());
+ Assert.assertNotNull(customizerClass);
+ Assert.assertTrue(customizerClass.isInterface(), "RuleSetCustomizer must
be an interface");
+
+ // RelOptRule is in org.apache.calcite.plan, also exported by
PluginManager.
+ Class<?> relOptRuleClass =
pluginRealm.loadClass("org.apache.calcite.plan.RelOptRule");
+ Assert.assertNotNull(relOptRuleClass);
+ } finally {
+ File[] files = tempDir.listFiles();
+ if (files != null) {
+ for (File f : files) {
+ f.delete();
+ }
+ }
+ tempDir.delete();
+ }
+ }
+}
diff --git a/pinot-query-planner/pom.xml b/pinot-query-planner/pom.xml
index 484fcd62524..f55f78b7b4b 100644
--- a/pinot-query-planner/pom.xml
+++ b/pinot-query-planner/pom.xml
@@ -36,6 +36,10 @@
</properties>
<dependencies>
+ <dependency>
+ <groupId>org.apache.pinot</groupId>
+ <artifactId>pinot-query-planner-spi</artifactId>
+ </dependency>
<dependency>
<groupId>org.apache.pinot</groupId>
<artifactId>pinot-core</artifactId>
@@ -52,6 +56,10 @@
<groupId>org.immutables</groupId>
<artifactId>value-annotations</artifactId>
</dependency>
+ <dependency>
+ <groupId>com.google.auto.service</groupId>
+ <artifactId>auto-service-annotations</artifactId>
+ </dependency>
<dependency>
<groupId>org.apache.pinot</groupId>
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
index 60923dce79c..19909597302 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
@@ -49,13 +49,14 @@ import
org.apache.pinot.spi.utils.CommonConstants.Broker.PlannerRuleNames;
/**
- * Default rule sets for Pinot query
+ * Default rule sets for Pinot query.
* Defaultly disabled rules are defined in
* {@link
org.apache.pinot.spi.utils.CommonConstants.Broker#DEFAULT_DISABLED_RULES}
*
- * TODO: This class started as a list of constant rule sets, but since then we
have added dynamic rule generation
- * to it as well. We should probably refactor the class to make it easier to
understand, maintain and change the rules
- * based on contextual information like query options.
+ * <p>TODO: Rule lists may be consolidated into
+ * {@link org.apache.pinot.query.planner.rules.DefaultRuleSetCustomizer} in a
future refactor once
+ * the {@link org.apache.pinot.query.planner.spi.RuleSetCustomizer} SPI is the
established
+ * extension point for broker rule customization.
*/
public class PinotQueryRuleSets {
private PinotQueryRuleSets() {
@@ -218,6 +219,36 @@ public class PinotQueryRuleSets {
.withDescription(PlannerRuleNames.PRUNE_EMPTY_UNION).toRule()
);
+ /// Pinot specific post-logical rules used when the physical optimizer is
<b>not</b> enabled.
+ /// Includes {@link PinotSortExchangeCopyRule#SORT_EXCHANGE_COPY} with the
default fetch-limit
+ /// threshold. Per-query overrides are applied by {@code
QueryEnvironment.getTraitProgram},
+ /// which swaps the configured rule on a per-query copy of this list.
+ public static final List<RelOptRule> POST_LOGICAL_RULES = List.of(
+ // TODO: Merge the following 2 rules into a single rule
+ // add an extra exchange for sort
+ PinotSortExchangeNodeInsertRule.INSTANCE,
+ PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY,
+
+ PinotSingleValueAggregateRemoveRule.INSTANCE,
+ PinotJoinExchangeNodeInsertRule.INSTANCE,
+ PinotAggregateExchangeNodeInsertRule.SortProjectAggregate.INSTANCE,
+ PinotAggregateExchangeNodeInsertRule.SortAggregate.INSTANCE,
+ PinotAggregateExchangeNodeInsertRule.WithoutSort.INSTANCE,
+ PinotWindowSplitRule.INSTANCE,
+ PinotWindowExchangeNodeInsertRule.INSTANCE,
+ PinotSetOpExchangeNodeInsertRule.INSTANCE,
+
+ // apply dynamic broadcast rule after exchange is inserted
+ PinotJoinToDynamicBroadcastRule.INSTANCE,
+
+ // remove exchanges when there's duplicates
+ PinotExchangeEliminationRule.INSTANCE,
+
+ // Evaluate the Literal filter nodes
+ CoreRules.FILTER_REDUCE_EXPRESSIONS,
+ PinotTableScanConverterRule.INSTANCE
+ );
+
public static final List<RelOptRule> PINOT_POST_RULES_V2 = List.of(
PinotTableScanConverterRule.INSTANCE,
PinotLogicalAggregateRule.SortProjectAggregate.INSTANCE,
@@ -228,45 +259,4 @@ public class PinotQueryRuleSets {
CoreRules.FILTER_REDUCE_EXPRESSIONS
);
//@formatter:on
-
- /// Pinot specific rules that should be run AFTER all other rules
- public static List<RelOptRule> getPinotPostRules(int sortExchangeCopyLimit) {
-
- // copy exchanges down, this must be done after SortExchangeNodeInsertRule
- PinotSortExchangeCopyRule sortExchangeCopyRule;
- if (sortExchangeCopyLimit !=
PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.config.getFetchLimitThreshold()) {
- sortExchangeCopyRule =
ImmutablePinotSortExchangeCopyRule.Config.builder()
- .from(PinotSortExchangeCopyRule.Config.DEFAULT)
- .fetchLimitThreshold(sortExchangeCopyLimit)
- .build()
- .toRule();
- } else {
- sortExchangeCopyRule = PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY;
- }
- return List.of(
- // TODO: Merge the following 2 rules into a single rule
- // add an extra exchange for sort
- PinotSortExchangeNodeInsertRule.INSTANCE,
- sortExchangeCopyRule,
-
- PinotSingleValueAggregateRemoveRule.INSTANCE,
- PinotJoinExchangeNodeInsertRule.INSTANCE,
- PinotAggregateExchangeNodeInsertRule.SortProjectAggregate.INSTANCE,
- PinotAggregateExchangeNodeInsertRule.SortAggregate.INSTANCE,
- PinotAggregateExchangeNodeInsertRule.WithoutSort.INSTANCE,
- PinotWindowSplitRule.INSTANCE,
- PinotWindowExchangeNodeInsertRule.INSTANCE,
- PinotSetOpExchangeNodeInsertRule.INSTANCE,
-
- // apply dynamic broadcast rule after exchange is inserted/
- PinotJoinToDynamicBroadcastRule.INSTANCE,
-
- // remove exchanges when there's duplicates
- PinotExchangeEliminationRule.INSTANCE,
-
- // Evaluate the Literal filter nodes
- CoreRules.FILTER_REDUCE_EXPRESSIONS,
- PinotTableScanConverterRule.INSTANCE
- );
- }
}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index e55566d1066..95be17357ac 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -57,10 +57,10 @@ import org.apache.calcite.tools.Frameworks;
import org.apache.calcite.tools.RelBuilder;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pinot.calcite.rel.rules.ImmutablePinotSortExchangeCopyRule;
import org.apache.pinot.calcite.rel.rules.PinotEnrichedJoinRule;
import org.apache.pinot.calcite.rel.rules.PinotImplicitTableHintRule;
import org.apache.pinot.calcite.rel.rules.PinotJoinToDynamicBroadcastRule;
-import org.apache.pinot.calcite.rel.rules.PinotQueryRuleSets;
import org.apache.pinot.calcite.rel.rules.PinotRelDistributionTraitRule;
import org.apache.pinot.calcite.rel.rules.PinotRuleUtils;
import org.apache.pinot.calcite.rel.rules.PinotSortExchangeCopyRule;
@@ -89,6 +89,9 @@ import
org.apache.pinot.query.planner.physical.v2.PRelNodeTreeValidator;
import
org.apache.pinot.query.planner.physical.v2.PlanFragmentAndMailboxAssignment;
import org.apache.pinot.query.planner.physical.v2.RelToPRelConverter;
import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.rules.PinotRuleSet;
+import org.apache.pinot.query.planner.spi.Phase;
+import org.apache.pinot.query.planner.spi.RuleSetCustomizer;
import org.apache.pinot.query.routing.WorkerManager;
import org.apache.pinot.query.type.TypeFactory;
import org.apache.pinot.query.validate.BytesCastVisitor;
@@ -168,7 +171,7 @@ public class QueryEnvironment {
rootSchema, List.of(database), _typeFactory, CONNECTION_CONFIG,
config.isCaseSensitive());
_defaultDisabledPlannerRules = _envConfig.defaultDisabledPlannerRules();
// default optProgram with no skip rule options and no use rule options
- _optProgram = getOptProgram(Set.of(), Set.of(),
_defaultDisabledPlannerRules);
+ _optProgram = getOptProgram(_envConfig.getRuleSet(), Set.of(), Set.of(),
_defaultDisabledPlannerRules);
_multiClusterRoutingContext = multiClusterRoutingContext;
}
@@ -206,7 +209,7 @@ public class QueryEnvironment {
Set<String> skipRuleSet = QueryOptionsUtils.getSkipPlannerRules(options);
if (!skipRuleSet.isEmpty() || !useRuleSet.isEmpty()) {
// dynamically create optProgram according to rule options
- optProgram = getOptProgram(skipRuleSet, useRuleSet,
_defaultDisabledPlannerRules);
+ optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet,
useRuleSet, _defaultDisabledPlannerRules);
}
}
int sortExchangeCopyLimit =
QueryOptionsUtils.getSortExchangeCopyThreshold(options,
@@ -536,7 +539,7 @@ public class QueryEnvironment {
* @param defaultDisabledRuleSet parsed default disabled rule set from
broker config
* @return HepProgram that performs logical transformations
*/
- private static HepProgram getOptProgram(Set<String> skipRuleSet, Set<String>
useRuleSet,
+ private static HepProgram getOptProgram(PinotRuleSet ruleSet, Set<String>
skipRuleSet, Set<String> useRuleSet,
Set<String> defaultDisabledRuleSet) {
HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
// Set the match order as DEPTH_FIRST. The default is arbitrary which
works the same as DEPTH_FIRST, but it's
@@ -544,16 +547,17 @@ public class QueryEnvironment {
hepProgramBuilder.addMatchOrder(HepMatchOrder.DEPTH_FIRST);
// ----
- // Rules are disabled if its corresponding value is set to false in
ruleFlags
- // construct filtered BASIC_RULES, FILTER_PUSHDOWN_RULES,
PROJECT_PUSHDOWN_RULES, PRUNE_RULES
+ // Rules are disabled if its corresponding value is set to false in
ruleFlags.
+ // Sources come from PinotRuleSet (after every RuleSetCustomizer ran);
per-query
+ // skip/use options are then applied by filterRuleList on a fresh copy.
List<RelOptRule> basicRules =
- filterRuleList(PinotQueryRuleSets.BASIC_RULES, skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
+ filterRuleList(ruleSet.rulesFor(Phase.BASIC), skipRuleSet, useRuleSet,
defaultDisabledRuleSet);
List<RelOptRule> filterPushdownRules =
- filterRuleList(PinotQueryRuleSets.FILTER_PUSHDOWN_RULES, skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
+ filterRuleList(ruleSet.rulesFor(Phase.FILTER_PUSHDOWN), skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
List<RelOptRule> projectPushdownRules =
- filterRuleList(PinotQueryRuleSets.PROJECT_PUSHDOWN_RULES, skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
+ filterRuleList(ruleSet.rulesFor(Phase.PROJECT_PUSHDOWN), skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
List<RelOptRule> pruneRules =
- filterRuleList(PinotQueryRuleSets.PRUNE_RULES, skipRuleSet,
useRuleSet, defaultDisabledRuleSet);
+ filterRuleList(ruleSet.rulesFor(Phase.PRUNE), skipRuleSet, useRuleSet,
defaultDisabledRuleSet);
// Run the Calcite CORE rules using 1 HepInstruction per rule. We use 1
HepInstruction per rule for simplicity:
// the rules used here can rest assured that they are the only ones
evaluated in a dedicated graph-traversal.
@@ -628,6 +632,7 @@ public class QueryEnvironment {
private static HepProgram getTraitProgram(@Nullable WorkerManager
workerManager, Config config,
boolean usePhysicalOptimizer, Set<String> useRuleSet, int
sortExchangeCopyLimit) {
HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ PinotRuleSet ruleSet = config.getRuleSet();
// Set the match order as BOTTOM_UP.
hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
@@ -635,25 +640,37 @@ public class QueryEnvironment {
// ----
// Run pinot specific rules that should run after all other rules, using 1
HepInstruction per rule.
if (!usePhysicalOptimizer) {
- for (RelOptRule relOptRule :
PinotQueryRuleSets.getPinotPostRules(sortExchangeCopyLimit)) {
+ // POST_LOGICAL list comes from PinotRuleSet; we copy it because we may
need to
+ // swap every PinotSortExchangeCopyRule with one configured for the
per-query
+ // (or broker-config) sortExchangeCopyLimit if it differs from the
rule's default.
+ List<RelOptRule> postLogical = new
ArrayList<>(ruleSet.rulesFor(Phase.POST_LOGICAL));
+ if (sortExchangeCopyLimit !=
PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.config.getFetchLimitThreshold()) {
+ PinotSortExchangeCopyRule overridden =
ImmutablePinotSortExchangeCopyRule.Config.builder()
+ .from(PinotSortExchangeCopyRule.Config.DEFAULT)
+ .fetchLimitThreshold(sortExchangeCopyLimit)
+ .build()
+ .toRule();
+ postLogical.replaceAll(r -> r instanceof PinotSortExchangeCopyRule ?
overridden : r);
+ }
+ for (RelOptRule relOptRule : postLogical) {
if (isEligibleQueryPostRule(relOptRule, config)) {
hepProgramBuilder.addRuleInstance(relOptRule);
}
}
if
(!isRuleSkipped(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN,
Set.of(), useRuleSet,
config.defaultDisabledPlannerRules())) {
- // push filter and project above join to enrichedJoin, does not work
with physical optimizer
hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
}
} else {
- for (RelOptRule relOptRule : PinotQueryRuleSets.PINOT_POST_RULES_V2) {
+ for (RelOptRule relOptRule :
ruleSet.rulesFor(Phase.POST_LOGICAL_PHYSICAL_OPT)) {
if (isEligibleQueryPostRule(relOptRule, config)) {
hepProgramBuilder.addRuleInstance(relOptRule);
}
}
}
if (!usePhysicalOptimizer) {
- // apply RelDistribution trait to all nodes
+ // apply RelDistribution trait to all nodes — these rules depend on the
+ // per-query WorkerManager, so they stay outside PinotRuleSet.
if (workerManager != null) {
hepProgramBuilder.addRuleInstance(PinotImplicitTableHintRule.withWorkerManager(workerManager));
}
@@ -695,6 +712,17 @@ public class QueryEnvironment {
@Nullable
TableCache getTableCache();
+ /**
+ * The multi-stage planner's per-phase Calcite rule lists. Defaults to the
+ * process-wide singleton built from {@link
java.util.ServiceLoader}-discovered
+ * {@link RuleSetCustomizer}s, so per-query {@link Config} instances do not
+ * repeat discovery work.
+ */
+ @Value.Default
+ default PinotRuleSet getRuleSet() {
+ return PinotRuleSet.defaultInstance();
+ }
+
@Value.Default
default boolean isNullHandlingEnabled() {
return false;
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/DefaultRuleSetCustomizer.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/DefaultRuleSetCustomizer.java
new file mode 100644
index 00000000000..323c8e3b1f5
--- /dev/null
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/DefaultRuleSetCustomizer.java
@@ -0,0 +1,70 @@
+/**
+ * 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.pinot.query.planner.rules;
+
+import com.google.auto.service.AutoService;
+import java.util.List;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.pinot.calcite.rel.rules.PinotQueryRuleSets;
+import org.apache.pinot.query.planner.spi.Phase;
+import org.apache.pinot.query.planner.spi.RuleSetCustomizer;
+
+
+/// [RuleSetCustomizer] that seeds every [Phase] with the OSS default Calcite
+/// rules for the multi-stage query planner. Registered as a
[java.util.ServiceLoader]
+/// service via [@AutoService], picked up automatically by [PinotRuleSet].
+///
+/// Rule lists are defined in [PinotQueryRuleSets] and may be consolidated here
+/// in a future refactor once the [RuleSetCustomizer] SPI is the established
+/// extension point for broker rule customization.
+@AutoService(RuleSetCustomizer.class)
+public final class DefaultRuleSetCustomizer implements RuleSetCustomizer {
+
+ /// No-arg constructor required by [java.util.ServiceLoader].
+ public DefaultRuleSetCustomizer() {
+ }
+
+ @Override
+ public void customize(Phase phase, List<RelOptRule> rules) {
+ switch (phase) {
+ case BASIC:
+ rules.addAll(PinotQueryRuleSets.BASIC_RULES);
+ return;
+ case FILTER_PUSHDOWN:
+ rules.addAll(PinotQueryRuleSets.FILTER_PUSHDOWN_RULES);
+ return;
+ case PROJECT_PUSHDOWN:
+ rules.addAll(PinotQueryRuleSets.PROJECT_PUSHDOWN_RULES);
+ return;
+ case PRUNE:
+ rules.addAll(PinotQueryRuleSets.PRUNE_RULES);
+ return;
+ case POST_LOGICAL:
+ rules.addAll(PinotQueryRuleSets.POST_LOGICAL_RULES);
+ return;
+ case POST_LOGICAL_PHYSICAL_OPT:
+ rules.addAll(PinotQueryRuleSets.PINOT_POST_RULES_V2);
+ return;
+ default:
+ throw new IllegalStateException(
+ "DefaultRuleSetCustomizer is missing OSS rule defaults for Phase."
+ phase
+ + "; extend the switch when adding a new Phase value.");
+ }
+ }
+}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/PinotRuleSet.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/PinotRuleSet.java
new file mode 100644
index 00000000000..12b152e1438
--- /dev/null
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/rules/PinotRuleSet.java
@@ -0,0 +1,127 @@
+/**
+ * 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.pinot.query.planner.rules;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.ServiceLoader;
+import java.util.Set;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.pinot.query.planner.spi.Phase;
+import org.apache.pinot.query.planner.spi.RuleSetCustomizer;
+import org.apache.pinot.spi.plugin.PluginManager;
+
+
+/// Owns the multi-stage planner's per-phase Calcite rule lists. Constructed
+/// once at broker startup; lists are immutable for the process lifetime.
+///
+/// Construction sequence:
+///
+/// 1. Allocate an empty mutable list per [Phase].
+/// 2. For every [RuleSetCustomizer] (in supplied order), call
+/// `customize(phase, list)` once per phase. The OSS defaults are themselves
+/// contributed by [DefaultRuleSetCustomizer] which is registered as a
+/// `ServiceLoader` entry; plugin customizers run after and observe a list
+/// pre-populated with the OSS defaults.
+/// 3. Defensively copy each per-phase list to an immutable [List] and freeze
+/// the map.
+///
+/// `QueryEnvironment` reads `rulesFor(phase)` and applies per-query
+/// `usePlannerRules` / `skipPlannerRules` filters on top.
+public final class PinotRuleSet {
+
+ private final Map<Phase, List<RelOptRule>> _rulesByPhase;
+
+ /// Builds a rule set from the supplied customizers. Customizers run in the
+ /// order of the iterable; the framework freezes the per-phase lists after
+ /// every customizer has run.
+ public PinotRuleSet(Iterable<RuleSetCustomizer> customizers) {
+ EnumMap<Phase, List<RelOptRule>> mutable = new EnumMap<>(Phase.class);
+ for (Phase phase : Phase.values()) {
+ mutable.put(phase, new ArrayList<>());
+ }
+ for (RuleSetCustomizer customizer : customizers) {
+ for (Phase phase : Phase.values()) {
+ customizer.customize(phase, mutable.get(phase));
+ }
+ }
+ EnumMap<Phase, List<RelOptRule>> frozen = new EnumMap<>(Phase.class);
+ for (Map.Entry<Phase, List<RelOptRule>> entry : mutable.entrySet()) {
+ frozen.put(entry.getKey(), List.copyOf(entry.getValue()));
+ }
+ _rulesByPhase = Collections.unmodifiableMap(frozen);
+ }
+
+ /// Discovers every [RuleSetCustomizer] via [ServiceLoader] and builds a
+ /// rule set from them. Used by [#defaultInstance()] and by callers that
+ /// don't have an externally-managed customizer list.
+ ///
+ /// Discovery order:
+ /// 1. Application-classpath customizers (context classloader) — picks up
+ /// [DefaultRuleSetCustomizer] and any customizer bundled with the broker.
+ /// 2. Plugin classloaders enumerated by
[PluginManager#getPluginClassLoaders()] —
+ /// picks up customizers registered in plugin JARs via
+ ///
`META-INF/services/org.apache.pinot.query.planner.spi.RuleSetCustomizer`.
+ ///
+ /// Call after all plugins have been loaded via [PluginManager]; plugins
loaded
+ /// after this method returns will not be included.
+ public static PinotRuleSet loadFromServiceLoader() {
+ List<RuleSetCustomizer> customizers = new ArrayList<>();
+ // Dedup by class name: the context classloader and a plugin classloader
may both see the same
+ // META-INF/services file if their classpaths overlap (e.g. fat-jar +
plugin realm).
+ Set<String> seen = new LinkedHashSet<>();
+ for (RuleSetCustomizer customizer :
ServiceLoader.load(RuleSetCustomizer.class)) {
+ if (seen.add(customizer.getClass().getName())) {
+ customizers.add(customizer);
+ }
+ }
+ for (ClassLoader pluginClassLoader :
PluginManager.get().getPluginClassLoaders()) {
+ for (RuleSetCustomizer customizer :
ServiceLoader.load(RuleSetCustomizer.class, pluginClassLoader)) {
+ if (seen.add(customizer.getClass().getName())) {
+ customizers.add(customizer);
+ }
+ }
+ }
+ return new PinotRuleSet(customizers);
+ }
+
+ /// Lazily-initialized process-wide singleton built from the
+ /// `ServiceLoader`-discovered customizers. Used as the `@Value.Default` of
+ /// `QueryEnvironment.Config#getRuleSet()` so per-query `Config` instances
+ /// don't repeat the discovery work.
+ public static PinotRuleSet defaultInstance() {
+ return DefaultInstanceHolder.INSTANCE;
+ }
+
+ /// Returns the rule list for the given phase, after every customization
+ /// was applied. Per-query filtering by `usePlannerRules` /
+ /// `skipPlannerRules` is the caller's responsibility (see
+ /// `QueryEnvironment#getOptProgram`).
+ public List<RelOptRule> rulesFor(Phase phase) {
+ return _rulesByPhase.getOrDefault(phase, List.of());
+ }
+
+ private static final class DefaultInstanceHolder {
+ static final PinotRuleSet INSTANCE = loadFromServiceLoader();
+ }
+}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/rules/PinotRuleSetTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/rules/PinotRuleSetTest.java
new file mode 100644
index 00000000000..3f7ce49acb9
--- /dev/null
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/rules/PinotRuleSetTest.java
@@ -0,0 +1,97 @@
+/**
+ * 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.pinot.query.planner.rules;
+
+import java.util.List;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.pinot.calcite.rel.rules.PinotQueryRuleSets;
+import org.apache.pinot.query.planner.spi.Phase;
+import org.apache.pinot.query.planner.spi.RuleSetCustomizer;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+
+
+/// Unit tests for [PinotRuleSet] and the [RuleSetCustomizer] SPI.
+public class PinotRuleSetTest {
+
+ @Test
+ public void defaultsSeedEveryPhaseFromDefaultRuleSetCustomizer() {
+ PinotRuleSet ruleSet = new PinotRuleSet(List.of(new
DefaultRuleSetCustomizer()));
+
+ assertEquals(ruleSet.rulesFor(Phase.BASIC),
PinotQueryRuleSets.BASIC_RULES);
+ assertEquals(ruleSet.rulesFor(Phase.FILTER_PUSHDOWN),
PinotQueryRuleSets.FILTER_PUSHDOWN_RULES);
+ assertEquals(ruleSet.rulesFor(Phase.PROJECT_PUSHDOWN),
PinotQueryRuleSets.PROJECT_PUSHDOWN_RULES);
+ assertEquals(ruleSet.rulesFor(Phase.PRUNE),
PinotQueryRuleSets.PRUNE_RULES);
+ assertEquals(ruleSet.rulesFor(Phase.POST_LOGICAL),
PinotQueryRuleSets.POST_LOGICAL_RULES);
+ assertEquals(ruleSet.rulesFor(Phase.POST_LOGICAL_PHYSICAL_OPT),
PinotQueryRuleSets.PINOT_POST_RULES_V2);
+ }
+
+ @Test
+ public void serviceLoaderDiscoveryFindsDefault() {
+ // The DefaultRuleSetCustomizer is registered via META-INF/services, so the
+ // default instance built by ServiceLoader exposes the OSS defaults.
+ PinotRuleSet ruleSet = PinotRuleSet.defaultInstance();
+ assertTrue(ruleSet.rulesFor(Phase.BASIC).size() > 0);
+ assertTrue(ruleSet.rulesFor(Phase.POST_LOGICAL).size() > 0);
+ }
+
+ @Test
+ public void customizerCanAppendRule() {
+ RelOptRule extraRule = PinotQueryRuleSets.BASIC_RULES.get(0);
+ int defaultSize = PinotQueryRuleSets.BASIC_RULES.size();
+
+ RuleSetCustomizer plugin = (phase, rules) -> {
+ if (phase == Phase.BASIC) {
+ rules.add(extraRule);
+ }
+ };
+ PinotRuleSet ruleSet = new PinotRuleSet(List.of(new
DefaultRuleSetCustomizer(), plugin));
+
+ assertEquals(ruleSet.rulesFor(Phase.BASIC).size(), defaultSize + 1);
+ assertSame(ruleSet.rulesFor(Phase.BASIC).get(defaultSize), extraRule);
+ }
+
+ @Test
+ public void customizerCanRemoveOssRuleByName() {
+ String firstRuleName = PinotQueryRuleSets.BASIC_RULES.get(0).toString();
+ int defaultSize = PinotQueryRuleSets.BASIC_RULES.size();
+
+ RuleSetCustomizer plugin = (phase, rules) -> {
+ if (phase == Phase.BASIC) {
+ rules.removeIf(r -> firstRuleName.equals(r.toString()));
+ }
+ };
+ PinotRuleSet ruleSet = new PinotRuleSet(List.of(new
DefaultRuleSetCustomizer(), plugin));
+
+ assertEquals(ruleSet.rulesFor(Phase.BASIC).size(), defaultSize - 1);
+ assertTrue(ruleSet.rulesFor(Phase.BASIC).stream().noneMatch(r ->
firstRuleName.equals(r.toString())));
+ }
+
+ @Test
+ public void rulesForUnseededPhaseReturnsEmpty() {
+ // No customizers — every phase comes back empty.
+ PinotRuleSet ruleSet = new PinotRuleSet(List.of());
+ for (Phase phase : Phase.values()) {
+ assertEquals(ruleSet.rulesFor(phase).size(), 0, phase + " should be
empty");
+ }
+ }
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/plugin/PluginManager.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/plugin/PluginManager.java
index a8ae0bdc70b..ec39f672dc4 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/plugin/PluginManager.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/plugin/PluginManager.java
@@ -34,9 +34,11 @@ import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
+import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.commons.io.FileUtils;
@@ -298,10 +300,19 @@ public class PluginManager {
ClassRealm pinotRealm = _classWorld.getClassRealm(PINOT_REALMID);
- // All packages to look up in pinot realm BEFORE itself
- Stream<String> importedPinotPackages =
- Stream.of("org.apache.pinot.spi"); // this works like a prefix, so
ALL spi classes will be accessible
- importedPinotPackages.forEach(p -> pluginRealm.importFrom(pinotRealm,
p));
+ // All packages to look up in pinot realm BEFORE itself.
+ // Acts like a prefix so all classes in (and below) the listed package
are accessible.
+ // Hardcoding the list here is cheap while it stays small. A cleaner
long-term approach
+ // would have each SPI module self-declare its exports in a
META-INF/pinot-realm-exports
+ // resource file (one package per line) and have PluginManager
discover them at init time
+ // via ClassLoader.getResources() — that would eliminate the layering
violation of
+ // pinot-spi naming packages from higher-level modules such as
pinot-query-planner-spi.
+ // TODO: implement the self-declaring META-INF/pinot-realm-exports
mechanism.
+ Stream.of(
+ "org.apache.pinot.spi",
+ "org.apache.pinot.query.planner.spi", // RuleSetCustomizer SPI
(pinot-query-planner-spi)
+ "org.apache.calcite.plan" // RelOptRule, used by
RuleSetCustomizer.customize
+ ).forEach(p -> pluginRealm.importFrom(pinotRealm, p));
// Additional importForm as specified by the plugin configuration
config.getImportsFromPerRealm().forEach((r, ifs) -> {
@@ -474,6 +485,32 @@ public class PluginManager {
return null;
}
+ /// Returns the set of classloaders for all new-style plugins, in load order.
+ /// New-style plugins are those packaged with a `pinot-plugin.properties`
file
+ /// and loaded into a dedicated [ClassRealm]. Legacy shaded plugins (loaded
via
+ /// [PluginClassLoader]) are excluded — new SPIs should only target
new-style plugins.
+ ///
+ /// Intended for `ServiceLoader` enumeration across plugin classloaders:
+ ///
+ /// ```java
+ /// for (ClassLoader cl : PluginManager.get().getPluginClassLoaders()) {
+ /// for (MyService svc : ServiceLoader.load(MyService.class, cl)) { ... }
+ /// }
+ /// ```
+ ///
+ /// Call after all plugins have been loaded; classloaders added after this
+ /// call returns will not appear in the snapshot.
+ public synchronized Set<ClassLoader> getPluginClassLoaders() {
+ Set<ClassLoader> result = new LinkedHashSet<>();
+ for (ClassRealm realm : _classWorld.getRealms()) {
+ String id = realm.getId();
+ if (!PINOT_REALMID.equals(id) && !DEFAULT_PLUGIN_NAME.equals(id)) {
+ result.add(realm);
+ }
+ }
+ return Collections.unmodifiableSet(result);
+ }
+
public static PluginManager get() {
return PLUGIN_MANAGER;
}
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/plugin/PluginManagerTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/plugin/PluginManagerTest.java
index b5f8ba2d8e6..89918acc209 100644
--- a/pinot-spi/src/test/java/org/apache/pinot/spi/plugin/PluginManagerTest.java
+++ b/pinot-spi/src/test/java/org/apache/pinot/spi/plugin/PluginManagerTest.java
@@ -22,8 +22,10 @@ import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URL;
+import java.nio.file.Files;
import java.util.HashMap;
import java.util.Map;
+import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import javax.tools.JavaCompiler;
@@ -339,6 +341,45 @@ public class PluginManagerTest {
Assert.assertSame(instance1, instance2, "PluginManager should be a
singleton");
}
+ @Test
+ public void testGetPluginClassLoadersEmptyOnFreshInstance() {
+ PluginManager pm = new PluginManager();
+ Set<ClassLoader> loaders = pm.getPluginClassLoaders();
+ Assert.assertNotNull(loaders, "getPluginClassLoaders() must never return
null");
+ Assert.assertTrue(loaders.isEmpty(),
+ "A fresh PluginManager with no loaded plugins should return an empty
set");
+ }
+
+ @Test
+ public void testGetPluginClassLoadersIncludesNewStylePlugin()
+ throws Exception {
+ // New-style plugins (with pinot-plugin.properties) are loaded into a
ClassRealm.
+ // Old-style plugins (without the file) are NOT returned by
getPluginClassLoaders().
+ File newStyleDir = new File(_tempDir, "new-style-plugin-for-cl-test");
+ newStyleDir.mkdirs();
+ Files.createFile(newStyleDir.toPath().resolve("pinot-plugin.properties"));
+ try (JarOutputStream jos = new JarOutputStream(new FileOutputStream(new
File(newStyleDir, "dummy.jar")))) {
+ // intentionally empty JAR
+ }
+ File oldStyleDir = new File(_tempDir, "old-style-plugin-for-cl-test");
+ oldStyleDir.mkdirs();
+ try (JarOutputStream jos = new JarOutputStream(new FileOutputStream(new
File(oldStyleDir, "dummy.jar")))) {
+ // intentionally empty JAR — no pinot-plugin.properties
+ }
+
+ PluginManager pm = new PluginManager();
+ Assert.assertTrue(pm.getPluginClassLoaders().isEmpty(), "Should start
empty");
+
+ pm.load("old-style-plugin-for-cl-test", oldStyleDir);
+ Assert.assertTrue(pm.getPluginClassLoaders().isEmpty(),
+ "Old-style plugins must not appear in getPluginClassLoaders()");
+
+ pm.load("new-style-plugin-for-cl-test", newStyleDir);
+ Set<ClassLoader> loaders = pm.getPluginClassLoaders();
+ Assert.assertEquals(loaders.size(), 1,
+ "New-style plugin must appear in getPluginClassLoaders()");
+ }
+
@Test
public void testRegisterRecordReaderClassCaseInsensitivity() {
// Test that registration works with different case formats
diff --git a/pom.xml b/pom.xml
index b906345eb77..6d143121469 100644
--- a/pom.xml
+++ b/pom.xml
@@ -58,6 +58,7 @@
<module>pinot-connectors</module>
<module>pinot-segment-local</module>
<module>pinot-compatibility-verifier</module>
+ <module>pinot-query-planner-spi</module>
<module>pinot-query-planner</module>
<module>pinot-query-runtime</module>
<module>pinot-timeseries</module>
@@ -643,6 +644,11 @@
<artifactId>pinot-materialized-view</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.pinot</groupId>
+ <artifactId>pinot-query-planner-spi</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.pinot</groupId>
<artifactId>pinot-query-planner</artifactId>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
