Use proper symbol type

trinodb · Oct 28, 2024 · d64fc7f · d64fc7f
1 parent cd45e5e
commit d64fc7f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/...trino-main/src/test/java/io/trino/sql/planner/assertions/AggregationFunctionProvider.java b/...trino-main/src/test/java/io/trino/sql/planner/assertions/AggregationFunctionProvider.java
@@ -18,14 +18,14 @@
 import com.google.common.collect.ImmutableMap;
 import io.trino.spi.connector.SortOrder;
 import io.trino.sql.ir.Expression;
+import io.trino.sql.ir.Reference;
 import io.trino.sql.planner.OrderingScheme;
 import io.trino.sql.planner.Symbol;
 
 import java.util.List;
 import java.util.Optional;
 
 import static io.trino.sql.planner.assertions.PlanMatchPattern.toSymbolReferences;
-import static io.trino.type.UnknownType.UNKNOWN;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -69,7 +69,8 @@ public AggregationFunction getExpectedValue(SymbolAliases aliases)
             ImmutableMap.Builder<Symbol, SortOrder> orders = ImmutableMap.builder();
 
             for (PlanMatchPattern.Ordering ordering : this.orderBy) {
-                Symbol symbol = new Symbol(UNKNOWN, aliases.get(ordering.getField()).name());
+                Reference reference = aliases.get(ordering.getField());
+                Symbol symbol = new Symbol(reference.type(), reference.name());
                 fields.add(symbol);
                 orders.put(symbol, ordering.getSortOrder());
             }

diff --git a/...main/src/test/java/io/trino/sql/planner/iterative/rule/TestPruneOrderByInAggregation.java b/...main/src/test/java/io/trino/sql/planner/iterative/rule/TestPruneOrderByInAggregation.java
@@ -38,7 +38,6 @@
 import static io.trino.sql.planner.plan.AggregationNode.Step.SINGLE;
 import static io.trino.sql.tree.SortItem.NullOrdering.LAST;
 import static io.trino.sql.tree.SortItem.Ordering.ASCENDING;
-import static io.trino.type.UnknownType.UNKNOWN;
 
 public class TestPruneOrderByInAggregation
         extends BaseRuleTest
@@ -81,16 +80,16 @@ private AggregationNode buildAggregation(PlanBuilder planBuilder)
                         "avg",
                         ImmutableList.of(new Reference(BIGINT, "input")),
                         new OrderingScheme(
-                                ImmutableList.of(new Symbol(UNKNOWN, "input")),
-                                ImmutableMap.of(new Symbol(UNKNOWN, "input"), SortOrder.ASC_NULLS_LAST))),
+                                ImmutableList.of(new Symbol(BIGINT, "input")),
+                                ImmutableMap.of(new Symbol(BIGINT, "input"), SortOrder.ASC_NULLS_LAST))),
                         ImmutableList.of(BIGINT),
                         mask)
                 .addAggregation(arrayAgg, PlanBuilder.aggregation(
                         "array_agg",
                         ImmutableList.of(new Reference(BIGINT, "input")),
                         new OrderingScheme(
-                                ImmutableList.of(new Symbol(UNKNOWN, "input")),
-                                ImmutableMap.of(new Symbol(UNKNOWN, "input"), SortOrder.ASC_NULLS_LAST))),
+                                ImmutableList.of(new Symbol(BIGINT, "input")),
+                                ImmutableMap.of(new Symbol(BIGINT, "input"), SortOrder.ASC_NULLS_LAST))),
                         ImmutableList.of(BIGINT),
                         mask)
                 .hashSymbol(keyHash)