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

Example 96 with Pair

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

the class SqlTesterImpl method buildQuery2.

/**
 * Builds a query that extracts all literals as columns in an underlying
 * select.
 *
 * <p>For example,</p>
 *
 * <blockquote>{@code 1 < 5}</blockquote>
 *
 * <p>becomes</p>
 *
 * <blockquote>{@code SELECT p0 < p1
 * FROM (VALUES (1, 5)) AS t(p0, p1)}</blockquote>
 *
 * <p>Null literals don't have enough type information to be extracted.
 * We push down {@code CAST(NULL AS type)} but raw nulls such as
 * {@code CASE 1 WHEN 2 THEN 'a' ELSE NULL END} are left as is.</p>
 *
 * @param expression Scalar expression
 * @return Query that evaluates a scalar expression
 */
private String buildQuery2(String expression) {
    // "values (1 < 5)"
    // becomes
    // "select p0 < p1 from (values (1, 5)) as t(p0, p1)"
    SqlNode x;
    final String sql = "values (" + expression + ")";
    try {
        x = parseQuery(sql);
    } catch (SqlParseException e) {
        throw new RuntimeException(e);
    }
    final Collection<SqlNode> literalSet = new LinkedHashSet<>();
    x.accept(new SqlShuttle() {

        private final List<SqlOperator> ops = ImmutableList.of(SqlStdOperatorTable.LITERAL_CHAIN, SqlStdOperatorTable.LOCALTIME, SqlStdOperatorTable.LOCALTIMESTAMP, SqlStdOperatorTable.CURRENT_TIME, SqlStdOperatorTable.CURRENT_TIMESTAMP);

        @Override
        public SqlNode visit(SqlLiteral literal) {
            if (!isNull(literal) && literal.getTypeName() != SqlTypeName.SYMBOL) {
                literalSet.add(literal);
            }
            return literal;
        }

        @Override
        public SqlNode visit(SqlCall call) {
            final SqlOperator operator = call.getOperator();
            if (operator == SqlStdOperatorTable.CAST && isNull(call.operand(0))) {
                literalSet.add(call);
                return call;
            } else if (ops.contains(operator)) {
                // literal"
                return call;
            } else {
                return super.visit(call);
            }
        }

        private boolean isNull(SqlNode sqlNode) {
            return sqlNode instanceof SqlLiteral && ((SqlLiteral) sqlNode).getTypeName() == SqlTypeName.NULL;
        }
    });
    final List<SqlNode> nodes = new ArrayList<>(literalSet);
    Collections.sort(nodes, new Comparator<SqlNode>() {

        public int compare(SqlNode o1, SqlNode o2) {
            final SqlParserPos pos0 = o1.getParserPosition();
            final SqlParserPos pos1 = o2.getParserPosition();
            int c = -Utilities.compare(pos0.getLineNum(), pos1.getLineNum());
            if (c != 0) {
                return c;
            }
            return -Utilities.compare(pos0.getColumnNum(), pos1.getColumnNum());
        }
    });
    String sql2 = sql;
    final List<Pair<String, String>> values = new ArrayList<>();
    int p = 0;
    for (SqlNode literal : nodes) {
        final SqlParserPos pos = literal.getParserPosition();
        final int start = SqlParserUtil.lineColToIndex(sql, pos.getLineNum(), pos.getColumnNum());
        final int end = SqlParserUtil.lineColToIndex(sql, pos.getEndLineNum(), pos.getEndColumnNum()) + 1;
        String param = "p" + (p++);
        values.add(Pair.of(sql2.substring(start, end), param));
        sql2 = sql2.substring(0, start) + param + sql2.substring(end);
    }
    if (values.isEmpty()) {
        values.add(Pair.of("1", "p0"));
    }
    return "select " + sql2.substring("values (".length(), sql2.length() - 1) + " from (values (" + Util.commaList(Pair.left(values)) + ")) as t(" + Util.commaList(Pair.right(values)) + ")";
}

Example 97 with Pair

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

the class DiffRepository method update.

/**
 * Creates a new document with a given resource.
 *
 * <p>This method is synchronized, in case two threads are running test
 * cases of this test at the same time.
 *
 * @param testCaseName Test case name
 * @param resourceName Resource name
 * @param value        New value of resource
 */
private synchronized void update(String testCaseName, String resourceName, String value) {
    final List<Pair<String, Element>> map = new ArrayList<>();
    Element testCaseElement = getTestCaseElement(testCaseName, true, map);
    if (testCaseElement == null) {
        testCaseElement = doc.createElement(TEST_CASE_TAG);
        testCaseElement.setAttribute(TEST_CASE_NAME_ATTR, testCaseName);
        Node refElement = ref(testCaseName, map);
        root.insertBefore(testCaseElement, refElement);
    }
    Element resourceElement = getResourceElement(testCaseElement, resourceName, true);
    if (resourceElement == null) {
        resourceElement = doc.createElement(RESOURCE_TAG);
        resourceElement.setAttribute(RESOURCE_NAME_ATTR, resourceName);
        testCaseElement.appendChild(resourceElement);
    } else {
        removeAllChildren(resourceElement);
    }
    if (!value.equals("")) {
        resourceElement.appendChild(doc.createCDATASection(value));
    }
    // Write out the document.
    flushDoc();
}

Example 98 with Pair

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

the class AggregateExpandDistinctAggregatesRule method createSelectDistinct.

/**
 * Given an {@link org.apache.calcite.rel.core.Aggregate}
 * and the ordinals of the arguments to a
 * particular call to an aggregate function, creates a 'select distinct'
 * relational expression which projects the group columns and those
 * arguments but nothing else.
 *
 * <p>For example, given
 *
 * <blockquote>
 * <pre>select f0, count(distinct f1), count(distinct f2)
 * from t group by f0</pre>
 * </blockquote>
 *
 * <p>and the argument list
 *
 * <blockquote>{2}</blockquote>
 *
 * <p>returns
 *
 * <blockquote>
 * <pre>select distinct f0, f2 from t</pre>
 * </blockquote>
 *
 * <p>The <code>sourceOf</code> map is populated with the source of each
 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.
 *
 * @param relBuilder Relational expression builder
 * @param aggregate Aggregate relational expression
 * @param argList   Ordinals of columns to make distinct
 * @param filterArg Ordinal of column to filter on, or -1
 * @param sourceOf  Out parameter, is populated with a map of where each
 *                  output field came from
 * @return Aggregate relational expression which projects the required
 * columns
 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder, Aggregate aggregate, List<Integer> argList, int filterArg, Map<Integer, Integer> sourceOf) {
    relBuilder.push(aggregate.getInput());
    final List<Pair<RexNode, String>> projects = new ArrayList<>();
    final List<RelDataTypeField> childFields = relBuilder.peek().getRowType().getFieldList();
    for (int i : aggregate.getGroupSet()) {
        sourceOf.put(i, projects.size());
        projects.add(RexInputRef.of2(i, childFields));
    }
    for (Integer arg : argList) {
        if (filterArg >= 0) {
            // Implement
            // agg(DISTINCT arg) FILTER $f
            // by generating
            // SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
            // and then applying
            // agg(arg)
            // as usual.
            // 
            // It works except for (rare) agg functions that need to see null
            // values.
            final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
            final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
            final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
            RexNode condition = rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef, argRef.left, rexBuilder.ensureType(argRef.left.getType(), rexBuilder.makeCast(argRef.left.getType(), rexBuilder.constantNull()), true));
            sourceOf.put(arg, projects.size());
            projects.add(Pair.of(condition, "i$" + argRef.right));
            continue;
        }
        if (sourceOf.get(arg) != null) {
            continue;
        }
        sourceOf.put(arg, projects.size());
        projects.add(RexInputRef.of2(arg, childFields));
    }
    relBuilder.project(Pair.left(projects), Pair.right(projects));
    // Get the distinct values of the GROUP BY fields and the arguments
    // to the agg functions.
    relBuilder.push(aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), false, ImmutableBitSet.range(projects.size()), null, ImmutableList.<AggregateCall>of()));
    return relBuilder;
}

Example 99 with Pair

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

the class AggregateProjectPullUpConstantsRule method onMatch.

// ~ Methods ----------------------------------------------------------------
public void onMatch(RelOptRuleCall call) {
    final Aggregate aggregate = call.rel(0);
    final RelNode input = call.rel(1);
    assert !aggregate.indicator : "predicate ensured no grouping sets";
    final int groupCount = aggregate.getGroupCount();
    if (groupCount == 1) {
        // GROUP BY list to the empty one.
        return;
    }
    final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
    final RelMetadataQuery mq = call.getMetadataQuery();
    final RelOptPredicateList predicates = mq.getPulledUpPredicates(aggregate.getInput());
    if (predicates == null) {
        return;
    }
    final NavigableMap<Integer, RexNode> map = new TreeMap<>();
    for (int key : aggregate.getGroupSet()) {
        final RexInputRef ref = rexBuilder.makeInputRef(aggregate.getInput(), key);
        if (predicates.constantMap.containsKey(ref)) {
            map.put(key, predicates.constantMap.get(ref));
        }
    }
    // None of the group expressions are constant. Nothing to do.
    if (map.isEmpty()) {
        return;
    }
    if (groupCount == map.size()) {
        // At least a single item in group by is required.
        // Otherwise "GROUP BY 1, 2" might be altered to "GROUP BY ()".
        // Removing of the first element is not optimal here,
        // however it will allow us to use fast path below (just trim
        // groupCount).
        map.remove(map.navigableKeySet().first());
    }
    ImmutableBitSet newGroupSet = aggregate.getGroupSet();
    for (int key : map.keySet()) {
        newGroupSet = newGroupSet.clear(key);
    }
    final int newGroupCount = newGroupSet.cardinality();
    // If the constants are on the trailing edge of the group list, we just
    // reduce the group count.
    final RelBuilder relBuilder = call.builder();
    relBuilder.push(input);
    // Clone aggregate calls.
    final List<AggregateCall> newAggCalls = new ArrayList<>();
    for (AggregateCall aggCall : aggregate.getAggCallList()) {
        newAggCalls.add(aggCall.adaptTo(input, aggCall.getArgList(), aggCall.filterArg, groupCount, newGroupCount));
    }
    relBuilder.aggregate(relBuilder.groupKey(newGroupSet, null), newAggCalls);
    // Create a projection back again.
    List<Pair<RexNode, String>> projects = new ArrayList<>();
    int source = 0;
    for (RelDataTypeField field : aggregate.getRowType().getFieldList()) {
        RexNode expr;
        final int i = field.getIndex();
        if (i >= groupCount) {
            // Aggregate expressions' names and positions are unchanged.
            expr = relBuilder.field(i - map.size());
        } else {
            int pos = aggregate.getGroupSet().nth(i);
            if (map.containsKey(pos)) {
                // Re-generate the constant expression in the project.
                RelDataType originalType = aggregate.getRowType().getFieldList().get(projects.size()).getType();
                if (!originalType.equals(map.get(pos).getType())) {
                    expr = rexBuilder.makeCast(originalType, map.get(pos), true);
                } else {
                    expr = map.get(pos);
                }
            } else {
                // Project the aggregation expression, in its original
                // position.
                expr = relBuilder.field(source);
                ++source;
            }
        }
        projects.add(Pair.of(expr, field.getName()));
    }
    // inverse
    relBuilder.project(Pair.left(projects), Pair.right(projects));
    call.transformTo(relBuilder.build());
}

Example 100 with Pair

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

the class MultiJoinOptimizeBushyRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final MultiJoin multiJoinRel = call.rel(0);
    final RexBuilder rexBuilder = multiJoinRel.getCluster().getRexBuilder();
    final RelBuilder relBuilder = call.builder();
    final RelMetadataQuery mq = call.getMetadataQuery();
    final LoptMultiJoin multiJoin = new LoptMultiJoin(multiJoinRel);
    final List<Vertex> vertexes = Lists.newArrayList();
    int x = 0;
    for (int i = 0; i < multiJoin.getNumJoinFactors(); i++) {
        final RelNode rel = multiJoin.getJoinFactor(i);
        double cost = mq.getRowCount(rel);
        vertexes.add(new LeafVertex(i, rel, cost, x));
        x += rel.getRowType().getFieldCount();
    }
    assert x == multiJoin.getNumTotalFields();
    final List<LoptMultiJoin.Edge> unusedEdges = Lists.newArrayList();
    for (RexNode node : multiJoin.getJoinFilters()) {
        unusedEdges.add(multiJoin.createEdge(node));
    }
    // Comparator that chooses the best edge. A "good edge" is one that has
    // a large difference in the number of rows on LHS and RHS.
    final Comparator<LoptMultiJoin.Edge> edgeComparator = new Comparator<LoptMultiJoin.Edge>() {

        public int compare(LoptMultiJoin.Edge e0, LoptMultiJoin.Edge e1) {
            return Double.compare(rowCountDiff(e0), rowCountDiff(e1));
        }

        private double rowCountDiff(LoptMultiJoin.Edge edge) {
            assert edge.factors.cardinality() == 2 : edge.factors;
            final int factor0 = edge.factors.nextSetBit(0);
            final int factor1 = edge.factors.nextSetBit(factor0 + 1);
            return Math.abs(vertexes.get(factor0).cost - vertexes.get(factor1).cost);
        }
    };
    final List<LoptMultiJoin.Edge> usedEdges = Lists.newArrayList();
    for (; ; ) {
        final int edgeOrdinal = chooseBestEdge(unusedEdges, edgeComparator);
        if (pw != null) {
            trace(vertexes, unusedEdges, usedEdges, edgeOrdinal, pw);
        }
        final int[] factors;
        if (edgeOrdinal == -1) {
            // No more edges. Are there any un-joined vertexes?
            final Vertex lastVertex = Util.last(vertexes);
            final int z = lastVertex.factors.previousClearBit(lastVertex.id - 1);
            if (z < 0) {
                break;
            }
            factors = new int[] { z, lastVertex.id };
        } else {
            final LoptMultiJoin.Edge bestEdge = unusedEdges.get(edgeOrdinal);
            // factors on this edge.
            assert bestEdge.factors.cardinality() == 2;
            factors = bestEdge.factors.toArray();
        }
        // Determine which factor is to be on the LHS of the join.
        final int majorFactor;
        final int minorFactor;
        if (vertexes.get(factors[0]).cost <= vertexes.get(factors[1]).cost) {
            majorFactor = factors[0];
            minorFactor = factors[1];
        } else {
            majorFactor = factors[1];
            minorFactor = factors[0];
        }
        final Vertex majorVertex = vertexes.get(majorFactor);
        final Vertex minorVertex = vertexes.get(minorFactor);
        // Find the join conditions. All conditions whose factors are now all in
        // the join can now be used.
        final int v = vertexes.size();
        final ImmutableBitSet newFactors = majorVertex.factors.rebuild().addAll(minorVertex.factors).set(v).build();
        final List<RexNode> conditions = Lists.newArrayList();
        final Iterator<LoptMultiJoin.Edge> edgeIterator = unusedEdges.iterator();
        while (edgeIterator.hasNext()) {
            LoptMultiJoin.Edge edge = edgeIterator.next();
            if (newFactors.contains(edge.factors)) {
                conditions.add(edge.condition);
                edgeIterator.remove();
                usedEdges.add(edge);
            }
        }
        double cost = majorVertex.cost * minorVertex.cost * RelMdUtil.guessSelectivity(RexUtil.composeConjunction(rexBuilder, conditions, false));
        final Vertex newVertex = new JoinVertex(v, majorFactor, minorFactor, newFactors, cost, ImmutableList.copyOf(conditions));
        vertexes.add(newVertex);
        // Re-compute selectivity of edges above the one just chosen.
        // Suppose that we just chose the edge between "product" (10k rows) and
        // "product_class" (10 rows).
        // Both of those vertices are now replaced by a new vertex "P-PC".
        // This vertex has fewer rows (1k rows) -- a fact that is critical to
        // decisions made later. (Hence "greedy" algorithm not "simple".)
        // The adjacent edges are modified.
        final ImmutableBitSet merged = ImmutableBitSet.of(minorFactor, majorFactor);
        for (int i = 0; i < unusedEdges.size(); i++) {
            final LoptMultiJoin.Edge edge = unusedEdges.get(i);
            if (edge.factors.intersects(merged)) {
                ImmutableBitSet newEdgeFactors = edge.factors.rebuild().removeAll(newFactors).set(v).build();
                assert newEdgeFactors.cardinality() == 2;
                final LoptMultiJoin.Edge newEdge = new LoptMultiJoin.Edge(edge.condition, newEdgeFactors, edge.columns);
                unusedEdges.set(i, newEdge);
            }
        }
    }
    // We have a winner!
    List<Pair<RelNode, Mappings.TargetMapping>> relNodes = Lists.newArrayList();
    for (Vertex vertex : vertexes) {
        if (vertex instanceof LeafVertex) {
            LeafVertex leafVertex = (LeafVertex) vertex;
            final Mappings.TargetMapping mapping = Mappings.offsetSource(Mappings.createIdentity(leafVertex.rel.getRowType().getFieldCount()), leafVertex.fieldOffset, multiJoin.getNumTotalFields());
            relNodes.add(Pair.of(leafVertex.rel, mapping));
        } else {
            JoinVertex joinVertex = (JoinVertex) vertex;
            final Pair<RelNode, Mappings.TargetMapping> leftPair = relNodes.get(joinVertex.leftFactor);
            RelNode left = leftPair.left;
            final Mappings.TargetMapping leftMapping = leftPair.right;
            final Pair<RelNode, Mappings.TargetMapping> rightPair = relNodes.get(joinVertex.rightFactor);
            RelNode right = rightPair.left;
            final Mappings.TargetMapping rightMapping = rightPair.right;
            final Mappings.TargetMapping mapping = Mappings.merge(leftMapping, Mappings.offsetTarget(rightMapping, left.getRowType().getFieldCount()));
            if (pw != null) {
                pw.println("left: " + leftMapping);
                pw.println("right: " + rightMapping);
                pw.println("combined: " + mapping);
                pw.println();
            }
            final RexVisitor<RexNode> shuttle = new RexPermuteInputsShuttle(mapping, left, right);
            final RexNode condition = RexUtil.composeConjunction(rexBuilder, joinVertex.conditions, false);
            final RelNode join = relBuilder.push(left).push(right).join(JoinRelType.INNER, condition.accept(shuttle)).build();
            relNodes.add(Pair.of(join, mapping));
        }
        if (pw != null) {
            pw.println(Util.last(relNodes));
        }
    }
    final Pair<RelNode, Mappings.TargetMapping> top = Util.last(relNodes);
    relBuilder.push(top.left).project(relBuilder.fields(top.right));
    call.transformTo(relBuilder.build());
}