rubenada commented on a change in pull request #1987:
URL: https://github.com/apache/calcite/pull/1987#discussion_r439257255



##########
File path: core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
##########
@@ -350,6 +354,100 @@ public TrimResult trimFields(
         Mappings.createIdentity(rel.getRowType().getFieldCount()));
   }
 
+  /**
+   * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
+   * {@link org.apache.calcite.rel.logical.LogicalCalc}.
+   */
+  public TrimResult trimFields(
+      Calc calc,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields) {
+    final RexProgram rexProgram = calc.getProgram();
+    final List<RexNode> projs = Lists.transform(rexProgram.getProjectList(),
+        rexProgram::expandLocalRef);
+
+    final RelDataType rowType = calc.getRowType();
+    final int fieldCount = rowType.getFieldCount();
+    final RelNode input = calc.getInput();
+
+    final Set<RelDataTypeField> inputExtraFields =
+        new LinkedHashSet<>(extraFields);
+    RelOptUtil.InputFinder inputFinder =
+        new RelOptUtil.InputFinder(inputExtraFields);
+    for (Ord<RexNode> ord : Ord.zip(projs)) {
+      if (fieldsUsed.get(ord.i)) {
+        ord.e.accept(inputFinder);
+      }
+    }
+    ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();
+
+    // Create input with trimmed columns.
+    TrimResult trimResult =
+        trimChild(calc, input, inputFieldsUsed, inputExtraFields);
+    RelNode newInput = trimResult.left;
+    final Mapping inputMapping = trimResult.right;
+
+    // If the input is unchanged, and we need to project all columns,
+    // there's nothing we can do.
+    if (newInput == input
+        && fieldsUsed.cardinality() == fieldCount) {
+      return result(calc, Mappings.createIdentity(fieldCount));
+    }
+
+    // Some parts of the system can't handle rows with zero fields, so
+    // pretend that one field is used.
+    if (fieldsUsed.cardinality() == 0) {
+      return dummyProject(fieldCount, newInput);
+    }
+
+    // Build new project expressions, and populate the mapping.
+    final List<RexNode> newProjects = new ArrayList<>();
+    final RexVisitor<RexNode> shuttle =
+        new RexPermuteInputsShuttle(
+            inputMapping, newInput);
+    final Mapping mapping =
+        Mappings.create(
+            MappingType.INVERSE_SURJECTION,
+            fieldCount,
+            fieldsUsed.cardinality());
+    for (Ord<RexNode> ord : Ord.zip(projs)) {
+      if (fieldsUsed.get(ord.i)) {
+        mapping.set(ord.i, newProjects.size());
+        RexNode newProjectExpr = ord.e.accept(shuttle);
+        newProjects.add(newProjectExpr);
+      }
+    }
+
+    final RelDataType newRowType =
+        RelOptUtil.permute(calc.getCluster().getTypeFactory(), rowType,
+            mapping);
+
+    final RelNode newInputRelNode = relBuilder.push(newInput).build();
+    if (rexProgram.getCondition() != null) {
+      final List<RexNode> filter = Lists.transform(
+          ImmutableList.of(
+          rexProgram.getCondition()), rexProgram::expandLocalRef);
+      assert filter.size() == 1;
+      final RexNode conditionExpr = filter.get(0);
+
+      final RexNode newConditionExpr =
+          conditionExpr.accept(shuttle);
+      final RexProgram newRexProgram = 
RexProgram.create(newInputRelNode.getRowType(),
+          newProjects, newConditionExpr, newRowType.getFieldNames(),
+          newInputRelNode.getCluster().getRexBuilder());
+      LogicalCalc logicalCalc = LogicalCalc.create(newInputRelNode, 
newRexProgram);
+
+      return result(logicalCalc, mapping);
+    } else {
+      final RexProgram newRexProgram = RexProgram
+          .create(newInputRelNode.getRowType(), newProjects, null,
+              newRowType.getFieldNames(), 
newInputRelNode.getCluster().getRexBuilder());
+
+      LogicalCalc logicalCalc = LogicalCalc.create(newInputRelNode, 
newRexProgram);
+      return result(logicalCalc, mapping);

Review comment:
       Hints need to be transferred from the original Calc to the new Calc (see 
CALCITE-4055).
   Also, IMHO this piece of code could be refactored, to avoid duplicating 
instructions at the end of the "if" and "else" block, something like:
   ```
   RexNode newConditionExpr = null;
   if (rexProgram.getCondition() != null) {
     ...
     newConditionExpr = conditionExpr.accept(shuttle);
   }
   final RexProgram newRexProgram = ...
   final LogicalCalc logicalCalc = LogicalCalc.create(newInputRelNode, 
newRexProgram);
   logicalCalc.withHints(calc.getHints()); // transfer hints
   return result(logicalCalc, mapping);
   ```




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to