Add support for 'size' in EQL Sample queries

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/BaseEqlSpecTestCase.java

@@ -46,6 +46,7 @@ public abstract class BaseEqlSpecTestCase extends RemoteClusterAwareEqlRestTestC
      * For now, every value will be converted to a String.
      */
     private final String[] joinKeys;
+    private final Integer size;
 
     @Before
     public void setup() throws Exception {
@@ -97,19 +98,20 @@ protected static List<Object[]> asArray(List<EqlSpec> specs) {
                 name = "" + (counter);
             }
 
-            results.add(new Object[] { spec.query(), name, spec.expectedEventIds(), spec.joinKeys() });
+            results.add(new Object[] { spec.query(), name, spec.expectedEventIds(), spec.joinKeys(), spec.size() });
         }
 
         return results;
     }
 
-    BaseEqlSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys) {
+    BaseEqlSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
         this.index = index;
 
         this.query = query;
         this.name = name;
         this.eventIds = eventIds;
         this.joinKeys = joinKeys;
+        this.size = size;
     }
 
     public void test() throws Exception {
@@ -139,7 +141,7 @@ protected ObjectPath runQuery(String index, String query) throws Exception {
         if (tiebreaker != null) {
             builder.field("tiebreaker_field", tiebreaker);
         }
-        builder.field("size", requestSize());
+        builder.field("size", this.size == null ? requestSize() : this.size);
         builder.field("fetch_size", requestFetchSize());
         builder.field("result_position", requestResultPosition());
         builder.endObject();

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlDateNanosSpecTestCase.java

@@ -21,13 +21,13 @@ public static List<Object[]> readTestSpecs() throws Exception {
     }
 
     // constructor for "local" rest tests
-    public EqlDateNanosSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        this(TEST_NANOS_INDEX, query, name, eventIds, joinKeys);
+    public EqlDateNanosSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_NANOS_INDEX, query, name, eventIds, joinKeys, size);
     }
 
     // constructor for multi-cluster tests
-    public EqlDateNanosSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlDateNanosSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 
     @Override

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlExtraSpecTestCase.java

@@ -21,13 +21,13 @@ public static List<Object[]> readTestSpecs() throws Exception {
     }
 
     // constructor for "local" rest tests
-    public EqlExtraSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        this(TEST_EXTRA_INDEX, query, name, eventIds, joinKeys);
+    public EqlExtraSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_EXTRA_INDEX, query, name, eventIds, joinKeys, size);
     }
 
     // constructor for multi-cluster tests
-    public EqlExtraSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlExtraSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 
     @Override

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlSampleTestCase.java

@@ -15,12 +15,12 @@
 
 public abstract class EqlSampleTestCase extends BaseEqlSpecTestCase {
 
-    public EqlSampleTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        this(TEST_SAMPLE, query, name, eventIds, joinKeys);
+    public EqlSampleTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_SAMPLE, query, name, eventIds, joinKeys, size);
     }
 
-    public EqlSampleTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlSampleTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 
     @ParametersFactory(shuffle = false, argumentFormatting = PARAM_FORMATTING)

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlSpec.java

@@ -28,6 +28,8 @@ public class EqlSpec {
     private List<long[]> expectedEventIds;
     private String[] joinKeys;
 
+    private Integer size;
+
     public String name() {
         return name;
     }
@@ -84,6 +86,14 @@ public void joinKeys(String[] joinKeys) {
         this.joinKeys = joinKeys;
     }
 
+    public Integer size() {
+        return size;
+    }
+
+    public void size(Integer size) {
+        this.size = size;
+    }
+
     @Override
     public String toString() {
         String str = "";
@@ -107,6 +117,10 @@ public String toString() {
         if (joinKeys != null) {
             str = appendWithComma(str, "join_keys", Arrays.toString(joinKeys));
         }
+        if (size != null) {
+            str = appendWithComma(str, "size", "" + size);
+        }
+
         return str;
     }
 
@@ -122,12 +136,12 @@ public boolean equals(Object other) {
 
         EqlSpec that = (EqlSpec) other;
 
-        return Objects.equals(this.query(), that.query());
+        return Objects.equals(this.query(), that.query()) && Objects.equals(size, that.size);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(this.query);
+        return Objects.hash(this.query, size);
     }
 
     private static String appendWithComma(String str, String name, String append) {

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlSpecLoader.java

@@ -68,6 +68,14 @@ private static String getTrimmedString(TomlTable table, String key) {
         return null;
     }
 
+    private static Integer getInteger(TomlTable table, String key) {
+        Long s = table.getLong(key);
+        if (s != null) {
+            return s.intValue();
+        }
+        return null;
+    }
+
     private static List<EqlSpec> readFromStream(InputStream is, Set<String> uniqueTestNames) throws Exception {
         List<EqlSpec> testSpecs = new ArrayList<>();
 
@@ -81,6 +89,7 @@ private static List<EqlSpec> readFromStream(InputStream is, Set<String> uniqueTe
             spec.name(getTrimmedString(table, "name"));
             spec.note(getTrimmedString(table, "note"));
             spec.description(getTrimmedString(table, "description"));
+            spec.size(getInteger(table, "size"));
 
             List<?> arr = table.getList("tags");
             if (arr != null) {

x-pack/plugin/eql/qa/common/src/main/java/org/elasticsearch/test/eql/EqlSpecTestCase.java

@@ -28,12 +28,12 @@ protected String tiebreaker() {
     }
 
     // constructor for "local" rest tests
-    public EqlSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        this(TEST_INDEX, query, name, eventIds, joinKeys);
+    public EqlSpecTestCase(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        this(TEST_INDEX, query, name, eventIds, joinKeys, size);
     }
 
     // constructor for multi-cluster tests
-    public EqlSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(index, query, name, eventIds, joinKeys);
+    public EqlSpecTestCase(String index, String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(index, query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/common/src/main/resources/test_sample.toml

@@ -286,3 +286,61 @@ expected_event_ids  = [18,11]
 join_keys = ["doom","win10"]
 
 
+[[queries]]
+name = "size0"
+size = 0
+query = '''
+  sample by host, ?os
+    [success where true]
+    [alert where true and id == 21 or id == 24 or id == 11]
+'''
+expected_event_ids  = []
+join_keys = []
+
+
+[[queries]]
+name = "size2"
+size = 2
+query = '''
+  sample by host, ?os
+    [success where true]
+    [alert where true and id == 21 or id == 24 or id == 11]
+'''
+expected_event_ids  = [29,24,
+    28,21]
+join_keys = ["GTA","null",
+    "doom","null"]
+
+
+[[queries]]
+name = "size3"
+size = 3
+query = '''
+  sample by host, ?os
+    [success where true]
+    [alert where true and id == 21 or id == 24 or id == 11]
+'''
+expected_event_ids  = [29,24,
+    28,21,
+    18,11]
+join_keys = ["GTA","null",
+    "doom","null",
+    "doom","win10"]
+
+
+[[queries]]
+name = "sizeBig"
+size = 500
+query = '''
+  sample by host, ?os
+    [success where true]
+    [alert where true and id == 21 or id == 24 or id == 11]
+'''
+expected_event_ids  = [29,24,
+    28,21,
+    18,11]
+join_keys = ["GTA","null",
+    "doom","null",
+    "doom","win10"]
+
+

x-pack/plugin/eql/qa/multi-cluster-with-security/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlDateNanosIT.java

@@ -16,7 +16,7 @@
 
 public class EqlDateNanosIT extends EqlDateNanosSpecTestCase {
 
-    public EqlDateNanosIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(remoteClusterIndex(TEST_NANOS_INDEX), query, name, eventIds, joinKeys);
+    public EqlDateNanosIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(remoteClusterIndex(TEST_NANOS_INDEX), query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/multi-cluster-with-security/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlExtraIT.java

@@ -16,7 +16,7 @@
 
 public class EqlExtraIT extends EqlExtraSpecTestCase {
 
-    public EqlExtraIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(remoteClusterIndex(TEST_EXTRA_INDEX), query, name, eventIds, joinKeys);
+    public EqlExtraIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(remoteClusterIndex(TEST_EXTRA_INDEX), query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/multi-cluster-with-security/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlSampleIT.java

@@ -16,7 +16,7 @@
 
 public class EqlSampleIT extends EqlSampleTestCase {
 
-    public EqlSampleIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(remoteClusterPattern(TEST_SAMPLE), query, name, eventIds, joinKeys);
+    public EqlSampleIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(remoteClusterPattern(TEST_SAMPLE), query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/multi-cluster-with-security/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlSpecIT.java

@@ -16,7 +16,7 @@
 
 public class EqlSpecIT extends EqlSpecTestCase {
 
-    public EqlSpecIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(remoteClusterIndex(TEST_INDEX), query, name, eventIds, joinKeys);
+    public EqlSpecIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(remoteClusterIndex(TEST_INDEX), query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/rest/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlDateNanosIT.java

@@ -13,7 +13,7 @@
 
 public class EqlDateNanosIT extends EqlDateNanosSpecTestCase {
 
-    public EqlDateNanosIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(query, name, eventIds, joinKeys);
+    public EqlDateNanosIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/rest/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlExtraIT.java

@@ -13,7 +13,7 @@
 
 public class EqlExtraIT extends EqlExtraSpecTestCase {
 
-    public EqlExtraIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(query, name, eventIds, joinKeys);
+    public EqlExtraIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/qa/rest/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlSampleIT.java

@@ -13,8 +13,8 @@
 
 public class EqlSampleIT extends EqlSampleTestCase {
 
-    public EqlSampleIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(query, name, eventIds, joinKeys);
+    public EqlSampleIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(query, name, eventIds, joinKeys, size);
     }
 
 }

x-pack/plugin/eql/qa/rest/src/javaRestTest/java/org/elasticsearch/xpack/eql/EqlSpecIT.java

@@ -13,7 +13,7 @@
 
 public class EqlSpecIT extends EqlSpecTestCase {
 
-    public EqlSpecIT(String query, String name, List<long[]> eventIds, String[] joinKeys) {
-        super(query, name, eventIds, joinKeys);
+    public EqlSpecIT(String query, String name, List<long[]> eventIds, String[] joinKeys, Integer size) {
+        super(query, name, eventIds, joinKeys, size);
     }
 }

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/assembler/ExecutionManager.java

@@ -158,7 +158,7 @@ public Executable assemble(
     /*
      * Sample assembler
      */
-    public Executable assemble(List<List<Attribute>> listOfKeys, List<PhysicalPlan> plans) {
+    public Executable assemble(List<List<Attribute>> listOfKeys, List<PhysicalPlan> plans, Limit limit) {
         if (cfg.fetchSize() > SAMPLE_MAX_PAGE_SIZE) {
             throw new EqlIllegalArgumentException("Fetch size cannot be greater than [{}]", SAMPLE_MAX_PAGE_SIZE);
         }
@@ -211,7 +211,7 @@ public Executable assemble(List<List<Attribute>> listOfKeys, List<PhysicalPlan>
             }
         }
 
-        return new SampleIterator(new PITAwareQueryClient(session), criteria, cfg.fetchSize(), session.circuitBreaker());
+        return new SampleIterator(new PITAwareQueryClient(session), criteria, cfg.fetchSize(), limit, session.circuitBreaker());
     }
 
     private HitExtractor timestampExtractor(HitExtractor hitExtractor) {

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/sample/SampleIterator.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.eql.execution.assembler.SampleCriterion;
 import org.elasticsearch.xpack.eql.execution.assembler.SampleQueryRequest;
 import org.elasticsearch.xpack.eql.execution.search.HitReference;
+import org.elasticsearch.xpack.eql.execution.search.Limit;
 import org.elasticsearch.xpack.eql.execution.search.QueryClient;
 import org.elasticsearch.xpack.eql.execution.search.RuntimeUtils;
 import org.elasticsearch.xpack.eql.execution.sequence.SequenceKey;
@@ -55,7 +56,7 @@ public class SampleIterator implements Executable {
     private final int maxCriteria;
     final List<Sample> samples;
     private final int fetchSize;
-
+    private final Limit limit;
     private long startTime;
 
     // ---------- CIRCUIT BREAKER -----------
@@ -83,12 +84,13 @@ public class SampleIterator implements Executable {
      */
     private long previousTotalPageSize = 0;
 
-    public SampleIterator(QueryClient client, List<SampleCriterion> criteria, int fetchSize, CircuitBreaker circuitBreaker) {
+    public SampleIterator(QueryClient client, List<SampleCriterion> criteria, int fetchSize, Limit limit, CircuitBreaker circuitBreaker) {
         this.client = client;
         this.criteria = criteria;
         this.maxCriteria = criteria.size();
         this.fetchSize = fetchSize;
         this.samples = new ArrayList<>();
+        this.limit = limit;
         this.circuitBreaker = circuitBreaker;
     }
 
@@ -213,7 +215,13 @@ private void finalStep(ActionListener<Payload> listener) {
                 if (docGroupsCounter == maxCriteria) {
                     List<SearchHit> match = matchSample(sample, maxCriteria);
                     if (match != null) {
-                        samples.add(new Sample(sampleKeys.get(responseIndex / maxCriteria), match));
+                        if (samples.size() < limit.limit()) {
+                            samples.add(new Sample(sampleKeys.get(responseIndex / maxCriteria), match));
+                        }
+                        if (samples.size() == limit.limit()) {
+                            payload(listener);
+                            return;
+                        }
                     }
                     docGroupsCounter = 1;
                     sample = new ArrayList<>(maxCriteria);

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/parser/LogicalPlanBuilder.java

@@ -102,7 +102,7 @@ public Object visitStatement(StatementContext ctx) {
             if (ctx.pipe().size() > 0) {
                 throw new ParsingException(source(ctx.pipe().get(0)), "Samples do not support pipes yet");
             }
-            return plan;
+            return new LimitWithOffset(plan.source(), new Literal(Source.EMPTY, params.size(), DataTypes.INTEGER), 0, plan);
         }
         //
         // Add implicit blocks