Added javadocs, minor PR revisions, and fixed jacoco errors by improv…

…ing coverage for OpenSearchIndexScan Signed-off-by: forestmvey <forestv@bitquilltech.com>
Bit-Quill · Jul 29, 2022 · 3db8788 · 3db8788
1 parent cd6f911
commit 3db8788
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 4 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/analysis/HighlightAnalyzer.java b/core/src/main/java/org/opensearch/sql/analysis/HighlightAnalyzer.java
@@ -14,8 +14,12 @@
 import org.opensearch.sql.planner.logical.LogicalHighlight;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 
+/**
+ * Analyze the highlight in the {@link AnalysisContext} to construct the {@link
+ * LogicalPlan}.
+ */
 @RequiredArgsConstructor
-public class HighlightAnalyzer  extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
+public class HighlightAnalyzer extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
   private final ExpressionAnalyzer expressionAnalyzer;
   private final LogicalPlan child;
 

diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/HighlightFunction.java b/core/src/main/java/org/opensearch/sql/ast/expression/HighlightFunction.java
@@ -12,6 +12,9 @@
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 
+/**
+ * Expression node of Highlight function.
+ */
 @AllArgsConstructor
 @EqualsAndHashCode(callSuper = false)
 @Getter

diff --git a/core/src/main/java/org/opensearch/sql/expression/HighlightExpression.java b/core/src/main/java/org/opensearch/sql/expression/HighlightExpression.java
@@ -22,6 +22,9 @@
 import org.opensearch.sql.expression.env.Environment;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 
+/**
+ * Highlight Expression.
+ */
 @Getter
 public class HighlightExpression extends FunctionExpression {
   private final Expression highlightField;
@@ -47,14 +50,15 @@ public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
       refName += "." + StringUtils.unquoteText(getHighlightField().toString());
     }
     ExprValue retVal = valueEnv.resolve(DSL.ref(refName, ExprCoreType.STRING));
-    ImmutableMap.Builder<String, ExprValue> builder = new ImmutableMap.Builder<>();
 
+    // If only one highlight returned, or no highlights can be parsed.
     if (retVal.isMissing() || retVal.type() != ExprCoreType.STRUCT) {
       return retVal;
     }
 
-    var hlBuilder = ImmutableMap.<String, ExprValue>builder();
-    hlBuilder.putAll(retVal.tupleValue());
+    var highlightMapBuilder = ImmutableMap.<String, ExprValue>builder();
+    highlightMapBuilder.putAll(retVal.tupleValue());
+    ImmutableMap.Builder<String, ExprValue> builder = new ImmutableMap.Builder<>();
     for (var entry : retVal.tupleValue().entrySet()) {
       String entryKey = "highlight(" + getHighlightField() + ")." + entry.getKey();
       builder.put(entryKey, ExprValueUtils.stringValue(entry.getValue().toString()));

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScanTest.java
@@ -31,6 +31,7 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.search.SearchHit;
+import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -110,6 +111,16 @@ void pushDownFilters() {
                 .filter(QueryBuilders.rangeQuery("balance").gte(10000)));
   }
 
+  @Test
+  void pushDownHighlight() {
+    assertThat()
+        .pushDown(QueryBuilders.termQuery("name", "John"))
+        .pushDownHighlight("Title")
+        .pushDownHighlight("Body")
+        .shouldQueryHighlight(QueryBuilders.termQuery("name", "John"),
+            new HighlightBuilder().field("Title").field("Body"));
+  }
+
   private PushDownAssertion assertThat() {
     return new PushDownAssertion(client, exprValueFactory, settings);
   }
@@ -135,6 +146,22 @@ PushDownAssertion pushDown(QueryBuilder query) {
       return this;
     }
 
+    PushDownAssertion pushDownHighlight(String query) {
+      indexScan.pushDownHighlight(query);
+      return this;
+    }
+
+    PushDownAssertion shouldQueryHighlight(QueryBuilder query, HighlightBuilder highlight) {
+      OpenSearchRequest request = new OpenSearchQueryRequest("test", 200, factory);
+      request.getSourceBuilder()
+          .query(query)
+          .highlighter(highlight)
+          .sort(DOC_FIELD_NAME, ASC);
+      when(client.search(request)).thenReturn(response);
+      indexScan.open();
+      return this;
+    }
+
     PushDownAssertion shouldQuery(QueryBuilder expected) {
       OpenSearchRequest request = new OpenSearchQueryRequest("test", 200, factory);
       request.getSourceBuilder()