ClickHouse/dbms/src/Storages/MergeTree/PKCondition.h

#pragma once

#include <sstream>

#include <Interpreters/Context.h>
#include <Core/SortDescription.h>
#include <Parsers/ASTExpressionList.h>
#include <Parsers/ASTSelectQuery.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTLiteral.h>


namespace DB
{

class IFunction;
using FunctionPtr = std::shared_ptr<IFunction>;


/** Диапазон с открытыми или закрытыми концами; возможно, неограниченный.
  */
struct Range
{
private:
    static bool equals(const Field & lhs, const Field & rhs);
    static bool less(const Field & lhs, const Field & rhs);

public:
    Field left;                /// левая граница, если есть
    Field right;            /// правая граница, если есть
    bool left_bounded = false;        /// ограничен ли слева
    bool right_bounded = false;     /// ограничен ли справа
    bool left_included = false;     /// включает левую границу, если есть
    bool right_included = false;    /// включает правую границу, если есть

    /// Всё множество.
    Range() {}

    /// Одна точка.
    Range(const Field & point)
        : left(point), right(point), left_bounded(true), right_bounded(true), left_included(true), right_included(true) {}

    /// Ограниченный с двух сторон диапазон.
    Range(const Field & left_, bool left_included_, const Field & right_, bool right_included_)
        : left(left_), right(right_),
        left_bounded(true), right_bounded(true),
        left_included(left_included_), right_included(right_included_)
    {
        shrinkToIncludedIfPossible();
    }

    static Range createRightBounded(const Field & right_point, bool right_included)
    {
        Range r;
        r.right = right_point;
        r.right_bounded = true;
        r.right_included = right_included;
        r.shrinkToIncludedIfPossible();
        return r;
    }

    static Range createLeftBounded(const Field & left_point, bool left_included)
    {
        Range r;
        r.left = left_point;
        r.left_bounded = true;
        r.left_included = left_included;
        r.shrinkToIncludedIfPossible();
        return r;
    }

    /** Оптимизировать диапазон. Если у него есть открытая граница и тип Field "неплотный"
      * - то преобразовать её в закрытую, сузив на единицу.
      * То есть, например, превратить (0,2) в [1].
      */
    void shrinkToIncludedIfPossible()
    {
        if (left_bounded && !left_included)
        {
            if (left.getType() == Field::Types::UInt64 && left.get<UInt64>() != std::numeric_limits<UInt64>::max())
            {
                ++left.get<UInt64 &>();
                left_included = true;
            }
            if (left.getType() == Field::Types::Int64 && left.get<Int64>() != std::numeric_limits<Int64>::max())
            {
                ++left.get<Int64 &>();
                left_included = true;
            }
        }
        if (right_bounded && !right_included)
        {
            if (right.getType() == Field::Types::UInt64 && right.get<UInt64>() != std::numeric_limits<UInt64>::min())
            {
                --right.get<UInt64 &>();
                right_included = true;
            }
            if (right.getType() == Field::Types::Int64 && right.get<Int64>() != std::numeric_limits<Int64>::min())
            {
                --right.get<Int64 &>();
                right_included = true;
            }
        }
    }

    bool empty() const
    {
        return left_bounded && right_bounded
            && (less(right, left)
                || ((!left_included || !right_included) && !less(left, right)));
    }

    /// x входит в range
    bool contains(const Field & x) const
    {
        return !leftThan(x) && !rightThan(x);
    }

    /// x находится левее
    bool rightThan(const Field & x) const
    {
        return (left_bounded
            ? !(less(left, x) || (left_included && equals(x, left)))
            : false);
    }

    /// x находится правее
    bool leftThan(const Field & x) const
    {
        return (right_bounded
            ? !(less(x, right) || (right_included && equals(x, right)))
            : false);
    }

    bool intersectsRange(const Range & r) const
    {
        /// r левее меня.
        if (r.right_bounded
            && left_bounded
            && (less(r.right, left)
                || ((!left_included || !r.right_included)
                    && equals(r.right, left))))
            return false;

        /// r правее меня.
        if (r.left_bounded
            && right_bounded
            && (less(right, r.left)                            /// ...} {...
                || ((!right_included || !r.left_included)    /// ...)[...  или ...](...
                    && equals(r.left, right))))
            return false;

        return true;
    }

    bool containsRange(const Range & r) const
    {
        /// r начинается левее меня.
        if (left_bounded
            && (!r.left_bounded
                || less(r.left, left)
                || (r.left_included
                    && !left_included
                    && equals(r.left, left))))
            return false;

        /// r заканчивается правее меня.
        if (right_bounded
            && (!r.right_bounded
                || less(right, r.right)
                || (r.right_included
                    && !right_included
                    && equals(r.right, right))))
            return false;

        return true;
    }

    void swapLeftAndRight()
    {
        std::swap(left, right);
        std::swap(left_bounded, right_bounded);
        std::swap(left_included, right_included);
    }

    String toString() const;
};


class ASTSet;


/** Условие на индекс.
  *
  * Состоит из условий на принадлежность ключа всевозможным диапазонам или множествам,
  *  а также логических связок AND/OR/NOT над этими условиями.
  *
  * Составляет reverse polish notation от этих условий
  *  и умеет вычислять (интерпретировать) её выполнимость над диапазонами ключа.
  */
class PKCondition
{
public:
    /// Не учитывает секцию SAMPLE. all_columns - набор всех столбцов таблицы.
    PKCondition(ASTPtr & query, const Context & context, const NamesAndTypesList & all_columns, const SortDescription & sort_descr,
        const Block & pk_sample_block);

    /// Выполнимо ли условие в диапазоне ключей.
    /// left_pk и right_pk должны содержать все поля из sort_descr в соответствующем порядке.
    /// data_types - типы столбцов первичного ключа.
    bool mayBeTrueInRange(size_t used_key_size, const Field * left_pk, const Field * right_pk, const DataTypes & data_types) const;

    /// Выполнимо ли условие в полубесконечном (не ограниченном справа) диапазоне ключей.
    /// left_pk должен содержать все поля из sort_descr в соответствующем порядке.
    bool mayBeTrueAfter(size_t used_key_size, const Field * left_pk, const DataTypes & data_types) const;

    /// Проверяет, что индекс не может быть использован.
    bool alwaysUnknownOrTrue() const;

    /// Получить максимальный номер используемого в условии элемента первичного ключа.
    size_t getMaxKeyColumn() const;

    /// Наложить дополнительное условие: значение в столбце column должно быть в диапазоне range.
    /// Возвращает, есть ли такой столбец в первичном ключе.
    bool addCondition(const String & column, const Range & range);

    String toString() const;


    /// Выражение хранится в виде обратной польской строки (Reverse Polish Notation).
    struct RPNElement
    {
        enum Function
        {
            /// Атомы логического выражения.
            FUNCTION_IN_RANGE,
            FUNCTION_NOT_IN_RANGE,
            FUNCTION_IN_SET,
            FUNCTION_NOT_IN_SET,
            FUNCTION_UNKNOWN, /// Может принимать любое значение.
            /// Операторы логического выражения.
            FUNCTION_NOT,
            FUNCTION_AND,
            FUNCTION_OR,
            /// Константы
            ALWAYS_FALSE,
            ALWAYS_TRUE,
        };

        RPNElement() {}
        RPNElement(Function function_) : function(function_) {}
        RPNElement(Function function_, size_t key_column_) : function(function_), key_column(key_column_) {}
        RPNElement(Function function_, size_t key_column_, const Range & range_)
            : function(function_), range(range_), key_column(key_column_) {}

        String toString() const;

        Function function = FUNCTION_UNKNOWN;

        /// Для FUNCTION_IN_RANGE и FUNCTION_NOT_IN_RANGE.
        Range range;
        size_t key_column;
        /// Для FUNCTION_IN_SET, FUNCTION_NOT_IN_SET
        ASTPtr in_function;

        /** Цепочка возможно-монотонных функций.
          * Если столбец первичного ключа завёрнут в функции, которые могут быть монотонными в некоторых диапазонах значений
          * (например: -toFloat64(toDayOfWeek(date))), то здесь будут расположены функции: toDayOfWeek, toFloat64, negate.
          */
        using MonotonicFunctionsChain = std::vector<FunctionPtr>;
        mutable MonotonicFunctionsChain monotonic_functions_chain;    /// Выполнение функции не нарушает константность.
    };

    static Block getBlockWithConstants(
        const ASTPtr & query, const Context & context, const NamesAndTypesList & all_columns);

    using AtomMap = std::unordered_map<std::string, bool(*)(RPNElement & out, const Field & value, ASTPtr & node)>;
    static const AtomMap atom_map;

private:
    using RPN = std::vector<RPNElement>;
    using ColumnIndices = std::map<String, size_t>;

    bool mayBeTrueInRange(
        size_t used_key_size,
        const Field * left_pk,
        const Field * right_pk,
        const DataTypes & data_types,
        bool right_bounded) const;

    bool mayBeTrueInRangeImpl(const std::vector<Range> & key_ranges, const DataTypes & data_types) const;

    void traverseAST(ASTPtr & node, const Context & context, Block & block_with_constants);
    bool atomFromAST(ASTPtr & node, const Context & context, Block & block_with_constants, RPNElement & out);
    bool operatorFromAST(const ASTFunction * func, RPNElement & out);

    /** Is node the primary key column
      *  or expression in which column of primary key is wrapped by chain of functions,
      *  that can be monotomic on certain ranges?
      * If these conditions are true, then returns number of column in primary key, type of resulting expression
      *  and fills chain of possibly-monotonic functions.
      */
    bool isPrimaryKeyPossiblyWrappedByMonotonicFunctions(
        const ASTPtr & node,
        const Context & context,
        size_t & out_primary_key_column_num,
        DataTypePtr & out_primary_key_res_column_type,
        RPNElement::MonotonicFunctionsChain & out_functions_chain);

    bool isPrimaryKeyPossiblyWrappedByMonotonicFunctionsImpl(
        const ASTPtr & node,
        size_t & out_primary_key_column_num,
        DataTypePtr & out_primary_key_column_type,
        std::vector<const ASTFunction *> & out_functions_chain);

    RPN rpn;

    SortDescription sort_descr;
    ColumnIndices pk_columns;
    const Block & pk_sample_block;
};

}