Examples with RexNode org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode used on opensource projects

Example 51 with RexNode

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode in project calcite by apache.

the class RelMdDistinctRowCount method getDistinctRowCount.

public Double getDistinctRowCount(Union rel, RelMetadataQuery mq, ImmutableBitSet groupKey, RexNode predicate) {
    Double rowCount = 0.0;
    int[] adjustments = new int[rel.getRowType().getFieldCount()];
    RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    for (RelNode input : rel.getInputs()) {
        // convert the predicate to reference the types of the union child
        RexNode modifiedPred;
        if (predicate == null) {
            modifiedPred = null;
        } else {
            modifiedPred = predicate.accept(new RelOptUtil.RexInputConverter(rexBuilder, null, input.getRowType().getFieldList(), adjustments));
        }
        Double partialRowCount = mq.getDistinctRowCount(input, groupKey, modifiedPred);
        if (partialRowCount == null) {
            return null;
        }
        rowCount += partialRowCount;
    }
    return rowCount;
}

Example 52 with RexNode

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode in project calcite by apache.

the class RelMdDistinctRowCount method getDistinctRowCount.

public Double getDistinctRowCount(Project rel, RelMetadataQuery mq, ImmutableBitSet groupKey, RexNode predicate) {
    if (predicate == null || predicate.isAlwaysTrue()) {
        if (groupKey.isEmpty()) {
            return 1D;
        }
    }
    ImmutableBitSet.Builder baseCols = ImmutableBitSet.builder();
    ImmutableBitSet.Builder projCols = ImmutableBitSet.builder();
    List<RexNode> projExprs = rel.getProjects();
    RelMdUtil.splitCols(projExprs, groupKey, baseCols, projCols);
    final List<RexNode> notPushable = new ArrayList<>();
    final List<RexNode> pushable = new ArrayList<>();
    RelOptUtil.splitFilters(ImmutableBitSet.range(rel.getRowType().getFieldCount()), predicate, pushable, notPushable);
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    // get the distinct row count of the child input, passing in the
    // columns and filters that only reference the child; convert the
    // filter to reference the children projection expressions
    RexNode childPred = RexUtil.composeConjunction(rexBuilder, pushable, true);
    RexNode modifiedPred;
    if (childPred == null) {
        modifiedPred = null;
    } else {
        modifiedPred = RelOptUtil.pushPastProject(childPred, rel);
    }
    Double distinctRowCount = mq.getDistinctRowCount(rel.getInput(), baseCols.build(), modifiedPred);
    if (distinctRowCount == null) {
        return null;
    } else if (!notPushable.isEmpty()) {
        RexNode preds = RexUtil.composeConjunction(rexBuilder, notPushable, true);
        distinctRowCount *= RelMdUtil.guessSelectivity(preds);
    }
    // are all column references
    if (projCols.cardinality() == 0) {
        return distinctRowCount;
    }
    // multiply by the cardinality of the non-child projection expressions
    for (int bit : projCols.build()) {
        Double subRowCount = RelMdUtil.cardOfProjExpr(mq, rel, projExprs.get(bit));
        if (subRowCount == null) {
            return null;
        }
        distinctRowCount *= subRowCount;
    }
    return RelMdUtil.numDistinctVals(distinctRowCount, mq.getRowCount(rel));
}

Example 53 with RexNode

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode in project calcite by apache.

the class RelMdExpressionLineage method getExpressionLineage.

/**
 * Expression lineage from {@link Join}.
 *
 * <p>We only extract the lineage for INNER joins.
 */
public Set<RexNode> getExpressionLineage(Join rel, RelMetadataQuery mq, RexNode outputExpression) {
    if (rel.getJoinType() != JoinRelType.INNER) {
        // We cannot map origin of this expression.
        return null;
    }
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    final RelNode leftInput = rel.getLeft();
    final RelNode rightInput = rel.getRight();
    final int nLeftColumns = leftInput.getRowType().getFieldList().size();
    // Infer column origin expressions for given references
    final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
    final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
    final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
    for (int idx = 0; idx < rel.getRowType().getFieldList().size(); idx++) {
        if (idx < nLeftColumns) {
            final RexInputRef inputRef = RexInputRef.of(idx, leftInput.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(leftInput, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            // Gather table references, left input references remain unchanged
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            for (RelTableRef leftRef : tableRefs) {
                qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef);
            }
            mapping.put(RexInputRef.of(idx, rel.getRowType().getFieldList()), originalExprs);
        } else {
            // Right input.
            final RexInputRef inputRef = RexInputRef.of(idx - nLeftColumns, rightInput.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(rightInput, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            // Gather table references, right input references might need to be
            // updated if there are table names clashes with left input
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            for (RelTableRef rightRef : tableRefs) {
                int shift = 0;
                Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(rightRef.getQualifiedName());
                if (lRefs != null) {
                    shift = lRefs.size();
                }
                currentTablesMapping.put(rightRef, RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber()));
            }
            final Set<RexNode> updatedExprs = Sets.newHashSet(Iterables.transform(originalExprs, new Function<RexNode, RexNode>() {

                @Override
                public RexNode apply(RexNode e) {
                    return RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping);
                }
            }));
            mapping.put(RexInputRef.of(idx, rel.getRowType().getFieldList()), updatedExprs);
        }
    }
    // Return result
    return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}

Example 54 with RexNode

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode in project calcite by apache.

the class RelMdExpressionLineage method getExpressionLineage.

/**
 * Expression lineage from {@link Union}.
 *
 * <p>For Union operator, we might be able to extract multiple origins for the
 * references in the given expression.
 */
public Set<RexNode> getExpressionLineage(Union rel, RelMetadataQuery mq, RexNode outputExpression) {
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    // Infer column origin expressions for given references
    final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
    final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
    for (RelNode input : rel.getInputs()) {
        final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
        for (int idx = 0; idx < input.getRowType().getFieldList().size(); idx++) {
            final RexInputRef inputRef = RexInputRef.of(idx, input.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(input, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
            // Gather table references, references might need to be
            // updated
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            for (RelTableRef tableRef : tableRefs) {
                int shift = 0;
                Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(tableRef.getQualifiedName());
                if (lRefs != null) {
                    shift = lRefs.size();
                }
                currentTablesMapping.put(tableRef, RelTableRef.of(tableRef.getTable(), shift + tableRef.getEntityNumber()));
            }
            final Set<RexNode> updatedExprs = Sets.newHashSet(Iterables.transform(originalExprs, new Function<RexNode, RexNode>() {

                @Override
                public RexNode apply(RexNode e) {
                    return RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping);
                }
            }));
            final Set<RexNode> set = mapping.get(ref);
            if (set != null) {
                set.addAll(updatedExprs);
            } else {
                mapping.put(ref, updatedExprs);
            }
        }
        // Add to existing qualified names
        for (RelTableRef newRef : currentTablesMapping.values()) {
            qualifiedNamesToRefs.put(newRef.getQualifiedName(), newRef);
        }
    }
    // Return result
    return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}

Example 55 with RexNode

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode in project calcite by apache.

the class RelJsonReader method readRel.

private void readRel(final Map<String, Object> jsonRel) {
    String id = (String) jsonRel.get("id");
    String type = (String) jsonRel.get("relOp");
    Constructor constructor = relJson.getConstructor(type);
    RelInput input = new RelInput() {

        public RelOptCluster getCluster() {
            return cluster;
        }

        public RelTraitSet getTraitSet() {
            return cluster.traitSetOf(Convention.NONE);
        }

        public RelOptTable getTable(String table) {
            final List<String> list = getStringList(table);
            return relOptSchema.getTableForMember(list);
        }

        public RelNode getInput() {
            final List<RelNode> inputs = getInputs();
            assert inputs.size() == 1;
            return inputs.get(0);
        }

        public List<RelNode> getInputs() {
            final List<String> jsonInputs = getStringList("inputs");
            if (jsonInputs == null) {
                return ImmutableList.of(lastRel);
            }
            final List<RelNode> inputs = new ArrayList<>();
            for (String jsonInput : jsonInputs) {
                inputs.add(lookupInput(jsonInput));
            }
            return inputs;
        }

        public RexNode getExpression(String tag) {
            return relJson.toRex(this, jsonRel.get(tag));
        }

        public ImmutableBitSet getBitSet(String tag) {
            return ImmutableBitSet.of(getIntegerList(tag));
        }

        public List<ImmutableBitSet> getBitSetList(String tag) {
            List<List<Integer>> list = getIntegerListList(tag);
            if (list == null) {
                return null;
            }
            final ImmutableList.Builder<ImmutableBitSet> builder = ImmutableList.builder();
            for (List<Integer> integers : list) {
                builder.add(ImmutableBitSet.of(integers));
            }
            return builder.build();
        }

        public List<String> getStringList(String tag) {
            // noinspection unchecked
            return (List<String>) jsonRel.get(tag);
        }

        public List<Integer> getIntegerList(String tag) {
            // noinspection unchecked
            return (List<Integer>) jsonRel.get(tag);
        }

        public List<List<Integer>> getIntegerListList(String tag) {
            // noinspection unchecked
            return (List<List<Integer>>) jsonRel.get(tag);
        }

        public List<AggregateCall> getAggregateCalls(String tag) {
            @SuppressWarnings("unchecked") final List<Map<String, Object>> jsonAggs = (List) jsonRel.get(tag);
            final List<AggregateCall> inputs = new ArrayList<>();
            for (Map<String, Object> jsonAggCall : jsonAggs) {
                inputs.add(toAggCall(jsonAggCall));
            }
            return inputs;
        }

        public Object get(String tag) {
            return jsonRel.get(tag);
        }

        public String getString(String tag) {
            return (String) jsonRel.get(tag);
        }

        public float getFloat(String tag) {
            return ((Number) jsonRel.get(tag)).floatValue();
        }

        public boolean getBoolean(String tag, boolean default_) {
            final Boolean b = (Boolean) jsonRel.get(tag);
            return b != null ? b : default_;
        }

        public <E extends Enum<E>> E getEnum(String tag, Class<E> enumClass) {
            return Util.enumVal(enumClass, getString(tag).toUpperCase(Locale.ROOT));
        }

        public List<RexNode> getExpressionList(String tag) {
            @SuppressWarnings("unchecked") final List<Object> jsonNodes = (List) jsonRel.get(tag);
            final List<RexNode> nodes = new ArrayList<>();
            for (Object jsonNode : jsonNodes) {
                nodes.add(relJson.toRex(this, jsonNode));
            }
            return nodes;
        }

        public RelDataType getRowType(String tag) {
            final Object o = jsonRel.get(tag);
            return relJson.toType(cluster.getTypeFactory(), o);
        }

        public RelDataType getRowType(String expressionsTag, String fieldsTag) {
            final List<RexNode> expressionList = getExpressionList(expressionsTag);
            @SuppressWarnings("unchecked") final List<String> names = (List<String>) get(fieldsTag);
            return cluster.getTypeFactory().createStructType(new AbstractList<Map.Entry<String, RelDataType>>() {

                @Override
                public Map.Entry<String, RelDataType> get(int index) {
                    return Pair.of(names.get(index), expressionList.get(index).getType());
                }

                @Override
                public int size() {
                    return names.size();
                }
            });
        }

        public RelCollation getCollation() {
            // noinspection unchecked
            return relJson.toCollation((List) get("collation"));
        }

        public RelDistribution getDistribution() {
            return relJson.toDistribution(get("distribution"));
        }

        public ImmutableList<ImmutableList<RexLiteral>> getTuples(String tag) {
            // noinspection unchecked
            final List<List> jsonTuples = (List) get(tag);
            final ImmutableList.Builder<ImmutableList<RexLiteral>> builder = ImmutableList.builder();
            for (List jsonTuple : jsonTuples) {
                builder.add(getTuple(jsonTuple));
            }
            return builder.build();
        }

        public ImmutableList<RexLiteral> getTuple(List jsonTuple) {
            final ImmutableList.Builder<RexLiteral> builder = ImmutableList.builder();
            for (Object jsonValue : jsonTuple) {
                builder.add((RexLiteral) relJson.toRex(this, jsonValue));
            }
            return builder.build();
        }
    };
    try {
        final RelNode rel = (RelNode) constructor.newInstance(input);
        relMap.put(id, rel);
        lastRel = rel;
    } catch (InstantiationException | IllegalAccessException e) {
        throw new RuntimeException(e);
    } catch (InvocationTargetException e) {
        final Throwable e2 = e.getCause();
        if (e2 instanceof RuntimeException) {
            throw (RuntimeException) e2;
        }
        throw new RuntimeException(e2);
    }
}