Skip to content

Commit 8c417f4

Browse files
LantaoJinLantaoJin
authored
Optimize count aggregation performance by utilizing native doc_count in v3 (opensearch-project#4337)
* Optimize bucket aggregation performance by utilizing native doc_count in v3 Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix UT Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix issue of count(FIELD) Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix comments Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix typo Signed-off-by: Lantao Jin <ltjin@amazon.com> * revert the doc_count pushdown for count(FIELD) by EXPR Signed-off-by: Lantao Jin <ltjin@amazon.com> * Support pushdown count aggregation in no bucket aggregation to hits.total.value Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 68d47cb commit 8c417f4

59 files changed

Lines changed: 890 additions & 69 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: 95 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_MAIN;
1919
import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_SUBSEARCH;
2020
import static org.opensearch.sql.calcite.utils.PlanUtils.getRelation;
21+
import static org.opensearch.sql.calcite.utils.PlanUtils.getRexCall;
2122
import static org.opensearch.sql.calcite.utils.PlanUtils.transformPlanToAttachChild;
2223

2324
import com.google.common.base.Strings;
@@ -53,6 +54,7 @@
5354
import org.apache.calcite.rex.RexNode;
5455
import org.apache.calcite.rex.RexVisitorImpl;
5556
import org.apache.calcite.rex.RexWindowBounds;
57+
import org.apache.calcite.sql.SqlKind;
5658
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
5759
import org.apache.calcite.sql.type.SqlTypeFamily;
5860
import org.apache.calcite.sql.type.SqlTypeName;
@@ -825,6 +827,23 @@ private void projectPlusOverriding(
825827
context.relBuilder.rename(expectedRenameFields);
826828
}
827829

830+
private List<List<RexInputRef>> extractInputRefList(List<RelBuilder.AggCall> aggCalls) {
831+
return aggCalls.stream()
832+
.map(RelBuilder.AggCall::over)
833+
.map(RelBuilder.OverCall::toRex)
834+
.map(node -> getRexCall(node, this::isCountField))
835+
.map(list -> list.isEmpty() ? null : list.getFirst())
836+
.map(PlanUtils::getInputRefs)
837+
.toList();
838+
}
839+
840+
/** Is count(FIELD) */
841+
private boolean isCountField(RexCall call) {
842+
return call.isA(SqlKind.COUNT)
843+
&& call.getOperands().size() == 1 // count(FIELD)
844+
&& call.getOperands().get(0) instanceof RexInputRef;
845+
}
846+
828847
/**
829848
* Resolve the aggregation with trimming unused fields to avoid bugs in {@link
830849
* org.apache.calcite.sql2rel.RelDecorrelator#decorrelateRel(Aggregate, boolean)}
@@ -838,6 +857,72 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
838857
List<UnresolvedExpression> groupExprList,
839858
List<UnresolvedExpression> aggExprList,
840859
CalcitePlanContext context) {
860+
Pair<List<RexNode>, List<AggCall>> resolved =
861+
resolveAttributesForAggregation(groupExprList, aggExprList, context);
862+
List<RexNode> resolvedGroupByList = resolved.getLeft();
863+
List<AggCall> resolvedAggCallList = resolved.getRight();
864+
865+
// `doc_count` optimization required a filter `isNotNull(RexInputRef)` for the
866+
// `count(FIELD)` aggregation which only can be applied to single FIELD without grouping:
867+
//
868+
// Example 1: source=t | stats count(a)
869+
// Before: Aggregate(count(a))
870+
// \- Scan t
871+
// After: Aggregate(count(a))
872+
// \- Filter(isNotNull(a))
873+
// \- Scan t
874+
//
875+
// Example 2: source=t | stats count(a), count(a)
876+
// Before: Aggregate(count(a), count(a))
877+
// \- Scan t
878+
// After: Aggregate(count(a), count(a))
879+
// \- Filter(isNotNull(a))
880+
// \- Scan t
881+
//
882+
// Example 3: source=t | stats count(a) by b
883+
// Before & After: Aggregate(count(a) by b)
884+
// \- Scan t
885+
//
886+
// Example 4: source=t | stats count()
887+
// Before & After: Aggregate(count())
888+
// \- Scan t
889+
//
890+
// Example 5: source=t | stats count(), count(a)
891+
// Before & After: Aggregate(count(), count(a))
892+
// \- Scan t
893+
//
894+
// Example 6: source=t | stats count(a), count(b)
895+
// Before & After: Aggregate(count(a), count(b))
896+
// \- Scan t
897+
//
898+
// Example 7: source=t | stats count(a+1)
899+
// Before & After: Aggregate(count(a+1))
900+
// \- Scan t
901+
if (resolvedGroupByList.isEmpty()) {
902+
List<List<RexInputRef>> refsPerCount = extractInputRefList(resolvedAggCallList);
903+
List<RexInputRef> distinctRefsOfCounts;
904+
if (context.relBuilder.peek() instanceof org.apache.calcite.rel.core.Project project) {
905+
List<RexNode> mappedInProject =
906+
refsPerCount.stream()
907+
.flatMap(List::stream)
908+
.map(ref -> project.getProjects().get(ref.getIndex()))
909+
.toList();
910+
if (mappedInProject.stream().allMatch(RexInputRef.class::isInstance)) {
911+
distinctRefsOfCounts =
912+
mappedInProject.stream().map(RexInputRef.class::cast).distinct().toList();
913+
} else {
914+
distinctRefsOfCounts = List.of();
915+
}
916+
} else {
917+
distinctRefsOfCounts = refsPerCount.stream().flatMap(List::stream).distinct().toList();
918+
}
919+
if (distinctRefsOfCounts.size() == 1 && refsPerCount.stream().noneMatch(List::isEmpty)) {
920+
context.relBuilder.filter(context.relBuilder.isNotNull(distinctRefsOfCounts.getFirst()));
921+
}
922+
}
923+
924+
// Add project before aggregate:
925+
//
841926
// Example 1: source=t | where a > 1 | stats avg(b + 1) by c
842927
// Before: Aggregate(avg(b + 1))
843928
// \- Filter(a > 1)
@@ -848,23 +933,22 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
848933
// \- Scan t
849934
//
850935
// Example 2: source=t | where a > 1 | top b by c
851-
// Before: Aggregate(count)
852-
// \-Filter(a > 1)
936+
// Before: Aggregate(count(b) by c)
937+
// \-Filter(a > 1 && isNotNull(b))
853938
// \- Scan t
854-
// After: Aggregate(count)
939+
// After: Aggregate(count(b) by c)
855940
// \- Project([c, b])
856-
// \- Filter(a > 1)
941+
// \- Filter(a > 1 && isNotNull(b))
857942
// \- Scan t
858-
// Example 3: source=t | stats count(): no project added for count()
859-
// Before: Aggregate(count)
943+
//
944+
// Example 3: source=t | stats count(): no change for count()
945+
// Before: Aggregate(count())
860946
// \- Scan t
861-
// After: Aggregate(count)
947+
// After: Aggregate(count())
862948
// \- Scan t
863-
Pair<List<RexNode>, List<AggCall>> resolved =
864-
resolveAttributesForAggregation(groupExprList, aggExprList, context);
865949
List<RexInputRef> trimmedRefs = new ArrayList<>();
866-
trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.getLeft())); // group-by keys first
867-
trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.getRight()));
950+
trimmedRefs.addAll(PlanUtils.getInputRefs(resolvedGroupByList)); // group-by keys first
951+
trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolvedAggCallList));
868952
context.relBuilder.project(trimmedRefs);
869953

870954
// Re-resolve all attributes based on adding trimmed Project.

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.util.ArrayList;
1616
import java.util.List;
1717
import java.util.Objects;
18+
import java.util.function.Predicate;
1819
import java.util.stream.Collectors;
1920
import javax.annotation.Nullable;
2021
import org.apache.calcite.plan.RelOptTable;
@@ -255,6 +256,9 @@ static RelBuilder.AggCall makeAggCall(
255256

256257
/** Get all uniq input references from a RexNode. */
257258
static List<RexInputRef> getInputRefs(RexNode node) {
259+
if (node == null) {
260+
return List.of();
261+
}
258262
List<RexInputRef> inputRefs = new ArrayList<>();
259263
node.accept(
260264
new RexVisitorImpl<Void>(true) {
@@ -274,6 +278,26 @@ static List<RexInputRef> getInputRefs(List<RexNode> nodes) {
274278
return nodes.stream().flatMap(node -> getInputRefs(node).stream()).toList();
275279
}
276280

281+
/** Get all uniq RexCall from RexNode with a predicate */
282+
static List<RexCall> getRexCall(RexNode node, Predicate<RexCall> predicate) {
283+
List<RexCall> list = new ArrayList<>();
284+
node.accept(
285+
new RexVisitorImpl<Void>(true) {
286+
@Override
287+
public Void visitCall(RexCall inputCall) {
288+
if (predicate.test(inputCall)) {
289+
if (!list.contains(inputCall)) {
290+
list.add(inputCall);
291+
}
292+
} else {
293+
inputCall.getOperands().forEach(call -> call.accept(this));
294+
}
295+
return null;
296+
}
297+
});
298+
return list;
299+
}
300+
277301
/** Get all uniq input references from a list of agg calls. */
278302
static List<RexInputRef> getInputRefsFromAggCall(List<RelBuilder.AggCall> aggCalls) {
279303
return aggCalls.stream()

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

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.opensearch.sql.calcite.remote;
77

8+
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
89
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
910
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_LOGS;
1011
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
@@ -645,6 +646,143 @@ public void testExplainMinOnStringField() throws IOException {
645646
explainQueryToString("source=opensearch-sql_test_index_account | stats min(firstname)"));
646647
}
647648

649+
@Test
650+
@Override
651+
public void testCountAggPushDownExplain() throws IOException {
652+
enabledOnlyWhenPushdownIsEnabled();
653+
// should be optimized by hits.total.value
654+
String expected = loadExpectedPlan("explain_count_agg_push1.yaml");
655+
assertYamlEqualsJsonIgnoreId(
656+
expected,
657+
explainQueryToString("source=opensearch-sql_test_index_account | stats count() as cnt"));
658+
659+
// should be optimized
660+
expected = loadExpectedPlan("explain_count_agg_push2.yaml");
661+
assertYamlEqualsJsonIgnoreId(
662+
expected,
663+
explainQueryToString(
664+
"source=opensearch-sql_test_index_account | stats count(lastname) as cnt"));
665+
666+
// should be optimized
667+
expected = loadExpectedPlan("explain_count_agg_push3.yaml");
668+
assertYamlEqualsJsonIgnoreId(
669+
expected,
670+
explainQueryToString(
671+
"source=opensearch-sql_test_index_account | eval name = lastname | stats count(name) as"
672+
+ " cnt"));
673+
674+
// should be optimized
675+
expected = loadExpectedPlan("explain_count_agg_push4.yaml");
676+
assertYamlEqualsJsonIgnoreId(
677+
expected,
678+
explainQueryToString(
679+
"source=opensearch-sql_test_index_account | stats count() as c1, count() as c2"));
680+
681+
// should be optimized
682+
expected = loadExpectedPlan("explain_count_agg_push5.yaml");
683+
assertYamlEqualsJsonIgnoreId(
684+
expected,
685+
explainQueryToString(
686+
"source=opensearch-sql_test_index_account | stats count(lastname) as c1,"
687+
+ " count(lastname) as c2"));
688+
689+
// should be optimized
690+
expected = loadExpectedPlan("explain_count_agg_push6.yaml");
691+
assertYamlEqualsJsonIgnoreId(
692+
expected,
693+
explainQueryToString(
694+
"source=opensearch-sql_test_index_account | eval name = lastname | stats"
695+
+ " count(lastname), count(name)"));
696+
697+
// should not be optimized
698+
expected = loadExpectedPlan("explain_count_agg_push7.yaml");
699+
assertYamlEqualsJsonIgnoreId(
700+
expected,
701+
explainQueryToString(
702+
"source=opensearch-sql_test_index_account | stats count(balance + 1) as cnt"));
703+
704+
// should not be optimized
705+
expected = loadExpectedPlan("explain_count_agg_push8.yaml");
706+
assertYamlEqualsJsonIgnoreId(
707+
expected,
708+
explainQueryToString(
709+
"source=opensearch-sql_test_index_account | stats count() as c1, count(lastname) as"
710+
+ " c2"));
711+
712+
// should not be optimized
713+
expected = loadExpectedPlan("explain_count_agg_push9.yaml");
714+
assertYamlEqualsJsonIgnoreId(
715+
expected,
716+
explainQueryToString(
717+
"source=opensearch-sql_test_index_account | stats count(firstname), count(lastname)"));
718+
719+
// should not be optimized
720+
expected = loadExpectedPlan("explain_count_agg_push10.yaml");
721+
assertYamlEqualsJsonIgnoreId(
722+
expected,
723+
explainQueryToString(
724+
"source=opensearch-sql_test_index_account | eval name = lastname | stats"
725+
+ " count(firstname), count(name)"));
726+
}
727+
728+
@Test
729+
public void testExplainCountsByAgg() throws IOException {
730+
enabledOnlyWhenPushdownIsEnabled();
731+
String expected = loadExpectedPlan("explain_agg_counts_by1.yaml");
732+
// case of only count(): doc_count works
733+
assertYamlEqualsJsonIgnoreId(
734+
expected,
735+
explainQueryToString(
736+
String.format(
737+
"source=%s | stats count(), count() as c1 by gender", TEST_INDEX_ACCOUNT)));
738+
739+
// count(FIELD) by: doc_count doesn't work
740+
expected = loadExpectedPlan("explain_agg_counts_by2.yaml");
741+
assertYamlEqualsJsonIgnoreId(
742+
expected,
743+
explainQueryToString(
744+
String.format(
745+
"source=%s | stats count(balance) as c1, count(balance) as c2 by gender",
746+
TEST_INDEX_ACCOUNT)));
747+
748+
// count(FIELD) by: doc_count doesn't work
749+
expected = loadExpectedPlan("explain_agg_counts_by3.yaml");
750+
assertYamlEqualsJsonIgnoreId(
751+
expected,
752+
explainQueryToString(
753+
String.format(
754+
"source=%s | eval account_number_alias = account_number"
755+
+ " | stats count(account_number), count(account_number_alias) as c2 by gender",
756+
TEST_INDEX_ACCOUNT)));
757+
758+
// count() + count(FIELD)): doc_count doesn't work
759+
expected = loadExpectedPlan("explain_agg_counts_by4.yaml");
760+
assertYamlEqualsJsonIgnoreId(
761+
expected,
762+
explainQueryToString(
763+
String.format(
764+
"source=%s | stats count(), count(account_number) by gender", TEST_INDEX_ACCOUNT)));
765+
766+
// count(FIELD1) + count(FIELD2)) by: doc_count doesn't work
767+
expected = loadExpectedPlan("explain_agg_counts_by5.yaml");
768+
assertYamlEqualsJsonIgnoreId(
769+
expected,
770+
explainQueryToString(
771+
String.format(
772+
"source=%s | stats count(balance), count(account_number) by gender",
773+
TEST_INDEX_ACCOUNT)));
774+
775+
// case of count(EXPRESSION) by: doc_count doesn't work
776+
expected = loadExpectedPlan("explain_agg_counts_by6.yaml");
777+
assertYamlEqualsJsonIgnoreId(
778+
expected,
779+
explainQueryToString(
780+
String.format(
781+
"source=%s | eval b_1 = balance + 1"
782+
+ " | stats count(b_1), count(pow(balance, 2)) as c3 by gender",
783+
TEST_INDEX_ACCOUNT)));
784+
}
785+
648786
@Test
649787
public void testExplainSortOnMetricsNoBucketNullable() throws IOException {
650788
// TODO enhancement later: https://github.com/opensearch-project/sql/issues/4282

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -772,4 +772,29 @@ public void testStatsBySpanTimeWithNullBucket() throws IOException {
772772
rows(8213, "2025-07-31 00:00:00"),
773773
rows(8490, "2025-07-31 12:00:00"));
774774
}
775+
776+
@Test
777+
public void testStatsByCounts() throws IOException {
778+
JSONObject response =
779+
executeQuery(
780+
String.format(
781+
"source=%s | eval b_1 = balance + 1 | stats count(), count() as c1,"
782+
+ " count(account_number), count(lastname) as c2, count(balance/10),"
783+
+ " count(pow(balance, 2)) as c3, count(b_1) by gender",
784+
TEST_INDEX_ACCOUNT));
785+
verifySchema(
786+
response,
787+
schema("count()", null, isCalciteEnabled() ? "bigint" : "int"),
788+
schema("c1", null, isCalciteEnabled() ? "bigint" : "int"),
789+
schema("count(account_number)", null, isCalciteEnabled() ? "bigint" : "int"),
790+
schema("c2", null, isCalciteEnabled() ? "bigint" : "int"),
791+
schema("count(balance/10)", null, isCalciteEnabled() ? "bigint" : "int"),
792+
schema("c3", null, isCalciteEnabled() ? "bigint" : "int"),
793+
schema("count(b_1)", null, isCalciteEnabled() ? "bigint" : "int"),
794+
schema("gender", null, "string"));
795+
verifyDataRows(
796+
response,
797+
rows(493, 493, 493, 493, 493, 493, 493, "F"),
798+
rows(507, 507, 507, 507, 507, 507, 507, "M"));
799+
}
775800
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(count()=[$1], c1=[$1], gender=[$0])
5+
LogicalAggregate(group=[{0}], count()=[COUNT()])
6+
LogicalProject(gender=[$4])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
8+
physical: |
9+
EnumerableCalc(expr#0..1=[{inputs}], count()=[$t0], count()0=[$t0], gender=[$t1])
10+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), LIMIT->10000, PROJECT->[count(), gender]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(c1=[$1], c2=[$1], gender=[$0])
5+
LogicalAggregate(group=[{0}], c1=[COUNT($1)])
6+
LogicalProject(gender=[$4], balance=[$3])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
8+
physical: |
9+
EnumerableCalc(expr#0..1=[{inputs}], c1=[$t0], c10=[$t0], gender=[$t1])
10+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},c1=COUNT($1)), LIMIT->10000, PROJECT->[c1, gender]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"c1":{"value_count":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

0 commit comments

Comments
 (0)