Counted keyword randomized testing (#121462)

jordan-powers · web-flow · commit 061920b5850f · 2025-02-04T13:06:26.000-08:00
This patch adds the needed data generator and source matcher to include counted_keyword fields in our randomized testing. This patch also updates the source matcher such that field-specific matchers are checked before the generic matcher is used. It seems that this is the correct behavior, and the only reason the generic matcher was checked first was as a workaround for issue #111916, which is now closed.
diff --git a/test/framework/build.gradle b/test/framework/build.gradle
@@ -33,6 +33,7 @@ dependencies {
   api "org.elasticsearch:mocksocket:${versions.mocksocket}"
 
   testImplementation project(':x-pack:plugin:mapper-unsigned-long')
+  testImplementation project(':x-pack:plugin:mapper-counted-keyword')
   testImplementation project(":modules:mapper-extras")
 }
 
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/FieldType.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/FieldType.java
@@ -11,6 +11,7 @@
 
 import org.elasticsearch.logsdb.datageneration.datasource.DataSource;
 import org.elasticsearch.logsdb.datageneration.fields.leaf.ByteFieldDataGenerator;
+import org.elasticsearch.logsdb.datageneration.fields.leaf.CountedKeywordFieldDataGenerator;
 import org.elasticsearch.logsdb.datageneration.fields.leaf.DoubleFieldDataGenerator;
 import org.elasticsearch.logsdb.datageneration.fields.leaf.FloatFieldDataGenerator;
 import org.elasticsearch.logsdb.datageneration.fields.leaf.HalfFloatFieldDataGenerator;
@@ -34,7 +35,8 @@ public enum FieldType {
     DOUBLE("double"),
     FLOAT("float"),
     HALF_FLOAT("half_float"),
-    SCALED_FLOAT("scaled_float");
+    SCALED_FLOAT("scaled_float"),
+    COUNTED_KEYWORD("counted_keyword");
 
     private final String name;
 
@@ -54,6 +56,7 @@ public FieldDataGenerator generator(String fieldName, DataSource dataSource) {
             case FLOAT -> new FloatFieldDataGenerator(fieldName, dataSource);
             case HALF_FLOAT -> new HalfFloatFieldDataGenerator(fieldName, dataSource);
             case SCALED_FLOAT -> new ScaledFloatFieldDataGenerator(fieldName, dataSource);
+            case COUNTED_KEYWORD -> new CountedKeywordFieldDataGenerator(fieldName, dataSource);
         };
     }
 
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceHandler.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceHandler.java
@@ -54,6 +54,10 @@ default DataSourceResponse.ArrayWrapper handle(DataSourceRequest.ArrayWrapper re
         return null;
     }
 
+    default DataSourceResponse.RepeatingWrapper handle(DataSourceRequest.RepeatingWrapper request) {
+        return null;
+    }
+
     default DataSourceResponse.ChildFieldGenerator handle(DataSourceRequest.ChildFieldGenerator request) {
         return null;
     }
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceRequest.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceRequest.java
@@ -85,6 +85,12 @@ public DataSourceResponse.ArrayWrapper accept(DataSourceHandler handler) {
         }
     }
 
+    record RepeatingWrapper() implements DataSourceRequest<DataSourceResponse.RepeatingWrapper> {
+        public DataSourceResponse.RepeatingWrapper accept(DataSourceHandler handler) {
+            return handler.handle(this);
+        }
+    }
+
     record ChildFieldGenerator(DataGeneratorSpecification specification)
         implements
             DataSourceRequest<DataSourceResponse.ChildFieldGenerator> {
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceResponse.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceResponse.java
@@ -39,6 +39,8 @@ record NullWrapper(Function<Supplier<Object>, Supplier<Object>> wrapper) impleme
 
     record ArrayWrapper(Function<Supplier<Object>, Supplier<Object>> wrapper) implements DataSourceResponse {}
 
+    record RepeatingWrapper(Function<Supplier<Object>, Supplier<Object>> wrapper) implements DataSourceResponse {}
+
     interface ChildFieldGenerator extends DataSourceResponse {
         int generateChildFieldCount();
 
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DefaultMappingParametersHandler.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DefaultMappingParametersHandler.java
@@ -33,6 +33,7 @@ public DataSourceResponse.LeafMappingParametersGenerator handle(DataSourceReques
             case KEYWORD -> keywordMapping(request, map);
             case LONG, INTEGER, SHORT, BYTE, DOUBLE, FLOAT, HALF_FLOAT, UNSIGNED_LONG -> plain(map);
             case SCALED_FLOAT -> scaledFloatMapping(map);
+            case COUNTED_KEYWORD -> plain(Map.of("index", ESTestCase.randomBoolean()));
         });
     }
 
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DefaultWrappersHandler.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DefaultWrappersHandler.java
@@ -11,6 +11,7 @@
 
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.HashSet;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
@@ -26,6 +27,11 @@ public DataSourceResponse.ArrayWrapper handle(DataSourceRequest.ArrayWrapper ign
         return new DataSourceResponse.ArrayWrapper(wrapInArray());
     }
 
+    @Override
+    public DataSourceResponse.RepeatingWrapper handle(DataSourceRequest.RepeatingWrapper ignored) {
+        return new DataSourceResponse.RepeatingWrapper(repeatValues());
+    }
+
     private static Function<Supplier<Object>, Supplier<Object>> injectNulls() {
         // Inject some nulls but majority of data should be non-null (as it likely is in reality).
         return (values) -> () -> ESTestCase.randomDouble() <= 0.05 ? null : values.get();
@@ -41,4 +47,19 @@ private static Function<Supplier<Object>, Supplier<Object>> wrapInArray() {
             return values.get();
         };
     }
+
+    private static Function<Supplier<Object>, Supplier<Object>> repeatValues() {
+        return (values) -> {
+            HashSet<Object> previousValues = new HashSet<>();
+            return () -> {
+                if (previousValues.size() > 0 && ESTestCase.randomBoolean()) {
+                    return ESTestCase.randomFrom(previousValues);
+                } else {
+                    var value = values.get();
+                    previousValues.add(value);
+                    return value;
+                }
+            };
+        };
+    }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/fields/leaf/CountedKeywordFieldDataGenerator.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/fields/leaf/CountedKeywordFieldDataGenerator.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.logsdb.datageneration.fields.leaf;
+
+import org.elasticsearch.logsdb.datageneration.FieldDataGenerator;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSource;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.function.Supplier;
+
+public class CountedKeywordFieldDataGenerator implements FieldDataGenerator {
+    private final Supplier<Object> valueGenerator;
+    private final Set<String> previousStrings = new HashSet<>();
+
+    public CountedKeywordFieldDataGenerator(String fieldName, DataSource dataSource) {
+        var strings = dataSource.get(new DataSourceRequest.StringGenerator());
+        var nulls = dataSource.get(new DataSourceRequest.NullWrapper());
+        var arrays = dataSource.get(new DataSourceRequest.ArrayWrapper());
+        var repeats = dataSource.get(new DataSourceRequest.RepeatingWrapper());
+
+        this.valueGenerator = arrays.wrapper().compose(nulls.wrapper().compose(repeats.wrapper())).apply(() -> strings.generator().get());
+    }
+
+    @Override
+    public Object generateValue() {
+        return valueGenerator.get();
+    }
+}
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/matchers/source/FieldSpecificMatcher.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/matchers/source/FieldSpecificMatcher.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.math.BigInteger;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -223,4 +224,68 @@ private static BigInteger toBigInteger(Object value) {
             return (BigInteger) value;
         }
     }
+
+    class CountedKeywordMatcher implements FieldSpecificMatcher {
+        private final XContentBuilder actualMappings;
+        private final Settings.Builder actualSettings;
+        private final XContentBuilder expectedMappings;
+        private final Settings.Builder expectedSettings;
+
+        CountedKeywordMatcher(
+            XContentBuilder actualMappings,
+            Settings.Builder actualSettings,
+            XContentBuilder expectedMappings,
+            Settings.Builder expectedSettings
+        ) {
+            this.actualMappings = actualMappings;
+            this.actualSettings = actualSettings;
+            this.expectedMappings = expectedMappings;
+            this.expectedSettings = expectedSettings;
+        }
+
+        private static List<String> normalize(List<Object> values) {
+            return values.stream().filter(Objects::nonNull).map(it -> (String) it).toList();
+        }
+
+        private static boolean matchCountsEqualExact(List<String> actualNormalized, List<String> expectedNormalized) {
+            HashMap<String, Integer> counts = new HashMap<>();
+            for (String value : actualNormalized) {
+                counts.put(value, counts.getOrDefault(value, 0) + 1);
+            }
+            for (String value : expectedNormalized) {
+                int newCount = counts.getOrDefault(value, 0) - 1;
+                if (newCount == 0) {
+                    counts.remove(value);
+                } else {
+                    counts.put(value, newCount);
+                }
+            }
+
+            return counts.isEmpty();
+        }
+
+        @Override
+        public MatchResult match(
+            List<Object> actual,
+            List<Object> expected,
+            Map<String, Object> actualMapping,
+            Map<String, Object> expectedMapping
+        ) {
+            var actualNormalized = normalize(actual);
+            var expectedNormalized = normalize(expected);
+
+            return matchCountsEqualExact(actualNormalized, expectedNormalized)
+                ? MatchResult.match()
+                : MatchResult.noMatch(
+                    formatErrorMessage(
+                        actualMappings,
+                        actualSettings,
+                        expectedMappings,
+                        expectedSettings,
+                        "Values of type [counted_keyword] don't match after normalization, normalized"
+                            + prettyPrintCollections(actualNormalized, expectedNormalized)
+                    )
+                );
+        }
+    }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/matchers/source/SourceMatcher.java b/test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/matchers/source/SourceMatcher.java
@@ -57,7 +57,9 @@ public SourceMatcher(
             "scaled_float",
             new FieldSpecificMatcher.ScaledFloatMatcher(actualMappings, actualSettings, expectedMappings, expectedSettings),
             "unsigned_long",
-            new FieldSpecificMatcher.UnsignedLongMatcher(actualMappings, actualSettings, expectedMappings, expectedSettings)
+            new FieldSpecificMatcher.UnsignedLongMatcher(actualMappings, actualSettings, expectedMappings, expectedSettings),
+            "counted_keyword",
+            new FieldSpecificMatcher.CountedKeywordMatcher(actualMappings, actualSettings, expectedMappings, expectedSettings)
         );
         this.dynamicFieldMatcher = new DynamicFieldMatcher(actualMappings, actualSettings, expectedMappings, expectedSettings);
     }
@@ -100,17 +102,8 @@ private MatchResult compareSource(Map<String, List<Object>> actual, Map<String,
             var actualValues = actual.get(name);
             var expectedValues = expectedFieldEntry.getValue();
 
-            // There are cases when field values are stored in ignored source
-            // so we try to match them as is first and then apply field specific matcher.
-            // This is temporary, we should be able to tell when source is exact using mappings.
-            // See #111916.
-            var genericMatchResult = matchWithGenericMatcher(actualValues, expectedValues);
-            if (genericMatchResult.isMatch()) {
-                continue;
-            }
-
-            var matchIncludingFieldSpecificMatchers = matchWithFieldSpecificMatcher(name, actualValues, expectedValues).orElse(
-                genericMatchResult
+            var matchIncludingFieldSpecificMatchers = matchWithFieldSpecificMatcher(name, actualValues, expectedValues).orElseGet(
+                () -> matchWithGenericMatcher(actualValues, expectedValues)
             );
             if (matchIncludingFieldSpecificMatchers.isMatch() == false) {
                 var message = "Source documents don't match for field [" + name + "]: " + matchIncludingFieldSpecificMatchers.getMessage();
diff --git a/test/framework/src/test/java/org/elasticsearch/logsdb/datageneration/DataGenerationTests.java b/test/framework/src/test/java/org/elasticsearch/logsdb/datageneration/DataGenerationTests.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.countedkeyword.CountedKeywordMapperPlugin;
 import org.elasticsearch.xpack.unsignedlong.UnsignedLongMapperPlugin;
 
 import java.io.IOException;
@@ -110,7 +111,7 @@ public DataSourceResponse.FieldTypeGenerator handle(DataSourceRequest.FieldTypeG
         var mappingService = new MapperServiceTestCase() {
             @Override
             protected Collection<? extends Plugin> getPlugins() {
-                return List.of(new UnsignedLongMapperPlugin(), new MapperExtrasPlugin());
+                return List.of(new UnsignedLongMapperPlugin(), new MapperExtrasPlugin(), new CountedKeywordMapperPlugin());
             }
         }.createMapperService(mappingXContent);
 
diff --git a/test/framework/src/test/java/org/elasticsearch/logsdb/datageneration/SourceMatcherTests.java b/test/framework/src/test/java/org/elasticsearch/logsdb/datageneration/SourceMatcherTests.java
@@ -105,4 +105,38 @@ public void testMappedMismatch() throws IOException {
         var sut = new SourceMatcher(mapping, Settings.builder(), mapping, Settings.builder(), actual, expected, false);
         assertFalse(sut.match().isMatch());
     }
+
+    public void testCountedKeywordMatch() throws IOException {
+        List<Map<String, Object>> actual = List.of(Map.of("field", List.of("a", "b", "a", "c", "b", "a")));
+        List<Map<String, Object>> expected = List.of(Map.of("field", List.of("a", "b", "a", "c", "b", "a")));
+
+        var mapping = XContentBuilder.builder(XContentType.JSON.xContent());
+        mapping.startObject();
+        mapping.startObject("_doc");
+        {
+            mapping.startObject("field").field("type", "counted_keyword").endObject();
+        }
+        mapping.endObject();
+        mapping.endObject();
+
+        var sut = new SourceMatcher(mapping, Settings.builder(), mapping, Settings.builder(), actual, expected, false);
+        assertTrue(sut.match().isMatch());
+    }
+
+    public void testCountedKeywordMismatch() throws IOException {
+        List<Map<String, Object>> actual = List.of(Map.of("field", List.of("a", "b", "a", "c", "b", "a")));
+        List<Map<String, Object>> expected = List.of(Map.of("field", List.of("a", "b", "c", "a")));
+
+        var mapping = XContentBuilder.builder(XContentType.JSON.xContent());
+        mapping.startObject();
+        mapping.startObject("_doc");
+        {
+            mapping.startObject("field").field("type", "counted_keyword").endObject();
+        }
+        mapping.endObject();
+        mapping.endObject();
+
+        var sut = new SourceMatcher(mapping, Settings.builder(), mapping, Settings.builder(), actual, expected, false);
+        assertFalse(sut.match().isMatch());
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@ dependencies {`
`33`	`33`	`api "org.elasticsearch:mocksocket:${versions.mocksocket}"`
`34`	`34`
`35`	`35`	`testImplementation project(':x-pack:plugin:mapper-unsigned-long')`
	`36`	`+ testImplementation project(':x-pack:plugin:mapper-counted-keyword')`
`36`	`37`	`testImplementation project(":modules:mapper-extras")`
`37`	`38`	`}`
`38`	`39`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,10 @@ default DataSourceResponse.ArrayWrapper handle(DataSourceRequest.ArrayWrapper re`
`54`	`54`	`return null;`
`55`	`55`	`}`
`56`	`56`
	`57`	`+ default DataSourceResponse.RepeatingWrapper handle(DataSourceRequest.RepeatingWrapper request) {`
	`58`	`+ return null;`
	`59`	`+ }`
	`60`	`+`
`57`	`61`	`default DataSourceResponse.ChildFieldGenerator handle(DataSourceRequest.ChildFieldGenerator request) {`
`58`	`62`	`return null;`
`59`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,12 @@ public DataSourceResponse.ArrayWrapper accept(DataSourceHandler handler) {`
`85`	`85`	`}`
`86`	`86`	`}`
`87`	`87`
	`88`	`+ record RepeatingWrapper() implements DataSourceRequest<DataSourceResponse.RepeatingWrapper> {`
	`89`	`+ public DataSourceResponse.RepeatingWrapper accept(DataSourceHandler handler) {`
	`90`	`+ return handler.handle(this);`
	`91`	`+ }`
	`92`	`+ }`
	`93`	`+`
`88`	`94`	`record ChildFieldGenerator(DataGeneratorSpecification specification)`
`89`	`95`	`implements`
`90`	`96`	`DataSourceRequest<DataSourceResponse.ChildFieldGenerator> {`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@ public DataSourceResponse.LeafMappingParametersGenerator handle(DataSourceReques`
`33`	`33`	`case KEYWORD -> keywordMapping(request, map);`
`34`	`34`	`case LONG, INTEGER, SHORT, BYTE, DOUBLE, FLOAT, HALF_FLOAT, UNSIGNED_LONG -> plain(map);`
`35`	`35`	`case SCALED_FLOAT -> scaledFloatMapping(map);`
	`36`	`+ case COUNTED_KEYWORD -> plain(Map.of("index", ESTestCase.randomBoolean()));`
`36`	`37`	`});`
`37`	`38`	`}`
`38`	`39`