Verify that TableScanNode's assignments match the output symbols

Before this change, only one-way match was verified: that each output symbol is backed by an assignment.
trinodb · Jan 22, 2025 · 417f8fe · 417f8fe
1 parent b17fa4e
commit 417f8fe
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/plan/TableScanNode.java b/core/trino-main/src/main/java/io/trino/sql/planner/plan/TableScanNode.java
@@ -19,6 +19,7 @@
 import com.google.common.base.VerifyException;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.DoNotCall;
 import com.google.errorprone.annotations.Immutable;
 import io.trino.cost.PlanNodeStatsEstimate;
@@ -89,7 +90,7 @@ private TableScanNode(
         this.table = requireNonNull(table, "table is null");
         this.outputSymbols = ImmutableList.copyOf(requireNonNull(outputs, "outputs is null"));
         this.assignments = ImmutableMap.copyOf(requireNonNull(assignments, "assignments is null"));
-        checkArgument(assignments.keySet().containsAll(outputs), "assignments does not cover all of outputs");
+        checkArgument(ImmutableSet.copyOf(outputs).equals(assignments.keySet()), "assignments and outputs do not match");
         this.enforcedConstraint = null;
         this.statistics = null;
         this.updateTarget = updateTarget;
@@ -110,7 +111,7 @@ public TableScanNode(
         this.table = requireNonNull(table, "table is null");
         this.outputSymbols = ImmutableList.copyOf(requireNonNull(outputs, "outputs is null"));
         this.assignments = ImmutableMap.copyOf(requireNonNull(assignments, "assignments is null"));
-        checkArgument(assignments.keySet().containsAll(outputs), "assignments does not cover all of outputs");
+        checkArgument(ImmutableSet.copyOf(outputs).equals(assignments.keySet()), "assignments and outputs do not match");
         requireNonNull(enforcedConstraint, "enforcedConstraint is null");
         validateEnforcedConstraint(enforcedConstraint, outputs, assignments);
         this.enforcedConstraint = enforcedConstraint;

diff --git a/...ava/io/trino/sql/planner/iterative/rule/TestMultipleDistinctAggregationsToSubqueries.java b/...ava/io/trino/sql/planner/iterative/rule/TestMultipleDistinctAggregationsToSubqueries.java
@@ -925,7 +925,7 @@ public void testGroupByWithDistinctToSubqueries()
                             .source(
                                     p.tableScan(
                                             testTableHandle(ruleTester),
-                                            ImmutableList.of(input1Symbol, input2Symbol),
+                                            ImmutableList.of(input1Symbol, input2Symbol, groupingKey),
                                             ImmutableMap.of(
                                                     input1Symbol, COLUMN_1_HANDLE,
                                                     input2Symbol, COLUMN_2_HANDLE,