Adds SpanGapQueryBuilder in the query DSL (#28636)

This change adds the support for a `span_gap` query inside the span query DSL.

Author: Chandan83

server/src/main/java/org/elasticsearch/index/query/SpanNearQueryBuilder.java

@@ -24,9 +24,11 @@
 import org.apache.lucene.search.spans.SpanQuery;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentLocation;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
@@ -203,18 +205,54 @@ public static SpanNearQueryBuilder fromXContent(XContentParser parser) throws IO
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
-        if (clauses.size() == 1) {
-            Query query = clauses.get(0).toQuery(context);
+        SpanQueryBuilder queryBuilder = clauses.get(0);
+        boolean isGap = queryBuilder instanceof SpanGapQueryBuilder;
+        Query query = null;
+        if (!isGap) {
+            query = queryBuilder.toQuery(context);
             assert query instanceof SpanQuery;
+        }
+        if (clauses.size() == 1) {
+            assert !isGap;
             return query;
         }
-        SpanQuery[] spanQueries = new SpanQuery[clauses.size()];
-        for (int i = 0; i < clauses.size(); i++) {
-            Query query = clauses.get(i).toQuery(context);
-            assert query instanceof SpanQuery;
-            spanQueries[i] = (SpanQuery) query;
+        String spanNearFieldName = null;
+        if (isGap) {
+            spanNearFieldName = ((SpanGapQueryBuilder) queryBuilder).fieldName();
+        } else {
+            spanNearFieldName = ((SpanQuery) query).getField();
         }
-        return new SpanNearQuery(spanQueries, slop, inOrder);
+
+        SpanNearQuery.Builder builder = new SpanNearQuery.Builder(spanNearFieldName, inOrder);
+        builder.setSlop(slop);
+        /*
+         * Lucene SpanNearQuery throws exceptions for certain use cases like adding gap to a
+         * unordered SpanNearQuery. Should ES have the same checks or wrap those thrown exceptions?
+         */
+        if (isGap) {
+            int gap = ((SpanGapQueryBuilder) queryBuilder).width();
+            builder.addGap(gap);
+        } else {
+            builder.addClause((SpanQuery) query);
+        }
+
+        for (int i = 1; i < clauses.size(); i++) {
+            queryBuilder = clauses.get(i);
+            isGap = queryBuilder instanceof SpanGapQueryBuilder;
+            if (isGap) {
+                String fieldName = ((SpanGapQueryBuilder) queryBuilder).fieldName();
+                if (!spanNearFieldName.equals(fieldName)) {
+                    throw new IllegalArgumentException("[span_near] clauses must have same field");
+                }
+                int gap = ((SpanGapQueryBuilder) queryBuilder).width();
+                builder.addGap(gap);
+            } else {
+                query = clauses.get(i).toQuery(context);
+                assert query instanceof SpanQuery;
+                builder.addClause((SpanQuery)query);
+            }
+        }
+        return builder.build();
     }
 
     @Override
@@ -233,4 +271,168 @@ protected boolean doEquals(SpanNearQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    /**
+     * SpanGapQueryBuilder enables gaps in a SpanNearQuery. 
+     * Since, SpanGapQuery is private to SpanNearQuery, SpanGapQueryBuilder cannot
+     * be used to generate a Query (SpanGapQuery) like another QueryBuilder.
+     * Instead, it just identifies a span_gap clause so that SpanNearQuery.addGap(int) 
+     * can be invoked for it. 
+     * This QueryBuilder is only applicable as a clause in SpanGapQueryBuilder but
+     * yet to enforce this restriction.
+     */
+    public static class SpanGapQueryBuilder implements SpanQueryBuilder {
+        public static final String NAME = "span_gap";
+
+        /** Name of field to match against. */
+        private final String fieldName;
+    
+        /** Width of the gap introduced. */
+        private final int width; 
+
+        /**
+         * Constructs a new SpanGapQueryBuilder term query.
+         *
+         * @param fieldName  The name of the field
+         * @param width The width of the gap introduced
+         */
+        public SpanGapQueryBuilder(String fieldName, int width) {
+            if (Strings.isEmpty(fieldName)) {
+                throw new IllegalArgumentException("[span_gap] field name is null or empty");
+            }
+            //lucene has not coded any restriction on value of width.
+            //to-do : find if theoretically it makes sense to apply restrictions. 
+            this.fieldName = fieldName;
+            this.width = width;
+        }
+
+        /**
+         * Read from a stream.
+         */
+        public SpanGapQueryBuilder(StreamInput in) throws IOException {
+            fieldName = in.readString();
+            width = in.readInt();
+        }
+
+        /**
+         * @return fieldName  The name of the field
+         */
+        public String fieldName() {
+            return fieldName;
+        }
+
+        /**
+         * @return width The width of the gap introduced
+         */
+        public int width() {
+            return width;
+        }
+
+        @Override
+        public Query toQuery(QueryShardContext context) throws IOException {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public Query toFilter(QueryShardContext context) throws IOException {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public String queryName() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public QueryBuilder queryName(String queryName) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public float boost() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public QueryBuilder boost(float boost) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public String getName() {
+            return NAME;
+        }
+
+        @Override
+        public String getWriteableName() {
+            return NAME;
+        }
+
+        @Override
+        public final void writeTo(StreamOutput out) throws IOException {
+            out.writeString(fieldName);
+            out.writeInt(width);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            builder.startObject();
+            builder.startObject(getName());
+            builder.field(fieldName, width);
+            builder.endObject();
+            builder.endObject();
+            return builder;
+        }
+
+        public static SpanGapQueryBuilder fromXContent(XContentParser parser) throws IOException {
+            String fieldName = null;
+            int width = 0;
+            String currentFieldName = null;
+            XContentParser.Token token;
+            while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                if (token == XContentParser.Token.FIELD_NAME) {
+                    currentFieldName = parser.currentName();
+                    throwParsingExceptionOnMultipleFields(NAME, parser.getTokenLocation(), fieldName, currentFieldName);
+                    fieldName = currentFieldName;
+                } else if (token.isValue()) {
+                    width = parser.intValue();
+                } 
+            }
+            SpanGapQueryBuilder result = new SpanGapQueryBuilder(fieldName, width);
+            return result;
+        }
+
+        @Override
+        public final boolean equals(Object obj) {
+            if (this == obj) {
+                return true;
+            }
+            if (obj == null || getClass() != obj.getClass()) {
+                return false;
+            }
+            SpanGapQueryBuilder other = (SpanGapQueryBuilder) obj;
+            return Objects.equals(fieldName, other.fieldName) &&
+                Objects.equals(width, other.width);
+        }
+
+        @Override
+        public final int hashCode() {
+            return Objects.hash(getClass(), fieldName, width);
+        }
+
+        
+        @Override
+        public final String toString() {
+            return Strings.toString(this, true, true);
+        }
+
+        //copied from AbstractQueryBuilder
+        protected static void throwParsingExceptionOnMultipleFields(String queryName, XContentLocation contentLocation,
+                String processedFieldName, String currentFieldName) {
+            if (processedFieldName != null) {
+                throw new ParsingException(contentLocation, "[" + queryName + "] query doesn't support multiple fields, found ["
+                        + processedFieldName + "] and [" + currentFieldName + "]");
+            }
+        }
+    }
 }

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -261,6 +261,7 @@
 
 import static java.util.Collections.unmodifiableMap;
 import static java.util.Objects.requireNonNull;
+import static org.elasticsearch.index.query.SpanNearQueryBuilder.SpanGapQueryBuilder;
 
 /**
  * Sets up things that can be done at search time like queries, aggregations, and suggesters.
@@ -743,6 +744,7 @@ private void registerQueryParsers(List<SearchPlugin> plugins) {
                 FieldMaskingSpanQueryBuilder::fromXContent));
         registerQuery(new QuerySpec<>(SpanFirstQueryBuilder.NAME, SpanFirstQueryBuilder::new, SpanFirstQueryBuilder::fromXContent));
         registerQuery(new QuerySpec<>(SpanNearQueryBuilder.NAME, SpanNearQueryBuilder::new, SpanNearQueryBuilder::fromXContent));
+        registerQuery(new QuerySpec<>(SpanGapQueryBuilder.NAME, SpanGapQueryBuilder::new, SpanGapQueryBuilder::fromXContent));
         registerQuery(new QuerySpec<>(SpanOrQueryBuilder.NAME, SpanOrQueryBuilder::new, SpanOrQueryBuilder::fromXContent));
         registerQuery(new QuerySpec<>(MoreLikeThisQueryBuilder.NAME, MoreLikeThisQueryBuilder::new,
                 MoreLikeThisQueryBuilder::fromXContent));

server/src/test/java/org/elasticsearch/index/query/SpanGapQueryBuilderTests.java

@@ -0,0 +1,127 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.index.query;
+
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.spans.SpanBoostQuery;
+import org.apache.lucene.search.spans.SpanNearQuery;
+import org.apache.lucene.search.spans.SpanQuery;
+import org.apache.lucene.search.spans.SpanTermQuery;
+import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.test.AbstractQueryTestCase;
+
+import java.io.IOException;
+import java.util.Iterator;
+
+import static org.elasticsearch.index.query.SpanNearQueryBuilder.SpanGapQueryBuilder;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.either;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.instanceOf;
+
+/*
+ * SpanGapQueryBuilder, unlike other QBs, is not used to build a Query. Therefore, it is not suited
+ * to test pattern of AbstractQueryTestCase. Since it is only used in SpanNearQueryBuilder, its test cases
+ * are same as those of later with SpanGapQueryBuilder included as clauses.
+ */
+
+public class SpanGapQueryBuilderTests extends AbstractQueryTestCase<SpanNearQueryBuilder> {
+    @Override
+    protected SpanNearQueryBuilder doCreateTestQueryBuilder() {
+        SpanTermQueryBuilder[] spanTermQueries = new SpanTermQueryBuilderTests().createSpanTermQueryBuilders(randomIntBetween(1, 6));
+        SpanNearQueryBuilder queryBuilder = new SpanNearQueryBuilder(spanTermQueries[0], randomIntBetween(-10, 10));
+        for (int i = 1; i < spanTermQueries.length; i++) {
+            SpanTermQueryBuilder termQB = spanTermQueries[i];
+            queryBuilder.addClause(termQB);
+            if (i % 2 == 1) {
+                SpanGapQueryBuilder gapQB = new SpanGapQueryBuilder(termQB.fieldName(), randomIntBetween(1,2));
+                queryBuilder.addClause(gapQB);
+            }
+        }
+        queryBuilder.inOrder(true);
+        return queryBuilder;
+    }
+
+    @Override
+    protected void doAssertLuceneQuery(SpanNearQueryBuilder queryBuilder, Query query, SearchContext context) throws IOException {
+        assertThat(query, either(instanceOf(SpanNearQuery.class))
+            .or(instanceOf(SpanTermQuery.class))
+            .or(instanceOf(SpanBoostQuery.class))
+            .or(instanceOf(MatchAllQueryBuilder.class)));
+        if (query instanceof SpanNearQuery) {
+            SpanNearQuery spanNearQuery = (SpanNearQuery) query;
+            assertThat(spanNearQuery.getSlop(), equalTo(queryBuilder.slop()));
+            assertThat(spanNearQuery.isInOrder(), equalTo(queryBuilder.inOrder()));
+            assertThat(spanNearQuery.getClauses().length, equalTo(queryBuilder.clauses().size()));
+            Iterator<SpanQueryBuilder> spanQueryBuilderIterator = queryBuilder.clauses().iterator();
+            for (SpanQuery spanQuery : spanNearQuery.getClauses()) {
+                SpanQueryBuilder spanQB = spanQueryBuilderIterator.next();
+                if (spanQB instanceof SpanGapQueryBuilder) continue;
+                assertThat(spanQuery, equalTo(spanQB.toQuery(context.getQueryShardContext())));
+            }
+        } else if (query instanceof SpanTermQuery || query instanceof SpanBoostQuery) {
+            assertThat(queryBuilder.clauses().size(), equalTo(1));
+            assertThat(query, equalTo(queryBuilder.clauses().get(0).toQuery(context.getQueryShardContext())));
+        }
+    }
+
+    public void testIllegalArguments() {
+            IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> new SpanGapQueryBuilder(null, 1));
+            assertEquals("[span_gap] field name is null or empty", e.getMessage());
+    }
+
+    public void testFromJson() throws IOException {
+        String json =
+                "{\n" +
+                "  \"span_near\" : {\n" +
+                "    \"clauses\" : [ {\n" +
+                "      \"span_term\" : {\n" +
+                "        \"field\" : {\n" +
+                "          \"value\" : \"value1\",\n" +
+                "          \"boost\" : 1.0\n" +
+                "        }\n" +
+                "      }\n" +
+                "    }, {\n" +
+                "      \"span_gap\" : {\n" +
+                "        \"field\" : 2" +
+                "      }\n" +
+                "    }, {\n" +
+                "      \"span_term\" : {\n" +
+                "        \"field\" : {\n" +
+                "          \"value\" : \"value3\",\n" +
+                "          \"boost\" : 1.0\n" +
+                "        }\n" +
+                "      }\n" +
+                "    } ],\n" +
+                "    \"slop\" : 12,\n" +
+                "    \"in_order\" : false,\n" +
+                "    \"boost\" : 1.0\n" +
+                "  }\n" +
+                "}";
+
+        SpanNearQueryBuilder parsed = (SpanNearQueryBuilder) parseQuery(json);
+        checkGeneratedJson(json, parsed);
+
+        assertEquals(json, 3, parsed.clauses().size());
+        assertEquals(json, 12, parsed.slop());
+        assertEquals(json, false, parsed.inOrder());
+    }
+}

server/src/test/java/org/elasticsearch/index/query/SpanNearQueryBuilderTests.java

@@ -184,4 +184,5 @@ public void testCollectPayloadsNoLongerSupported() throws Exception {
                 () -> parseQuery(json));
         assertThat(e.getMessage(), containsString("[span_near] query does not support [collect_payloads]"));
     }
+
 }