Examples with RelFieldCollation org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation used on opensource projects

Example 11 with RelFieldCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation in project calcite by apache.

the class SqlToRelConverter method convertMatchRecognize.

protected void convertMatchRecognize(Blackboard bb, SqlCall call) {
    final SqlMatchRecognize matchRecognize = (SqlMatchRecognize) call;
    final SqlValidatorNamespace ns = validator.getNamespace(matchRecognize);
    final SqlValidatorScope scope = validator.getMatchRecognizeScope(matchRecognize);
    final Blackboard matchBb = createBlackboard(scope, null, false);
    final RelDataType rowType = ns.getRowType();
    // convert inner query, could be a table name or a derived table
    SqlNode expr = matchRecognize.getTableRef();
    convertFrom(matchBb, expr);
    final RelNode input = matchBb.root;
    // PARTITION BY
    final SqlNodeList partitionList = matchRecognize.getPartitionList();
    final List<RexNode> partitionKeys = new ArrayList<>();
    for (SqlNode partition : partitionList) {
        RexNode e = matchBb.convertExpression(partition);
        partitionKeys.add(e);
    }
    // ORDER BY
    final SqlNodeList orderList = matchRecognize.getOrderList();
    final List<RelFieldCollation> orderKeys = new ArrayList<>();
    for (SqlNode order : orderList) {
        final RelFieldCollation.Direction direction;
        switch(order.getKind()) {
            case DESCENDING:
                direction = RelFieldCollation.Direction.DESCENDING;
                order = ((SqlCall) order).operand(0);
                break;
            case NULLS_FIRST:
            case NULLS_LAST:
                throw new AssertionError();
            default:
                direction = RelFieldCollation.Direction.ASCENDING;
                break;
        }
        final RelFieldCollation.NullDirection nullDirection = validator.getDefaultNullCollation().last(desc(direction)) ? RelFieldCollation.NullDirection.LAST : RelFieldCollation.NullDirection.FIRST;
        RexNode e = matchBb.convertExpression(order);
        orderKeys.add(new RelFieldCollation(((RexInputRef) e).getIndex(), direction, nullDirection));
    }
    final RelCollation orders = cluster.traitSet().canonize(RelCollations.of(orderKeys));
    // convert pattern
    final Set<String> patternVarsSet = new HashSet<>();
    SqlNode pattern = matchRecognize.getPattern();
    final SqlBasicVisitor<RexNode> patternVarVisitor = new SqlBasicVisitor<RexNode>() {

        @Override
        public RexNode visit(SqlCall call) {
            List<SqlNode> operands = call.getOperandList();
            List<RexNode> newOperands = Lists.newArrayList();
            for (SqlNode node : operands) {
                newOperands.add(node.accept(this));
            }
            return rexBuilder.makeCall(validator.getUnknownType(), call.getOperator(), newOperands);
        }

        @Override
        public RexNode visit(SqlIdentifier id) {
            assert id.isSimple();
            patternVarsSet.add(id.getSimple());
            return rexBuilder.makeLiteral(id.getSimple());
        }

        @Override
        public RexNode visit(SqlLiteral literal) {
            if (literal instanceof SqlNumericLiteral) {
                return rexBuilder.makeExactLiteral(BigDecimal.valueOf(literal.intValue(true)));
            } else {
                return rexBuilder.makeLiteral(literal.booleanValue());
            }
        }
    };
    final RexNode patternNode = pattern.accept(patternVarVisitor);
    SqlLiteral interval = matchRecognize.getInterval();
    RexNode intervalNode = null;
    if (interval != null) {
        intervalNode = matchBb.convertLiteral(interval);
    }
    // convert subset
    final SqlNodeList subsets = matchRecognize.getSubsetList();
    final Map<String, TreeSet<String>> subsetMap = Maps.newHashMap();
    for (SqlNode node : subsets) {
        List<SqlNode> operands = ((SqlCall) node).getOperandList();
        SqlIdentifier left = (SqlIdentifier) operands.get(0);
        patternVarsSet.add(left.getSimple());
        SqlNodeList rights = (SqlNodeList) operands.get(1);
        final TreeSet<String> list = new TreeSet<String>();
        for (SqlNode right : rights) {
            assert right instanceof SqlIdentifier;
            list.add(((SqlIdentifier) right).getSimple());
        }
        subsetMap.put(left.getSimple(), list);
    }
    SqlNode afterMatch = matchRecognize.getAfter();
    if (afterMatch == null) {
        afterMatch = SqlMatchRecognize.AfterOption.SKIP_TO_NEXT_ROW.symbol(SqlParserPos.ZERO);
    }
    final RexNode after;
    if (afterMatch instanceof SqlCall) {
        List<SqlNode> operands = ((SqlCall) afterMatch).getOperandList();
        SqlOperator operator = ((SqlCall) afterMatch).getOperator();
        assert operands.size() == 1;
        SqlIdentifier id = (SqlIdentifier) operands.get(0);
        assert patternVarsSet.contains(id.getSimple()) : id.getSimple() + " not defined in pattern";
        RexNode rex = rexBuilder.makeLiteral(id.getSimple());
        after = rexBuilder.makeCall(validator.getUnknownType(), operator, ImmutableList.of(rex));
    } else {
        after = matchBb.convertExpression(afterMatch);
    }
    matchBb.setPatternVarRef(true);
    // convert measures
    final ImmutableMap.Builder<String, RexNode> measureNodes = ImmutableMap.builder();
    for (SqlNode measure : matchRecognize.getMeasureList()) {
        List<SqlNode> operands = ((SqlCall) measure).getOperandList();
        String alias = ((SqlIdentifier) operands.get(1)).getSimple();
        RexNode rex = matchBb.convertExpression(operands.get(0));
        measureNodes.put(alias, rex);
    }
    // convert definitions
    final ImmutableMap.Builder<String, RexNode> definitionNodes = ImmutableMap.builder();
    for (SqlNode def : matchRecognize.getPatternDefList()) {
        List<SqlNode> operands = ((SqlCall) def).getOperandList();
        String alias = ((SqlIdentifier) operands.get(1)).getSimple();
        RexNode rex = matchBb.convertExpression(operands.get(0));
        definitionNodes.put(alias, rex);
    }
    final SqlLiteral rowsPerMatch = matchRecognize.getRowsPerMatch();
    final boolean allRows = rowsPerMatch != null && rowsPerMatch.getValue() == SqlMatchRecognize.RowsPerMatchOption.ALL_ROWS;
    matchBb.setPatternVarRef(false);
    final RelFactories.MatchFactory factory = RelFactories.DEFAULT_MATCH_FACTORY;
    final RelNode rel = factory.createMatch(input, patternNode, rowType, matchRecognize.getStrictStart().booleanValue(), matchRecognize.getStrictEnd().booleanValue(), definitionNodes.build(), measureNodes.build(), after, subsetMap, allRows, partitionKeys, orders, intervalNode);
    bb.setRoot(rel, false);
}

Example 12 with RelFieldCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation in project calcite by apache.

the class RelBuilder method fields.

/**
 * Returns references to fields for a given collation.
 */
public ImmutableList<RexNode> fields(RelCollation collation) {
    final ImmutableList.Builder<RexNode> nodes = ImmutableList.builder();
    for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
        RexNode node = field(fieldCollation.getFieldIndex());
        switch(fieldCollation.direction) {
            case DESCENDING:
                node = desc(node);
        }
        switch(fieldCollation.nullDirection) {
            case FIRST:
                node = nullsFirst(node);
                break;
            case LAST:
                node = nullsLast(node);
                break;
        }
        nodes.add(node);
    }
    return nodes.build();
}

Example 13 with RelFieldCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation in project calcite by apache.

the class RelBuilder method sortLimit.

/**
 * Creates a {@link Sort} by a list of expressions, with limit and offset.
 *
 * @param offset Number of rows to skip; non-positive means don't skip any
 * @param fetch Maximum number of rows to fetch; negative means no limit
 * @param nodes Sort expressions
 */
public RelBuilder sortLimit(int offset, int fetch, Iterable<? extends RexNode> nodes) {
    final List<RelFieldCollation> fieldCollations = new ArrayList<>();
    final List<RexNode> originalExtraNodes = fields();
    final List<RexNode> extraNodes = new ArrayList<>(originalExtraNodes);
    for (RexNode node : nodes) {
        final RelFieldCollation collation = collation(node, RelFieldCollation.Direction.ASCENDING, null, extraNodes);
        if (!RelCollations.ordinals(fieldCollations).contains(collation.getFieldIndex())) {
            fieldCollations.add(collation);
        }
    }
    final RexNode offsetNode = offset <= 0 ? null : literal(offset);
    final RexNode fetchNode = fetch < 0 ? null : literal(fetch);
    if (offsetNode == null && fetch == 0) {
        return empty();
    }
    if (offsetNode == null && fetchNode == null && fieldCollations.isEmpty()) {
        // sort is trivial
        return this;
    }
    final boolean addedFields = extraNodes.size() > originalExtraNodes.size();
    if (fieldCollations.isEmpty()) {
        assert !addedFields;
        RelNode top = peek();
        if (top instanceof Sort) {
            final Sort sort2 = (Sort) top;
            if (sort2.offset == null && sort2.fetch == null) {
                replaceTop(sort2.getInput());
                final RelNode sort = sortFactory.createSort(peek(), sort2.collation, offsetNode, fetchNode);
                replaceTop(sort);
                return this;
            }
        }
        if (top instanceof Project) {
            final Project project = (Project) top;
            if (project.getInput() instanceof Sort) {
                final Sort sort2 = (Sort) project.getInput();
                if (sort2.offset == null && sort2.fetch == null) {
                    final RelNode sort = sortFactory.createSort(sort2.getInput(), sort2.collation, offsetNode, fetchNode);
                    replaceTop(projectFactory.createProject(sort, project.getProjects(), Pair.right(project.getNamedProjects())));
                    return this;
                }
            }
        }
    }
    if (addedFields) {
        project(extraNodes);
    }
    final RelNode sort = sortFactory.createSort(peek(), RelCollations.of(fieldCollations), offsetNode, fetchNode);
    replaceTop(sort);
    if (addedFields) {
        project(originalExtraNodes);
    }
    return this;
}

Example 14 with RelFieldCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation in project calcite by apache.

the class RelBuilder method match.

/**
 * Creates a {@link org.apache.calcite.rel.core.Match}.
 */
public RelBuilder match(RexNode pattern, boolean strictStart, boolean strictEnd, Map<String, RexNode> patternDefinitions, Iterable<? extends RexNode> measureList, RexNode after, Map<String, ? extends SortedSet<String>> subsets, boolean allRows, Iterable<? extends RexNode> partitionKeys, Iterable<? extends RexNode> orderKeys, RexNode interval) {
    final List<RelFieldCollation> fieldCollations = new ArrayList<>();
    for (RexNode orderKey : orderKeys) {
        final RelFieldCollation.Direction direction;
        switch(orderKey.getKind()) {
            case DESCENDING:
                direction = RelFieldCollation.Direction.DESCENDING;
                orderKey = ((RexCall) orderKey).getOperands().get(0);
                break;
            case NULLS_FIRST:
            case NULLS_LAST:
                throw new AssertionError();
            default:
                direction = RelFieldCollation.Direction.ASCENDING;
                break;
        }
        final RelFieldCollation.NullDirection nullDirection = direction.defaultNullDirection();
        final RexInputRef ref = (RexInputRef) orderKey;
        fieldCollations.add(new RelFieldCollation(ref.getIndex(), direction, nullDirection));
    }
    final RelDataTypeFactory.Builder typeBuilder = cluster.getTypeFactory().builder();
    for (RexNode partitionKey : partitionKeys) {
        typeBuilder.add(partitionKey.toString(), partitionKey.getType());
    }
    if (allRows) {
        for (RexNode orderKey : orderKeys) {
            if (!typeBuilder.nameExists(orderKey.toString())) {
                typeBuilder.add(orderKey.toString(), orderKey.getType());
            }
        }
        final RelDataType inputRowType = peek().getRowType();
        for (RelDataTypeField fs : inputRowType.getFieldList()) {
            if (!typeBuilder.nameExists(fs.getName())) {
                typeBuilder.add(fs);
            }
        }
    }
    final ImmutableMap.Builder<String, RexNode> measures = ImmutableMap.builder();
    for (RexNode measure : measureList) {
        List<RexNode> operands = ((RexCall) measure).getOperands();
        String alias = operands.get(1).toString();
        typeBuilder.add(alias, operands.get(0).getType());
        measures.put(alias, operands.get(0));
    }
    final RelNode match = matchFactory.createMatch(peek(), pattern, typeBuilder.build(), strictStart, strictEnd, patternDefinitions, measures.build(), after, subsets, allRows, ImmutableList.copyOf(partitionKeys), RelCollations.of(fieldCollations), interval);
    stack.push(new Frame(match));
    return this;
}

Example 15 with RelFieldCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelFieldCollation 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);
}