ClickHouse/dbms/src/Common/VirtualColumnUtils.cpp

#include <numeric>

#include <DB/Interpreters/Context.h>
#include <DB/Interpreters/ExpressionAnalyzer.h>
#include <DB/Interpreters/ExpressionActions.h>

#include <DB/Parsers/ASTIdentifier.h>
#include <DB/Parsers/ASTExpressionList.h>
#include <DB/Parsers/ASTLiteral.h>
#include <DB/Parsers/ASTFunction.h>
#include <DB/Parsers/ASTSelectQuery.h>

#include <DB/Columns/ColumnsNumber.h>

#include <DB/Common/VirtualColumnUtils.h>


namespace DB
{

namespace VirtualColumnUtils
{

String chooseSuffix(const NamesAndTypesList & columns, const String & name)
{
    int id = 0;
    String current_suffix;
    while (true)
    {
        bool done = true;
        for (const auto & it : columns)
            if (it.name == name + current_suffix)
            {
                done = false;
                break;
            }
        if (done) break;
        ++id;
        current_suffix = toString<Int32>(id);
    }
    return current_suffix;
}

String chooseSuffixForSet(const NamesAndTypesList & columns, const std::vector<String> & names)
{
    int id = 0;
    String current_suffix;
    while (true)
    {
        bool done = true;
        for (const auto & it : columns)
        {
            for (size_t i = 0; i < names.size(); ++i)
            {
                if (it.name == names[i] + current_suffix)
                {
                    done = false;
                    break;
                }
            }
            if (!done)
                break;
        }
        if (done)
            break;
        ++id;
        current_suffix = toString<Int32>(id);
    }
    return current_suffix;
}

void rewriteEntityInAst(ASTPtr ast, const String & column_name, const Field & value)
{
    ASTSelectQuery & select = typeid_cast<ASTSelectQuery &>(*ast);
    ASTExpressionList & node = typeid_cast<ASTExpressionList &>(*select.select_expression_list);
    ASTs & asts = node.children;
    auto cur = std::make_shared<ASTLiteral>(StringRange(), value);
    cur->alias = column_name;
    ASTPtr column_value = cur;
    bool is_replaced = false;
    for (size_t i = 0; i < asts.size(); ++i)
    {
        if (const ASTIdentifier * identifier = typeid_cast<const ASTIdentifier *>(&* asts[i]))
        {
            if (identifier->kind == ASTIdentifier::Kind::Column && identifier->name == column_name)
            {
                asts[i] = column_value;
                is_replaced = true;
            }
        }
    }
    if (!is_replaced)
        asts.insert(asts.begin(), column_value);
}

/// Verifying that the function depends only on the specified columns
static bool isValidFunction(ASTPtr expression, const NameSet & columns)
{
    for (size_t i = 0; i < expression->children.size(); ++i)
        if (!isValidFunction(expression->children[i], columns))
            return false;

    if (const ASTIdentifier * identifier = typeid_cast<const ASTIdentifier *>(&*expression))
    {
        if (identifier->kind == ASTIdentifier::Kind::Column)
            return columns.count(identifier->name);
    }
    return true;
}

/// Extract all subfunctions of the main conjunction, but depending only on the specified columns
static void extractFunctions(ASTPtr expression, const NameSet & columns, std::vector<ASTPtr> & result)
{
    const ASTFunction * function = typeid_cast<const ASTFunction *>(&*expression);
    if (function && function->name == "and")
    {
        for (size_t i = 0; i < function->arguments->children.size(); ++i)
            extractFunctions(function->arguments->children[i], columns, result);
    }
    else if (isValidFunction(expression, columns))
    {
        result.push_back(expression->clone());
    }
}

/// Construct a conjunction from given functions
static ASTPtr buildWhereExpression(const ASTs & functions)
{
    if (functions.size() == 0) return nullptr;
    if (functions.size() == 1) return functions[0];
    ASTPtr new_query = std::make_shared<ASTFunction>();
    ASTFunction & new_function = typeid_cast<ASTFunction & >(*new_query);
    new_function.name = "and";
    new_function.arguments = std::make_shared<ASTExpressionList>();
    new_function.arguments->children = functions;
    new_function.children.push_back(new_function.arguments);
    return new_query;
}

bool filterBlockWithQuery(ASTPtr query, Block & block, const Context & context)
{
    query = query->clone();
    const ASTSelectQuery & select = typeid_cast<ASTSelectQuery & >(*query);
    if (!select.where_expression && !select.prewhere_expression)
        return false;

    NameSet columns;
    for (const auto & it : block.getColumnsList())
        columns.insert(it.name);

    /// We will create an expression that evaluates the expressions in WHERE and PREWHERE, depending only on the existing columns.
    std::vector<ASTPtr> functions;
    if (select.where_expression)
        extractFunctions(select.where_expression, columns, functions);
    if (select.prewhere_expression)
        extractFunctions(select.prewhere_expression, columns, functions);
    ASTPtr expression_ast = buildWhereExpression(functions);
    if (!expression_ast)
        return false;

    /// Let's parse and calculate the expression.
    ExpressionAnalyzer analyzer(expression_ast, context, {}, block.getColumnsList());
    ExpressionActionsPtr actions = analyzer.getActions(false);
    actions->execute(block);

    /// Filter the block.
    String filter_column_name = expression_ast->getColumnName();
    ColumnPtr filter_column = block.getByName(filter_column_name).column;
    if (auto converted = filter_column->convertToFullColumnIfConst())
        filter_column = converted;
    const IColumn::Filter & filter = dynamic_cast<ColumnUInt8 &>(*filter_column).getData();

    if (std::accumulate(filter.begin(), filter.end(), 0ul) == filter.size())
        return false;

    for (size_t i = 0; i < block.columns(); ++i)
    {
        ColumnPtr & column = block.safeGetByPosition(i).column;
        column = column->filter(filter, -1);
    }

    return true;
}

}

}