Java Code Examples for org.apache.calcite.rex.RexBuilder#makeInputRef()

/**
 * 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;
    }
}
 

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

@Override
public RexNode visitCall(RexCall call) {

    String functionName = call.getOperator().getName();

    List<RexNode> newOps = new ArrayList<>();
    for (RexNode operand : call.getOperands()) {
        newOps.add(operand.accept(this));
    }
    if (funcReg.isFunctionComplexOutput(functionName)) {
        RexBuilder builder = new RexBuilder(factory);
        RexNode ret = builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory),
                lastUsedIndex);
        lastUsedIndex++;
        complexExprs.add(call.clone(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), newOps));
        return ret;
    }
    return call.clone(call.getType(), newOps);
}
 

private Prel visit(ExchangePrel hashPrel, List<DistributionTrait.DistributionField> fields, Prel child) {
  final List<String> childFields = child.getRowType().getFieldNames();


  // Insert Project SqlOperatorImpl with new column that will be a hash for HashToRandomExchange fields
  final ProjectPrel addColumnprojectPrel = HashPrelUtil.addHashProject(fields, child, null);
  final Prel newPrel = (Prel) hashPrel.copy(addColumnprojectPrel.getTraitSet(), Collections.<RelNode>singletonList(addColumnprojectPrel));

  int validRows = newPrel.getRowType().getFieldCount() - 1;
  final List<RelDataTypeField> all = newPrel.getRowType().getFieldList();
  final List<RexNode> keptExprs = new ArrayList<>(validRows);

  final RexBuilder rexBuilder = newPrel.getCluster().getRexBuilder();
  for(int i = 0; i < validRows; i++){
    RexNode rex = rexBuilder.makeInputRef(all.get(i).getType(), i);
    keptExprs.add(rex);
  }

  // remove earlier inserted Project SqlOperatorImpl - since it creates issues down the road in HashJoin
  RelDataType removeRowType = RexUtil.createStructType(newPrel.getCluster().getTypeFactory(), keptExprs, childFields);
  return ProjectPrel.create(newPrel.getCluster(), newPrel.getTraitSet(), newPrel, keptExprs, removeRowType);
}
 

@Override
public RelNode visit(TableScan tableScan) {
  if (!(tableScan instanceof IncrementallyUpdateable)) {
    return tableScan;
  }

  final RelNode newScan = updateScan((IncrementallyUpdateable) tableScan);

  // build new filter to apply refresh condition.
  final RexBuilder rexBuilder = tableScan.getCluster().getRexBuilder();
  RelDataTypeField field = newScan.getRowType().getField(UPDATE_COLUMN, false, false);
  final RexNode inputRef = rexBuilder.makeInputRef(newScan, field.getIndex());
  final Optional<RexNode> literal = generateLiteral(rexBuilder, tableScan.getCluster().getTypeFactory(), field.getType().getSqlTypeName());
  if (literal.isPresent()) {
    RexNode condition = tableScan.getCluster().getRexBuilder().makeCall(SqlStdOperatorTable.GREATER_THAN, ImmutableList.of(inputRef, literal.get()));
    return LogicalFilter.create(newScan, condition);
  }
  return newScan;
}
 

@Override public RexNode singleton(RexBuilder rexBuilder,
    RelDataType inputRowType, AggregateCall aggregateCall) {
  final int arg = aggregateCall.getArgList().get(0);
  final RelDataType type = inputRowType.getFieldList().get(arg).getType();
  final RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
  final RelDataType type1 = typeFactory.getTypeSystem().deriveSumType(typeFactory, type);
  final RexNode inputRef = rexBuilder.makeInputRef(type1, arg);
  if (type.isNullable()) {
    return rexBuilder.makeCall(SqlStdOperatorTable.COALESCE, inputRef,
        rexBuilder.makeExactLiteral(BigDecimal.ZERO, type));
  } else {
    return inputRef;
  }
}
 

@Override
public RexNode visitInputRef(RexInputRef inputRef) {
    final RelDataType leftRowType = left.getRowType();
    final RexBuilder rexBuilder = getRexBuilder();
    final int leftCount = leftRowType.getFieldCount();
    if (inputRef.getIndex() < leftCount) {
        final RexNode v = rexBuilder.makeCorrel(leftRowType, id);
        return rexBuilder.makeFieldAccess(v, inputRef.getIndex());
    } else {
        return rexBuilder.makeInputRef(right, inputRef.getIndex() - leftCount);
    }
}
 

@Override
public Prel visitProject(ProjectPrel node, Object unused) throws RelConversionException {
  ProjectPrel project = node;
  List<RexNode> exprList = new ArrayList<>();
  boolean rewrite = false;

  List<RelDataTypeField> relDataTypes = new ArrayList<>();
  int i = 0;
  RexNode flatttenExpr = null;
  for (RexNode rex : project.getChildExps()) {
    RexNode newExpr = rex;
    if (rex instanceof RexCall) {
      RexCall function = (RexCall) rex;
      String functionName = function.getOperator().getName();

      if (functionName.equalsIgnoreCase("flatten") ) {
        rewrite = true;
        if (function.getOperands().size() != 1) {
          throw new RelConversionException("Flatten expression expects a single input.");
        }
        newExpr = function.getOperands().get(0);
        RexBuilder builder = new RexBuilder(factory);
        flatttenExpr = builder.makeInputRef( new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), i);
      }
    }
    relDataTypes.add(project.getRowType().getFieldList().get(i));
    i++;
    exprList.add(newExpr);
  }
  if (rewrite == true) {
    // TODO - figure out what is the right setting for the traits
    Prel newChild = ((Prel)project.getInput(0)).accept(this, null);
    ProjectPrel newProject = new ProjectPrel(node.getCluster(), project.getTraitSet(), newChild, exprList, new RelRecordType(relDataTypes));
    FlattenPrel flatten = new FlattenPrel(project.getCluster(), project.getTraitSet(), newProject, flatttenExpr);
    return flatten;
  }

  Prel child = ((Prel)project.getInput()).accept(this, null);
  return (Prel) project.copy(project.getTraitSet(), child, exprList, new RelRecordType(relDataTypes));
}
 

public RexNode singleton(RexBuilder rexBuilder,
    RelDataType inputRowType, AggregateCall aggregateCall) {
  final int arg = aggregateCall.getArgList().get(0);
  final RelDataTypeField field = inputRowType.getFieldList().get(arg);
  final RelDataType fieldType = field.getType();
  final RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
  RelDataType type = typeFactory.getTypeSystem().deriveSumType(typeFactory, fieldType);
  return rexBuilder.makeInputRef(type, arg);
}
 

public RexNode visitInputRef(RexInputRef inputRef) {
  final RelDataType leftRowType = left.getRowType();
  final RexBuilder rexBuilder = getRexBuilder();
  final int leftCount = leftRowType.getFieldCount();
  if (inputRef.getIndex() < leftCount) {
    final RexNode v = rexBuilder.makeCorrel(leftRowType, id);
    return rexBuilder.makeFieldAccess(v, inputRef.getIndex());
  } else {
    return rexBuilder.makeInputRef(right, inputRef.getIndex() - leftCount);
  }
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Union union = call.rel(0);

  final RexBuilder rexBuilder = union.getCluster().getRexBuilder();
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(union);
  if (predicates == null) {
    return;
  }

  final Map<Integer, RexNode> constants = new HashMap<>();
  for (Map.Entry<RexNode, RexNode> e : predicates.constantMap.entrySet()) {
    if (e.getKey() instanceof RexInputRef) {
      constants.put(((RexInputRef) e.getKey()).getIndex(), e.getValue());
    }
  }

  // None of the expressions are constant. Nothing to do.
  if (constants.isEmpty()) {
    return;
  }

  // Create expressions for Project operators before and after the Union
  List<RelDataTypeField> fields = union.getRowType().getFieldList();
  List<RexNode> topChildExprs = new ArrayList<>();
  List<String> topChildExprsFields = new ArrayList<>();
  List<RexNode> refs = new ArrayList<>();
  ImmutableBitSet.Builder refsIndexBuilder = ImmutableBitSet.builder();
  for (RelDataTypeField field : fields) {
    final RexNode constant = constants.get(field.getIndex());
    if (constant != null) {
      topChildExprs.add(constant);
      topChildExprsFields.add(field.getName());
    } else {
      final RexNode expr = rexBuilder.makeInputRef(union, field.getIndex());
      topChildExprs.add(expr);
      topChildExprsFields.add(field.getName());
      refs.add(expr);
      refsIndexBuilder.set(field.getIndex());
    }
  }
  ImmutableBitSet refsIndex = refsIndexBuilder.build();

  // Update top Project positions
  final Mappings.TargetMapping mapping =
      RelOptUtil.permutation(refs, union.getInput(0).getRowType()).inverse();
  topChildExprs = RexUtil.apply(mapping, topChildExprs);

  // Create new Project-Union-Project sequences
  final RelBuilder relBuilder = call.builder();
  for (RelNode input : union.getInputs()) {
    List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
    for (int j : refsIndex) {
      newChildExprs.add(
          Pair.of(rexBuilder.makeInputRef(input, j),
              input.getRowType().getFieldList().get(j).getName()));
    }
    if (newChildExprs.isEmpty()) {
      // At least a single item in project is required.
      newChildExprs.add(
          Pair.of(topChildExprs.get(0), topChildExprsFields.get(0)));
    }
    // Add the input with project on top
    relBuilder.push(input);
    relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
  }
  relBuilder.union(union.all, union.getInputs().size());
  // Create top Project fixing nullability of fields
  relBuilder.project(topChildExprs, topChildExprsFields);
  relBuilder.convert(union.getRowType(), false);

  call.transformTo(relBuilder.build());
}
 

/**
 * Pulls the project above the semijoin and returns the resulting semijoin
 * condition. As a result, the semijoin condition should be modified such
 * that references to the LHS of a semijoin should now reference the
 * children of the project that's on the LHS.
 *
 * @param project  LogicalProject on the LHS of the semijoin
 * @param semiJoin the semijoin
 * @return the modified semijoin condition
 */
private RexNode adjustCondition(LogicalProject project, LogicalJoin semiJoin) {
  // create two RexPrograms -- the bottom one representing a
  // concatenation of the project and the RHS of the semijoin and the
  // top one representing the semijoin condition

  RexBuilder rexBuilder = project.getCluster().getRexBuilder();
  RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
  RelNode rightChild = semiJoin.getRight();

  // for the bottom RexProgram, the input is a concatenation of the
  // child of the project and the RHS of the semijoin
  RelDataType bottomInputRowType =
      SqlValidatorUtil.deriveJoinRowType(
          project.getInput().getRowType(),
          rightChild.getRowType(),
          JoinRelType.INNER,
          typeFactory,
          null,
          semiJoin.getSystemFieldList());
  RexProgramBuilder bottomProgramBuilder =
      new RexProgramBuilder(bottomInputRowType, rexBuilder);

  // add the project expressions, then add input references for the RHS
  // of the semijoin
  for (Pair<RexNode, String> pair : project.getNamedProjects()) {
    bottomProgramBuilder.addProject(pair.left, pair.right);
  }
  int nLeftFields = project.getInput().getRowType().getFieldCount();
  List<RelDataTypeField> rightFields =
      rightChild.getRowType().getFieldList();
  int nRightFields = rightFields.size();
  for (int i = 0; i < nRightFields; i++) {
    final RelDataTypeField field = rightFields.get(i);
    RexNode inputRef =
        rexBuilder.makeInputRef(
            field.getType(), i + nLeftFields);
    bottomProgramBuilder.addProject(inputRef, field.getName());
  }
  RexProgram bottomProgram = bottomProgramBuilder.getProgram();

  // input rowtype into the top program is the concatenation of the
  // project and the RHS of the semijoin
  RelDataType topInputRowType =
      SqlValidatorUtil.deriveJoinRowType(
          project.getRowType(),
          rightChild.getRowType(),
          JoinRelType.INNER,
          typeFactory,
          null,
          semiJoin.getSystemFieldList());
  RexProgramBuilder topProgramBuilder =
      new RexProgramBuilder(
          topInputRowType,
          rexBuilder);
  topProgramBuilder.addIdentity();
  topProgramBuilder.addCondition(semiJoin.getCondition());
  RexProgram topProgram = topProgramBuilder.getProgram();

  // merge the programs and expand out the local references to form
  // the new semijoin condition; it now references a concatenation of
  // the project's child and the RHS of the semijoin
  RexProgram mergedProgram =
      RexProgramBuilder.mergePrograms(
          topProgram,
          bottomProgram,
          rexBuilder);

  return mergedProgram.expandLocalRef(
      mergedProgram.getCondition());
}
 

@Override protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder,
    RelNode topProject, RelNode unionInputQuery, RelNode unionInputView) {
  // Union
  relBuilder.push(unionInputQuery);
  relBuilder.push(unionInputView);
  relBuilder.union(true);
  List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
    // We can take unionInputQuery as it is query based.
    RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
    exprList.add(
        rexBuilder.ensureType(
            field.getType(),
            rexBuilder.makeInputRef(relBuilder.peek(), i),
            true));
    nameList.add(field.getName());
  }
  relBuilder.project(exprList, nameList);
  // Rollup aggregate
  Aggregate aggregate = (Aggregate) unionInputQuery;
  final ImmutableBitSet groupSet = ImmutableBitSet.range(aggregate.getGroupCount());
  final List<AggCall> aggregateCalls = new ArrayList<>();
  for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
    AggregateCall aggCall = aggregate.getAggCallList().get(i);
    if (aggCall.isDistinct()) {
      // Cannot ROLLUP distinct
      return null;
    }
    SqlAggFunction rollupAgg =
        getRollup(aggCall.getAggregation());
    if (rollupAgg == null) {
      // Cannot rollup this aggregate, bail out
      return null;
    }
    final RexInputRef operand =
        rexBuilder.makeInputRef(relBuilder.peek(),
            aggregate.getGroupCount() + i);
    aggregateCalls.add(
        relBuilder.aggregateCall(rollupAgg, operand)
            .distinct(aggCall.isDistinct())
            .approximate(aggCall.isApproximate())
            .as(aggCall.name));
  }
  RelNode prevNode = relBuilder.peek();
  RelNode result = relBuilder
      .aggregate(relBuilder.groupKey(groupSet), aggregateCalls)
      .build();
  if (prevNode == result && groupSet.cardinality() != result.getRowType().getFieldCount()) {
    // Aggregate was not inserted but we need to prune columns
    result = relBuilder
        .push(result)
        .project(relBuilder.fields(groupSet))
        .build();
  }
  if (topProject != null) {
    // Top project
    return topProject.copy(topProject.getTraitSet(), ImmutableList.of(result));
  }
  // Result
  return result;
}
 

public RexNode singleton(RexBuilder rexBuilder,
    RelDataType inputRowType, AggregateCall aggregateCall) {
  final int arg = aggregateCall.getArgList().get(0);
  final RelDataTypeField field = inputRowType.getFieldList().get(arg);
  return rexBuilder.makeInputRef(field.getType(), arg);
}
 

/**
 * Pulls the project above the semi/anti join and returns the resulting semi/anti join
 * condition. As a result, the semi/anti join condition should be modified such
 * that references to the LHS of a semi/anti join should now reference the
 * children of the project that's on the LHS.
 *
 * @param project LogicalProject on the LHS of the semi/anti join
 * @param join the semi/anti join
 * @return the modified semi/anti join condition
 */
private RexNode adjustCondition(LogicalProject project, Join join) {
	// create two RexPrograms -- the bottom one representing a
	// concatenation of the project and the RHS of the semi/anti join and the
	// top one representing the semi/anti join condition

	RexBuilder rexBuilder = project.getCluster().getRexBuilder();
	RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
	RelNode rightChild = join.getRight();

	// for the bottom RexProgram, the input is a concatenation of the
	// child of the project and the RHS of the semi/anti join
	RelDataType bottomInputRowType =
			SqlValidatorUtil.deriveJoinRowType(
					project.getInput().getRowType(),
					rightChild.getRowType(),
					JoinRelType.INNER,
					typeFactory,
					null,
					join.getSystemFieldList());
	RexProgramBuilder bottomProgramBuilder =
			new RexProgramBuilder(bottomInputRowType, rexBuilder);

	// add the project expressions, then add input references for the RHS
	// of the semi/anti join
	for (Pair<RexNode, String> pair : project.getNamedProjects()) {
		bottomProgramBuilder.addProject(pair.left, pair.right);
	}
	int nLeftFields = project.getInput().getRowType().getFieldCount();
	List<RelDataTypeField> rightFields =
			rightChild.getRowType().getFieldList();
	int nRightFields = rightFields.size();
	for (int i = 0; i < nRightFields; i++) {
		final RelDataTypeField field = rightFields.get(i);
		RexNode inputRef =
				rexBuilder.makeInputRef(
						field.getType(), i + nLeftFields);
		bottomProgramBuilder.addProject(inputRef, field.getName());
	}
	RexProgram bottomProgram = bottomProgramBuilder.getProgram();

	// input rowtype into the top program is the concatenation of the
	// project and the RHS of the semi/anti join
	RelDataType topInputRowType =
			SqlValidatorUtil.deriveJoinRowType(
					project.getRowType(),
					rightChild.getRowType(),
					JoinRelType.INNER,
					typeFactory,
					null,
					join.getSystemFieldList());
	RexProgramBuilder topProgramBuilder =
			new RexProgramBuilder(
					topInputRowType,
					rexBuilder);
	topProgramBuilder.addIdentity();
	topProgramBuilder.addCondition(join.getCondition());
	RexProgram topProgram = topProgramBuilder.getProgram();

	// merge the programs and expand out the local references to form
	// the new semi/anti join condition; it now references a concatenation of
	// the project's child and the RHS of the semi/anti join
	RexProgram mergedProgram =
			RexProgramBuilder.mergePrograms(
					topProgram,
					bottomProgram,
					rexBuilder);

	return mergedProgram.expandLocalRef(
			mergedProgram.getCondition());
}
 

/**
 * This method creates a new UnorderedMux and Demux exchanges if mux operators are enabled.
 * @param child input to the new Unordered[Mux/Demux]Prel or new HashToRandomExchange node.
 * @param options options manager to check if mux is enabled.
 */
@Override
public Prel constructMuxPrel(Prel child, OptionManager options) {
  boolean isMuxEnabled = options.getOption(PlannerSettings.MUX_EXCHANGE.getOptionName()).bool_val;
  Prel newPrel = child;

  final List<String> childFields = child.getRowType().getFieldNames();

  List <RexNode> removeUpdatedExpr = null;

  if (isMuxEnabled) {
    // Insert Project Operator with new column that will be a hash for HashToRandomExchange fields
    final List<DistributionField> distFields = getFields();
    final List<String> outputFieldNames = Lists.newArrayList(childFields);
    final RexBuilder rexBuilder = getCluster().getRexBuilder();
    final List<RelDataTypeField> childRowTypeFields = child.getRowType().getFieldList();

    final HashPrelUtil.HashExpressionCreatorHelper<RexNode> hashHelper =
                                  new HashPrelUtil.RexNodeBasedHashExpressionCreatorHelper(rexBuilder);

    final List<RexNode> distFieldRefs = Lists.newArrayListWithExpectedSize(distFields.size());
    for (DistributionField distField : distFields) {
      final int fieldId = distField.getFieldId();
      distFieldRefs.add(rexBuilder.makeInputRef(childRowTypeFields.get(fieldId).getType(), fieldId));
    }

    final List <RexNode> updatedExpr = Lists.newArrayListWithExpectedSize(childRowTypeFields.size());
    removeUpdatedExpr = Lists.newArrayListWithExpectedSize(childRowTypeFields.size());
    for (RelDataTypeField field : childRowTypeFields) {
      RexNode rex = rexBuilder.makeInputRef(field.getType(), field.getIndex());
      updatedExpr.add(rex);
      removeUpdatedExpr.add(rex);
    }

    outputFieldNames.add(HashPrelUtil.HASH_EXPR_NAME);
    final RexNode distSeed = rexBuilder.makeBigintLiteral(BigDecimal.valueOf(HashPrelUtil.DIST_SEED)); // distribution seed
    updatedExpr.add(HashPrelUtil.createHashBasedPartitionExpression(distFieldRefs, distSeed, hashHelper));

    RelDataType rowType = RexUtil.createStructType(getCluster().getTypeFactory(),
        updatedExpr, outputFieldNames, null);

    ProjectPrel addColumnprojectPrel = new ProjectPrel(child.getCluster(), child.getTraitSet(), child, updatedExpr, rowType);

    newPrel = new UnorderedMuxExchangePrel(addColumnprojectPrel.getCluster(), addColumnprojectPrel.getTraitSet(),
            addColumnprojectPrel);
  }

  newPrel = new HashToRandomExchangePrel(getCluster(), getTraitSet(), newPrel, getFields());

  if (options.getOption(PlannerSettings.DEMUX_EXCHANGE.getOptionName()).bool_val) {
    HashToRandomExchangePrel hashExchangePrel = (HashToRandomExchangePrel) newPrel;
    // Insert a DeMuxExchange to narrow down the number of receivers
    newPrel = new UnorderedDeMuxExchangePrel(getCluster(), getTraitSet(), hashExchangePrel, hashExchangePrel.getFields());
  }

  if (isMuxEnabled) {
    // remove earlier inserted Project Operator - since it creates issues down the road in HashJoin
    RelDataType removeRowType = RexUtil.createStructType(newPrel.getCluster().getTypeFactory(),
        removeUpdatedExpr, childFields, null);

    ProjectPrel removeColumnProjectPrel = new ProjectPrel(newPrel.getCluster(), newPrel.getTraitSet(),
                                                          newPrel, removeUpdatedExpr, removeRowType);
    return removeColumnProjectPrel;
  }
  return newPrel;
}
 

public RexNode singleton(RexBuilder rexBuilder,
    RelDataType inputRowType, AggregateCall aggregateCall) {
  final int arg = aggregateCall.getArgList().get(0);
  final RelDataTypeField field = inputRowType.getFieldList().get(arg);
  return rexBuilder.makeInputRef(field.getType(), arg);
}
 

public RexNode singleton(RexBuilder rexBuilder,
    RelDataType inputRowType, AggregateCall aggregateCall) {
  final int arg = aggregateCall.getArgList().get(0);
  final RelDataTypeField field = inputRowType.getFieldList().get(arg);
  return rexBuilder.makeInputRef(field.getType(), arg);
}
 

/**
 * Pulls the project above the semijoin and returns the resulting semijoin
 * condition. As a result, the semijoin condition should be modified such
 * that references to the LHS of a semijoin should now reference the
 * children of the project that's on the LHS.
 *
 * @param project  LogicalProject on the LHS of the semijoin
 * @param semiJoin the semijoin
 * @return the modified semijoin condition
 */
private RexNode adjustCondition(LogicalProject project, SemiJoin semiJoin) {
  // create two RexPrograms -- the bottom one representing a
  // concatenation of the project and the RHS of the semijoin and the
  // top one representing the semijoin condition

  RexBuilder rexBuilder = project.getCluster().getRexBuilder();
  RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
  RelNode rightChild = semiJoin.getRight();

  // for the bottom RexProgram, the input is a concatenation of the
  // child of the project and the RHS of the semijoin
  RelDataType bottomInputRowType =
      SqlValidatorUtil.deriveJoinRowType(
          project.getInput().getRowType(),
          rightChild.getRowType(),
          JoinRelType.INNER,
          typeFactory,
          null,
          semiJoin.getSystemFieldList());
  RexProgramBuilder bottomProgramBuilder =
      new RexProgramBuilder(bottomInputRowType, rexBuilder);

  // add the project expressions, then add input references for the RHS
  // of the semijoin
  for (Pair<RexNode, String> pair : project.getNamedProjects()) {
    bottomProgramBuilder.addProject(pair.left, pair.right);
  }
  int nLeftFields = project.getInput().getRowType().getFieldCount();
  List<RelDataTypeField> rightFields =
      rightChild.getRowType().getFieldList();
  int nRightFields = rightFields.size();
  for (int i = 0; i < nRightFields; i++) {
    final RelDataTypeField field = rightFields.get(i);
    RexNode inputRef =
        rexBuilder.makeInputRef(
            field.getType(), i + nLeftFields);
    bottomProgramBuilder.addProject(inputRef, field.getName());
  }
  RexProgram bottomProgram = bottomProgramBuilder.getProgram();

  // input rowtype into the top program is the concatenation of the
  // project and the RHS of the semijoin
  RelDataType topInputRowType =
      SqlValidatorUtil.deriveJoinRowType(
          project.getRowType(),
          rightChild.getRowType(),
          JoinRelType.INNER,
          typeFactory,
          null,
          semiJoin.getSystemFieldList());
  RexProgramBuilder topProgramBuilder =
      new RexProgramBuilder(
          topInputRowType,
          rexBuilder);
  topProgramBuilder.addIdentity();
  topProgramBuilder.addCondition(semiJoin.getCondition());
  RexProgram topProgram = topProgramBuilder.getProgram();

  // merge the programs and expand out the local references to form
  // the new semijoin condition; it now references a concatenation of
  // the project's child and the RHS of the semijoin
  RexProgram mergedProgram =
      RexProgramBuilder.mergePrograms(
          topProgram,
          bottomProgram,
          rexBuilder);

  return mergedProgram.expandLocalRef(
      mergedProgram.getCondition());
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Union union = call.rel(0);

  final int count = union.getRowType().getFieldCount();
  if (count == 1) {
    // No room for optimization since we cannot create an empty Project
    // operator. If we created a Project with one column, this rule would
    // cycle.
    return;
  }

  final RexBuilder rexBuilder = union.getCluster().getRexBuilder();
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(union);
  if (predicates == null) {
    return;
  }

  final Map<Integer, RexNode> constants = new HashMap<>();
  for (Map.Entry<RexNode, RexNode> e : predicates.constantMap.entrySet()) {
    if (e.getKey() instanceof RexInputRef) {
      constants.put(((RexInputRef) e.getKey()).getIndex(), e.getValue());
    }
  }

  // None of the expressions are constant. Nothing to do.
  if (constants.isEmpty()) {
    return;
  }

  // Create expressions for Project operators before and after the Union
  List<RelDataTypeField> fields = union.getRowType().getFieldList();
  List<RexNode> topChildExprs = new ArrayList<>();
  List<String> topChildExprsFields = new ArrayList<>();
  List<RexNode> refs = new ArrayList<>();
  ImmutableBitSet.Builder refsIndexBuilder = ImmutableBitSet.builder();
  for (RelDataTypeField field : fields) {
    final RexNode constant = constants.get(field.getIndex());
    if (constant != null) {
      topChildExprs.add(constant);
      topChildExprsFields.add(field.getName());
    } else {
      final RexNode expr = rexBuilder.makeInputRef(union, field.getIndex());
      topChildExprs.add(expr);
      topChildExprsFields.add(field.getName());
      refs.add(expr);
      refsIndexBuilder.set(field.getIndex());
    }
  }
  ImmutableBitSet refsIndex = refsIndexBuilder.build();

  // Update top Project positions
  final Mappings.TargetMapping mapping =
      RelOptUtil.permutation(refs, union.getInput(0).getRowType()).inverse();
  topChildExprs = ImmutableList.copyOf(RexUtil.apply(mapping, topChildExprs));

  // Create new Project-Union-Project sequences
  final RelBuilder relBuilder = call.builder();
  for (RelNode input : union.getInputs()) {
    List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
    for (int j : refsIndex) {
      newChildExprs.add(
          Pair.of(rexBuilder.makeInputRef(input, j),
              input.getRowType().getFieldList().get(j).getName()));
    }
    if (newChildExprs.isEmpty()) {
      // At least a single item in project is required.
      newChildExprs.add(
          Pair.of(topChildExprs.get(0), topChildExprsFields.get(0)));
    }
    // Add the input with project on top
    relBuilder.push(input);
    relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
  }
  relBuilder.union(union.all, union.getInputs().size());
  // Create top Project fixing nullability of fields
  relBuilder.project(topChildExprs, topChildExprsFields);
  relBuilder.convert(union.getRowType(), false);

  call.transformTo(relBuilder.build());
}