org.apache.calcite.rel.RelCollations Java Examples

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

/** Creates a relational expression that reads from an input and throws
 *  all of the rows away.
 */
@Override
public RelBuilder empty() {
  final Frame frame = stack.pop();
  final RelNode input;
  // If the rel that we are limiting the output of a rel, we should just add a limit 0 on top.
  // If the rel that we are limiting is a Filter replace it as well since Filter does not
  // change the row type.
  if (!(frame.rel instanceof Filter)) {
    input = frame.rel;
  } else {
    input = frame.rel.getInput(0);
  }
  final RelNode sort = sortFactory.createSort(input, RelCollations.EMPTY,
    frame.rel.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(0)),
    frame.rel.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(0)));
  push(sort);
  return this;
}
 

@Test void testTraitSetDefault() {
  RelTraitSet traits = RelTraitSet.createEmpty();
  traits = traits.plus(Convention.NONE).plus(RelCollations.EMPTY);
  assertEquals(traits.size(), 2);
  assertTrue(traits.isDefault());
  traits = traits.replace(EnumerableConvention.INSTANCE);
  assertFalse(traits.isDefault());
  assertTrue(traits.isDefaultSansConvention());
  traits = traits.replace(RelCollations.of(0));
  assertFalse(traits.isDefault());
  assertFalse(traits.replace(Convention.NONE).isDefaultSansConvention());
  assertTrue(traits.getDefault().isDefault());
  traits = traits.getDefaultSansConvention();
  assertFalse(traits.isDefault());
  assertEquals(traits.getConvention(), EnumerableConvention.INSTANCE);
  assertTrue(traits.isDefaultSansConvention());
  assertEquals(traits.toString(), "ENUMERABLE.[]");
}
 

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

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

@Test void testSortExchangeFieldTrimmerWithSingletonDistribution() {
  final RelBuilder builder = RelBuilder.create(config().build());
  final RelNode root =
      builder.scan("EMP")
          .project(builder.field("EMPNO"), builder.field("ENAME"), builder.field("DEPTNO"))
          .sortExchange(RelDistributions.SINGLETON, RelCollations.of(0))
          .project(builder.field("EMPNO"), builder.field("ENAME"))
          .build();

  RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
  RelNode trimmed = fieldTrimmer.trim(root);

  final String expected = ""
      + "LogicalSortExchange(distribution=[single], collation=[[0]])\n"
      + "  LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
      + "    LogicalTableScan(table=[[scott, EMP]])\n";
  assertThat(trimmed, hasTree(expected));
}
 

/**
 * Presents a view of the {@link RexWinAggCall} list as a list of
 * {@link AggregateCall}.
 */
public List<AggregateCall> getAggregateCalls(Window windowRel) {
    final List<String> fieldNames = Util.skip(windowRel.getRowType().getFieldNames(),
            windowRel.getInput().getRowType().getFieldCount());
    return new AbstractList<AggregateCall>() {
        @Override
        public int size() {
            return aggCalls.size();
        }

        @Override
        public AggregateCall get(int index) {
            final RexWinAggCall aggCall = aggCalls.get(index);
            final SqlAggFunction op = (SqlAggFunction) aggCall.getOperator();
            return AggregateCall.create(op, aggCall.isDistinct(), false,
                    getProjectOrdinals(aggCall.getOperands()), -1, RelCollations.EMPTY, aggCall.getType(),
                    fieldNames.get(aggCall.getOrdinal()));
        }
    };
}
 

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

@Test void testSortExchangeFieldTrimmerWithEmptyCollation() {
  final RelBuilder builder = RelBuilder.create(config().build());
  final RelNode root =
      builder.scan("EMP")
          .project(builder.field("EMPNO"), builder.field("ENAME"), builder.field("DEPTNO"))
          .sortExchange(RelDistributions.hash(Lists.newArrayList(1)), RelCollations.EMPTY)
          .project(builder.field("EMPNO"), builder.field("ENAME"))
          .build();

  RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
  RelNode trimmed = fieldTrimmer.trim(root);

  final String expected = ""
      + "LogicalSortExchange(distribution=[hash[1]], collation=[[]])\n"
      + "  LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
      + "    LogicalTableScan(table=[[scott, EMP]])\n";
  assertThat(trimmed, hasTree(expected));
}
 

public RelNode convert(RelNode rel) {
  final LogicalAggregate agg = (LogicalAggregate) rel;
  if (!Aggregate.isSimple(agg)) {
    return null;
  }
  final RelTraitSet inputTraits = rel.getCluster()
      .traitSet().replace(EnumerableConvention.INSTANCE)
      .replace(
          RelCollations.of(
              ImmutableIntList.copyOf(
          agg.getGroupSet().asList())));
  final RelTraitSet selfTraits = inputTraits.replace(
      RelCollations.of(
      ImmutableIntList.identity(agg.getGroupSet().cardinality())));
  return new EnumerableSortedAggregate(
      rel.getCluster(),
      selfTraits,
      convert(agg.getInput(), inputTraits),
      agg.getGroupSet(),
      agg.getGroupSets(),
      agg.getAggCallList());
}
 

public String toString() {
  StringBuilder buf = new StringBuilder(aggFunction.toString());
  buf.append("(");
  if (distinct) {
    buf.append((argList.size() == 0) ? "DISTINCT" : "DISTINCT ");
  }
  int i = -1;
  for (Integer arg : argList) {
    if (++i > 0) {
      buf.append(", ");
    }
    buf.append("$");
    buf.append(arg);
  }
  buf.append(")");
  if (!collation.equals(RelCollations.EMPTY)) {
    buf.append(" WITHIN GROUP (");
    buf.append(collation);
    buf.append(")");
  }
  if (hasFilter()) {
    buf.append(" FILTER $");
    buf.append(filterArg);
  }
  return buf.toString();
}
 

static Pair<RelTraitSet, List<RelTraitSet>> passThroughTraitsForProject(
    RelTraitSet required,
    List<RexNode> exps,
    RelDataType inputRowType,
    RelDataTypeFactory typeFactory,
    RelTraitSet currentTraits) {
  final RelCollation collation = required.getCollation();
  if (collation == null || collation == RelCollations.EMPTY) {
    return null;
  }

  final Mappings.TargetMapping map =
      RelOptUtil.permutationIgnoreCast(
          exps, inputRowType);

  if (collation.getFieldCollations().stream().anyMatch(
      rc -> !isCollationOnTrivialExpr(exps, typeFactory,
          map, rc, true))) {
    return null;
  }

  final RelCollation newCollation = collation.apply(map);
  return Pair.of(currentTraits.replace(collation),
      ImmutableList.of(currentTraits.replace(newCollation)));
}
 

private AggregateCall toAggCall(RelInput relInput, Map<String, Object> jsonAggCall) {
  final Map<String, Object> aggMap = (Map) jsonAggCall.get("agg");
  final SqlAggFunction aggregation =
      relJson.toAggregation(aggMap);
  final Boolean distinct = (Boolean) jsonAggCall.get("distinct");
  @SuppressWarnings("unchecked")
  final List<Integer> operands = (List<Integer>) jsonAggCall.get("operands");
  final Integer filterOperand = (Integer) jsonAggCall.get("filter");
  final RelDataType type =
      relJson.toType(cluster.getTypeFactory(), jsonAggCall.get("type"));
  final String name = (String) jsonAggCall.get("name");
  return AggregateCall.create(aggregation, distinct, false, false, operands,
      filterOperand == null ? -1 : filterOperand,
      RelCollations.EMPTY,
      type, name);
}
 

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

  final RexBuilder rexBuilder = sort.getCluster().getRexBuilder();
  final List<RelFieldCollation> collationsList =
      sort.getCollation().getFieldCollations().stream()
          .filter(fc ->
              !predicates.constantMap.containsKey(
                  rexBuilder.makeInputRef(input, fc.getFieldIndex())))
          .collect(Collectors.toList());

  if (collationsList.size() == sort.collation.getFieldCollations().size()) {
    return;
  }

  // No active collations. Remove the sort completely
  if (collationsList.isEmpty() && sort.offset == null && sort.fetch == null) {
    call.transformTo(input);
    call.getPlanner().prune(sort);
    return;
  }

  final Sort result =
      sort.copy(sort.getTraitSet(), input, RelCollations.of(collationsList));
  call.transformTo(result);
  call.getPlanner().prune(sort);
}
 

private Statistic buildStatistic() {
  if (stats != null || primaryKey == -1) {
    return stats;
  }
  Direction dir = primaryKeyMonotonicity == INCREASING ? ASCENDING : DESCENDING;
  RelFieldCollation collation = new RelFieldCollation(primaryKey, dir, NullDirection.UNSPECIFIED);
  return Statistics.of(fields.size(), ImmutableList.of(ImmutableBitSet.of(primaryKey)),
      ImmutableList.of(RelCollations.of(collation)));
}
 

@Deprecated // to be removed before 2.0
public static AggregateCall create(SqlAggFunction aggFunction,
    boolean distinct, boolean approximate, List<Integer> argList,
    int filterArg, int groupCount,
    RelNode input, RelDataType type, String name) {
  return create(aggFunction, distinct, approximate, false, argList,
      filterArg, RelCollations.EMPTY, groupCount, input, type, name);
}
 

@Deprecated // to be removed before 2.0
public static AggregateCall create(SqlAggFunction aggFunction,
    boolean distinct, List<Integer> argList, int filterArg, RelDataType type,
    String name) {
  return create(aggFunction, distinct, false, false, argList, filterArg,
      RelCollations.EMPTY, type, name);
}
 

/**
 * Creates an AggregateCall.
 *
 * @param aggFunction Aggregate function
 * @param distinct    Whether distinct
 * @param argList     List of ordinals of arguments
 * @param type        Result type
 * @param name        Name (may be null)
 */
@Deprecated // to be removed before 2.0
public AggregateCall(
    SqlAggFunction aggFunction,
    boolean distinct,
    List<Integer> argList,
    RelDataType type,
    String name) {
  this(aggFunction, distinct, false, false,
      argList, -1, RelCollations.EMPTY, type, name);
}
 

@Deprecated // to be removed before 2.0
public static AggregateCall create(SqlAggFunction aggFunction,
    boolean distinct, List<Integer> argList, int groupCount, RelNode input,
    RelDataType type, String name) {
  return create(aggFunction, distinct, false, false, argList, -1,
      RelCollations.EMPTY, groupCount, input, type, name);
}
 

/**
 * Unit test for {@link org.apache.calcite.rel.externalize.RelJsonWriter} on
 * a simple tree of relational expressions, consisting of a table and a
 * project including window expressions.
 */
@Test void testWriter() {
  String s =
      Frameworks.withPlanner((cluster, relOptSchema, rootSchema) -> {
        rootSchema.add("hr",
            new ReflectiveSchema(new JdbcTest.HrSchema()));
        LogicalTableScan scan =
            LogicalTableScan.create(cluster,
                relOptSchema.getTableForMember(
                    Arrays.asList("hr", "emps")),
                ImmutableList.of());
        final RexBuilder rexBuilder = cluster.getRexBuilder();
        LogicalFilter filter =
            LogicalFilter.create(scan,
                rexBuilder.makeCall(
                    SqlStdOperatorTable.EQUALS,
                    rexBuilder.makeFieldAccess(
                        rexBuilder.makeRangeReference(scan),
                        "deptno", true),
                    rexBuilder.makeExactLiteral(BigDecimal.TEN)));
        final RelJsonWriter writer = new RelJsonWriter();
        final RelDataType bigIntType =
            cluster.getTypeFactory().createSqlType(SqlTypeName.BIGINT);
        LogicalAggregate aggregate =
            LogicalAggregate.create(filter,
                ImmutableList.of(),
                ImmutableBitSet.of(0),
                null,
                ImmutableList.of(
                    AggregateCall.create(SqlStdOperatorTable.COUNT,
                        true, false, false, ImmutableList.of(1), -1,
                        RelCollations.EMPTY, bigIntType, "c"),
                    AggregateCall.create(SqlStdOperatorTable.COUNT,
                        false, false, false, ImmutableList.of(), -1,
                        RelCollations.EMPTY, bigIntType, "d")));
        aggregate.explain(writer);
        return writer.asString();
      });
  assertThat(s, is(XX));
}
 

public RelNode convert(RelNode rel) {
  final Sort sort = (Sort) rel;
  final RelTraitSet traitSet =
      sort.getTraitSet().replace(out)
          .replace(sort.getCollation());
  return new MongoSort(rel.getCluster(), traitSet,
      convert(sort.getInput(), traitSet.replace(RelCollations.EMPTY)),
      sort.getCollation(), sort.offset, sort.fetch);
}
 

@Override
public RelNode convert(RelNode rel) {
  final Sort sort = (Sort) rel;
  final RelTraitSet traitSet = sort.getTraitSet().replace(out).replace(sort.getCollation());
  return new SolrSort(
      rel.getCluster(),
      traitSet,
      convert(sort.getInput(), traitSet.replace(RelCollations.EMPTY)),
      sort.getCollation(),
      sort.offset,
      sort.fetch);
}
 

private RelNode createSortPlan(RelDistribution distribution) {
  final FrameworkConfig config = RelBuilderTest.config().build();
  final RelBuilder builder = RelBuilder.create(config);
  return builder.scan("EMP")
          .sortExchange(distribution,
              RelCollations.of(0))
          .build();
}
 

@Override
protected RelNode visitChild(RelNode parent, int i, RelNode child) {
  RelNode newParent = parent;
  if (parent instanceof JdbcRelImpl) {
    transformer.setTraitSet(parent.getTraitSet().plus(DistributionTrait.ANY).plus(RelCollations.EMPTY));
    newParent = parent.accept(transformer);
  }
  return super.visitChild(newParent, i, newParent.getInput(i));
}
 

private static RelCollation getCollation(RelTraitSet set, List<Integer> keys) {
  return set.canonize(RelCollations.of(FluentIterable.from(keys)
      .transform(new Function<Integer, RelFieldCollation>() {
        @Override
        public RelFieldCollation apply(Integer input) {
          return new RelFieldCollation(input);
        }
      }).toList()));
}
 

/**
 * Creates collation based on streaming agg groupset
 *
 * @param groupSet
 * @return
 */
public static RelCollation collation(ImmutableBitSet groupSet) {
  // The collation of a streaming agg is based on its groups
  List<RelFieldCollation> fields = IntStream.range(0, groupSet.cardinality())
      .mapToObj(RelFieldCollation::new)
      .collect(Collectors.toList());
  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);
  }
 

@Override
public T read(final Kryo kryo, final Input input, final Class<T> type) {
  final boolean isKnown = kryo.readObject(input, Boolean.class);
  final T result;
  if (isKnown) {
    final Integer pos = kryo.readObject(input, Integer.class);
    result = (T) (pos == 0 ? RelCollations.EMPTY:RelCollations.PRESERVE);
  } else {
    result = super.read(kryo, input, type);
  }

  final T normalized = (T) result.getTraitDef().canonize(result);
  kryo.reference(normalized);
  return normalized;
}