Add filter function for AbstractQueryBuilder, BoolQueryBuilder, ConstantScoreQueryBuilder. (#17409) (#17409)

* The filter function will combine a filter with the query builder. If the query builder itself has a filter we will combine the filter and return the query builder itself. If no we will use a bool query builder to combine the query builder and the filter and then return the bool query builder.

Signed-off-by: Chloe Gao <[email protected]>

Author: chloewqg

CHANGELOG-3.0.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add execution_hint to cardinality aggregator request (#[17312](https://github.com/opensearch-project/OpenSearch/pull/17312))
 - Arrow Flight RPC plugin with Flight server bootstrap logic and client for internode communication ([#16962](https://github.com/opensearch-project/OpenSearch/pull/16962))
 - Added offset management for the pull-based Ingestion ([#17354](https://github.com/opensearch-project/OpenSearch/pull/17354))
+- Add filter function for AbstractQueryBuilder, BoolQueryBuilder, ConstantScoreQueryBuilder([#17409](https://github.com/opensearch-project/OpenSearch/pull/17409))
 
 ### Dependencies
 - Update Apache Lucene to 10.1.0 ([#16366](https://github.com/opensearch-project/OpenSearch/pull/16366))

server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java

@@ -86,6 +86,30 @@ protected AbstractQueryBuilder(StreamInput in) throws IOException {
         queryName = in.readOptionalString();
     }
 
+    /**
+     * Check the input parameters of filter function.
+     * @param filter filter to combine with current query builder
+     * @return true if parameters are valid. Returns false when the filter is null.
+     */
+    public static boolean validateFilterParams(QueryBuilder filter) {
+        return filter != null;
+    }
+
+    /**
+     * Combine filter with current query builder
+     * @param filter filter to combine with current query builder
+     * @return query builder with filter combined
+     */
+    public QueryBuilder filter(QueryBuilder filter) {
+        if (validateFilterParams(filter) == false) {
+            return this;
+        }
+        final BoolQueryBuilder modifiedQB = new BoolQueryBuilder();
+        modifiedQB.must(this);
+        modifiedQB.filter(filter);
+        return modifiedQB;
+    }
+
     @Override
     public final void writeTo(StreamOutput out) throws IOException {
         out.writeFloat(boost);

server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java

@@ -135,13 +135,15 @@ public List<QueryBuilder> must() {
 
     /**
      * Adds a query that <b>must</b> appear in the matching documents but will
-     * not contribute to scoring. No {@code null} value allowed.
+     * not contribute to scoring. If null value passed, then do nothing and return.
+     * @param filter the filter to add to the current ConstantScoreQuery
+     * @return query builder with filter combined
      */
-    public BoolQueryBuilder filter(QueryBuilder queryBuilder) {
-        if (queryBuilder == null) {
-            throw new IllegalArgumentException("inner bool query clause cannot be null");
+    public BoolQueryBuilder filter(QueryBuilder filter) {
+        if (validateFilterParams(filter) == false) {
+            return this;
         }
-        filterClauses.add(queryBuilder);
+        filterClauses.add(filter);
         return this;
     }
 

server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java

@@ -101,6 +101,22 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
         builder.endObject();
     }
 
+    /**
+     * Adds a filter to the current ConstantScoreQuery.
+     * @param filter the filter to add to the current ConstantScoreQuery
+     * @return query builder with filter combined
+     */
+    public ConstantScoreQueryBuilder filter(QueryBuilder filter) {
+        if (validateFilterParams(filter) == false) {
+            return this;
+        }
+        QueryBuilder filteredFilterBuilder = filterBuilder.filter(filter);
+        if (filteredFilterBuilder != filterBuilder) {
+            return new ConstantScoreQueryBuilder(filteredFilterBuilder);
+        }
+        return this;
+    }
+
     public static ConstantScoreQueryBuilder fromXContent(XContentParser parser) throws IOException {
         QueryBuilder query = null;
         boolean queryFound = false;

server/src/main/java/org/opensearch/index/query/QueryBuilder.java

@@ -47,6 +47,18 @@
 @PublicApi(since = "1.0.0")
 public interface QueryBuilder extends NamedWriteable, ToXContentObject, Rewriteable<QueryBuilder> {
 
+    /**
+     * This function combines a filter with a query builder. If the query builder itself has
+     * a filter we will combine the filter and return the query builder itself.
+     * If not we will use a bool query builder to combine the query builder and
+     * the filter and then return the bool query builder.
+     * If the filter is null we simply return the query builder without any operation.
+     *
+     * @param filter The null filter to be added to the existing filter.
+     * @return A QueryBuilder with the filter added to the existing filter.
+     */
+    QueryBuilder filter(QueryBuilder filter);
+
     /**
      * Converts this QueryBuilder to a lucene {@link Query}.
      * Returns {@code null} if this query should be ignored in the context of

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

@@ -375,6 +375,11 @@ public Query toQuery(QueryShardContext context) throws IOException {
             throw new UnsupportedOperationException();
         }
 
+        @Override
+        public QueryBuilder filter(QueryBuilder filter) {
+            throw new UnsupportedOperationException("You can't add a filter to a SpanGapQueryBuilder");
+        }
+
         @Override
         public String queryName() {
             throw new UnsupportedOperationException();

server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java

@@ -178,7 +178,6 @@ public void testIllegalArguments() {
         BoolQueryBuilder booleanQuery = new BoolQueryBuilder();
         expectThrows(IllegalArgumentException.class, () -> booleanQuery.must(null));
         expectThrows(IllegalArgumentException.class, () -> booleanQuery.mustNot(null));
-        expectThrows(IllegalArgumentException.class, () -> booleanQuery.filter(null));
         expectThrows(IllegalArgumentException.class, () -> booleanQuery.should(null));
     }
 
@@ -326,6 +325,23 @@ public void testFilterNull() throws IOException {
         assertTrue(builder.filter().isEmpty());
     }
 
+    /**
+     * Check if a filter can be applied to the BoolQuery
+     * @throws IOException
+     */
+    public void testFilter() throws IOException {
+        // Test for non null filter
+        String query = "{\"bool\" : {\"filter\" : null } }";
+        QueryBuilder filter = QueryBuilders.matchAllQuery();
+        BoolQueryBuilder builder = (BoolQueryBuilder) parseQuery(query);
+        assertFalse(builder.filter(filter).filter().isEmpty());
+        assertEquals(builder.filter(filter).filter().get(0), filter);
+
+        // Test for null filter case
+        builder = (BoolQueryBuilder) parseQuery(query);
+        assertTrue(builder.filter(null).filter().isEmpty());
+    }
+
     /**
      * test that unknown query names in the clauses throw an error
      */

server/src/test/java/org/opensearch/index/query/ConstantScoreQueryBuilderTests.java

@@ -143,4 +143,21 @@ public void testVisit() {
 
         assertEquals(2, visitorQueries.size());
     }
+
+    public void testFilter() {
+        // Test for non null filter
+        BoolQueryBuilder filterBuilder = new BoolQueryBuilder();
+        ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(filterBuilder);
+        QueryBuilder filter = QueryBuilders.matchAllQuery();
+        constantScoreQueryBuilder.filter(filter);
+        assertEquals(1, filterBuilder.filter().size());
+        assertEquals(filter, filterBuilder.filter().get(0));
+
+        // Test for null filter
+        filterBuilder = new BoolQueryBuilder();
+        constantScoreQueryBuilder = new ConstantScoreQueryBuilder(filterBuilder);
+        constantScoreQueryBuilder.filter(null);
+        assertEquals(0, filterBuilder.filter().size());
+
+    }
 }

server/src/test/java/org/opensearch/index/query/SpanMultiTermQueryBuilderTests.java

@@ -182,6 +182,11 @@ public void writeTo(StreamOutput out) throws IOException {
         public String fieldName() {
             return "foo";
         }
+
+        @Override
+        public QueryBuilder filter(QueryBuilder filter) {
+            return this;
+        }
     }
 
     @Override

test/framework/src/main/java/org/opensearch/test/AbstractQueryTestCase.java

@@ -63,7 +63,9 @@
 import org.opensearch.core.xcontent.XContentParseException;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.query.AbstractQueryBuilder;
+import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.Rewriteable;
@@ -868,4 +870,21 @@ public void testCacheability() throws IOException {
         assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
+    /**
+     * Check if a filter can be applied to the abstract query builder.
+     * @throws UnsupportedOperationException
+     */
+    public void testFilter() throws IOException {
+        QB queryBuilder = createTestQueryBuilder();
+        QueryBuilder filter = QueryBuilders.matchAllQuery();
+        // Test for Null Filter case
+        QueryBuilder returnedQuerybuilder = queryBuilder.filter(null);
+        assertEquals(queryBuilder, returnedQuerybuilder);
+
+        // Test for non null filter
+        returnedQuerybuilder = queryBuilder.filter(filter);
+        assertTrue(returnedQuerybuilder instanceof BoolQueryBuilder);
+        assertTrue(((BoolQueryBuilder) returnedQuerybuilder).filter().size() == 1);
+        assertEquals(filter, ((BoolQueryBuilder) returnedQuerybuilder).filter().get(0));
+    }
 }