Skip to content

Commit 71c9c4f

Browse files
xiedeyantuxiedeyantu
committed
support count(*) filter (where ...)
1 parent 3be39b7 commit 71c9c4f

File tree

2 files changed

+54
-42
lines changed

2 files changed

+54
-42
lines changed

‎core/src/main/java/org/apache/calcite/rel/rules/IntersectToDistinctRule.java

+26-16
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.calcite.rex.RexBuilder;
2626
import org.apache.calcite.rex.RexNode;
2727
import org.apache.calcite.tools.RelBuilder;
28+
import org.apache.calcite.tools.RelBuilder.AggCall;
2829
import org.apache.calcite.tools.RelBuilderFactory;
2930
import org.apache.calcite.util.ImmutableBitSet;
3031
import org.apache.calcite.util.Util;
@@ -90,38 +91,47 @@ public IntersectToDistinctRule(Class<? extends Intersect> intersectClass,
9091
final RelOptCluster cluster = intersect.getCluster();
9192
final RexBuilder rexBuilder = cluster.getRexBuilder();
9293
final RelBuilder relBuilder = call.builder();
94+
final String countPrefix = "count_i";
95+
final String indexAlias = "i";
96+
final int oriFieldCount = intersect.getRowType().getFieldCount();
97+
final int inputCount = intersect.getInputs().size();
9398

94-
for (int i = 0; i < intersect.getInputs().size(); ++i) {
99+
AggCall[] aggCalls = new AggCall[inputCount];
100+
for (int i = 0; i < inputCount; ++i) {
95101
relBuilder.push(intersect.getInputs().get(i));
96102
List<RexNode> fields = new ArrayList<>(relBuilder.fields());
97103
// Add a constant column "i" for each input
98104
fields.add(
99105
relBuilder.alias(
100-
rexBuilder.makeBigintLiteral(new BigDecimal(i)), "i"));
106+
rexBuilder.makeBigintLiteral(new BigDecimal(i)), indexAlias));
101107
relBuilder.project(fields);
108+
// Generate count(*) filter (where i = n) as count_i{n}
109+
aggCalls[i] = relBuilder.countStar(null)
110+
.filter(
111+
relBuilder.equals(relBuilder.field(indexAlias),
112+
relBuilder.literal(i))).as(countPrefix + i);
102113
}
103114

104115
// create a union above all the branches
105-
final int branchCount = intersect.getInputs().size();
106-
relBuilder.union(true, branchCount);
116+
relBuilder.union(true, inputCount);
107117
final RelNode union = relBuilder.peek();
108-
109118
final int fieldCount = union.getRowType().getFieldCount();
110119

111-
final ImmutableBitSet groupSet =
112-
ImmutableBitSet.range(fieldCount);
120+
// Add aggCalls
121+
relBuilder.aggregate(relBuilder.groupKey(ImmutableBitSet.range(fieldCount)), aggCalls);
113122

114-
// Perform count(distinct) on column "i"
115-
relBuilder.aggregate(relBuilder.groupKey(groupSet),
116-
relBuilder.count(true, "i", relBuilder.field("i")));
123+
// Generate filter count_i{n} > 0 for each branch
124+
RexNode[] filters = new RexNode[inputCount];
125+
for (int i = 0; i < inputCount; i++) {
126+
filters[i] =
127+
relBuilder.greaterThan(relBuilder.field(countPrefix + i), relBuilder.literal(0));
128+
}
117129

118-
// add a filter count(c) = #branches
119-
relBuilder.filter(
120-
relBuilder.equals(relBuilder.field(fieldCount),
121-
rexBuilder.makeBigintLiteral(new BigDecimal(branchCount))));
130+
// Add filter count_i{n} > 0 for each branch
131+
relBuilder.filter(relBuilder.and(filters));
122132

123-
// Project all but the last two field(i and c)
124-
relBuilder.project(Util.skipLast(relBuilder.fields(), 2));
133+
// Project all but the last added field(e.g. i and count_i{n})
134+
relBuilder.project(Util.skipLast(relBuilder.fields(), relBuilder.fields().size() - oriFieldCount));
125135

126136
call.transformTo(relBuilder.build());
127137
}

‎core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml

+28-26
Original file line numberDiff line numberDiff line change
@@ -5917,21 +5917,22 @@ LogicalIntersect(all=[false])
59175917
<Resource name="planAfter">
59185918
<![CDATA[
59195919
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5920-
LogicalFilter(condition=[=($10, 3)])
5921-
LogicalAggregate(group=[{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}], i=[COUNT(DISTINCT $9)])
5922-
LogicalUnion(all=[true])
5923-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[0:BIGINT])
5924-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5925-
LogicalFilter(condition=[=($7, 10)])
5926-
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5927-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[1:BIGINT])
5928-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5929-
LogicalFilter(condition=[=($7, 20)])
5930-
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5931-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[2:BIGINT])
5932-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5933-
LogicalFilter(condition=[=($7, 30)])
5934-
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5920+
LogicalFilter(condition=[AND(>($10, 0), >($11, 0), >($12, 0))])
5921+
LogicalAggregate(group=[{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}], count_i0=[COUNT() FILTER $10], count_i1=[COUNT() FILTER $11], count_i2=[COUNT() FILTER $12])
5922+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[$9], $f10=[=($9, 0)], $f11=[=($9, 1)], $f12=[=($9, 2)])
5923+
LogicalUnion(all=[true])
5924+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[0:BIGINT])
5925+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5926+
LogicalFilter(condition=[=($7, 10)])
5927+
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5928+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[1:BIGINT])
5929+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5930+
LogicalFilter(condition=[=($7, 20)])
5931+
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5932+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[2:BIGINT])
5933+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5934+
LogicalFilter(condition=[=($7, 30)])
5935+
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
59355936
]]>
59365937
</Resource>
59375938
</TestCase>
@@ -5963,17 +5964,18 @@ LogicalIntersect(all=[true])
59635964
<![CDATA[
59645965
LogicalIntersect(all=[true])
59655966
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5966-
LogicalFilter(condition=[=($10, 2)])
5967-
LogicalAggregate(group=[{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}], i=[COUNT(DISTINCT $9)])
5968-
LogicalUnion(all=[true])
5969-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[0:BIGINT])
5970-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5971-
LogicalFilter(condition=[=($7, 10)])
5972-
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5973-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[1:BIGINT])
5974-
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5975-
LogicalFilter(condition=[=($7, 20)])
5976-
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5967+
LogicalFilter(condition=[AND(>($10, 0), >($11, 0))])
5968+
LogicalAggregate(group=[{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}], count_i0=[COUNT() FILTER $10], count_i1=[COUNT() FILTER $11])
5969+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[$9], $f10=[=($9, 0)], $f11=[=($9, 1)])
5970+
LogicalUnion(all=[true])
5971+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[0:BIGINT])
5972+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5973+
LogicalFilter(condition=[=($7, 10)])
5974+
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
5975+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], i=[1:BIGINT])
5976+
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
5977+
LogicalFilter(condition=[=($7, 20)])
5978+
LogicalTableScan(table=[[CATALOG, SALES, EMP]])
59775979
LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
59785980
LogicalFilter(condition=[=($7, 30)])
59795981
LogicalTableScan(table=[[CATALOG, SALES, EMP]])

0 commit comments

Comments
 (0)