Expr.java

// 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.
// This file is copied from
// https://github.com/apache/impala/blob/branch-2.9.0/fe/src/main/java/org/apache/impala/Expr.java
// and modified by Doris

package org.apache.doris.analysis;

import org.apache.doris.catalog.AggStateType;
import org.apache.doris.catalog.Function;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.catalog.TableIf.TableType;
import org.apache.doris.catalog.Type;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.common.FormatOptions;
import org.apache.doris.common.TreeNode;
import org.apache.doris.nereids.util.Utils;
import org.apache.doris.planner.normalize.Normalizer;
import org.apache.doris.thrift.TExpr;
import org.apache.doris.thrift.TExprNode;
import org.apache.doris.thrift.TExprOpcode;

import com.google.common.base.Joiner;
import com.google.common.base.MoreObjects;
import com.google.common.base.Objects;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.base.Suppliers;
import com.google.common.collect.Lists;
import com.google.gson.annotations.SerializedName;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;

/**
 * Root of the expr node hierarchy.
 */
public abstract class Expr extends TreeNode<Expr> implements Cloneable {

    public static final String AGG_STATE_SUFFIX = "_state";
    public static final String AGG_UNION_SUFFIX = "_union";
    public static final String AGG_MERGE_SUFFIX = "_merge";
    public static final String AGG_FOREACH_SUFFIX = "_foreach";
    public static final String DEFAULT_EXPR_NAME = "expr";

    protected boolean disableTableName = false;

    protected Optional<Boolean> nullableFromNereids = Optional.empty();
    protected Optional<Boolean> originCastNullable = Optional.empty();

    @SerializedName("type")
    protected Type type;  // result of analysis

    protected boolean isAnalyzed = false;  // true after analyze() has been called

    @SerializedName("opcode")
    protected TExprOpcode opcode;  // opcode for this expr

    // The function to call. This can either be a scalar or aggregate function.
    // Set in analyze().
    protected Function fn;

    // Cached value of IsConstant(), set during analyze() and valid if isAnalyzed_ is true.
    private Supplier<Boolean> isConstant = Suppliers.memoize(() -> false);

    protected Optional<String> exprName = Optional.empty();

    protected Expr() {
        super();
        type = Type.INVALID;
        opcode = TExprOpcode.INVALID_OPCODE;
    }

    protected Expr(Expr other) {
        super();
        type = other.type;
        isAnalyzed = other.isAnalyzed;
        opcode = other.opcode;
        isConstant = other.isConstant;
        fn = other.fn;
        children = Expr.cloneList(other.children);
        nullableFromNereids = other.nullableFromNereids;
    }

    public boolean isAnalyzed() {
        return isAnalyzed;
    }

    public void checkValueValid() throws AnalysisException {
    }

    // Name of expr, this is used by generating column name automatically when there is no
    // alias or is not slotRef
    public String getExprName() {
        if (!this.exprName.isPresent()) {
            this.exprName = Optional.of(Utils.normalizeName(this.getClass().getSimpleName(), DEFAULT_EXPR_NAME));
        }
        return this.exprName.get();
    }

    public Type getType() {
        return type;
    }

    // add by cmy. for restoring
    public void setType(Type type) {
        this.type = type;
    }

    public TExprOpcode getOpcode() {
        return opcode;
    }

    public Function getFn() {
        return fn;
    }

    /**
     * Set the expr to be analyzed and computes isConstant_.
     */
    protected void analysisDone() {
        Preconditions.checkState(!isAnalyzed);
        // We need to compute the const-ness as the last step, since analysis may change
        // the result, e.g. by resolving function.
        isConstant = Suppliers.memoize(this::isConstantImpl);
        isAnalyzed = true;
    }

    /**
     * Collects the returns types of the child nodes in an array.
     */
    protected Type[] collectChildReturnTypes() {
        Type[] childTypes = new Type[children.size()];
        for (int i = 0; i < children.size(); ++i) {
            childTypes[i] = children.get(i).type;
        }
        return childTypes;
    }

    protected Boolean[] collectChildReturnNullables() {
        Boolean[] childNullables = new Boolean[children.size()];
        for (int i = 0; i < children.size(); ++i) {
            childNullables[i] = children.get(i).isNullable();
        }
        return childNullables;
    }

    public List<Expr> getChildrenWithoutCast() {
        List<Expr> result = new ArrayList<>();
        for (int i = 0; i < children.size(); ++i) {
            if (children.get(i) instanceof CastExpr) {
                CastExpr castExpr = (CastExpr) children.get(i);
                result.add(castExpr.getChild(0));
            } else {
                result.add(children.get(i));
            }
        }
        return result;
    }

    public Expr getChildWithoutCast(int i) {
        Preconditions.checkArgument(i < children.size(), "child index {0} out of range {1}", i, children.size());
        Expr child = children.get(i);
        return child instanceof CastExpr ? child.children.get(0) : child;
    }

    public static List<TExpr> treesToThrift(List<? extends Expr> exprs) {
        List<TExpr> result = Lists.newArrayList();
        for (Expr expr : exprs) {
            result.add(expr.treeToThrift());
        }
        return result;
    }

    public static String debugString(List<? extends Expr> exprs) {
        if (exprs == null || exprs.isEmpty()) {
            return "";
        }
        List<String> strings = Lists.newArrayList();
        for (Expr expr : exprs) {
            strings.add(Strings.nullToEmpty(expr.debugString()));
        }
        return "(" + Joiner.on(" ").join(strings) + ")";
    }

    /**
     * Return true if l1 equals l2 when both lists are interpreted as sets.
     */
    public static <C extends Expr> boolean equalSets(List<C> l1, List<C> l2) {
        if (l1.size() != l2.size()) {
            return false;
        }
        Map cMap1 = toCountMap(l1);
        Map cMap2 = toCountMap(l2);
        if (cMap1.size() != cMap2.size()) {
            return false;
        }
        Iterator it = cMap1.keySet().iterator();
        while (it.hasNext()) {
            C obj = (C) it.next();
            Integer count1 = (Integer) cMap1.get(obj);
            Integer count2 = (Integer) cMap2.get(obj);
            if (count2 == null || count1 != count2) {
                return false;
            }
        }
        return true;
    }

    public static <C extends Expr> HashMap<C, Integer> toCountMap(List<C> list) {
        HashMap countMap = new HashMap<C, Integer>();
        for (int i = 0; i < list.size(); i++) {
            C obj = list.get(i);
            Integer count = (Integer) countMap.get(obj);
            if (count == null) {
                countMap.put(obj, 1);
            } else {
                countMap.put(obj, count + 1);
            }
        }
        return countMap;
    }

    /**
     * Create a deep copy of 'l'. If sMap is non-null, use it to substitute the
     * elements of l.
     */
    public static <C extends Expr> ArrayList<C> cloneList(List<C> l, ExprSubstitutionMap sMap) {
        Preconditions.checkNotNull(l);
        ArrayList<C> result = new ArrayList<C>();
        for (C element : l) {
            result.add((C) element.clone(sMap));
        }
        return result;
    }

    /**
     * Create a deep copy of 'l'. If sMap is non-null, use it to substitute the
     * elements of l.
     */
    public static <C extends Expr> ArrayList<C> cloneList(List<C> l) {
        Preconditions.checkNotNull(l);
        ArrayList<C> result = new ArrayList<C>();
        for (C element : l) {
            result.add((C) element.clone());
        }
        return result;
    }

    /**
     * Collect all unique Expr nodes of type 'cl' present in 'input' and add them to
     * 'output' if they do not exist in 'output'.
     * This can't go into TreeNode<>, because we'd be using the template param
     * NodeType.
     */
    public static <C extends Expr> void collectList(List<? extends Expr> input, Class<C> cl, List<C> output) {
        Preconditions.checkNotNull(input);
        for (Expr e : input) {
            e.collect(cl, output);
        }
    }

    /**
     * get the expr which in l1 and l2 in the same time.
     * Return the intersection of l1 and l2
     */
    public static <C extends Expr> List<C> intersect(List<C> l1, List<C> l2) {
        List<C> result = new ArrayList<C>();

        for (C element : l1) {
            if (l2.contains(element)) {
                result.add(element);
            }
        }

        return result;
    }

    public static void extractSlots(Expr root, Set<SlotId> slotIdSet) {
        if (root instanceof SlotRef) {
            slotIdSet.add(((SlotRef) root).getDesc().getId());
            return;
        }
        for (Expr child : root.getChildren()) {
            extractSlots(child, slotIdSet);
        }
    }

    /**
     * Removes duplicate exprs (according to equals()).
     */
    public static <C extends Expr> void removeDuplicates(List<C> l) {
        if (l == null) {
            return;
        }
        ListIterator<C> it1 = l.listIterator();
        while (it1.hasNext()) {
            C e1 = it1.next();
            ListIterator<C> it2 = l.listIterator();
            boolean duplicate = false;
            while (it2.hasNext()) {
                C e2 = it2.next();
                if (e1 == e2) {
                    // only check up to but excluding e1
                    break;
                }
                if (e1.equals(e2)) {
                    duplicate = true;
                    break;
                }
            }
            if (duplicate) {
                it1.remove();
            }
        }
    }

    public String toSql() {
        if (disableTableName) {
            return toSqlWithoutTbl();
        }
        return toSqlImpl();
    }

    public String toSql(boolean disableTableName, boolean needExternalSql, TableType tableType, TableIf table) {
        return toSqlImpl(disableTableName, needExternalSql, tableType, table);
    }

    public void disableTableName() {
        disableTableName = true;
        for (Expr child : children) {
            child.disableTableName();
        }
    }

    public String toSqlWithoutTbl() {
        return toSql(true, false, null, null);
    }

    /**
     * Returns a SQL string representing this expr. Subclasses should override this method
     * instead of toSql() to ensure that parenthesis are properly added around the toSql().
     */
    protected abstract String toSqlImpl();

    protected abstract String toSqlImpl(boolean disableTableName, boolean needExternalSql, TableType tableType,
            TableIf table);

    public String toExternalSql(TableType tableType, TableIf table) {
        return toSql(false, true, tableType, table);
    }

    /**
     * Return a column label for the expression
     */
    public String toColumnLabel() {
        return toSql();
    }

    // Convert this expr, including all children, to its Thrift representation.
    public TExpr treeToThrift() {
        TExpr result = new TExpr();
        treeToThriftHelper(result);
        return result;
    }

    protected void treeToThriftHelper(TExpr container) {
        treeToThriftHelper(container, ((expr, exprNode) -> expr.toThrift(exprNode)));
    }

    // Append a flattened version of this expr, including all children, to 'container'.
    protected void treeToThriftHelper(TExpr container, ExprVisitor visitor) {
        TExprNode msg = new TExprNode();
        msg.type = type.toThrift();
        msg.num_children = children.size();
        if (fn != null) {
            msg.setFn(fn.toThrift(type, collectChildReturnTypes(), collectChildReturnNullables()));
            if (fn.hasVarArgs()) {
                msg.setVarargStartIdx(fn.getNumArgs() - 1);
            }
        }
        // useless parameter, just give a number
        msg.output_scale = -1;
        msg.setIsNullable(nullableFromNereids.isPresent() ? nullableFromNereids.get() : isNullable());
        visitor.visit(this, msg);
        container.addToNodes(msg);
        for (Expr child : children) {
            child.treeToThriftHelper(container, visitor);
        }
    }

    public interface ExprVisitor {
        void visit(Expr expr, TExprNode exprNode);
    }

    // Convert this expr into msg (excluding children), which requires setting
    // msg.op as well as the expr-specific field.
    protected abstract void toThrift(TExprNode msg);

    public String debugString() {
        return debugString(children);
    }

    /**
     * Creates a deep copy of this expr including its analysis state. The method is
     * abstract in this class to force new Exprs to implement it.
     */
    @Override
    public abstract Expr clone();

    public boolean notCheckDescIdEquals(Object obj) {
        return equals(obj);
    }

    @Override
    public boolean equals(Object obj) {
        if (obj == null) {
            return false;
        }
        if (obj.getClass() != this.getClass()) {
            return false;
        }
        // don't compare type, this could be called pre-analysis
        Expr expr = (Expr) obj;
        if (children.size() != expr.children.size()) {
            return false;
        }
        for (int i = 0; i < children.size(); ++i) {
            if (!children.get(i).equals(expr.children.get(i))) {
                return false;
            }
        }
        if (fn == null && expr.fn == null) {
            return true;
        }
        if (fn == null || expr.fn == null) {
            return false;
        }
        // Both fn_'s are not null
        return fn.equals(expr.fn);
    }

    @Override
    public int hashCode() {
        int result = 31 * Objects.hashCode(type) + Objects.hashCode(opcode);
        for (Expr child : children) {
            result = 31 * result + Objects.hashCode(child);
        }
        return result;
    }

    /**
     * Gather conjuncts from this expr and return them in a list.
     * A conjunct is an expr that returns a boolean, e.g., Predicates, function calls,
     * SlotRefs, etc. Hence, this method is placed here and not in Predicate.
     */
    public List<Expr> getConjuncts() {
        List<Expr> list = Lists.newArrayList();
        if (this instanceof CompoundPredicate
                && ((CompoundPredicate) this).getOp() == CompoundPredicate.Operator.AND) {
            // TODO: we have to convert CompoundPredicate.AND to two expr trees for
            // conjuncts because NULLs are handled differently for CompoundPredicate.AND
            // and conjunct evaluation.  This is not optimal for jitted exprs because it
            // will result in two functions instead of one. Create a new CompoundPredicate
            // Operator (i.e. CONJUNCT_AND) with the right NULL semantics and use that
            // instead
            list.addAll((getChild(0)).getConjuncts());
            list.addAll((getChild(1)).getConjuncts());
        } else {
            list.add(this);
        }
        return list;
    }

    /**
     * Create a deep copy of 'this'. If sMap is non-null,
     * use it to substitute 'this' or its subnodes.
     * <p/>
     * Expr subclasses that add non-value-type members must override this.
     */
    public Expr clone(ExprSubstitutionMap sMap) {
        if (sMap != null) {
            for (int i = 0; i < sMap.getLhs().size(); ++i) {
                if (this.equals(sMap.getLhs().get(i))) {
                    return sMap.getRhs().get(i).clone(null);
                }
            }
        }
        Expr result = (Expr) this.clone();
        result.children = Lists.newArrayList();
        for (Expr child : children) {
            result.children.add(((Expr) child).clone(sMap));
        }
        return result;
    }

    /**
     * Returns true if expr is fully bound by slotId, otherwise false.
     */
    public boolean isBound(SlotId slotId) {
        for (Expr child : children) {
            if (!child.isBound(slotId)) {
                return false;
            }
        }
        return true;
    }

    public Map<Long, Set<String>> getTableIdToColumnNames() {
        Map<Long, Set<String>> tableIdToColumnNames = new HashMap<Long, Set<String>>();
        getTableIdToColumnNames(tableIdToColumnNames);
        return tableIdToColumnNames;
    }

    public void getTableIdToColumnNames(Map<Long, Set<String>> tableIdToColumnNames) {
        Preconditions.checkState(tableIdToColumnNames != null);
        for (Expr child : children) {
            child.getTableIdToColumnNames(tableIdToColumnNames);
        }
    }

    /**
     * @return true if this is an instance of LiteralExpr
     */
    public boolean isLiteral() {
        return this instanceof LiteralExpr;
    }

    /**
     * Returns true if this expression should be treated as constant. I.e. if the frontend
     * and backend should assume that two evaluations of the expression within a query will
     * return the same value. Examples of constant expressions include:
     * - Literal values like 1, "foo", or NULL
     * - Deterministic operators applied to constant arguments, e.g. 1 + 2, or
     *   concat("foo", "bar")
     * - Functions that should be always return the same value within a query but may
     *   return different values for different queries. E.g. now(), which we want to
     *   evaluate only once during planning.
     * May incorrectly return true if the expression is not analyzed.
     * TODO: isAnalyzed_ should be a precondition for isConstant(), since it is not always
     * possible to correctly determine const-ness before analysis (e.g. see
     * FunctionCallExpr.isConstant()).
     */
    public final boolean isConstant() {
        if (isAnalyzed) {
            return isConstant.get();
        }
        return isConstantImpl();
    }

    /**
     * Implements isConstant() - computes the value without using 'isConstant_'.
     */
    protected boolean isConstantImpl() {
        for (Expr expr : children) {
            if (!expr.isConstant()) {
                return false;
            }
        }
        return true;
    }

    @Override
    public String toString() {
        return MoreObjects.toStringHelper(this.getClass()).add("type", type).toString();
    }

    /**
     * If 'this' is a SlotRef or a Cast that wraps a SlotRef, returns that SlotRef.
     * Otherwise returns null.
     */
    public SlotRef unwrapSlotRef() {
        if (this instanceof SlotRef) {
            return (SlotRef) this;
        } else if (this instanceof CastExpr && getChild(0) instanceof SlotRef) {
            return (SlotRef) getChild(0);
        } else {
            return null;
        }
    }

    /**
     * Returns the first child if this Expr is a CastExpr. Otherwise, returns 'this'.
     */
    public Expr unwrapExpr(boolean implicitOnly) {
        if (this instanceof CastExpr
                && (!implicitOnly || ((CastExpr) this).isImplicit())) {
            return children.get(0);
        }
        return this;
    }

    public boolean isImplicitCast() {
        return this instanceof CastExpr && ((CastExpr) this).isImplicit();
    }

    public boolean contains(Expr expr) {
        if (this.equals(expr)) {
            return true;
        }

        for (Expr child : getChildren()) {
            if (child.contains(expr)) {
                return true;
            }
        }

        return false;
    }

    public Expr findEqual(List<Expr> exprs) {
        if (exprs.isEmpty()) {
            return null;
        }
        for (Expr expr : exprs) {
            if (contains(expr)) {
                return expr;
            }
        }
        return null;
    }

    public String getStringValue() {
        return "";
    }

    /**
     * This method is used for constant fold of query in FE,
     * for different serde dialect(hive, presto, doris).
     */
    public String getStringValueForQuery(FormatOptions options) {
        return getStringValue();
    }

    /**
     * This method is to return the string value of this expr in a complex type for query
     * It is only used for "getStringValueForQuery()"
     * For most of the integer types, it is same as getStringValueForQuery().
     * But for others like StringLiteral and DateLiteral, it should be wrapped with quotations.
     * eg: 1,2,abc,[1,2,3],["abc","def"],{10:20},{"abc":20}
     */
    protected String getStringValueInComplexTypeForQuery(FormatOptions options) {
        return getStringValueForQuery(options);
    }

    /**
     * This method is to return the string value of this expr for stream load.
     * so there is a little different from "getStringValueForQuery()".
     * eg, for NullLiteral, it should be "\N" for stream load, but "null" for FE constant
     * for StructLiteral, the value should not contain sub column's name.
     */
    public String getStringValueForStreamLoad(FormatOptions options) {
        return getStringValueForQuery(options);
    }

    public final TExpr normalize(Normalizer normalizer) {
        TExpr result = new TExpr();
        treeToThriftHelper(result, (expr, texprNode) -> expr.normalize(texprNode, normalizer));
        return result;
    }

    protected void normalize(TExprNode msg, Normalizer normalizer) {
        this.toThrift(msg);
    }

    protected boolean hasNullableChild() {
        return hasNullableChild(children);
    }

    protected static boolean hasNullableChild(List<Expr> children) {
        for (Expr expr : children) {
            if (expr.isNullable()) {
                return true;
            }
        }
        return false;
    }

    /**
     * For excute expr the result is nullable
     * TODO: Now only SlotRef and LiteralExpr overwrite the method, each child of Expr should
     * overwrite this method to plan correct
     */
    public boolean isNullable() {
        return isNullable(fn, children);
    }

    public static boolean isNullable(Function fn, List<Expr> children) {
        if (fn == null) {
            return true;
        }
        switch (fn.getNullableMode()) {
            case DEPEND_ON_ARGUMENT:
                return hasNullableChild(children);
            case ALWAYS_NOT_NULLABLE:
                return false;
            case ALWAYS_NULLABLE:
            default:
                return true;
        }
    }

    public static AggStateType createAggStateType(String name, List<Type> typeList,
            List<Boolean> nullableList, boolean resultNullable) {
        return new AggStateType(name, resultNullable, typeList, nullableList);
    }

    public static List<Expr> getMockedExprs(List<Type> typeList, List<Boolean> nullableList) {
        List<Expr> mockedExprs = Lists.newArrayList();
        for (int i = 0; i < typeList.size(); i++) {
            mockedExprs.add(new SlotRef(typeList.get(i), nullableList.get(i)));
        }
        return mockedExprs;
    }

    public static List<Expr> getMockedExprs(AggStateType type) {
        return getMockedExprs(type.getSubTypes(), type.getSubTypeNullables());
    }

    // This is only for transactional insert operation,
    // to check it the given value in insert stmt is LiteralExpr.
    // And if we write "1" to a boolean column, there will be a cast(1 as boolean) expr,
    // which is also accepted.
    public boolean isLiteralOrCastExpr() {
        if (this instanceof CastExpr) {
            return children.get(0) instanceof LiteralExpr;
        } else {
            return this instanceof LiteralExpr;
        }
    }

    public boolean isNullLiteral() {
        return this instanceof NullLiteral;
    }

    public void setNullableFromNereids(boolean nullable) {
        nullableFromNereids = Optional.of(nullable);
    }

    public Optional<Boolean> getNullableFromNereids() {
        return nullableFromNereids;
    }

    public void setOriginCastNullable(boolean nullable) {
        originCastNullable = Optional.of(nullable);
    }

    public Set<SlotRef> getInputSlotRef() {
        Set<SlotRef> slots = new HashSet<>();
        if (this instanceof SlotRef) {
            slots.add((SlotRef) this);
            return slots;
        } else {
            for (Expr expr : children) {
                slots.addAll(expr.getInputSlotRef());
            }
        }
        return slots;
    }
}