Add filter function for AbstractQueryBuilder, BoolQueryBuilder, Const…

…antScoreQueryBuilder. (#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]>
opensearch-project · Feb 26, 2025 · a2983d4 · a2983d4
1 parent e397903
commit a2983d4
Show file tree

Hide file tree

Showing 11 changed files with 181 additions and 0 deletions.
diff --git a/CHANGELOG-3.0.md b/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))

diff --git a/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java
@@ -86,6 +86,40 @@ 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
+     * @param filterCombinationMode filter combination mode.
+     * @return true if parameters are valid. Returns false when the filter is null. Throws IllegalArgumentException on
+     * null filterCombinationMode.
+     */
+    public static boolean validateFilterParams(QueryBuilder filter, FilterCombinationMode filterCombinationMode) {
+        if (filter == null) {
+            return false;
+        }
+        if (filterCombinationMode == null) {
+            // Unexpected FilterCombinationMode
+            throw new IllegalArgumentException("FilterCombinationMode cannot be null.");
+        }
+        return true;
+    }
+
+    /**
+     * Combine filter with current query builder based on filterCombinationMode
+     * @param filter filter to combine with current query builder
+     * @param filterCombinationMode filter combination mode.
+     * @return query builder with filter combined
+     */
+    public QueryBuilder filter(QueryBuilder filter, FilterCombinationMode filterCombinationMode) {
+        if (!validateFilterParams(filter, filterCombinationMode)) {
+            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);

diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java
@@ -152,6 +152,20 @@ public List<QueryBuilder> filter() {
         return this.filterClauses;
     }
 
+    /**
+     * Add query filter to filterClauses.
+     * @param filter the query filter
+     * @param filterCombinationMode the filter combination mode
+     * @return BoolQueryBuilder
+     */
+    public BoolQueryBuilder filter(QueryBuilder filter, FilterCombinationMode filterCombinationMode) {
+        if (!validateFilterParams(filter, filterCombinationMode)) {
+            return this;
+        }
+        filter(filter);
+        return this;
+    }
+
     /**
      * Adds a query that <b>must not</b> appear in the matching documents.
      * No {@code null} value allowed.

diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/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
+     * @param filterCombinationMode the combination mode for the filter. Currently support {@link FilterCombinationMode#AND}
+     */
+    public ConstantScoreQueryBuilder filter(QueryBuilder filter, FilterCombinationMode filterCombinationMode) {
+        if (!validateFilterParams(filter, filterCombinationMode)) {
+            return this;
+        }
+        QueryBuilder filteredFilterBuilder = filterBuilder.filter(filter, filterCombinationMode);
+        if (filteredFilterBuilder != filterBuilder) {
+            return new ConstantScoreQueryBuilder(filteredFilterBuilder);
+        }
+        return this;
+    }
+
     public static ConstantScoreQueryBuilder fromXContent(XContentParser parser) throws IOException {
         QueryBuilder query = null;
         boolean queryFound = false;

diff --git a/server/src/main/java/org/opensearch/index/query/FilterCombinationMode.java b/server/src/main/java/org/opensearch/index/query/FilterCombinationMode.java
@@ -0,0 +1,20 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.query;
+
+import org.opensearch.common.annotation.PublicApi;
+
+/**
+ * Currently support AND mode.
+ * May extend for supporting more modes in the future (OR, IGNORE_IF_EXISTS, etc.).
+ */
+@PublicApi(since = "3.0.0")
+public enum FilterCombinationMode {
+    AND
+}
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilder.java b/server/src/main/java/org/opensearch/index/query/QueryBuilder.java
@@ -47,6 +47,19 @@
 @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.
+     * @param filterCombinationMode The non-null mode to combine the existing filter with the new filter.
+     * @return A QueryBuilder with the filter added to the existing filter.
+     */
+    QueryBuilder filter(QueryBuilder filter, FilterCombinationMode filterCombinationMode);
+
     /**
      * Converts this QueryBuilder to a lucene {@link Query}.
      * Returns {@code null} if this query should be ignored in the context of

diff --git a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java b/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, FilterCombinationMode filterCombinationMode) {
+            throw new UnsupportedOperationException("You can't add a filter to a SpanGapQueryBuilder");
+        }
+
         @Override
         public String queryName() {
             throw new UnsupportedOperationException();

diff --git a/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java
@@ -326,6 +326,31 @@ 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 AND filter case
+        String query = "{\"bool\" : {\"filter\" : null } }";
+        QueryBuilder filter = QueryBuilders.matchAllQuery();
+        BoolQueryBuilder builder = (BoolQueryBuilder) parseQuery(query);
+        assertFalse(builder.filter(filter, FilterCombinationMode.AND).filter().isEmpty());
+        assertEquals(builder.filter(filter, FilterCombinationMode.AND).filter().get(0), filter);
+
+        // Test for null case
+        builder = (BoolQueryBuilder) parseQuery(query);
+        assertTrue(builder.filter(null, FilterCombinationMode.AND).filter().isEmpty());
+
+        // Test for null filterCombinationMode
+        final QueryBuilder builderForUnsupportedFilter = (BoolQueryBuilder) parseQuery(query);
+        assertThrows(
+            "FilterCombinationMode is null.",
+            IllegalArgumentException.class,
+            () -> builderForUnsupportedFilter.filter(filter, null)
+        );
+    }
+
     /**
      * test that unknown query names in the clauses throw an error
      */

diff --git a/server/src/test/java/org/opensearch/index/query/ConstantScoreQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/ConstantScoreQueryBuilderTests.java
@@ -143,4 +143,29 @@ public void testVisit() {
 
         assertEquals(2, visitorQueries.size());
     }
+
+    public void testFilter() {
+        // Test for for AND Filter Combination Mode
+        BoolQueryBuilder filterBuilder = new BoolQueryBuilder();
+        ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(filterBuilder);
+        QueryBuilder filter = QueryBuilders.matchAllQuery();
+        constantScoreQueryBuilder.filter(filter, FilterCombinationMode.AND);
+        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, FilterCombinationMode.AND);
+        assertEquals(0, filterBuilder.filter().size());
+
+        // Test for null Filter Combination Mode
+        filterBuilder = new BoolQueryBuilder();
+        final ConstantScoreQueryBuilder queryBuilderForUnsupportedMode = new ConstantScoreQueryBuilder(filterBuilder);
+        assertThrows(
+            "FilterCombinationMode is null.",
+            IllegalArgumentException.class,
+            () -> queryBuilderForUnsupportedMode.filter(filter, null)
+        );
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/SpanMultiTermQueryBuilderTests.java b/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, FilterCombinationMode filterCombinationMode) {
+            return this;
+        }
     }
 
     @Override

diff --git a/test/framework/src/main/java/org/opensearch/test/AbstractQueryTestCase.java b/test/framework/src/main/java/org/opensearch/test/AbstractQueryTestCase.java
@@ -63,7 +63,10 @@
 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.FilterCombinationMode;
 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 +871,24 @@ 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, FilterCombinationMode.AND);
+        assertEquals(queryBuilder, returnedQuerybuilder);
+
+        // Test for unspported Filter Combination Mode
+        assertThrows("FilterCombinationMode is null.", IllegalArgumentException.class, () -> queryBuilder.filter(filter, null));
+
+        // Test for AND case
+        returnedQuerybuilder = queryBuilder.filter(filter, FilterCombinationMode.AND);
+        assertTrue(returnedQuerybuilder instanceof BoolQueryBuilder);
+        assertTrue(((BoolQueryBuilder) returnedQuerybuilder).filter().size() == 1);
+        assertEquals(filter, ((BoolQueryBuilder) returnedQuerybuilder).filter().get(0));
+    }
 }