org.apache.calcite.rel.RelFieldCollation Java Examples

/** Converts a RexFieldCollation to an ORDER BY item. */
private void addOrderItem(List<SqlNode> orderByList,
                          RexProgram program, RexFieldCollation field) {
  SqlNode node = toSql(program, field.left);
  SqlNode nullDirectionNode = null;
  if (field.getNullDirection() != RelFieldCollation.NullDirection.UNSPECIFIED) {
    final boolean first =
              field.getNullDirection() == RelFieldCollation.NullDirection.FIRST;
    nullDirectionNode = dialect.emulateNullDirection(
            node, first, field.getDirection().isDescending());
  }
  if (nullDirectionNode != null) {
    orderByList.add(nullDirectionNode);
    switch (field.getDirection()) {
    case DESCENDING:
    case STRICTLY_DESCENDING:
      node = SqlStdOperatorTable.DESC.createCall(POS, node);
      break;
    default:
      break;
    }
    orderByList.add(node);
  } else {
    orderByList.add(toSql(program, field));
  }
}
 

/**
 * @see #dispatch
 */
public Result visit(Sort e) {
  Result x = visitChild(0, e.getInput());
  Builder builder = x.builder(e, Clause.ORDER_BY);
  List<SqlNode> orderByList = Expressions.list();
  for (RelFieldCollation field : e.getCollation().getFieldCollations()) {
    builder.addOrderItem(orderByList, field);
  }
  if (!orderByList.isEmpty()) {
    builder.setOrderBy(new SqlNodeList(orderByList, POS));
    x = builder.result();
  }
  if (e.fetch != null) {
    builder = x.builder(e, Clause.FETCH);
    builder.setFetch(builder.context.toSql(null, e.fetch));
    x = builder.result();
  }
  if (e.offset != null) {
    builder = x.builder(e, Clause.OFFSET);
    builder.setOffset(builder.context.toSql(null, e.offset));
    x = builder.result();
  }
  return x;
}
 

/**
 * Converts a collation to an ORDER BY item.
 */
public SqlNode toSql(RelFieldCollation collation) {
  SqlNode node = field(collation.getFieldIndex());
  switch (collation.getDirection()) {
    case DESCENDING:
    case STRICTLY_DESCENDING:
      node = SqlStdOperatorTable.DESC.createCall(POS, node);
  }
  if (collation.nullDirection != dialect.defaultNullDirection(collation.direction)) {
    switch (collation.nullDirection) {
      case FIRST:
        node = SqlStdOperatorTable.NULLS_FIRST.createCall(POS, node);
        break;
      case LAST:
        node = SqlStdOperatorTable.NULLS_LAST.createCall(POS, node);
        break;
    }
  }
  return node;
}
 

/** Converts a collation to an ORDER BY item. */
public SqlNode toSql(RelFieldCollation collation) {
    SqlNode node = field(collation.getFieldIndex());
    switch (collation.getDirection()) {
    case DESCENDING:
    case STRICTLY_DESCENDING:
        node = SqlStdOperatorTable.DESC.createCall(POS, node);
    }
    if (collation.nullDirection != dialect.defaultNullDirection(collation.direction)) {
        switch (collation.nullDirection) {
        case FIRST:
            node = SqlStdOperatorTable.NULLS_FIRST.createCall(POS, node);
            break;
        case LAST:
            node = SqlStdOperatorTable.NULLS_LAST.createCall(POS, node);
            break;
        }
    }
    return node;
}
 

@Override public void implement(GeodeImplementContext geodeImplementContext) {
  geodeImplementContext.visitChild(getInput());

  List<RelFieldCollation> sortCollations = collation.getFieldCollations();

  if (!sortCollations.isEmpty()) {

    List<String> orderByFields = new ArrayList<>();

    for (RelFieldCollation fieldCollation : sortCollations) {
      final String name = fieldName(fieldCollation.getFieldIndex());
      orderByFields.add(name + " " + direction(fieldCollation.getDirection()));
    }
    geodeImplementContext.addOrderByFields(orderByFields);
  }

  if (fetch != null) {
    geodeImplementContext.setLimit(((RexLiteral) fetch).getValueAs(Long.class));
  }
}
 

ColumnRowType buildColumnRowType() {
    OLAPRel olapChild = (OLAPRel) getInput(0);
    ColumnRowType inputColumnRowType = olapChild.getColumnRowType();

    List<TblColRef> columns = new ArrayList<>();
    // the input col always be collected by left
    columns.addAll(inputColumnRowType.getAllColumns());

    // add window aggregate calls column
    for (Group group : groups) {
        List<TupleExpression> sourceColOuter = Lists.newArrayList();
        group.keys.asSet().stream().map(inputColumnRowType::getTupleExpressionByIndex).forEach(sourceColOuter::add);
        group.orderKeys.getFieldCollations().stream().map(RelFieldCollation::getFieldIndex)
                .map(inputColumnRowType::getTupleExpressionByIndex).forEach(sourceColOuter::add);
        for (AggregateCall aggrCall : group.getAggregateCalls(this)) {
            TblColRef aggrCallCol = TblColRef.newInnerColumn(aggrCall.getName(),
                    TblColRef.InnerDataTypeEnum.LITERAL);
            List<TupleExpression> sourceColInner = Lists.newArrayList(sourceColOuter.iterator());
            aggrCall.getArgList().stream().filter(i -> i < inputColumnRowType.size())
                    .map(inputColumnRowType::getTupleExpressionByIndex).forEach(sourceColInner::add);
            aggrCallCol.setSubTupleExps(sourceColInner);
            columns.add(aggrCallCol);
        }
    }
    return new ColumnRowType(columns);
}
 

/** Returns whether NULL values are sorted first or last, in this dialect,
 * in an ORDER BY item of a given direction. */
public @Nonnull RelFieldCollation.NullDirection defaultNullDirection(
    RelFieldCollation.Direction direction) {
  switch (direction) {
  case ASCENDING:
  case STRICTLY_ASCENDING:
    return getNullCollation().last(false)
        ? RelFieldCollation.NullDirection.LAST
        : RelFieldCollation.NullDirection.FIRST;
  case DESCENDING:
  case STRICTLY_DESCENDING:
    return getNullCollation().last(true)
        ? RelFieldCollation.NullDirection.LAST
        : RelFieldCollation.NullDirection.FIRST;
  default:
    return RelFieldCollation.NullDirection.UNSPECIFIED;
  }
}
 

private static RelFieldCollation collation(RexNode node,
    RelFieldCollation.Direction direction,
    RelFieldCollation.NullDirection nullDirection, List<RexNode> extraNodes) {
  switch (node.getKind()) {
  case INPUT_REF:
    return new RelFieldCollation(((RexInputRef) node).getIndex(), direction,
        Util.first(nullDirection, direction.defaultNullDirection()));
  case DESCENDING:
    return collation(((RexCall) node).getOperands().get(0),
        RelFieldCollation.Direction.DESCENDING,
        nullDirection, extraNodes);
  case NULLS_FIRST:
    return collation(((RexCall) node).getOperands().get(0), direction,
        RelFieldCollation.NullDirection.FIRST, extraNodes);
  case NULLS_LAST:
    return collation(((RexCall) node).getOperands().get(0), direction,
        RelFieldCollation.NullDirection.LAST, extraNodes);
  default:
    final int fieldIndex = extraNodes.size();
    extraNodes.add(node);
    return new RelFieldCollation(fieldIndex, direction,
        Util.first(nullDirection, direction.defaultNullDirection()));
  }
}
 

public void rewriteRel(Sort rel) {
    RelCollation oldCollation = rel.getCollation();
    final RelNode oldChild = rel.getInput();
    final RelNode newChild = getNewForOldRel(oldChild);
    final Mappings.TargetMapping mapping = getNewForOldInputMapping(oldChild);

    // validate
    for (RelFieldCollation field : oldCollation.getFieldCollations()) {
        int oldInput = field.getFieldIndex();
        RelDataType sortFieldType = oldChild.getRowType().getFieldList().get(oldInput).getType();
        if (sortFieldType.isStruct()) {
            // TODO jvs 10-Feb-2005
            throw Util.needToImplement("sorting on structured types");
        }
    }
    RelCollation newCollation = RexUtil.apply(mapping, oldCollation);
    Sort newRel = LogicalSort.create(newChild, newCollation, rel.offset, rel.fetch);
    setNewForOldRel(rel, newRel);
}
 

public static void validateCollation(RelOptCluster cluster, RelNode child, ImmutableBitSet groupSet) {
  if (groupSet.isEmpty()) {
    // If no groups, no collation is required
    return;
  }

  final RelCollation requiredCollation = RelCollations.of(
      StreamSupport.stream(groupSet.spliterator(), false).map(RelFieldCollation::new).collect(Collectors.toList()));

  final RelMetadataQuery mq = cluster.getMetadataQuery();
  final List<RelCollation> collations = mq.collations(child);

  for(RelCollation collation: collations) {
    if (collation.satisfies(requiredCollation)) {
      return;
    }
  }

  throw new AssertionError("child collations [" + collations + "] does not match expected collation [" + requiredCollation + "]");
}
 

/** Returns references to fields for a given collation. */
public ImmutableList<RexNode> fields(RelCollation collation) {
    final ImmutableList.Builder<RexNode> nodes = ImmutableList.builder();
    for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
        RexNode node = field(fieldCollation.getFieldIndex());
        switch (fieldCollation.direction) {
        case DESCENDING:
            node = desc(node);
        }
        switch (fieldCollation.nullDirection) {
        case FIRST:
            node = nullsFirst(node);
            break;
        case LAST:
            node = nullsLast(node);
            break;
        }
        nodes.add(node);
    }
    return nodes.build();
}
 

public static RelCollation getCollation(
    final List<RexFieldCollation> collations) {
  return RelCollations.of(
      new AbstractList<RelFieldCollation>() {
        public RelFieldCollation get(int index) {
          final RexFieldCollation collation = collations.get(index);
          return new RelFieldCollation(
              ((RexLocalRef) collation.left).getIndex(),
              collation.getDirection(),
              collation.getNullDirection());
        }

        public int size() {
          return collations.size();
        }
      });
}
 

private static RelFieldCollation collation(RexNode node, RelFieldCollation.Direction direction,
        RelFieldCollation.NullDirection nullDirection, List<RexNode> extraNodes) {
    switch (node.getKind()) {
    case INPUT_REF:
        return new RelFieldCollation(((RexInputRef) node).getIndex(), direction,
                Util.first(nullDirection, direction.defaultNullDirection()));
    case DESCENDING:
        return collation(((RexCall) node).getOperands().get(0), RelFieldCollation.Direction.DESCENDING,
                nullDirection, extraNodes);
    case NULLS_FIRST:
        return collation(((RexCall) node).getOperands().get(0), direction, RelFieldCollation.NullDirection.FIRST,
                extraNodes);
    case NULLS_LAST:
        return collation(((RexCall) node).getOperands().get(0), direction, RelFieldCollation.NullDirection.LAST,
                extraNodes);
    default:
        final int fieldIndex = extraNodes.size();
        extraNodes.add(node);
        return new RelFieldCollation(fieldIndex, direction,
                Util.first(nullDirection, direction.defaultNullDirection()));
    }
}
 

@Override public SqlMonotonicity getOperandMonotonicity(int ordinal) {
  RexNode operand = operands.get(ordinal);

  if (operand instanceof RexInputRef) {
    for (RelCollation ic : inputCollations) {
      if (ic.getFieldCollations().isEmpty()) {
        continue;
      }

      for (RelFieldCollation rfc : ic.getFieldCollations()) {
        if (rfc.getFieldIndex() == ((RexInputRef) operand).getIndex()) {
          return rfc.direction.monotonicity();
          // TODO: Is it possible to have more than one RelFieldCollation for a RexInputRef?
        }
      }
    }
  } else if (operand instanceof RexCall) {
    final RexCallBinding binding =
        RexCallBinding.create(typeFactory, (RexCall) operand, inputCollations);
    return ((RexCall) operand).getOperator().getMonotonicity(binding);
  }

  return SqlMonotonicity.NOT_MONOTONIC;
}
 

@Override public void implement(Implementor implementor) {
  implementor.visitChild(0, getInput());
  final List<RelDataTypeField> fields = getRowType().getFieldList();

  for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
    final String name = fields.get(fieldCollation.getFieldIndex()).getName();
    final String rawName = implementor.expressionItemMap.getOrDefault(name, name);
    implementor.addSort(rawName, fieldCollation.getDirection());
  }

  if (offset != null) {
    implementor.offset(((RexLiteral) offset).getValueAs(Long.class));
  }

  if (fetch != null) {
    implementor.fetch(((RexLiteral) fetch).getValueAs(Long.class));
  }
}
 

public static RelNode convert(Order order, ConversionContext context) throws InvalidRelException{

    // if there are compound expressions in the order by, we need to convert into projects on either side.
    RelNode input = context.toRel(order.getInput());
    List<String> fields = input.getRowType().getFieldNames();

    // build a map of field names to indices.
    Map<String, Integer> fieldMap = Maps.newHashMap();
    int i =0;
    for(String field : fields){
      fieldMap.put(field, i);
      i++;
    }

    List<RelFieldCollation> collations = Lists.newArrayList();

    for(Ordering o : order.getOrderings()){
      String fieldName = ExprHelper.getFieldName(o.getExpr());
      int fieldId = fieldMap.get(fieldName);
      RelFieldCollation c = new RelFieldCollation(fieldId, o.getDirection(), o.getNullDirection());
      collations.add(c);
    }
    return new DrillSortRel(context.getCluster(), context.getLogicalTraits(), input, RelCollations.of(collations));
  }
 

private static List<RelCollation> enumerableJoin0(RelMetadataQuery mq,
    RelNode left, RelNode right, JoinRelType joinType) {
  // The current implementation can preserve the sort order of the left input if one of the
  // following conditions hold:
  // (i) join type is INNER or LEFT;
  // (ii) RelCollation always orders nulls last.
  final ImmutableList<RelCollation> leftCollations = mq.collations(left);
  switch (joinType) {
  case SEMI:
  case ANTI:
  case INNER:
  case LEFT:
    return leftCollations;
  case RIGHT:
  case FULL:
    for (RelCollation collation : leftCollations) {
      for (RelFieldCollation field : collation.getFieldCollations()) {
        if (!(RelFieldCollation.NullDirection.LAST == field.nullDirection)) {
          return ImmutableList.of();
        }
      }
    }
    return leftCollations;
  }
  return ImmutableList.of();
}
 

public void rewriteRel(Sort rel) {
  RelCollation oldCollation = rel.getCollation();
  final RelNode oldChild = rel.getInput();
  final RelNode newChild = getNewForOldRel(oldChild);
  final Mappings.TargetMapping mapping =
      getNewForOldInputMapping(oldChild);

  // validate
  for (RelFieldCollation field : oldCollation.getFieldCollations()) {
    int oldInput = field.getFieldIndex();
    RelDataType sortFieldType =
        oldChild.getRowType().getFieldList().get(oldInput).getType();
    if (sortFieldType.isStruct()) {
      // TODO jvs 10-Feb-2005
      throw Util.needToImplement("sorting on structured types");
    }
  }
  RelCollation newCollation = RexUtil.apply(mapping, oldCollation);
  Sort newRel =
      LogicalSort.create(newChild, newCollation, rel.offset, rel.fetch);
  setNewForOldRel(rel, newRel);
}
 

private List<RelCollation> deduceMonotonicity(SqlValidatorTable table) {
  final RelDataType rowType = table.getRowType();
  final List<RelCollation> collationList = new ArrayList<>();

  // Deduce which fields the table is sorted on.
  int i = -1;
  for (RelDataTypeField field : rowType.getFieldList()) {
    ++i;
    final SqlMonotonicity monotonicity =
        table.getMonotonicity(field.getName());
    if (monotonicity != SqlMonotonicity.NOT_MONOTONIC) {
      final RelFieldCollation.Direction direction =
          monotonicity.isDecreasing()
              ? RelFieldCollation.Direction.DESCENDING
              : RelFieldCollation.Direction.ASCENDING;
      collationList.add(
          RelCollations.of(new RelFieldCollation(i, direction)));
    }
  }
  return collationList;
}
 

/** Infer the implicit correlation from the unrestricted clustering keys.
 *
 * @return The collation of the filtered results
 */
public RelCollation getImplicitCollation() {
  // No collation applies if we aren't restricted to a single partition
  if (!isSinglePartition()) {
    return RelCollations.EMPTY;
  }

  // Pull out the correct fields along with their original collations
  List<RelFieldCollation> fieldCollations = new ArrayList<>();
  for (int i = restrictedClusteringKeys; i < clusteringKeys.size(); i++) {
    int fieldIndex = fieldNames.indexOf(clusteringKeys.get(i));
    RelFieldCollation.Direction direction = implicitFieldCollations.get(i).getDirection();
    fieldCollations.add(new RelFieldCollation(fieldIndex, direction));
  }

  return RelCollations.of(fieldCollations);
}
 

@Override
public Prel prepareForLateralUnnestPipeline(List<RelNode> children) {
  List<RelFieldCollation> relFieldCollations = Lists.newArrayList();
  relFieldCollations.add(new RelFieldCollation(0,
                          RelFieldCollation.Direction.ASCENDING, RelFieldCollation.NullDirection.FIRST));
  for (RelFieldCollation fieldCollation : this.collation.getFieldCollations()) {
    relFieldCollations.add(new RelFieldCollation(fieldCollation.getFieldIndex() + 1,
            fieldCollation.direction, fieldCollation.nullDirection));
  }

  RelCollation collationTrait = RelCollationImpl.of(relFieldCollations);
  RelTraitSet traits = RelTraitSet.createEmpty()
                                  .replace(this.getTraitSet().getTrait(DrillDistributionTraitDef.INSTANCE))
                                  .replace(collationTrait)
                                  .replace(DRILL_PHYSICAL);

  return this.copy(traits, children.get(0), collationTrait, this.offset, this.fetch);
}
 

/**
 * Determine whether there is mapping between project input and output fields.
 * Bail out if sort relies on non-trivial expressions.
 */
private static boolean isCollationOnTrivialExpr(
    List<RexNode> projects, RelDataTypeFactory typeFactory,
    Mappings.TargetMapping map, RelFieldCollation fc, boolean passDown) {
  final int index = fc.getFieldIndex();
  int target = map.getTargetOpt(index);
  if (target < 0) {
    return false;
  }

  final RexNode node = passDown ? projects.get(index) : projects.get(target);
  if (node.isA(SqlKind.CAST)) {
    // Check whether it is a monotonic preserving cast
    final RexCall cast = (RexCall) node;
    RelFieldCollation newFieldCollation = Objects.requireNonNull(RexUtil.apply(map, fc));
    final RexCallBinding binding =
        RexCallBinding.create(typeFactory, cast,
            ImmutableList.of(RelCollations.of(newFieldCollation)));
    if (cast.getOperator().getMonotonicity(binding)
        == SqlMonotonicity.NOT_MONOTONIC) {
      return false;
    }
  }

  return true;
}
 

/**
 * Create a RelCollation that has partition-by as the leading keys followed by order-by keys
 * @param window The window specification
 * @return a RelCollation with {partition-by keys, order-by keys}
 */
private RelCollation getCollation(Window.Group window) {
    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : BitSets.toIter(window.keys)) {
        fields.add(new RelFieldCollation(group));
    }

    fields.addAll(window.orderKeys.getFieldCollations());

    return RelCollations.of(fields);
}
 

private List<DistributionField> getDistributionFieldsFromCollation(Window.Group window) {
    List<DistributionField> distFields = Lists.newArrayList();

    for (RelFieldCollation relField : window.collation().getFieldCollations()) {
        DistributionField field = new DistributionField(relField.getFieldIndex());
        distFields.add(field);
    }

    return distFields;
}
 

@Override
public RelWriter explainTerms(RelWriter pw) {
  super.explainTerms(pw);
  for (Ord<RelFieldCollation> ord : Ord.zip(this.fieldCollation.getFieldCollations())) {
    pw.item("sort" + ord.i, ord.e);
  }
  return pw;
}
 

/**
 * Applies a mapping to a list of field collations.
 *
 * @param mapping         Mapping
 * @param fieldCollations Field collations
 * @return collations with mapping applied
 */
public static List<RelFieldCollation> applyFields(
    Mappings.TargetMapping mapping,
    List<RelFieldCollation> fieldCollations) {
  final List<RelFieldCollation> newFieldCollations = new ArrayList<>();
  for (RelFieldCollation fieldCollation : fieldCollations) {
    RelFieldCollation newFieldCollation = apply(mapping, fieldCollation);
    if (newFieldCollation == null) {
      break;
    }
    newFieldCollations.add(newFieldCollation);
  }
  return newFieldCollations;
}
 

@Override
public void addOrderItem(List<SqlNode> orderByList,
                         RelFieldCollation field) {
  // Do not attempt to order by literals as databases such as SQL Server
  // disallow ORDER BY on const expressions.
  if (!(context.toSql(field) instanceof SqlLiteral)) {
    super.addOrderItem(orderByList, field);
  }
}
 

private RexNode createCollation(RexNode node, RelFieldCollation.Direction direction,
		RelFieldCollation.NullDirection nullDirection, Set<SqlKind> kinds) {
	switch (node.getKind()) {
		case DESCENDING:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), RelFieldCollation.Direction.DESCENDING,
					nullDirection, kinds);
		case NULLS_FIRST:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), direction,
					RelFieldCollation.NullDirection.FIRST, kinds);
		case NULLS_LAST:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), direction,
					RelFieldCollation.NullDirection.LAST, kinds);
		default:
			if (nullDirection == null) {
				// Set the null direction if not specified.
				// Consistent with HIVE/SPARK/MYSQL/BLINK-RUNTIME.
				if (FlinkPlannerImpl.defaultNullCollation()
						.last(direction.equals(RelFieldCollation.Direction.DESCENDING))) {
					kinds.add(SqlKind.NULLS_LAST);
				} else {
					kinds.add(SqlKind.NULLS_FIRST);
				}
			}
			return node;
	}
}
 

private int direction(RelFieldCollation fieldCollation) {
  switch (fieldCollation.getDirection()) {
  case DESCENDING:
  case STRICTLY_DESCENDING:
    return -1;
  case ASCENDING:
  case STRICTLY_ASCENDING:
  default:
    return 1;
  }
}
 

public RelFieldCollation.NullDirection getNullDirection() {
  return right.contains(SqlKind.NULLS_LAST)
      ? RelFieldCollation.NullDirection.LAST
      : right.contains(SqlKind.NULLS_FIRST)
          ? RelFieldCollation.NullDirection.FIRST
          : getDirection().defaultNullDirection();
}