From 6cd8c6be22c40688df7513a7c6c51768a5b24a10 Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Wed, 7 Aug 2024 23:29:56 -0700
Subject: [PATCH] fix IndexedStringDruidPredicateIndexes to not needlessly
 lookup index of values (#16860)

---
 .../IndexedStringDruidPredicateIndexes.java   | 15 ++----
 .../NestedFieldColumnIndexSupplier.java       | 46 ++++++-------------
 .../ScalarDoubleColumnAndIndexSupplier.java   | 10 ++--
 .../ScalarLongColumnAndIndexSupplier.java     | 10 ++--
 4 files changed, 25 insertions(+), 56 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/segment/index/IndexedStringDruidPredicateIndexes.java b/processing/src/main/java/org/apache/druid/segment/index/IndexedStringDruidPredicateIndexes.java
index 8429b8b6a9f..42a8982d7ba 100644
--- a/processing/src/main/java/org/apache/druid/segment/index/IndexedStringDruidPredicateIndexes.java
+++ b/processing/src/main/java/org/apache/druid/segment/index/IndexedStringDruidPredicateIndexes.java
@@ -62,9 +62,8 @@ public final class IndexedStringDruidPredicateIndexes<TDictionary extends Indexe
         return () -> new Iterator<ImmutableBitmap>()
         {
           final Iterator<String> iterator = dictionary.iterator();
-          @Nullable
-          String next = null;
           boolean nextSet = false;
+          int index = -1;
 
           @Override
           public boolean hasNext()
@@ -85,23 +84,17 @@ public final class IndexedStringDruidPredicateIndexes<TDictionary extends Indexe
               }
             }
             nextSet = false;
-            final int idx = dictionary.indexOf(next);
-            if (idx < 0) {
-              return bitmapFactory.makeEmptyImmutableBitmap();
-            }
 
-            final ImmutableBitmap bitmap = bitmaps.get(idx);
+            final ImmutableBitmap bitmap = bitmaps.get(index);
             return bitmap == null ? bitmapFactory.makeEmptyImmutableBitmap() : bitmap;
           }
 
           private void findNext()
           {
             while (!nextSet && iterator.hasNext()) {
-              String nextValue = iterator.next();
+              final String nextValue = iterator.next();
+              index++;
               nextSet = stringPredicate.apply(nextValue).matches(includeUnknown);
-              if (nextSet) {
-                next = nextValue;
-              }
             }
           }
         };
diff --git a/processing/src/main/java/org/apache/druid/segment/nested/NestedFieldColumnIndexSupplier.java b/processing/src/main/java/org/apache/druid/segment/nested/NestedFieldColumnIndexSupplier.java
index a07306fb53d..c784eaa8d13 100644
--- a/processing/src/main/java/org/apache/druid/segment/nested/NestedFieldColumnIndexSupplier.java
+++ b/processing/src/main/java/org/apache/druid/segment/nested/NestedFieldColumnIndexSupplier.java
@@ -612,8 +612,7 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
 
             // in the future, this could use an int iterator
             final Iterator<Integer> iterator = localDictionary.iterator();
-            int next;
-            int index = 0;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -635,19 +634,16 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
                 }
               }
               nextSet = false;
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
                 Integer nextValue = iterator.next();
+                index++;
                 nextSet = stringPredicate.apply(StringUtils.fromUtf8Nullable(stringDictionary.get(nextValue)))
                                          .matches(includeUnknown);
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };
@@ -892,8 +888,7 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
 
             // in the future, this could use an int iterator
             final Iterator<Integer> iterator = localDictionary.iterator();
-            int next;
-            int index = 0;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -916,22 +911,19 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
               }
               nextSet = false;
 
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
-                Integer nextValue = iterator.next();
+                final Integer nextValue = iterator.next();
+                index++;
                 if (nextValue == 0) {
                   nextSet = longPredicate.applyNull().matches(includeUnknown);
                 } else {
                   nextSet = longPredicate.applyLong(longDictionary.get(nextValue - adjustLongId)).matches(includeUnknown);
                 }
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };
@@ -1158,8 +1150,7 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
 
             // in the future, this could use an int iterator
             final Iterator<Integer> iterator = localDictionary.iterator();
-            int next;
-            int index = 0;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -1181,23 +1172,20 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
                 }
               }
               nextSet = false;
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
-                Integer nextValue = iterator.next();
+                final Integer nextValue = iterator.next();
+                index++;
                 if (nextValue == 0) {
                   nextSet = doublePredicate.applyNull().matches(includeUnknown);
                 } else {
                   nextSet = doublePredicate.applyDouble(doubleDictionary.get(nextValue - adjustDoubleId))
                                            .matches(includeUnknown);
                 }
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };
@@ -1384,8 +1372,7 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
 
             // in the future, this could use an int iterator
             final Iterator<Integer> iterator = localDictionary.iterator();
-            int next;
-            int index;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -1407,13 +1394,14 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
                 }
               }
               nextSet = false;
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
-                Integer nextValue = iterator.next();
+                final Integer nextValue = iterator.next();
+                index++;
                 if (nextValue >= adjustArrayId) {
                   // this shouldn't be possible since arrayIds will only exist if array dictionary is not null
                   // v4 columns however have a null array dictionary
@@ -1442,10 +1430,6 @@ public class NestedFieldColumnIndexSupplier<TStringDictionary extends Indexed<By
                   nextSet = stringPredicate.apply(StringUtils.fromUtf8Nullable(stringDictionary.get(nextValue)))
                                            .matches(includeUnknown);
                 }
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };
diff --git a/processing/src/main/java/org/apache/druid/segment/nested/ScalarDoubleColumnAndIndexSupplier.java b/processing/src/main/java/org/apache/druid/segment/nested/ScalarDoubleColumnAndIndexSupplier.java
index ff0019a2b68..1c9d5a6aa35 100644
--- a/processing/src/main/java/org/apache/druid/segment/nested/ScalarDoubleColumnAndIndexSupplier.java
+++ b/processing/src/main/java/org/apache/druid/segment/nested/ScalarDoubleColumnAndIndexSupplier.java
@@ -544,8 +544,7 @@ public class ScalarDoubleColumnAndIndexSupplier implements Supplier<NestedCommon
             final Iterator<Double> iterator = doubleDictionarySupplier.get().iterator();
             final DruidDoublePredicate doublePredicate = matcherFactory.makeDoublePredicate();
 
-            int next;
-            int index = 0;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -567,13 +566,14 @@ public class ScalarDoubleColumnAndIndexSupplier implements Supplier<NestedCommon
                 }
               }
               nextSet = false;
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
                 Double nextValue = iterator.next();
+                index++;
                 if (nextValue == null) {
                   if (NullHandling.sqlCompatible()) {
                     nextSet = doublePredicate.applyNull().matches(includeUnknown);
@@ -583,10 +583,6 @@ public class ScalarDoubleColumnAndIndexSupplier implements Supplier<NestedCommon
                 } else {
                   nextSet = doublePredicate.applyDouble(nextValue).matches(includeUnknown);
                 }
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };
diff --git a/processing/src/main/java/org/apache/druid/segment/nested/ScalarLongColumnAndIndexSupplier.java b/processing/src/main/java/org/apache/druid/segment/nested/ScalarLongColumnAndIndexSupplier.java
index 5f0f3a61efe..068613aba68 100644
--- a/processing/src/main/java/org/apache/druid/segment/nested/ScalarLongColumnAndIndexSupplier.java
+++ b/processing/src/main/java/org/apache/druid/segment/nested/ScalarLongColumnAndIndexSupplier.java
@@ -554,8 +554,7 @@ public class ScalarLongColumnAndIndexSupplier implements Supplier<NestedCommonFo
             final Iterator<Long> iterator = dictionary.iterator();
             final DruidLongPredicate longPredicate = matcherFactory.makeLongPredicate();
 
-            int next;
-            int index = 0;
+            int index = -1;
             boolean nextSet = false;
 
             @Override
@@ -578,13 +577,14 @@ public class ScalarLongColumnAndIndexSupplier implements Supplier<NestedCommonFo
               }
               nextSet = false;
 
-              return getBitmap(next);
+              return getBitmap(index);
             }
 
             private void findNext()
             {
               while (!nextSet && iterator.hasNext()) {
                 Long nextValue = iterator.next();
+                index++;
                 if (nextValue == null) {
                   if (NullHandling.sqlCompatible()) {
                     nextSet = longPredicate.applyNull().matches(includeUnknown);
@@ -594,10 +594,6 @@ public class ScalarLongColumnAndIndexSupplier implements Supplier<NestedCommonFo
                 } else {
                   nextSet = longPredicate.applyLong(nextValue).matches(includeUnknown);
                 }
-                if (nextSet) {
-                  next = index;
-                }
-                index++;
               }
             }
           };