Improve the output of SQL explain message

docs/configuration/index.md

@@ -1834,6 +1834,7 @@ The Druid SQL server is configured through the following properties on the Broke
 |`druid.sql.planner.metadataSegmentCacheEnable`|Whether to keep a cache of published segments in broker. If true, broker polls coordinator in background to get segments from metadata store and maintains a local cache. If false, coordinator's REST API will be invoked when broker needs published segments info.|false|
 |`druid.sql.planner.metadataSegmentPollPeriod`|How often to poll coordinator for published segments list if `druid.sql.planner.metadataSegmentCacheEnable` is set to true. Poll period is in milliseconds. |60000|
 |`druid.sql.planner.authorizeSystemTablesDirectly`|If true, Druid authorizes queries against any of the system schema tables (`sys` in SQL) as `SYSTEM_TABLE` resources which require `READ` access, in addition to permissions based content filtering.|false|
+|`druid.sql.planner.useNativeQueryExplain`|If true, `EXPLAIN PLAN FOR` will return the explain plan as a JSON representation of equivalent native query(s), else it will return the original version of explain plan generated by Calcite. It can be overridden per query with `useNativeQueryExplain` context key.|false|
 |`druid.sql.approxCountDistinct.function`|Implementation to use for the [`APPROX_COUNT_DISTINCT` function](../querying/sql.md#aggregation-functions). Without extensions loaded, the only valid value is `APPROX_COUNT_DISTINCT_BUILTIN` (a HyperLogLog, or HLL, based implementation). If the [DataSketches extension](../development/extensions-core/datasketches-extension.md) is loaded, this can also be `APPROX_COUNT_DISTINCT_DS_HLL` (alternative HLL implementation) or `APPROX_COUNT_DISTINCT_DS_THETA`.<br><br>Theta sketches use significantly more memory than HLL sketches, so you should prefer one of the two HLL implementations.|APPROX_COUNT_DISTINCT_BUILTIN|
 
 > Previous versions of Druid had properties named `druid.sql.planner.maxQueryCount` and `druid.sql.planner.maxSemiJoinRowsInMemory`.

sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java

@@ -20,6 +20,9 @@
 package org.apache.druid.sql.calcite.planner;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
@@ -63,22 +66,27 @@
 import org.apache.calcite.tools.RelConversionException;
 import org.apache.calcite.tools.ValidationException;
 import org.apache.calcite.util.Pair;
+import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.guava.BaseSequence;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.query.Query;
 import org.apache.druid.segment.DimensionHandlerUtils;
 import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.Resource;
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.server.security.ResourceType;
 import org.apache.druid.sql.calcite.rel.DruidConvention;
+import org.apache.druid.sql.calcite.rel.DruidQuery;
 import org.apache.druid.sql.calcite.rel.DruidRel;
+import org.apache.druid.sql.calcite.rel.DruidUnionRel;
 import org.apache.druid.sql.calcite.run.QueryMaker;
 import org.apache.druid.sql.calcite.run.QueryMakerFactory;
 
 import javax.annotation.Nullable;
+
 import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -270,7 +278,7 @@ private PlannerResult planWithDruidConvention(
     );
 
     if (explain != null) {
-      return planExplanation(druidRel, explain);
+      return planExplanation(druidRel, explain, true);
     } else {
       final Supplier<Sequence<Object[]>> resultsSupplier = () -> {
         // sanity check
@@ -331,7 +339,7 @@ private PlannerResult planWithBindableConvention(
     }
 
     if (explain != null) {
-      return planExplanation(bindableRel, explain);
+      return planExplanation(bindableRel, explain, false);
     } else {
       final BindableRel theRel = bindableRel;
       final DataContext dataContext = plannerContext.createDataContext(
@@ -380,12 +388,20 @@ public void cleanup(EnumeratorIterator<Object[]> iterFromMake)
    */
   private PlannerResult planExplanation(
       final RelNode rel,
-      final SqlExplain explain
+      final SqlExplain explain,
+      final boolean isDruidConventionExplanation
   )
   {
-    final String explanation = RelOptUtil.dumpPlan("", rel, explain.getFormat(), explain.getDetailLevel());
+    String explanation = RelOptUtil.dumpPlan("", rel, explain.getFormat(), explain.getDetailLevel());
     String resourcesString;
     try {
+      if (isDruidConventionExplanation && rel instanceof DruidRel) {
+        // Show the native queries instead of Calcite's explain if the legacy flag is turned off
+        if (plannerContext.getPlannerConfig().isUseNativeQueryExplain()) {
+          DruidRel<?> druidRel = (DruidRel<?>) rel;
+          explanation = explainSqlPlanAsNativeQueries(druidRel);
+        }
+      }
       final Set<Resource> resources =
           plannerContext.getResourceActions().stream().map(ResourceAction::getResource).collect(Collectors.toSet());
       resourcesString = plannerContext.getJsonMapper().writeValueAsString(resources);
@@ -395,11 +411,58 @@ private PlannerResult planExplanation(
       log.error(jpe, "Encountered exception while serializing Resources for explain output");
       resourcesString = null;
     }
+    catch (ISE ise) {
+      log.error(ise, "Unable to translate to a native Druid query. Resorting to legacy Druid explain plan");
+      resourcesString = null;
+    }
     final Supplier<Sequence<Object[]>> resultsSupplier = Suppliers.ofInstance(
         Sequences.simple(ImmutableList.of(new Object[]{explanation, resourcesString})));
     return new PlannerResult(resultsSupplier, getExplainStructType(rel.getCluster().getTypeFactory()));
   }
 
+  /**
+   * This method doesn't utilize the Calcite's internal {@link RelOptUtil#dumpPlan} since that tends to be verbose
+   * and not indicative of the native Druid Queries which will get executed
+   * This method assumes that the Planner has converted the RelNodes to DruidRels, and thereby we can implictly cast it
+   *
+   * @param rel Instance of the root {@link DruidRel} which is formed by running the planner transformations on it
+   * @return A string representing an array of native queries that correspond to the given SQL query, in JSON format
+   * @throws JsonProcessingException
+   */
+  private String explainSqlPlanAsNativeQueries(DruidRel<?> rel) throws JsonProcessingException
+  {
+    // Only if rel is an instance of DruidUnionRel, do we run multiple native queries corresponding to single SQL query
+    // Also, DruidUnionRel can only be a top level node, so we don't need to check for this condition in the subsequent
+    // child nodes
+    ObjectMapper jsonMapper = plannerContext.getJsonMapper();
+    List<DruidQuery> druidQueryList;
+    if (rel instanceof DruidUnionRel) {
+      druidQueryList = rel.getInputs().stream().map(childRel -> (DruidRel<?>) childRel).map(childRel -> {
+        if (childRel instanceof DruidUnionRel) {
+          log.error("DruidUnionRel can only be the outermost RelNode. This error shouldn't be encountered");
+          throw new ISE("DruidUnionRel is only supported at the outermost RelNode.");
+        }
+        return childRel.toDruidQuery(false);
+      }).collect(Collectors.toList());
+    } else {
+      druidQueryList = ImmutableList.of(rel.toDruidQuery(false));
+    }
+
+    // Putting the queries as object node in an ArrayNode, since directly returning a list causes issues when
+    // serializing the "queryType"
+    ArrayNode nativeQueriesArrayNode = jsonMapper.createArrayNode();
+
+    for (DruidQuery druidQuery : druidQueryList) {
+      Query<?> nativeQuery = druidQuery.getQuery();
+      ObjectNode objectNode = jsonMapper.createObjectNode();
+      objectNode.put("query", jsonMapper.convertValue(nativeQuery, ObjectNode.class));
+      objectNode.put("signature", jsonMapper.convertValue(druidQuery.getOutputRowSignature(), ArrayNode.class));
+      nativeQueriesArrayNode.add(objectNode);
+    }
+
+    return jsonMapper.writeValueAsString(nativeQueriesArrayNode);
+  }
+
   /**
    * This method wraps the root with a {@link LogicalSort} that applies a limit (no ordering change). If the outer rel
    * is already a {@link Sort}, we can merge our outerLimit into it, similar to what is going on in
@@ -409,7 +472,6 @@ private PlannerResult planExplanation(
    * the web console, allowing it to apply a limit to queries without rewriting the original SQL.
    *
    * @param root root node
-   *
    * @return root node wrapped with a limiting logical sort if a limit is specified in the query context.
    */
   @Nullable

sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java

@@ -33,6 +33,7 @@ public class PlannerConfig
   public static final String CTX_KEY_USE_GROUPING_SET_FOR_EXACT_DISTINCT = "useGroupingSetForExactDistinct";
   public static final String CTX_KEY_USE_APPROXIMATE_TOPN = "useApproximateTopN";
   public static final String CTX_COMPUTE_INNER_JOIN_COST_AS_FILTER = "computeInnerJoinCostAsFilter";
+  public static final String CTX_KEY_USE_NATIVE_QUERY_EXPLAIN = "useNativeQueryExplain";
 
   @JsonProperty
   private Period metadataRefreshPeriod = new Period("PT1M");
@@ -70,6 +71,9 @@ public class PlannerConfig
   @JsonProperty
   private boolean authorizeSystemTablesDirectly = false;
 
+  @JsonProperty
+  private boolean useNativeQueryExplain = false;
+
   public long getMetadataSegmentPollPeriod()
   {
     return metadataSegmentPollPeriod;
@@ -137,6 +141,11 @@ public boolean isAuthorizeSystemTablesDirectly()
     return authorizeSystemTablesDirectly;
   }
 
+  public boolean isUseNativeQueryExplain()
+  {
+    return useNativeQueryExplain;
+  }
+
   public PlannerConfig withOverrides(final Map<String, Object> context)
   {
     if (context == null) {
@@ -166,6 +175,11 @@ public PlannerConfig withOverrides(final Map<String, Object> context)
         CTX_COMPUTE_INNER_JOIN_COST_AS_FILTER,
         computeInnerJoinCostAsFilter
     );
+    newConfig.useNativeQueryExplain = getContextBoolean(
+        context,
+        CTX_KEY_USE_NATIVE_QUERY_EXPLAIN,
+        isUseNativeQueryExplain()
+    );
     newConfig.requireTimeCondition = isRequireTimeCondition();
     newConfig.sqlTimeZone = getSqlTimeZone();
     newConfig.awaitInitializationOnStart = isAwaitInitializationOnStart();
@@ -213,7 +227,8 @@ public boolean equals(final Object o)
            metadataSegmentPollPeriod == that.metadataSegmentPollPeriod &&
            serializeComplexValues == that.serializeComplexValues &&
            Objects.equals(metadataRefreshPeriod, that.metadataRefreshPeriod) &&
-           Objects.equals(sqlTimeZone, that.sqlTimeZone);
+           Objects.equals(sqlTimeZone, that.sqlTimeZone) &&
+           useNativeQueryExplain == that.useNativeQueryExplain;
   }
 
   @Override
@@ -230,7 +245,8 @@ public int hashCode()
         sqlTimeZone,
         metadataSegmentCacheEnable,
         metadataSegmentPollPeriod,
-        serializeComplexValues
+        serializeComplexValues,
+        useNativeQueryExplain
     );
   }
 
@@ -248,6 +264,7 @@ public String toString()
            ", metadataSegmentPollPeriod=" + metadataSegmentPollPeriod +
            ", sqlTimeZone=" + sqlTimeZone +
            ", serializeComplexValues=" + serializeComplexValues +
+           ", useNativeQueryExplain=" + useNativeQueryExplain +
            '}';
   }
 }

sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java

@@ -191,6 +191,15 @@ public boolean isAuthorizeSystemTablesDirectly()
     }
   };
 
+  public static final PlannerConfig PLANNER_CONFIG_NATIVE_QUERY_EXPLAIN = new PlannerConfig()
+  {
+    @Override
+    public boolean isUseNativeQueryExplain()
+    {
+      return true;
+    }
+  };
+
   public static final String DUMMY_SQL_ID = "dummy";
   public static final String LOS_ANGELES = "America/Los_Angeles";