Reuse createValueGenerator

Author: iwanttobepowerful

core/src/main/java/org/apache/calcite/sql2rel/RelDecorrelator.java

@@ -1781,98 +1781,92 @@ private static boolean isWidening(RelDataType type, RelDataType type1) {
     }
 
     RexNode cond = decorrelateExpr(castNonNull(currentRel), map, cm, rel.getCondition());
-    boolean contained = RexUtil.containsFieldAccess(cond);
-    int groupKeySize = 0;
-    RelNode join = null;
+    int valueGenFieldCount = 0;
+    RelNode newLeft = null;
     RexNode replacedCond = null;
-    final NavigableMap<CorDef, Integer> corDefOutputs = new TreeMap<>();
-
-    if (contained && isCorVarDefined) {
-      final Pair<CorrelationId, Frame> outerFramePair = requireNonNull(this.frameStack.peek());
-      final CorrelationId outerFrameCorrId = outerFramePair.left;
-      final Frame outerFrame = outerFramePair.right;
-
-      // Collect CorDef from decorrelated condition.
-      RexUtil.FieldAccessFinder finder = new RexUtil.FieldAccessFinder();
-      cond.accept(finder);
-      List<RexFieldAccess> correlatedKeyList = finder.getFieldAccessList();
-      ImmutableBitSet.Builder corFieldBuilder = ImmutableBitSet.builder();
-
-      for (RexFieldAccess rfa : correlatedKeyList) {
-        RexCorrelVariable correlVariable = (RexCorrelVariable) rfa.getReferenceExpr();
-        if (correlVariable.id.equals(outerFrameCorrId)) {
-          int idx = rfa.getField().getIndex();
-          int newIdx = requireNonNull(outerFrame.oldToNewOutputs.get(idx));
-          corFieldBuilder.set(newIdx);
-        }
-      }
-
-      ImmutableBitSet groupSet = corFieldBuilder.build();
-      groupKeySize = groupSet.cardinality();
-
-      // Build LogicalAggregate to obtain the distinct set of corVar from outerFrame.
-      relBuilder.push(outerFrame.r).aggregate(relBuilder.groupKey(groupSet));
-      join = relBuilder.push(leftFrame.r)
-          .join(JoinRelType.INNER, relBuilder.literal(true)).build();
+    final NavigableMap<CorDef, Integer> valueGenCorDefOutputs = new TreeMap<>();
+    final Pair<CorrelationId, Frame> outerFramePair = requireNonNull(this.frameStack.peek());
+    final CorrelationId outerFrameCorrId = outerFramePair.left;
+    final List<CorRef> corVarList = new ArrayList<>();
+    for (CorRef ref : cm.mapRefRelToCorRef.get(rel)) {
+      if (ref != null && ref.corr.equals(outerFrameCorrId)) {
+        corVarList.add(ref);
+      }
+    }
+    Collections.sort(corVarList);
+
+    if (!corVarList.isEmpty()) {
+      final RelNode valueGen = createValueGenerator(corVarList, 0, valueGenCorDefOutputs);
+      requireNonNull(valueGen, "valueGen");
+      valueGenFieldCount = valueGen.getRowType().getFieldCount();
+      newLeft = relBuilder
+          .push(valueGen)
+          .push(leftFrame.r)
+          .join(JoinRelType.INNER, relBuilder.literal(true))
+          .build();
 
       final Map<RexFieldAccess, RexInputRef> replacementMap = new HashMap<>();
-      for (RexFieldAccess rfa : correlatedKeyList) {
-        RexCorrelVariable correlVariable = (RexCorrelVariable) rfa.getReferenceExpr();
-        if (correlVariable.id.equals(outerFrameCorrId)) {
-          int idx = rfa.getField().getIndex();
-          int newIdx = requireNonNull(outerFrame.oldToNewOutputs.get(idx));
-          int pos = groupSet.indexOf(newIdx);
-          corDefOutputs.put(new CorDef(outerFrameCorrId, newIdx), pos);
-
-          RelDataType type = outerFrame.r.getRowType().getFieldList().get(newIdx).getType();
-          RexInputRef replace = new RexInputRef(pos, type);
-          replacementMap.put(rfa, replace);
+      RexUtil.FieldAccessFinder finder = new RexUtil.FieldAccessFinder();
+      rel.getCondition().accept(finder);
+      for (RexFieldAccess rfa : finder.getFieldAccessList()) {
+        final RexNode ref = rfa.getReferenceExpr();
+        if (ref instanceof RexCorrelVariable) {
+          final RexCorrelVariable corrVar = (RexCorrelVariable) ref;
+          if (corrVar.id.equals(outerFrameCorrId)) {
+            final CorDef def = new CorDef(corrVar.id, rfa.getField().getIndex());
+            final Integer pos = valueGenCorDefOutputs.get(def);
+            if (pos != null) {
+              final RelDataType type = valueGen.getRowType().getFieldList().get(pos).getType();
+              replacementMap.put(rfa, new RexInputRef(pos, type));
+            }
+          }
         }
       }
 
-      final RexNode condShifted = RexUtil.shift(cond, groupKeySize);
+      RexNode condShifted = RexUtil.shift(cond, valueGenFieldCount);
       replacedCond = RexUtil.replaceRexFieldAccessToInputRef(condShifted, replacementMap);
     }
 
     RelNode newJoin = relBuilder
-        .push(contained ? requireNonNull(join, "join") : leftFrame.r)
+        .push(!corVarList.isEmpty() ? requireNonNull(newLeft, "newLeft") : leftFrame.r)
         .push(rightFrame.r)
         .join(rel.getJoinType(),
-            contained ? requireNonNull(replacedCond, "replacedCond") : cond,
+            !corVarList.isEmpty() ? requireNonNull(replacedCond, "replacedCond") : cond,
             ImmutableSet.of())
         .build();
 
     // Create the mapping between the output of the old correlation rel
     // and the new join rel
     Map<Integer, Integer> mapOldToNewOutputs = new HashMap<>();
 
+    final NavigableMap<CorDef, Integer> corDefOutputs = new TreeMap<>();
+    corDefOutputs.putAll(valueGenCorDefOutputs);
+
     int oldLeftFieldCount = oldLeft.getRowType().getFieldCount();
-    int newLeftFieldCount = leftFrame.r.getRowType().getFieldCount() + groupKeySize;
+    int newLeftFieldCount = leftFrame.r.getRowType().getFieldCount() + valueGenFieldCount;
 
     int oldRightFieldCount = oldRight.getRowType().getFieldCount();
     //noinspection AssertWithSideEffects
     assert rel.getRowType().getFieldCount()
         == oldLeftFieldCount + oldRightFieldCount;
 
-    // Left input positions are shifted by groupKeySize.
+    // Left input positions are shifted by valueGenFieldCount.
     for (Map.Entry<Integer, Integer> entry : leftFrame.oldToNewOutputs.entrySet()) {
-      mapOldToNewOutputs.put(entry.getKey(), groupKeySize + entry.getValue());
+      mapOldToNewOutputs.put(entry.getKey(), valueGenFieldCount + entry.getValue());
     }
     // Right input positions are shifted by newLeftFieldCount.
     for (int i = 0; i < oldRightFieldCount; i++) {
       mapOldToNewOutputs.put(i + oldLeftFieldCount,
           requireNonNull(rightFrame.oldToNewOutputs.get(i)) + newLeftFieldCount);
     }
 
-    // Left input positions are shifted by groupKeySize.
+    // Left input positions are shifted by valueGenFieldCount.
     for (Map.Entry<CorDef, Integer> entry : leftFrame.corDefOutputs.entrySet()) {
-      corDefOutputs.put(entry.getKey(), groupKeySize + entry.getValue());
+      corDefOutputs.put(entry.getKey(), valueGenFieldCount + entry.getValue());
     }
     // Right input positions are shifted by newLeftFieldCount.
-    for (Map.Entry<CorDef, Integer> entry
-        : rightFrame.corDefOutputs.entrySet()) {
-      corDefOutputs.put(entry.getKey(),
-          entry.getValue() + newLeftFieldCount);
+    for (Map.Entry<CorDef, Integer> entry : rightFrame.corDefOutputs.entrySet()) {
+      corDefOutputs.put(entry.getKey(), entry.getValue() + newLeftFieldCount);
     }
     return register(rel, newJoin, mapOldToNewOutputs, corDefOutputs);
   }