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 segments are queried, classified by whether a historical replica exists. `include` queries all segments. `exclusive` queries only segments served solely by realtime servers; any segment with at least one historical replica (including segments mid-handoff) is excluded. `exclude` is the inverse: segments served solely by realtime servers are skipped, but segments mid-handoff that have both a realtime and a historical replica are still queried. |
+|`realtimeSegmentsOnly` |`false`| **Deprecated.** Use `realtimeSegmentsMode=exclusive` instead. When set to `true`, 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,8 +782,34 @@ public boolean isPrePlanned()
     return getBoolean(QueryContexts.CTX_PREPLANNED, QueryContexts.DEFAULT_PREPLANNED);
   }
 
+  /**
+   * Returns the realtime segments mode for this query. If {@code realtimeSegmentsMode} is absent
+   * or null, falls back to the deprecated {@code realtimeSegmentsOnly} boolean: {@code true} maps
+   * to {@link RealtimeSegmentsMode#EXCLUSIVE}; otherwise returns {@link RealtimeSegmentsMode#INCLUDE}.
+   */
+  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);
+    return getRealtimeSegmentsMode() == RealtimeSegmentsMode.EXCLUSIVE;
   }
 }

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,39 @@ public String toString()
     }
   }
 
+  /**
+   * Classifies segments by whether a historical replica exists
+   * (see {@link org.apache.druid.client.selector.ServerSelector#isRealtimeSegment()}: a segment is
+   * "realtime" only when it has realtime servers and zero historical servers).
+   */
+  public enum RealtimeSegmentsMode
+  {
+    /** Include all segments (default). */
+    INCLUDE,
+    /** Include only segments served solely by realtime servers; any segment with a historical replica
+     *  (including segments mid-handoff) is excluded. */
+    EXCLUSIVE,
+    /** Exclude segments served solely by realtime servers; segments mid-handoff with both realtime
+     *  and historical replicas are still included. */
+    EXCLUDE;
+
+    @JsonCreator
+    public static RealtimeSegmentsMode fromString(String str)
+    {
+      if (str == null) {
+        return null;
+      }
+      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,66 @@ 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 testGetRealtimeSegmentsModeInvalidValue()
+  {
+    BadQueryContextException e = assertThrows(
+        BadQueryContextException.class,
+        () -> QueryContext.of(ImmutableMap.of(QueryContexts.REALTIME_SEGMENTS_MODE, "badvalue"))
+                         .getRealtimeSegmentsMode()
+    );
+    assertTrue(e.getMessage().contains("realtimeSegmentsMode"));
+  }
+
   @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,19 @@ 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
+          switch (realtimeSegmentsMode) {
+            case EXCLUSIVE:
+              if (!server.isRealtimeSegment()) {
+                continue;
+              }
+              break;
+            case EXCLUDE:
+              if (server.isRealtimeSegment()) {
+                continue;
+              }
+              break;
+            case INCLUDE:
+              break;
           }
           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")