refactor(sql): reduce copies of query context in PlannerContext

processing/src/main/java/org/apache/druid/query/QueryContext.java

@@ -70,13 +70,22 @@
 
   @JsonCreator
   public QueryContext(Map<String, Object> context)
+  {
+    this(context, true);
+  }
+
+  /**
+   * @param copy whether to copy {@code context}.
+   *             Pass false only for a {@link TreeMap} this class has just built and will not modify afterwards.
+   */
+  private QueryContext(@Nullable final Map<String, Object> context, final boolean copy)
   {
     // There is no semantic difference between an empty and a null context.
     // Ensure that a context always exists to avoid the need to check for
     // a null context. Jackson serialization will omit empty contexts.
     this.context = context == null
         ? Collections.emptyMap()
-        : Collections.unmodifiableMap(new TreeMap<>(context));
+        : Collections.unmodifiableMap(copy ? new TreeMap<>(context) : context);
   }
 
   public static QueryContext empty()
@@ -235,7 +244,7 @@
     return context.isEmpty();
   }
 
   public Map<String, Object> asMap()
   {
     return context;
   }
@@ -850,7 +859,17 @@
     if (contextOverride == null || contextOverride.isEmpty()) {
       return this;
     }
-    return QueryContext.of(QueryContexts.override(asMap(), contextOverride));
+
+    final TreeMap<String, Object> overridden = new TreeMap<>(context);
+    // A null override unsets the key, like QueryContexts.override(Map, Map).
+    contextOverride.forEach((key, value) -> {
+      if (value == null) {
+        overridden.remove(key);
+      } else {
+        overridden.put(key, value);
+      }
+    });
+    return new QueryContext(overridden, false);
   }
 
   public QueryContext override(QueryContext queryContext)
@@ -861,6 +880,18 @@
     return override(queryContext.asMap());
   }
 
+  /**
+   * Returns a new context with one typed parameter inserted, updated or, if {@code value} is null, removed. Leaves this context unchanged.
+   *
+   * @throws BadQueryContextException if the value is invalid
+   */
+  public <T> QueryContext override(final QueryContextParameter<T> parameter, @Nullable final T value)
+  {
+    final TreeMap<String, Object> overridden = new TreeMap<>(context);
+    parameter.set(overridden, value);
+    return new QueryContext(overridden, false);
+  }
+
   public boolean isPrePlanned()
   {
     return getOrDefault(QueryContextParameters.PREPLANNED);

processing/src/test/java/org/apache/druid/query/QueryContextTest.java

@@ -125,6 +125,67 @@ public void testOfTypedParameters()
     );
   }
 
+  @Test
+  public void testOverrideTypedParameter()
+  {
+    final QueryContext original = QueryContext.of(
+        ImmutableMap.of(
+            "k", "v",
+            QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING.getName(), 10
+        )
+    );
+
+    assertEquals(
+        ImmutableMap.of(
+            "k", "v",
+            QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING.getName(), 20,
+            QueryContextParameters.USE_RESULT_LEVEL_CACHE.getName(), false
+        ),
+        original.override(QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING, 20)
+                .override(QueryContextParameters.USE_RESULT_LEVEL_CACHE, false)
+                .asMap()
+    );
+    assertEquals(
+        ImmutableMap.of("k", "v"),
+        original.override(QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING, null).asMap()
+    );
+    assertThrows(
+        UnsupportedOperationException.class,
+        () -> original.override(QueryContextParameters.USE_RESULT_LEVEL_CACHE, false).asMap().put("k", "v2")
+    );
+    assertThrows(
+        BadQueryContextException.class,
+        () -> original.override(QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING, 0)
+    );
+
+    // The original context is unchanged.
+    assertEquals(
+        ImmutableMap.of("k", "v", QueryContextParameters.MAX_ROWS_QUEUED_FOR_ORDERING.getName(), 10),
+        original.asMap()
+    );
+  }
+
+  @Test
+  public void testOverrideKeepsKeysSorted()
+  {
+    // A HashMap iterates these keys out of order, so the assertions fail unless the result is sorted.
+    final QueryContext original = QueryContext.of(ImmutableMap.of("sqlQueryId", 1, "priority", 2, "lane", 3));
+    final Map<String, Object> overrides = new HashMap<>();
+    overrides.put("timeout", 4);
+    overrides.put("finalize", 5);
+    overrides.put("vectorize", 6);
+    overrides.put("lane", null);
+
+    assertEquals(
+        List.of("finalize", "priority", "sqlQueryId", "timeout", "vectorize"),
+        List.copyOf(original.override(overrides).asMap().keySet())
+    );
+    assertEquals(
+        List.of("lane", "priority", "queryId", "sqlQueryId"),
+        List.copyOf(original.override(QueryContext.of(ImmutableMap.of("queryId", 7))).asMap().keySet())
+    );
+  }
+
   @Test
   public void testIsEmpty()
   {

sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlParser.java

@@ -88,9 +88,8 @@ public static StatementAndSetContext parse(final String sql, final boolean allow
 
   /**
    * 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
+   * single {@link SqlNode} which is NOT a {@link SqlSetOption}. All {@link SqlSetOption} will be converted into a context parameters {@link 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.
@@ -100,18 +99,16 @@ private static StatementAndSetContext processStatementList(
       final boolean allowSetStatements
   )
   {
-    if (root instanceof SqlNodeList) {
+    if (root instanceof SqlNodeList nodeList) {
       final Map<String, Object> setContext = new LinkedHashMap<>();
-      final SqlNodeList nodeList = (SqlNodeList) root;
       if (!allowSetStatements && nodeList.size() > 1) {
         throw InvalidSqlInput.exception("SQL query string must contain only a single statement");
       }
       boolean isMissingDruidStatementNode = true;
       // convert 0 or more SET statements into a Map of stuff to add to the query context
       for (int i = 0; i < nodeList.size(); i++) {
         SqlNode sqlNode = nodeList.get(i);
-        if (sqlNode instanceof SqlSetOption) {
-          final SqlSetOption sqlSetOption = (SqlSetOption) sqlNode;
+        if (sqlNode instanceof SqlSetOption sqlSetOption) {
           if (!(sqlSetOption.getValue() instanceof SqlLiteral)) {
             throw InvalidSqlInput.exception(
                 "Assigned value must be a literal for SET statement[%s]",

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

@@ -209,8 +209,8 @@ public void validateInsert(final SqlInsert insert)
       // during conversion, however, we've just worked out the granularity, so we
       // do it here instead.
       try {
-        QueryContextParameters.SQL_INSERT_SEGMENT_GRANULARITY.set(
-            plannerContext.queryContextMap(),
+        plannerContext.setQueryContextParameter(
+            QueryContextParameters.SQL_INSERT_SEGMENT_GRANULARITY,
             plannerContext.getPlannerToolbox().jsonMapper().writeValueAsString(effectiveGranularity)
         );
       }
@@ -231,8 +231,7 @@ public void validateInsert(final SqlInsert insert)
     // know names and we match by name.) Thus, we'd have to validate (to know names and types)
     // to get the target types, but we need the target types to validate. Catch-22. So, we punt.
     final SqlValidatorScope scope;
-    if (source instanceof SqlSelect) {
-      final SqlSelect sqlSelect = (SqlSelect) source;
+    if (source instanceof SqlSelect sqlSelect) {
       validateSelect(sqlSelect, unknownType);
       scope = null;
     } else {
@@ -267,7 +266,7 @@ public void validateInsert(final SqlInsert insert)
     if (tableMetadata != null && !plannerContext.queryContext().has(QueryContextParameters.ROWS_PER_SEGMENT)) {
       final Integer targetSegmentRows = tableMetadata.targetSegmentRows();
       if (targetSegmentRows != null) {
-        QueryContextParameters.ROWS_PER_SEGMENT.set(plannerContext.queryContextMap(), targetSegmentRows);
+        plannerContext.setQueryContextParameter(QueryContextParameters.ROWS_PER_SEGMENT, targetSegmentRows);
       }
     }
   }

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

@@ -128,8 +128,8 @@ private void validateExport()
           operationName()
       );
     } else {
-      QueryContextParameters.EXPORT_FILE_FORMAT.set(
-          handlerContext.plannerContext().queryContextMap(),
+      handlerContext.plannerContext().setQueryContextParameter(
+          QueryContextParameters.EXPORT_FILE_FORMAT,
           exportFileFormat.toString()
       );
     }
@@ -407,7 +407,10 @@ public void validate()
       );
       if (replaceIntervalsList != null) {
         replaceIntervals = String.join(",", replaceIntervalsList);
-        QueryContextParameters.SQL_REPLACE_TIME_CHUNKS.set(handlerContext.queryContextMap(), replaceIntervals);
+        handlerContext.plannerContext().setQueryContextParameter(
+            QueryContextParameters.SQL_REPLACE_TIME_CHUNKS,
+            replaceIntervals
+        );
       }
     }
 

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

@@ -74,7 +74,6 @@
 import javax.annotation.Nullable;
 import java.util.Collections;
 import java.util.HashSet;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -215,7 +214,11 @@ public static JoinAlgorithm getJoinAlgorithm(Map<String, Object> queryContext)
   private final DruidSchemaCatalog rootSchema;
   private final AuthenticationResult authenticationResult;
   private final Set<String> authContextKeys;
-  private final Map<String, Object> queryContext;
+  /**
+   * Immutable; replaced, never modified in place, when planning records a value with
+   * {@link #setQueryContextParameter}.
+   */
+  private volatile QueryContext queryContext;
   private final CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>();
   private final PlannerHook hook;
   private final Set<String> lookupsToLoad = new HashSet<>();
@@ -282,9 +285,8 @@ private PlannerContext(
     this.rootSchema = rootSchema;
     this.authenticationResult = authenticationResult;
     this.authContextKeys = authContextKeys;
-    this.queryContext = new LinkedHashMap<>(queryContext);
+    this.queryContext = initializeContextFieldsAndPlannerConfig(queryContext);
     this.hook = hook == null ? NoOpPlannerHook.INSTANCE : hook;
-    initializeContextFieldsAndPlannerConfig();
   }
 
   public static PlannerContext create(
@@ -449,12 +451,12 @@ public LookupLoadingSpec getLookupLoadingSpec()
   }
 
   /**
-   * Return the query context as a mutable map. Use this form when
-   * modifying the context during planning.
+   * Return the query context as an immutable map. The map is a snapshot: it does not reflect values recorded later
+   * with {@link #setQueryContextParameter}. Use that method to modify the context during planning.
    */
   public Map<String, Object> queryContextMap()
   {
-    return queryContext;
+    return queryContext.asMap();
   }
 
   /**
@@ -471,7 +473,19 @@ public Set<String> authContextKeys()
    */
   public QueryContext queryContext()
   {
-    return QueryContext.of(queryContext);
+    return queryContext;
+  }
+
+  /**
+   * Record a value computed during planning, such as the segment granularity of an INSERT, so that later planning
+   * steps and the generated query or task see it. A null value removes the parameter.
+   *
+   * Fields derived from the context when this object is created, such as the time zone or {@link PlannerConfig},
+   * are not recomputed, so do not use this method for parameters that feed them.
+   */
+  public <T> void setQueryContextParameter(final QueryContextParameter<T> parameter, @Nullable final T value)
+  {
+    queryContext = queryContext.override(parameter, value);
   }
 
   public boolean isStringifyArrays()
@@ -547,7 +561,7 @@ public AuthenticationResult getAuthenticationResult()
 
   public JoinAlgorithm getJoinAlgorithm()
   {
-    return QueryContext.of(queryContext).getOrDefault(QueryContextParameters.SQL_JOIN_ALGORITHM);
+    return queryContext.getOrDefault(QueryContextParameters.SQL_JOIN_ALGORITHM);
   }
 
   public String getSql()
@@ -604,21 +618,15 @@ public DataContext createDataContext(final JavaTypeFactory typeFactory, List<Typ
   {
     class DruidDataContext implements DataContext
     {
-      private final Map<String, Object> base_context = ImmutableMap.of(
-          DataContext.Variable.UTC_TIMESTAMP.camelName, localNow.getMillis(),
-          DataContext.Variable.CURRENT_TIMESTAMP.camelName, localNow.getMillis(),
-          DataContext.Variable.LOCAL_TIMESTAMP.camelName, new Interval(
-              new DateTime("1970-01-01T00:00:00.000", localNow.getZone()),
-              localNow
-          ).toDurationMillis(),
-          DataContext.Variable.TIME_ZONE.camelName, localNow.getZone().toTimeZone().clone()
-      );
       private final Map<String, Object> context;
 
       DruidDataContext()
       {
-        ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
-        builder.putAll(base_context);
+        ImmutableMap.Builder<String, Object> builder = ImmutableMap.<String, Object>builder()
+            .put(Variable.UTC_TIMESTAMP.camelName, localNow.getMillis())
+            .put(Variable.CURRENT_TIMESTAMP.camelName, localNow.getMillis())
+            .put(Variable.LOCAL_TIMESTAMP.camelName, new Interval(new DateTime("1970-01-01T00:00:00.000", localNow.getZone()), localNow).toDurationMillis())
+            .put(Variable.TIME_ZONE.camelName, localNow.getZone().toTimeZone().clone());
         int i = 0;
         for (TypedValue parameter : parameters) {
           builder.put("?" + i, parameter.value);
@@ -704,15 +712,6 @@ public void setQueryMaker(QueryMaker queryMaker)
     this.queryMaker = Preconditions.checkNotNull(queryMaker, "queryMaker");
   }
 
-  /**
-   * Add additional query context parameters, overriding any existing values.
-   */
-  public void addAllToQueryContext(Map<String, Object> toAdd)
-  {
-    this.queryContext.putAll(toAdd);
-    initializeContextFieldsAndPlannerConfig();
-  }
-
   public SqlEngine getEngine()
   {
     return engine;
@@ -787,11 +786,9 @@ public <T> void dispatchHook(HookKey<T> key, T object)
     plannerToolbox.getHookDispatcher().dispatch(key, object);
   }
 
-
-
-  private void initializeContextFieldsAndPlannerConfig()
+  private QueryContext initializeContextFieldsAndPlannerConfig(final Map<String, Object> queryContext)
   {
-    final QueryContext context = QueryContext.of(queryContext);
+    QueryContext context = QueryContext.of(queryContext);
     final String tsParam = context.get(QueryContextParameters.SQL_CURRENT_TIMESTAMP);
     final DateTime utcNow = tsParam != null ? DateTimes.of(tsParam) : DateTimes.nowUtc();
 
@@ -816,13 +813,15 @@ private void initializeContextFieldsAndPlannerConfig()
     // special handling for DruidViewMacro, normal client will allocate sqlid in SqlLifecyle
     if (Strings.isNullOrEmpty(sqlQueryId)) {
       sqlQueryId = UUID.randomUUID().toString();
-      QueryContextParameters.SQL_QUERY_ID.set(this.queryContext, sqlQueryId);
+      context = context.override(QueryContextParameters.SQL_QUERY_ID, sqlQueryId);
     }
 
     if (plannerConfig != null) {
-      plannerConfig = plannerConfig.withOverrides(queryContext);
+      plannerConfig = plannerConfig.withOverrides(context.asMap());
     } else {
-      plannerConfig = getPlannerToolbox().plannerConfig.withOverrides(queryContext);
+      plannerConfig = getPlannerToolbox().plannerConfig.withOverrides(context.asMap());
     }
+
+    return context;
   }
 }