asolimando commented on code in PR #4585:
URL: https://github.com/apache/calcite/pull/4585#discussion_r2439896339
##########
core/src/main/java/org/apache/calcite/plan/volcano/VolcanoPlanner.java:
##########
@@ -522,26 +523,33 @@ protected void registerMaterializations() {
ensureRootConverters();
registerMaterializations();
- ruleDriver.drive();
-
- if (LOGGER.isTraceEnabled()) {
- StringWriter sw = new StringWriter();
- final PrintWriter pw = new PrintWriter(sw);
- dump(pw);
- pw.flush();
- LOGGER.info(sw.toString());
- }
- dumpRuleAttemptsInfo();
- RelNode cheapest = root.buildCheapestPlan(this);
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(
- "Cheapest plan:\n{}", RelOptUtil.toString(cheapest,
SqlExplainLevel.ALL_ATTRIBUTES));
-
- if (!provenanceMap.isEmpty()) {
- LOGGER.debug("Provenance:\n{}", Dumpers.provenance(provenanceMap,
cheapest));
+ // Set the current planner context for optimization history tracking
+ RelCollationTraitDef.setPlannerContext(this);
+ try {
+ ruleDriver.drive();
+
+ if (LOGGER.isTraceEnabled()) {
+ StringWriter sw = new StringWriter();
+ final PrintWriter pw = new PrintWriter(sw);
+ dump(pw);
+ pw.flush();
+ LOGGER.info(sw.toString());
+ }
+ dumpRuleAttemptsInfo();
+ RelNode cheapest = root.buildCheapestPlan(this);
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(
+ "Cheapest plan:\n{}", RelOptUtil.toString(cheapest,
SqlExplainLevel.ALL_ATTRIBUTES));
+
+ if (!provenanceMap.isEmpty()) {
+ LOGGER.debug("Provenance:\n{}", Dumpers.provenance(provenanceMap,
cheapest));
+ }
}
+ return cheapest;
+ } finally {
+ // Clear the planner context when optimization completes
Review Comment:
nit: this comment can be removed, pretty clear from the code
##########
core/src/main/java/org/apache/calcite/rel/RelCollationTraitDef.java:
##########
@@ -87,4 +141,45 @@ private RelCollationTraitDef() {
RelOptPlanner planner, RelCollation fromTrait, RelCollation toTrait) {
return true;
}
+
+ /**
+ * Associates an optimized collation with its original collation.
+ * Used to preserve optimization history for trait satisfaction checking.
+ * The association is scoped to the planner's lifecycle.
+ */
+ public static void setOriginalCollation(
+ RelOptPlanner planner,
+ RelCollation optimized,
+ List<RelFieldCollation> original) {
+ Map<RelCollation, List<RelFieldCollation>> map =
+ getOriginalCollationsMap(planner);
+ map.put(optimized, original);
+ }
+
+ /**
+ * Gets the original collation associated with an optimized collation, if
any.
+ * Returns null if no original collation was recorded.
+ */
+ public static @Nullable List<RelFieldCollation> getOriginalCollation(
+ RelOptPlanner planner,
+ RelCollation collation) {
+ Map<RelCollation, List<RelFieldCollation>> map =
+ getOriginalCollationsMap(planner);
+ return map.get(collation);
+ }
+
+ /**
+ * Gets the original collation associated with an optimized collation using
the
+ * current planner context. Convenience method for trait satisfaction
checking.
+ * Returns null if no current planner or no original collation was recorded.
Review Comment:
Javadoc has `@returns`, you should probably use that, especially since those
are public methods, it applies to all (public) methods
##########
core/src/main/java/org/apache/calcite/rel/RelCollationTraitDef.java:
##########
@@ -87,4 +141,45 @@ private RelCollationTraitDef() {
RelOptPlanner planner, RelCollation fromTrait, RelCollation toTrait) {
return true;
}
+
+ /**
+ * Associates an optimized collation with its original collation.
+ * Used to preserve optimization history for trait satisfaction checking.
+ * The association is scoped to the planner's lifecycle.
+ */
+ public static void setOriginalCollation(
+ RelOptPlanner planner,
+ RelCollation optimized,
+ List<RelFieldCollation> original) {
+ Map<RelCollation, List<RelFieldCollation>> map =
+ getOriginalCollationsMap(planner);
+ map.put(optimized, original);
+ }
+
+ /**
+ * Gets the original collation associated with an optimized collation, if
any.
+ * Returns null if no original collation was recorded.
+ */
+ public static @Nullable List<RelFieldCollation> getOriginalCollation(
+ RelOptPlanner planner,
+ RelCollation collation) {
+ Map<RelCollation, List<RelFieldCollation>> map =
+ getOriginalCollationsMap(planner);
+ return map.get(collation);
+ }
+
+ /**
+ * Gets the original collation associated with an optimized collation using
the
+ * current planner context. Convenience method for trait satisfaction
checking.
+ * Returns null if no current planner or no original collation was recorded.
+ */
+ public static @Nullable List<RelFieldCollation>
getOriginalCollationFromContext(
+ RelCollation collation) {
+ RelOptPlanner planner = getCurrentPlanner();
+ if (planner == null) {
Review Comment:
I'd put this check in the `getOriginalCollation` method since that one is
public too and here you don't use the planner, better to guard against a `null`
there
##########
core/src/main/java/org/apache/calcite/plan/volcano/VolcanoPlanner.java:
##########
@@ -522,26 +523,33 @@ protected void registerMaterializations() {
ensureRootConverters();
registerMaterializations();
- ruleDriver.drive();
-
- if (LOGGER.isTraceEnabled()) {
- StringWriter sw = new StringWriter();
- final PrintWriter pw = new PrintWriter(sw);
- dump(pw);
- pw.flush();
- LOGGER.info(sw.toString());
- }
- dumpRuleAttemptsInfo();
- RelNode cheapest = root.buildCheapestPlan(this);
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(
- "Cheapest plan:\n{}", RelOptUtil.toString(cheapest,
SqlExplainLevel.ALL_ATTRIBUTES));
-
- if (!provenanceMap.isEmpty()) {
- LOGGER.debug("Provenance:\n{}", Dumpers.provenance(provenanceMap,
cheapest));
+ // Set the current planner context for optimization history tracking
+ RelCollationTraitDef.setPlannerContext(this);
Review Comment:
Q: can you explain what is the effect of this? what would happen without it?
##########
core/src/main/java/org/apache/calcite/rel/RelCollationTraitDef.java:
##########
@@ -42,6 +48,54 @@ public class RelCollationTraitDef extends
RelTraitDef<RelCollation> {
public static final RelCollationTraitDef INSTANCE =
new RelCollationTraitDef();
+ /** ThreadLocal to store the current planner during optimization.
+ * This allows trait satisfaction checking to access the planner's collation
map
+ * without modifying the trait interface. */
+ private static final ThreadLocal<RelOptPlanner> CURRENT_PLANNER =
+ new ThreadLocal<>();
+
+ /** Maps planners to their original collation maps using weak references
+ * for planner keys to avoid memory leaks when planners are garbage
collected.
+ * Each planner gets its own map to store optimization history for that
query.
+ * The inner map uses HashMap because RelCollations are typically kept alive
+ * by the planner's RelSet and won't be GC'd during optimization. */
+ private static final Map<RelOptPlanner, Map<RelCollation,
List<RelFieldCollation>>>
+ PLANNER_COLLATIONS_MAP = Collections.synchronizedMap(new
WeakHashMap<>());
+
+ /**
+ * Sets the current planner for the optimization pass.
+ * Should be called at the beginning of optimization to enable
+ * proper trait satisfaction checking with original collations.
+ */
+ public static void setPlannerContext(RelOptPlanner planner) {
+ CURRENT_PLANNER.set(planner);
+ }
+
+ /**
+ * Clears the current planner context.
+ * Should be called at the end of optimization.
+ */
+ public static void clearPlannerContext() {
Review Comment:
Q: this kind of usage reminds me of the `AutoCloseable` interface +
`try-with-resources`, but making this class implement that would be too
invasive probably?
##########
core/src/main/java/org/apache/calcite/rel/RelCollationTraitDef.java:
##########
@@ -87,4 +141,45 @@ private RelCollationTraitDef() {
RelOptPlanner planner, RelCollation fromTrait, RelCollation toTrait) {
return true;
}
+
+ /**
+ * Associates an optimized collation with its original collation.
+ * Used to preserve optimization history for trait satisfaction checking.
+ * The association is scoped to the planner's lifecycle.
+ */
+ public static void setOriginalCollation(
Review Comment:
Nit: I was a bit confused by the name of the method since we are passing
both original and optimized, maybe something `registerCollectionOptimization`
or something that conveys the message that we are storing a correspondence
between original and optimized collations?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
