org.apache.calcite.util.mapping.Mappings Java Examples

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

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

/** Equivalent to
 * {@link RelOptUtil#createProject(org.apache.calcite.rel.RelNode, java.util.List)}
 * for {@link MutableRel}. */
public static MutableRel createProject(final MutableRel child,
    final List<Integer> posList) {
  final RelDataType rowType = child.rowType;
  if (Mappings.isIdentity(posList, rowType.getFieldCount())) {
    return child;
  }
  return MutableProject.of(
      RelOptUtil.permute(child.cluster.getTypeFactory(), rowType,
          Mappings.bijection(posList)),
      child,
      new AbstractList<RexNode>() {
        public int size() {
          return posList.size();
        }

        public RexNode get(int index) {
          final int pos = posList.get(index);
          return RexInputRef.of(pos, rowType);
        }
      });
}
 

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 a mapping of a set of project expressions.
 *
 * <p>The mapping is an inverse surjection.
 * Every target has a source field, but
 * a source field may appear as zero, one, or more target fields.
 * Thus you can safely call
 * {@link org.apache.calcite.util.mapping.Mappings.TargetMapping#getTarget(int)}.
 *
 * @param inputFieldCount Number of input fields
 * @param projects Project expressions
 * @return Mapping of a set of project expressions, or null if projection is
 * not a mapping
 */
public static Mappings.TargetMapping getMapping(int inputFieldCount,
    List<? extends RexNode> projects) {
  if (inputFieldCount < projects.size()) {
    return null; // surjection is not possible
  }
  Mappings.TargetMapping mapping =
      Mappings.create(MappingType.INVERSE_SURJECTION,
          inputFieldCount, projects.size());
  for (Ord<RexNode> exp : Ord.<RexNode>zip(projects)) {
    if (!(exp.e instanceof RexInputRef)) {
      return null;
    }
    mapping.set(((RexInputRef) exp.e).getIndex(), exp.i);
  }
  return mapping;
}
 

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 a mapping of a set of project expressions.
 *
 * <p>The mapping is an inverse surjection.
 * Every target has a source field, but no
 * source has more than one target.
 * Thus you can safely call
 * {@link org.apache.calcite.util.mapping.Mappings.TargetMapping#getSourceOpt(int)}.
 *
 * @param inputFieldCount Number of input fields
 * @param projects Project expressions
 * @return Mapping of a set of project expressions, or null if projection is
 * not a mapping
 */
public static Mappings.TargetMapping getMapping(int inputFieldCount,
    List<? extends RexNode> projects) {
  if (inputFieldCount < projects.size()) {
    return null; // surjection is not possible
  }
  Mappings.TargetMapping mapping =
      Mappings.create(MappingType.INVERSE_SURJECTION,
          inputFieldCount, projects.size());
  for (Ord<RexNode> exp : Ord.<RexNode>zip(projects)) {
    if (!(exp.e instanceof RexInputRef)) {
      return null;
    }

    int source = ((RexInputRef) exp.e).getIndex();
    if (mapping.getTargetOpt(source) != -1) {
      return null;
    }
    mapping.set(source, exp.i);
  }
  return mapping;
}
 

/**
 * Returns a permutation describing where output fields come from. In
 * the returned map, value of {@code map.getTargetOpt(i)} is {@code n} if
 * field {@code i} projects input field {@code n} or applies a cast on
 * {@code n}, -1 if it is another expression.
 */
public static Mappings.TargetMapping permutationIgnoreCast(
    List<RexNode> nodes,
    RelDataType inputRowType) {
  final Mappings.TargetMapping mapping =
      Mappings.create(
          MappingType.PARTIAL_FUNCTION,
          nodes.size(),
          inputRowType.getFieldCount());
  for (Ord<RexNode> node : Ord.zip(nodes)) {
    if (node.e instanceof RexInputRef) {
      mapping.set(
          node.i,
          ((RexInputRef) node.e).getIndex());
    } else if (node.e.isA(SqlKind.CAST)) {
      final RexNode operand = ((RexCall) node.e).getOperands().get(0);
      if (operand instanceof RexInputRef) {
        mapping.set(node.i, ((RexInputRef) operand).getIndex());
      }
    }
  }
  return mapping;
}
 

public TrimResult trimFields(
  LimitRel limit,
  ImmutableBitSet fieldsUsed,
  Set<RelDataTypeField> extraFields
) {
  final RelDataType rowType = limit.getRowType();
  final int fieldCount = rowType.getFieldCount();
  final RelNode input = limit.getInput();

  final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
  TrimResult trimResult =
    trimChild(limit, input, fieldsUsed, inputExtraFields);
  RelNode newInput = trimResult.left;
  final Mapping inputMapping = trimResult.right;

  if (newInput == input
    && inputMapping.isIdentity()
    && fieldsUsed.cardinality() == fieldCount) {
    return result(limit, Mappings.createIdentity(fieldCount));
  }
  return result(limit.copy(newInput.getTraitSet(), ImmutableList.of(newInput)), inputMapping);
}
 

/**
 * Returns a permutation describing where output fields come from. In
 * the returned map, value of {@code map.getTargetOpt(i)} is {@code n} if
 * field {@code i} projects input field {@code n}, -1 if it is an
 * expression.
 */
public static Mappings.TargetMapping permutation(
    List<RexNode> nodes,
    RelDataType inputRowType) {
  final Mappings.TargetMapping mapping =
      Mappings.create(
          MappingType.PARTIAL_FUNCTION,
          nodes.size(),
          inputRowType.getFieldCount());
  for (Ord<RexNode> node : Ord.zip(nodes)) {
    if (node.e instanceof RexInputRef) {
      mapping.set(
          node.i,
          ((RexInputRef) node.e).getIndex());
    }
  }
  return mapping;
}
 

private List<RelBuilder.AggCall> getNewAggCallList(
	Aggregate oldAggregate,
	RelBuilder relBuilder,
	Mapping mapping) {

	final List<RelBuilder.AggCall> newAggCallList = new ArrayList<>();

	for (AggregateCall aggCall : oldAggregate.getAggCallList()) {
		final RexNode filterArg = aggCall.filterArg < 0 ? null
			: relBuilder.field(Mappings.apply(mapping, aggCall.filterArg));
		newAggCallList.add(
			relBuilder
				.aggregateCall(
					aggCall.getAggregation(),
					relBuilder.fields(Mappings.apply2(mapping, aggCall.getArgList())))
				.distinct(aggCall.isDistinct())
				.filter(filterArg)
				.approximate(aggCall.isApproximate())
				.sort(relBuilder.fields(aggCall.collation))
				.as(aggCall.name));
	}
	return newAggCallList;
}
 

private void initializeMapping() {
  nextMapping = Mappings.create(MappingType.PARTIAL_FUNCTION,
      nSysFields + nFieldsLeft + nFieldsRight,
      nSysFields + nFieldsLeft + nFieldsRight);
  for (int i = 0; i < columnSets.length; i++) {
    BitSet c = columnSets[i];
    int t = c.nextSetBit(iterationIdx[i]);
    if (t < 0) {
      nextMapping = null;
      return;
    }
    nextMapping.set(columns[i], t);
    iterationIdx[i] = t + 1;
  }
}
 

private void initializeMapping() {
  nextMapping = Mappings.create(MappingType.PARTIAL_FUNCTION,
      nSysFields + nFieldsLeft + nFieldsRight,
      nSysFields + nFieldsLeft + nFieldsRight);
  for (int i = 0; i < columnSets.length; i++) {
    BitSet c = columnSets[i];
    int t = c.nextSetBit(iterationIdx[i]);
    if (t < 0) {
      nextMapping = null;
      return;
    }
    nextMapping.set(columns[i], t);
    iterationIdx[i] = t + 1;
  }
}
 

@Test void testProjectLeadingEdge() {
  final RelBuilder builder = RelBuilder.create(config().build());
  RelNode root =
      builder.scan("EMP")
          .project(builder.fields(Mappings.bijection(Arrays.asList(0, 1, 2))))
          .build();
  final String expected = "LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2])\n"
      + "  LogicalTableScan(table=[[scott, EMP]])\n";
  assertThat(root, hasTree(expected));
}
 

/** Equivalent to
 * {@link RelOptUtil#createProject(org.apache.calcite.rel.RelNode, java.util.List)}
 * for {@link MutableRel}. */
public static MutableRel createProject(final MutableRel child,
    final List<Integer> posList) {
  final RelDataType rowType = child.rowType;
  if (Mappings.isIdentity(posList, rowType.getFieldCount())) {
    return child;
  }
  final Mapping mapping =
      Mappings.create(
          MappingType.INVERSE_SURJECTION,
          rowType.getFieldCount(),
          posList.size());
  for (int i = 0; i < posList.size(); i++) {
    mapping.set(posList.get(i), i);
  }
  return MutableProject.of(
      RelOptUtil.permute(child.cluster.getTypeFactory(), rowType, mapping),
      child,
      new AbstractList<RexNode>() {
        public int size() {
          return posList.size();
        }

        public RexNode get(int index) {
          final int pos = posList.get(index);
          return RexInputRef.of(pos, rowType);
        }
      });
}
 

/** Creates a copy of this collation that changes the ordinals of input
 * fields. */
public static RelCollation permute(RelCollation collation,
    Mappings.TargetMapping mapping) {
  return of(
      Util.transform(collation.getFieldCollations(),
          fc -> fc.withFieldIndex(mapping.getTarget(fc.getFieldIndex()))));
}
 

protected Mapping createMapping(ImmutableBitSet fieldsUsed, int fieldCount) {
  final Mapping mapping =
      Mappings.create(
          MappingType.INVERSE_SURJECTION,
          fieldCount,
          fieldsUsed.cardinality());
  int i = 0;
  for (int field : fieldsUsed) {
    mapping.set(field, i++);
  }
  return mapping;
}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

/**
 * Extract projects from the Aggregate and return the index mapping between the new projects
 * and it's input.
 */
private Mapping extractProjectsAndMapping(
	Aggregate aggregate,
	RelNode input,
	RelBuilder relBuilder) {

	// Compute which input fields are used.
	final ImmutableBitSet.Builder inputFieldsUsed = getInputFieldUsed(aggregate, input);

	final List<RexNode> projects = new ArrayList<>();
	final Mapping mapping =
		Mappings.create(MappingType.INVERSE_SURJECTION,
			aggregate.getInput().getRowType().getFieldCount(),
			inputFieldsUsed.cardinality());
	int j = 0;
	for (int i : inputFieldsUsed.build()) {
		projects.add(relBuilder.field(i));
		mapping.set(i, j++);
	}

	if (input instanceof Project) {
		// this will not create trivial projects
		relBuilder.project(projects);
	} else {
		relBuilder.project(projects, Collections.emptyList(), true);
	}

	return mapping;
}
 

/** Creates a project with a dummy column, to protect the parts of the system
 * that cannot handle a relational expression with no columns.
 *
 * @param fieldCount Number of fields in the original relational expression
 * @param input Trimmed input
 * @return Dummy project, or null if no dummy is required
 */
protected TrimResult dummyProject(int fieldCount, RelNode input) {
    final RelOptCluster cluster = input.getCluster();
    final Mapping mapping = Mappings.create(MappingType.INVERSE_SURJECTION, fieldCount, 1);
    if (input.getRowType().getFieldCount() == 1) {
        // Input already has one field (and may in fact be a dummy project we
        // created for the child). We can't do better.
        return result(input, mapping);
    }
    final RexLiteral expr = cluster.getRexBuilder().makeExactLiteral(BigDecimal.ZERO);
    relBuilder.push(input);
    relBuilder.project(ImmutableList.<RexNode> of(expr), ImmutableList.of("DUMMY"));
    return result(relBuilder.build(), mapping);
}
 

public boolean isRefTrivial() {
  if (SqlKind.DML.contains(kind)) {
    // DML statements return a single count column.
    // The validated type is of the SELECT.
    // Still, we regard the mapping as trivial.
    return true;
  }
  final RelDataType inputRowType = rel.getRowType();
  return Mappings.isIdentity(Pair.left(fields), inputRowType.getFieldCount());
}
 

@Override
public TrimResult trimFields(
    SetOp setOp,
    ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {
  if(!setOp.all) {
    return result(setOp, Mappings.createIdentity(setOp.getRowType().getFieldCount()));
  }
  return super.trimFields(setOp, fieldsUsed, extraFields);
}
 

/**
 * Applies a mapping to this traitSet.
 *
 * @param mapping   Mapping
 * @return traitSet with mapping applied
 */
public RelTraitSet apply(Mappings.TargetMapping mapping) {
  RelTrait[] newTraits = new RelTrait[traits.length];
  for (int i = 0; i < traits.length; i++) {
    newTraits[i] = traits[i].apply(mapping);
  }
  return cache.getOrAdd(new RelTraitSet(cache, newTraits));
}
 

private static ProjectFilterTable of3(RexNode condition,
    Mappings.TargetMapping mapping, RelNode node) {
  if (node instanceof TableScan) {
    return new ProjectFilterTable(condition, mapping,
        (TableScan) node);
  } else {
    return null;
  }
}
 

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for {@link ScanCrel}.
 */
@SuppressWarnings("unused")
public TrimResult trimFields(
    ScanCrel crel,
    ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {

  if(fieldsUsed.cardinality() == crel.getRowType().getFieldCount()) {
    return result(crel, Mappings.createIdentity(crel.getRowType().getFieldCount()));
  }

  if(fieldsUsed.cardinality() == 0) {
    // do something similar to dummy project but avoid using a scan field. This ensures the scan
    // does a skipAll operation rather than projectin a useless column.
    final RelOptCluster cluster = crel.getCluster();
    final Mapping mapping = Mappings.create(MappingType.INVERSE_SURJECTION, crel.getRowType().getFieldCount(), 1);
    final RexLiteral expr = cluster.getRexBuilder().makeExactLiteral(BigDecimal.ZERO);
    builder.push(crel);
    builder.project(ImmutableList.<RexNode>of(expr), ImmutableList.of("DUMMY"));
    return result(builder.build(), mapping);
  }

  final List<SchemaPath> paths = new ArrayList<>();
  final Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION, crel.getRowType().getFieldCount(), fieldsUsed.cardinality());
  int index = 0;
  for(int i : fieldsUsed) {
    paths.add(SchemaPath.getSimplePath(crel.getRowType().getFieldList().get(i).getName()));
    m.set(i, index);
    index++;
  }

  ScanCrel newCrel = crel.cloneWithProject(paths);
  return result(newCrel, m);
}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

/**
 * 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) {
        newFieldCollations.add(apply(mapping, fieldCollation));
    }
    return newFieldCollations;
}
 

/**
 * Returns a mapping between old and new fields.
 *
 * @param oldRel Old relational expression
 * @return Mapping between fields of old and new
 */
private Mappings.TargetMapping getNewForOldInputMapping(RelNode oldRel) {
  final RelNode newRel = getNewForOldRel(oldRel);
  return Mappings.target(
      this::getNewForOldInput,
      oldRel.getRowType().getFieldCount(),
      newRel.getRowType().getFieldCount());
}
 

/** Creates a copy of this collation that changes the ordinals of input
 * fields. */
public static RelCollation permute(RelCollation collation,
    Mappings.TargetMapping mapping) {
  return of(
      Util.transform(collation.getFieldCollations(),
          fc -> fc.copy(mapping.getTarget(fc.getFieldIndex()))));
}
 

private static ProjectFilterTable of3(RexNode condition,
    Mappings.TargetMapping mapping, RelNode node) {
  if (node instanceof TableScan) {
    return new ProjectFilterTable(condition, mapping,
        (TableScan) node);
  } else {
    return null;
  }
}