Skip to content

Commit 64a3794

Browse files
dai-chendai-chen
authored
Improve pushdown optimization and logical to physical transformation (#1091)
* Add new table scan builder and optimizer rules Signed-off-by: Chen Dai <[email protected]> * Fix jacoco test coverage Signed-off-by: Chen Dai <[email protected]> * Update javadoc with more details Signed-off-by: Chen Dai <[email protected]> * Fix highlight pushdown issue Signed-off-by: Chen Dai <[email protected]> * Rename new class more properly Signed-off-by: Chen Dai <[email protected]> * Fix default sort by doc issue Signed-off-by: Chen Dai <[email protected]> * Rename visit method and javadoc Signed-off-by: Chen Dai <[email protected]> * Move table scan builder and optimize rule to read package Signed-off-by: Chen Dai <[email protected]> * Fix sort push down issue Signed-off-by: Chen Dai <[email protected]> * Move sortByFields to parent scan builder Signed-off-by: Chen Dai <[email protected]> * Add back old test Signed-off-by: Chen Dai <[email protected]> Signed-off-by: Chen Dai <[email protected]>
1 parent f530770 commit 64a3794

File tree

44 files changed

+1810
-2050
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

44 files changed

+1810
-2050
lines changed

‎core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.opensearch.sql.planner.physical.SortOperator;
3535
import org.opensearch.sql.planner.physical.ValuesOperator;
3636
import org.opensearch.sql.planner.physical.WindowOperator;
37+
import org.opensearch.sql.storage.read.TableScanBuilder;
3738

3839
/**
3940
* Default implementor for implementing logical to physical translation. "Default" here means all
@@ -123,6 +124,11 @@ public PhysicalPlan visitLimit(LogicalLimit node, C context) {
123124
return new LimitOperator(visitChild(node, context), node.getLimit(), node.getOffset());
124125
}
125126

127+
@Override
128+
public PhysicalPlan visitTableScanBuilder(TableScanBuilder plan, C context) {
129+
return plan.build();
130+
}
131+
126132
@Override
127133
public PhysicalPlan visitRelation(LogicalRelation node, C context) {
128134
throw new UnsupportedOperationException("Storage engine is responsible for "

‎core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
package org.opensearch.sql.planner.logical;
88

9+
import org.opensearch.sql.storage.read.TableScanBuilder;
10+
911
/**
1012
* The visitor of {@link LogicalPlan}.
1113
*
@@ -22,6 +24,10 @@ public R visitRelation(LogicalRelation plan, C context) {
2224
return visitNode(plan, context);
2325
}
2426

27+
public R visitTableScanBuilder(TableScanBuilder plan, C context) {
28+
return visitNode(plan, context);
29+
}
30+
2531
public R visitFilter(LogicalFilter plan, C context) {
2632
return visitNode(plan, context);
2733
}

‎core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
import org.opensearch.sql.planner.logical.LogicalPlan;
1616
import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter;
1717
import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
18+
import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
19+
import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown;
1820

1921
/**
2022
* {@link LogicalPlan} Optimizer.
@@ -39,8 +41,21 @@ public LogicalPlanOptimizer(List<Rule<?>> rules) {
3941
*/
4042
public static LogicalPlanOptimizer create() {
4143
return new LogicalPlanOptimizer(Arrays.asList(
44+
/*
45+
* Phase 1: Transformations that rely on relational algebra equivalence
46+
*/
4247
new MergeFilterAndFilter(),
43-
new PushFilterUnderSort()));
48+
new PushFilterUnderSort(),
49+
/*
50+
* Phase 2: Transformations that rely on data source push down capability
51+
*/
52+
new CreateTableScanBuilder(),
53+
TableScanPushDown.PUSH_DOWN_FILTER,
54+
TableScanPushDown.PUSH_DOWN_AGGREGATION,
55+
TableScanPushDown.PUSH_DOWN_SORT,
56+
TableScanPushDown.PUSH_DOWN_LIMIT,
57+
TableScanPushDown.PUSH_DOWN_HIGHLIGHT,
58+
TableScanPushDown.PUSH_DOWN_PROJECT));
4459
}
4560

4661
/**
@@ -63,7 +78,14 @@ private LogicalPlan internalOptimize(LogicalPlan plan) {
6378
Match match = DEFAULT_MATCHER.match(rule.pattern(), node);
6479
if (match.isPresent()) {
6580
node = rule.apply(match.value(), match.captures());
66-
done = false;
81+
82+
// For new TableScanPushDown impl, pattern match doesn't necessarily cause
83+
// push down to happen. So reiterate all rules against the node only if the node
84+
// is actually replaced by any rule.
85+
// TODO: may need to introduce fixed point or maximum iteration limit in future
86+
if (node != match.value()) {
87+
done = false;
88+
}
6789
}
6890
}
6991
}

‎core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,78 @@
66

77
package org.opensearch.sql.planner.optimizer.pattern;
88

9+
import com.facebook.presto.matching.Capture;
10+
import com.facebook.presto.matching.Pattern;
911
import com.facebook.presto.matching.Property;
12+
import com.facebook.presto.matching.PropertyPattern;
1013
import java.util.Optional;
1114
import lombok.experimental.UtilityClass;
15+
import org.opensearch.sql.planner.logical.LogicalAggregation;
16+
import org.opensearch.sql.planner.logical.LogicalFilter;
17+
import org.opensearch.sql.planner.logical.LogicalHighlight;
18+
import org.opensearch.sql.planner.logical.LogicalLimit;
1219
import org.opensearch.sql.planner.logical.LogicalPlan;
20+
import org.opensearch.sql.planner.logical.LogicalProject;
21+
import org.opensearch.sql.planner.logical.LogicalRelation;
22+
import org.opensearch.sql.planner.logical.LogicalSort;
23+
import org.opensearch.sql.storage.Table;
24+
import org.opensearch.sql.storage.read.TableScanBuilder;
1325

1426
/**
1527
* Pattern helper class.
1628
*/
1729
@UtilityClass
1830
public class Patterns {
1931

32+
/**
33+
* Logical filter with a given pattern on inner field.
34+
*/
35+
public static <T extends LogicalPlan> Pattern<LogicalFilter> filter(Pattern<T> pattern) {
36+
return Pattern.typeOf(LogicalFilter.class).with(source(pattern));
37+
}
38+
39+
/**
40+
* Logical aggregate operator with a given pattern on inner field.
41+
*/
42+
public static <T extends LogicalPlan> Pattern<LogicalAggregation> aggregate(Pattern<T> pattern) {
43+
return Pattern.typeOf(LogicalAggregation.class).with(source(pattern));
44+
}
45+
46+
/**
47+
* Logical sort operator with a given pattern on inner field.
48+
*/
49+
public static <T extends LogicalPlan> Pattern<LogicalSort> sort(Pattern<T> pattern) {
50+
return Pattern.typeOf(LogicalSort.class).with(source(pattern));
51+
}
52+
53+
/**
54+
* Logical limit operator with a given pattern on inner field.
55+
*/
56+
public static <T extends LogicalPlan> Pattern<LogicalLimit> limit(Pattern<T> pattern) {
57+
return Pattern.typeOf(LogicalLimit.class).with(source(pattern));
58+
}
59+
60+
/**
61+
* Logical highlight operator with a given pattern on inner field.
62+
*/
63+
public static <T extends LogicalPlan> Pattern<LogicalHighlight> highlight(Pattern<T> pattern) {
64+
return Pattern.typeOf(LogicalHighlight.class).with(source(pattern));
65+
}
66+
67+
/**
68+
* Logical project operator with a given pattern on inner field.
69+
*/
70+
public static <T extends LogicalPlan> Pattern<LogicalProject> project(Pattern<T> pattern) {
71+
return Pattern.typeOf(LogicalProject.class).with(source(pattern));
72+
}
73+
74+
/**
75+
* Pattern for {@link TableScanBuilder} and capture it meanwhile.
76+
*/
77+
public static Pattern<TableScanBuilder> scanBuilder() {
78+
return Pattern.typeOf(TableScanBuilder.class).capturedAs(Capture.newCapture());
79+
}
80+
2081
/**
2182
* LogicalPlan source {@link Property}.
2283
*/
@@ -25,4 +86,28 @@ public static Property<LogicalPlan, LogicalPlan> source() {
2586
? Optional.of(plan.getChild().get(0))
2687
: Optional.empty());
2788
}
89+
90+
/**
91+
* Source (children field) with a given pattern.
92+
*/
93+
@SuppressWarnings("unchecked")
94+
public static <T extends LogicalPlan>
95+
PropertyPattern<LogicalPlan, T> source(Pattern<T> pattern) {
96+
Property<LogicalPlan, T> property = Property.optionalProperty("source",
97+
plan -> plan.getChild().size() == 1
98+
? Optional.of((T) plan.getChild().get(0))
99+
: Optional.empty());
100+
101+
return property.matching(pattern);
102+
}
103+
104+
/**
105+
* Logical relation with table field.
106+
*/
107+
public static Property<LogicalPlan, Table> table() {
108+
return Property.optionalProperty("table",
109+
plan -> plan instanceof LogicalRelation
110+
? Optional.of(((LogicalRelation) plan).getTable())
111+
: Optional.empty());
112+
}
28113
}

‎core/src/main/java/org/opensearch/sql/planner/optimizer/rule/read/CreateTableScanBuilder.java

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.planner.optimizer.rule.read;
7+
8+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.table;
9+
10+
import com.facebook.presto.matching.Capture;
11+
import com.facebook.presto.matching.Captures;
12+
import com.facebook.presto.matching.Pattern;
13+
import lombok.Getter;
14+
import lombok.experimental.Accessors;
15+
import org.opensearch.sql.planner.logical.LogicalPlan;
16+
import org.opensearch.sql.planner.logical.LogicalRelation;
17+
import org.opensearch.sql.planner.optimizer.Rule;
18+
import org.opensearch.sql.storage.Table;
19+
import org.opensearch.sql.storage.read.TableScanBuilder;
20+
21+
/**
22+
* Rule that replace logical relation operator to {@link TableScanBuilder} for later
23+
* push down optimization. All push down optimization rules that depends on table scan
24+
* builder needs to run after this.
25+
*/
26+
public class CreateTableScanBuilder implements Rule<LogicalRelation> {
27+
28+
/** Capture the table inside matched logical relation operator. */
29+
private final Capture<Table> capture;
30+
31+
/** Pattern that matches logical relation operator. */
32+
@Accessors(fluent = true)
33+
@Getter
34+
private final Pattern<LogicalRelation> pattern;
35+
36+
/**
37+
* Construct create table scan builder rule.
38+
*/
39+
public CreateTableScanBuilder() {
40+
this.capture = Capture.newCapture();
41+
this.pattern = Pattern.typeOf(LogicalRelation.class)
42+
.with(table().capturedAs(capture));
43+
}
44+
45+
@Override
46+
public LogicalPlan apply(LogicalRelation plan, Captures captures) {
47+
TableScanBuilder scanBuilder = captures.get(capture).createScanBuilder();
48+
// TODO: Remove this after Prometheus refactored to new table scan builder too
49+
return (scanBuilder == null) ? plan : scanBuilder;
50+
}
51+
}

‎core/src/main/java/org/opensearch/sql/planner/optimizer/rule/read/TableScanPushDown.java

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.planner.optimizer.rule.read;
7+
8+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.aggregate;
9+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.filter;
10+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.highlight;
11+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.limit;
12+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.project;
13+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.scanBuilder;
14+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.sort;
15+
import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.TableScanPushDownBuilder.match;
16+
17+
import com.facebook.presto.matching.Capture;
18+
import com.facebook.presto.matching.Captures;
19+
import com.facebook.presto.matching.Pattern;
20+
import com.facebook.presto.matching.pattern.CapturePattern;
21+
import com.facebook.presto.matching.pattern.WithPattern;
22+
import java.util.function.BiFunction;
23+
import org.opensearch.sql.planner.logical.LogicalPlan;
24+
import org.opensearch.sql.planner.optimizer.Rule;
25+
import org.opensearch.sql.storage.read.TableScanBuilder;
26+
27+
/**
28+
* Rule template for all table scan push down rules. Because all push down optimization rules
29+
* have similar workflow in common, such as a pattern that match an operator on top of table scan
30+
* builder, and action that eliminates the original operator if pushed down, this class helps
31+
* remove redundant code and improve readability.
32+
*
33+
* @param <T> logical plan node type
34+
*/
35+
public class TableScanPushDown<T extends LogicalPlan> implements Rule<T> {
36+
37+
/** Push down optimize rule for filtering condition. */
38+
public static final Rule<?> PUSH_DOWN_FILTER =
39+
match(
40+
filter(
41+
scanBuilder()))
42+
.apply((filter, scanBuilder) -> scanBuilder.pushDownFilter(filter));
43+
44+
/** Push down optimize rule for aggregate operator. */
45+
public static final Rule<?> PUSH_DOWN_AGGREGATION =
46+
match(
47+
aggregate(
48+
scanBuilder()))
49+
.apply((agg, scanBuilder) -> scanBuilder.pushDownAggregation(agg));
50+
51+
/** Push down optimize rule for sort operator. */
52+
public static final Rule<?> PUSH_DOWN_SORT =
53+
match(
54+
sort(
55+
scanBuilder()))
56+
.apply((sort, scanBuilder) -> scanBuilder.pushDownSort(sort));
57+
58+
/** Push down optimize rule for limit operator. */
59+
public static final Rule<?> PUSH_DOWN_LIMIT =
60+
match(
61+
limit(
62+
scanBuilder()))
63+
.apply((limit, scanBuilder) -> scanBuilder.pushDownLimit(limit));
64+
65+
public static final Rule<?> PUSH_DOWN_PROJECT =
66+
match(
67+
project(
68+
scanBuilder()))
69+
.apply((project, scanBuilder) -> scanBuilder.pushDownProject(project));
70+
71+
public static final Rule<?> PUSH_DOWN_HIGHLIGHT =
72+
match(
73+
highlight(
74+
scanBuilder()))
75+
.apply((highlight, scanBuilder) -> scanBuilder.pushDownHighlight(highlight));
76+
77+
78+
/** Pattern that matches a plan node. */
79+
private final WithPattern<T> pattern;
80+
81+
/** Capture table scan builder inside a plan node. */
82+
private final Capture<TableScanBuilder> capture;
83+
84+
/** Push down function applied to the plan node and captured table scan builder. */
85+
private final BiFunction<T, TableScanBuilder, Boolean> pushDownFunction;
86+
87+
88+
@SuppressWarnings("unchecked")
89+
private TableScanPushDown(WithPattern<T> pattern,
90+
BiFunction<T, TableScanBuilder, Boolean> pushDownFunction) {
91+
this.pattern = pattern;
92+
this.capture = ((CapturePattern<TableScanBuilder>) pattern.getPattern()).capture();
93+
this.pushDownFunction = pushDownFunction;
94+
}
95+
96+
@Override
97+
public Pattern<T> pattern() {
98+
return pattern;
99+
}
100+
101+
@Override
102+
public LogicalPlan apply(T plan, Captures captures) {
103+
TableScanBuilder scanBuilder = captures.get(capture);
104+
if (pushDownFunction.apply(plan, scanBuilder)) {
105+
return scanBuilder;
106+
}
107+
return plan;
108+
}
109+
110+
/**
111+
* Custom builder class other than generated by Lombok to provide more readable code.
112+
*/
113+
static class TableScanPushDownBuilder<T extends LogicalPlan> {
114+
115+
private WithPattern<T> pattern;
116+
117+
public static <T extends LogicalPlan>
118+
TableScanPushDownBuilder<T> match(Pattern<T> pattern) {
119+
TableScanPushDownBuilder<T> builder = new TableScanPushDownBuilder<>();
120+
builder.pattern = (WithPattern<T>) pattern;
121+
return builder;
122+
}
123+
124+
public TableScanPushDown<T> apply(
125+
BiFunction<T, TableScanBuilder, Boolean> pushDownFunction) {
126+
return new TableScanPushDown<>(pattern, pushDownFunction);
127+
}
128+
}
129+
}

0 commit comments

Comments
 (0)