You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by ab...@apache.org on 2025/05/02 16:40:08 UTC
(druid) 01/01: Make explain plan queries honor query context from SET statements.
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 30cca6d66f741e5743bce8d5eab7faa23d6ed6be
Author: Abhishek Balaji Radhakrishnan <ab...@gmail.com>
AuthorDate: Fri May 2 09:27:08 2025 -0700
Make explain plan queries honor query context from SET statements.
- EXPLAIN PLAN queries were ignoring the query context parameters
from the SET statement.
- This was because the planner config was eagerly built before the
set statement list from the sql query is processed and the query
context is extracted from it.
- This patch now overrides the planner config in processStatementList.
In order to do this override, the field is made non-final (similar to
other non-final mutable members in PlannerContext).
---
.../druid/sql/calcite/planner/DruidPlanner.java | 6 +-
.../druid/sql/calcite/planner/PlannerContext.java | 10 ++-
.../druid/sql/calcite/CalciteExplainQueryTest.java | 75 ++++++++++++++++++++++
.../druid/sql/calcite/CalciteInsertDmlTest.java | 44 +++++++++++++
4 files changed, 132 insertions(+), 3 deletions(-)
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..984cc303666 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
@@ -297,8 +297,9 @@ public class DruidPlanner implements Closeable
* If an {@link SqlNode} is a {@link SqlNodeList}, it must consist of 0 or more {@link SqlSetOption} followed by a
* single {@link SqlNode} which is NOT a {@link SqlSetOption}. All {@link SqlSetOption} will be converted into a
* context parameters {@link Map} and added to the {@link PlannerContext} with
- * {@link PlannerContext#addAllToQueryContext(Map)}. The final {@link SqlNode} of the {@link SqlNodeList} is returned
- * by this method as the {@link SqlNode} which should actually be validated and executed, and will have access to the
+ * {@link PlannerContext#addAllToQueryContext(Map)} and {@link PlannerContext#addAllToPlannerConfig(Map)}.
+ * The final {@link SqlNode} of the {@link SqlNodeList} is returned by this method as the {@link SqlNode} which
+ * should actually be validated and executed, and will have access to the
* modified query context through the {@link PlannerContext}. {@link SqlSetOption} override any existing query
* context parameter values.
*/
@@ -349,6 +350,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..8f884d26bf2 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,13 +128,13 @@ public class PlannerContext
private final PlannerToolbox plannerToolbox;
private final ExpressionParser expressionParser;
private final String sql;
- private final PlannerConfig plannerConfig;
private final SqlEngine engine;
private final Map<String, Object> queryContext;
private final CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>();
private final PlannerHook hook;
private final Set<String> lookupsToLoad = new HashSet<>();
+ private PlannerConfig plannerConfig;
private String sqlQueryId;
private boolean stringifyArrays;
private boolean useBoundsAndSelectors;
@@ -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: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org