Java Code Examples for org.apache.calcite.rex.RexUtil#createStructType()

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

/**
 * Adds top project to ensure final output field names are preserved.
 * In case of duplicated column names, will rename duplicates.
 * Top project will be added only if top project is non-trivial and
 * child physical relational node is not project.
 *
 * @param prel physical relational node
 * @param validatedRowType final output row type
 * @return physical relational node with top project if necessary
 */
private Prel addTopProjectPrel(Prel prel, RelDataType validatedRowType) {
  RelDataType rowType = prel.getRowType();
  if (rowType.getFieldCount() != validatedRowType.getFieldCount()) {
    return prel;
  }

  RexBuilder rexBuilder = prel.getCluster().getRexBuilder();
  List<RexNode> projections = new ArrayList<>();
  int projectCount = rowType.getFieldList().size();

  for (int i = 0; i < projectCount; i++) {
    projections.add(rexBuilder.makeInputRef(prel, i));
  }

  List<String> fieldNames = SqlValidatorUtil.uniquify(
      validatedRowType.getFieldNames(),
      SqlValidatorUtil.EXPR_SUGGESTER,
      prel.getCluster().getTypeFactory().getTypeSystem().isSchemaCaseSensitive());

  RelDataType newRowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(), projections, fieldNames, null);
  ProjectPrel topProject = new ProjectPrel(prel.getCluster(),
      prel.getTraitSet(),
      prel,
      projections,
      newRowType,
      true);  //outputProj = true : NONE -> OK_NEW_SCHEMA, also handle expression with NULL type.

  if (prel instanceof Project && DrillRelOptUtil.isTrivialProject(topProject, true)) {
    return new ProjectPrel(prel.getCluster(),
        prel.getTraitSet(),
        ((Project) prel).getInput(),
        ((Project) prel).getProjects(),
        prel.getRowType(),
        true); //outputProj = true : NONE -> OK_NEW_SCHEMA, also handle expression with NULL type.
  } else {
    return topProject;
  }
}
 

public static ProjectPrel addHashProject(List<DistributionField> distFields, Prel input, Integer ringCount){

    // Insert Project SqlOperatorImpl with new column that will be a hash for HashToRandomExchange fields

    final List<String> outputFieldNames = Lists.newArrayList(input.getRowType().getFieldNames());
    final String fieldName = ringCount == null ? HashPrelUtil.HASH_EXPR_NAME : WriterPrel.BUCKET_NUMBER_FIELD;
    outputFieldNames.add(fieldName);

    final RexBuilder rexBuilder = input.getCluster().getRexBuilder();
    final List<RelDataTypeField> childRowTypeFields = input.getRowType().getFieldList();

    // create new hashed field.
    final HashExpressionCreatorHelper<RexNode> hashHelper = new RexNodeBasedHashExpressionCreatorHelper(rexBuilder);
    final List<RexNode> distFieldRefs = Lists.newArrayListWithExpectedSize(distFields.size());
    for(int i = 0; i < distFields.size(); i++) {
      final int fieldId = distFields.get(i).getFieldId();
      distFieldRefs.add(rexBuilder.makeInputRef(childRowTypeFields.get(fieldId).getType(), fieldId));
    }

    final List <RexNode> updatedExpr = Lists.newArrayListWithExpectedSize(childRowTypeFields.size() + 1);
    for ( RelDataTypeField field : childRowTypeFields) {
      RexNode rex = rexBuilder.makeInputRef(field.getType(), field.getIndex());
      updatedExpr.add(rex);
    }
    RexNode hashExpression = HashPrelUtil.createHashBasedPartitionExpression(distFieldRefs, hashHelper);

    if(ringCount != null){
      RelDataType intType = input.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER);
      hashExpression = rexBuilder.makeCall(SqlStdOperatorTable.MOD, ImmutableList.of(hashExpression, rexBuilder.makeExactLiteral(BigDecimal.valueOf(ringCount), intType)));
      hashExpression = rexBuilder.makeCall(SqlStdOperatorTable.ABS, Collections.singletonList(hashExpression));
    }
    updatedExpr.add(hashExpression);

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

    ProjectPrel addColumnprojectPrel = ProjectPrel.create(input.getCluster(), input.getTraitSet(), input, updatedExpr, rowType);
    return addColumnprojectPrel;
  }
 

@Override
public RelNode createProject(RelNode child,
                             List<? extends RexNode> childExprs, List<String> fieldNames) {
  final RelOptCluster cluster = child.getCluster();
  final RelDataType rowType = RexUtil.createStructType(cluster.getTypeFactory(), childExprs, fieldNames, SqlValidatorUtil.F_SUGGESTER);
  final RelNode project = ProjectRel.create(
      cluster,
      child.getTraitSet().plus(Rel.LOGICAL),
      RelOptRule.convert(child, child.getTraitSet().plus(Rel.LOGICAL).simplify()),
      childExprs,
      rowType);

  return project;
}
 

@Override
public RelNode createProject(RelNode child,
                             List<? extends RexNode> childExprs, List<String> fieldNames) {
  final RelOptCluster cluster = child.getCluster();
  final RelDataType rowType = RexUtil.createStructType(cluster.getTypeFactory(), childExprs, fieldNames, SqlValidatorUtil.F_SUGGESTER);
  final RelNode project = ProjectRel.create(cluster, child.getTraitSet().plus(Rel.LOGICAL), child, childExprs, rowType);

  return project;
}
 

protected DrillRel addRenamedProject(DrillRel rel, RelDataType validatedRowType) {
  RelDataType t = rel.getRowType();

  RexBuilder b = rel.getCluster().getRexBuilder();
  List<RexNode> projections = Lists.newArrayList();
  int projectCount = t.getFieldList().size();

  for (int i =0; i < projectCount; i++) {
    projections.add(b.makeInputRef(rel, i));
  }

  final List<String> fieldNames2 = SqlValidatorUtil.uniquify(
      validatedRowType.getFieldNames(),
      SqlValidatorUtil.EXPR_SUGGESTER,
      rel.getCluster().getTypeFactory().getTypeSystem().isSchemaCaseSensitive());

  RelDataType newRowType = RexUtil.createStructType(rel.getCluster().getTypeFactory(),
      projections, fieldNames2, null);

  DrillProjectRel topProj = DrillProjectRel.create(rel.getCluster(), rel.getTraitSet(), rel, projections, newRowType);

  // Add a final non-trivial Project to get the validatedRowType, if child is not project.
  if (rel instanceof Project && DrillRelOptUtil.isTrivialProject(topProj, true)) {
    return rel;
  } else{
    return topProj;
  }
}
 

/**
 * Creates a MutableProject.
 *
 * @param input         Input relational expression
 * @param exprList      List of expressions for the input columns
 * @param fieldNameList Aliases of the expressions, or null to generate
 */
public static MutableRel of(MutableRel input, List<RexNode> exprList,
    List<String> fieldNameList) {
  final RelDataType rowType =
      RexUtil.createStructType(input.cluster.getTypeFactory(), exprList,
          fieldNameList, SqlValidatorUtil.F_SUGGESTER);
  return of(rowType, input, exprList);
}
 

private Prel prefixTabNameToStar(Prel prel, Void value) throws RuntimeException {
  if (StarColumnHelper.containsStarColumn(prel.getRowType()) && prefixedForStar) {

    List<RexNode> exprs = Lists.newArrayList();

    for (RelDataTypeField field : prel.getRowType().getFieldList()) {
      RexNode expr = prel.getCluster().getRexBuilder().makeInputRef(field.getType(), field.getIndex());
      exprs.add(expr);
    }

    List<String> fieldNames = Lists.newArrayList();

    long tableId = tableNumber.getAndIncrement();

    for (String name : prel.getRowType().getFieldNames()) {
      if (StarColumnHelper.isNonPrefixedStarColumn(name)) {
        fieldNames.add("T" +  tableId + StarColumnHelper.PREFIX_DELIMITER + name);  // Add prefix to * column.
      } else {
        fieldNames.add(name);  // Keep regular column as it is.
      }
    }
    RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(),
        exprs, fieldNames, null);

    // insert a PAS.
    ProjectPrel proj = new ProjectPrel(prel.getCluster(), prel.getTraitSet(), prel, exprs, rowType);

    return proj;
  } else {
    return visitPrel(prel, value);
  }
}
 

@Override
public Prel visitProject(ProjectPrel prel, Void value) throws RuntimeException {
  // Require prefix rename : there exists other expression, in addition to a star column.
  if (!prefixedForStar  // not set yet.
      && StarColumnHelper.containsStarColumnInProject(prel.getInput().getRowType(), prel.getProjects())
      && prel.getRowType().getFieldNames().size() > 1) {
    prefixedForStar = true;
  }

  // For project, we need make sure that the project's field name is same as the input,
  // when the project expression is RexInPutRef, since we may insert a PAS which will
  // rename the projected fields.

  RelNode child = ((Prel) prel.getInput(0)).accept(this, null);

  List<String> fieldNames = Lists.newArrayList();

  for (Pair<String, RexNode> pair : Pair.zip(prel.getRowType().getFieldNames(), prel.getProjects())) {
    if (pair.right instanceof RexInputRef) {
      String name = child.getRowType().getFieldNames().get(((RexInputRef) pair.right).getIndex());
      fieldNames.add(name);
    } else {
      fieldNames.add(pair.left);
    }
  }

  // Make sure the field names are unique : no allow of duplicate field names in a rowType.
  fieldNames = makeUniqueNames(fieldNames);

  RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(),
      prel.getProjects(), fieldNames, null);

  ProjectPrel newProj = (ProjectPrel) prel.copy(prel.getTraitSet(), child, prel.getProjects(), rowType);

  if (ProjectRemoveRule.isTrivial(newProj)) {
    return (Prel) child;
  } else {
    return newProj;
  }
}
 

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

private RelNode rename(RelNode input, List<RelDataTypeField> inputFields, List<String> outputFieldNames) {
  List<RexNode> exprs = Lists.newArrayList();

  for (RelDataTypeField field : inputFields) {
    RexNode expr = input.getCluster().getRexBuilder().makeInputRef(field.getType(), field.getIndex());
    exprs.add(expr);
  }

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

  ProjectPrel proj = new ProjectPrel(input.getCluster(), input.getTraitSet(), input, exprs, rowType);

  return proj;
}
 

private RelNode rename(RelNode input, List<RelDataTypeField> inputFields, List<String> outputFieldNames) {
  List<RexNode> exprs = Lists.newArrayList();

  for (RelDataTypeField field : inputFields) {
    RexNode expr = input.getCluster().getRexBuilder().makeInputRef(field.getType(), field.getIndex());
    exprs.add(expr);
  }

  RelDataType rowType = RexUtil.createStructType(input.getCluster().getTypeFactory(), exprs, outputFieldNames);

  ProjectPrel proj = ProjectPrel.create(input.getCluster(), input.getTraitSet(), input, exprs, rowType);

  return proj;
}
 

@Override
public RelNode createProject(RelNode child,
                             List<? extends RexNode> childExprs, List<String> fieldNames) {
  final RelOptCluster cluster = child.getCluster();
  final RelDataType rowType = RexUtil.createStructType(cluster.getTypeFactory(), childExprs, fieldNames);
  final RelNode project = new ProjectPrel(cluster, child.getTraitSet().plus(Prel.DRILL_PHYSICAL),
      child, Lists.newArrayList(childExprs), rowType);

  return project;
}
 

@Override
public RelNode createProject(RelNode child,
                             List<? extends RexNode> childExprs, List<String> fieldNames) {
  final RelOptCluster cluster = child.getCluster();
  final RelDataType rowType =
      RexUtil.createStructType(cluster.getTypeFactory(), childExprs, fieldNames, null);

  return DrillProjectRel.create(cluster, child.getTraitSet().plus(DRILL_LOGICAL), child, childExprs, rowType);
}
 

@Override
public Prel visitJoin(JoinPrel prel, Void value) throws RuntimeException {

  List<RelNode> children = Lists.newArrayList();

  for(Prel child : prel){
    child = child.accept(this, null);
    children.add(child);
  }

  if ((prel.getRowType().getFieldCount() != prel.getInputRowType().getFieldCount())
      || !(prel.getRowType().getFieldNames().equals(prel.getInputRowType().getFieldNames()))) {
    List<String> outputFieldNames = new ArrayList<>();
    List<RexNode> exprs = new ArrayList<>();
    List<RelDataTypeField> fields = prel.getRowType().getFieldList();
    for (RelDataTypeField field : fields) {
      RexNode expr = prel.getCluster().getRexBuilder().makeInputRef(field.getType(), field.getIndex());
      exprs.add(expr);
      outputFieldNames.add(field.getName());
    }

    RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(), exprs, outputFieldNames);

    JoinPrel newJoin = (JoinPrel) prel.copy(prel.getTraitSet(), children);
    ProjectPrel project = ProjectPrel.create(newJoin.getCluster(), newJoin.getTraitSet(), newJoin, exprs, rowType);
    return (Prel) project;
  }

  return (Prel) prel.copy(prel.getTraitSet(), children);
}
 

public RelNode createProject(RelNode input, List<RelHint> hints,
                      List<? extends RexNode> childExprs, List<String> fieldNames) {
  final RelDataType rowType =
      RexUtil.createStructType(input.getCluster().getTypeFactory(), childExprs,
          fieldNames, SqlValidatorUtil.F_SUGGESTER);
  return EnumerableProject.create(input, childExprs, rowType);
}
 

@Deprecated // to be removed before 2.0
public LogicalProject(RelOptCluster cluster, RelNode input,
    List<RexNode> projects, List<String> fieldNames, int flags) {
  this(cluster, cluster.traitSetOf(RelCollations.EMPTY),
      input, projects,
      RexUtil.createStructType(cluster.getTypeFactory(), projects,
          fieldNames, null));
  Util.discard(flags);
}
 

@Override
public Prel visitProject(ProjectPrel prel, Void value) throws RuntimeException {
  ProjectPrel proj = prel;

  // Require prefix rename : there exists other expression, in addition to a star column.
  if (!prefixedForStar  // not set yet.
      && StarColumnHelper.containsStarColumnInProject(prel.getInput().getRowType(), proj.getProjects())
      && prel.getRowType().getFieldNames().size() > 1) {
    prefixedForStar = true;
  }

  // For project, we need make sure that the project's field name is same as the input,
  // when the project expression is RexInPutRef, since we may insert a PAS which will
  // rename the projected fields.

  RelNode child = ((Prel) prel.getInput(0)).accept(this, null);

  List<String> fieldNames = Lists.newArrayList();

  for (Pair<String, RexNode> pair : Pair.zip(prel.getRowType().getFieldNames(), proj.getProjects())) {
    if (pair.right instanceof RexInputRef) {
      String name = child.getRowType().getFieldNames().get(((RexInputRef) pair.right).getIndex());
      fieldNames.add(name);
    } else {
      fieldNames.add(pair.left);
    }
  }

  // Make sure the field names are unique : no allow of duplicate field names in a rowType.
  fieldNames = makeUniqueNames(fieldNames);

  RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(), proj.getProjects(), fieldNames);

  ProjectPrel newProj = (ProjectPrel) proj.copy(proj.getTraitSet(), child, proj.getProjects(), rowType);

  if (ProjectRemoveRule.isTrivial(newProj)) {
    return (Prel) child;
  } else {
    return newProj;
  }
}
 

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

@Override
public void onMatch(RelOptRuleCall call) {
  NestedLoopJoinPrel nlj = call.rel(0);
  RexBuilder rexBuilder = nlj.getCluster().getRexBuilder();

  RexNode condition = nlj.getCondition();

  condition = RexRewriter.rewrite(condition, ImmutableList.of(new SimplifyNLJConditionRule.TrigFractionRewriteRule(rexBuilder), new SimplifyNLJConditionRule.SinCosPlusMinusRewriteRule(rexBuilder)));

  int leftCount = nlj.getLeft().getRowType().getFieldCount();
  List<RexNode> identityProjects = identityProjects(nlj.getRowType(), nlj.getProjectedFields());
  SimplifyNLJConditionRule.ExpressionPusher pusher = new SimplifyNLJConditionRule.ExpressionPusher(rexBuilder, identityProjects, leftCount);

  Pointer<SimplifyNLJConditionRule.Side> side = new Pointer<>();

  identityProjects.forEach(p -> p.accept(pusher, side));

  side.value = SimplifyNLJConditionRule.Side.EMPTY;
  RexNode tempCondition = condition.accept(pusher, side);

  if (pusher.exprs.stream().allMatch(e -> e instanceof RexInputRef)) {
    return;
  }


  List<RexNode> leftExprs = new ArrayList<>();
  pusher.leftExprs.forEach(i -> leftExprs.add(pusher.exprs.get(i)));

  RelDataType leftType = RexUtil.createStructType(nlj.getCluster().getTypeFactory(), leftExprs,
    null, SqlValidatorUtil.F_SUGGESTER);

  ProjectPrel newLeft = ProjectPrel.create(nlj.getCluster(), nlj.getLeft().getTraitSet(), nlj.getLeft(), leftExprs, leftType);

  List<RexNode> rightExprs = new ArrayList<>();
  pusher.rightExprs.forEach(i -> rightExprs.add(RexUtil.shift(pusher.exprs.get(i), -leftCount)));

  RelDataType rightType = RexUtil.createStructType(nlj.getCluster().getTypeFactory(), rightExprs,
    null, SqlValidatorUtil.F_SUGGESTER);

  ProjectPrel newRight = ProjectPrel.create(nlj.getCluster(), nlj.getRight().getTraitSet(), nlj.getRight(), rightExprs, rightType);

  Mapping mapping = Mappings.bijection(ImmutableList.<Integer>builder().addAll(pusher.leftExprs).addAll(pusher.rightExprs).build());

  RexShuttle shuttle = new RexPermuteInputsShuttle(mapping.inverse());
  RexNode newCondition = tempCondition.accept(shuttle);

  ImmutableBitSet.Builder projectedBuilder = ImmutableBitSet.builder();
  for (RexNode r : pusher.identityProjects()) {
    RexInputRef newRef = (RexInputRef) r.accept(shuttle);
    projectedBuilder.set(newRef.getIndex());
  }

  if (newCondition.toString().equals(nlj.getCondition().toString())) {
    return;
  }

  NestedLoopJoinPrel newJoin = NestedLoopJoinPrel.create(nlj.getCluster(), nlj.getTraitSet(), newLeft, newRight, nlj.getJoinType(), newCondition, projectedBuilder.build());

  call.transformTo(newJoin);
}