Redesign QueryContext class

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java

@@ -29,7 +29,6 @@
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.logger.Logger;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.aggregation.datasketches.hll.sql.HllSketchApproxCountDistinctSqlAggregator;
@@ -516,7 +515,7 @@ public void querySql(Blackhole blackhole) throws Exception
         QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize
     );
     final String sql = QUERIES.get(Integer.parseInt(query));
-    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
+    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, context)) {
       final PlannerResult plannerResult = planner.plan();
       final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
@@ -534,7 +533,7 @@ public void planSql(Blackhole blackhole) throws Exception
         QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize
     );
     final String sql = QUERIES.get(Integer.parseInt(query));
-    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
+    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, context)) {
       final PlannerResult plannerResult = planner.plan();
       blackhole.consume(plannerResult);
     }

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java

@@ -29,7 +29,6 @@
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.query.DruidProcessingConfig;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.segment.QueryableIndex;
@@ -352,7 +351,7 @@ public void querySql(Blackhole blackhole) throws Exception
         QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize
     );
     final String sql = QUERIES.get(Integer.parseInt(query));
-    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
+    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, context)) {
       final PlannerResult plannerResult = planner.plan();
       final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlNestedDataBenchmark.java

@@ -31,7 +31,6 @@
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.query.DruidProcessingConfig;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.expression.TestExprMacroTable;
@@ -318,7 +317,7 @@ public void querySql(Blackhole blackhole) throws Exception
         QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize
     );
     final String sql = QUERIES.get(Integer.parseInt(query));
-    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
+    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, context)) {
       final PlannerResult plannerResult = planner.plan();
       final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java

@@ -26,7 +26,6 @@
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.logger.Logger;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryPlus;
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
@@ -66,6 +65,7 @@
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
+import java.util.Collections;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -167,7 +167,7 @@ public void queryNative(Blackhole blackhole)
   @OutputTimeUnit(TimeUnit.MILLISECONDS)
   public void queryPlanner(Blackhole blackhole) throws Exception
   {
-    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sqlQuery, new QueryContext())) {
+    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sqlQuery, Collections.emptyMap())) {
       final PlannerResult plannerResult = planner.plan();
       final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);

core/src/main/java/org/apache/druid/utils/CollectionUtils.java

@@ -29,9 +29,11 @@
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.Spliterator;
 import java.util.TreeSet;
 import java.util.function.Function;
@@ -148,6 +150,40 @@ public static boolean isNullOrEmpty(@Nullable Collection<?> list)
     return list == null || list.isEmpty();
   }
 
+  /**
+   * Subtract one set from another: {@code C = A - B}.
+   */
+  public static <T> Set<T> subtract(Set<T> left, Set<T> right)
+  {
+    Set<T> result = new HashSet<>(left);
+    result.removeAll(right);
+    return result;
+  }
+
+  /**
+   * Intersection of two sets: {@code C = A ∩ B}.
+   */
+  public static <T> Set<T> intersect(Set<T> left, Set<T> right)
+  {
+    Set<T> result = new HashSet<>();
+    for (T key : left) {
+      if (right.contains(key)) {
+        result.add(key);
+      }
+    }
+    return result;
+  }
+
+  /**
+   * Intersection of two sets: {@code C = A ∪ B}.
+   */
+  public static <T> Set<T> union(Set<T> left, Set<T> right)
+  {
+    Set<T> result = new HashSet<>(left);
+    result.addAll(right);
+    return result;
+  }
+
   private CollectionUtils()
   {
   }

core/src/test/java/org/apache/druid/utils/CollectionUtilsTest.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.utils;
+
+import com.google.common.collect.ImmutableSet;
+import org.junit.Test;
+
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+
+public class CollectionUtilsTest
+{
+  // When Java 9 is allowed, use Set.of().
+  Set<String> empty = ImmutableSet.of();
+  Set<String> abc = ImmutableSet.of("a", "b", "c");
+  Set<String> bcd = ImmutableSet.of("b", "c", "d");
+  Set<String> efg = ImmutableSet.of("e", "f", "g");
+
+  @Test
+  public void testSubtract()
+  {
+    assertEquals(empty, CollectionUtils.subtract(empty, empty));
+    assertEquals(abc, CollectionUtils.subtract(abc, empty));
+    assertEquals(empty, CollectionUtils.subtract(abc, abc));
+    assertEquals(abc, CollectionUtils.subtract(abc, efg));
+    assertEquals(ImmutableSet.of("a"), CollectionUtils.subtract(abc, bcd));
+  }
+
+  @Test
+  public void testIntersect()
+  {
+    assertEquals(empty, CollectionUtils.intersect(empty, empty));
+    assertEquals(abc, CollectionUtils.intersect(abc, abc));
+    assertEquals(empty, CollectionUtils.intersect(abc, efg));
+    assertEquals(ImmutableSet.of("b", "c"), CollectionUtils.intersect(abc, bcd));
+  }
+
+  @Test
+  public void testUnion()
+  {
+    assertEquals(empty, CollectionUtils.union(empty, empty));
+    assertEquals(abc, CollectionUtils.union(abc, abc));
+    assertEquals(ImmutableSet.of("a", "b", "c", "e", "f", "g"), CollectionUtils.union(abc, efg));
+    assertEquals(ImmutableSet.of("a", "b", "c", "d"), CollectionUtils.union(abc, bcd));
+  }
+}

extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java

@@ -28,7 +28,6 @@
 import org.apache.druid.query.BaseQuery;
 import org.apache.druid.query.DataSource;
 import org.apache.druid.query.Query;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryRunner;
 import org.apache.druid.query.QuerySegmentWalker;
 import org.apache.druid.query.filter.DimFilter;
@@ -41,6 +40,7 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -146,12 +146,6 @@ public Map<String, Object> getContext()
     return query.getContext();
   }
 
-  @Override
-  public QueryContext getQueryContext()
-  {
-    return query.getQueryContext();
-  }
-
   @Override
   public boolean isDescending()
   {

extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java

@@ -121,7 +121,6 @@ public void testGetContextHumanReadableBytes()
         .postAggregators(QueryRunnerTestHelper.ADD_ROWS_INDEX_CONSTANT)
         .build();
     MaterializedViewQuery query = new MaterializedViewQuery(topNQuery, optimizer);
-    Assert.assertEquals(20_000_000, query.getContextAsHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes());
-
+    Assert.assertEquals(20_000_000, query.context().getHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes());
   }
 }

extensions-contrib/moving-average-query/src/main/java/org/apache/druid/query/movingaverage/MovingAverageQuery.java

@@ -237,7 +237,7 @@ public String getType()
   @JsonIgnore
   public boolean getContextSortByDimsFirst()
   {
-    return getContextBoolean(CTX_KEY_SORT_BY_DIMS_FIRST, false);
+    return context().getBoolean(CTX_KEY_SORT_BY_DIMS_FIRST, false);
   }
 
   @Override

extensions-contrib/moving-average-query/src/main/java/org/apache/druid/query/movingaverage/MovingAverageQueryRunner.java

@@ -30,7 +30,6 @@
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.query.DataSource;
-import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryDataSource;
 import org.apache.druid.query.QueryPlus;
 import org.apache.druid.query.QueryRunner;
@@ -52,6 +51,7 @@
 import org.joda.time.Period;
 
 import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -124,7 +124,7 @@ public Sequence<Row> run(QueryPlus<Row> query, ResponseContext responseContext)
       ResponseContext gbqResponseContext = ResponseContext.createEmpty();
       gbqResponseContext.merge(responseContext);
       gbqResponseContext.putQueryFailDeadlineMs(
-          System.currentTimeMillis() + QueryContexts.getTimeout(gbq)
+          System.currentTimeMillis() + gbq.context().getTimeout()
       );
 
       Sequence<ResultRow> results = gbq.getRunner(walker).run(QueryPlus.wrap(gbq), gbqResponseContext);
@@ -164,7 +164,7 @@ public Sequence<Row> run(QueryPlus<Row> query, ResponseContext responseContext)
       ResponseContext tsqResponseContext = ResponseContext.createEmpty();
       tsqResponseContext.merge(responseContext);
       tsqResponseContext.putQueryFailDeadlineMs(
-          System.currentTimeMillis() + QueryContexts.getTimeout(tsq)
+          System.currentTimeMillis() + tsq.context().getTimeout()
       );
 
       Sequence<Result<TimeseriesResultValue>> results = tsq.getRunner(walker).run(QueryPlus.wrap(tsq), tsqResponseContext);

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchApproxQuantileSqlAggregator.java

@@ -49,6 +49,7 @@
 import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
 
 import javax.annotation.Nullable;
+
 import java.util.List;
 
 public class DoublesSketchApproxQuantileSqlAggregator implements SqlAggregator
@@ -171,7 +172,7 @@ public Aggregation toDruidAggregation(
           histogramName,
           input.getDirectColumn(),
           k,
-          getMaxStreamLengthFromQueryContext(plannerContext.getQueryContext())
+          getMaxStreamLengthFromQueryContext(plannerContext.queryContext())
       );
     } else {
       String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(
@@ -182,7 +183,7 @@ public Aggregation toDruidAggregation(
           histogramName,
           virtualColumnName,
           k,
-          getMaxStreamLengthFromQueryContext(plannerContext.getQueryContext())
+          getMaxStreamLengthFromQueryContext(plannerContext.queryContext())
       );
     }
 
@@ -201,7 +202,7 @@ public Aggregation toDruidAggregation(
 
   static long getMaxStreamLengthFromQueryContext(QueryContext queryContext)
   {
-    return queryContext.getAsLong(
+    return queryContext.getLong(
         CTX_APPROX_QUANTILE_DS_MAX_STREAM_LENGTH,
         DoublesSketchAggregatorFactory.DEFAULT_MAX_STREAM_LENGTH
     );

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchObjectSqlAggregator.java

@@ -46,6 +46,7 @@
 import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
 
 import javax.annotation.Nullable;
+
 import java.util.List;
 
 public class DoublesSketchObjectSqlAggregator implements SqlAggregator
@@ -113,7 +114,7 @@ public Aggregation toDruidAggregation(
           histogramName,
           input.getDirectColumn(),
           k,
-          DoublesSketchApproxQuantileSqlAggregator.getMaxStreamLengthFromQueryContext(plannerContext.getQueryContext())
+          DoublesSketchApproxQuantileSqlAggregator.getMaxStreamLengthFromQueryContext(plannerContext.queryContext())
       );
     } else {
       String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(
@@ -124,7 +125,7 @@ public Aggregation toDruidAggregation(
           histogramName,
           virtualColumnName,
           k,
-          DoublesSketchApproxQuantileSqlAggregator.getMaxStreamLengthFromQueryContext(plannerContext.getQueryContext())
+          DoublesSketchApproxQuantileSqlAggregator.getMaxStreamLengthFromQueryContext(plannerContext.queryContext())
       );
     }
 
@@ -136,7 +137,6 @@ public Aggregation toDruidAggregation(
 
   private static class DoublesSketchSqlAggFunction extends SqlAggFunction
   {
-    private static final String SIGNATURE1 = "'" + NAME + "(column)'\n";
     private static final String SIGNATURE2 = "'" + NAME + "(column, k)'\n";
 
     DoublesSketchSqlAggFunction()

extensions-core/histogram/src/test/java/org/apache/druid/query/aggregation/histogram/sql/FixedBucketsHistogramQuantileSqlAggregatorTest.java

@@ -27,6 +27,7 @@
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.query.Druids;
+import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryDataSource;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
 import org.apache.druid.query.aggregation.DoubleSumAggregatorFactory;
@@ -53,7 +54,6 @@
 import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
 import org.apache.druid.sql.calcite.filtration.Filtration;
 import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.util.CalciteTests;
 import org.apache.druid.sql.calcite.util.SpecificSegmentsQuerySegmentWalker;
 import org.apache.druid.timeline.DataSegment;
@@ -324,7 +324,7 @@ public void testQuantileOnCastedString()
                       new QuantilePostAggregator("a6", "a6:agg", 0.999f),
                       new QuantilePostAggregator("a7", "a5:agg", 0.999f)
                   )
-                  .context(ImmutableMap.of(PlannerContext.CTX_SQL_QUERY_ID, "dummy"))
+                  .context(ImmutableMap.of(QueryContexts.CTX_SQL_QUERY_ID, "dummy"))
                   .build()
         ),
         ImmutableList.of(