Java Code Examples for org.apache.calcite.util.ImmutableBitSet#rebuild()

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

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

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

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

  // Create input with trimmed columns.
  final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
  final TrimResult trimResult =
      trimChild(exchange, input, inputFieldsUsed.build(), inputExtraFields);
  final 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(exchange, Mappings.createIdentity(fieldCount));
  }

  relBuilder.push(newInput);
  final RelDistribution newDistribution = distribution.apply(inputMapping);
  relBuilder.exchange(newDistribution);

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

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

public InputFinder(Set<RelDataTypeField> extraFields,
    ImmutableBitSet initialBits) {
  this(extraFields, initialBits.rebuild());
}