Skip to content

Commit 1b2ed31

Browse files
LantaoJinLantaoJin
committed
Convert dedup pushdown to composite + top_hits (opensearch-project#4844)
* Enable dedup pushdown Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix doctest Signed-off-by: Lantao Jin <ltjin@amazon.com> * refactor Signed-off-by: Lantao Jin <ltjin@amazon.com> * Disable dedup expr Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix IT Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix yaml test Signed-off-by: Lantao Jin <ltjin@amazon.com> * add more comments in code Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix conflicts Signed-off-by: Lantao Jin <ltjin@amazon.com> * Address comments Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com> (cherry picked from commit 5ceacb6)
1 parent f04f606 commit 1b2ed31

74 files changed

Lines changed: 1224 additions & 493 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/CalciteRelNodeVisitor.java‎

Lines changed: 24 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
1616
import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
1717
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;
1819
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_MAIN;
1920
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_RARE_TOP;
2021
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_STREAMSTATS;
@@ -48,9 +49,6 @@
4849
import org.apache.calcite.rel.RelNode;
4950
import org.apache.calcite.rel.core.Aggregate;
5051
import org.apache.calcite.rel.core.JoinRelType;
51-
import org.apache.calcite.rel.hint.HintStrategyTable;
52-
import org.apache.calcite.rel.hint.RelHint;
53-
import org.apache.calcite.rel.logical.LogicalAggregate;
5452
import org.apache.calcite.rel.logical.LogicalValues;
5553
import org.apache.calcite.rel.type.RelDataType;
5654
import org.apache.calcite.rel.type.RelDataTypeFamily;
@@ -1055,7 +1053,7 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
10551053
List<String> intendedGroupKeyAliases = getGroupKeyNamesAfterAggregation(reResolved.getLeft());
10561054
context.relBuilder.aggregate(
10571055
context.relBuilder.groupKey(reResolved.getLeft()), reResolved.getRight());
1058-
if (hintBucketNonNull) addIgnoreNullBucketHintToAggregate(context);
1056+
if (hintBucketNonNull) PlanUtils.addIgnoreNullBucketHintToAggregate(context.relBuilder);
10591057
// During aggregation, Calcite projects both input dependencies and output group-by fields.
10601058
// When names conflict, Calcite adds numeric suffixes (e.g., "value0").
10611059
// Apply explicit renaming to restore the intended aliases.
@@ -1323,7 +1321,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
13231321
: duplicatedFieldNames.stream()
13241322
.map(a -> (RexNode) context.relBuilder.field(a))
13251323
.collect(Collectors.toList());
1326-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1324+
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
13271325
}
13281326
context.relBuilder.join(
13291327
JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1379,7 +1377,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
13791377
List<RexNode> dedupeFields =
13801378
getRightColumnsInJoinCriteria(context.relBuilder, joinCondition);
13811379

1382-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1380+
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
13831381
}
13841382
context.relBuilder.join(
13851383
JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1544,24 +1542,20 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
15441542
if (keepEmpty) {
15451543
buildDedupOrNull(context, dedupeFields, allowedDuplication);
15461544
} else {
1547-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1545+
buildDedupNotNull(context, dedupeFields, allowedDuplication, false);
15481546
}
15491547
return context.relBuilder.peek();
15501548
}
15511549

15521550
private static void buildDedupOrNull(
15531551
CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
15541552
/*
1555-
* | dedup 2 a, b keepempty=false
1556-
* DropColumns('_row_number_dedup_)
1557-
* +- Filter ('_row_number_dedup_ <= n OR isnull('a) OR isnull('b))
1558-
* +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
1553+
* | dedup 2 a, b keepempty=true
1554+
* LogicalProject(...)
1555+
* +- LogicalFilter(condition=[OR(IS NULL(a), IS NULL(b), <=(_row_number_dedup_, 1))])
1556+
* +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
15591557
* +- ...
15601558
*/
1561-
// Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
1562-
// specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a
1563-
// ASC
1564-
// NULLS FIRST, 'b ASC NULLS FIRST]
15651559
RexNode rowNumber =
15661560
context
15671561
.relBuilder
@@ -1584,16 +1578,21 @@ private static void buildDedupOrNull(
15841578
}
15851579

15861580
private static void buildDedupNotNull(
1587-
CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
1581+
CalcitePlanContext context,
1582+
List<RexNode> dedupeFields,
1583+
Integer allowedDuplication,
1584+
boolean fromJoinMaxOption) {
15881585
/*
15891586
* | dedup 2 a, b keepempty=false
1590-
* DropColumns('_row_number_dedup_)
1591-
* +- Filter ('_row_number_dedup_ <= n)
1592-
* +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
1593-
* +- Filter (isnotnull('a) AND isnotnull('b))
1594-
* +- ...
1587+
* LogicalProject(...)
1588+
* +- LogicalFilter(condition=[<=(_row_number_dedup_, n)]))
1589+
* +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
1590+
* +- LogicalFilter(condition=[AND(IS NOT NULL(a), IS NOT NULL(b))])
1591+
* +- ...
15951592
*/
15961593
// Filter (isnotnull('a) AND isnotnull('b))
1594+
String rowNumberAlias =
1595+
fromJoinMaxOption ? ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP : ROW_NUMBER_COLUMN_FOR_DEDUP;
15971596
context.relBuilder.filter(
15981597
context.relBuilder.and(dedupeFields.stream().map(context.relBuilder::isNotNull).collect(Collectors.toList())));
15991598
// Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
@@ -1607,15 +1606,15 @@ private static void buildDedupNotNull(
16071606
.partitionBy(dedupeFields)
16081607
.orderBy(dedupeFields)
16091608
.rowsTo(RexWindowBounds.CURRENT_ROW)
1610-
.as(ROW_NUMBER_COLUMN_FOR_DEDUP);
1609+
.as(rowNumberAlias);
16111610
context.relBuilder.projectPlus(rowNumber);
1612-
RexNode _row_number_dedup_ = context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_DEDUP);
1611+
RexNode rowNumberField = context.relBuilder.field(rowNumberAlias);
16131612
// Filter ('_row_number_dedup_ <= n)
16141613
context.relBuilder.filter(
16151614
context.relBuilder.lessThanOrEqual(
1616-
_row_number_dedup_, context.relBuilder.literal(allowedDuplication)));
1615+
rowNumberField, context.relBuilder.literal(allowedDuplication)));
16171616
// DropColumns('_row_number_dedup_)
1618-
context.relBuilder.projectExcept(_row_number_dedup_);
1617+
context.relBuilder.projectExcept(rowNumberField);
16191618
}
16201619

16211620
@Override
@@ -2399,25 +2398,6 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
23992398
return context.relBuilder.peek();
24002399
}
24012400

2402-
private static void addIgnoreNullBucketHintToAggregate(CalcitePlanContext context) {
2403-
final RelHint statHits =
2404-
RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
2405-
assert context.relBuilder.peek() instanceof LogicalAggregate
2406-
: "Stats hits should be added to LogicalAggregate";
2407-
context.relBuilder.hints(statHits);
2408-
context
2409-
.relBuilder
2410-
.getCluster()
2411-
.setHintStrategies(
2412-
HintStrategyTable.builder()
2413-
.hintStrategy(
2414-
"stats_args",
2415-
(hint, rel) -> {
2416-
return rel instanceof LogicalAggregate;
2417-
})
2418-
.build());
2419-
}
2420-
24212401
@Override
24222402
public RelNode visitTableFunction(TableFunction node, CalcitePlanContext context) {
24232403
throw new CalciteUnsupportedException("Table function is unsupported in Calcite");

‎core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java‎

Lines changed: 72 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,21 @@
2525
import org.apache.calcite.rel.RelHomogeneousShuttle;
2626
import org.apache.calcite.rel.RelNode;
2727
import org.apache.calcite.rel.RelShuttle;
28+
import org.apache.calcite.rel.core.AggregateCall;
2829
import org.apache.calcite.rel.core.Project;
2930
import org.apache.calcite.rel.core.Sort;
3031
import org.apache.calcite.rel.core.TableScan;
32+
import org.apache.calcite.rel.hint.HintStrategyTable;
33+
import org.apache.calcite.rel.hint.RelHint;
34+
import org.apache.calcite.rel.logical.LogicalAggregate;
35+
import org.apache.calcite.rel.logical.LogicalFilter;
3136
import org.apache.calcite.rel.logical.LogicalProject;
3237
import org.apache.calcite.rel.logical.LogicalSort;
3338
import org.apache.calcite.rel.type.RelDataType;
3439
import org.apache.calcite.rex.RexCall;
3540
import org.apache.calcite.rex.RexCorrelVariable;
3641
import org.apache.calcite.rex.RexInputRef;
42+
import org.apache.calcite.rex.RexLiteral;
3743
import org.apache.calcite.rex.RexNode;
3844
import org.apache.calcite.rex.RexOver;
3945
import org.apache.calcite.rex.RexVisitorImpl;
@@ -45,8 +51,11 @@
4551
import org.apache.calcite.tools.RelBuilder;
4652
import org.apache.calcite.util.Pair;
4753
import org.apache.calcite.util.Util;
54+
import org.apache.calcite.util.mapping.Mapping;
55+
import org.apache.calcite.util.mapping.Mappings;
4856
import org.opensearch.sql.ast.AbstractNodeVisitor;
4957
import org.opensearch.sql.ast.Node;
58+
import org.opensearch.sql.ast.expression.Argument;
5059
import org.opensearch.sql.ast.expression.IntervalUnit;
5160
import org.opensearch.sql.ast.expression.SpanUnit;
5261
import org.opensearch.sql.ast.expression.WindowBound;
@@ -62,6 +71,7 @@ public interface PlanUtils {
6271
/** this is only for dedup command, do not reuse it in other command */
6372
String ROW_NUMBER_COLUMN_FOR_DEDUP = "_row_number_dedup_";
6473

74+
String ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP = "_row_number_join_max_dedup_";
6575
String ROW_NUMBER_COLUMN_FOR_RARE_TOP = "_row_number_rare_top_";
6676
String ROW_NUMBER_COLUMN_FOR_MAIN = "_row_number_main_";
6777
String ROW_NUMBER_COLUMN_FOR_SUBSEARCH = "_row_number_subsearch_";
@@ -477,18 +487,15 @@ static RexNode derefMapCall(RexNode rexNode) {
477487
return rexNode;
478488
}
479489

480-
/** Check if contains RexOver introduced by dedup */
481-
static boolean containsRowNumberDedup(LogicalProject project) {
482-
return project.getProjects().stream()
483-
.anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER)
484-
&& project.getRowType().getFieldNames().contains(ROW_NUMBER_COLUMN_FOR_DEDUP);
490+
/** Check if contains dedup */
491+
static boolean containsRowNumberDedup(RelNode node) {
492+
return node.getRowType().getFieldNames().stream().anyMatch(ROW_NUMBER_COLUMN_FOR_DEDUP::equals);
485493
}
486494

487-
/** Check if contains RexOver introduced by dedup top/rare */
488-
static boolean containsRowNumberRareTop(LogicalProject project) {
489-
return project.getProjects().stream()
490-
.anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER)
491-
&& project.getRowType().getFieldNames().contains(ROW_NUMBER_COLUMN_FOR_RARE_TOP);
495+
/** Check if contains dedup for top/rare */
496+
static boolean containsRowNumberRareTop(RelNode node) {
497+
return node.getRowType().getFieldNames().stream()
498+
.anyMatch(ROW_NUMBER_COLUMN_FOR_RARE_TOP::equals);
492499
}
493500

494501
/** Get all RexWindow list from LogicalProject */
@@ -536,10 +543,6 @@ static boolean distinctProjectList(LogicalProject project) {
536543
return project.getNamedProjects().stream().allMatch(rexSet::add);
537544
}
538545

539-
static boolean containsRexOver(LogicalProject project) {
540-
return project.getProjects().stream().anyMatch(RexOver::containsOver);
541-
}
542-
543546
/**
544547
* The LogicalSort is a LIMIT that should be pushed down when its fetch field is not null and its
545548
* collation is empty. For example: <code>sort name | head 5</code> should not be pushed down
@@ -552,7 +555,7 @@ static boolean isLogicalSortLimit(LogicalSort sort) {
552555
return sort.fetch != null;
553556
}
554557

555-
static boolean projectContainsExpr(Project project) {
558+
static boolean containsRexCall(Project project) {
556559
return project.getProjects().stream().anyMatch(p -> p instanceof RexCall);
557560
}
558561

@@ -623,4 +626,58 @@ static void replaceTop(RelBuilder relBuilder, RelNode relNode) {
623626
throw new IllegalStateException("Unable to invoke RelBuilder.replaceTop", e);
624627
}
625628
}
629+
630+
static void addIgnoreNullBucketHintToAggregate(RelBuilder relBuilder) {
631+
final RelHint statHits =
632+
RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
633+
assert relBuilder.peek() instanceof LogicalAggregate
634+
: "Stats hits should be added to LogicalAggregate";
635+
relBuilder.hints(statHits);
636+
relBuilder
637+
.getCluster()
638+
.setHintStrategies(
639+
HintStrategyTable.builder()
640+
.hintStrategy(
641+
"stats_args",
642+
(hint, rel) -> {
643+
return rel instanceof LogicalAggregate;
644+
})
645+
.build());
646+
}
647+
648+
/** Extract the RexLiteral from the aggregate call if the aggregate call is a LITERAL_AGG. */
649+
static @Nullable RexLiteral getObjectFromLiteralAgg(AggregateCall aggCall) {
650+
if (aggCall.getAggregation().kind == SqlKind.LITERAL_AGG) {
651+
return (RexLiteral)
652+
aggCall.rexList.stream().filter(rex -> rex instanceof RexLiteral).findAny().orElse(null);
653+
} else {
654+
return null;
655+
}
656+
}
657+
658+
/**
659+
* This is a helper method to create a target mapping easily for replacing calling {@link
660+
* Mappings#target(List, int)}
661+
*
662+
* @param rexNodes the rex list in schema
663+
* @param schema the schema which contains the rex list
664+
* @return the target mapping
665+
*/
666+
static Mapping mapping(List<RexNode> rexNodes, RelDataType schema) {
667+
return Mappings.target(getSelectColumns(rexNodes), schema.getFieldCount());
668+
}
669+
670+
static boolean mayBeFilterFromBucketNonNull(LogicalFilter filter) {
671+
RexNode condition = filter.getCondition();
672+
return isNotNullOnRef(condition)
673+
|| ((condition instanceof RexCall)
674+
&& ((RexCall) condition).getOperator().equals(SqlStdOperatorTable.AND)
675+
&& ((RexCall) condition).getOperands().stream().allMatch(PlanUtils::isNotNullOnRef));
676+
}
677+
678+
private static boolean isNotNullOnRef(RexNode rex) {
679+
return rex instanceof RexCall
680+
&& rex.isA(SqlKind.IS_NOT_NULL)
681+
&& ((RexCall) rex).getOperands().get(0) instanceof RexInputRef;
682+
}
626683
}

‎core/src/main/java/org/opensearch/sql/data/type/ExprType.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,18 @@ default Optional<String> getOriginalPath() {
6262
}
6363

6464
/**
65-
* Get the original path. Types like alias type should be derived from the type of the original
66-
* field.
65+
* Get the original expr path. Types like alias type should be derived from the type of the
66+
* original field.
6767
*/
6868
default ExprType getOriginalExprType() {
6969
return this;
7070
}
71+
72+
/**
73+
* Get the original data type. Types like alias type should be derived from the type of the
74+
* original field.
75+
*/
76+
default ExprType getOriginalType() {
77+
return this;
78+
}
7179
}

0 commit comments

Comments
 (0)