Long key bucket ords key iterator (#95809)

Relates to #89437

In order to merge aggregations using LongKeyedBucketOrds, we need a way to align the buckets by key. This adds an iterator to return all keys for a given owning bucket ordinal, in natural sorted order. This can then be used to merge without generating and sorting bucket objects directly.

This makes two big assumptions:

1 - the total set of keys is small enough that we don't need to put the key set in a Big Arrays backed data structure
2 - Building the tree set of keys (at reduce time) will not be excessively expensive.

We can mitigate 1 if we build a big arrays backed balanced binary tree. That can be done in a follow up PR to this one.

Author: not-napoleon

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java

@@ -14,21 +14,32 @@
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
 
+import java.util.Iterator;
 import java.util.Locale;
+import java.util.NoSuchElementException;
+import java.util.TreeSet;
 
 /**
  * Maps owning bucket ordinals and long bucket keys to bucket ordinals.
  */
 public abstract class LongKeyedBucketOrds implements Releasable {
     /**
      * Build a {@link LongKeyedBucketOrds} who's values have unknown bounds.
+     *
+     * @param cardinality - This should come from the owning aggregation, and is used as an upper bound on the
+     *                    owning bucket ordinals.
      */
     public static LongKeyedBucketOrds build(BigArrays bigArrays, CardinalityUpperBound cardinality) {
         return cardinality.map(estimate -> estimate < 2 ? new FromSingle(bigArrays) : new FromMany(bigArrays));
     }
 
     /**
      * Build a {@link LongKeyedBucketOrds} who's values have known bounds.
+     *
+     * @param cardinality - This should come from the owning aggregation, and is used as an upper bound on the
+     *                    owning bucket ordinals.
+     * @param  min - The minimum key value for this aggregation
+     * @param max - The maximum key value for this aggregation
      */
     public static LongKeyedBucketOrds buildForValueRange(BigArrays bigArrays, CardinalityUpperBound cardinality, long min, long max) {
         return cardinality.map((int cardinalityUpperBound) -> {
@@ -49,6 +60,8 @@ public static LongKeyedBucketOrds buildForValueRange(BigArrays bigArrays, Cardin
         });
     }
 
+    private TreeSet<Long> keySet = null;
+
     private LongKeyedBucketOrds() {}
 
     /**
@@ -101,6 +114,58 @@ private LongKeyedBucketOrds() {}
      */
     public abstract BucketOrdsEnum ordsEnum(long owningBucketOrd);
 
+    /**
+     * Return an iterator for all keys in the given owning bucket, ordered in natural sort order.
+     * This is suitable for aligning buckets across different instances of an aggregation.
+     *
+     * @param owningBucketOrd Only return keys that occured under this owning bucket
+     * @return a sorted iterator of long key values
+     */
+    public Iterator<Long> keyOrderedIterator(long owningBucketOrd) {
+        if (keySet == null) {
+            // TreeSet's contract includes a naturally ordered iterator
+            keySet = new TreeSet<>();
+            for (long ord = 0; ord < size(); ord++) {
+                keySet.add(this.get(ord));
+            }
+        }
+        Iterator<Long> toReturn = new Iterator<>() {
+            Iterator<Long> wrapped = keySet.iterator();
+            long filterOrd = owningBucketOrd;
+            long next;
+            boolean hasNext = true;
+
+            @Override
+            public boolean hasNext() {
+                return hasNext;
+            }
+
+            @Override
+            public Long next() {
+                if (hasNext == false) {
+                    throw new NoSuchElementException();
+                }
+                long toReturn = next;
+                hasNext = false;
+                while (wrapped.hasNext()) {
+                    long candidate = wrapped.next();
+                    if (find(filterOrd, candidate) != -1) {
+                        next = candidate;
+                        hasNext = true;
+                        break;
+                    }
+                }
+                return toReturn;
+            }
+        };
+        toReturn.next(); // Prime the first actual value
+        return toReturn;
+    }
+
+    public void close() {
+        keySet = null;
+    }
+
     /**
      * An iterator for buckets inside a particular {@code owningBucketOrd}.
      */
@@ -223,6 +288,7 @@ public long ord() {
 
         @Override
         public void close() {
+            super.close();
             ords.close();
         }
     }
@@ -320,6 +386,7 @@ public long ord() {
 
         @Override
         public void close() {
+            super.close();
             ords.close();
         }
     }
@@ -455,6 +522,7 @@ public long ord() {
 
         @Override
         public void close() {
+            super.close();
             ords.close();
         }
     }

server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/LongKeyedBucketOrdsTests.java

@@ -15,11 +15,16 @@
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.TreeSet;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 
 public class LongKeyedBucketOrdsTests extends ESTestCase {
     private final MockBigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService());
@@ -32,6 +37,21 @@ public void testSurpriseCollectsFromSingleBucket() {
         collectsFromSingleBucketCase(LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.MANY));
     }
 
+    public void testCollectsFromManyBuckets() {
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.MANY)) {
+            assertCollectsFromManyBuckets(ords, scaledRandomIntBetween(1, 10000), Long.MIN_VALUE, Long.MAX_VALUE);
+        }
+    }
+
+    public void testCollectsFromManyBucketsSmall() {
+        int owningBucketOrds = scaledRandomIntBetween(2, 10000);
+        long maxValue = randomLongBetween(10000 / owningBucketOrds, 2 << (16 * 3));
+        CardinalityUpperBound cardinality = CardinalityUpperBound.ONE.multiply(owningBucketOrds);
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.buildForValueRange(bigArrays, cardinality, 0, maxValue)) {
+            assertCollectsFromManyBuckets(ords, owningBucketOrds, 0, maxValue);
+        }
+    }
+
     private void collectsFromSingleBucketCase(LongKeyedBucketOrds ords) {
         try {
             // Test a few explicit values
@@ -93,21 +113,6 @@ private void collectsFromSingleBucketCase(LongKeyedBucketOrds ords) {
         }
     }
 
-    public void testCollectsFromManyBuckets() {
-        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.MANY)) {
-            assertCollectsFromManyBuckets(ords, scaledRandomIntBetween(1, 10000), Long.MIN_VALUE, Long.MAX_VALUE);
-        }
-    }
-
-    public void testCollectsFromManyBucketsSmall() {
-        int owningBucketOrds = scaledRandomIntBetween(2, 10000);
-        long maxValue = randomLongBetween(10000 / owningBucketOrds, 2 << (16 * 3));
-        CardinalityUpperBound cardinality = CardinalityUpperBound.ONE.multiply(owningBucketOrds);
-        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.buildForValueRange(bigArrays, cardinality, 0, maxValue)) {
-            assertCollectsFromManyBuckets(ords, owningBucketOrds, 0, maxValue);
-        }
-    }
-
     private void assertCollectsFromManyBuckets(LongKeyedBucketOrds ords, int maxAllowedOwningBucketOrd, long minValue, long maxValue) {
         // Test a few explicit values
         assertThat(ords.add(0, 0), equalTo(0L));
@@ -176,6 +181,173 @@ private void assertCollectsFromManyBuckets(LongKeyedBucketOrds ords, int maxAllo
         assertThat(ords.maxOwningBucketOrd(), equalTo(maxOwningBucketOrd));
     }
 
+    public void testKeyIteratorSingleValue() {
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.ONE)) {
+            // Test a few explicit values
+            assertThat(ords.add(0, 0), equalTo(0L));
+            assertThat(ords.add(0, 1000), equalTo(1L));
+            assertThat(ords.add(0, 0), equalTo(-1L));
+            assertThat(ords.add(0, 1000), equalTo(-2L));
+            assertThat(ords.find(0, 0), equalTo(0L));
+            assertThat(ords.find(0, 1000), equalTo(1L));
+
+            // And some random values
+            Set<Long> seen = new TreeSet<>();
+            seen.add(0L);
+            seen.add(1000L);
+            assertThat(ords.size(), equalTo(2L));
+            long[] values = new long[scaledRandomIntBetween(1, 10000)];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = randomValueOtherThanMany(seen::contains, ESTestCase::randomLong);
+                seen.add(values[i]);
+            }
+            for (int i = 0; i < values.length; i++) {
+                assertThat(ords.find(0, values[i]), equalTo(-1L));
+                assertThat(ords.add(0, values[i]), equalTo(i + 2L));
+                assertThat(ords.find(0, values[i]), equalTo(i + 2L));
+                assertThat(ords.size(), equalTo(i + 3L));
+            }
+
+            // For the single value case, the sorted iterator should exactly equal the values tree set iterator
+            Iterator<Long> expected = seen.iterator();
+            Iterator<Long> actual = ords.keyOrderedIterator(0);
+            assertNotSame(expected, actual);
+            while (expected.hasNext()) {
+                assertThat(actual.hasNext(), is(true));
+                long actualNext = actual.next();
+                long expectedNext = expected.next();
+                assertThat(actualNext, equalTo(expectedNext));
+            }
+            assertThat(actual.hasNext(), is(false));
+        }
+    }
+
+    public void testKeyIteratormanyBuckets() {
+        long maxAllowedOwningBucketOrd = scaledRandomIntBetween(1, 10000);
+        long minValue = Long.MIN_VALUE;
+        long maxValue = Long.MAX_VALUE;
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.MANY)) {
+            Map<Long, TreeSet<Long>> expected = new HashMap<>();
+            // Test a few explicit values
+            assertThat(ords.add(0, 0), equalTo(0L));
+            assertThat(ords.add(1, 0), equalTo(1L));
+            assertThat(ords.add(0, 0), equalTo(-1L));
+            assertThat(ords.add(1, 0), equalTo(-2L));
+            assertThat(ords.size(), equalTo(2L));
+            assertThat(ords.find(0, 0), equalTo(0L));
+            assertThat(ords.find(1, 0), equalTo(1L));
+
+            Set<OwningBucketOrdAndValue> seen = new HashSet<>();
+            seen.add(new OwningBucketOrdAndValue(0, 0));
+            seen.add(new OwningBucketOrdAndValue(1, 0));
+
+            expected.put(0L, new TreeSet<>());
+            expected.get(0L).add(0L);
+            expected.put(1L, new TreeSet<>());
+            expected.get(1L).add(0L);
+
+            OwningBucketOrdAndValue[] values = new OwningBucketOrdAndValue[scaledRandomIntBetween(1, 10000)];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = randomValueOtherThanMany(
+                    seen::contains,
+                    () -> new OwningBucketOrdAndValue(
+                        randomLongBetween(0, maxAllowedOwningBucketOrd),
+                        randomLongBetween(minValue, maxValue)
+                    )
+                );
+                seen.add(values[i]);
+                if (expected.containsKey(values[i].owningBucketOrd) == false) {
+                    expected.put(values[i].owningBucketOrd, new TreeSet<>());
+                }
+                expected.get(values[i].owningBucketOrd).add(values[i].value);
+            }
+            for (int i = 0; i < values.length; i++) {
+                assertThat(ords.find(values[i].owningBucketOrd, values[i].value), equalTo(-1L));
+                assertThat(ords.add(values[i].owningBucketOrd, values[i].value), equalTo(i + 2L));
+                assertThat(ords.find(values[i].owningBucketOrd, values[i].value), equalTo(i + 2L));
+                assertThat(ords.size(), equalTo(i + 3L));
+            }
+
+            for (Long owningBucketOrd : expected.keySet()) {
+                Iterator<Long> expectedIterator = expected.get(owningBucketOrd).iterator();
+                Iterator<Long> actualIterator = ords.keyOrderedIterator(owningBucketOrd);
+
+                assertNotSame(expectedIterator, actualIterator);
+                while (expectedIterator.hasNext()) {
+                    assertThat(actualIterator.hasNext(), is(true));
+                    long actualNext = actualIterator.next();
+                    long expectedNext = expectedIterator.next();
+                    assertThat(actualNext, equalTo(expectedNext));
+                }
+                assertThat(actualIterator.hasNext(), is(false));
+            }
+        }
+    }
+
+    public void testKeyIteratorManyBucketsSmall() {
+        int maxAllowedOwningBucketOrd = scaledRandomIntBetween(2, 10000);
+        long minValue = 0;
+        long maxValue = randomLongBetween(10000 / maxAllowedOwningBucketOrd, 2 << (16 * 3));
+        CardinalityUpperBound cardinality = CardinalityUpperBound.ONE.multiply(maxAllowedOwningBucketOrd);
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.buildForValueRange(bigArrays, cardinality, minValue, maxValue)) {
+            assertTrue(ords instanceof LongKeyedBucketOrds.FromManySmall);
+            Map<Long, TreeSet<Long>> expected = new HashMap<>();
+            // Test a few explicit values
+            assertThat(ords.add(0, 0), equalTo(0L));
+            assertThat(ords.add(1, 0), equalTo(1L));
+            assertThat(ords.add(0, 0), equalTo(-1L));
+            assertThat(ords.add(1, 0), equalTo(-2L));
+            assertThat(ords.size(), equalTo(2L));
+            assertThat(ords.find(0, 0), equalTo(0L));
+            assertThat(ords.find(1, 0), equalTo(1L));
+
+            Set<OwningBucketOrdAndValue> seen = new HashSet<>();
+            seen.add(new OwningBucketOrdAndValue(0, 0));
+            seen.add(new OwningBucketOrdAndValue(1, 0));
+
+            expected.put(0L, new TreeSet<>());
+            expected.get(0L).add(0L);
+            expected.put(1L, new TreeSet<>());
+            expected.get(1L).add(0L);
+
+            OwningBucketOrdAndValue[] values = new OwningBucketOrdAndValue[scaledRandomIntBetween(1, 10000)];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = randomValueOtherThanMany(
+                    seen::contains,
+                    () -> new OwningBucketOrdAndValue(
+                        randomLongBetween(0, maxAllowedOwningBucketOrd),
+                        randomLongBetween(minValue, maxValue)
+                    )
+                );
+                seen.add(values[i]);
+                if (expected.containsKey(values[i].owningBucketOrd) == false) {
+                    expected.put(values[i].owningBucketOrd, new TreeSet<>());
+                }
+                expected.get(values[i].owningBucketOrd).add(values[i].value);
+            }
+            for (int i = 0; i < values.length; i++) {
+                assertThat(ords.find(values[i].owningBucketOrd, values[i].value), equalTo(-1L));
+                assertThat(ords.add(values[i].owningBucketOrd, values[i].value), equalTo(i + 2L));
+                assertThat(ords.find(values[i].owningBucketOrd, values[i].value), equalTo(i + 2L));
+                assertThat(ords.size(), equalTo(i + 3L));
+            }
+
+            for (Long owningBucketOrd : expected.keySet()) {
+                Iterator<Long> expectedIterator = expected.get(owningBucketOrd).iterator();
+                Iterator<Long> actualIterator = ords.keyOrderedIterator(owningBucketOrd);
+
+                assertNotSame(expectedIterator, actualIterator);
+                while (expectedIterator.hasNext()) {
+                    assertThat(actualIterator.hasNext(), is(true));
+                    long actualNext = actualIterator.next();
+                    long expectedNext = expectedIterator.next();
+                    assertThat(actualNext, equalTo(expectedNext));
+                }
+                assertThat(actualIterator.hasNext(), is(false));
+            }
+        }
+    }
+
     private class OwningBucketOrdAndValue {
         private final long owningBucketOrd;
         private final long value;