Java Code Examples for org.apache.calcite.util.mapping.Mapping#set()

@Override
public void onMatch(RelOptRuleCall call) {
  FilterPrel filter = call.rel(0);
  NestedLoopJoinPrel join = call.rel(1);

  if ((join.getProjectedFields() == null) || join.getProjectedFields().cardinality() == join.getInputRowType().getFieldCount()) {
    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), filter.getCondition()), join.getProjectedFields()));
  } else {
    // Current filter condition is written based on projected fields on join. In order to push this filter down we need to rewrite filter condition
    final ImmutableBitSet topProjectedColumns = RelOptUtil.InputFinder.bits(filter.getCondition());
    final ImmutableBitSet bottomProjectedColumns = join.getProjectedFields();

    Mapping mapping = Mappings.create(MappingType.SURJECTION, join.getRowType().getFieldCount(), join.getInputRowType().getFieldCount());
    for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
      if (topProjectedColumns.get(ord.i)) {
        mapping.set(ord.i, ord.e);
      }
    }

    RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
    RexNode updatedCondition = shuttle.apply(filter.getCondition());

    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), updatedCondition), join.getProjectedFields()));
  }
}
 

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
  final RelNode input = agg.getInput();
  final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> aggPullUpPredicates = new ArrayList<>();

  ImmutableBitSet groupKeys = agg.getGroupSet();
  if (groupKeys.isEmpty()) {
    // "GROUP BY ()" can convert an empty relation to a non-empty relation, so
    // it is not valid to pull up predicates. In particular, consider the
    // predicate "false": it is valid on all input rows (trivially - there are
    // no rows!) but not on the output (there is one row).
    return RelOptPredicateList.EMPTY;
  }
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());

  int i = 0;
  for (int j : groupKeys) {
    m.set(j, i++);
  }

  for (RexNode r : inputInfo.pulledUpPredicates) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
    if (groupKeys.contains(rCols)) {
      r = r.accept(new RexPermuteInputsShuttle(m, input));
      aggPullUpPredicates.add(r);
    }
  }
  return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}
 

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

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

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

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

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

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
  final RelNode input = agg.getInput();
  final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> aggPullUpPredicates = new ArrayList<>();

  ImmutableBitSet groupKeys = agg.getGroupSet();
  if (groupKeys.isEmpty()) {
    // "GROUP BY ()" can convert an empty relation to a non-empty relation, so
    // it is not valid to pull up predicates. In particular, consider the
    // predicate "false": it is valid on all input rows (trivially - there are
    // no rows!) but not on the output (there is one row).
    return RelOptPredicateList.EMPTY;
  }
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());

  int i = 0;
  for (int j : groupKeys) {
    m.set(j, i++);
  }

  for (RexNode r : inputInfo.pulledUpPredicates) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
    if (groupKeys.contains(rCols)) {
      r = r.accept(new RexPermuteInputsShuttle(m, input));
      aggPullUpPredicates.add(r);
    }
  }
  return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}
 

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

/**
 * Infers predicates for a project.
 *
 * <ol>
 * <li>create a mapping from input to projection. Map only positions that
 * directly reference an input column.
 * <li>Expressions that only contain above columns are retained in the
 * Project's pullExpressions list.
 * <li>For e.g. expression 'a + e = 9' below will not be pulled up because 'e'
 * is not in the projection list.
 *
 * <blockquote><pre>
 * inputPullUpExprs:      {a &gt; 7, b + c &lt; 10, a + e = 9}
 * projectionExprs:       {a, b, c, e / 2}
 * projectionPullupExprs: {a &gt; 7, b + c &lt; 10}
 * </pre></blockquote>
 *
 * </ol>
 */
public RelOptPredicateList getPredicates(Project project,
    RelMetadataQuery mq) {
  final RelNode input = project.getInput();
  final RexBuilder rexBuilder = project.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> projectPullUpPredicates = new ArrayList<>();

  ImmutableBitSet.Builder columnsMappedBuilder = ImmutableBitSet.builder();
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(),
      project.getRowType().getFieldCount());

  for (Ord<RexNode> expr : Ord.zip(project.getProjects())) {
    if (expr.e instanceof RexInputRef) {
      int sIdx = ((RexInputRef) expr.e).getIndex();
      m.set(sIdx, expr.i);
      columnsMappedBuilder.set(sIdx);
    // Project can also generate constants. We need to include them.
    } else if (RexLiteral.isNullLiteral(expr.e)) {
      projectPullUpPredicates.add(
          rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL,
              rexBuilder.makeInputRef(project, expr.i)));
    } else if (RexUtil.isConstant(expr.e)) {
      final List<RexNode> args =
          ImmutableList.of(rexBuilder.makeInputRef(project, expr.i), expr.e);
      final SqlOperator op = args.get(0).getType().isNullable()
          || args.get(1).getType().isNullable()
          ? SqlStdOperatorTable.IS_NOT_DISTINCT_FROM
          : SqlStdOperatorTable.EQUALS;
      projectPullUpPredicates.add(rexBuilder.makeCall(op, args));
    }
  }

  // Go over childPullUpPredicates. If a predicate only contains columns in
  // 'columnsMapped' construct a new predicate based on mapping.
  final ImmutableBitSet columnsMapped = columnsMappedBuilder.build();
  for (RexNode r : inputInfo.pulledUpPredicates) {
    RexNode r2 = projectPredicate(rexBuilder, input, r, columnsMapped);
    if (!r2.isAlwaysTrue()) {
      r2 = r2.accept(new RexPermuteInputsShuttle(m, input));
      projectPullUpPredicates.add(r2);
    }
  }
  return RelOptPredicateList.of(rexBuilder, projectPullUpPredicates);
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  ProjectPrel project = call.rel(0);
  NestedLoopJoinPrel nlj = call.rel(1);

  ImmutableBitSet topProjectedColumns = InputFinder.bits(project.getProjects(), null);

  ImmutableBitSet bottomProjectedColumns = null;
  if (nlj.getProjectedFields() == null) {
    bottomProjectedColumns = ImmutableBitSet.range(nlj.getRowType().getFieldCount());
  } else {
    bottomProjectedColumns = nlj.getProjectedFields();
  }

  ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
  int field = 0;
  Mapping mapping = Mappings.create(MappingType.SURJECTION, bottomProjectedColumns.cardinality(), topProjectedColumns.cardinality());
  for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
    if (topProjectedColumns.get(ord.i)) {
      builder.set(ord.e);
      mapping.set(ord.i, field);
      field++;
    }
  }

  if (builder.cardinality() == 0) {
    //project at least one column
    builder.set(0);
  }

  ImmutableBitSet newJoinProjectedFields = builder.build();

  if (newJoinProjectedFields.equals(nlj.getProjectedFields())) {
    return;
  }

  RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
  List<RexNode> newProjects = shuttle.apply(project.getProjects());

  NestedLoopJoinPrel newJoin = (NestedLoopJoinPrel) nlj.copy(newJoinProjectedFields);
  ProjectPrel newProject = ProjectPrel.create(nlj.getCluster(), project.getTraitSet(), newJoin, newProjects, project.getRowType());
  call.transformTo(newProject);
}
 

/**
 * Infers predicates for a project.
 *
 * <ol>
 * <li>create a mapping from input to projection. Map only positions that
 * directly reference an input column.
 * <li>Expressions that only contain above columns are retained in the
 * Project's pullExpressions list.
 * <li>For e.g. expression 'a + e = 9' below will not be pulled up because 'e'
 * is not in the projection list.
 *
 * <blockquote><pre>
 * inputPullUpExprs:      {a &gt; 7, b + c &lt; 10, a + e = 9}
 * projectionExprs:       {a, b, c, e / 2}
 * projectionPullupExprs: {a &gt; 7, b + c &lt; 10}
 * </pre></blockquote>
 *
 * </ol>
 */
public RelOptPredicateList getPredicates(Project project,
    RelMetadataQuery mq) {
  final RelNode input = project.getInput();
  final RexBuilder rexBuilder = project.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> projectPullUpPredicates = new ArrayList<>();

  ImmutableBitSet.Builder columnsMappedBuilder = ImmutableBitSet.builder();
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(),
      project.getRowType().getFieldCount());

  for (Ord<RexNode> expr : Ord.zip(project.getProjects())) {
    if (expr.e instanceof RexInputRef) {
      int sIdx = ((RexInputRef) expr.e).getIndex();
      m.set(sIdx, expr.i);
      columnsMappedBuilder.set(sIdx);
    // Project can also generate constants. We need to include them.
    } else if (RexLiteral.isNullLiteral(expr.e)) {
      projectPullUpPredicates.add(
          rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL,
              rexBuilder.makeInputRef(project, expr.i)));
    } else if (RexUtil.isConstant(expr.e)) {
      final List<RexNode> args =
          ImmutableList.of(rexBuilder.makeInputRef(project, expr.i), expr.e);
      final SqlOperator op = args.get(0).getType().isNullable()
          || args.get(1).getType().isNullable()
          ? SqlStdOperatorTable.IS_NOT_DISTINCT_FROM
          : SqlStdOperatorTable.EQUALS;
      projectPullUpPredicates.add(rexBuilder.makeCall(op, args));
    }
  }

  // Go over childPullUpPredicates. If a predicate only contains columns in
  // 'columnsMapped' construct a new predicate based on mapping.
  final ImmutableBitSet columnsMapped = columnsMappedBuilder.build();
  for (RexNode r : inputInfo.pulledUpPredicates) {
    RexNode r2 = projectPredicate(rexBuilder, input, r, columnsMapped);
    if (!r2.isAlwaysTrue()) {
      r2 = r2.accept(new RexPermuteInputsShuttle(m, input));
      projectPullUpPredicates.add(r2);
    }
  }
  return RelOptPredicateList.of(rexBuilder, projectPullUpPredicates);
}