From ed2f27bc33324af48aa838045f6f3d16bd63fde5 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Wed, 1 Nov 2023 14:46:15 -0700
Subject: [PATCH] Update.

---
 .../druid/client/CachingClusteredClient.java  | 22 +++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/client/CachingClusteredClient.java b/server/src/main/java/org/apache/druid/client/CachingClusteredClient.java
index 2eea9aa594db..892e39e90079 100644
--- a/server/src/main/java/org/apache/druid/client/CachingClusteredClient.java
+++ b/server/src/main/java/org/apache/druid/client/CachingClusteredClient.java
@@ -436,20 +436,34 @@ private Set<SegmentServerSelector> computeSegmentsToQuery(
       );
 
       final Set<SegmentServerSelector> segments = new LinkedHashSet<>();
-      final Map<String, Optional<RangeSet<String>>> dimensionRangeCache = new HashMap<>();
+      final Map<String, Optional<RangeSet<String>>> dimensionRangeCache;
+      final Set<String> filterFieldsForPruning;
+
+      final boolean trySecondaryPartititionPruning =
+          query.getFilter() != null && query.context().isSecondaryPartitionPruningEnabled();
+
+      if (trySecondaryPartititionPruning) {
+        dimensionRangeCache = new HashMap<>();
+        filterFieldsForPruning =
+            DimFilterUtils.onlyBaseFields(query.getFilter().getRequiredColumns(), dataSourceAnalysis);
+      } else {
+        dimensionRangeCache = null;
+        filterFieldsForPruning = null;
+      }
+
       // Filter unneeded chunks based on partition dimension
       for (TimelineObjectHolder<String, ServerSelector> holder : serversLookup) {
         final Set<PartitionChunk<ServerSelector>> filteredChunks;
-        if (query.getFilter() != null && query.context().isSecondaryPartitionPruningEnabled()) {
+        if (trySecondaryPartititionPruning) {
           filteredChunks = DimFilterUtils.filterShards(
               query.getFilter(),
-              DimFilterUtils.onlyBaseFields(query.getFilter().getRequiredColumns(), dataSourceAnalysis),
+              filterFieldsForPruning,
               holder.getObject(),
               partitionChunk -> partitionChunk.getObject().getSegment().getShardSpec(),
               dimensionRangeCache
           );
         } else {
-          filteredChunks = Sets.newHashSet(holder.getObject());
+          filteredChunks = Sets.newLinkedHashSet(holder.getObject());
         }
         for (PartitionChunk<ServerSelector> chunk : filteredChunks) {
           ServerSelector server = chunk.getObject();