feat: add realtimeSegmentsMode query context param

Author: jtuglu1

docs/querying/query-context-reference.md

@@ -71,7 +71,8 @@ Unless otherwise noted, the following parameters apply to all query types, and t
 |`setProcessingThreadNames`|`true`| Whether processing thread names will be set to `queryType_dataSource_intervals` while processing a query. This aids in interpreting thread dumps, and is on by default. Query overhead can be reduced slightly by setting this to `false`. This has a tiny effect in most scenarios, but can be meaningful in high-QPS, low-per-segment-processing-time scenarios. |
 |`sqlPlannerBloat`|`1000`|Calcite parameter which controls whether to merge two Project operators when inlining expressions causes complexity to increase. Implemented as a workaround to exception `There are not enough rules to produce a node with desired properties: convention=DRUID, sort=[]` thrown after rejecting the merge of two projects.|
 |`cloneQueryMode`|`excludeClones`| Indicates whether clone Historicals should be queried by brokers. Clone servers are created by the `cloneServers` Coordinator dynamic configuration. Possible values are `excludeClones`, `includeClones` and `preferClones`. `excludeClones` means that clone Historicals are not queried by the broker. `preferClones` indicates that when given a choice between the clone Historical and the original Historical which is being cloned, the broker chooses the clones. Historicals which are not involved in the cloning process will still be queried. `includeClones` means that broker queries any Historical without regarding clone status. This parameter only affects native queries. MSQ does not query Historicals directly.|
-|`realtimeSegmentsOnly` |`false`| When set to true, only query realtime segments. Historical segments are excluded. |
+|`realtimeSegmentsMode` |`include`| Controls which segment types are queried. `include` queries both realtime and historical segments. `exclusive` queries only realtime segments (historical segments are excluded). `exclude` queries only historical segments (realtime segments are excluded). |
+|`realtimeSegmentsOnly` |`false`| **Deprecated.** Use `realtimeSegmentsMode=exclusive` instead. When set to `true`, only realtime segments are queried; equivalent to `realtimeSegmentsMode=exclusive`. |
 
 ## Parameters by query type
 

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

@@ -27,6 +27,7 @@
 import org.apache.druid.java.util.common.HumanReadableBytes;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularity;
+import org.apache.druid.query.QueryContexts.RealtimeSegmentsMode;
 import org.apache.druid.query.QueryContexts.Vectorize;
 import org.apache.druid.query.filter.InDimFilter;
 import org.apache.druid.query.filter.TypedInFilter;
@@ -781,6 +782,32 @@ public boolean isPrePlanned()
     return getBoolean(QueryContexts.CTX_PREPLANNED, QueryContexts.DEFAULT_PREPLANNED);
   }
 
+  /**
+   * Returns the realtime segments mode for this query. If {@code realtimeSegmentsMode} is not set,
+   * falls back to the deprecated {@code realtimeSegmentsOnly} boolean: {@code true} maps to
+   * {@link RealtimeSegmentsMode#EXCLUSIVE}.
+   */
+  public RealtimeSegmentsMode getRealtimeSegmentsMode()
+  {
+    RealtimeSegmentsMode mode = getEnum(
+        QueryContexts.REALTIME_SEGMENTS_MODE,
+        RealtimeSegmentsMode.class,
+        null
+    );
+    if (mode != null) {
+      return mode;
+    }
+    // Backward-compat: honour the deprecated realtimeSegmentsOnly flag.
+    if (getBoolean(QueryContexts.REALTIME_SEGMENTS_ONLY, QueryContexts.DEFAULT_REALTIME_SEGMENTS_ONLY)) {
+      return RealtimeSegmentsMode.EXCLUSIVE;
+    }
+    return QueryContexts.DEFAULT_REALTIME_SEGMENTS_MODE;
+  }
+
+  /**
+   * @deprecated Use {@link #getRealtimeSegmentsMode()} instead.
+   */
+  @Deprecated
   public boolean isRealtimeSegmentsOnly()
   {
     return getBoolean(QueryContexts.REALTIME_SEGMENTS_ONLY, QueryContexts.DEFAULT_REALTIME_SEGMENTS_ONLY);

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

@@ -146,9 +146,20 @@ public class QueryContexts
   public static final String NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED = "COUPLED";
   public static final String NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED = "DECOUPLED";
 
+  /**
+   * @deprecated Use {@link #REALTIME_SEGMENTS_MODE} instead.
+   */
+  @Deprecated
   public static final String REALTIME_SEGMENTS_ONLY = "realtimeSegmentsOnly";
+  /**
+   * @deprecated Use {@link #DEFAULT_REALTIME_SEGMENTS_MODE} instead.
+   */
+  @Deprecated
   public static final boolean DEFAULT_REALTIME_SEGMENTS_ONLY = false;
 
+  public static final String REALTIME_SEGMENTS_MODE = "realtimeSegmentsMode";
+  public static final RealtimeSegmentsMode DEFAULT_REALTIME_SEGMENTS_MODE = RealtimeSegmentsMode.INCLUDE;
+
   public static final String CTX_PREPLANNED = "prePlanned";
   public static final boolean DEFAULT_PREPLANNED = true;
 
@@ -233,6 +244,29 @@ public String toString()
     }
   }
 
+  public enum RealtimeSegmentsMode
+  {
+    /** Include both realtime and historical segments (default). */
+    INCLUDE,
+    /** Query only realtime segments; historical segments are excluded. */
+    EXCLUSIVE,
+    /** Query only historical segments; realtime segments are excluded. */
+    EXCLUDE;
+
+    @JsonCreator
+    public static RealtimeSegmentsMode fromString(String str)
+    {
+      return RealtimeSegmentsMode.valueOf(StringUtils.toUpperCase(str));
+    }
+
+    @Override
+    @JsonValue
+    public String toString()
+    {
+      return StringUtils.toLowerCase(name());
+    }
+  }
+
   private QueryContexts()
   {
   }

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

@@ -430,6 +430,7 @@ public void testNonLegacyIsNotLegacyContext()
   }
 
   @Test
+  @SuppressWarnings("deprecation")
   public void testIsRealtimeSegmentsOnly()
   {
     assertFalse(QueryContext.empty().isRealtimeSegmentsOnly());
@@ -440,6 +441,55 @@ public void testIsRealtimeSegmentsOnly()
     );
   }
 
+  @Test
+  public void testGetRealtimeSegmentsMode()
+  {
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.INCLUDE,
+        QueryContext.empty().getRealtimeSegmentsMode()
+    );
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.EXCLUSIVE,
+        QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "exclusive"))
+                   .getRealtimeSegmentsMode()
+    );
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.EXCLUDE,
+        QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "exclude"))
+                   .getRealtimeSegmentsMode()
+    );
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.INCLUDE,
+        QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "include"))
+                   .getRealtimeSegmentsMode()
+    );
+  }
+
+  @Test
+  public void testGetRealtimeSegmentsModeBackwardCompat()
+  {
+    // realtimeSegmentsOnly=true maps to EXCLUSIVE
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.EXCLUSIVE,
+        QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_ONLY, true))
+                   .getRealtimeSegmentsMode()
+    );
+    // realtimeSegmentsOnly=false maps to INCLUDE (default)
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.INCLUDE,
+        QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_ONLY, false))
+                   .getRealtimeSegmentsMode()
+    );
+    // realtimeSegmentsMode takes precedence over realtimeSegmentsOnly
+    assertEquals(
+        QueryContexts.RealtimeSegmentsMode.EXCLUDE,
+        QueryContext.of(ImmutableMap.of(
+            QueryContexts.REALTIME_SEGMENTS_ONLY, true,
+            QueryContexts.REALTIME_SEGMENTS_MODE, "exclude"
+        )).getRealtimeSegmentsMode()
+    );
+  }
+
   @Test
   public void testSerialization() throws Exception
   {

server/src/main/java/org/apache/druid/client/CachingClusteredClient.java

@@ -61,6 +61,7 @@
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryContexts.RealtimeSegmentsMode;
 import org.apache.druid.query.QueryMetrics;
 import org.apache.druid.query.QueryPlus;
 import org.apache.druid.query.QueryRunner;
@@ -444,7 +445,7 @@ private Set<SegmentServerSelector> computeSegmentsToQuery(
       final Set<SegmentServerSelector> segments = new LinkedHashSet<>();
       final SegmentPruner segmentPruner = ev.getSegmentPruner();
 
-      boolean isRealtimeSegmentOnly = query.context().isRealtimeSegmentsOnly();
+      RealtimeSegmentsMode realtimeSegmentsMode = query.context().getRealtimeSegmentsMode();
       // Filter unneeded chunks based on partition dimension
       for (TimelineObjectHolder<String, ServerSelector> holder : serversLookup) {
         final Collection<PartitionChunk<ServerSelector>> filteredChunks;
@@ -458,8 +459,11 @@ private Set<SegmentServerSelector> computeSegmentsToQuery(
         }
         for (PartitionChunk<ServerSelector> chunk : filteredChunks) {
           ServerSelector server = chunk.getObject();
-          if (isRealtimeSegmentOnly && !server.isRealtimeSegment()) {
-            continue; // Skip historical segments when only realtime segments are requested
+          if (realtimeSegmentsMode == RealtimeSegmentsMode.EXCLUSIVE && !server.isRealtimeSegment()) {
+            continue;
+          }
+          if (realtimeSegmentsMode == RealtimeSegmentsMode.EXCLUDE && server.isRealtimeSegment()) {
+            continue;
           }
           final SegmentDescriptor segment = new SegmentDescriptor(
               holder.getInterval(),

server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java

@@ -3132,27 +3132,94 @@ public void testRealtimeSegmentsQueryContext()
     selector.addServerAndUpdateSegment(new QueryableDruidServer(servers[0], null), dataSegment);
     timeline.add(interval, "ver", new SingleElementPartitionChunk<>(selector));
 
-    final TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder()
+    // include (default): historical segment is included
+    final TimeBoundaryQuery queryInclude = Druids.newTimeBoundaryQueryBuilder()
             .dataSource(DATA_SOURCE)
             .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(queryInterval)))
-            .context(ImmutableMap.of("realtimeSegmentsOnly", false))
+            .context(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "include"))
             .randomQueryId()
             .build();
 
-    final TimeBoundaryQuery query2 = Druids.newTimeBoundaryQueryBuilder()
+    // exclusive: only realtime segments — historical segment is excluded
+    final TimeBoundaryQuery queryExclusive = Druids.newTimeBoundaryQueryBuilder()
+            .dataSource(DATA_SOURCE)
+            .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(queryInterval)))
+            .context(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "exclusive"))
+            .randomQueryId()
+            .build();
+
+    // backward compat: realtimeSegmentsOnly=true maps to EXCLUSIVE
+    final TimeBoundaryQuery queryLegacyTrue = Druids.newTimeBoundaryQueryBuilder()
             .dataSource(DATA_SOURCE)
             .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(queryInterval)))
-            .context(ImmutableMap.of("realtimeSegmentsOnly", true))
+            .context(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_ONLY, true))
             .randomQueryId()
             .build();
 
     final ResponseContext responseContext = initializeResponseContext();
 
-    getDefaultQueryRunner().run(QueryPlus.wrap(query), responseContext);
-    getDefaultQueryRunner().run(QueryPlus.wrap(query2), responseContext);
+    getDefaultQueryRunner().run(QueryPlus.wrap(queryInclude), responseContext);
+    getDefaultQueryRunner().run(QueryPlus.wrap(queryExclusive), responseContext);
+    getDefaultQueryRunner().run(QueryPlus.wrap(queryLegacyTrue), responseContext);
+
+    final Map<String, Integer> remainingResponseMap = (Map<String, Integer>) responseContext.get(ResponseContext.Keys.REMAINING_RESPONSES_FROM_QUERY_SERVERS);
+    Assert.assertEquals(1, remainingResponseMap.get(queryInclude.getId()).intValue());
+    Assert.assertEquals(0, remainingResponseMap.get(queryExclusive.getId()).intValue());
+    Assert.assertEquals(0, remainingResponseMap.get(queryLegacyTrue.getId()).intValue());
+  }
+
+  @Test
+  public void testRealtimeSegmentsModeExclude()
+  {
+    final Interval interval = Intervals.of("2016-01-01/2016-01-02");
+    final Interval queryInterval = Intervals.of("2016-01-01T14:00:00/2016-01-02T14:00:00");
+    final DataSegment dataSegment = new DataSegment(
+            "dataSource",
+            interval,
+            "ver",
+            ImmutableMap.of("type", "hdfs", "path", "/tmp"),
+            ImmutableList.of("product"),
+            ImmutableList.of("visited_sum"),
+            NoneShardSpec.instance(),
+            9,
+            12334
+    );
+
+    // selector backed only by a realtime server — isRealtimeSegment() == true
+    final DruidServer realtimeServer = new DruidServer(
+            "rt1", "rt1", null, 10, null, ServerType.REALTIME, DruidServer.DEFAULT_TIER, 0
+    );
+    final ServerSelector realtimeSelector = new ServerSelector(
+            dataSegment,
+            new HighestPriorityTierSelectorStrategy(new RandomServerSelectorStrategy()),
+            HistoricalFilter.IDENTITY_FILTER
+    );
+    realtimeSelector.addServerAndUpdateSegment(new QueryableDruidServer(realtimeServer, null), dataSegment);
+    timeline.add(interval, "ver", new SingleElementPartitionChunk<>(realtimeSelector));
+
+    // exclude: realtime-only segment is skipped
+    final TimeBoundaryQuery queryExclude = Druids.newTimeBoundaryQueryBuilder()
+            .dataSource(DATA_SOURCE)
+            .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(queryInterval)))
+            .context(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "exclude"))
+            .randomQueryId()
+            .build();
+
+    // include: realtime-only segment is included
+    final TimeBoundaryQuery queryInclude = Druids.newTimeBoundaryQueryBuilder()
+            .dataSource(DATA_SOURCE)
+            .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(queryInterval)))
+            .context(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "include"))
+            .randomQueryId()
+            .build();
+
+    final ResponseContext responseContext = initializeResponseContext();
+    getDefaultQueryRunner().run(QueryPlus.wrap(queryExclude), responseContext);
+    getDefaultQueryRunner().run(QueryPlus.wrap(queryInclude), responseContext);
+
     final Map<String, Integer> remainingResponseMap = (Map<String, Integer>) responseContext.get(ResponseContext.Keys.REMAINING_RESPONSES_FROM_QUERY_SERVERS);
-    Assert.assertEquals(1, remainingResponseMap.get(query.getId()).intValue());
-    Assert.assertEquals(0, remainingResponseMap.get(query2.getId()).intValue());
+    Assert.assertEquals(0, remainingResponseMap.get(queryExclude.getId()).intValue());
+    Assert.assertEquals(1, remainingResponseMap.get(queryInclude.getId()).intValue());
   }
 
   @SuppressWarnings("unchecked")