Apply boolean must_not rewrite to numeric match, term, and terms queries

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Rename WorkloadGroupTestUtil to WorkloadManagementTestUtil ([#18709](https://github.com/opensearch-project/OpenSearch/pull/18709))
 - Disallow resize for Warm Index, add Parameterized ITs for close in remote store ([#18686](https://github.com/opensearch-project/OpenSearch/pull/18686))
 - Ability to run Code Coverage with Gradle and produce the jacoco reports locally ([#18509](https://github.com/opensearch-project/OpenSearch/issues/18509))
+- Extend BooleanQuery must_not rewrite to numeric must, term, and terms queries ([#18498](https://github.com/opensearch-project/OpenSearch/pull/18498))
 - [Workload Management] Update logging and Javadoc, rename QueryGroup to WorkloadGroup ([#18711](https://github.com/opensearch-project/OpenSearch/issues/18711))
 - Add NodeResourceUsageStats to ClusterInfo ([#18480](https://github.com/opensearch-project/OpenSearch/issues/18472))
 - Introduce SecureHttpTransportParameters experimental API (to complement SecureTransportParameters counterpart) ([#18572](https://github.com/opensearch-project/OpenSearch/issues/18572))

server/src/internalClusterTest/java/org/opensearch/search/query/BooleanQueryIT.java

@@ -13,13 +13,18 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
+import java.util.Map;
 
 import static org.opensearch.index.query.QueryBuilders.boolQuery;
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
+import static org.opensearch.index.query.QueryBuilders.matchQuery;
 import static org.opensearch.index.query.QueryBuilders.rangeQuery;
 import static org.opensearch.index.query.QueryBuilders.termQuery;
+import static org.opensearch.index.query.QueryBuilders.termsQuery;
 import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
 
@@ -229,6 +234,53 @@ public void testMustNotRangeRewriteWithMoreThanOneValue() throws Exception {
         assertHitCount(client().prepareSearch().setQuery(matchAllQuery()).get(), numDocs);
     }
 
+    public void testMustNotNumericMatchOrTermQueryRewrite() throws Exception {
+        Map<Integer, Integer> statusToDocCountMap = Map.of(200, 1000, 404, 30, 500, 1, 400, 1293);
+        String statusField = "status";
+        createIndex("test");
+        int totalDocs = 0;
+        for (Map.Entry<Integer, Integer> entry : statusToDocCountMap.entrySet()) {
+            for (int i = 0; i < entry.getValue(); i++) {
+                client().prepareIndex("test").setSource(statusField, entry.getKey()).get();
+            }
+            totalDocs += entry.getValue();
+        }
+        ensureGreen();
+        waitForRelocation();
+        forceMerge();
+        refresh();
+
+        int excludedValue = randomFrom(statusToDocCountMap.keySet());
+        int expectedHitCount = totalDocs - statusToDocCountMap.get(excludedValue);
+
+        // Check the rewritten match query behaves as expected
+        assertHitCount(
+            client().prepareSearch("test").setQuery(boolQuery().mustNot(matchQuery(statusField, excludedValue))).get(),
+            expectedHitCount
+        );
+
+        // Check the rewritten term query behaves as expected
+        assertHitCount(
+            client().prepareSearch("test").setQuery(boolQuery().mustNot(termQuery(statusField, excludedValue))).get(),
+            expectedHitCount
+        );
+
+        // Check a rewritten terms query behaves as expected
+        List<Integer> excludedValues = new ArrayList<>();
+        excludedValues.add(excludedValue);
+        int secondExcludedValue = randomFrom(statusToDocCountMap.keySet());
+        expectedHitCount = totalDocs - statusToDocCountMap.get(excludedValue);
+        if (secondExcludedValue != excludedValue) {
+            excludedValues.add(secondExcludedValue);
+            expectedHitCount -= statusToDocCountMap.get(secondExcludedValue);
+        }
+
+        assertHitCount(
+            client().prepareSearch("test").setQuery(boolQuery().mustNot(termsQuery(statusField, excludedValues))).get(),
+            expectedHitCount
+        );
+    }
+
     public void testMustToFilterRewrite() throws Exception {
         // Check we still get expected behavior after rewriting must clauses --> filter clauses.
         String intField = "int_field";

server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java

@@ -1855,6 +1855,10 @@ public byte[] encodePoint(Number value) {
         public double toDoubleValue(long value) {
             return type.toDoubleValue(value);
         }
+
+        public Number parse(Object value) {
+            return type.parse(value, coerce);
+        }
     }
 
     private final NumberType type;

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

@@ -485,31 +485,33 @@ private boolean rewriteMustNotRangeClausesToShould(BoolQueryBuilder newBuilder,
             return false;
         }
 
+        QueryShardContext shardContext = getQueryShardContext(queryRewriteContext);
+
         boolean changed = false;
-        // For now, only handle the case where there's exactly 1 range query for this field.
+        // For now, only handle the case where there's exactly 1 complement-aware query for this field.
         Map<String, Integer> fieldCounts = new HashMap<>();
-        Set<RangeQueryBuilder> rangeQueries = new HashSet<>();
+        Set<ComplementAwareQueryBuilder> complementAwareQueries = new HashSet<>();
         for (QueryBuilder clause : mustNotClauses) {
-            if (clause instanceof RangeQueryBuilder rq) {
-                fieldCounts.merge(rq.fieldName(), 1, Integer::sum);
-                rangeQueries.add(rq);
+            if (clause instanceof ComplementAwareQueryBuilder && clause instanceof WithFieldName wfn) {
+                fieldCounts.merge(wfn.fieldName(), 1, Integer::sum);
+                complementAwareQueries.add((ComplementAwareQueryBuilder) wfn);
             }
         }
 
-        for (RangeQueryBuilder rq : rangeQueries) {
-            String fieldName = rq.fieldName();
+        for (ComplementAwareQueryBuilder caq : complementAwareQueries) {
+            String fieldName = ((WithFieldName) caq).fieldName();
             if (fieldCounts.getOrDefault(fieldName, 0) == 1) {
                 // Check that all docs on this field have exactly 1 value, otherwise we can't perform this rewrite
                 if (checkAllDocsHaveOneValue(leafReaderContexts, fieldName)) {
-                    List<RangeQueryBuilder> complement = rq.getComplement();
+                    List<? extends QueryBuilder> complement = caq.getComplement(shardContext);
                     if (complement != null) {
                         BoolQueryBuilder nestedBoolQuery = new BoolQueryBuilder();
                         nestedBoolQuery.minimumShouldMatch(1);
-                        for (RangeQueryBuilder complementComponent : complement) {
+                        for (QueryBuilder complementComponent : complement) {
                             nestedBoolQuery.should(complementComponent);
                         }
                         newBuilder.must(nestedBoolQuery);
-                        newBuilder.mustNotClauses.remove(rq);
+                        newBuilder.mustNotClauses.remove(caq);
                         changed = true;
                     }
                 }
@@ -537,6 +539,10 @@ private List<LeafReaderContext> getLeafReaderContexts(QueryRewriteContext queryR
         return indexSearcher.getIndexReader().leaves();
     }
 
+    private QueryShardContext getQueryShardContext(QueryRewriteContext queryRewriteContext) {
+        return queryRewriteContext == null ? null : queryRewriteContext.convertToShardContext(); // Note this can still be null
+    }
+
     private boolean checkAllDocsHaveOneValue(List<LeafReaderContext> contexts, String fieldName) {
         for (LeafReaderContext lrc : contexts) {
             PointValues values;

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

@@ -0,0 +1,23 @@
+/*
+ * 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 java.util.List;
+
+/**
+ * A QueryBuilder which can provide QueryBuilders that make up the complement of the original query.
+ */
+public interface ComplementAwareQueryBuilder {
+    /**
+     * Returns a list of RangeQueryBuilder whose elements, when combined, form the complement of this range query.
+     * May be null, in which case the complement couldn't be determined.
+     * @return the complement
+     */
+    List<? extends QueryBuilder> getComplement(QueryShardContext context);
+}

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

@@ -0,0 +1,81 @@
+/*
+ * 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.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.NumberFieldMapper;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A class with helper functions to construct complements for some queries.
+ */
+public class ComplementHelperUtils {
+    /**
+     * Get the NumberFieldType for this fieldName from the context, or null if it isn't a NumberFieldType.
+     */
+    public static NumberFieldMapper.NumberFieldType getNumberFieldType(QueryShardContext context, String fieldName) {
+        if (context == null) return null;
+        MappedFieldType fieldType = context.fieldMapper(fieldName);
+        if (!(fieldType instanceof NumberFieldMapper.NumberFieldType nft)) return null;
+        return nft;
+    }
+
+    /**
+     * Returns a list of 2 RangeQueryBuilders matching everything but value.
+     */
+    public static List<QueryBuilder> numberValueToComplement(String fieldName, Number value) {
+        List<QueryBuilder> complement = new ArrayList<>();
+        RangeQueryBuilder belowRange = new RangeQueryBuilder(fieldName);
+        belowRange.to(value);
+        belowRange.includeUpper(false);
+        complement.add(belowRange);
+
+        RangeQueryBuilder aboveRange = new RangeQueryBuilder(fieldName);
+        aboveRange.from(value);
+        aboveRange.includeLower(false);
+        complement.add(aboveRange);
+        return complement;
+    }
+
+    /**
+     * Returns a list of RangeQueryBuilders matching everything except the provided sorted values.
+     * if isWholeNumber == true, and two sorted values are off by 1, the range between the two of them won't appear in
+     * the complement since no value could match it.
+     */
+    public static List<QueryBuilder> numberValuesToComplement(String fieldName, List<Number> sortedValues, boolean isWholeNumber) {
+        if (sortedValues.isEmpty()) return null;
+        List<QueryBuilder> complement = new ArrayList<>();
+        Number lastValue = null;
+        for (Number value : sortedValues) {
+            RangeQueryBuilder range = new RangeQueryBuilder(fieldName);
+            range.includeUpper(false);
+            range.to(value);
+            if (lastValue != null) {
+                // If this is a whole number field and the last value is 1 less than the current value, we can skip this part of the
+                // complement
+                if (isWholeNumber && value.longValue() - lastValue.longValue() == 1) {
+                    continue;
+                }
+                range.includeLower(false);
+                range.from(lastValue);
+            }
+            complement.add(range);
+            lastValue = value;
+        }
+        // Finally add the last range query
+        RangeQueryBuilder lastRange = new RangeQueryBuilder(fieldName);
+        lastRange.from(sortedValues.get(sortedValues.size() - 1));
+        lastRange.includeLower(false);
+        lastRange.includeUpper(true);
+        complement.add(lastRange);
+        return complement;
+    }
+}

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

@@ -43,11 +43,13 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.index.mapper.NumberFieldMapper;
 import org.opensearch.index.query.support.QueryParsers;
 import org.opensearch.index.search.MatchQuery;
 import org.opensearch.index.search.MatchQuery.ZeroTermsQuery;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -56,7 +58,7 @@
  *
  * @opensearch.internal
  */
-public class MatchQueryBuilder extends AbstractQueryBuilder<MatchQueryBuilder> implements WithFieldName {
+public class MatchQueryBuilder extends AbstractQueryBuilder<MatchQueryBuilder> implements ComplementAwareQueryBuilder, WithFieldName {
 
     private static final String CUTOFF_FREQUENCY_DEPRECATION_MSG = "you can omit this option, "
         + "the [match] query can skip block of documents efficiently if the total number of hits is not tracked";
@@ -589,4 +591,12 @@ public static MatchQueryBuilder fromXContent(XContentParser parser) throws IOExc
         return matchQuery;
     }
 
+    @Override
+    public List<QueryBuilder> getComplement(QueryShardContext context) {
+        // If this is a match query on a numeric field, we can provide the complement using RangeQueryBuilder.
+        NumberFieldMapper.NumberFieldType nft = ComplementHelperUtils.getNumberFieldType(context, fieldName);
+        if (nft == null) return null;
+        Number numberValue = nft.parse(value);
+        return ComplementHelperUtils.numberValueToComplement(fieldName, numberValue);
+    }
 }

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

@@ -60,7 +60,10 @@
  *
  * @opensearch.internal
  */
-public class RangeQueryBuilder extends AbstractQueryBuilder<RangeQueryBuilder> implements MultiTermQueryBuilder {
+public class RangeQueryBuilder extends AbstractQueryBuilder<RangeQueryBuilder>
+    implements
+        MultiTermQueryBuilder,
+        ComplementAwareQueryBuilder {
     public static final String NAME = "range";
 
     public static final boolean DEFAULT_INCLUDE_UPPER = true;
@@ -545,12 +548,9 @@ protected boolean doEquals(RangeQueryBuilder other) {
             && Objects.equals(format, other.format);
     }
 
-    /**
-     * Returns a list of RangeQueryBuilder whose elements, when combined, form the complement of this range query.
-     * May be null.
-     * @return the complement
-     */
-    public List<RangeQueryBuilder> getComplement() {
+    @Override
+    public List<? extends QueryBuilder> getComplement(QueryShardContext context) {
+        // This implementation doesn't need info from QueryShardContext
         if (relation != null && relation != ShapeRelation.INTERSECTS) {
             return null;
         }

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

@@ -43,16 +43,18 @@
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.mapper.ConstantFieldType;
 import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.NumberFieldMapper;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Objects;
 
 /**
  * A Query that matches documents containing a term.
  *
  * @opensearch.internal
  */
-public class TermQueryBuilder extends BaseTermQueryBuilder<TermQueryBuilder> {
+public class TermQueryBuilder extends BaseTermQueryBuilder<TermQueryBuilder> implements ComplementAwareQueryBuilder {
     public static final String NAME = "term";
     public static final boolean DEFAULT_CASE_INSENSITIVITY = false;
     private static final ParseField CASE_INSENSITIVE_FIELD = new ParseField("case_insensitive");
@@ -238,4 +240,12 @@ protected final boolean doEquals(TermQueryBuilder other) {
         return super.doEquals(other) && Objects.equals(caseInsensitive, other.caseInsensitive);
     }
 
+    @Override
+    public List<QueryBuilder> getComplement(QueryShardContext context) {
+        // If this is a term query on a numeric field, we can provide the complement using RangeQueryBuilder.
+        NumberFieldMapper.NumberFieldType nft = ComplementHelperUtils.getNumberFieldType(context, fieldName);
+        if (nft == null) return null;
+        Number numberValue = nft.parse(value);
+        return ComplementHelperUtils.numberValueToComplement(fieldName, numberValue);
+    }
 }

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

@@ -66,6 +66,7 @@
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
@@ -79,7 +80,7 @@
  *
  * @opensearch.internal
  */
-public class TermsQueryBuilder extends AbstractQueryBuilder<TermsQueryBuilder> implements WithFieldName {
+public class TermsQueryBuilder extends AbstractQueryBuilder<TermsQueryBuilder> implements ComplementAwareQueryBuilder, WithFieldName {
     public static final String NAME = "terms";
 
     private final String fieldName;
@@ -639,4 +640,26 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
 
         return this;
     }
+
+    @Override
+    public List<QueryBuilder> getComplement(QueryShardContext context) {
+        // If this uses BITMAP value type, or if we're using termsLookup, we can't provide the complement.
+        if (valueType.equals(ValueType.BITMAP)) return null;
+        if (values == null || termsLookup != null) return null;
+        // If this is a terms query on a numeric field, we can provide the complement using RangeQueryBuilder.
+        NumberFieldMapper.NumberFieldType nft = ComplementHelperUtils.getNumberFieldType(context, fieldName);
+        if (nft == null) return null;
+        List<Number> numberValues = new ArrayList<>();
+        for (Object value : values) {
+            numberValues.add(nft.parse(value));
+        }
+        numberValues.sort(Comparator.comparingDouble(Number::doubleValue)); // For sorting purposes, use double value.
+        NumberFieldMapper.NumberType numberType = nft.numberType();
+        // If there is some other field type that's a whole number, this will still be correct, the complement may just have some
+        // unnecessary components like "x < value < x + 1"
+        boolean isWholeNumber = numberType.equals(NumberFieldMapper.NumberType.INTEGER)
+            || numberType.equals(NumberFieldMapper.NumberType.LONG)
+            || numberType.equals(NumberFieldMapper.NumberType.SHORT);
+        return ComplementHelperUtils.numberValuesToComplement(fieldName, numberValues, isWholeNumber);
+    }
 }