Java Code Examples for org.apache.calcite.rex.RexProgram#create()

public void onMatch(RelOptRuleCall call) {
  final LogicalProject project = call.rel(0);
  final RelNode input = project.getInput();
  final RexProgram program =
      RexProgram.create(
          input.getRowType(),
          project.getProjects(),
          null,
          project.getRowType(),
          project.getCluster().getRexBuilder());
  final LogicalCalc calc = LogicalCalc.create(input, program);
  call.transformTo(calc);
}
 

@Nonnull
private Function<String, RelNode> convertProjectAsCalc() {
  return s -> {
    Project project = (Project) convertSql(s);
    RexProgram program = RexProgram.create(
        project.getInput().getRowType(),
        project.getProjects(),
        null,
        project.getRowType(),
        project.getCluster().getRexBuilder());
    return LogicalCalc.create(project.getInput(), program);
  };
}
 

@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);
      }
    }
 

public void onMatch(RelOptRuleCall call) {
  final EnumerableProject project = call.rel(0);
  final RelNode input = project.getInput();
  final RexProgram program =
      RexProgram.create(input.getRowType(),
          project.getProjects(),
          null,
          project.getRowType(),
          project.getCluster().getRexBuilder());
  final EnumerableCalc calc = EnumerableCalc.create(input, program);
  call.transformTo(calc);
}
 

public void onMatch(RelOptRuleCall call) {
  final LogicalProject project = call.rel(0);
  final RelNode input = project.getInput();
  final RexProgram program =
      RexProgram.create(
          input.getRowType(),
          project.getProjects(),
          null,
          project.getRowType(),
          project.getCluster().getRexBuilder());
  final LogicalCalc calc = LogicalCalc.create(input, program);
  call.transformTo(calc);
}
 

/**
 * Creates a LogicalCalc by parsing serialized output.
 */
public LogicalCalc(RelInput input) {
  this(input.getCluster(),
      input.getTraitSet(),
      ImmutableList.of(),
      input.getInput(),
      RexProgram.create(input));
}
 

private FlinkLogicalCalc createTopCalc(
	int primitiveLeftFieldCount,
	RexBuilder rexBuilder,
	ArrayBuffer<RexNode> extractedRexNodes,
	RelDataType calcRowType,
	FlinkLogicalCorrelate newCorrelate) {
	RexProgram rexProgram = new RexProgramBuilder(newCorrelate.getRowType(), rexBuilder).getProgram();
	int offset = extractedRexNodes.size() + primitiveLeftFieldCount;

	// extract correlate output RexNode.
	List<RexNode> newTopCalcProjects = rexProgram
		.getExprList()
		.stream()
		.filter(x -> x instanceof RexInputRef)
		.filter(x -> {
			int index = ((RexInputRef) x).getIndex();
			return index < primitiveLeftFieldCount || index >= offset;
		})
		.collect(Collectors.toList());

	return new FlinkLogicalCalc(
		newCorrelate.getCluster(),
		newCorrelate.getTraitSet(),
		newCorrelate,
		RexProgram.create(
			newCorrelate.getRowType(),
			newTopCalcProjects,
			null,
			calcRowType,
			rexBuilder));
}
 

private FlinkLogicalCalc createNewLeftCalc(
	RelNode left,
	RexBuilder rexBuilder,
	ArrayBuffer<RexNode> extractedRexNodes,
	FlinkLogicalCorrelate correlate) {
	// add the fields of the primitive left input.
	List<RexNode> leftCalcProjects = new LinkedList<>();
	RelDataType leftRowType = left.getRowType();
	List<String> leftCalcCalcFieldNames = createNewFieldNames(
		leftRowType,
		rexBuilder,
		leftRowType.getFieldCount(),
		extractedRexNodes,
		leftCalcProjects);

	// create a new calc
	return new FlinkLogicalCalc(
		correlate.getCluster(),
		correlate.getTraitSet(),
		left,
		RexProgram.create(
			leftRowType,
			leftCalcProjects,
			null,
			leftCalcCalcFieldNames,
			rexBuilder));
}
 

private FlinkLogicalCalc createTopCalc(
	int primitiveLeftFieldCount,
	RexBuilder rexBuilder,
	ArrayBuffer<RexNode> extractedRexNodes,
	RelDataType calcRowType,
	FlinkLogicalCorrelate newCorrelate) {
	RexProgram rexProgram = new RexProgramBuilder(newCorrelate.getRowType(), rexBuilder).getProgram();
	int offset = extractedRexNodes.size() + primitiveLeftFieldCount;

	// extract correlate output RexNode.
	List<RexNode> newTopCalcProjects = rexProgram
		.getExprList()
		.stream()
		.filter(x -> x instanceof RexInputRef)
		.filter(x -> {
			int index = ((RexInputRef) x).getIndex();
			return index < primitiveLeftFieldCount || index >= offset;
		})
		.collect(Collectors.toList());

	return new FlinkLogicalCalc(
		newCorrelate.getCluster(),
		newCorrelate.getTraitSet(),
		newCorrelate,
		RexProgram.create(
			newCorrelate.getRowType(),
			newTopCalcProjects,
			null,
			calcRowType,
			rexBuilder));
}
 

private FlinkLogicalCalc createNewLeftCalc(
	RelNode left,
	RexBuilder rexBuilder,
	ArrayBuffer<RexNode> extractedRexNodes,
	FlinkLogicalCorrelate correlate) {
	// add the fields of the primitive left input.
	List<RexNode> leftCalcProjects = new LinkedList<>();
	RelDataType leftRowType = left.getRowType();
	List<String> leftCalcCalcFieldNames = createNewFieldNames(
		leftRowType,
		rexBuilder,
		leftRowType.getFieldCount(),
		extractedRexNodes,
		leftCalcProjects);

	// create a new calc
	return new FlinkLogicalCalc(
		correlate.getCluster(),
		correlate.getTraitSet(),
		left,
		RexProgram.create(
			leftRowType,
			leftCalcProjects,
			null,
			leftCalcCalcFieldNames,
			rexBuilder));
}
 

@Override
public void onMatch(RelOptRuleCall call) {
	FlinkLogicalCorrelate correlate = call.rel(0);
	FlinkLogicalCalc right = call.rel(2);
	RexBuilder rexBuilder = call.builder().getRexBuilder();
	FlinkLogicalCalc mergedCalc = StreamExecCorrelateRule.getMergedCalc(right);
	FlinkLogicalTableFunctionScan tableScan = StreamExecCorrelateRule.getTableScan(mergedCalc);
	RexProgram mergedCalcProgram = mergedCalc.getProgram();

	InputRefRewriter inputRefRewriter = new InputRefRewriter(
		correlate.getRowType().getFieldCount() - mergedCalc.getRowType().getFieldCount());
	List<RexNode> correlateFilters = RelOptUtil
		.conjunctions(mergedCalcProgram.expandLocalRef(mergedCalcProgram.getCondition()))
		.stream()
		.map(x -> x.accept(inputRefRewriter))
		.collect(Collectors.toList());

	FlinkLogicalCorrelate newCorrelate = new FlinkLogicalCorrelate(
		correlate.getCluster(),
		correlate.getTraitSet(),
		correlate.getLeft(),
		tableScan,
		correlate.getCorrelationId(),
		correlate.getRequiredColumns(),
		correlate.getJoinType());

	RexNode topCalcCondition = RexUtil.composeConjunction(rexBuilder, correlateFilters);
	RexProgram rexProgram = new RexProgramBuilder(
		newCorrelate.getRowType(), rexBuilder).getProgram();
	FlinkLogicalCalc newTopCalc = new FlinkLogicalCalc(
		newCorrelate.getCluster(),
		newCorrelate.getTraitSet(),
		newCorrelate,
		RexProgram.create(
			newCorrelate.getRowType(),
			rexProgram.getExprList(),
			topCalcCondition,
			newCorrelate.getRowType(),
			rexBuilder));

	call.transformTo(newTopCalc);
}
 

public void onMatch(RelOptRuleCall call) {
  final LogicalProject project = call.rel(0);
  final LogicalCalc calc = call.rel(1);

  // Don't merge a project which contains windowed aggregates onto a
  // calc. That would effectively be pushing a windowed aggregate down
  // through a filter. Transform the project into an identical calc,
  // which we'll have chance to merge later, after the over is
  // expanded.
  final RelOptCluster cluster = project.getCluster();
  RexProgram program =
      RexProgram.create(
          calc.getRowType(),
          project.getProjects(),
          null,
          project.getRowType(),
          cluster.getRexBuilder());
  if (RexOver.containsOver(program)) {
    LogicalCalc projectAsCalc = LogicalCalc.create(calc, program);
    call.transformTo(projectAsCalc);
    return;
  }

  // Create a program containing the project node's expressions.
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  final RexProgramBuilder progBuilder =
      new RexProgramBuilder(
          calc.getRowType(),
          rexBuilder);
  for (Pair<RexNode, String> field : project.getNamedProjects()) {
    progBuilder.addProject(field.left, field.right);
  }
  RexProgram topProgram = progBuilder.getProgram();
  RexProgram bottomProgram = calc.getProgram();

  // Merge the programs together.
  RexProgram mergedProgram =
      RexProgramBuilder.mergePrograms(
          topProgram,
          bottomProgram,
          rexBuilder);
  final LogicalCalc newCalc =
      LogicalCalc.create(calc.getInput(), mergedProgram);
  call.transformTo(newCalc);
}
 

@Override protected PreparedResult implement(RelRoot root) {
  Hook.PLAN_BEFORE_IMPLEMENTATION.run(root);
  RelDataType resultType = root.rel.getRowType();
  boolean isDml = root.kind.belongsTo(SqlKind.DML);
  final Bindable bindable;
  if (resultConvention == BindableConvention.INSTANCE) {
    bindable = Interpreters.bindable(root.rel);
  } else {
    EnumerableRel enumerable = (EnumerableRel) root.rel;
    if (!root.isRefTrivial()) {
      final List<RexNode> projects = new ArrayList<>();
      final RexBuilder rexBuilder = enumerable.getCluster().getRexBuilder();
      for (int field : Pair.left(root.fields)) {
        projects.add(rexBuilder.makeInputRef(enumerable, field));
      }
      RexProgram program = RexProgram.create(enumerable.getRowType(),
          projects, null, root.validatedRowType, rexBuilder);
      enumerable = EnumerableCalc.create(enumerable, program);
    }

    try {
      CatalogReader.THREAD_LOCAL.set(catalogReader);
      final SqlConformance conformance = context.config().conformance();
      internalParameters.put("_conformance", conformance);
      bindable = EnumerableInterpretable.toBindable(internalParameters,
          context.spark(), enumerable, prefer);
    } finally {
      CatalogReader.THREAD_LOCAL.remove();
    }
  }

  if (timingTracer != null) {
    timingTracer.traceTime("end codegen");
  }

  if (timingTracer != null) {
    timingTracer.traceTime("end compilation");
  }

  return new PreparedResultImpl(
      resultType,
      parameterRowType,
      fieldOrigins,
      root.collation.getFieldCollations().isEmpty()
          ? ImmutableList.of()
          : ImmutableList.of(root.collation),
      root.rel,
      mapTableModOp(isDml, root.kind),
      isDml) {
    public String getCode() {
      throw new UnsupportedOperationException();
    }

    public Bindable getBindable(Meta.CursorFactory cursorFactory) {
      return bindable;
    }

    public Type getElementType() {
      return ((Typed) bindable).getElementType();
    }
  };
}
 

@Override
public void onMatch(RelOptRuleCall call) {
	FlinkLogicalCorrelate correlate = call.rel(0);
	FlinkLogicalCalc right = call.rel(2);
	RexBuilder rexBuilder = call.builder().getRexBuilder();
	FlinkLogicalCalc mergedCalc = CorrelateUtil.getMergedCalc(right);
	FlinkLogicalTableFunctionScan tableScan = CorrelateUtil.getTableFunctionScan(mergedCalc).get();
	RexProgram mergedCalcProgram = mergedCalc.getProgram();

	InputRefRewriter inputRefRewriter = new InputRefRewriter(
		correlate.getRowType().getFieldCount() - mergedCalc.getRowType().getFieldCount());
	List<RexNode> correlateFilters = RelOptUtil
		.conjunctions(mergedCalcProgram.expandLocalRef(mergedCalcProgram.getCondition()))
		.stream()
		.map(x -> x.accept(inputRefRewriter))
		.collect(Collectors.toList());

	FlinkLogicalCorrelate newCorrelate = new FlinkLogicalCorrelate(
		correlate.getCluster(),
		correlate.getTraitSet(),
		correlate.getLeft(),
		tableScan,
		correlate.getCorrelationId(),
		correlate.getRequiredColumns(),
		correlate.getJoinType());

	RexNode topCalcCondition = RexUtil.composeConjunction(rexBuilder, correlateFilters);
	RexProgram rexProgram = new RexProgramBuilder(
		newCorrelate.getRowType(), rexBuilder).getProgram();
	FlinkLogicalCalc newTopCalc = new FlinkLogicalCalc(
		newCorrelate.getCluster(),
		newCorrelate.getTraitSet(),
		newCorrelate,
		RexProgram.create(
			newCorrelate.getRowType(),
			rexProgram.getExprList(),
			topCalcCondition,
			newCorrelate.getRowType(),
			rexBuilder));

	call.transformTo(newTopCalc);
}
 

public UnifyResult apply(UnifyRuleCall call) {
  final MutableCalc query = (MutableCalc) call.query;
  final Pair<RexNode, List<RexNode>> queryExplained = explainCalc(query);
  final RexNode queryCond = queryExplained.left;
  final List<RexNode> queryProjs = queryExplained.right;

  final MutableCalc target = (MutableCalc) call.target;
  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();

  try {
    final RexShuttle shuttle = getRexShuttle(targetProjs);
    final RexNode splitted =
        splitFilter(call.getSimplify(), queryCond, targetCond);

    final RexNode compenCond;
    if (splitted != null) {
      if (splitted.isAlwaysTrue()) {
        compenCond = null;
      } else {
        // Compensate the residual filtering condition.
        compenCond = shuttle.apply(splitted);
      }
    } else if (implies(
        call.getCluster(), queryCond, targetCond, query.getInput().rowType)) {
      // Fail to split filtering condition, but implies that target contains
      // all lines of query, thus just set compensating filtering condition
      // as the filtering condition of query.
      compenCond = shuttle.apply(queryCond);
    } else {
      return null;
    }

    final List<RexNode> compenProjs = shuttle.apply(queryProjs);
    if (compenCond == null
        && RexUtil.isIdentity(compenProjs, target.rowType)) {
      return call.result(target);
    } else {
      final RexProgram compenRexProgram = RexProgram.create(
          target.rowType, compenProjs, compenCond,
          query.rowType, rexBuilder);
      final MutableCalc compenCalc = MutableCalc.of(target, compenRexProgram);
      return tryMergeParentCalcAndGenResult(call, compenCalc);
    }
  } catch (MatchFailed e) {
    return null;
  }
}
 

@Override protected UnifyResult apply(UnifyRuleCall call) {
  final MutableJoin query = (MutableJoin) call.query;
  final MutableRel qInput0 = query.getLeft();
  final MutableCalc qInput1 = (MutableCalc) query.getRight();
  final Pair<RexNode, List<RexNode>> qInput1Explained = explainCalc(qInput1);
  final RexNode qInput1Cond = qInput1Explained.left;
  final List<RexNode> qInput1Projs = qInput1Explained.right;

  final MutableJoin target = (MutableJoin) call.target;

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

  // Try pulling up MutableCalc only when:
  // 1. it's inner join.
  // 2. it's outer join but no filtering condition from MutableCalc.
  final JoinRelType joinRelType = sameJoinType(query.joinType, target.joinType);
  if (joinRelType == null) {
    return null;
  }
  if (joinRelType != JoinRelType.INNER
      && !(joinRelType.isOuterJoin() && qInput1Cond.isAlwaysTrue())) {
    return null;
  }
  // Try pulling up MutableCalc only when Join condition references mapping.
  final List<RexNode> identityProjects =
      (List<RexNode>) rexBuilder.identityProjects(qInput0.rowType);
  if (!referenceByMapping(query.condition, identityProjects, qInput1Projs)) {
    return null;
  }

  final RexNode newQueryJoinCond = new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef inputRef) {
      final int idx = inputRef.getIndex();
      if (idx < fieldCnt(qInput0)) {
        return inputRef;
      } else {
        final int newIdx = ((RexInputRef) qInput1Projs.get(idx - fieldCnt(qInput0)))
            .getIndex() + fieldCnt(qInput0);
        return new RexInputRef(newIdx, inputRef.getType());
      }
    }
  }.apply(query.condition);

  final RexNode splitted =
      splitFilter(call.getSimplify(), newQueryJoinCond, target.condition);
  // MutableJoin matches only when the conditions are analyzed to be same.
  if (splitted != null && splitted.isAlwaysTrue()) {
    final RexNode compenCond =
        RexUtil.shift(qInput1Cond, qInput0.rowType.getFieldCount());
    final List<RexNode> compenProjs = new ArrayList<>();
    for (int i = 0; i < query.rowType.getFieldCount(); i++) {
      if (i < fieldCnt(qInput0)) {
        compenProjs.add(
            new RexInputRef(i, query.rowType.getFieldList().get(i).getType()));
      } else {
        final RexNode shifted = RexUtil.shift(qInput1Projs.get(i - fieldCnt(qInput0)),
            qInput0.rowType.getFieldCount());
        compenProjs.add(shifted);
      }
    }
    final RexProgram compensatingRexProgram = RexProgram.create(
        target.rowType, compenProjs, compenCond,
        query.rowType, rexBuilder);
    final MutableCalc compenCalc = MutableCalc.of(target, compensatingRexProgram);
    return tryMergeParentCalcAndGenResult(call, compenCalc);
  }
  return null;
}
 

@Override protected UnifyResult apply(UnifyRuleCall call) {
  final MutableAggregate query = (MutableAggregate) call.query;
  final MutableCalc qInput = (MutableCalc) query.getInput();
  final Pair<RexNode, List<RexNode>> qInputExplained = explainCalc(qInput);
  final RexNode qInputCond = qInputExplained.left;
  final List<RexNode> qInputProjs = qInputExplained.right;

  final MutableAggregate target = (MutableAggregate) call.target;

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

  final Mappings.TargetMapping mapping =
      Project.getMapping(fieldCnt(qInput.getInput()), qInputProjs);
  if (mapping == null) {
    return null;
  }

  if (!qInputCond.isAlwaysTrue()) {
    try {
      // Fail the matching when filtering condition references
      // non-grouping columns in target.
      qInputCond.accept(new RexVisitorImpl<Void>(true) {
        @Override public Void visitInputRef(RexInputRef inputRef) {
          if (!target.groupSets.stream()
              .allMatch(groupSet -> groupSet.get(inputRef.getIndex()))) {
            throw Util.FoundOne.NULL;
          }
          return super.visitInputRef(inputRef);
        }
      });
    } catch (Util.FoundOne one) {
      return null;
    }
  }

  final Mapping inverseMapping = mapping.inverse();
  final MutableAggregate aggregate2 =
      permute(query, qInput.getInput(), inverseMapping);

  final Mappings.TargetMapping mappingForQueryCond = Mappings.target(
      target.groupSet::indexOf,
      target.getInput().rowType.getFieldCount(),
      target.groupSet.cardinality());
  final RexNode targetCond = RexUtil.apply(mappingForQueryCond, qInputCond);

  final MutableRel unifiedAggregate =
      unifyAggregates(aggregate2, targetCond, target);
  if (unifiedAggregate == null) {
    return null;
  }
  // Add Project if the mapping breaks order of fields in GroupSet
  if (!Mappings.keepsOrdering(mapping)) {
    final List<Integer> posList = new ArrayList<>();
    final int fieldCount = aggregate2.rowType.getFieldCount();
    final List<Pair<Integer, Integer>> pairs = new ArrayList<>();
    final List<Integer> groupings = aggregate2.groupSet.toList();
    for (int i = 0; i < groupings.size(); i++) {
      pairs.add(Pair.of(mapping.getTarget(groupings.get(i)), i));
    }
    Collections.sort(pairs);
    pairs.forEach(pair -> posList.add(pair.right));
    for (int i = posList.size(); i < fieldCount; i++) {
      posList.add(i);
    }
    final List<RexNode> compenProjs =
        MutableRels.createProjectExprs(unifiedAggregate, posList);
    final RexProgram compensatingRexProgram = RexProgram.create(
        unifiedAggregate.rowType, compenProjs, null,
        query.rowType, rexBuilder);
    final MutableCalc compenCalc =
        MutableCalc.of(unifiedAggregate, compensatingRexProgram);
    if (unifiedAggregate instanceof MutableCalc) {
      final MutableCalc newCompenCalc =
          mergeCalc(rexBuilder, compenCalc, (MutableCalc) unifiedAggregate);
      return tryMergeParentCalcAndGenResult(call, newCompenCalc);
    } else {
      return tryMergeParentCalcAndGenResult(call, compenCalc);
    }
  } else {
    return tryMergeParentCalcAndGenResult(call, unifiedAggregate);
  }
}
 

public void onMatch(RelOptRuleCall call) {
  final LogicalProject project = call.rel(0);
  final LogicalCalc calc = call.rel(1);

  // Don't merge a project which contains windowed aggregates onto a
  // calc. That would effectively be pushing a windowed aggregate down
  // through a filter. Transform the project into an identical calc,
  // which we'll have chance to merge later, after the over is
  // expanded.
  final RelOptCluster cluster = project.getCluster();
  RexProgram program =
      RexProgram.create(
          calc.getRowType(),
          project.getProjects(),
          null,
          project.getRowType(),
          cluster.getRexBuilder());
  if (RexOver.containsOver(program)) {
    LogicalCalc projectAsCalc = LogicalCalc.create(calc, program);
    call.transformTo(projectAsCalc);
    return;
  }

  // Create a program containing the project node's expressions.
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  final RexProgramBuilder progBuilder =
      new RexProgramBuilder(
          calc.getRowType(),
          rexBuilder);
  for (Pair<RexNode, String> field : project.getNamedProjects()) {
    progBuilder.addProject(field.left, field.right);
  }
  RexProgram topProgram = progBuilder.getProgram();
  RexProgram bottomProgram = calc.getProgram();

  // Merge the programs together.
  RexProgram mergedProgram =
      RexProgramBuilder.mergePrograms(
          topProgram,
          bottomProgram,
          rexBuilder);
  final LogicalCalc newCalc =
      LogicalCalc.create(calc.getInput(), mergedProgram);
  call.transformTo(newCalc);
}