Java Code Examples for org.apache.calcite.rel.RelCollations#of()

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

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 + "]");
}
 

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

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

public RelCollation toCollation(
    List<Map<String, Object>> jsonFieldCollations) {
  final List<RelFieldCollation> fieldCollations = new ArrayList<>();
  for (Map<String, Object> map : jsonFieldCollations) {
    fieldCollations.add(toFieldCollation(map));
  }
  return RelCollations.of(fieldCollations);
}
 

/**
 * 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 RelCollation getInputCollation(AggregateRel rel){

    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : rel.getGroupSet()) {
      fields.add(new RelFieldCollation(group));
    }
    return RelCollations.of(fields);
  }
 

private RelCollation getOutputCollation(AggregateRel rel){

    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group = 0; group < rel.getGroupSet().cardinality(); group++) {
      fields.add(new RelFieldCollation(group));
    }
    return RelCollations.of(fields);
  }
 

/**
 * Applies a mapping to a collation.
 *
 * @param mapping   Mapping
 * @param collation Collation
 * @return collation with mapping applied
 */
public static RelCollation apply(
    Mappings.TargetMapping mapping,
    RelCollation collation) {
  List<RelFieldCollation> fieldCollations =
      applyFields(mapping, collation.getFieldCollations());
  return fieldCollations.equals(collation.getFieldCollations())
      ? collation
      : RelCollations.of(fieldCollations);
}
 

private RelCollation getCollation(List<Integer> keys) {
  List<RelFieldCollation> fields = Lists.newArrayList();
  for (int key : keys) {
    fields.add(new RelFieldCollation(key));
  }
  return RelCollations.of(fields);
}
 

public RelCollation toCollation(
    List<Map<String, Object>> jsonFieldCollations) {
  final List<RelFieldCollation> fieldCollations = new ArrayList<>();
  for (Map<String, Object> map : jsonFieldCollations) {
    fieldCollations.add(toFieldCollation(map));
  }
  return RelCollations.of(fieldCollations);
}
 

private RelCollation getCollation(DrillAggregateRel rel) {

    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : BitSets.toIter(rel.getGroupSet())) {
      fields.add(new RelFieldCollation(group));
    }
    return RelCollations.of(fields);
  }
 

private RelCollation getCollation(List<Integer> keys){
  List<RelFieldCollation> fields = Lists.newArrayList();
  for (int key : keys) {
    fields.add(new RelFieldCollation(key));
  }
  return RelCollations.of(fields);
}
 

@Override
@JsonIgnore
public RelCollation getCollation() {
  if (indexCollationContext != null) {
    return RelCollations.of(indexCollationContext.relFieldCollations);
  }
  return null;
}
 

/**
 * Build collation property for the 'lower' project, the one closer to the Scan
 * @param projectRexs
 * @param input
 * @param indexInfo
 * @return the output RelCollation
 */
public static RelCollation buildCollationLowerProject(List<RexNode> projectRexs, RelNode input, FunctionalIndexInfo indexInfo) {
  // if leading fields of index are here, add them to RelCollation
  List<RelFieldCollation> newFields = Lists.newArrayList();
  if (!indexInfo.hasFunctional()) {
    Map<LogicalExpression, Integer> projectExprs = Maps.newLinkedHashMap();
    DrillParseContext parserContext = new DrillParseContext(PrelUtil.getPlannerSettings(input.getCluster()));
    int idx=0;
    for (RexNode rex : projectRexs) {
      projectExprs.put(DrillOptiq.toDrill(parserContext, input, rex), idx);
      idx++;
    }
    int idxFieldCount = 0;
    for (LogicalExpression expr : indexInfo.getIndexDesc().getIndexColumns()) {
      if (!projectExprs.containsKey(expr)) {
        break;
      }
      RelFieldCollation.Direction dir = indexInfo.getIndexDesc().getCollation().getFieldCollations().get(idxFieldCount).direction;
      if ( dir == null) {
        break;
      }
      newFields.add(new RelFieldCollation(projectExprs.get(expr), dir,
          RelFieldCollation.NullDirection.UNSPECIFIED));
    }
    idxFieldCount++;
  } else {
    // TODO: handle functional index
  }

  return RelCollations.of(newFields);
}
 

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

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

@Override public void onMatch(RelOptRuleCall call) {
  final SortExchange sortExchange = call.rel(0);
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelNode input = sortExchange.getInput();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(input);
  if (predicates == null) {
    return;
  }

  final Set<Integer> constants = new HashSet<>();
  predicates.constantMap.keySet().forEach(key -> {
    if (key instanceof RexInputRef) {
      constants.add(((RexInputRef) key).getIndex());
    }
  });

  if (constants.isEmpty()) {
    return;
  }

  List<Integer> distributionKeys = new ArrayList<>();
  boolean distributionSimplified = false;
  boolean hashDistribution = sortExchange.getDistribution().getType()
      == RelDistribution.Type.HASH_DISTRIBUTED;
  if (hashDistribution) {
    distributionKeys = simplifyDistributionKeys(
        sortExchange.getDistribution(), constants);
    distributionSimplified =
        distributionKeys.size() != sortExchange.getDistribution().getKeys()
            .size();
  }

  final List<RelFieldCollation> fieldCollations = sortExchange
      .getCollation().getFieldCollations().stream().filter(
          fc -> !constants.contains(fc.getFieldIndex()))
       .collect(Collectors.toList());

  boolean collationSimplified =
       fieldCollations.size() != sortExchange.getCollation()
           .getFieldCollations().size();
  if (distributionSimplified
       || collationSimplified) {
    RelDistribution distribution = distributionSimplified
        ? (distributionKeys.isEmpty()
            ? RelDistributions.SINGLETON
            : RelDistributions.hash(distributionKeys))
        : sortExchange.getDistribution();
    RelCollation collation = collationSimplified
        ? RelCollations.of(fieldCollations)
        : sortExchange.getCollation();

    call.transformTo(call.builder()
        .push(sortExchange.getInput())
        .sortExchange(distribution, collation)
        .build());
    call.getPlanner().prune(sortExchange);
  }
}
 

/**
 * Applies a mapping to a collation.
 *
 * @param mapping   Mapping
 * @param collation Collation
 * @return collation with mapping applied
 */
public static RelCollation apply(Mappings.TargetMapping mapping, RelCollation collation) {
    List<RelFieldCollation> fieldCollations = applyFields(mapping, collation.getFieldCollations());
    return fieldCollations.equals(collation.getFieldCollations()) ? collation : RelCollations.of(fieldCollations);
}