fix bugs

Signed-off-by: shuming.li <[email protected]>
StarRocks · Feb 19, 2025 · 411ed70 · 411ed70
1 parent a201c8a
commit 411ed70
Showing 1 changed file with 11 additions and 23 deletions.
diff --git a/.../starrocks/sql/optimizer/rule/transformation/materialization/OptExpressionDuplicator.java b/.../starrocks/sql/optimizer/rule/transformation/materialization/OptExpressionDuplicator.java
@@ -209,8 +209,7 @@ public OptExpression visitLogicalTableScan(OptExpression optExpression, Void con
                 LogicalOlapScanOperator.Builder olapScanBuilder = (LogicalOlapScanOperator.Builder) scanBuilder;
                 if (olapScan.getDistributionSpec() instanceof HashDistributionSpec) {
                     HashDistributionSpec newHashDistributionSpec =
-                            processHashDistributionSpec((HashDistributionSpec) olapScan.getDistributionSpec(),
-                            newColumnMetaToColRefMap);
+                            processHashDistributionSpec((HashDistributionSpec) olapScan.getDistributionSpec());
                     olapScanBuilder.setDistributionSpec(newHashDistributionSpec);
                 }
 
@@ -659,13 +658,11 @@ public OptExpression visit(OptExpression optExpression, Void context) {
 
         // rewrite shuffle columns and joinEquivalentColumns in HashDistributionSpec
         // because the columns ids have changed
-        private HashDistributionSpec processHashDistributionSpec(
-                HashDistributionSpec originSpec,
-                Map<Column, ColumnRefOperator> newColumnMetaToColRefMap) {
+        private HashDistributionSpec processHashDistributionSpec(HashDistributionSpec originSpec) {
             // HashDistributionDesc
             final List<DistributionCol> newColumns = Lists.newArrayList();
             for (DistributionCol distributionCol : originSpec.getShuffleColumns()) {
-                final ColumnRefOperator newRefOperator = getNewDistributionColRef(distributionCol, newColumnMetaToColRefMap);
+                final ColumnRefOperator newRefOperator = getNewDistributionColRef(distributionCol);
                 Preconditions.checkNotNull(newRefOperator);
                 newColumns.add(new DistributionCol(newRefOperator.getId(), distributionCol.isNullStrict()));
             }
@@ -676,43 +673,34 @@ private HashDistributionSpec processHashDistributionSpec(
             final EquivalentDescriptor equivDesc = originSpec.getEquivDesc();
             final EquivalentDescriptor newEquivDesc = new EquivalentDescriptor(equivDesc.getTableId(),
                     equivDesc.getPartitionIds());
-            updateDistributionUnionFind(newEquivDesc.getNullRelaxUnionFind(), equivDesc.getNullStrictUnionFind(),
-                    newColumnMetaToColRefMap);
-            updateDistributionUnionFind(newEquivDesc.getNullStrictUnionFind(), equivDesc.getNullRelaxUnionFind(),
-                    newColumnMetaToColRefMap);
+            updateDistributionUnionFind(newEquivDesc.getNullRelaxUnionFind(), equivDesc.getNullStrictUnionFind());
+            updateDistributionUnionFind(newEquivDesc.getNullStrictUnionFind(), equivDesc.getNullRelaxUnionFind());
             return new HashDistributionSpec(hashDistributionDesc, newEquivDesc);
         }
 
         private void updateDistributionUnionFind(UnionFind<DistributionCol> newUnionFind,
-                                                UnionFind<DistributionCol> oldUnionFind,
-                                                 Map<Column, ColumnRefOperator> newColumnMetaToColRefMap) {
+                                                UnionFind<DistributionCol> oldUnionFind) {
             for (Set<DistributionCol> distributionColSet : oldUnionFind.getAllGroups()) {
                 DistributionCol first = null;
                 for (DistributionCol next : distributionColSet) {
                     if (first == null) {
                         first = next;
                     }
-                    final ColumnRefOperator newFirstCol = getNewDistributionColRef(first, newColumnMetaToColRefMap);
-                    final ColumnRefOperator newNextCol = getNewDistributionColRef(next, newColumnMetaToColRefMap);
+                    final ColumnRefOperator newFirstCol = getNewDistributionColRef(first);
+                    final ColumnRefOperator newNextCol = getNewDistributionColRef(next);
                     newUnionFind.union(first.updateColId(newFirstCol.getId()), next.updateColId(newNextCol.getId()));
                 }
             }
         }
 
-        private ColumnRefOperator getNewDistributionColRef(DistributionCol col,
-                                                           Map<Column, ColumnRefOperator> newColumnMetaToColRefMap) {
+        private ColumnRefOperator getNewDistributionColRef(DistributionCol col) {
             int colId = col.getColId();
             final ColumnRefOperator oldRefOperator = columnRefFactory.getColumnRef(colId);
             Preconditions.checkArgument(oldRefOperator != null);
             // use column mapping to find the new ColumnRefOperator
             ColumnRefOperator newRefOperator = columnMapping.get(oldRefOperator);
-            if (newRefOperator != null) {
-                return newRefOperator;
-            }
-            // use column to meta to find the new ColumnRefOperator if not found in column mapping
-            final Column column = columnRefFactory.getColumn(oldRefOperator);
-            Preconditions.checkArgument(column != null);
-            return newColumnMetaToColRefMap.get(column);
+            Preconditions.checkArgument(newRefOperator != null);
+            return newRefOperator;
         }
 
         private void processCommon(Operator.Builder opBuilder) {