Examples with ImmutableBitSet org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet used on opensource projects

Example 36 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class RelFieldTrimmer method trimFields.

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
 * {@link org.apache.calcite.rel.logical.LogicalFilter}.
 */
public TrimResult trimFields(Filter filter, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
    final RelDataType rowType = filter.getRowType();
    final int fieldCount = rowType.getFieldCount();
    final RexNode conditionExpr = filter.getCondition();
    final RelNode input = filter.getInput();
    // We use the fields used by the consumer, plus any fields used in the
    // filter.
    final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>(extraFields);
    RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
    inputFinder.inputBitSet.addAll(fieldsUsed);
    conditionExpr.accept(inputFinder);
    final ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();
    // Create input with trimmed columns.
    TrimResult trimResult = trimChild(filter, input, inputFieldsUsed, inputExtraFields);
    RelNode newInput = trimResult.left;
    final Mapping inputMapping = trimResult.right;
    // there's nothing we can do.
    if (newInput == input && fieldsUsed.cardinality() == fieldCount) {
        return result(filter, Mappings.createIdentity(fieldCount));
    }
    // Build new project expressions, and populate the mapping.
    final RexVisitor<RexNode> shuttle = new RexPermuteInputsShuttle(inputMapping, newInput);
    RexNode newConditionExpr = conditionExpr.accept(shuttle);
    // Use copy rather than relBuilder so that correlating variables get set.
    relBuilder.push(filter.copy(filter.getTraitSet(), newInput, newConditionExpr));
    // needs them for its condition.
    return result(relBuilder.build(), inputMapping);
}

Example 37 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class RelFieldTrimmer method trimChild.

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

Example 38 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class RelFieldTrimmer method trimFields.

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
 * {@link org.apache.calcite.rel.core.Sort}.
 */
public TrimResult trimFields(Sort sort, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
    final RelDataType rowType = sort.getRowType();
    final int fieldCount = rowType.getFieldCount();
    final RelCollation collation = sort.getCollation();
    final RelNode input = sort.getInput();
    // We use the fields used by the consumer, plus any fields used as sort
    // keys.
    final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
    for (RelFieldCollation field : collation.getFieldCollations()) {
        inputFieldsUsed.set(field.getFieldIndex());
    }
    // Create input with trimmed columns.
    final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
    TrimResult trimResult = trimChild(sort, input, inputFieldsUsed.build(), inputExtraFields);
    RelNode newInput = trimResult.left;
    final Mapping inputMapping = trimResult.right;
    // there's nothing we can do.
    if (newInput == input && inputMapping.isIdentity() && fieldsUsed.cardinality() == fieldCount) {
        return result(sort, Mappings.createIdentity(fieldCount));
    }
    // leave the Sort unchanged in case we have dynamic limits
    if (sort.offset instanceof RexDynamicParam || sort.fetch instanceof RexDynamicParam) {
        return result(sort, inputMapping);
    }
    relBuilder.push(newInput);
    final int offset = sort.offset == null ? 0 : RexLiteral.intValue(sort.offset);
    final int fetch = sort.fetch == null ? -1 : RexLiteral.intValue(sort.fetch);
    final ImmutableList<RexNode> fields = relBuilder.fields(RexUtil.apply(inputMapping, collation));
    relBuilder.sortLimit(offset, fetch, fields);
    // needs them for its condition.
    return result(relBuilder.build(), inputMapping);
}

Example 39 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class SqlValidatorUtil method convertGroupSet.

/**
 * Analyzes a GROUPING SETS item in a GROUP BY clause.
 */
private static void convertGroupSet(SqlValidatorScope scope, GroupAnalyzer groupAnalyzer, ImmutableList.Builder<ImmutableBitSet> builder, SqlNode groupExpr) {
    switch(groupExpr.getKind()) {
        case GROUPING_SETS:
            final SqlCall call = (SqlCall) groupExpr;
            for (SqlNode node : call.getOperandList()) {
                convertGroupSet(scope, groupAnalyzer, builder, node);
            }
            return;
        case ROW:
            final List<ImmutableBitSet> bitSets = analyzeGroupTuple(scope, groupAnalyzer, ((SqlCall) groupExpr).getOperandList());
            builder.add(ImmutableBitSet.union(bitSets));
            return;
        case ROLLUP:
        case CUBE:
            {
                // GROUPING SETS ( (a), ROLLUP(c,b), CUBE(d,e) )
                // is EQUIVALENT to
                // GROUPING SETS ( (a), (c,b), (b) ,(), (d,e), (d), (e) ).
                // Expand all ROLLUP/CUBE nodes
                List<ImmutableBitSet> operandBitSet = analyzeGroupTuple(scope, groupAnalyzer, ((SqlCall) groupExpr).getOperandList());
                switch(groupExpr.getKind()) {
                    case ROLLUP:
                        builder.addAll(rollup(operandBitSet));
                        return;
                    default:
                        builder.addAll(cube(operandBitSet));
                        return;
                }
            }
        default:
            builder.add(analyzeGroupExpr(scope, groupAnalyzer, groupExpr));
            return;
    }
}

Example 40 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class SqlValidatorUtil method getOrdinalBitSet.

/**
 * Gets the bit-set to the column ordinals in the source for columns that
 * intersect in the target.
 *
 * @param sourceRowType The source upon which to ordinate the bit set.
 * @param indexToField  The map of ordinals to target fields.
 */
public static ImmutableBitSet getOrdinalBitSet(RelDataType sourceRowType, Map<Integer, RelDataTypeField> indexToField) {
    ImmutableBitSet source = ImmutableBitSet.of(Lists.transform(sourceRowType.getFieldList(), new RelDataTypeField.ToFieldIndex()));
    ImmutableBitSet target = ImmutableBitSet.of(indexToField.keySet());
    return source.intersect(target);
}