[Feature] Visitor design pattern in QueryBuilder (opensearch-project#…

…10110) * Implementation of Visitor Design Pattern in QueryBuilder Signed-off-by: Varun Jain <varunudr@amazon.com> * Adding Changelog Signed-off-by: Varun Jain <varunudr@amazon.com> * Adding Changelog Signed-off-by: Varun Jain <varunudr@amazon.com> * Comments addressed of michael Signed-off-by: Varun Jain <varunudr@amazon.com> * Adding test case and some minor fixes in DisMaxQueryBuilder Signed-off-by: Varun Jain <varunudr@amazon.com> * Modifying test cases and remove exeception Signed-off-by: Varun Jain <varunudr@amazon.com> * Adding bool query builder test cases at clause level Signed-off-by: Varun Jain <varunudr@amazon.com> * Fixing test case Signed-off-by: Varun Jain <varunudr@amazon.com> * Additional test cases Signed-off-by: Varun Jain <varunudr@amazon.com> --------- Signed-off-by: Varun Jain <varunudr@amazon.com>
kotwanikunal · Sep 21, 2023 · a63f8de · a63f8de
1 parent 7dca3ca
commit a63f8de
Show file tree

Hide file tree

Showing 29 changed files with 377 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add events correlation engine plugin ([#6854](https://github.com/opensearch-project/OpenSearch/issues/6854))
 - Introduce new dynamic cluster setting to control slice computation for concurrent segment search ([#9107](https://github.com/opensearch-project/OpenSearch/pull/9107))
 - Implement on behalf of token passing for extensions ([#8679](https://github.com/opensearch-project/OpenSearch/pull/8679))
+- Implement Visitor Design pattern in QueryBuilder to enable the capability to traverse through the complex QueryBuilder tree. ([#10110](https://github.com/opensearch-project/OpenSearch/pull/10110))
 
 ### Dependencies
 - Bump `log4j-core` from 2.18.0 to 2.19.0

diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java
@@ -427,4 +427,35 @@ private static boolean rewriteClauses(
         }
         return changed;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (mustClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST);
+            for (QueryBuilder mustClause : mustClauses) {
+                mustClause.visit(subVisitor);
+            }
+        }
+        if (shouldClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.SHOULD);
+            for (QueryBuilder shouldClause : shouldClauses) {
+                shouldClause.visit(subVisitor);
+            }
+        }
+        if (mustNotClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST_NOT);
+            for (QueryBuilder mustNotClause : mustNotClauses) {
+                mustNotClause.visit(subVisitor);
+            }
+        }
+        if (filterClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.FILTER);
+            for (QueryBuilder filterClause : filterClauses) {
+                filterClause.visit(subVisitor);
+            }
+        }
+
+    }
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java
@@ -33,6 +33,7 @@
 package org.opensearch.index.query;
 
 import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -252,4 +253,15 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
         InnerHitContextBuilder.extractInnerHits(positiveQuery, innerHits);
         InnerHitContextBuilder.extractInnerHits(negativeQuery, innerHits);
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (positiveQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery);
+        }
+        if (negativeQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
@@ -183,4 +184,11 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
     protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> innerHits) {
         InnerHitContextBuilder.extractInnerHits(filterBuilder, innerHits);
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder);
+    }
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.DisjunctionMaxQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.common.lucene.search.Queries;
@@ -246,4 +247,15 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
             InnerHitContextBuilder.extractInnerHits(query, innerHits);
         }
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (queries.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : queries) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.FieldMaskingSpanQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -207,4 +208,10 @@ protected boolean doEquals(FieldMaskingSpanQueryBuilder other) {
     public String getWriteableName() {
         return SPAN_FIELD_MASKING_FIELD.getPreferredName();
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilder.java b/server/src/main/java/org/opensearch/index/query/QueryBuilder.java
@@ -95,4 +95,13 @@ public interface QueryBuilder extends NamedWriteable, ToXContentObject, Rewritea
     default QueryBuilder rewrite(QueryRewriteContext queryShardContext) throws IOException {
         return this;
     }
+
+    /**
+     * Recurse through the QueryBuilder tree, visiting any child QueryBuilder.
+     * @param visitor a query builder visitor to be called by each query builder in the tree.
+     */
+    default void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+    };
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java b/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java
@@ -0,0 +1,48 @@
+/*
+ * 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.apache.lucene.search.BooleanClause;
+
+/**
+ * QueryBuilderVisitor is an interface to define Visitor Object to be traversed in QueryBuilder tree.
+ */
+public interface QueryBuilderVisitor {
+
+    /**
+     * Accept method is called when the visitor accepts the queryBuilder object to be traversed in the query tree.
+     * @param qb is a queryBuilder object which is accepeted by the visitor.
+     */
+    void accept(QueryBuilder qb);
+
+    /**
+     * Fetches the child sub visitor from the main QueryBuilderVisitor Object.
+     * @param occur defines the occurrence of the result fetched from the search query in the final search result.
+     * @return a child queryBuilder Visitor Object.
+     */
+    QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur);
+
+    /**
+     *  NoopQueryVisitor is a default implementation of QueryBuilderVisitor.
+     *  When a user does not want to implement QueryBuilderVisitor and have to just pass an empty object then this class will be used.
+     *
+     */
+    QueryBuilderVisitor NO_OP_VISITOR = new QueryBuilderVisitor() {
+        @Override
+        public void accept(QueryBuilder qb) {
+            // Do nothing
+        }
+
+        @Override
+        public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+            return this;
+        }
+    };
+
+}
diff --git a/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanContainingQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -188,4 +189,11 @@ protected boolean doEquals(SpanContainingQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanFirstQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -186,4 +187,10 @@ protected boolean doEquals(SpanFirstQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(matchBuilder);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java
@@ -33,6 +33,7 @@
 
 import org.apache.lucene.queries.SpanMatchNoDocsQuery;
 import org.apache.lucene.queries.spans.SpanMultiTermQueryWrapper;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
@@ -213,4 +214,12 @@ protected boolean doEquals(SpanMultiTermQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (multiTermQueryBuilder != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(multiTermQueryBuilder);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanNearQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -299,6 +300,17 @@ public String getWriteableName() {
         return NAME;
     }
 
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (this.clauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.MUST);
+            for (QueryBuilder subQb : this.clauses) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
+
     /**
      * SpanGapQueryBuilder enables gaps in a SpanNearQuery.
      * Since, SpanGapQuery is private to SpanNearQuery, SpanGapQueryBuilder cannot

diff --git a/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanNotQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -284,4 +285,16 @@ protected boolean doEquals(SpanNotQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (include != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(include);
+        }
+
+        if (exclude != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST_NOT).accept(exclude);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanOrQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -188,4 +189,15 @@ protected boolean doEquals(SpanOrQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (clauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : this.clauses) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanQuery;
 import org.apache.lucene.queries.spans.SpanWithinQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -197,4 +198,11 @@ protected boolean doEquals(SpanWithinQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java
@@ -47,10 +47,13 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.opensearch.index.query.QueryBuilders.boolQuery;
 import static org.opensearch.index.query.QueryBuilders.termQuery;
@@ -456,4 +459,26 @@ public void testMustRewrite() throws IOException {
         IllegalStateException e = expectThrows(IllegalStateException.class, () -> boolQuery.toQuery(context));
         assertEquals("Rewrite first", e.getMessage());
     }
+
+    public void testVisit() {
+        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
+        boolQueryBuilder.should(new TermQueryBuilder(TEXT_FIELD_NAME, "should"));
+        boolQueryBuilder.must(new TermQueryBuilder(TEXT_FIELD_NAME, "must1"));
+        boolQueryBuilder.must(new TermQueryBuilder(TEXT_FIELD_NAME, "must2")); // Add a second one to confirm that they both get visited
+        boolQueryBuilder.mustNot(new TermQueryBuilder(TEXT_FIELD_NAME, "mustNot"));
+        boolQueryBuilder.filter(new TermQueryBuilder(TEXT_FIELD_NAME, "filter"));
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        boolQueryBuilder.visit(createTestVisitor(visitedQueries));
+        assertEquals(6, visitedQueries.size());
+        Set<Object> set = new HashSet<>(Arrays.asList("should", "must1", "must2", "mustNot", "filter"));
+
+        for (QueryBuilder qb : visitedQueries) {
+            if (qb instanceof TermQueryBuilder) {
+                set.remove(((TermQueryBuilder) qb).value());
+            }
+        }
+
+        assertEquals(0, set.size());
+
+    }
 }