From 54a664a462f1caa0b3f78a43fe96a9db54f95aaf Mon Sep 17 00:00:00 2001 From: Laksh Singla Date: Fri, 15 Sep 2023 02:09:46 +0530 Subject: [PATCH] review comments --- .../apache/druid/msq/exec/MSQFaultsTest.java | 21 ++++++++++------- .../druid/sql/calcite/rule/DruidRules.java | 7 ++++-- .../sql/calcite/rule/DruidSortUnionRule.java | 23 +++++++------------ .../sql/calcite/rule/DruidUnionRule.java | 8 ------- 4 files changed, 26 insertions(+), 33 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java index 43e5019ddc82..6174518e18aa 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java @@ -361,27 +361,32 @@ public void testTooManyInputFiles() throws IOException } @Test - public void testUnionAllIsDisallowed() + public void testUnionAllUsingUnionDataSourceDisallowed() { final RowSignature rowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build(); - testIngestQuery() + // This plans the query using DruidUnionDataSourceRule since the DruidUnionDataSourceRule#isCompatible + // returns true (column names, types match, and it is a union on the table data sources). + // It gets planned correctly, however MSQ engine cannot plan the query correctly + testSelectQuery() .setSql("SELECT * FROM foo\n" + "UNION ALL\n" + "SELECT * FROM foo\n") .setExpectedRowSignature(rowSignature) - .setExpectedDataSource("foo1") .setExpectedMSQFault(QueryNotSupportedFault.instance()) .verifyResults(); } @Test - public void testUnionAllIsDisallowedWhilePlanning() + public void testUnionAllUsingTopLevelUnionDisallowedWhilePlanning() { - // This results in a planning error however the planning error isn't an accurate representation of the actual error - // because Calcite rewrites it using CoreRules.UNION_TO_DISTINCT, which plans it using Union Datasource. - // However, this fails since the column names mismatch. Once MSQ is able to support Union datasources, the planning - // error would become an accurate representation of the error. + // This results in a planning error however the planning error isn't an accurate representation of the actual error. + // Calcite tries to plan the query using DruidUnionRule, which passes with native, however fails with MSQ (due to engine + // feature ALLOW_TOP_LEVEL_UNION_ALL being absent) + // Calcite then tries to write it using DruidUnionDataSourceRule. However, the isCompatible check fails because column + // names mismatch. But it sets the planning error with this mismatch, which can be misleading since native queries can + // plan fine using the DruidUnionRule (that's disabled in MSQ) + // Once MSQ is able to support union datasources, we'd be able to execute this query fine in MSQ testIngestQuery() .setSql( "INSERT INTO druid.dst " diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java index b1a6a557f761..82ac218b5097 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java @@ -95,13 +95,16 @@ public static List rules(PlannerContext plannerContext) DruidOuterQueryRule.WHERE_FILTER, DruidOuterQueryRule.SELECT_PROJECT, DruidOuterQueryRule.SORT, - new DruidUnionRule(plannerContext), new DruidUnionDataSourceRule(plannerContext), - new DruidSortUnionRule(plannerContext), DruidJoinRule.instance(plannerContext) ) ); + if (plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) { + retVal.add(new DruidUnionRule(plannerContext)); + retVal.add(DruidSortUnionRule.instance()); + } + if (plannerContext.featureAvailable(EngineFeature.WINDOW_FUNCTIONS)) { retVal.add(new DruidQueryRule<>(Window.class, PartialDruidQuery.Stage.WINDOW, PartialDruidQuery::withWindow)); retVal.add( diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidSortUnionRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidSortUnionRule.java index f316ab9f0487..d06c39d72b5b 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidSortUnionRule.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidSortUnionRule.java @@ -23,9 +23,7 @@ import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.core.Sort; import org.apache.calcite.rex.RexLiteral; -import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.rel.DruidUnionRel; -import org.apache.druid.sql.calcite.run.EngineFeature; import java.util.Collections; @@ -34,27 +32,22 @@ */ public class DruidSortUnionRule extends RelOptRule { - private final PlannerContext plannerContext; - public DruidSortUnionRule(PlannerContext plannerContext) + private static final DruidSortUnionRule INSTANCE = new DruidSortUnionRule(); + + private DruidSortUnionRule() { super(operand(Sort.class, operand(DruidUnionRel.class, any()))); - this.plannerContext = plannerContext; + } + + public static DruidSortUnionRule instance() + { + return INSTANCE; } @Override public boolean matches(final RelOptRuleCall call) { - // Defensive check. If the planner disallows top level union all, then the DruidUnionRule would have prevented - // creating the DruidUnionRel in the first place - if (!plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) { - plannerContext.setPlanningError( - "Top level 'UNION ALL' is unsupported by SQL engine [%s].", - plannerContext.getEngine().name() - ); - return false; - } - // LIMIT, no ORDER BY final Sort sort = call.rel(0); return sort.collation.getFieldCollations().isEmpty() && sort.fetch != null; diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java index c7de9eb0e643..40cb2161c155 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java @@ -26,7 +26,6 @@ import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.rel.DruidRel; import org.apache.druid.sql.calcite.rel.DruidUnionRel; -import org.apache.druid.sql.calcite.run.EngineFeature; import java.util.List; @@ -52,13 +51,6 @@ public DruidUnionRule(PlannerContext plannerContext) @Override public boolean matches(RelOptRuleCall call) { - if (!plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) { - plannerContext.setPlanningError( - "Top level 'UNION ALL' is unsupported by SQL engine [%s].", - plannerContext.getEngine().name() - ); - return false; - } // Make DruidUnionRule and DruidUnionDataSourceRule mutually exclusive. final Union unionRel = call.rel(0); final DruidRel firstDruidRel = call.rel(1);