org.apache.calcite.rex.RexShuttle Java Examples

public RelNode accept(RexShuttle shuttle) {
  RexNode joinFilter = shuttle.apply(this.joinFilter);
  List<RexNode> outerJoinConditions = shuttle.apply(this.outerJoinConditions);
  RexNode postJoinFilter = shuttle.apply(this.postJoinFilter);

  if (joinFilter == this.joinFilter
      && outerJoinConditions == this.outerJoinConditions
      && postJoinFilter == this.postJoinFilter) {
    return this;
  }

  return new MultiJoin(
      getCluster(),
      inputs,
      joinFilter,
      rowType,
      isFullOuterJoin,
      outerJoinConditions,
      joinTypes,
      projFields,
      joinFieldRefCountsMap,
      postJoinFilter);
}
 

private static Pair<List<RexNode>, List<RexNode>> splitProjects(
    final RexBuilder rexBuilder, final RelNode input, List<RexNode> nodes) {
  final RelOptUtil.InputReferencedVisitor visitor =
      new RelOptUtil.InputReferencedVisitor();
  visitor.visitEach(nodes);
  if (visitor.inputPosReferenced.size() == input.getRowType().getFieldCount()) {
    // All inputs are referenced
    return null;
  }
  final List<RexNode> belowNodes = new ArrayList<>();
  final List<RelDataType> belowTypes = new ArrayList<>();
  final List<Integer> positions = Lists.newArrayList(visitor.inputPosReferenced);
  for (int i : positions) {
    final RexNode node = rexBuilder.makeInputRef(input, i);
    belowNodes.add(node);
    belowTypes.add(node.getType());
  }
  final List<RexNode> aboveNodes = new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef ref) {
      final int index = positions.indexOf(ref.getIndex());
      return rexBuilder.makeInputRef(belowTypes.get(index), index);
    }
  }.visitList(nodes);
  return Pair.of(aboveNodes, belowNodes);
}
 

public Set<RelColumnOrigin> getColumnOrigins(Calc rel,
    final RelMetadataQuery mq, int iOutputColumn) {
  final RelNode input = rel.getInput();
  final RexShuttle rexShuttle = new RexShuttle() {
    @Override public RexNode visitLocalRef(RexLocalRef localRef) {
      return rel.getProgram().expandLocalRef(localRef);
    }
  };
  final List<RexNode> projects = new ArrayList<>();
  for (RexNode rex: rexShuttle.apply(rel.getProgram().getProjectList())) {
    projects.add(rex);
  }
  final RexNode rexNode = projects.get(iOutputColumn);
  if (rexNode instanceof RexInputRef) {
    // Direct reference:  no derivation added.
    RexInputRef inputRef = (RexInputRef) rexNode;
    return mq.getColumnOrigins(input, inputRef.getIndex());
  }
  // Anything else is a derivation, possibly from multiple columns.
  final Set<RelColumnOrigin> set = getMultipleColumns(rexNode, input, mq);
  return createDerivedColumnOrigins(set);
}
 

/**
 * Replaces all the input references by the position in the
 * input column set. If a reference index cannot be found in
 * the input set, then we return null.
 */
protected RexNode shuttleReferences(final RexBuilder rexBuilder,
    final RexNode node, final Mapping mapping) {
  try {
    RexShuttle visitor =
        new RexShuttle() {
          @Override public RexNode visitInputRef(RexInputRef inputRef) {
            int pos = mapping.getTargetOpt(inputRef.getIndex());
            if (pos != -1) {
              // Found it
              return rexBuilder.makeInputRef(inputRef.getType(), pos);
            }
            throw Util.FoundOne.NULL;
          }
        };
    return visitor.apply(node);
  } catch (Util.FoundOne ex) {
    Util.swallow(ex, null);
    return null;
  }
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  FilterPrel filter = call.rel(0);
  NestedLoopJoinPrel join = call.rel(1);

  if ((join.getProjectedFields() == null) || join.getProjectedFields().cardinality() == join.getInputRowType().getFieldCount()) {
    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), filter.getCondition()), join.getProjectedFields()));
  } else {
    // Current filter condition is written based on projected fields on join. In order to push this filter down we need to rewrite filter condition
    final ImmutableBitSet topProjectedColumns = RelOptUtil.InputFinder.bits(filter.getCondition());
    final ImmutableBitSet bottomProjectedColumns = join.getProjectedFields();

    Mapping mapping = Mappings.create(MappingType.SURJECTION, join.getRowType().getFieldCount(), join.getInputRowType().getFieldCount());
    for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
      if (topProjectedColumns.get(ord.i)) {
        mapping.set(ord.i, ord.e);
      }
    }

    RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
    RexNode updatedCondition = shuttle.apply(filter.getCondition());

    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), updatedCondition), join.getProjectedFields()));
  }
}
 

@Override
public RelNode visit(final LogicalFilter filter) {
  final RelBuilder relBuilder = newCalciteRelBuilderWithoutContext(filter.getCluster());
  RelNode input = filter.getInput().accept(this);
  relBuilder.push(input);

  RexNode newCondition = filter.getCondition().accept(new RexShuttle() {
    @Override
    public RexNode visitInputRef(RexInputRef inputRef) {
      return relBuilder.field(filter.getRowType().getFieldNames().get(inputRef.getIndex()));
    }
  });

  relBuilder.filter(newCondition);
  return relBuilder.build();
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode joinFilter = shuttle.apply(this.joinFilter);
  List<RexNode> outerJoinConditions = shuttle.apply(this.outerJoinConditions);
  RexNode postJoinFilter = shuttle.apply(this.postJoinFilter);

  if (joinFilter == this.joinFilter
      && outerJoinConditions == this.outerJoinConditions
      && postJoinFilter == this.postJoinFilter) {
    return this;
  }

  return new MultiJoin(
      getCluster(),
      inputs,
      joinFilter,
      rowType,
      isFullOuterJoin,
      outerJoinConditions,
      joinTypes,
      projFields,
      joinFieldRefCountsMap,
      postJoinFilter);
}
 

@Override
public UnifyResult apply(UnifyRuleCall call) {
    final MutableProject target = (MutableProject) call.target;
    final MutableScan query = (MutableScan) call.query;
    // We do not need to check query's parent type to avoid duplication
    // of ProjectToProjectUnifyRule or FilterToProjectUnifyRule, since
    // SubstitutionVisitor performs a top-down match.
    if (!query.equals(target.getInput())) {
        return null;
    }
    final RexShuttle shuttle = getRexShuttle(target);
    final RexBuilder rexBuilder = target.cluster.getRexBuilder();
    final List<RexNode> newProjects;
    try {
        newProjects = (List<RexNode>) shuttle.apply(rexBuilder.identityProjects(query.rowType));
    } catch (MatchFailed e) {
        return null;
    }
    final MutableProject newProject = MutableProject.of(query.rowType, target, newProjects);
    final MutableRel newProject2 = MutableRels.strip(newProject);
    return call.result(newProject2);
}
 

private static List<RexNode> transformRex(List<RexNode> nodes, final List<RelDataTypeField> oldFields,
        final List<RelDataTypeField> newFields) {
    RexShuttle shuttle = new RexShuttle() {
        @Override
        public RexNode visitInputRef(RexInputRef ref) {
            RelDataTypeField f = oldFields.get(ref.getIndex());
            for (int index = 0; index < newFields.size(); index++) {
                RelDataTypeField newf = newFields.get(index);
                if (f.getKey().equals(newf.getKey()) && f.getValue() == newf.getValue()) {
                    return RexBuilder.getRexFactory().makeInputRef(index, f.getValue());
                }
            }
            throw MatchFailed.INSTANCE;
        }
    };
    return shuttle.apply(nodes);
}
 

@Override
public void onMatch(RelOptRuleCall call) {
    LogicalCalc calc = call.rel(0);

    // Expand decimals in every expression in this program. If no
    // expression changes, don't apply the rule.
    final RexProgram program = calc.getProgram();
    if (!RexUtil.requiresDecimalExpansion(program, true)) {
        return;
    }

    final RexBuilder rexBuilder = calc.getCluster().getRexBuilder();
    final RexShuttle shuttle = new DecimalShuttle(rexBuilder);
    RexProgramBuilder programBuilder = RexProgramBuilder.create(rexBuilder, calc.getInput().getRowType(),
            program.getExprList(), program.getProjectList(), program.getCondition(), program.getOutputRowType(),
            shuttle, true);

    final RexProgram newProgram = programBuilder.getProgram();
    LogicalCalc newCalc = LogicalCalc.create(calc.getInput(), newProgram);
    call.transformTo(newCalc);
}
 

/**
 * Replaces all the input references by the position in the
 * input column set. If a reference index cannot be found in
 * the input set, then we return null.
 */
private static RexNode shuttleReferences(final RexBuilder rexBuilder, final RexNode node, final Mapping mapping) {
    try {
        RexShuttle visitor = new RexShuttle() {
            @Override
            public RexNode visitInputRef(RexInputRef inputRef) {
                int pos = mapping.getTargetOpt(inputRef.getIndex());
                if (pos != -1) {
                    // Found it
                    return rexBuilder.makeInputRef(inputRef.getType(), pos);
                }
                throw Util.FoundOne.NULL;
            }
        };
        return visitor.apply(node);
    } catch (Util.FoundOne ex) {
        Util.swallow(ex, null);
        return null;
    }
}
 

public RelNode accept(RexShuttle shuttle) {
  List<RexNode> exps = shuttle.apply(this.exps);
  if (this.exps == exps) {
    return this;
  }
  final RelDataType rowType =
      RexUtil.createStructType(
          getInput().getCluster().getTypeFactory(),
          exps,
          this.rowType.getFieldNames(),
          null);
  return copy(traitSet, getInput(), exps, rowType);
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode condition = shuttle.apply(this.period);
  if (this.period == condition) {
    return this;
  }
  return copy(traitSet, getInput(), condition);
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode offset = shuttle.apply(this.offset);
  RexNode fetch = shuttle.apply(this.fetch);
  List<RexNode> fieldExps = shuttle.apply(this.fieldExps);
  assert fieldExps == this.fieldExps
      : "Sort node does not support modification of input field expressions."
        + " Old expressions: " + this.fieldExps + ", new ones: " + fieldExps;
  if (offset == this.offset
      && fetch == this.fetch) {
    return this;
  }
  return copy(traitSet, getInput(), collation, offset, fetch);
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode condition = shuttle.apply(this.condition);
  if (this.condition == condition) {
    return this;
  }
  return copy(traitSet, getInput(), condition);
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode rexCall = shuttle.apply(this.rexCall);
  if (rexCall == this.rexCall) {
    return this;
  }
  return copy(traitSet, inputs, rexCall, elementType, rowType,
      columnMappings);
}
 

@Override public RelNode accept(RexShuttle shuttle) {
  RexNode condition = shuttle.apply(this.condition);
  if (this.condition == condition) {
    return this;
  }
  return copy(traitSet, condition, left, right, joinType, isSemiJoinDone());
}
 

private static RexShuttle pushShuttle(final Project project) {
  return new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef ref) {
      return project.getProjects().get(ref.getIndex());
    }
  };
}
 

public RelNode accept(RexShuttle shuttle) {
  List<RexNode> oldExprs = program.getExprList();
  List<RexNode> exprs = shuttle.apply(oldExprs);
  List<RexLocalRef> oldProjects = program.getProjectList();
  List<RexLocalRef> projects = shuttle.apply(oldProjects);
  RexLocalRef oldCondition = program.getCondition();
  RexNode condition;
  if (oldCondition != null) {
    condition = shuttle.apply(oldCondition);
    assert condition instanceof RexLocalRef
        : "Invalid condition after rewrite. Expected RexLocalRef, got "
        + condition;
  } else {
    condition = null;
  }
  if (exprs == oldExprs
      && projects == oldProjects
      && condition == oldCondition) {
    return this;
  }

  final RexBuilder rexBuilder = getCluster().getRexBuilder();
  final RelDataType rowType =
      RexUtil.createStructType(
          rexBuilder.getTypeFactory(),
          projects,
          this.rowType.getFieldNames(),
          null);
  final RexProgram newProgram =
      RexProgramBuilder.create(
          rexBuilder, program.getInputRowType(), exprs, projects,
          condition, rowType, true, null)
      .getProgram(false);
  return copy(traitSet, getInput(), newProgram);
}
 

public void onMatch(RelOptRuleCall call) {
  LogicalCalc calc = call.rel(0);

  // Expand decimals in every expression in this program. If no
  // expression changes, don't apply the rule.
  final RexProgram program = calc.getProgram();
  if (!RexUtil.requiresDecimalExpansion(program, true)) {
    return;
  }

  final RexBuilder rexBuilder = calc.getCluster().getRexBuilder();
  final RexShuttle shuttle = new DecimalShuttle(rexBuilder);
  RexProgramBuilder programBuilder =
      RexProgramBuilder.create(
          rexBuilder,
          calc.getInput().getRowType(),
          program.getExprList(),
          program.getProjectList(),
          program.getCondition(),
          program.getOutputRowType(),
          shuttle,
          true);

  final RexProgram newProgram = programBuilder.getProgram();
  LogicalCalc newCalc = LogicalCalc.create(calc.getInput(), newProgram);
  call.transformTo(newCalc);
}
 

@Override protected UnifyResult apply(UnifyRuleCall call) {

      final MutableScan query = (MutableScan) call.query;

      final MutableCalc target = (MutableCalc) call.target;
      final MutableScan targetInput = (MutableScan) target.getInput();
      final Pair<RexNode, List<RexNode>> targetExplained = explainCalc(target);
      final RexNode targetCond = targetExplained.left;
      final List<RexNode> targetProjs = targetExplained.right;

      final RexBuilder rexBuilder = call.getCluster().getRexBuilder();

      if (!query.equals(targetInput) || !targetCond.isAlwaysTrue()) {
        return null;
      }
      final RexShuttle shuttle = getRexShuttle(targetProjs);
      final List<RexNode> compenProjs;
      try {
        compenProjs = (List<RexNode>) shuttle.apply(
            rexBuilder.identityProjects(query.rowType));
      } catch (MatchFailed e) {
        return null;
      }
      if (RexUtil.isIdentity(compenProjs, target.rowType)) {
        return call.result(target);
      } else {
        RexProgram compenRexProgram = RexProgram.create(
            target.rowType, compenProjs, null, query.rowType, rexBuilder);
        MutableCalc compenCalc = MutableCalc.of(target, compenRexProgram);
        return tryMergeParentCalcAndGenResult(call, compenCalc);
      }
    }
 

/** Explain filtering condition and projections from MutableCalc. */
private static Pair<RexNode, List<RexNode>> explainCalc(MutableCalc calc) {
  final RexShuttle shuttle = getExpandShuttle(calc.program);
  final RexNode condition = shuttle.apply(calc.program.getCondition());
  final List<RexNode> projects = new ArrayList<>();
  for (RexNode rex: shuttle.apply(calc.program.getProjectList())) {
    projects.add(rex);
  }
  if (condition == null) {
    return Pair.of(calc.cluster.getRexBuilder().makeLiteral(true), projects);
  } else {
    return Pair.of(condition, projects);
  }
}
 

private static RexShuttle getExpandShuttle(RexProgram rexProgram) {
  return new RexShuttle() {
    @Override public RexNode visitLocalRef(RexLocalRef localRef) {
      return rexProgram.expandLocalRef(localRef);
    }
  };
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-3165">[CALCITE-3165]
 * Project#accept(RexShuttle shuttle) does not update rowType</a>. */
@Test void testProjectUpdatesRowType() {
  final RelBuilder builder = RelBuilder.create(RelBuilderTest.config().build());

  // Equivalent SQL: SELECT deptno, sal FROM emp
  final RelNode root =
      builder
          .scan("EMP")
          .project(
              builder.field("DEPTNO"),
              builder.field("SAL"))
          .build();

  // Equivalent SQL: SELECT CAST(deptno AS VARCHAR), CAST(sal AS VARCHAR) FROM emp
  final RelNode rootWithCast =
      builder
          .scan("EMP")
          .project(
              builder.cast(builder.field("DEPTNO"), SqlTypeName.VARCHAR),
              builder.cast(builder.field("SAL"), SqlTypeName.VARCHAR))
          .build();
  final RelDataType type = rootWithCast.getRowType();

  // Transform the first expression into the second one, by using a RexShuttle
  // that converts every RexInputRef into a 'CAST(RexInputRef AS VARCHAR)'
  final RelNode rootWithCastViaRexShuttle = root.accept(new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef inputRef) {
      return  builder.cast(inputRef, SqlTypeName.VARCHAR);
    }
  });
  final RelDataType type2 = rootWithCastViaRexShuttle.getRowType();

  assertThat(type, is(type2));
}
 

private static RexShuttle pushShuttle(final Project project) {
  return new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef ref) {
      return project.getProjects().get(ref.getIndex());
    }
  };
}
 

public RelNode accept(RexShuttle shuttle) {
  List<RexNode> oldExprs = program.getExprList();
  List<RexNode> exprs = shuttle.apply(oldExprs);
  List<RexLocalRef> oldProjects = program.getProjectList();
  List<RexLocalRef> projects = shuttle.apply(oldProjects);
  RexLocalRef oldCondition = program.getCondition();
  RexNode condition;
  if (oldCondition != null) {
    condition = shuttle.apply(oldCondition);
    assert condition instanceof RexLocalRef
        : "Invalid condition after rewrite. Expected RexLocalRef, got "
        + condition;
  } else {
    condition = null;
  }
  if (exprs == oldExprs
      && projects == oldProjects
      && condition == oldCondition) {
    return this;
  }
  return copy(traitSet, getInput(),
      new RexProgram(program.getInputRowType(),
          exprs,
          projects,
          (RexLocalRef) condition,
          program.getOutputRowType()));
}
 

@Override public RelNode accept(RexShuttle shuttle) {
  RexNode condition = shuttle.apply(this.condition);
  if (this.condition == condition) {
    return this;
  }
  return copy(traitSet, condition, left, right, joinType, isSemiJoinDone());
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode condition = shuttle.apply(this.condition);
  if (this.condition == condition) {
    return this;
  }
  return copy(traitSet, getInput(), condition);
}
 

@Override
public RelNode accept(RexShuttle shuttle) {
  RexNode ref = shuttle.apply(this.ref);
  if (this.ref == ref) {
    return this;
  }
  return new UnnestPrel(getCluster(), traitSet, rowType, ref);
}
 

public RelNode accept(RexShuttle shuttle) {
  RexNode offset = shuttle.apply(this.offset);
  RexNode fetch = shuttle.apply(this.fetch);
  List<RexNode> fieldExps = shuttle.apply(this.fieldExps);
  assert fieldExps == this.fieldExps
      : "Sort node does not support modification of input field expressions."
        + " Old expressions: " + this.fieldExps + ", new ones: " + fieldExps;
  if (offset == this.offset
      && fetch == this.fetch) {
    return this;
  }
  return copy(traitSet, getInput(), collation, offset, fetch);
}