All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.hazelcast.org.apache.calcite.rel.metadata.RelMdExpressionLineage Maven / Gradle / Ivy

There is a newer version: 5.5.0
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to you under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.hazelcast.org.apache.calcite.rel.metadata;

import com.hazelcast.org.apache.calcite.plan.RelOptUtil;
import com.hazelcast.org.apache.calcite.plan.volcano.RelSubset;
import com.hazelcast.org.apache.calcite.rel.RelNode;
import com.hazelcast.org.apache.calcite.rel.core.Aggregate;
import com.hazelcast.org.apache.calcite.rel.core.Calc;
import com.hazelcast.org.apache.calcite.rel.core.Exchange;
import com.hazelcast.org.apache.calcite.rel.core.Filter;
import com.hazelcast.org.apache.calcite.rel.core.Join;
import com.hazelcast.org.apache.calcite.rel.core.JoinRelType;
import com.hazelcast.org.apache.calcite.rel.core.Project;
import com.hazelcast.org.apache.calcite.rel.core.Sort;
import com.hazelcast.org.apache.calcite.rel.core.TableModify;
import com.hazelcast.org.apache.calcite.rel.core.TableScan;
import com.hazelcast.org.apache.calcite.rel.core.Union;
import com.hazelcast.org.apache.calcite.rel.type.RelDataType;
import com.hazelcast.org.apache.calcite.rel.type.RelDataTypeField;
import com.hazelcast.org.apache.calcite.rex.RexBuilder;
import com.hazelcast.org.apache.calcite.rex.RexInputRef;
import com.hazelcast.org.apache.calcite.rex.RexNode;
import com.hazelcast.org.apache.calcite.rex.RexShuttle;
import com.hazelcast.org.apache.calcite.rex.RexTableInputRef;
import com.hazelcast.org.apache.calcite.rex.RexTableInputRef.RelTableRef;
import com.hazelcast.org.apache.calcite.rex.RexUtil;
import com.hazelcast.org.apache.calcite.sql.validate.SqlValidatorUtil;
import com.hazelcast.org.apache.calcite.util.ImmutableBitSet;
import com.hazelcast.org.apache.calcite.util.Pair;
import com.hazelcast.org.apache.calcite.util.Util;

import com.hazelcast.com.google.common.collect.HashMultimap;
import com.hazelcast.com.google.common.collect.ImmutableList;
import com.hazelcast.com.google.common.collect.ImmutableSet;
import com.hazelcast.com.google.common.collect.Multimap;

import com.hazelcast.org.checkerframework.checker.nullness.qual.KeyFor;
import com.hazelcast.org.checkerframework.checker.nullness.qual.Nullable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;

/**
 * Default implementation of
 * {@link RelMetadataQuery#getExpressionLineage} for the standard logical
 * algebra.
 *
 * 

The goal of this provider is to infer the lineage for the given expression. * *

The output expressions might contain references to columns produced by * {@link TableScan} operators ({@link RexTableInputRef}). In turn, each * TableScan operator is identified uniquely by a {@link RelTableRef} containing * its qualified name and an identifier. * *

If the lineage cannot be inferred, we return null. */ public class RelMdExpressionLineage implements MetadataHandler { public static final RelMetadataProvider SOURCE = ReflectiveRelMetadataProvider.reflectiveSource( new RelMdExpressionLineage(), BuiltInMetadata.ExpressionLineage.Handler.class); //~ Constructors ----------------------------------------------------------- protected RelMdExpressionLineage() {} //~ Methods ---------------------------------------------------------------- @Override public MetadataDef getDef() { return BuiltInMetadata.ExpressionLineage.DEF; } // Catch-all rule when none of the others apply. public @Nullable Set getExpressionLineage(RelNode rel, RelMetadataQuery mq, RexNode outputExpression) { return null; } public @Nullable Set getExpressionLineage(RelSubset rel, RelMetadataQuery mq, RexNode outputExpression) { RelNode bestOrOriginal = Util.first(rel.getBest(), rel.getOriginal()); if (bestOrOriginal == null) { return null; } return mq.getExpressionLineage(bestOrOriginal, outputExpression); } /** * Expression lineage from {@link TableScan}. * *

We extract the fields referenced by the expression and we express them * using {@link RexTableInputRef}. */ public @Nullable Set getExpressionLineage(TableScan rel, RelMetadataQuery mq, RexNode outputExpression) { final RexBuilder rexBuilder = rel.getCluster().getRexBuilder(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); // Infer column origin expressions for given references final Map> mapping = new LinkedHashMap<>(); for (int idx : inputFieldsUsed) { final RexNode inputRef = RexTableInputRef.of( RelTableRef.of(rel.getTable(), 0), RexInputRef.of(idx, rel.getRowType().getFieldList())); final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList()); mapping.put(ref, ImmutableSet.of(inputRef)); } // Return result return createAllPossibleExpressions(rexBuilder, outputExpression, mapping); } /** * Expression lineage from {@link Aggregate}. * *

If the expression references grouping sets or aggregate function * results, we cannot extract the lineage and we return null. */ public @Nullable Set getExpressionLineage(Aggregate rel, RelMetadataQuery mq, RexNode outputExpression) { final RelNode input = rel.getInput(); final RexBuilder rexBuilder = rel.getCluster().getRexBuilder(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); for (int idx : inputFieldsUsed) { if (idx >= rel.getGroupCount()) { // We cannot map origin of this expression. return null; } } // Infer column origin expressions for given references final Map> mapping = new LinkedHashMap<>(); for (int idx : inputFieldsUsed) { final RexInputRef inputRef = RexInputRef.of(rel.getGroupSet().nth(idx), input.getRowType().getFieldList()); final Set originalExprs = mq.getExpressionLineage(input, inputRef); if (originalExprs == null) { // Bail out return null; } final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList()); mapping.put(ref, originalExprs); } // Return result return createAllPossibleExpressions(rexBuilder, outputExpression, mapping); } /** * Expression lineage from {@link Join}. * *

We only extract the lineage for INNER joins. */ public @Nullable Set getExpressionLineage(Join rel, RelMetadataQuery mq, RexNode outputExpression) { final RexBuilder rexBuilder = rel.getCluster().getRexBuilder(); final RelNode leftInput = rel.getLeft(); final RelNode rightInput = rel.getRight(); final int nLeftColumns = leftInput.getRowType().getFieldList().size(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); if (rel.getJoinType().isOuterJoin()) { // If we reference the inner side, we will bail out if (rel.getJoinType() == JoinRelType.LEFT) { ImmutableBitSet rightFields = ImmutableBitSet.range( nLeftColumns, rel.getRowType().getFieldCount()); if (inputFieldsUsed.intersects(rightFields)) { // We cannot map origin of this expression. return null; } } else if (rel.getJoinType() == JoinRelType.RIGHT) { ImmutableBitSet leftFields = ImmutableBitSet.range( 0, nLeftColumns); if (inputFieldsUsed.intersects(leftFields)) { // We cannot map origin of this expression. return null; } } else { // We cannot map origin of this expression. return null; } } // Gather table references final Set leftTableRefs = mq.getTableReferences(leftInput); if (leftTableRefs == null) { // Bail out return null; } final Set rightTableRefs = mq.getTableReferences(rightInput); if (rightTableRefs == null) { // Bail out return null; } final Multimap, RelTableRef> qualifiedNamesToRefs = HashMultimap.create(); final Map currentTablesMapping = new HashMap<>(); for (RelTableRef leftRef : leftTableRefs) { qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef); } for (RelTableRef rightRef : rightTableRefs) { int shift = 0; Collection lRefs = qualifiedNamesToRefs.get( rightRef.getQualifiedName()); if (lRefs != null) { shift = lRefs.size(); } currentTablesMapping.put(rightRef, RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber())); } // Infer column origin expressions for given references final Map> mapping = new LinkedHashMap<>(); for (int idx : inputFieldsUsed) { if (idx < nLeftColumns) { final RexInputRef inputRef = RexInputRef.of(idx, leftInput.getRowType().getFieldList()); final Set originalExprs = mq.getExpressionLineage(leftInput, inputRef); if (originalExprs == null) { // Bail out return null; } // Left input references remain unchanged mapping.put(RexInputRef.of(idx, rel.getRowType().getFieldList()), originalExprs); } else { // Right input. final RexInputRef inputRef = RexInputRef.of(idx - nLeftColumns, rightInput.getRowType().getFieldList()); final Set originalExprs = mq.getExpressionLineage(rightInput, inputRef); if (originalExprs == null) { // Bail out return null; } // Right input references might need to be updated if there are // table names clashes with left input final RelDataType fullRowType = SqlValidatorUtil.createJoinType( rexBuilder.getTypeFactory(), rel.getLeft().getRowType(), rel.getRight().getRowType(), null, ImmutableList.of()); final Set updatedExprs = ImmutableSet.copyOf( Util.transform(originalExprs, e -> RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping))); mapping.put(RexInputRef.of(idx, fullRowType), updatedExprs); } } // Return result return createAllPossibleExpressions(rexBuilder, outputExpression, mapping); } /** * Expression lineage from {@link Union}. * *

For Union operator, we might be able to extract multiple origins for the * references in the given expression. */ public @Nullable Set getExpressionLineage(Union rel, RelMetadataQuery mq, RexNode outputExpression) { final RexBuilder rexBuilder = rel.getCluster().getRexBuilder(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); // Infer column origin expressions for given references final Multimap, RelTableRef> qualifiedNamesToRefs = HashMultimap.create(); final Map> mapping = new LinkedHashMap<>(); for (RelNode input : rel.getInputs()) { // Gather table references final Map currentTablesMapping = new HashMap<>(); final Set tableRefs = mq.getTableReferences(input); if (tableRefs == null) { // Bail out return null; } for (RelTableRef tableRef : tableRefs) { int shift = 0; Collection lRefs = qualifiedNamesToRefs.get( tableRef.getQualifiedName()); if (lRefs != null) { shift = lRefs.size(); } currentTablesMapping.put(tableRef, RelTableRef.of(tableRef.getTable(), shift + tableRef.getEntityNumber())); } // Map references for (int idx : inputFieldsUsed) { final RexInputRef inputRef = RexInputRef.of(idx, input.getRowType().getFieldList()); final Set originalExprs = mq.getExpressionLineage(input, inputRef); if (originalExprs == null) { // Bail out return null; } // References might need to be updated final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList()); final Set updatedExprs = originalExprs.stream() .map(e -> RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping)) .collect(Collectors.toSet()); final Set 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); } /** * Expression lineage from Project. */ public @Nullable Set getExpressionLineage(Project rel, final RelMetadataQuery mq, RexNode outputExpression) { final RelNode input = rel.getInput(); final RexBuilder rexBuilder = rel.getCluster().getRexBuilder(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); // Infer column origin expressions for given references final Map> mapping = new LinkedHashMap<>(); for (int idx : inputFieldsUsed) { final RexNode inputExpr = rel.getProjects().get(idx); final Set originalExprs = mq.getExpressionLineage(input, inputExpr); if (originalExprs == null) { // Bail out return null; } final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList()); mapping.put(ref, originalExprs); } // Return result return createAllPossibleExpressions(rexBuilder, outputExpression, mapping); } /** * Expression lineage from Filter. */ public @Nullable Set getExpressionLineage(Filter rel, RelMetadataQuery mq, RexNode outputExpression) { return mq.getExpressionLineage(rel.getInput(), outputExpression); } /** * Expression lineage from Sort. */ public @Nullable Set getExpressionLineage(Sort rel, RelMetadataQuery mq, RexNode outputExpression) { return mq.getExpressionLineage(rel.getInput(), outputExpression); } /** * Expression lineage from TableModify. */ public @Nullable Set getExpressionLineage(TableModify rel, RelMetadataQuery mq, RexNode outputExpression) { return mq.getExpressionLineage(rel.getInput(), outputExpression); } /** * Expression lineage from Exchange. */ public @Nullable Set getExpressionLineage(Exchange rel, RelMetadataQuery mq, RexNode outputExpression) { return mq.getExpressionLineage(rel.getInput(), outputExpression); } /** * Expression lineage from Calc. */ public @Nullable Set getExpressionLineage(Calc calc, RelMetadataQuery mq, RexNode outputExpression) { final RelNode input = calc.getInput(); final RexBuilder rexBuilder = calc.getCluster().getRexBuilder(); // Extract input fields referenced by expression final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression); // Infer column origin expressions for given references final Map> mapping = new LinkedHashMap<>(); Pair, ImmutableList> calcProjectsAndFilter = calc.getProgram().split(); for (int idx : inputFieldsUsed) { final RexNode inputExpr = calcProjectsAndFilter.getKey().get(idx); final Set originalExprs = mq.getExpressionLineage(input, inputExpr); if (originalExprs == null) { // Bail out return null; } final RexInputRef ref = RexInputRef.of(idx, calc.getRowType().getFieldList()); mapping.put(ref, originalExprs); } // Return result return createAllPossibleExpressions(rexBuilder, outputExpression, mapping); } /** * Given an expression, it will create all equivalent expressions resulting * from replacing all possible combinations of references in the mapping by * the corresponding expressions. * * @param rexBuilder rexBuilder * @param expr expression * @param mapping mapping * @return set of resulting expressions equivalent to the input expression */ protected static @Nullable Set createAllPossibleExpressions(RexBuilder rexBuilder, RexNode expr, Map> mapping) { // Extract input fields referenced by expression final ImmutableBitSet predFieldsUsed = extractInputRefs(expr); if (predFieldsUsed.isEmpty()) { // The unique expression is the input expression return ImmutableSet.of(expr); } try { return createAllPossibleExpressions(rexBuilder, expr, predFieldsUsed, mapping, new HashMap<>()); } catch (UnsupportedOperationException e) { // There may be a RexNode unsupported by RexCopier, just return null return null; } } private static Set createAllPossibleExpressions(RexBuilder rexBuilder, RexNode expr, ImmutableBitSet predFieldsUsed, Map> mapping, Map singleMapping) { final @KeyFor("mapping") RexInputRef inputRef = mapping.keySet().iterator().next(); final Set replacements = requireNonNull(mapping.remove(inputRef), () -> "mapping.remove(inputRef) is null for " + inputRef); Set result = new HashSet<>(); assert !replacements.isEmpty(); if (predFieldsUsed.indexOf(inputRef.getIndex()) != -1) { for (RexNode replacement : replacements) { singleMapping.put(inputRef, replacement); createExpressions(rexBuilder, expr, predFieldsUsed, mapping, singleMapping, result); singleMapping.remove(inputRef); } } else { createExpressions(rexBuilder, expr, predFieldsUsed, mapping, singleMapping, result); } mapping.put(inputRef, replacements); return result; } private static void createExpressions(RexBuilder rexBuilder, RexNode expr, ImmutableBitSet predFieldsUsed, Map> mapping, Map singleMapping, Set result) { if (mapping.isEmpty()) { final RexReplacer replacer = new RexReplacer(singleMapping); final List updatedPreds = new ArrayList<>(1); updatedPreds.add(rexBuilder.copy(expr)); replacer.mutate(updatedPreds); result.addAll(updatedPreds); } else { result.addAll( createAllPossibleExpressions( rexBuilder, expr, predFieldsUsed, mapping, singleMapping)); } } /** * Replaces expressions with their equivalences. Note that we only have to * look for RexInputRef. */ private static class RexReplacer extends RexShuttle { private final Map replacementValues; RexReplacer(Map replacementValues) { this.replacementValues = replacementValues; } @Override public RexNode visitInputRef(RexInputRef inputRef) { return requireNonNull( replacementValues.get(inputRef), () -> "no replacement found for inputRef " + inputRef); } } private static ImmutableBitSet extractInputRefs(RexNode expr) { final Set inputExtraFields = new LinkedHashSet<>(); final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields); expr.accept(inputFinder); return inputFinder.build(); } }





© 2015 - 2024 Weber Informatics LLC | Privacy Policy