| /* |
| * 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 org.apache.calcite.rex; |
| |
| import org.apache.calcite.plan.RelOptPredicateList; |
| import org.apache.calcite.plan.RelOptUtil; |
| import org.apache.calcite.rel.type.RelDataType; |
| import org.apache.calcite.rel.type.RelDataTypeField; |
| import org.apache.calcite.sql.fun.SqlStdOperatorTable; |
| import org.apache.calcite.util.Litmus; |
| import org.apache.calcite.util.Pair; |
| |
| import org.checkerframework.checker.nullness.qual.Nullable; |
| |
| import java.util.ArrayList; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| |
| import static java.util.Objects.requireNonNull; |
| |
| /** |
| * Workspace for constructing a {@link RexProgram}. |
| * |
| * <p>RexProgramBuilder is necessary because a {@link RexProgram} is immutable. |
| * (The {@link String} class has the same problem: it is immutable, so they |
| * introduced {@link StringBuilder}.) |
| */ |
| public class RexProgramBuilder { |
| //~ Instance fields -------------------------------------------------------- |
| |
| private final RexBuilder rexBuilder; |
| private final RelDataType inputRowType; |
| private final List<RexNode> exprList = new ArrayList<>(); |
| private final Map<Pair<RexNode, String>, RexLocalRef> exprMap = |
| new HashMap<>(); |
| private final List<RexLocalRef> localRefList = new ArrayList<>(); |
| private final List<RexLocalRef> projectRefList = new ArrayList<>(); |
| private final List<@Nullable String> projectNameList = new ArrayList<>(); |
| @SuppressWarnings("unused") |
| private final @Nullable RexSimplify simplify; |
| private @Nullable RexLocalRef conditionRef = null; |
| private boolean validating; |
| |
| //~ Constructors ----------------------------------------------------------- |
| |
| /** |
| * Creates a program-builder that will not simplify. |
| */ |
| public RexProgramBuilder(RelDataType inputRowType, RexBuilder rexBuilder) { |
| this(inputRowType, rexBuilder, null); |
| } |
| |
| /** |
| * Creates a program-builder. |
| */ |
| @SuppressWarnings("method.invocation.invalid") |
| private RexProgramBuilder(RelDataType inputRowType, RexBuilder rexBuilder, |
| @Nullable RexSimplify simplify) { |
| this.inputRowType = requireNonNull(inputRowType, "inputRowType"); |
| this.rexBuilder = requireNonNull(rexBuilder, "rexBuilder"); |
| this.simplify = simplify; // may be null |
| this.validating = assertionsAreEnabled(); |
| |
| // Pre-create an expression for each input field. |
| if (inputRowType.isStruct()) { |
| final List<RelDataTypeField> fields = inputRowType.getFieldList(); |
| for (int i = 0; i < fields.size(); i++) { |
| registerInternal(RexInputRef.of(i, fields)); |
| } |
| } |
| } |
| |
| /** |
| * Creates a program builder with the same contents as a program. |
| * |
| * @param rexBuilder Rex builder |
| * @param inputRowType Input row type |
| * @param exprList Common expressions |
| * @param projectList Projections |
| * @param condition Condition, or null |
| * @param outputRowType Output row type |
| * @param normalize Whether to normalize |
| * @param simplify Simplifier, or null to not simplify |
| */ |
| @SuppressWarnings("method.invocation.invalid") |
| private RexProgramBuilder( |
| RexBuilder rexBuilder, |
| final RelDataType inputRowType, |
| final List<RexNode> exprList, |
| final Iterable<? extends RexNode> projectList, |
| @Nullable RexNode condition, |
| final RelDataType outputRowType, |
| boolean normalize, |
| @Nullable RexSimplify simplify) { |
| this(inputRowType, rexBuilder, simplify); |
| |
| // Create a shuttle for registering input expressions. |
| final RexShuttle shuttle = |
| new RegisterMidputShuttle(true, exprList); |
| |
| // If we are not normalizing, register all internal expressions. If we |
| // are normalizing, expressions will be registered if and when they are |
| // first used. |
| if (!normalize) { |
| shuttle.visitEach(exprList); |
| } |
| |
| final RexShuttle expander = new RexProgram.ExpansionShuttle(exprList); |
| |
| // Register project expressions |
| // and create a named project item. |
| final List<RelDataTypeField> fieldList = outputRowType.getFieldList(); |
| for (Pair<? extends RexNode, RelDataTypeField> pair |
| : Pair.zip(projectList, fieldList)) { |
| final RexNode project; |
| if (simplify != null) { |
| project = simplify.simplify(pair.left.accept(expander)); |
| } else { |
| project = pair.left; |
| } |
| final String name = pair.right.getName(); |
| final RexLocalRef ref = (RexLocalRef) project.accept(shuttle); |
| addProject(ref.getIndex(), name); |
| } |
| |
| // Register the condition, if there is one. |
| if (condition != null) { |
| if (simplify != null) { |
| condition = |
| simplify.simplify( |
| rexBuilder.makeCall(SqlStdOperatorTable.IS_TRUE, |
| condition.accept(expander))); |
| if (condition.isAlwaysTrue()) { |
| condition = null; |
| } |
| } |
| if (condition != null) { |
| final RexLocalRef ref = (RexLocalRef) condition.accept(shuttle); |
| addCondition(ref); |
| } |
| } |
| } |
| |
| //~ Methods ---------------------------------------------------------------- |
| |
| /** |
| * Returns whether assertions are enabled in this class. |
| */ |
| private static boolean assertionsAreEnabled() { |
| boolean assertionsEnabled = false; |
| //noinspection AssertWithSideEffects |
| assert assertionsEnabled = true; |
| return assertionsEnabled; |
| } |
| |
| private void validate(final RexNode expr, final int fieldOrdinal) { |
| final RexVisitor<Void> validator = |
| new RexVisitorImpl<Void>(true) { |
| @Override public Void visitInputRef(RexInputRef input) { |
| final int index = input.getIndex(); |
| final List<RelDataTypeField> fields = |
| inputRowType.getFieldList(); |
| if (index < fields.size()) { |
| final RelDataTypeField inputField = fields.get(index); |
| if (input.getType() != inputField.getType()) { |
| throw new AssertionError("in expression " + expr |
| + ", field reference " + input + " has inconsistent type"); |
| } |
| } else { |
| if (index >= fieldOrdinal) { |
| throw new AssertionError("in expression " + expr |
| + ", field reference " + input + " is out of bounds"); |
| } |
| RexNode refExpr = exprList.get(index); |
| if (refExpr.getType() != input.getType()) { |
| throw new AssertionError("in expression " + expr |
| + ", field reference " + input + " has inconsistent type"); |
| } |
| } |
| return null; |
| } |
| }; |
| expr.accept(validator); |
| } |
| |
| /** |
| * Adds a project expression to the program. |
| * |
| * <p>The expression specified in terms of the input fields. If not, call |
| * {@link #registerOutput(RexNode)} first. |
| * |
| * @param expr Expression to add |
| * @param name Name of field in output row type; if null, a unique name will |
| * be generated when the program is created |
| * @return the ref created |
| */ |
| public RexLocalRef addProject(RexNode expr, @Nullable String name) { |
| final RexLocalRef ref = registerInput(expr); |
| return addProject(ref.getIndex(), name); |
| } |
| |
| /** |
| * Adds a projection based upon the <code>index</code>th expression. |
| * |
| * @param ordinal Index of expression to project |
| * @param name Name of field in output row type; if null, a unique name |
| * will be generated when the program is created |
| * @return the ref created |
| */ |
| public RexLocalRef addProject(int ordinal, final @Nullable String name) { |
| final RexLocalRef ref = localRefList.get(ordinal); |
| projectRefList.add(ref); |
| projectNameList.add(name); |
| return ref; |
| } |
| |
| /** |
| * Adds a project expression to the program at a given position. |
| * |
| * <p>The expression specified in terms of the input fields. If not, call |
| * {@link #registerOutput(RexNode)} first. |
| * |
| * @param at Position in project list to add expression |
| * @param expr Expression to add |
| * @param name Name of field in output row type; if null, a unique name will |
| * be generated when the program is created |
| * @return the ref created |
| */ |
| public RexLocalRef addProject(int at, RexNode expr, String name) { |
| final RexLocalRef ref = registerInput(expr); |
| projectRefList.add(at, ref); |
| projectNameList.add(at, name); |
| return ref; |
| } |
| |
| /** |
| * Adds a projection based upon the <code>index</code>th expression at a |
| * given position. |
| * |
| * @param at Position in project list to add expression |
| * @param ordinal Index of expression to project |
| * @param name Name of field in output row type; if null, a unique name |
| * will be generated when the program is created |
| * @return the ref created |
| */ |
| public RexLocalRef addProject(int at, int ordinal, final String name) { |
| return addProject( |
| at, |
| localRefList.get(ordinal), |
| name); |
| } |
| |
| /** |
| * Sets the condition of the program. |
| * |
| * <p>The expression must be specified in terms of the input fields. If |
| * not, call {@link #registerOutput(RexNode)} first. |
| */ |
| public void addCondition(RexNode expr) { |
| assert expr != null; |
| RexLocalRef conditionRef = this.conditionRef; |
| if (conditionRef == null) { |
| this.conditionRef = conditionRef = registerInput(expr); |
| } else { |
| // AND the new condition with the existing condition. |
| // If the new condition is identical to the existing condition, skip it. |
| RexLocalRef ref = registerInput(expr); |
| if (!ref.equals(conditionRef)) { |
| this.conditionRef = |
| registerInput( |
| rexBuilder.makeCall( |
| SqlStdOperatorTable.AND, |
| conditionRef, |
| ref)); |
| } |
| } |
| } |
| |
| /** |
| * Registers an expression in the list of common sub-expressions, and |
| * returns a reference to that expression. |
| * |
| * <p>The expression must be expressed in terms of the <em>inputs</em> of |
| * this program. |
| */ |
| public RexLocalRef registerInput(RexNode expr) { |
| final RexShuttle shuttle = new RegisterInputShuttle(true); |
| final RexNode ref = expr.accept(shuttle); |
| return (RexLocalRef) ref; |
| } |
| |
| /** |
| * Converts an expression expressed in terms of the <em>outputs</em> of this |
| * program into an expression expressed in terms of the <em>inputs</em>, |
| * registers it in the list of common sub-expressions, and returns a |
| * reference to that expression. |
| * |
| * @param expr Expression to register |
| */ |
| public RexLocalRef registerOutput(RexNode expr) { |
| final RexShuttle shuttle = new RegisterOutputShuttle(exprList); |
| final RexNode ref = expr.accept(shuttle); |
| return (RexLocalRef) ref; |
| } |
| |
| /** |
| * Registers an expression in the list of common sub-expressions, and |
| * returns a reference to that expression. |
| * |
| * @param expr Expression to register |
| */ |
| private RexLocalRef registerInternal(RexNode expr) { |
| final RexSimplify simplify = |
| new RexSimplify(rexBuilder, RelOptPredicateList.EMPTY, RexUtil.EXECUTOR); |
| expr = simplify.simplifyPreservingType(expr); |
| |
| RexLocalRef ref; |
| final Pair<RexNode, String> key; |
| if (expr instanceof RexLocalRef) { |
| key = null; |
| ref = (RexLocalRef) expr; |
| } else { |
| key = RexUtil.makeKey(expr); |
| ref = exprMap.get(key); |
| } |
| if (ref == null) { |
| if (validating) { |
| validate( |
| expr, |
| exprList.size()); |
| } |
| |
| // Add expression to list, and return a new reference to it. |
| ref = addExpr(expr); |
| exprMap.put(requireNonNull(key, "key"), ref); |
| } |
| |
| for (;;) { |
| int index = ref.index; |
| final RexNode expr2 = exprList.get(index); |
| if (expr2 instanceof RexLocalRef) { |
| ref = (RexLocalRef) expr2; |
| } else { |
| return ref; |
| } |
| } |
| } |
| |
| /** |
| * Adds an expression to the list of common expressions, and returns a |
| * reference to the expression. <b>DOES NOT CHECK WHETHER THE EXPRESSION |
| * ALREADY EXISTS</b>. |
| * |
| * @param expr Expression |
| * @return Reference to expression |
| */ |
| public RexLocalRef addExpr(RexNode expr) { |
| RexLocalRef ref; |
| final int index = exprList.size(); |
| exprList.add(expr); |
| ref = |
| new RexLocalRef( |
| index, |
| expr.getType()); |
| localRefList.add(ref); |
| return ref; |
| } |
| |
| /** |
| * Converts the state of the program builder to an immutable program, |
| * normalizing in the process. |
| * |
| * <p>It is OK to call this method, modify the program specification (by |
| * adding projections, and so forth), and call this method again. |
| */ |
| public RexProgram getProgram() { |
| return getProgram(true); |
| } |
| |
| /** |
| * Converts the state of the program builder to an immutable program. |
| * |
| * <p>It is OK to call this method, modify the program specification (by |
| * adding projections, and so forth), and call this method again. |
| * |
| * @param normalize Whether to normalize |
| */ |
| public RexProgram getProgram(boolean normalize) { |
| assert projectRefList.size() == projectNameList.size(); |
| |
| // Make sure all fields have a name. |
| generateMissingNames(); |
| RelDataType outputRowType = computeOutputRowType(); |
| |
| if (normalize) { |
| return create( |
| rexBuilder, |
| inputRowType, |
| exprList, |
| projectRefList, |
| conditionRef, |
| outputRowType, |
| true) |
| .getProgram(false); |
| } |
| |
| return new RexProgram( |
| inputRowType, |
| exprList, |
| projectRefList, |
| conditionRef, |
| outputRowType); |
| } |
| |
| private RelDataType computeOutputRowType() { |
| return RexUtil.createStructType(rexBuilder.typeFactory, projectRefList, |
| projectNameList, null); |
| } |
| |
| private void generateMissingNames() { |
| int i = -1; |
| int j = 0; |
| for (String projectName : projectNameList) { |
| ++i; |
| if (projectName == null) { |
| while (true) { |
| final String candidateName = "$" + j++; |
| if (!projectNameList.contains(candidateName)) { |
| projectNameList.set(i, candidateName); |
| break; |
| } |
| } |
| } |
| } |
| } |
| |
| /** |
| * Creates a program builder and initializes it from an existing program. |
| * |
| * <p>Calling {@link #getProgram()} immediately after creation will return a |
| * program equivalent (in terms of external behavior) to the existing |
| * program. |
| * |
| * <p>The existing program will not be changed. (It cannot: programs are |
| * immutable.) |
| * |
| * @param program Existing program |
| * @param rexBuilder Rex builder |
| * @param normalize Whether to normalize |
| * @return A program builder initialized with an equivalent program |
| */ |
| public static RexProgramBuilder forProgram( |
| RexProgram program, |
| RexBuilder rexBuilder, |
| boolean normalize) { |
| assert program.isValid(Litmus.THROW, null); |
| final RelDataType inputRowType = program.getInputRowType(); |
| final List<RexLocalRef> projectRefs = program.getProjectList(); |
| final RexLocalRef conditionRef = program.getCondition(); |
| final List<RexNode> exprs = program.getExprList(); |
| final RelDataType outputRowType = program.getOutputRowType(); |
| return create( |
| rexBuilder, |
| inputRowType, |
| exprs, |
| projectRefs, |
| conditionRef, |
| outputRowType, |
| normalize, |
| false); |
| } |
| |
| /** |
| * Creates a program builder with the same contents as a program. |
| * |
| * <p>If {@code normalize}, converts the program to canonical form. In |
| * canonical form, in addition to the usual constraints: |
| * |
| * <ul> |
| * <li>The first N internal expressions are {@link RexInputRef}s to the N |
| * input fields; |
| * <li>Subsequent internal expressions reference only preceding expressions; |
| * <li>Arguments to {@link RexCall}s must be {@link RexLocalRef}s (that is, |
| * expressions must have maximum depth 1) |
| * </ul> |
| * |
| * <p>there are additional constraints: |
| * |
| * <ul> |
| * <li>Expressions appear in the left-deep order they are needed by |
| * the projections and (if present) the condition. Thus, expression N+1 |
| * is the leftmost argument (literal or or call) in the expansion of |
| * projection #0. |
| * <li>There are no duplicate expressions |
| * <li>There are no unused expressions |
| * </ul> |
| * |
| * @param rexBuilder Rex builder |
| * @param inputRowType Input row type |
| * @param exprList Common expressions |
| * @param projectList Projections |
| * @param condition Condition, or null |
| * @param outputRowType Output row type |
| * @param normalize Whether to normalize |
| * @param simplify Whether to simplify expressions |
| * @return A program builder |
| */ |
| public static RexProgramBuilder create( |
| RexBuilder rexBuilder, |
| final RelDataType inputRowType, |
| final List<RexNode> exprList, |
| final List<? extends RexNode> projectList, |
| final @Nullable RexNode condition, |
| final RelDataType outputRowType, |
| boolean normalize, |
| @Nullable RexSimplify simplify) { |
| return new RexProgramBuilder(rexBuilder, inputRowType, exprList, |
| projectList, condition, outputRowType, normalize, simplify); |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public static RexProgramBuilder create( |
| RexBuilder rexBuilder, |
| final RelDataType inputRowType, |
| final List<RexNode> exprList, |
| final List<? extends RexNode> projectList, |
| final @Nullable RexNode condition, |
| final RelDataType outputRowType, |
| boolean normalize, |
| boolean simplify_) { |
| RexSimplify simplify = null; |
| if (simplify_) { |
| simplify = |
| new RexSimplify(rexBuilder, RelOptPredicateList.EMPTY, RexUtil.EXECUTOR); |
| } |
| return new RexProgramBuilder(rexBuilder, inputRowType, exprList, |
| projectList, condition, outputRowType, normalize, simplify); |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public static RexProgramBuilder create( |
| RexBuilder rexBuilder, |
| final RelDataType inputRowType, |
| final List<RexNode> exprList, |
| final List<? extends RexNode> projectList, |
| final @Nullable RexNode condition, |
| final RelDataType outputRowType, |
| boolean normalize) { |
| return create(rexBuilder, inputRowType, exprList, projectList, condition, |
| outputRowType, normalize, null); |
| } |
| |
| /** |
| * Creates a program builder with the same contents as a program, applying a |
| * shuttle first. |
| * |
| * <p>TODO: Refactor the above create method in terms of this one. |
| * |
| * @param rexBuilder Rex builder |
| * @param inputRowType Input row type |
| * @param exprList Common expressions |
| * @param projectRefList Projections |
| * @param conditionRef Condition, or null |
| * @param outputRowType Output row type |
| * @param shuttle Shuttle to apply to each expression before adding it |
| * to the program builder |
| * @param updateRefs Whether to update references that changes as a result |
| * of rewrites made by the shuttle |
| * @return A program builder |
| */ |
| public static RexProgramBuilder create( |
| RexBuilder rexBuilder, |
| final RelDataType inputRowType, |
| final List<RexNode> exprList, |
| final List<RexLocalRef> projectRefList, |
| final @Nullable RexLocalRef conditionRef, |
| final RelDataType outputRowType, |
| final RexShuttle shuttle, |
| final boolean updateRefs) { |
| final RexProgramBuilder progBuilder = |
| new RexProgramBuilder(inputRowType, rexBuilder); |
| progBuilder.add( |
| exprList, |
| projectRefList, |
| conditionRef, |
| outputRowType, |
| shuttle, |
| updateRefs); |
| return progBuilder; |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public static RexProgram normalize( |
| RexBuilder rexBuilder, |
| RexProgram program) { |
| return program.normalize(rexBuilder, null); |
| } |
| |
| /** |
| * Adds a set of expressions, projections and filters, applying a shuttle |
| * first. |
| * |
| * @param exprList Common expressions |
| * @param projectRefList Projections |
| * @param conditionRef Condition, or null |
| * @param outputRowType Output row type |
| * @param shuttle Shuttle to apply to each expression before adding it |
| * to the program builder |
| * @param updateRefs Whether to update references that changes as a result |
| * of rewrites made by the shuttle |
| */ |
| private void add( |
| List<RexNode> exprList, |
| List<RexLocalRef> projectRefList, |
| @Nullable RexLocalRef conditionRef, |
| final RelDataType outputRowType, |
| RexShuttle shuttle, |
| boolean updateRefs) { |
| final List<RelDataTypeField> outFields = outputRowType.getFieldList(); |
| final RexShuttle registerInputShuttle = new RegisterInputShuttle(false); |
| |
| // For each common expression, first apply the user's shuttle, then |
| // register the result. |
| // REVIEW jpham 28-Apr-2006: if the user shuttle rewrites an input |
| // expression, then input references may change |
| List<RexLocalRef> newRefs = new ArrayList<>(exprList.size()); |
| RexShuttle refShuttle = new UpdateRefShuttle(newRefs); |
| int i = 0; |
| for (RexNode expr : exprList) { |
| RexNode newExpr = expr; |
| if (updateRefs) { |
| newExpr = expr.accept(refShuttle); |
| } |
| newExpr = newExpr.accept(shuttle); |
| newRefs.add( |
| i++, |
| (RexLocalRef) newExpr.accept(registerInputShuttle)); |
| } |
| i = -1; |
| for (RexLocalRef oldRef : projectRefList) { |
| ++i; |
| RexLocalRef ref = oldRef; |
| if (updateRefs) { |
| ref = (RexLocalRef) oldRef.accept(refShuttle); |
| } |
| ref = (RexLocalRef) ref.accept(shuttle); |
| this.projectRefList.add(ref); |
| final String name = outFields.get(i).getName(); |
| assert name != null; |
| projectNameList.add(name); |
| } |
| if (conditionRef != null) { |
| if (updateRefs) { |
| conditionRef = (RexLocalRef) conditionRef.accept(refShuttle); |
| } |
| conditionRef = (RexLocalRef) conditionRef.accept(shuttle); |
| addCondition(conditionRef); |
| } |
| } |
| |
| /** |
| * Merges two programs together, and normalizes the result. |
| * |
| * @param topProgram Top program. Its expressions are in terms of the |
| * outputs of the bottom program. |
| * @param bottomProgram Bottom program. Its expressions are in terms of the |
| * result fields of the relational expression's input |
| * @param rexBuilder Rex builder |
| * @return Merged program |
| * @see #mergePrograms(RexProgram, RexProgram, RexBuilder, boolean) |
| */ |
| public static RexProgram mergePrograms( |
| RexProgram topProgram, |
| RexProgram bottomProgram, |
| RexBuilder rexBuilder) { |
| return mergePrograms(topProgram, bottomProgram, rexBuilder, true); |
| } |
| |
| /** |
| * Merges two programs together. |
| * |
| * <p>All expressions become common sub-expressions. For example, the query |
| * |
| * <blockquote><pre>SELECT x + 1 AS p, x + y AS q FROM ( |
| * SELECT a + b AS x, c AS y |
| * FROM t |
| * WHERE c = 6)}</pre></blockquote> |
| * |
| * <p>would be represented as the programs |
| * |
| * <blockquote><pre> |
| * Calc: |
| * Projects={$2, $3}, |
| * Condition=null, |
| * Exprs={$0, $1, $0 + 1, $0 + $1}) |
| * Calc( |
| * Projects={$3, $2}, |
| * Condition={$4} |
| * Exprs={$0, $1, $2, $0 + $1, $2 = 6} |
| * </pre></blockquote> |
| * |
| * <p>The merged program is |
| * |
| * <blockquote><pre> |
| * Calc( |
| * Projects={$4, $5} |
| * Condition=$6 |
| * Exprs={0: $0 // a |
| * 1: $1 // b |
| * 2: $2 // c |
| * 3: ($0 + $1) // x = a + b |
| * 4: ($3 + 1) // p = x + 1 |
| * 5: ($3 + $2) // q = x + y |
| * 6: ($2 = 6) // c = 6 |
| * </pre></blockquote> |
| * |
| * <p>Another example: |
| * |
| * <blockquote> |
| * <pre>SELECT * |
| * FROM ( |
| * SELECT a + b AS x, c AS y |
| * FROM t |
| * WHERE c = 6) |
| * WHERE x = 5</pre> |
| * </blockquote> |
| * |
| * <p>becomes |
| * |
| * <blockquote> |
| * <pre>SELECT a + b AS x, c AS y |
| * FROM t |
| * WHERE c = 6 AND (a + b) = 5</pre> |
| * </blockquote> |
| * |
| * @param topProgram Top program. Its expressions are in terms of the |
| * outputs of the bottom program. |
| * @param bottomProgram Bottom program. Its expressions are in terms of the |
| * result fields of the relational expression's input |
| * @param rexBuilder Rex builder |
| * @param normalize Whether to convert program to canonical form |
| * @return Merged program |
| */ |
| public static RexProgram mergePrograms( |
| RexProgram topProgram, |
| RexProgram bottomProgram, |
| RexBuilder rexBuilder, |
| boolean normalize) { |
| // Initialize a program builder with the same expressions, outputs |
| // and condition as the bottom program. |
| assert bottomProgram.isValid(Litmus.THROW, null); |
| assert topProgram.isValid(Litmus.THROW, null); |
| final RexProgramBuilder progBuilder = |
| RexProgramBuilder.forProgram(bottomProgram, rexBuilder, false); |
| |
| // Drive from the outputs of the top program. Register each expression |
| // used as an output. |
| final List<RexLocalRef> projectRefList = |
| progBuilder.registerProjectsAndCondition(topProgram); |
| |
| // Switch to the projects needed by the top program. The original |
| // projects of the bottom program are no longer needed. |
| progBuilder.clearProjects(); |
| final RelDataType outputRowType = topProgram.getOutputRowType(); |
| for (Pair<RexLocalRef, String> pair |
| : Pair.zip(projectRefList, outputRowType.getFieldNames(), true)) { |
| progBuilder.addProject(pair.left, pair.right); |
| } |
| RexProgram mergedProg = progBuilder.getProgram(normalize); |
| assert mergedProg.isValid(Litmus.THROW, null); |
| assert mergedProg.getOutputRowType() == topProgram.getOutputRowType(); |
| return mergedProg; |
| } |
| |
| private List<RexLocalRef> registerProjectsAndCondition(RexProgram program) { |
| final List<RexNode> exprList = program.getExprList(); |
| final List<RexLocalRef> projectRefList = new ArrayList<>(); |
| final RexShuttle shuttle = new RegisterOutputShuttle(exprList); |
| |
| // For each project, lookup the expr and expand it so it is in terms of |
| // bottomCalc's input fields |
| for (RexLocalRef topProject : program.getProjectList()) { |
| final RexNode topExpr = exprList.get(topProject.getIndex()); |
| final RexLocalRef expanded = (RexLocalRef) topExpr.accept(shuttle); |
| |
| // Remember the expr, but don't add to the project list yet. |
| projectRefList.add(expanded); |
| } |
| |
| // Similarly for the condition. |
| final RexLocalRef topCondition = program.getCondition(); |
| if (topCondition != null) { |
| final RexNode topExpr = exprList.get(topCondition.getIndex()); |
| final RexLocalRef expanded = (RexLocalRef) topExpr.accept(shuttle); |
| |
| addCondition(registerInput(expanded)); |
| } |
| return projectRefList; |
| } |
| |
| /** |
| * Removes all project items. |
| * |
| * <p>After calling this method, you may need to re-normalize. |
| */ |
| public void clearProjects() { |
| projectRefList.clear(); |
| projectNameList.clear(); |
| } |
| |
| /** |
| * Clears the condition. |
| * |
| * <p>After calling this method, you may need to re-normalize. |
| */ |
| public void clearCondition() { |
| conditionRef = null; |
| } |
| |
| /** |
| * Adds a project item for every input field. |
| * |
| * <p>You cannot call this method if there are other project items. |
| */ |
| public void addIdentity() { |
| assert projectRefList.isEmpty(); |
| for (RelDataTypeField field : inputRowType.getFieldList()) { |
| addProject( |
| new RexInputRef( |
| field.getIndex(), |
| field.getType()), |
| field.getName()); |
| } |
| } |
| |
| /** |
| * Creates a reference to a given input field. |
| * |
| * @param index Ordinal of input field, must be less than the number of |
| * fields in the input type |
| * @return Reference to input field |
| */ |
| public RexLocalRef makeInputRef(int index) { |
| final List<RelDataTypeField> fields = inputRowType.getFieldList(); |
| assert index < fields.size(); |
| final RelDataTypeField field = fields.get(index); |
| return new RexLocalRef( |
| index, |
| field.getType()); |
| } |
| |
| /** |
| * Returns the row type of the input to the program. |
| */ |
| public RelDataType getInputRowType() { |
| return inputRowType; |
| } |
| |
| /** |
| * Returns the list of project expressions. |
| */ |
| public List<RexLocalRef> getProjectList() { |
| return projectRefList; |
| } |
| |
| //~ Inner Classes ---------------------------------------------------------- |
| |
| /** Shuttle that visits a tree of {@link RexNode} and registers them |
| * in a program. */ |
| private abstract class RegisterShuttle extends RexShuttle { |
| @Override public RexNode visitCall(RexCall call) { |
| final RexNode expr = super.visitCall(call); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitOver(RexOver over) { |
| final RexNode expr = super.visitOver(over); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitLiteral(RexLiteral literal) { |
| final RexNode expr = super.visitLiteral(literal); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitFieldAccess(RexFieldAccess fieldAccess) { |
| final RexNode expr = super.visitFieldAccess(fieldAccess); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitDynamicParam(RexDynamicParam dynamicParam) { |
| final RexNode expr = super.visitDynamicParam(dynamicParam); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitCorrelVariable(RexCorrelVariable variable) { |
| final RexNode expr = super.visitCorrelVariable(variable); |
| return registerInternal(expr); |
| } |
| |
| @Override public RexNode visitLambda(RexLambda lambda) { |
| super.visitLambda(lambda); |
| return registerInternal(lambda); |
| } |
| } |
| |
| /** |
| * Shuttle which walks over an expression, registering each sub-expression. |
| * Each {@link RexInputRef} is assumed to refer to an <em>input</em> of the |
| * program. |
| */ |
| private class RegisterInputShuttle extends RegisterShuttle { |
| private final boolean valid; |
| |
| protected RegisterInputShuttle(boolean valid) { |
| this.valid = valid; |
| } |
| |
| @Override public RexNode visitInputRef(RexInputRef input) { |
| final int index = input.getIndex(); |
| if (valid) { |
| // The expression should already be valid. Check that its |
| // index is within bounds. |
| if ((index < 0) || (index >= inputRowType.getFieldCount())) { |
| assert false |
| : "RexInputRef index " + index + " out of range 0.." |
| + (inputRowType.getFieldCount() - 1); |
| } |
| |
| // Check that the type is consistent with the referenced |
| // field. If it is an object type, the rules are different, so |
| // skip the check. |
| assert input.getType().isStruct() |
| || RelOptUtil.eq("type1", input.getType(), |
| "type2", inputRowType.getFieldList().get(index).getType(), |
| Litmus.THROW); |
| } |
| |
| // Return a reference to the N'th expression, which should be |
| // equivalent. |
| final RexLocalRef ref = localRefList.get(index); |
| return ref; |
| } |
| |
| @Override public RexNode visitLocalRef(RexLocalRef local) { |
| if (valid) { |
| // The expression should already be valid. |
| final int index = local.getIndex(); |
| assert index >= 0 : index; |
| assert index < exprList.size() |
| : "index=" + index + ", exprList=" + exprList; |
| assert RelOptUtil.eq( |
| "expr type", |
| exprList.get(index).getType(), |
| "ref type", |
| local.getType(), |
| Litmus.THROW); |
| } |
| |
| // Resolve the expression to an input. |
| while (true) { |
| final int index = local.getIndex(); |
| final RexNode expr = exprList.get(index); |
| if (expr instanceof RexLocalRef) { |
| local = (RexLocalRef) expr; |
| if (local.index >= index) { |
| throw new AssertionError( |
| "expr " + local + " references later expr " + local.index); |
| } |
| } else { |
| // Add expression to the list, just so that subsequent |
| // expressions don't get screwed up. This expression is |
| // unused, so will be eliminated soon. |
| return registerInternal(local); |
| } |
| } |
| } |
| } |
| |
| /** |
| * Extension to {@link RegisterInputShuttle} which allows expressions to be |
| * in terms of inputs or previous common sub-expressions. |
| */ |
| private class RegisterMidputShuttle extends RegisterInputShuttle { |
| private final List<RexNode> localExprList; |
| |
| protected RegisterMidputShuttle( |
| boolean valid, |
| List<RexNode> localExprList) { |
| super(valid); |
| this.localExprList = localExprList; |
| } |
| |
| @Override public RexNode visitLocalRef(RexLocalRef local) { |
| // Convert a local ref into the common-subexpression it references. |
| final int index = local.getIndex(); |
| return localExprList.get(index).accept(this); |
| } |
| } |
| |
| /** |
| * Shuttle which walks over an expression, registering each sub-expression. |
| * Each {@link RexInputRef} is assumed to refer to an <em>output</em> of the |
| * program. |
| */ |
| private class RegisterOutputShuttle extends RegisterShuttle { |
| private final List<RexNode> localExprList; |
| |
| RegisterOutputShuttle(List<RexNode> localExprList) { |
| super(); |
| this.localExprList = localExprList; |
| } |
| |
| @Override public RexNode visitInputRef(RexInputRef input) { |
| // This expression refers to the Nth project column. Lookup that |
| // column and find out what common sub-expression IT refers to. |
| final int index = input.getIndex(); |
| final RexLocalRef local = projectRefList.get(index); |
| assert RelOptUtil.eq( |
| "type1", |
| local.getType(), |
| "type2", |
| input.getType(), |
| Litmus.THROW); |
| return local; |
| } |
| |
| @Override public RexNode visitLocalRef(RexLocalRef local) { |
| // Convert a local ref into the common-subexpression it references. |
| final int index = local.getIndex(); |
| return localExprList.get(index).accept(this); |
| } |
| } |
| |
| /** |
| * Shuttle that rewires {@link RexLocalRef} using a list of updated |
| * references. |
| */ |
| private static class UpdateRefShuttle extends RexShuttle { |
| private List<RexLocalRef> newRefs; |
| |
| private UpdateRefShuttle(List<RexLocalRef> newRefs) { |
| this.newRefs = newRefs; |
| } |
| |
| @Override public RexNode visitLocalRef(RexLocalRef localRef) { |
| return newRefs.get(localRef.getIndex()); |
| } |
| } |
| } |
