Add DataSource RepeatingWrapper

Author: jordan-powers

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;
     }

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> {

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();
 

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 (ESTestCase.randomBoolean()) {
+                    return ESTestCase.randomFrom(previousValues);
+                } else {
+                    var value = values.get();
+                    previousValues.add(value);
+                    return value;
+                }
+            };
+        };
+    }
 }

test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/fields/leaf/CountedKeywordFieldDataGenerator.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.logsdb.datageneration.FieldDataGenerator;
 import org.elasticsearch.logsdb.datageneration.datasource.DataSource;
 import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
-import org.elasticsearch.test.ESTestCase;
 
 import java.util.HashSet;
 import java.util.Set;
@@ -26,16 +25,9 @@ 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()).apply(() -> {
-            if (previousStrings.size() > 0 && ESTestCase.randomBoolean()) {
-                return ESTestCase.randomFrom(previousStrings);
-            } else {
-                String value = strings.generator().get();
-                previousStrings.add(value);
-                return value;
-            }
-        });
+        this.valueGenerator = arrays.wrapper().compose(nulls.wrapper().compose(repeats.wrapper())).apply(() -> strings.generator().get());
     }
 
     @Override