Address comments

Signed-off-by: Peter Alfonsi <petealft@amazon.com>

Author: Peter Alfonsi

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

@@ -485,22 +485,22 @@ private boolean rewriteMustNotRangeClausesToShould(BoolQueryBuilder newBuilder,
         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<ComplementAwareQueryBuilder> complementAwareQueries = new HashSet<>();
         for (QueryBuilder clause : mustNotClauses) {
-            if (clause instanceof ComplementAwareQueryBuilder caq) {
-                fieldCounts.merge(caq.fieldName(), 1, Integer::sum);
-                complementAwareQueries.add(caq);
+            if (clause instanceof ComplementAwareQueryBuilder && clause instanceof WithFieldName wfn) {
+                fieldCounts.merge(wfn.fieldName(), 1, Integer::sum);
+                complementAwareQueries.add((ComplementAwareQueryBuilder) wfn);
             }
         }
 
         for (ComplementAwareQueryBuilder caq : complementAwareQueries) {
-            String fieldName = caq.fieldName();
+            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<QueryBuilder> complement = caq.getComplement(shardContext);
+                    List<? extends QueryBuilder> complement = caq.getComplement(shardContext);
                     if (complement != null) {
                         BoolQueryBuilder nestedBoolQuery = new BoolQueryBuilder();
                         nestedBoolQuery.minimumShouldMatch(1);

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

@@ -13,11 +13,11 @@
 /**
  * A QueryBuilder which can provide QueryBuilders that make up the complement of the original query.
  */
-public interface ComplementAwareQueryBuilder extends WithFieldName {
+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<QueryBuilder> getComplement(QueryShardContext context);
+    List<? extends QueryBuilder> getComplement(QueryShardContext context);
 }

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

@@ -58,7 +58,7 @@
  *
  * @opensearch.internal
  */
-public class MatchQueryBuilder extends AbstractQueryBuilder<MatchQueryBuilder> implements ComplementAwareQueryBuilder {
+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";

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

@@ -549,12 +549,12 @@ protected boolean doEquals(RangeQueryBuilder other) {
     }
 
     @Override
-    public List<QueryBuilder> getComplement(QueryShardContext context) {
+    public List<? extends QueryBuilder> getComplement(QueryShardContext context) {
         // This implementation doesn't need info from QueryShardContext
         if (relation != null && relation != ShapeRelation.INTERSECTS) {
             return null;
         }
-        List<QueryBuilder> complement = new ArrayList<>();
+        List<RangeQueryBuilder> complement = new ArrayList<>();
         if (from != null) {
             RangeQueryBuilder belowRange = new RangeQueryBuilder(fieldName);
             belowRange.to(from);
@@ -570,13 +570,13 @@ public List<QueryBuilder> getComplement(QueryShardContext context) {
         }
 
         if (format != null) {
-            for (QueryBuilder rq : complement) {
-                ((RangeQueryBuilder) rq).format(format);
+            for (RangeQueryBuilder rq : complement) {
+                rq.format(format);
             }
         }
         if (timeZone != null) {
-            for (QueryBuilder rq : complement) {
-                ((RangeQueryBuilder) rq).timeZone = timeZone;
+            for (RangeQueryBuilder rq : complement) {
+                rq.timeZone = timeZone;
             }
         }
 

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

@@ -80,7 +80,7 @@
  *
  * @opensearch.internal
  */
-public class TermsQueryBuilder extends AbstractQueryBuilder<TermsQueryBuilder> implements ComplementAwareQueryBuilder {
+public class TermsQueryBuilder extends AbstractQueryBuilder<TermsQueryBuilder> implements ComplementAwareQueryBuilder, WithFieldName {
     public static final String NAME = "terms";
 
     private final String fieldName;

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

@@ -669,7 +669,7 @@ static void testGetComplementNumericField(
         String fieldName,
         IndexSearcher searcher
     ) {
-        List<QueryBuilder> complement = queryBuilder.getComplement(createShardContext(searcher));
+        List<? extends QueryBuilder> complement = queryBuilder.getComplement(createShardContext(searcher));
         RangeQueryBuilder expectedLower = new RangeQueryBuilder(fieldName).to(value).includeLower(true).includeUpper(false);
         RangeQueryBuilder expectedUpper = new RangeQueryBuilder(fieldName).from(value).includeLower(false).includeUpper(true);
         assertEquals(2, complement.size());