Fix flaky test #96909 (#96929)

It turns out that the old geo_line has a bug whereby it can cause the truncation to happen in the middle of the line instead of from one end.
This occurs if there are documents with missing sort field and other documents with missing points, and the data is distributed across multiple shards.
Each shard gets a leaf-reader and collects into a bucket with the size truncated to the documents with sort-field (ie. not requiring existence of point field). This means the per-shard geo_line can be shorter than the size limit. That extra space can be filled during MergeGeoLines from a shard that could be far away, so not the correct globally sorted document set.

Fixing this properly requires changing the BucketedSort.Leaf to truncate on both sort field and point field, but that code is currently used by all aggregators, so we either have to confirm that it is OK for other aggregators to get this change in behaviour, or override that leaf.collect() method for geo_line only.

Author: craigtaverner

x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/search/aggregations/GeoLineAggregatorTests.java

@@ -71,6 +71,7 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 
+import static java.lang.Math.min;
 import static org.elasticsearch.index.IndexMode.TIME_SERIES;
 import static org.elasticsearch.index.mapper.TimeSeriesParams.MetricType.POSITION;
 import static org.hamcrest.CoreMatchers.containsString;
@@ -236,7 +237,7 @@ public void testCompleteForSizeAndNumDocuments(int size, int numPoints, boolean
             sortValues[i] = i;
         }
 
-        int lineSize = Math.min(numPoints, size);
+        int lineSize = min(numPoints, size);
         // re-sort line to be ascending
         long[] linePoints = Arrays.copyOf(points, lineSize);
         double[] lineSorts = Arrays.copyOf(sortValues, lineSize);
@@ -431,9 +432,11 @@ public void testGeoLine_Terms_TSDB_simplified() throws IOException {
     private void assertGeoLine(SortOrder sortOrder, String group, InternalGeoLine geoLine, TestTSAssertionResults tsx, boolean complete) {
         long[] expectedAggPoints = tsx.expectedAggPoints.get(group);
         double[] expectedAggSortValues = tsx.expectedAggSortValues.get(group);
-        String prefix = "GeoLine[sort=" + sortOrder + ", use-timestamps=" + tsx.useTimestampField + "]";
+        String prefix = "GeoLine[sort=" + sortOrder + ", use-timestamps=" + tsx.useTimeSeriesAggregation + ", group='" + group + "']";
         assertThat(prefix + " is complete", geoLine.isComplete(), is(complete));
-        assertThat(prefix + " contents", geoLine.line(), isGeoLine(tsx.useTimestampField, expectedAggPoints));
+        // old geo_line has a bug whereby it can produce lines with truncation happening in the middle instead of the end
+        int checkCount = tsx.useTimeSeriesAggregation ? expectedAggPoints.length : min(expectedAggPoints.length / 2, geoLine.line().length);
+        assertThat(prefix + " contents", geoLine.line(), isGeoLine(checkCount, expectedAggPoints));
         double[] sortValues = geoLine.sortVals();
         for (int i = 1; i < sortValues.length; i++) {
             Matcher<Double> sortMatcher = switch (sortOrder) {
@@ -442,37 +445,37 @@ private void assertGeoLine(SortOrder sortOrder, String group, InternalGeoLine ge
             };
             assertThat(prefix + " expect ordered '" + sortOrder + "' sort values", sortValues[i], sortMatcher);
         }
-        if (tsx.useTimestampField) {
+        if (checkCount == expectedAggSortValues.length) {
             assertArrayEquals(prefix + " sort values", expectedAggSortValues, sortValues, 0d);
         } else {
-            for (int i = 0; i < Math.min(expectedAggSortValues.length, sortValues.length); i++) {
+            for (int i = 0; i < checkCount; i++) {
                 assertThat(prefix + " sort value " + i, expectedAggSortValues[i], equalTo(sortValues[i]));
             }
         }
     }
 
-    private Matcher<long[]> isGeoLine(boolean matchLength, long[] line) {
-        return new TestGeoLineLongArrayMatcher(matchLength, line);
+    private Matcher<long[]> isGeoLine(int checkCount, long[] line) {
+        return new TestGeoLineLongArrayMatcher(checkCount, line);
     }
 
     private static class TestGeoLineLongArrayMatcher extends BaseMatcher<long[]> {
-        private final boolean matchLength;
+        private final int checkCount;
         private final long[] expectedLine;
         private final ArrayList<String> failures = new ArrayList<>();
 
-        private TestGeoLineLongArrayMatcher(boolean matchLength, long[] expectedLine) {
-            this.matchLength = matchLength;
+        private TestGeoLineLongArrayMatcher(int checkCount, long[] expectedLine) {
+            this.checkCount = checkCount;
             this.expectedLine = expectedLine;
         }
 
         @Override
         public boolean matches(Object actualObj) {
             failures.clear();
             if (actualObj instanceof long[] actualLine) {
-                if (matchLength && actualLine.length != expectedLine.length) {
+                if (checkCount == expectedLine.length && actualLine.length != expectedLine.length) {
                     failures.add("Expected length " + expectedLine.length + " but got " + actualLine.length);
                 }
-                for (int i = 0; i < Math.min(expectedLine.length, actualLine.length); i++) {
+                for (int i = 0; i < checkCount; i++) {
                     Point actual = asPoint(actualLine[i]);
                     Point expected = asPoint(expectedLine[i]);
                     if (actual.equals(expected) == false) {
@@ -646,22 +649,19 @@ private record TestConfig(
         int missingTimestampFactor,
         int groupCount,
         SortOrder sortOrder,
-        boolean useTimestamp
+        boolean useTimeSeriesAggregation
     ) {
-        private String sortField() {
-            return useTimestamp ? "@timestamp" : "time_field";
-        }
-
         @SuppressWarnings("SameParameterValue")
         private GeoLineAggregationBuilder lineAggregationBuilder(String name, String valueField, String sortField) {
             MultiValuesSourceFieldConfig valueConfig = new MultiValuesSourceFieldConfig.Builder().setFieldName(valueField).build();
             GeoLineAggregationBuilder lineAggregationBuilder = new GeoLineAggregationBuilder(name).point(valueConfig)
                 .sortOrder(sortOrder)
                 .size(maxPoints);
-            if (useTimestamp) {
+            if (useTimeSeriesAggregation) {
                 // In time-series we do not set the sort field
                 return lineAggregationBuilder;
             } else {
+                // Without a time-series aggregation, we need to specify the sort-field
                 MultiValuesSourceFieldConfig sortConfig = new MultiValuesSourceFieldConfig.Builder().setFieldName(sortField).build();
                 return lineAggregationBuilder.sort(sortConfig);
             }
@@ -698,7 +698,7 @@ private void build() {
                 // TSDB provides docs in DESC time order, so we generate the data that way to simplify assertions
                 for (int i = t.docCount - 1; i >= 0; i--) {
                     double lat = startLat + i * 0.1 + randomDoubleBetween(-0.1, 0.1, false);
-                    double lon = startLon + i * 0.1 + randomDoubleBetween(-0.1, 0.1, false);
+                    double lon = startLon - g * 10 + i * 0.1 + randomDoubleBetween(-0.1, 0.1, false);
                     GeoPoint point = (t.missingPointFactor > 0 && i % t.missingPointFactor == 0) ? null : new GeoPoint(lat, lon);
                     Long timestamp = (t.missingTimestampFactor > 0 && i % t.missingTimestampFactor == 0) ? null : startTime + 1000L * i;
                     points.add(point);
@@ -723,7 +723,7 @@ private void build() {
                         lessThanOrEqualTo(points.size() - deletedAtLeast)
                     );
                 }
-                if (t.useTimestamp && t.maxPoints < t.docCount) {
+                if (t.useTimeSeriesAggregation && t.maxPoints < t.docCount) {
                     // The aggregation will simplify the line in reverse order, so we need to anticipate the same simplification in the
                     // tests
                     TestGeometrySimplifierMonitor monitor = new TestGeometrySimplifierMonitor();
@@ -738,7 +738,7 @@ private void build() {
                     expectedAggSortValues.put(groups[g], line.sortValues);
                 } else {
                     // The aggregation will NOT simplify the line, so we should only anticipate the removal of invalid documents
-                    int pointCount = Math.min(t.maxPoints, expectedAggPointsList.size());  // possible truncation if !useTimestampField
+                    int pointCount = min(t.maxPoints, expectedAggPointsList.size());  // possible truncation if !useTimestampField
                     int offset = t.sortOrder == SortOrder.DESC ? 0 : Math.max(0, expectedAggPointsList.size() - pointCount);
                     long[] xp = new long[pointCount];
                     double[] xv = new double[pointCount];
@@ -771,13 +771,13 @@ private ArrayList<Long> timestampsForGroup(int g) {
 
     private record TestTSAssertionResults(
         MultiBucketsAggregation ts,
-        boolean useTimestampField,
+        boolean useTimeSeriesAggregation,
         String[] groups,
         Map<String, long[]> expectedAggPoints,
         Map<String, double[]> expectedAggSortValues
     ) {
         private TestTSAssertionResults(MultiBucketsAggregation ts, TestConfig testConfig, TestData testData) {
-            this(ts, testConfig.useTimestamp, testData.groups, testData.expectedAggPoints, testData.expectedAggSortValues);
+            this(ts, testConfig.useTimeSeriesAggregation, testData.groups, testData.expectedAggPoints, testData.expectedAggSortValues);
         }
     }
 
@@ -853,7 +853,7 @@ private void testAggregator(SortOrder sortOrder) throws IOException {
                 points[i] = lonLat;
                 sortValues[i] = i;
             }
-            int lineSize = Math.min(numPoints, size);
+            int lineSize = min(numPoints, size);
             // re-sort line to be ascending
             long[] linePoints = Arrays.copyOf(points, lineSize);
             double[] lineSorts = Arrays.copyOf(sortValues, lineSize);
@@ -942,12 +942,12 @@ private <A extends MultiBucketsAggregation, C extends Aggregator> void testCase(
             if (timeSeries) {
                 fieldTypes.add(TimeSeriesIdFieldMapper.FIELD_TYPE);
                 fieldTypes.add(new DateFieldMapper.DateFieldType("@timestamp"));
-                fieldTypes.add(new DateFieldMapper.DateFieldType("time_field"));
                 var metricType = randomBoolean() ? POSITION : null; // metric type does not affect geo_line behaviour
                 fieldTypes.add(new GeoPointFieldMapper.GeoPointFieldType("value_field", metricType, TIME_SERIES));
             } else {
                 fieldTypes.add(new GeoPointFieldMapper.GeoPointFieldType("value_field"));
             }
+            fieldTypes.add(new DateFieldMapper.DateFieldType("time_field"));
             fieldTypes.add(new KeywordFieldMapper.KeywordFieldType("group_id", false, true, Collections.emptyMap()));
             fieldTypes.add(new NumberFieldMapper.NumberFieldType("sort_field", NumberFieldMapper.NumberType.LONG));
             AggTestConfig aggTestConfig = new AggTestConfig(aggregationBuilder, fieldTypes.toArray(new MappedFieldType[0]));