This is an automated email from the ASF dual-hosted git repository. abhishekrb pushed a commit to branch explain_plan_set_statement in repository https://gitbox.apache.org/repos/asf/druid.git
commit e4f9d59b02d12dd2f4d8d3c7d9b42faa87e30ee3 Author: Abhishek Balaji Radhakrishnan <[email protected]> AuthorDate: Fri May 2 09:27:08 2025 -0700 Make explain plan queries honor query context from SET statements. Please enter the commit message for your changes. Lines starting --- .../druid/sql/calcite/planner/DruidPlanner.java | 1 + .../druid/sql/calcite/planner/PlannerContext.java | 10 ++- .../druid/sql/calcite/CalciteExplainQueryTest.java | 75 ++++++++++++++++++++++ .../druid/sql/calcite/CalciteInsertDmlTest.java | 44 +++++++++++++ 4 files changed, 129 insertions(+), 1 deletion(-) 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 2842eb76e94..aeea9c6f4e4 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 @@ -349,6 +349,7 @@ public class DruidPlanner implements Closeable throw InvalidSqlInput.exception("Statement list is missing a non-SET statement to execute"); } plannerContext.addAllToQueryContext(contextMap); + plannerContext.addAllToPlannerConfig(contextMap); } return root; } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java index a884a3895b7..00ee3422605 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java @@ -128,7 +128,7 @@ public class PlannerContext private final PlannerToolbox plannerToolbox; private final ExpressionParser expressionParser; private final String sql; - private final PlannerConfig plannerConfig; + private PlannerConfig plannerConfig; private final SqlEngine engine; private final Map<String, Object> queryContext; private final CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>(); @@ -559,6 +559,14 @@ public class PlannerContext initializeContextFields(); } + /** + * Add additional query context parameters to planner config, overriding any existing values. + */ + public void addAllToPlannerConfig(Map<String, Object> toAdd) + { + this.plannerConfig = this.plannerConfig.withOverrides(toAdd); + } + public SqlEngine getEngine() { return engine; diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteExplainQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteExplainQueryTest.java index b5c3061e503..7a8733f5908 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteExplainQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteExplainQueryTest.java @@ -20,6 +20,7 @@ package org.apache.druid.sql.calcite; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import org.apache.druid.sql.calcite.planner.PlannerConfig; import org.apache.druid.sql.calcite.util.CalciteTests; import org.junit.jupiter.api.Test; @@ -113,6 +114,80 @@ public class CalciteExplainQueryTest extends BaseCalciteQueryTest ); } + @Test + public void testSetStatementWithExplainSanity() + { + // Skip vectorization since otherwise the "context" will change for each subtest. + skipVectorize(); + + final String query = "SET plannerStrategy = 'DECOUPLED';\n" + + " EXPLAIN PLAN FOR SELECT COUNT(*)\n" + + "FROM (\n" + + " SELECT DISTINCT dim2\n" + + " FROM druid.foo\n" + + " WHERE SUBSTRING(dim2, 1, 1) IN (\n" + + " SELECT SUBSTRING(dim1, 1, 1) FROM druid.foo WHERE dim1 IS NOT NULL\n" + + " )\n" + + ")"; + + final String explanation = "DruidAggregate(group=[{}], EXPR$0=[COUNT()], druid=[logical])\n" + + " DruidAggregate(group=[{0}], druid=[logical])\n" + + " DruidJoin(condition=[=($1, $2)], joinType=[inner])\n" + + " DruidProject(dim2=[$2], $f1=[SUBSTRING($2, 1, 1)], druid=[logical])\n" + + " DruidTableScan(table=[[druid, foo]], druid=[logical])\n" + + " DruidAggregate(group=[{0}], druid=[logical])\n" + + " DruidProject(EXPR$0=[SUBSTRING($1, 1, 1)], druid=[logical])\n" + + " DruidFilter(condition=[IS NOT NULL($1)])\n" + + " DruidTableScan(table=[[druid, foo]], druid=[logical])\n"; + final String resources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]"; + final String attributes = "{\"statementType\":\"SELECT\"}"; + + testQuery( + query, + ImmutableList.of(), + ImmutableList.of(new Object[]{explanation, resources, attributes}) + ); + } + + @Test + public void testMultiStatementSetsContextOverridesQueryContext() + { + // Skip vectorization since otherwise the "context" will change for each subtest. + skipVectorize(); + + final String query = "SET plannerStrategy = 'DECOUPLED';\n" + + " SET timeout = 90000;\n" + + " EXPLAIN PLAN FOR \n" + + "SELECT COUNT(*)\n" + + "FROM (\n" + + " SELECT DISTINCT dim2\n" + + " FROM druid.foo\n" + + " WHERE SUBSTRING(dim2, 1, 1) IN (\n" + + " SELECT SUBSTRING(dim1, 1, 1) FROM druid.foo WHERE dim1 IS NOT NULL\n" + + " )\n" + + ")"; + + final String explanation = "DruidAggregate(group=[{}], EXPR$0=[COUNT()], druid=[logical])\n" + + " DruidAggregate(group=[{0}], druid=[logical])\n" + + " DruidJoin(condition=[=($1, $2)], joinType=[inner])\n" + + " DruidProject(dim2=[$2], $f1=[SUBSTRING($2, 1, 1)], druid=[logical])\n" + + " DruidTableScan(table=[[druid, foo]], druid=[logical])\n" + + " DruidAggregate(group=[{0}], druid=[logical])\n" + + " DruidProject(EXPR$0=[SUBSTRING($1, 1, 1)], druid=[logical])\n" + + " DruidFilter(condition=[IS NOT NULL($1)])\n" + + " DruidTableScan(table=[[druid, foo]], druid=[logical])\n"; + + final String resources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]"; + final String attributes = "{\"statementType\":\"SELECT\"}"; + + testQuery( + query, + ImmutableMap.of("plannerStrategy", "COUPLED"), + ImmutableList.of(), + ImmutableList.of(new Object[]{explanation, resources, attributes}) + ); + } + @Test public void testExplainSelectStar() { diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteInsertDmlTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteInsertDmlTest.java index 6a078e68fb7..7fffbc5216c 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteInsertDmlTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteInsertDmlTest.java @@ -901,6 +901,50 @@ public class CalciteInsertDmlTest extends CalciteIngestionDmlTest didTest = true; } + @Test + public void testSetStatementWithExplainPlanInsertJoinQuery() + { + skipVectorize(); + + final String resources = "[{\"name\":\"dst\",\"type\":\"DATASOURCE\"},{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]"; + final String attributes = "{\"statementType\":\"INSERT\",\"targetDataSource\":\"dst\",\"partitionedBy\":\"DAY\",\"clusteredBy\":[\"floor_m1\",\"dim1\",\"CEIL(\\\"m2\\\")\"]}"; + + final String sql = "SET plannerStrategy = 'DECOUPLED';\n" + + " SET timeout = 9000;\n" + + "EXPLAIN PLAN FOR \n" + + "INSERT INTO druid.dst \n" + + "SELECT __time, FLOOR(m1) as floor_m1, dim1, CEIL(m2) as ceil_m2 FROM foo \n" + + "PARTITIONED BY FLOOR(__time TO DAY) CLUSTERED BY 2, dim1, CEIL(m2)"; + + // Test correctness of the query when only the CLUSTERED BY clause is present + final String explanation = + "DruidSort(sort0=[$1], sort1=[$2], sort2=[$3], dir0=[ASC], dir1=[ASC], dir2=[ASC], druid=[logical])\n" + + " DruidProject(__time=[$0], floor_m1=[FLOOR($5)], dim1=[$1], ceil_m2=[CEIL($6)], druid=[logical])\n" + + " DruidTableScan(table=[[druid, foo]], druid=[logical])\n"; + + testQuery( + PLANNER_CONFIG_NATIVE_QUERY_EXPLAIN, + ImmutableMap.of("sqlQueryId", "dummy"), + Collections.emptyList(), + sql, + CalciteTests.SUPER_USER_AUTH_RESULT, + ImmutableList.of(), + new DefaultResultsVerifier( + ImmutableList.of( + new Object[]{ + explanation, + resources, + attributes + } + ), + null + ) + ); + + // Not using testIngestionQuery, so must set didTest manually to satisfy the check in tearDown. + didTest = true; + } + @Test public void testExplainPlanInsertWithClusteredByDescThrowsException() { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
