Skip to content

Commit 8a1bf11

Browse files
LantaoJinLantaoJin
committed
Enable dedup pushdown
Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 3a3c8c8 commit 8a1bf11

15 files changed

Lines changed: 155 additions & 48 deletions

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

Lines changed: 15 additions & 9 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;
@@ -1317,7 +1318,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
13171318
: duplicatedFieldNames.stream()
13181319
.map(a -> (RexNode) context.relBuilder.field(a))
13191320
.toList();
1320-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1321+
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
13211322
}
13221323
context.relBuilder.join(
13231324
JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1373,7 +1374,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
13731374
List<RexNode> dedupeFields =
13741375
getRightColumnsInJoinCriteria(context.relBuilder, joinCondition);
13751376

1376-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1377+
buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
13771378
}
13781379
context.relBuilder.join(
13791380
JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1538,15 +1539,15 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
15381539
if (keepEmpty) {
15391540
buildDedupOrNull(context, dedupeFields, allowedDuplication);
15401541
} else {
1541-
buildDedupNotNull(context, dedupeFields, allowedDuplication);
1542+
buildDedupNotNull(context, dedupeFields, allowedDuplication, false);
15421543
}
15431544
return context.relBuilder.peek();
15441545
}
15451546

15461547
private static void buildDedupOrNull(
15471548
CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
15481549
/*
1549-
* | dedup 2 a, b keepempty=false
1550+
* | dedup 2 a, b keepempty=true
15501551
* DropColumns('_row_number_dedup_)
15511552
* +- Filter ('_row_number_dedup_ <= n OR isnull('a) OR isnull('b))
15521553
* +- 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]
@@ -1578,7 +1579,10 @@ private static void buildDedupOrNull(
15781579
}
15791580

15801581
private static void buildDedupNotNull(
1581-
CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
1582+
CalcitePlanContext context,
1583+
List<RexNode> dedupeFields,
1584+
Integer allowedDuplication,
1585+
boolean fromJoinMaxOption) {
15821586
/*
15831587
* | dedup 2 a, b keepempty=false
15841588
* DropColumns('_row_number_dedup_)
@@ -1588,6 +1592,8 @@ private static void buildDedupNotNull(
15881592
* +- ...
15891593
*/
15901594
// Filter (isnotnull('a) AND isnotnull('b))
1595+
String rowNumberAlias =
1596+
fromJoinMaxOption ? ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP : ROW_NUMBER_COLUMN_FOR_DEDUP;
15911597
context.relBuilder.filter(
15921598
context.relBuilder.and(dedupeFields.stream().map(context.relBuilder::isNotNull).toList()));
15931599
// Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
@@ -1601,15 +1607,15 @@ private static void buildDedupNotNull(
16011607
.partitionBy(dedupeFields)
16021608
.orderBy(dedupeFields)
16031609
.rowsTo(RexWindowBounds.CURRENT_ROW)
1604-
.as(ROW_NUMBER_COLUMN_FOR_DEDUP);
1610+
.as(rowNumberAlias);
16051611
context.relBuilder.projectPlus(rowNumber);
1606-
RexNode _row_number_dedup_ = context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_DEDUP);
1612+
RexNode rowNumberField = context.relBuilder.field(rowNumberAlias);
16071613
// Filter ('_row_number_dedup_ <= n)
16081614
context.relBuilder.filter(
16091615
context.relBuilder.lessThanOrEqual(
1610-
_row_number_dedup_, context.relBuilder.literal(allowedDuplication)));
1616+
rowNumberField, context.relBuilder.literal(allowedDuplication)));
16111617
// DropColumns('_row_number_dedup_)
1612-
context.relBuilder.projectExcept(_row_number_dedup_);
1618+
context.relBuilder.projectExcept(rowNumberField);
16131619
}
16141620

16151621
@Override

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ public interface PlanUtils {
6262
/** this is only for dedup command, do not reuse it in other command */
6363
String ROW_NUMBER_COLUMN_FOR_DEDUP = "_row_number_dedup_";
6464

65+
String ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP = "_row_number_join_max_dedup_";
6566
String ROW_NUMBER_COLUMN_FOR_RARE_TOP = "_row_number_rare_top_";
6667
String ROW_NUMBER_COLUMN_FOR_MAIN = "_row_number_main_";
6768
String ROW_NUMBER_COLUMN_FOR_SUBSEARCH = "_row_number_subsearch_";

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -82,25 +82,29 @@ public void supportSearchSargPushDown_timeRange() throws IOException {
8282
}
8383

8484
// Only for Calcite
85-
@Ignore("https://github.com/opensearch-project/OpenSearch/issues/3725")
85+
@Test
8686
public void testJoinWithCriteriaAndMaxOption() throws IOException {
87+
// TODO could be optimized with https://github.com/opensearch-project/OpenSearch/issues/3725
88+
enabledOnlyWhenPushdownIsEnabled();
8789
String query =
8890
"source=opensearch-sql_test_index_bank | join max=1 left=l right=r on"
8991
+ " l.account_number=r.account_number opensearch-sql_test_index_bank";
90-
var result = explainQueryToString(query);
91-
String expected = loadExpectedPlan("explain_join_with_criteria_max_option.json");
92-
assertJsonEqualsIgnoreId(expected, result);
92+
var result = explainQueryYaml(query);
93+
String expected = loadExpectedPlan("explain_join_with_criteria_max_option.yaml");
94+
assertYamlEqualsIgnoreId(expected, result);
9395
}
9496

9597
// Only for Calcite
96-
@Ignore("https://github.com/opensearch-project/OpenSearch/issues/3725")
98+
@Test
9799
public void testJoinWithFieldListAndMaxOption() throws IOException {
100+
// TODO could be optimized with https://github.com/opensearch-project/OpenSearch/issues/3725
101+
enabledOnlyWhenPushdownIsEnabled();
98102
String query =
99103
"source=opensearch-sql_test_index_bank | join type=inner max=1 account_number"
100104
+ " opensearch-sql_test_index_bank";
101-
var result = explainQueryToString(query);
102-
String expected = loadExpectedPlan("explain_join_with_fields_max_option.json");
103-
assertJsonEqualsIgnoreId(expected, result);
105+
var result = explainQueryYaml(query);
106+
String expected = loadExpectedPlan("explain_join_with_fields_max_option.yaml");
107+
assertYamlEqualsIgnoreId(expected, result);
104108
}
105109

106110
// Only for Calcite

‎integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,7 +484,7 @@ public void testStatsByTimeSpan() throws IOException {
484484
TEST_INDEX_BANK)));
485485
}
486486

487-
@Ignore("https://github.com/opensearch-project/OpenSearch/issues/3725")
487+
@Test
488488
public void testDedupPushdown() throws IOException {
489489
String expected = loadExpectedPlan("explain_dedup_push.json");
490490
assertJsonEqualsIgnoreId(
@@ -504,7 +504,7 @@ public void testDedupKeepEmptyTruePushdown() throws IOException {
504504
+ " | dedup gender KEEPEMPTY=true"));
505505
}
506506

507-
@Ignore("https://github.com/opensearch-project/OpenSearch/issues/3725")
507+
@Test
508508
public void testDedupKeepEmptyFalsePushdown() throws IOException {
509509
String expected = loadExpectedPlan("explain_dedup_keepempty_false_push.json");
510510
assertJsonEqualsIgnoreId(
@@ -514,6 +514,13 @@ public void testDedupKeepEmptyFalsePushdown() throws IOException {
514514
+ " | dedup gender KEEPEMPTY=false"));
515515
}
516516

517+
@Test
518+
public void testDedupUnsupportedTypeNotPushdown() throws IOException {
519+
String expected = loadExpectedPlan("explain_dedup_unsupported_type_no_push.yaml");
520+
assertYamlEqualsIgnoreId(
521+
expected, explainQueryYaml(String.format("source=%s | dedup male", TEST_INDEX_BANK)));
522+
}
523+
517524
@Test
518525
public void testSingleFieldRelevanceQueryFunctionExplain() throws IOException {
519526
enabledOnlyWhenPushdownIsEnabled();

‎integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_unsupported_type_no_push.yaml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12])
5+
LogicalFilter(condition=[<=($19, 1)])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _id=[$13], _index=[$14], _score=[$15], _maxscore=[$16], _sort=[$17], _routing=[$18], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $12 ORDER BY $12)])
7+
LogicalFilter(condition=[IS NOT NULL($12)])
8+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
9+
physical: |
10+
EnumerableLimit(fetch=[10000])
11+
EnumerableCalc(expr#0..13=[{inputs}], expr#14=[1], expr#15=[<=($t13, $t14)], proj#0..12=[{exprs}], $condition=[$t15])
12+
EnumerableWindow(window#0=[window(partition {12} order by [12] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
13+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], FILTER->IS NOT NULL($12)], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"exists":{"field":"male","boost":1.0}},"_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_join_with_criteria_max_option.json‎

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎integ-test/src/test/resources/expectedOutput/calcite/explain_join_with_criteria_max_option.yaml‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], r.account_number=[$13], r.firstname=[$14], r.address=[$15], r.birthdate=[$16], r.gender=[$17], r.city=[$18], r.lastname=[$19], r.balance=[$20], r.employer=[$21], r.state=[$22], r.age=[$23], r.email=[$24], r.male=[$25])
5+
LogicalJoin(condition=[=($0, $13)], joinType=[inner])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
8+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12])
9+
LogicalFilter(condition=[<=($13, 1)])
10+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _row_number_join_max_dedup_=[ROW_NUMBER() OVER (PARTITION BY $0 ORDER BY $0)])
11+
LogicalFilter(condition=[IS NOT NULL($0)])
12+
LogicalSystemLimit(fetch=[50000], type=[JOIN_SUBSEARCH_MAXOUT])
13+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12])
14+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
15+
physical: |
16+
EnumerableLimit(fetch=[10000])
17+
EnumerableMergeJoin(condition=[=($0, $13)], joinType=[inner])
18+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], SORT->[{
19+
"account_number" : {
20+
"order" : "asc",
21+
"missing" : "_last"
22+
}
23+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
24+
EnumerableSort(sort0=[$0], dir0=[ASC])
25+
EnumerableCalc(expr#0..13=[{inputs}], expr#14=[1], expr#15=[<=($t13, $t14)], proj#0..12=[{exprs}], $condition=[$t15])
26+
EnumerableWindow(window#0=[window(partition {0} order by [0] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
27+
EnumerableCalc(expr#0..12=[{inputs}], expr#13=[IS NOT NULL($t0)], proj#0..12=[{exprs}], $condition=[$t13])
28+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], LIMIT->50000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":50000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]}}, requestedTotalSize=50000, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_join_with_fields_max_option.json‎

Lines changed: 0 additions & 6 deletions
This file was deleted.

0 commit comments

Comments
 (0)