From 1bc728a462d64bb3b4a7e59adaf62c8bebbcbd08 Mon Sep 17 00:00:00 2001
From: Michael Froh <froh@amazon.com>
Date: Mon, 19 Aug 2024 18:01:35 -0700
Subject: [PATCH] Address comments from @mch2

Signed-off-by: Michael Froh <froh@amazon.com>
---
 .../bucket/terms/IncludeExclude.java          | 10 +-
 .../bucket/terms/IncludeExcludeTests.java     | 96 ++++++++++++++++++-
 2 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/IncludeExclude.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/IncludeExclude.java
index de83c7221835f..20f294bc7199b 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/IncludeExclude.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/IncludeExclude.java
@@ -429,16 +429,12 @@ private static BytesRef nextBytesRef(BytesRef bytesRef) {
             return next;
         }
 
-        private interface LongLongBiconsumer {
+        private interface LongBiConsumer {
             void accept(long a, long b);
         }
 
-        private static void process(
-            SortedSetDocValues globalOrdinals,
-            long length,
-            SortedSet<BytesRef> prefixes,
-            LongLongBiconsumer consumer
-        ) throws IOException {
+        private static void process(SortedSetDocValues globalOrdinals, long length, SortedSet<BytesRef> prefixes, LongBiConsumer consumer)
+            throws IOException {
             for (BytesRef prefix : prefixes) {
                 long startOrd = globalOrdinals.lookupTerm(prefix);
                 if (startOrd < 0) {
diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/IncludeExcludeTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/IncludeExcludeTests.java
index 20d5f84c1a406..2356e2dc6d855 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/IncludeExcludeTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/IncludeExcludeTests.java
@@ -355,7 +355,6 @@ public void testPrefixOrds() throws IOException {
             false };
 
         LongAdder lookupCount = new LongAdder();
-        LongAdder termsEnumAcceptCount = new LongAdder();
         SortedSetDocValues sortedSetDocValues = new AbstractSortedSetDocValues() {
             @Override
             public boolean advanceExact(int target) {
@@ -391,7 +390,6 @@ public long getValueCount() {
         for (int i = 0; i < expectedFilter.length; i++) {
             assertEquals(expectedFilter[i], acceptedOrds.get(i));
         }
-        assertEquals(0, termsEnumAcceptCount.longValue());
 
         // Now repeat, but this time, the prefix optimization won't work (because of the .+ on the exclude filter)
         includeExclude = new IncludeExclude("(color|length|size):.*", "color:g.+");
@@ -453,4 +451,98 @@ private static SortedSet<BytesRef> bytesRefs(String... strings) {
         }
         return bytesRefs;
     }
+
+    private static SortedSetDocValues toDocValues(BytesRef[] bytesRefs) {
+        return new AbstractSortedSetDocValues() {
+            @Override
+            public boolean advanceExact(int target) {
+                return false;
+            }
+
+            @Override
+            public long nextOrd() {
+                return 0;
+            }
+
+            @Override
+            public int docValueCount() {
+                return 1;
+            }
+
+            @Override
+            public BytesRef lookupOrd(long ord) {
+                int ordIndex = Math.toIntExact(ord);
+                return bytesRefs[ordIndex];
+            }
+
+            @Override
+            public long getValueCount() {
+                return bytesRefs.length;
+            }
+        };
+    }
+
+    public void testOnlyIncludeExcludePrefix() throws IOException {
+        String incExcString = "(color:g|width:).*";
+        IncludeExclude excludeOnly = new IncludeExclude(null, incExcString);
+
+        OrdinalsFilter ordinalsFilter = excludeOnly.convertToOrdinalsFilter(DocValueFormat.RAW);
+        // Which of the following match the filter or not?
+        BytesRef[] bytesRefs = toBytesRefArray(
+            "color:blue", // true
+            "color:crimson", // true
+            "color:green", // false (excluded)
+            "color:gray", // false (excluded)
+            "depth:10in", // true
+            "depth:12in", // true
+            "depth:17in", // true
+            "length:long", // true
+            "length:medium", // true
+            "length:short", // true
+            "material:cotton", // true
+            "material:linen", // true
+            "material:polyester", // true
+            "size:large", // true
+            "size:medium", // true
+            "size:small", // true
+            "width:13in", // false
+            "width:27in", // false
+            "width:55in" // false
+        );
+        boolean[] expectedFilter = new boolean[] {
+            true,
+            true,
+            false,
+            false,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            true,
+            false,
+            false,
+            false };
+        LongBitSet longBitSet = ordinalsFilter.acceptedGlobalOrdinals(toDocValues(bytesRefs));
+
+        for (int i = 0; i < expectedFilter.length; i++) {
+            assertEquals(expectedFilter[i], longBitSet.get(i));
+        }
+
+        // Now repeat, but this time include only.
+        IncludeExclude includeOnly = new IncludeExclude(incExcString, null);
+        ordinalsFilter = includeOnly.convertToOrdinalsFilter(DocValueFormat.RAW);
+        longBitSet = ordinalsFilter.acceptedGlobalOrdinals(toDocValues(bytesRefs));
+
+        // The previously excluded ordinals should be included and vice versa.
+        for (int i = 0; i < expectedFilter.length; i++) {
+            assertEquals(!expectedFilter[i], longBitSet.get(i));
+        }
+    }
 }