diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 6f24346c20e7..bdc1eb8eb4a0 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -70,13 +70,22 @@ public class QueryContext @JsonCreator public QueryContext(Map 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 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() @@ -850,7 +859,17 @@ public QueryContext override(Map contextOverride) if (contextOverride == null || contextOverride.isEmpty()) { return this; } - return QueryContext.of(QueryContexts.override(asMap(), contextOverride)); + + final TreeMap 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 @@ public QueryContext override(QueryContext queryContext) 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 QueryContext override(final QueryContextParameter parameter, @Nullable final T value) + { + final TreeMap overridden = new TreeMap<>(context); + parameter.set(overridden, value); + return new QueryContext(overridden, false); + } + public boolean isPrePlanned() { return getOrDefault(QueryContextParameters.PREPLANNED); diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java index a61cd53a230e..d758c3ac478d 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java +++ b/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 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() { diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlParser.java b/sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlParser.java index da2fccb62d41..a4ff4544b105 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlParser.java +++ b/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,9 +99,8 @@ private static StatementAndSetContext processStatementList( final boolean allowSetStatements ) { - if (root instanceof SqlNodeList) { + if (root instanceof SqlNodeList nodeList) { final Map 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"); } @@ -110,8 +108,7 @@ private static StatementAndSetContext processStatementList( // 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]", diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java index e5b753490588..529fc0fff01b 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java +++ b/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); } } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/IngestHandler.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/IngestHandler.java index d7e5864e5dea..b9c4827b742b 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/IngestHandler.java +++ b/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 + ); } } 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 0e8c596befd4..15a8ae65b55a 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 @@ -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 queryContext) private final DruidSchemaCatalog rootSchema; private final AuthenticationResult authenticationResult; private final Set authContextKeys; - private final Map queryContext; + /** + * Immutable; replaced, never modified in place, when planning records a value with + * {@link #setQueryContextParameter}. + */ + private volatile QueryContext queryContext; private final CopyOnWriteArrayList nativeQueryIds = new CopyOnWriteArrayList<>(); private final PlannerHook hook; private final Set 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 queryContextMap() { - return queryContext; + return queryContext.asMap(); } /** @@ -471,7 +473,19 @@ public Set 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 void setQueryContextParameter(final QueryContextParameter 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 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 context; DruidDataContext() { - ImmutableMap.Builder builder = ImmutableMap.builder(); - builder.putAll(base_context); + ImmutableMap.Builder builder = ImmutableMap.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 toAdd) - { - this.queryContext.putAll(toAdd); - initializeContextFieldsAndPlannerConfig(); - } - public SqlEngine getEngine() { return engine; @@ -787,11 +786,9 @@ public void dispatchHook(HookKey key, T object) plannerToolbox.getHookDispatcher().dispatch(key, object); } - - - private void initializeContextFieldsAndPlannerConfig() + private QueryContext initializeContextFieldsAndPlannerConfig(final Map 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; } }