From ed4560fc327cae4ba505a15b2dc213539398e048 Mon Sep 17 00:00:00 2001 From: Mihai Budiu Date: Fri, 3 Jul 2026 12:10:45 -0700 Subject: [PATCH] [SQL] Generalize optimization rule for MAX(CASE WHEN expression THEN 1 END) Signed-off-by: Mihai Budiu --- .../optimizer/MaxCaseToCountRule.java | 85 ++++++++++++------- .../compiler/sql/simple/Regression2Tests.java | 80 ++++++++++++++++- 2 files changed, 133 insertions(+), 32 deletions(-) diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/calciteCompiler/optimizer/MaxCaseToCountRule.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/calciteCompiler/optimizer/MaxCaseToCountRule.java index 93ac1fb682a..693879d6fa1 100644 --- a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/calciteCompiler/optimizer/MaxCaseToCountRule.java +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/calciteCompiler/optimizer/MaxCaseToCountRule.java @@ -30,15 +30,34 @@ /** * Rule that converts MAX(CASE ... THEN 1 ELSE 0 END) - * into COUNT(...) FILTER (WHERE ...) - * Rewrite: + * or MAX(CASE ... THEN 1 ELSE NULL END) (equivalently, no ELSE clause) + * into COUNT(...) FILTER (WHERE ...). + * + *

Rewrite for ELSE 0: + * 

  * LogicalAggregate(group=[{0}], [MAX($1)]
  *   LogicalProject([$0], $f1=[CASE(=(...), 1, 0)])
- * into
+ *
+ * becomes + * 
  * LogicalProject([$0], [CASE(>($1, 0), 1, 0)])
  *   LogicalAggregate(group=[{0}], agg#0=[COUNT(*) FILTER $1])
  *     LogicalProject([$0], $f1=[IS TRUE(...)])
- * This works only if there is a GROUP BY.
+ *
+ * + *

Rewrite for ELSE NULL (or no ELSE): + * 

+ * LogicalAggregate(group=[{0}], [MAX($1)]
+ *   LogicalProject([$0], $f1=[CASE(=(...), 1, null)])
+ *
+ * becomes + * 
+ * LogicalProject([$0], [CASE(>($1, 0), 1, null)])
+ *   LogicalAggregate(group=[{0}], agg#0=[COUNT(*) FILTER $1])
+ *     LogicalProject([$0], $f1=[IS TRUE(...)])
+ *
+ * + *

This works only if there is a GROUP BY. */ public class MaxCaseToCountRule extends RelRule> @@ -66,16 +85,17 @@ public void onMatch(RelOptRuleCall call) { RelDataType bigInt = cluster.getTypeFactory().createSqlType(SqlTypeName.BIGINT); bigInt = cluster.getTypeFactory().createTypeWithNullability(bigInt, true); - RexLiteral one = rexBuilder.makeLiteral(BigDecimal.ONE, bigInt); - newProjects.add(one); for (int i = 0; i < aggregate.getGroupCount(); i++) // Preserve the keys in the postProjects postProjects.add(new RexInputRef(i, resultType.getFieldList().get(i).getType())); + final int groupCount = aggregate.getGroupCount(); for (AggregateCall aggregateCall : aggregate.getAggCallList()) { AggregateCall newCall = transform(aggregateCall, project, newProjects, postProjects, bigInt); if (newCall == null) { - postProjects.add(rexBuilder.makeInputRef(aggregateCall.getType(), newAggregates.size())); + // Aggregate output: [group_keys (0..groupCount-1), agg_0, agg_1, ...] + postProjects.add(rexBuilder.makeInputRef( + aggregateCall.getType(), groupCount + newAggregates.size())); newAggregates.add(aggregateCall); } else { newAggregates.add(newCall); @@ -147,11 +167,17 @@ public static Long getLiteralValue(RexNode node) { if (op1Value == null) return null; Long op2Value = getLiteralValue(op2); - if (op2Value == null) + if (op2Value == null && !(op2 instanceof RexLiteral)) return null; + boolean elseIsNull = op2Value == null; - boolean flip; - if (op1Value == 0L && op2Value == 1L) { + final boolean flip; + if (elseIsNull) { + // Only THEN 1 ELSE NULL is supported: MAX = 1 if any row matches, NULL otherwise. + if (op1Value != 1L) + return null; + flip = false; + } else if (op1Value == 0L && op2Value == 1L) { flip = true; } else if (op1Value == 1L && op2Value == 0L) { flip = false; @@ -179,27 +205,26 @@ public static Long getLiteralValue(RexNode node) { false, call.rexList, ImmutableList.of(), newProjects.size() - 1, null, RelCollations.EMPTY, bigInt, call.getName()); - // CASE(WHEN $1 IS NULL THEN NULL WHEN $1 = 0 THEN 0 ELSE 1 END) - RelDataType nullable = - cluster.getTypeFactory().createTypeWithNullability(op1.getType(), call.type.isNullable()); - List postOperands = new ArrayList<>(); - // WHEN $1 IS NULL RexInputRef ref = rexBuilder.makeInputRef(bigInt, postProjects.size()); - RexNode isNull = rexBuilder.makeCall(pos, SqlStdOperatorTable.IS_NULL, ref); - postOperands.add(isNull); - // THEN NULL - RexLiteral nullLit = rexBuilder.makeNullLiteral(nullable); - postOperands.add(nullLit); - // WHEN $1 = 0 - RexLiteral biggerZero = rexBuilder.makeLiteral(BigDecimal.ZERO, bigInt); - postOperands.add(rexBuilder.makeCall(pos, SqlStdOperatorTable.EQUALS, ref, biggerZero)); - // THEN 0 - RexNode zero = rexBuilder.makeLiteral(BigDecimal.ZERO, nullable, true); - postOperands.add(zero); - // ELSE one - RexNode one = rexBuilder.makeLiteral(BigDecimal.ONE, nullable, true); - postOperands.add(one); - final RexNode postCondition = rexBuilder.makeCall(pos, SqlStdOperatorTable.CASE, postOperands); + List postOperands = new ArrayList<>(); + RelDataType nullable = cluster.getTypeFactory().createTypeWithNullability( + op1.getType(), call.type.isNullable()); + RexLiteral zero = rexBuilder.makeLiteral(BigDecimal.ZERO, bigInt); + final RexNode postCondition; + if (elseIsNull) { + // CASE WHEN count > 0 THEN 1 ELSE NULL END + postOperands.add(rexBuilder.makeCall(pos, SqlStdOperatorTable.GREATER_THAN, ref, zero)); + postOperands.add(rexBuilder.makeLiteral(BigDecimal.ONE, nullable, true)); + postOperands.add(rexBuilder.makeNullLiteral(nullable)); + } else { + // CASE WHEN count IS NULL THEN NULL WHEN count = 0 THEN 0 ELSE 1 END + postOperands.add(rexBuilder.makeCall(pos, SqlStdOperatorTable.IS_NULL, ref)); + postOperands.add(rexBuilder.makeNullLiteral(nullable)); + postOperands.add(rexBuilder.makeCall(pos, SqlStdOperatorTable.EQUALS, ref, zero)); + postOperands.add(rexBuilder.makeLiteral(BigDecimal.ZERO, nullable, true)); + postOperands.add(rexBuilder.makeLiteral(BigDecimal.ONE, nullable, true)); + } + postCondition = rexBuilder.makeCall(pos, SqlStdOperatorTable.CASE, postOperands); postProjects.add(postCondition); return result; } diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/simple/Regression2Tests.java b/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/simple/Regression2Tests.java index cd913406d64..20c025fd6c7 100644 --- a/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/simple/Regression2Tests.java +++ b/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/simple/Regression2Tests.java @@ -980,13 +980,14 @@ public void issue5927() { ccs.stepWeightOne("", """ g | max ---------"""); - ccs.stepWeightOne("INSERT INTO T VALUES ('a', 1), ('b', 3), ('c', 3), ('c', 5), ('d', 1), ('d', 5);", """ + ccs.stepWeightOne("INSERT INTO T VALUES ('a', 1), ('b', 3), ('c', 3), ('c', 5), ('d', 1), ('d', 5), ('e', NULL);", """ g | max --------- a| 0 b| 1 c| 1 - d| 1"""); + d| 1 + e| 0"""); ccs.visit(this.findLinear(ccs.compiler)); } @@ -1042,6 +1043,81 @@ public void issue5927a() { ccs.visit(this.findLinear(ccs.compiler)); } + @Test + public void issue6590() { + // MAX(CASE WHEN cond THEN 1 ELSE NULL END) — ELSE NULL variant + var ccs = this.getCCS(""" + CREATE TABLE T(g VARCHAR, x INT); + CREATE VIEW V AS SELECT g, MAX(CASE WHEN x > 2 THEN 1 ELSE NULL END) FROM T GROUP BY g;"""); + // Validated on Postgres: NULL when no row satisfies cond, 1 otherwise. + ccs.stepWeightOne("", """ + g | max + ---------"""); + ccs.stepWeightOne(""" + INSERT INTO T VALUES + ('a', 1), ('b', 3), ('c', 3), ('c', 5), ('d', 1), ('d', 5), ('e', NULL); + """, """ + g | max + --------- + a| NULL + b| 1 + c| 1 + d| 1 + e| NULL"""); + ccs.visit(this.findLinear(ccs.compiler)); + } + + @Test + public void issue6590a() { + // MAX(CASE WHEN cond THEN 1 END) — no ELSE clause (equivalent to ELSE NULL) + var ccs = this.getCCS(""" + CREATE TABLE T(g VARCHAR, x INT); + CREATE VIEW V AS SELECT g, MAX(CASE WHEN x > 2 THEN 1 END) FROM T GROUP BY g;"""); + // Validated on Postgres: NULL when no row satisfies cond, 1 otherwise. + ccs.stepWeightOne("", """ + g | max + ---------"""); + ccs.stepWeightOne(""" + INSERT INTO T VALUES + ('a', 1), ('b', 3), ('c', 3), ('c', 5), ('d', 1), ('d', 5), + ('e', NULL); + """, """ + g | max + --------- + a| NULL + b| 1 + c| 1 + d| 1 + e| NULL"""); + ccs.visit(this.findLinear(ccs.compiler)); + } + + @Test + public void issue6590b() { + // MAX(CASE...) mixed with SUM in the same aggregate. + // Exercises the post-project index path for untransformed aggregate calls. + var ccs = this.getCCS(""" + CREATE TABLE T(g VARCHAR, x INT); + CREATE VIEW V AS SELECT g, MAX(CASE WHEN x > 2 THEN 1 ELSE 0 END), SUM(x) FROM T GROUP BY g;"""); + // Validated on Postgres + ccs.stepWeightOne("", """ + g | max | sum + --------------"""); + ccs.stepWeightOne(""" + INSERT INTO T VALUES + ('a', 1), ('b', 3), ('c', 3), ('c', 5), ('d', 1), ('d', 5), + ('e', NULL); + """, """ + g | max | sum + -------------- + a| 0| 1 + b| 1| 3 + c| 1| 8 + d| 1| 6 + e| 0| NULL"""); + ccs.visit(this.findLinear(ccs.compiler)); + } + @Test public void calciteIssue7501() { this.getCC("""