Java Code Examples for org.apache.calcite.rel.RelCollation#getFieldCollations()

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

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

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

/**
 * This function can be reused when a Join's traits pass-down shall only
 * pass through collation to left input.
 *
 * @param required required trait set for the join
 * @param joinType the join type
 * @param leftInputFieldCount number of field count of left join input
 * @param joinTraitSet trait set of the join
 */
static Pair<RelTraitSet, List<RelTraitSet>> passThroughTraitsForJoin(
    RelTraitSet required, JoinRelType joinType,
    int leftInputFieldCount, RelTraitSet joinTraitSet) {
  RelCollation collation = required.getCollation();
  if (collation == null
      || collation == RelCollations.EMPTY
      || joinType == JoinRelType.FULL
      || joinType == JoinRelType.RIGHT) {
    return null;
  }

  for (RelFieldCollation fc : collation.getFieldCollations()) {
    // If field collation belongs to right input: cannot push down collation.
    if (fc.getFieldIndex() >= leftInputFieldCount) {
      return null;
    }
  }

  RelTraitSet passthroughTraitSet = joinTraitSet.replace(collation);
  return Pair.of(passthroughTraitSet,
      ImmutableList.of(
          passthroughTraitSet,
          passthroughTraitSet.replace(RelCollations.EMPTY)));
}
 

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

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

private PlannerOp planSort(EnumerableSort op, RelDataType rowType) {
    PlannerOp input = convertRelNode(op.getInput(), rowType, false, false);
    RelCollation collation = op.getCollation();
    List<RelFieldCollation> fieldCollations = collation.getFieldCollations();
    boolean[] directions = new boolean[fieldCollations.size()];
    int[] fields = new int[fieldCollations.size()];
    int i = 0;
    for (RelFieldCollation col : fieldCollations) {
        RelFieldCollation.Direction direction = col.getDirection();
        int index = col.getFieldIndex();
        directions[i] = direction == RelFieldCollation.Direction.ASCENDING
                || direction == RelFieldCollation.Direction.STRICTLY_ASCENDING;
        fields[i++] = index;
    }
    return new SortOp(input, directions, fields);

}
 

/**
 * Trims the fields of an input relational expression.
 *
 * @param rel        Relational expression
 * @param input      Input relational expression, whose fields to trim
 * @param fieldsUsed Bitmap of fields needed by the consumer
 * @return New relational expression and its field mapping
 */
protected TrimResult trimChild(
    RelNode rel,
    RelNode input,
    final ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {
  final ImmutableBitSet.Builder fieldsUsedBuilder = fieldsUsed.rebuild();

  // Fields that define the collation cannot be discarded.
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
  final ImmutableList<RelCollation> collations = mq.collations(input);
  for (RelCollation collation : collations) {
    for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
      fieldsUsedBuilder.set(fieldCollation.getFieldIndex());
    }
  }

  // Correlating variables are a means for other relational expressions to use
  // fields.
  for (final CorrelationId correlation : rel.getVariablesSet()) {
    rel.accept(
        new CorrelationReferenceFinder() {
          protected RexNode handle(RexFieldAccess fieldAccess) {
            final RexCorrelVariable v =
                (RexCorrelVariable) fieldAccess.getReferenceExpr();
            if (v.id.equals(correlation)) {
              fieldsUsedBuilder.set(fieldAccess.getField().getIndex());
            }
            return fieldAccess;
          }
        });
  }

  return dispatchTrimFields(input, fieldsUsedBuilder.build(), extraFields);
}
 

@Override public RelNode passThrough(final RelTraitSet required) {
  RelCollation collation = required.getCollation();
  if (collation == null || collation.isDefault()) {
    return null;
  }

  // A Values with 0 or 1 rows can be ordered by any collation.
  if (tuples.size() > 1) {
    Ordering<List<RexLiteral>> ordering = null;
    // Generate ordering comparator according to the required collations.
    for (RelFieldCollation fc : collation.getFieldCollations()) {
      Ordering<List<RexLiteral>> comparator = RelMdCollation.comparator(fc);
      if (ordering == null) {
        ordering = comparator;
      } else {
        ordering = ordering.compound(comparator);
      }
    }
    // Check whether the tuples are sorted by required collations.
    if (!ordering.isOrdered(tuples)) {
      return null;
    }
  }

  // The tuples order satisfies the collation, we just create a new
  // relnode with required collation info.
  return copy(traitSet.replace(collation), ImmutableList.of());
}
 

/** Check if it is possible to exploit native CQL sorting for a given collation.
 *
 * @return True if it is possible to achieve this sort in Cassandra
 */
private boolean collationsCompatible(RelCollation sortCollation,
    RelCollation implicitCollation) {
  List<RelFieldCollation> sortFieldCollations = sortCollation.getFieldCollations();
  List<RelFieldCollation> implicitFieldCollations = implicitCollation.getFieldCollations();

  if (sortFieldCollations.size() > implicitFieldCollations.size()) {
    return false;
  }
  if (sortFieldCollations.size() == 0) {
    return true;
  }

  // Check if we need to reverse the order of the implicit collation
  boolean reversed = reverseDirection(sortFieldCollations.get(0).getDirection())
      == implicitFieldCollations.get(0).getDirection();

  for (int i = 0; i < sortFieldCollations.size(); i++) {
    RelFieldCollation sorted = sortFieldCollations.get(i);
    RelFieldCollation implied = implicitFieldCollations.get(i);

    // Check that the fields being sorted match
    if (sorted.getFieldIndex() != implied.getFieldIndex()) {
      return false;
    }

    // Either all fields must be sorted in the same direction
    // or the opposite direction based on whether we decided
    // if the sort direction should be reversed above
    RelFieldCollation.Direction sortDirection = sorted.getDirection();
    RelFieldCollation.Direction implicitDirection = implied.getDirection();
    if ((!reversed && sortDirection != implicitDirection)
        || (reversed && reverseDirection(sortDirection) != implicitDirection)) {
      return false;
    }
  }

  return true;
}
 

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 com.google.common.collect.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 com.google.common.collect.ImmutableList.of();
					}
				}
			}
			return leftCollations;
	}
	return com.google.common.collect.ImmutableList.of();
}
 

public static List<Ordering> getOrdering(RelCollation collation, RelDataType rowType) {
  List<Ordering> orderExpr = Lists.newArrayList();

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

  for (RelFieldCollation fc: collation.getFieldCollations() ) {
    FieldReference fr = new FieldReference(childFields.get(fc.getFieldIndex()));
    orderExpr.add(new Ordering(fc.getDirection(), fr, fc.nullDirection));
  }

  return orderExpr;
}
 

public static List<Ordering> getOrdering(RelCollation collation, RelDataType rowType) {
    List<Ordering> orderExpr = Lists.newArrayList();

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

    for (RelFieldCollation fc : collation.getFieldCollations()) {
        FieldReference fr = new FieldReference(childFields.get(fc.getFieldIndex()), ExpressionPosition.UNKNOWN,
                false);
        orderExpr.add(new Ordering(fc.getDirection(), fr, fc.nullDirection));
    }

    return orderExpr;
}
 

private RelCollation convertRelCollation(RelCollation src, Map<Integer, Integer> inToOut) {
  List<RelFieldCollation> newFields = Lists.newArrayList();

  for ( RelFieldCollation field : src.getFieldCollations()) {
    if (inToOut.containsKey(field.getFieldIndex())) {
      newFields.add(new RelFieldCollation(inToOut.get(field.getFieldIndex()), field.getDirection(), field.nullDirection));
    }
  }

  if (newFields.isEmpty()) {
    return RelCollations.of();
  } else {
    return RelCollations.of(newFields);
  }
}
 

/**
 * Trims the fields of an input relational expression.
 *
 * @param rel        Relational expression
 * @param input      Input relational expression, whose fields to trim
 * @param fieldsUsed Bitmap of fields needed by the consumer
 * @return New relational expression and its field mapping
 */
protected TrimResult trimChild(RelNode rel, RelNode input, final ImmutableBitSet fieldsUsed,
        Set<RelDataTypeField> extraFields) {
    final ImmutableBitSet.Builder fieldsUsedBuilder = fieldsUsed.rebuild();

    // Fields that define the collation cannot be discarded.
    final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
    final ImmutableList<RelCollation> collations = mq.collations(input);
    for (RelCollation collation : collations) {
        for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
            fieldsUsedBuilder.set(fieldCollation.getFieldIndex());
        }
    }

    // Correlating variables are a means for other relational expressions to use
    // fields.
    for (final CorrelationId correlation : rel.getVariablesSet()) {
        rel.accept(new CorrelationReferenceFinder() {
            @Override
            protected RexNode handle(RexFieldAccess fieldAccess) {
                final RexCorrelVariable v = (RexCorrelVariable) fieldAccess.getReferenceExpr();
                if (v.getCorrelationId().equals(correlation)) {
                    fieldsUsedBuilder.set(fieldAccess.getField().getIndex());
                }
                return fieldAccess;
            }
        });
    }

    return dispatchTrimFields(input, fieldsUsedBuilder.build(), extraFields);
}
 

/**
 * Applies a mapping to a collation list.
 *
 * @param mapping       Mapping
 * @param collationList Collation list
 * @return collation list with mapping applied to each field
 */
public static List<RelCollation> apply(
    Mappings.TargetMapping mapping,
    List<RelCollation> collationList) {
  final List<RelCollation> newCollationList = new ArrayList<>();
  for (RelCollation collation : collationList) {
    final List<RelFieldCollation> newFieldCollationList = new ArrayList<>();
    for (RelFieldCollation fieldCollation
        : collation.getFieldCollations()) {
      final RelFieldCollation newFieldCollation =
          apply(mapping, fieldCollation);
      if (newFieldCollation == null) {
        // This field is not mapped. Stop here. The leading edge
        // of the collation is still valid (although it's useless
        // if it's empty).
        break;
      }
      newFieldCollationList.add(newFieldCollation);
    }
    // Truncation to collations to their leading edge creates empty
    // and duplicate collations. Ignore these.
    if (!newFieldCollationList.isEmpty()) {
      final RelCollation newCollation =
          RelCollations.of(newFieldCollationList);
      if (!newCollationList.contains(newCollation)) {
        newCollationList.add(newCollation);
      }
    }
  }

  // REVIEW: There might be redundant collations in the list. For example,
  // in {(x), (x, y)}, (x) is redundant because it is a leading edge of
  // another collation in the list. Could remove redundant collations.

  return newCollationList;
}
 

public Expression generateComparator(RelCollation collation) {
  // int c;
  // c = Utilities.compare(v0, v1);
  // if (c != 0) return c; // or -c if descending
  // ...
  // return 0;
  BlockBuilder body = new BlockBuilder();
  final Type javaRowClass = Primitive.box(this.javaRowClass);
  final ParameterExpression parameterV0 =
      Expressions.parameter(javaRowClass, "v0");
  final ParameterExpression parameterV1 =
      Expressions.parameter(javaRowClass, "v1");
  final ParameterExpression parameterC =
      Expressions.parameter(int.class, "c");
  final int mod =
      collation.getFieldCollations().size() == 1 ? Modifier.FINAL : 0;
  body.add(Expressions.declare(mod, parameterC, null));
  for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
    final int index = fieldCollation.getFieldIndex();
    final RelDataType fieldType = rowType.getFieldList().get(index).getType();
    final Expression fieldComparator = generateCollatorExpression(fieldType.getCollation());
    Expression arg0 = fieldReference(parameterV0, index);
    Expression arg1 = fieldReference(parameterV1, index);
    switch (Primitive.flavor(fieldClass(index))) {
    case OBJECT:
      arg0 = EnumUtils.convert(arg0, Comparable.class);
      arg1 = EnumUtils.convert(arg1, Comparable.class);
    }
    final boolean nullsFirst =
        fieldCollation.nullDirection
            == RelFieldCollation.NullDirection.FIRST;
    final boolean descending =
        fieldCollation.getDirection()
            == RelFieldCollation.Direction.DESCENDING;
    body.add(
        Expressions.statement(
            Expressions.assign(
                parameterC,
                Expressions.call(
                    Utilities.class,
                    fieldNullable(index)
                        ? (nullsFirst != descending
                        ? "compareNullsFirst"
                        : "compareNullsLast")
                        : "compare",
                    Expressions.list(
                        arg0,
                        arg1)
                        .appendIfNotNull(fieldComparator)))));
    body.add(
        Expressions.ifThen(
            Expressions.notEqual(
                parameterC, Expressions.constant(0)),
            Expressions.return_(
                null,
                descending
                    ? Expressions.negate(parameterC)
                    : parameterC)));
  }
  body.add(
      Expressions.return_(null, Expressions.constant(0)));

  final List<MemberDeclaration> memberDeclarations =
      Expressions.list(
          Expressions.methodDecl(
              Modifier.PUBLIC,
              int.class,
              "compare",
              ImmutableList.of(parameterV0, parameterV1),
              body.toBlock()));

  if (EnumerableRules.BRIDGE_METHODS) {
    final ParameterExpression parameterO0 =
        Expressions.parameter(Object.class, "o0");
    final ParameterExpression parameterO1 =
        Expressions.parameter(Object.class, "o1");
    BlockBuilder bridgeBody = new BlockBuilder();
    bridgeBody.add(
        Expressions.return_(
            null,
            Expressions.call(
                Expressions.parameter(
                    Comparable.class, "this"),
                BuiltInMethod.COMPARATOR_COMPARE.method,
                Expressions.convert_(
                    parameterO0,
                    javaRowClass),
                Expressions.convert_(
                    parameterO1,
                    javaRowClass))));
    memberDeclarations.add(
        overridingMethodDecl(
            BuiltInMethod.COMPARATOR_COMPARE.method,
            ImmutableList.of(parameterO0, parameterO1),
            bridgeBody.toBlock()));
  }
  return Expressions.new_(
      Comparator.class,
      ImmutableList.of(),
      memberDeclarations);
}
 

/**
 * Build collation property for the 'upper' project, the one above the filter
 * @param projectRexs
 * @param inputCollation
 * @param indexInfo
 * @param collationFilterMap
 * @return the output RelCollation
 */
public static RelCollation buildCollationUpperProject(List<RexNode> projectRexs,
                                                      RelCollation inputCollation, FunctionalIndexInfo indexInfo,
                                                      Map<Integer, List<RexNode>> collationFilterMap) {
  List<RelFieldCollation> outputFieldCollations = Lists.newArrayList();

  if (inputCollation != null) {
    List<RelFieldCollation> inputFieldCollations = inputCollation.getFieldCollations();
    if (!indexInfo.hasFunctional()) {
      for (int projectExprIdx = 0; projectExprIdx < projectRexs.size(); projectExprIdx++) {
        RexNode n = projectRexs.get(projectExprIdx);
        if (n instanceof RexInputRef) {
          RexInputRef ref = (RexInputRef)n;
          boolean eligibleForCollation = true;
          int maxIndex = getIndexFromCollation(ref.getIndex(), inputFieldCollations);
          if (maxIndex < 0) {
            eligibleForCollation = false;
            continue;
          }
          // check if the prefix has equality conditions
          for (int i = 0; i < maxIndex; i++) {
            int fieldIdx = inputFieldCollations.get(i).getFieldIndex();
            List<RexNode> conditions = collationFilterMap != null ? collationFilterMap.get(fieldIdx) : null;
            if ((conditions == null || conditions.size() == 0) &&
                i < maxIndex-1) {
              // if an intermediate column has no filter condition, it would select all values
              // of that column, so a subsequent column cannot be eligible for collation
              eligibleForCollation = false;
              break;
            } else {
              for (RexNode r : conditions) {
                if (!(r.getKind() == SqlKind.EQUALS)) {
                  eligibleForCollation = false;
                  break;
                }
              }
            }
          }
          // for every projected expr, if it is eligible for collation, get the
          // corresponding field collation from the input
          if (eligibleForCollation) {
            for (RelFieldCollation c : inputFieldCollations) {
              if (ref.getIndex() == c.getFieldIndex()) {
                RelFieldCollation outFieldCollation = new RelFieldCollation(projectExprIdx, c.getDirection(), c.nullDirection);
                outputFieldCollations.add(outFieldCollation);
              }
            }
          }
        }
      }
    } else {
      // TODO: handle functional index
    }
  }
  return RelCollations.of(outputFieldCollations);
}
 

public TrimResult trimFields(
    SortExchange sortExchange,
    ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {
  final RelDataType rowType = sortExchange.getRowType();
  final int fieldCount = rowType.getFieldCount();
  final RelCollation collation = sortExchange.getCollation();
  final RelDistribution distribution = sortExchange.getDistribution();
  final RelNode input = sortExchange.getInput();

  // We use the fields used by the consumer, plus any fields used as sortExchange
  // keys.
  final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
  for (RelFieldCollation field : collation.getFieldCollations()) {
    inputFieldsUsed.set(field.getFieldIndex());
  }
  for (int keyIndex : distribution.getKeys()) {
    inputFieldsUsed.set(keyIndex);
  }

  // Create input with trimmed columns.
  final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
  TrimResult trimResult =
      trimChild(sortExchange, input, inputFieldsUsed.build(), inputExtraFields);
  RelNode newInput = trimResult.left;
  final Mapping inputMapping = trimResult.right;

  // If the input is unchanged, and we need to project all columns,
  // there's nothing we can do.
  if (newInput == input
      && inputMapping.isIdentity()
      && fieldsUsed.cardinality() == fieldCount) {
    return result(sortExchange, Mappings.createIdentity(fieldCount));
  }

  relBuilder.push(newInput);
  RelCollation newCollation = RexUtil.apply(inputMapping, collation);
  RelDistribution newDistribution = distribution.apply(inputMapping);
  relBuilder.sortExchange(newDistribution, newCollation);

  return result(relBuilder.build(), inputMapping);
}