Skip to content

Commit 087c5f0

Browse files
RyanL1997RyanL1997
committed
pushdown fix (similar to v2)
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
1 parent 2e901a3 commit 087c5f0

6 files changed

Lines changed: 141 additions & 19 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,6 @@ public class CalcitePlanContext {
4545
private static final ThreadLocal<Boolean> legacyPreferredFlag =
4646
ThreadLocal.withInitial(() -> true);
4747

48-
/** Highlight field patterns set by the {@code | highlight} command. */
49-
@Getter @Setter private List<String> highlightArgs;
50-
5148
@Getter @Setter private boolean isResolvingJoinCondition = false;
5249
@Getter @Setter private boolean isResolvingSubquery = false;
5350
@Getter @Setter private boolean inCoalesceFunction = false;

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -157,9 +157,9 @@
157157
import org.opensearch.sql.ast.tree.Values;
158158
import org.opensearch.sql.ast.tree.Window;
159159
import org.opensearch.sql.calcite.plan.AliasFieldsWrappable;
160-
import org.opensearch.sql.calcite.plan.HighlightPushable;
161160
import org.opensearch.sql.calcite.plan.OpenSearchConstants;
162161
import org.opensearch.sql.calcite.plan.rel.LogicalGraphLookup;
162+
import org.opensearch.sql.calcite.plan.rel.LogicalHighlight;
163163
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
164164
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType;
165165
import org.opensearch.sql.calcite.utils.BinUtils;
@@ -229,15 +229,6 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) {
229229
context.relBuilder.scan(node.getTableQualifiedName().getParts());
230230
RelNode scan = context.relBuilder.peek();
231231

232-
if (context.getHighlightArgs() != null && scan instanceof HighlightPushable highlightPushable) {
233-
RelNode newScan = highlightPushable.pushDownHighlight(context.getHighlightArgs());
234-
context.relBuilder.build(); // pop old scan
235-
context.relBuilder.push(newScan);
236-
scan = newScan;
237-
// Clear so that join's right-side scan is not affected
238-
context.setHighlightArgs(null);
239-
}
240-
241232
if (scan instanceof AliasFieldsWrappable) {
242233
return ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
243234
}
@@ -3201,9 +3192,11 @@ static ChartConfig fromArguments(ArgumentMap argMap) {
32013192

32023193
@Override
32033194
public RelNode visitHighlight(Highlight node, CalcitePlanContext context) {
3204-
context.setHighlightArgs(node.getHighlightArgs());
32053195
visitChildren(node, context);
3206-
return context.relBuilder.peek();
3196+
RelNode input = context.relBuilder.build();
3197+
LogicalHighlight highlight = LogicalHighlight.create(input, node.getHighlightArgs());
3198+
context.relBuilder.push(highlight);
3199+
return highlight;
32073200
}
32083201

32093202
@Override

‎core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalHighlight.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.plan.rel;
7+
8+
import java.util.List;
9+
import lombok.Getter;
10+
import org.apache.calcite.plan.Convention;
11+
import org.apache.calcite.plan.RelOptCluster;
12+
import org.apache.calcite.plan.RelTraitSet;
13+
import org.apache.calcite.rel.RelNode;
14+
import org.apache.calcite.rel.RelWriter;
15+
import org.apache.calcite.rel.SingleRel;
16+
import org.apache.calcite.rel.type.RelDataType;
17+
import org.apache.calcite.rel.type.RelDataTypeFactory;
18+
import org.apache.calcite.sql.type.SqlTypeName;
19+
import org.opensearch.sql.expression.HighlightExpression;
20+
21+
/**
22+
* Logical relational node representing a PPL {@code | highlight} command. Stores the highlight
23+
* arguments (terms to highlight or {@code *} for wildcard) and adds a {@code _highlight} column to
24+
* the output row type. An optimizer rule ({@code HighlightIndexScanRule}) pushes this node down
25+
* into the OpenSearch index scan.
26+
*/
27+
public class LogicalHighlight extends SingleRel {
28+
29+
@Getter private final List<String> highlightArgs;
30+
private final RelDataType highlightRowType;
31+
32+
protected LogicalHighlight(
33+
RelOptCluster cluster,
34+
RelTraitSet traitSet,
35+
RelNode input,
36+
List<String> highlightArgs,
37+
RelDataType highlightRowType) {
38+
super(cluster, traitSet, input);
39+
this.highlightArgs = highlightArgs;
40+
this.highlightRowType = highlightRowType;
41+
}
42+
43+
public static LogicalHighlight create(RelNode input, List<String> highlightArgs) {
44+
final RelOptCluster cluster = input.getCluster();
45+
RelTraitSet traitSet = cluster.traitSetOf(Convention.NONE);
46+
47+
// Add _highlight column to the output row type so that downstream operators
48+
// (e.g. visitProject) can reference it before the optimizer rule fires.
49+
RelDataTypeFactory typeFactory = cluster.getTypeFactory();
50+
RelDataTypeFactory.Builder schemaBuilder = typeFactory.builder();
51+
schemaBuilder.addAll(input.getRowType().getFieldList());
52+
if (!input.getRowType().getFieldNames().contains(HighlightExpression.HIGHLIGHT_FIELD)) {
53+
schemaBuilder.add(
54+
HighlightExpression.HIGHLIGHT_FIELD, typeFactory.createSqlType(SqlTypeName.ANY));
55+
}
56+
57+
return new LogicalHighlight(cluster, traitSet, input, highlightArgs, schemaBuilder.build());
58+
}
59+
60+
@Override
61+
protected RelDataType deriveRowType() {
62+
return highlightRowType;
63+
}
64+
65+
@Override
66+
public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
67+
assert traitSet.containsIfApplicable(Convention.NONE);
68+
return new LogicalHighlight(
69+
getCluster(), traitSet, sole(inputs), highlightArgs, highlightRowType);
70+
}
71+
72+
@Override
73+
public RelWriter explainTerms(RelWriter pw) {
74+
return super.explainTerms(pw).item("highlightArgs", highlightArgs);
75+
}
76+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/HighlightIndexScanRule.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.planner.rules;
7+
8+
import org.apache.calcite.plan.RelOptRuleCall;
9+
import org.apache.calcite.rel.RelNode;
10+
import org.immutables.value.Value;
11+
import org.opensearch.sql.calcite.plan.rel.LogicalHighlight;
12+
import org.opensearch.sql.calcite.plan.rule.OpenSearchRuleConfig;
13+
import org.opensearch.sql.calcite.utils.PlanUtils;
14+
import org.opensearch.sql.opensearch.storage.scan.CalciteLogicalIndexScan;
15+
16+
/**
17+
* Planner rule that pushes a {@link LogicalHighlight} down to {@link CalciteLogicalIndexScan}. This
18+
* adds the {@code _highlight} column to the scan schema and configures the OpenSearch {@code
19+
* HighlightBuilder} for the search request.
20+
*/
21+
@Value.Enclosing
22+
public class HighlightIndexScanRule extends InterruptibleRelRule<HighlightIndexScanRule.Config> {
23+
24+
protected HighlightIndexScanRule(Config config) {
25+
super(config);
26+
}
27+
28+
@Override
29+
protected void onMatchImpl(RelOptRuleCall call) {
30+
final LogicalHighlight highlight = call.rel(0);
31+
final CalciteLogicalIndexScan scan = call.rel(1);
32+
RelNode newScan = scan.pushDownHighlight(highlight.getHighlightArgs());
33+
if (newScan != null) {
34+
call.transformTo(newScan);
35+
PlanUtils.tryPruneRelNodes(call);
36+
}
37+
}
38+
39+
/** Rule configuration. */
40+
@Value.Immutable
41+
public interface Config extends OpenSearchRuleConfig {
42+
Config DEFAULT =
43+
ImmutableHighlightIndexScanRule.Config.builder()
44+
.build()
45+
.withOperandSupplier(
46+
b0 ->
47+
b0.operand(LogicalHighlight.class)
48+
.oneInput(b1 -> b1.operand(CalciteLogicalIndexScan.class).noInputs()));
49+
50+
@Override
51+
default HighlightIndexScanRule toRule() {
52+
return new HighlightIndexScanRule(this);
53+
}
54+
}
55+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/OpenSearchIndexRules.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ public class OpenSearchIndexRules {
6262
SortExprIndexScanRule.Config.DEFAULT.toRule();
6363
private static final EnumerableTopKMergeRule ENUMERABLE_TOP_K_MERGE_RULE =
6464
EnumerableTopKMergeRule.Config.DEFAULT.toRule();
65+
private static final HighlightIndexScanRule HIGHLIGHT_INDEX_SCAN =
66+
HighlightIndexScanRule.Config.DEFAULT.toRule();
6567

6668
/** The rules will apply only when the pushdown is enabled. */
6769
public static final List<RelOptRule> OPEN_SEARCH_PUSHDOWN_RULES =
@@ -80,7 +82,8 @@ public class OpenSearchIndexRules {
8082
RARE_TOP_PUSH_DOWN,
8183
ENUMERABLE_TOP_K_MERGE_RULE,
8284
EXPAND_COLLATION_ON_PROJECT_EXPR,
83-
SORT_EXPR_INDEX_SCAN);
85+
SORT_EXPR_INDEX_SCAN,
86+
HIGHLIGHT_INDEX_SCAN);
8487

8588
// prevent instantiation
8689
private OpenSearchIndexRules() {}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -508,7 +508,7 @@ protected static void applyHighlightPushDown(
508508
// Wildcard: highlight search query matches in all fields
509509
hb.field("*");
510510
} else {
511-
// Splunk-like: highlight specific terms across all fields
511+
// Highlight specific terms across all fields
512512
String queryStr =
513513
highlightArgs.stream()
514514
.map(term -> "\"" + term + "\"")
@@ -521,8 +521,6 @@ protected static void applyHighlightPushDown(
521521
hb.field(field);
522522
}
523523
hb.fragmentSize(Integer.MAX_VALUE);
524-
hb.preTags("@opensearch-dashboards-highlighted-field@");
525-
hb.postTags("@/opensearch-dashboards-highlighted-field@");
526524
requestBuilder.getSourceBuilder().highlighter(hb);
527525
}
528526
}

0 commit comments

Comments
 (0)