LUCENE-7309: Speed up Polygon2D creation.

lucene/core/src/java/org/apache/lucene/geo/Polygon2D.java

@@ -17,9 +17,11 @@
 package org.apache.lucene.geo;
 
 import java.util.Arrays;
+import java.util.Comparator;
 
 import org.apache.lucene.geo.Polygon;
 import org.apache.lucene.index.PointValues.Relation;
+import org.apache.lucene.util.ArrayUtil;
 
 /**
  * 2D polygon implementation represented as a balanced interval tree of edges.
@@ -209,28 +211,30 @@ public final class Polygon2D {
   private static Polygon2D createTree(Polygon2D components[], int low, int high, boolean splitX) {
     if (low > high) {
       return null;
-    } else if (low < high) {
-      // TODO: do one sort instead! there are better algorithms!
+    }
+    final int mid = (low + high) >>> 1;
+    if (low < high) {
+      Comparator<Polygon2D> comparator;
       if (splitX) {
-        Arrays.sort(components, low, high+1, (left, right) -> {
+        comparator = (left, right) -> {
           int ret = Double.compare(left.minLon, right.minLon);
           if (ret == 0) {
             ret = Double.compare(left.maxX, right.maxX);
           }
           return ret;
-        });
+        };
       } else {
-        Arrays.sort(components, low, high+1, (left, right) -> {
+        comparator = (left, right) -> {
           int ret = Double.compare(left.minLat, right.minLat);
           if (ret == 0) {
             ret = Double.compare(left.maxY, right.maxY);
           }
           return ret;
-        });
+        };
       }
+      ArrayUtil.select(components, low, high + 1, mid, comparator);
     }
     // add midpoint
-    int mid = (low + high) >>> 1;
     Polygon2D newNode = components[mid];
     newNode.splitX = splitX;
     // add children

lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java

@@ -453,4 +453,75 @@ public final class ArrayUtil {
     timSort(a, 0, a.length);
   }
 
+  /** Reorganize {@code arr[from:to[} so that the element at offset k is at the
+   *  same position as if {@code arr[from:to[} was sorted, and all elements on
+   *  its left are less than or equal to it, and all elements on its right are
+   *  greater than or equal to it.
+   *  This runs in linear time on average and in {@code n log(n)} time in the
+   *  worst case.*/
+  public static <T> void select(T[] arr, int from, int to, int k, Comparator<T> comparator) {
+    if (k < from) {
+      throw new IllegalArgumentException("k must be >= from");
+    }
+    if (k >= to) {
+      throw new IllegalArgumentException("k must be < to");
+    }
+    final int maxDepth = 2 * MathUtil.log(to - from, 2);
+    quickSelect(arr, from, to, k, comparator, maxDepth);
+  }
+
+  private static <T> void quickSelect(T[] arr, int from, int to, int k, Comparator<T> comparator, int maxDepth) {
+    assert from <= k;
+    assert k < to;
+    if (to - from == 1) {
+      return;
+    }
+    if (--maxDepth < 0) {
+      Arrays.sort(arr, from, to, comparator);
+      return;
+    }
+
+    final int mid = (from + to) >>> 1;
+    // heuristic: we use the median of the values at from, to-1 and mid as a pivot
+    if (comparator.compare(arr[from], arr[to - 1]) > 0) {
+      swap(arr, from, to - 1);
+    }
+    if (comparator.compare(arr[to - 1], arr[mid]) > 0) {
+      swap(arr, to - 1, mid);
+      if (comparator.compare(arr[from], arr[to - 1]) > 0) {
+        swap(arr, from, to - 1);
+      }
+    }
+
+    T pivot = arr[to - 1];
+
+    int left = from + 1;
+    int right = to - 2;
+
+    for (;;) {
+      while (comparator.compare(pivot, arr[left]) > 0) {
+        ++left;
+      }
+
+      while (left < right && comparator.compare(pivot, arr[right]) <= 0) {
+        --right;
+      }
+
+      if (left < right) {
+        swap(arr, left, right);
+        --right;
+      } else {
+        break;
+      }
+    }
+    swap(arr, left, to - 1);
+
+    if (left == k) {
+      return;
+    } else if (left < k) {
+      quickSelect(arr, left + 1, to, k, comparator, maxDepth);
+    } else {
+      quickSelect(arr, from, left, k, comparator, maxDepth);
+    }
+  }
 }

lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java

@@ -20,6 +20,7 @@ package org.apache.lucene.util;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Random;
 
 public class TestArrayUtil extends LuceneTestCase {
@@ -275,5 +276,38 @@ public class TestArrayUtil extends LuceneTestCase {
     ArrayUtil.introSort(a, Collections.reverseOrder());
     ArrayUtil.timSort(a, Collections.reverseOrder());
   }
-  
+
+  public void testSelect() {
+    for (int iter = 0; iter < 100; ++iter) {
+      doTestSelect();
+    }
+  }
+
+  private void doTestSelect() {
+    final int from = random().nextInt(5);
+    final int to = from + TestUtil.nextInt(random(), 1, 10000);
+    final int max = random().nextBoolean() ? random().nextInt(100) : random().nextInt(100000);
+    Integer[] arr = new Integer[from + to + random().nextInt(5)];
+    for (int i = 0; i < arr.length; ++i) {
+      arr[i] = TestUtil.nextInt(random(), 0, max);
+    }
+    final int k = TestUtil.nextInt(random(), from, to - 1);
+
+    Integer[] expected = arr.clone();
+    Arrays.sort(expected, from, to);
+
+    Integer[] actual = arr.clone();
+    ArrayUtil.select(actual, from, to, k, Comparator.naturalOrder());
+
+    assertEquals(expected[k], actual[k]);
+    for (int i = 0; i < actual.length; ++i) {
+      if (i < from || i >= to) {
+        assertSame(arr[i], actual[i]);
+      } else if (i <= k) {
+        assertTrue(actual[i].intValue() <= actual[k].intValue());
+      } else {
+        assertTrue(actual[i].intValue() >= actual[k].intValue());
+      }
+    }
+  }
 }