(calcite) branch main updated: [CALCITE-6734] RelFieldTrimmer should trim Aggregate's input fields which are arguments of unused aggregate functions

[zhehu]

This is an automated email from the ASF dual-hosted git repository.

zhehu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/calcite.git


The following commit(s) were added to refs/heads/main by this push:
     new 137fed23f9 [CALCITE-6734] RelFieldTrimmer should trim Aggregate's 
input fields which are arguments of unused aggregate functions
137fed23f9 is described below

commit 137fed23f91ba687f6c4b029e4b654e616191e47
Author: ViggoC <viggo...@gmail.com>
AuthorDate: Tue Dec 17 23:27:58 2024 +0800

    [CALCITE-6734] RelFieldTrimmer should trim Aggregate's input fields which 
are arguments of unused aggregate functions
---
 .../apache/calcite/sql2rel/RelFieldTrimmer.java    | 20 +++++++++------
 .../calcite/sql2rel/RelFieldTrimmerTest.java       | 29 ++++++++++++++++++++++
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java 
b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
index 69e9ae90bb..effb5078f1 100644
--- a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
+++ b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
@@ -1054,16 +1054,20 @@ public class RelFieldTrimmer implements 
ReflectiveVisitor {
     // 1. group fields are always used
     final ImmutableBitSet.Builder inputFieldsUsed =
         aggregate.getGroupSet().rebuild();
-    // 2. agg functions
+    // 2. agg functions: consider only the ones that are needed according to 
fieldsUsed
+    int aggCallIndex = aggregate.getGroupCount();
     for (AggregateCall aggCall : aggregate.getAggCallList()) {
-      inputFieldsUsed.addAll(aggCall.getArgList());
-      if (aggCall.filterArg >= 0) {
-        inputFieldsUsed.set(aggCall.filterArg);
-      }
-      if (aggCall.distinctKeys != null) {
-        inputFieldsUsed.addAll(aggCall.distinctKeys);
+      if (fieldsUsed.get(aggCallIndex)) {
+        inputFieldsUsed.addAll(aggCall.getArgList());
+        if (aggCall.filterArg >= 0) {
+          inputFieldsUsed.set(aggCall.filterArg);
+        }
+        if (aggCall.distinctKeys != null) {
+          inputFieldsUsed.addAll(aggCall.distinctKeys);
+        }
+        inputFieldsUsed.addAll(RelCollations.ordinals(aggCall.collation));
       }
-      inputFieldsUsed.addAll(RelCollations.ordinals(aggCall.collation));
+      aggCallIndex++;
     }
 
     // Create input with trimmed columns.
diff --git 
a/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java 
b/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
index ebe65d93d0..34d94f990c 100644
--- a/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
+++ b/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
@@ -281,6 +281,35 @@ class RelFieldTrimmerTest {
     assertTrue(aggregate.getHints().contains(aggHint));
   }
 
+  /** Test case for
+   * <a 
href="https://issues.apache.org/jira/browse/CALCITE-6734";>[CALCITE-6734]
+   * RelFieldTrimmer should trim Aggregate's input fields which are arguments 
of
+   * unused aggregate functions</a>. */
+  @Test void testTrimUnusedAggregateInput() {
+    final RelBuilder builder = RelBuilder.create(config().build());
+    final RelNode original =
+        builder.scan("EMP")
+            .filter(
+                builder.greaterThan(builder.field("DEPTNO"),
+                    builder.literal(100)))
+            .aggregate(
+                builder.groupKey(builder.field("DEPTNO")),
+                builder.sum(false, "SAL", builder.field("SAL")),
+                builder.count(false, "ENAME", builder.field("ENAME")))
+            .project(builder.field("DEPTNO"), builder.field("SAL"))
+            .build();
+
+    final RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+    final RelNode trimmed = fieldTrimmer.trim(original);
+
+    final String expected = ""
+        + "LogicalAggregate(group=[{2}], SAL=[SUM($1)])\n"
+        + "  LogicalFilter(condition=[>($2, 100)])\n"
+        + "    LogicalProject(EMPNO=[$0], SAL=[$5], DEPTNO=[$7])\n"
+        + "      LogicalTableScan(table=[[scott, EMP]])\n";
+    assertThat(trimmed, hasTree(expected));
+  }
+
   /** Test case for
    * <a 
href="https://issues.apache.org/jira/browse/CALCITE-4055";>[CALCITE-4055]
    * RelFieldTrimmer loses hints</a>. */