Skip to content

Commit 8833504

Browse files
qianheng-awsaaarone90
qianheng-aws
authored and
aaarone90
committed
Introduce logical dedup operators for PPL (opensearch-project#5014)
* Add new operator for dedup Signed-off-by: Heng Qian <qianheng@amazon.com> * Add customized field trimmer for dedup Signed-off-by: Heng Qian <qianheng@amazon.com> * Support convert dedup to basic sql operators Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix UT Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix IT Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix CI Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix IT Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix IT Signed-off-by: Heng Qian <qianheng@amazon.com> * Add java doc Signed-off-by: Heng Qian <qianheng@amazon.com> * Revert some code and add dedup simplify rule Signed-off-by: Heng Qian <qianheng@amazon.com> * Fix IT Signed-off-by: Heng Qian <qianheng@amazon.com> * Address comments Signed-off-by: Heng Qian <qianheng@amazon.com> * Address comments Signed-off-by: Heng Qian <qianheng@amazon.com> * Address comments Signed-off-by: Heng Qian <qianheng@amazon.com> * Address comments Signed-off-by: Heng Qian <qianheng@amazon.com> --------- Signed-off-by: Heng Qian <qianheng@amazon.com> Signed-off-by: Aaron Alvarez <aaarone@amazon.com>
1 parent 18a9075 commit 8833504

124 files changed

Lines changed: 911 additions & 414 deletions

File tree

Some content is hidden

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

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@
2121
import org.apache.calcite.rex.RexLambdaRef;
2222
import org.apache.calcite.rex.RexNode;
2323
import org.apache.calcite.tools.FrameworkConfig;
24-
import org.apache.calcite.tools.RelBuilder;
2524
import org.opensearch.sql.ast.expression.UnresolvedExpression;
2625
import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
26+
import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelBuilder;
2727
import org.opensearch.sql.common.setting.Settings;
2828
import org.opensearch.sql.executor.QueryType;
2929
import org.opensearch.sql.expression.function.FunctionProperties;
@@ -32,7 +32,7 @@ public class CalcitePlanContext {
3232

3333
public FrameworkConfig config;
3434
public final Connection connection;
35-
public final RelBuilder relBuilder;
35+
public final OpenSearchRelBuilder relBuilder;
3636
public final ExtendedRexBuilder rexBuilder;
3737
public final FunctionProperties functionProperties;
3838
public final QueryType queryType;

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

Lines changed: 8 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
import static org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC;
1515
import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
1616
import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
17-
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_DEDUP;
18-
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP;
17+
import static org.opensearch.sql.calcite.plan.rule.PPLDedupConvertRule.buildDedupNotNull;
18+
import static org.opensearch.sql.calcite.plan.rule.PPLDedupConvertRule.buildDedupOrNull;
1919
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_MAIN;
2020
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_RARE_TOP;
2121
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_STREAMSTATS;
@@ -146,9 +146,9 @@
146146
import org.opensearch.sql.ast.tree.Values;
147147
import org.opensearch.sql.ast.tree.Window;
148148
import org.opensearch.sql.calcite.plan.AliasFieldsWrappable;
149-
import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
150-
import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
151149
import org.opensearch.sql.calcite.plan.OpenSearchConstants;
150+
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
151+
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType;
152152
import org.opensearch.sql.calcite.utils.BinUtils;
153153
import org.opensearch.sql.calcite.utils.JoinAndLookupUtils;
154154
import org.opensearch.sql.calcite.utils.PPLHintUtils;
@@ -1387,7 +1387,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
13871387
: duplicatedFieldNames.stream()
13881388
.map(a -> (RexNode) context.relBuilder.field(a))
13891389
.toList();
1390-
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
1390+
buildDedupNotNull(context.relBuilder, dedupeFields, allowedDuplication);
13911391
}
13921392
// add LogicalSystemLimit after dedup
13931393
addSysLimitForJoinSubsearch(context);
@@ -1445,7 +1445,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
14451445
List<RexNode> dedupeFields =
14461446
getRightColumnsInJoinCriteria(context.relBuilder, joinCondition);
14471447

1448-
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
1448+
buildDedupNotNull(context.relBuilder, dedupeFields, allowedDuplication);
14491449
}
14501450
// add LogicalSystemLimit after dedup
14511451
addSysLimitForJoinSubsearch(context);
@@ -1622,81 +1622,13 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
16221622
List<RexNode> dedupeFields =
16231623
node.getFields().stream().map(f -> rexVisitor.analyze(f, context)).toList();
16241624
if (keepEmpty) {
1625-
buildDedupOrNull(context, dedupeFields, allowedDuplication);
1625+
buildDedupOrNull(context.relBuilder, dedupeFields, allowedDuplication);
16261626
} else {
1627-
buildDedupNotNull(context, dedupeFields, allowedDuplication, false);
1627+
buildDedupNotNull(context.relBuilder, dedupeFields, allowedDuplication);
16281628
}
16291629
return context.relBuilder.peek();
16301630
}
16311631

1632-
private static void buildDedupOrNull(
1633-
CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
1634-
/*
1635-
* | dedup 2 a, b keepempty=true
1636-
* LogicalProject(...)
1637-
* +- LogicalFilter(condition=[OR(IS NULL(a), IS NULL(b), <=(_row_number_dedup_, 1))])
1638-
* +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
1639-
* +- ...
1640-
*/
1641-
RexNode rowNumber =
1642-
context
1643-
.relBuilder
1644-
.aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
1645-
.over()
1646-
.partitionBy(dedupeFields)
1647-
.rowsTo(RexWindowBounds.CURRENT_ROW)
1648-
.as(ROW_NUMBER_COLUMN_FOR_DEDUP);
1649-
context.relBuilder.projectPlus(rowNumber);
1650-
RexNode _row_number_dedup_ = context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_DEDUP);
1651-
// Filter (isnull('a) OR isnull('b) OR '_row_number_dedup_ <= n)
1652-
context.relBuilder.filter(
1653-
context.relBuilder.or(
1654-
context.relBuilder.or(dedupeFields.stream().map(context.relBuilder::isNull).toList()),
1655-
context.relBuilder.lessThanOrEqual(
1656-
_row_number_dedup_, context.relBuilder.literal(allowedDuplication))));
1657-
// DropColumns('_row_number_dedup_)
1658-
context.relBuilder.projectExcept(_row_number_dedup_);
1659-
}
1660-
1661-
private static void buildDedupNotNull(
1662-
CalcitePlanContext context,
1663-
List<RexNode> dedupeFields,
1664-
Integer allowedDuplication,
1665-
boolean fromJoinMaxOption) {
1666-
/*
1667-
* | dedup 2 a, b keepempty=false
1668-
* LogicalProject(...)
1669-
* +- LogicalFilter(condition=[<=(_row_number_dedup_, n)]))
1670-
* +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
1671-
* +- LogicalFilter(condition=[AND(IS NOT NULL(a), IS NOT NULL(b))])
1672-
* +- ...
1673-
*/
1674-
// Filter (isnotnull('a) AND isnotnull('b))
1675-
String rowNumberAlias =
1676-
fromJoinMaxOption ? ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP : ROW_NUMBER_COLUMN_FOR_DEDUP;
1677-
context.relBuilder.filter(
1678-
context.relBuilder.and(dedupeFields.stream().map(context.relBuilder::isNotNull).toList()));
1679-
// Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
1680-
// specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a ASC
1681-
// NULLS FIRST, 'b ASC NULLS FIRST]
1682-
RexNode rowNumber =
1683-
context
1684-
.relBuilder
1685-
.aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
1686-
.over()
1687-
.partitionBy(dedupeFields)
1688-
.rowsTo(RexWindowBounds.CURRENT_ROW)
1689-
.as(rowNumberAlias);
1690-
context.relBuilder.projectPlus(rowNumber);
1691-
RexNode rowNumberField = context.relBuilder.field(rowNumberAlias);
1692-
// Filter ('_row_number_dedup_ <= n)
1693-
context.relBuilder.filter(
1694-
context.relBuilder.lessThanOrEqual(
1695-
rowNumberField, context.relBuilder.literal(allowedDuplication)));
1696-
// DropColumns('_row_number_dedup_)
1697-
context.relBuilder.projectExcept(rowNumberField);
1698-
}
1699-
17001632
@Override
17011633
public RelNode visitWindow(Window node, CalcitePlanContext context) {
17021634
visitChildren(node, context);

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,8 @@
7070
import org.opensearch.sql.ast.expression.subquery.ScalarSubquery;
7171
import org.opensearch.sql.ast.expression.subquery.SubqueryExpression;
7272
import org.opensearch.sql.ast.tree.UnresolvedPlan;
73-
import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
74-
import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
73+
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
74+
import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType;
7575
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
7676
import org.opensearch.sql.calcite.utils.PlanUtils;
7777
import org.opensearch.sql.calcite.utils.SubsearchUtils;

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

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
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.RelOptCluster;
11+
import org.apache.calcite.plan.RelOptPlanner;
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.rex.RexNode;
17+
import org.opensearch.sql.exception.CalciteUnsupportedException;
18+
19+
/** Relational expression representing a dedup command. */
20+
@Getter
21+
public abstract class Dedup extends SingleRel {
22+
final List<RexNode> dedupeFields;
23+
final Integer allowedDuplication;
24+
final Boolean keepEmpty;
25+
final Boolean consecutive;
26+
27+
/** */
28+
protected Dedup(
29+
RelOptCluster cluster,
30+
RelTraitSet traitSet,
31+
RelNode input,
32+
List<RexNode> dedupeFields,
33+
Integer allowedDuplication,
34+
Boolean keepEmpty,
35+
Boolean consecutive) {
36+
super(cluster, traitSet, input);
37+
if (allowedDuplication <= 0) {
38+
throw new IllegalArgumentException("Number of duplicate events must be greater than 0");
39+
}
40+
if (consecutive) {
41+
throw new CalciteUnsupportedException("Consecutive deduplication is unsupported in Calcite");
42+
}
43+
this.dedupeFields = dedupeFields;
44+
this.allowedDuplication = allowedDuplication;
45+
this.keepEmpty = keepEmpty;
46+
this.consecutive = consecutive;
47+
}
48+
49+
@Override
50+
public final RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
51+
return copy(
52+
traitSet,
53+
sole(inputs),
54+
this.dedupeFields,
55+
this.allowedDuplication,
56+
this.keepEmpty,
57+
this.consecutive);
58+
}
59+
60+
public abstract Dedup copy(
61+
RelTraitSet traitSet,
62+
RelNode input,
63+
List<RexNode> dedupeFields,
64+
Integer allowedDuplication,
65+
Boolean keepEmpty,
66+
Boolean consecutive);
67+
68+
public Dedup copy(RelNode input, List<RexNode> dedupeFields) {
69+
return this.copy(
70+
this.getTraitSet(),
71+
input,
72+
dedupeFields,
73+
this.allowedDuplication,
74+
this.keepEmpty,
75+
this.consecutive);
76+
}
77+
78+
@Override
79+
public RelWriter explainTerms(RelWriter pw) {
80+
return super.explainTerms(pw)
81+
.item("dedup_fields", dedupeFields)
82+
.item("allowed_dedup", allowedDuplication)
83+
.item("keepEmpty", keepEmpty)
84+
.item("consecutive", consecutive);
85+
}
86+
87+
@Override
88+
public void register(RelOptPlanner planner) {}
89+
}

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
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 static org.opensearch.sql.calcite.plan.rule.PPLDedupConvertRule.DEDUP_CONVERT_RULE;
9+
10+
import java.util.List;
11+
import org.apache.calcite.plan.Convention;
12+
import org.apache.calcite.plan.RelOptCluster;
13+
import org.apache.calcite.plan.RelOptPlanner;
14+
import org.apache.calcite.plan.RelTraitSet;
15+
import org.apache.calcite.rel.RelNode;
16+
import org.apache.calcite.rex.RexNode;
17+
18+
public class LogicalDedup extends Dedup {
19+
20+
protected LogicalDedup(
21+
RelOptCluster cluster,
22+
RelTraitSet traitSet,
23+
RelNode input,
24+
List<RexNode> dedupeFields,
25+
Integer allowedDuplication,
26+
Boolean keepEmpty,
27+
Boolean consecutive) {
28+
super(cluster, traitSet, input, dedupeFields, allowedDuplication, keepEmpty, consecutive);
29+
}
30+
31+
@Override
32+
public Dedup copy(
33+
RelTraitSet traitSet,
34+
RelNode input,
35+
List<RexNode> dedupeFields,
36+
Integer allowedDuplication,
37+
Boolean keepEmpty,
38+
Boolean consecutive) {
39+
assert traitSet.containsIfApplicable(Convention.NONE);
40+
return new LogicalDedup(
41+
getCluster(), traitSet, input, dedupeFields, allowedDuplication, keepEmpty, consecutive);
42+
}
43+
44+
public static LogicalDedup create(
45+
RelNode input,
46+
List<RexNode> dedupeFields,
47+
Integer allowedDuplication,
48+
Boolean keepEmpty,
49+
Boolean consecutive) {
50+
final RelOptCluster cluster = input.getCluster();
51+
RelTraitSet traitSet = cluster.traitSetOf(Convention.NONE);
52+
return new LogicalDedup(
53+
cluster, traitSet, input, dedupeFields, allowedDuplication, keepEmpty, consecutive);
54+
}
55+
56+
@Override
57+
public void register(RelOptPlanner planner) {
58+
planner.addRule(DEDUP_CONVERT_RULE);
59+
}
60+
}

core/src/main/java/org/opensearch/sql/calcite/plan/LogicalSystemLimit.java renamed to core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalSystemLimit.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.plan;
6+
package org.opensearch.sql.calcite.plan.rel;
77

88
import java.util.Collections;
99
import java.util.List;

core/src/main/java/org/opensearch/sql/calcite/plan/OpenSearchTableScan.java renamed to core/src/main/java/org/opensearch/sql/calcite/plan/rel/OpenSearchTableScan.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.plan;
6+
package org.opensearch.sql.calcite.plan.rel;
77

88
import com.google.common.collect.ImmutableList;
99
import org.apache.calcite.adapter.enumerable.EnumerableConvention;
@@ -14,6 +14,7 @@
1414
import org.apache.calcite.plan.RelOptTable;
1515
import org.apache.calcite.rel.core.TableScan;
1616
import org.apache.calcite.rel.rules.CoreRules;
17+
import org.opensearch.sql.calcite.plan.rule.OpenSearchRules;
1718

1819
/** Relational expression representing a scan of an OpenSearch type. */
1920
public abstract class OpenSearchTableScan extends TableScan implements EnumerableRel {

core/src/main/java/org/opensearch/sql/calcite/plan/OpenSearchRuleConfig.java renamed to core/src/main/java/org/opensearch/sql/calcite/plan/rule/OpenSearchRuleConfig.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.plan;
6+
package org.opensearch.sql.calcite.plan.rule;
77

88
import org.apache.calcite.plan.Contexts;
99
import org.apache.calcite.plan.RelRule;

core/src/main/java/org/opensearch/sql/calcite/plan/OpenSearchRules.java renamed to core/src/main/java/org/opensearch/sql/calcite/plan/rule/OpenSearchRules.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.plan;
6+
package org.opensearch.sql.calcite.plan.rule;
77

88
import com.google.common.collect.ImmutableList;
99
import java.util.List;

core/src/main/java/org/opensearch/sql/calcite/plan/PPLAggGroupMergeRule.java renamed to core/src/main/java/org/opensearch/sql/calcite/plan/rule/PPLAggGroupMergeRule.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.plan;
6+
package org.opensearch.sql.calcite.plan.rule;
77

88
import java.util.ArrayList;
99
import java.util.Collection;

0 commit comments

Comments
 (0)