Refactor constant folding and make it reusable for primary_key_expr

.gitignore

@@ -248,3 +248,6 @@ website/package-lock.json
 
 # Ignore files for locally disabled tests
 /dbms/tests/queries/**/*.disabled
+
+# cquery cache
+/.cquery-cache

dbms/src/Core/Block.h

@@ -94,8 +94,8 @@ class Block
     /// Approximate number of allocated bytes in memory - for profiling and limits.
     size_t allocatedBytes() const;
 
-    operator bool() const { return !data.empty(); }
-    bool operator!() const { return data.empty(); }
+    operator bool() const { return !!columns(); }
+    bool operator!() const { return !this->operator bool(); }
 
     /** Get a list of column names separated by commas. */
     std::string dumpNames() const;

dbms/src/Interpreters/evaluateConstantExpression.cpp

@@ -1,18 +1,20 @@
-#include <Core/Block.h>
+#include <Interpreters/evaluateConstantExpression.h>
+
 #include <Columns/ColumnConst.h>
 #include <Columns/ColumnsNumber.h>
-#include <Parsers/ASTIdentifier.h>
-#include <Parsers/ASTLiteral.h>
-#include <Parsers/ASTFunction.h>
-#include <Parsers/ExpressionElementParsers.h>
+#include <Core/Block.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <Interpreters/Context.h>
-#include <Interpreters/SyntaxAnalyzer.h>
-#include <Interpreters/ExpressionAnalyzer.h>
+#include <Interpreters/convertFieldToType.h>
 #include <Interpreters/ExpressionActions.h>
-#include <Interpreters/evaluateConstantExpression.h>
-#include <Common/typeid_cast.h>
+#include <Interpreters/ExpressionAnalyzer.h>
+#include <Interpreters/SyntaxAnalyzer.h>
+#include <Parsers/ASTFunction.h>
+#include <Parsers/ASTIdentifier.h>
+#include <Parsers/ASTLiteral.h>
+#include <Parsers/ExpressionElementParsers.h>
 #include <TableFunctions/TableFunctionFactory.h>
+#include <Common/typeid_cast.h>
 
 
 namespace DB
@@ -57,7 +59,7 @@ std::pair<Field, std::shared_ptr<const IDataType>> evaluateConstantExpression(co
 
 ASTPtr evaluateConstantExpressionAsLiteral(const ASTPtr & node, const Context & context)
 {
-    /// Branch with string in qery.
+    /// Branch with string in query.
     if (typeid_cast<const ASTLiteral *>(node.get()))
         return node;
 
@@ -77,4 +79,238 @@ ASTPtr evaluateConstantExpressionOrIdentifierAsLiteral(const ASTPtr & node, cons
     return evaluateConstantExpressionAsLiteral(node, context);
 }
 
+namespace
+{
+    using Conjunction = ColumnsWithTypeAndName;
+    using Disjunction = std::vector<Conjunction>;
+
+    Disjunction analyzeEquals(const ASTIdentifier * identifier, const ASTLiteral * literal, const ExpressionActionsPtr & expr)
+    {
+        if (!identifier || !literal)
+        {
+            return {};
+        }
+
+        for (const auto & name_and_type : expr->getRequiredColumnsWithTypes())
+        {
+            const auto & name = name_and_type.name;
+            const auto & type = name_and_type.type;
+
+            if (name == identifier->name)
+            {
+                ColumnWithTypeAndName column;
+                // FIXME: what to do if field is not convertable?
+                column.column = type->createColumnConst(1, convertFieldToType(literal->value, *type));
+                column.name = name;
+                column.type = type;
+                return {{std::move(column)}};
+            }
+        }
+
+        return {};
+    }
+
+    Disjunction andDNF(const Disjunction & left, const Disjunction & right)
+    {
+        if (left.empty())
+        {
+            return right;
+        }
+
+        Disjunction result;
+
+        for (const auto & conjunct1 : left)
+        {
+            for (const auto & conjunct2 : right)
+            {
+                Conjunction new_conjunct{conjunct1};
+                new_conjunct.insert(new_conjunct.end(), conjunct2.begin(), conjunct2.end());
+                result.emplace_back(new_conjunct);
+            }
+        }
+
+        return result;
+    }
+
+    Disjunction analyzeFunction(const ASTFunction * fn, const ExpressionActionsPtr & expr)
+    {
+        if (!fn)
+        {
+            return {};
+        }
+
+        // TODO: enumerate all possible function names!
+
+        if (fn->name == "equals")
+        {
+            const auto * left = fn->arguments->children.front().get();
+            const auto * right = fn->arguments->children.back().get();
+            const auto * identifier = typeid_cast<const ASTIdentifier *>(left) ? typeid_cast<const ASTIdentifier *>(left)
+                                                                               : typeid_cast<const ASTIdentifier *>(right);
+            const auto * literal = typeid_cast<const ASTLiteral *>(left) ? typeid_cast<const ASTLiteral *>(left)
+                                                                         : typeid_cast<const ASTLiteral *>(right);
+
+            return analyzeEquals(identifier, literal, expr);
+        }
+        else if (fn->name == "in")
+        {
+            const auto * left = fn->arguments->children.front().get();
+            const auto * right = fn->arguments->children.back().get();
+            const auto * identifier = typeid_cast<const ASTIdentifier *>(left);
+            const auto * inner_fn = typeid_cast<const ASTFunction *>(right);
+
+            if (!inner_fn)
+            {
+                return {};
+            }
+
+            const auto * tuple = typeid_cast<const ASTExpressionList *>(inner_fn->children.front().get());
+
+            if (!tuple)
+            {
+                return {};
+            }
+
+            Disjunction result;
+
+            for (const auto & child : tuple->children)
+            {
+                const auto * literal = typeid_cast<const ASTLiteral *>(child.get());
+                const auto dnf = analyzeEquals(identifier, literal, expr);
+
+                if (dnf.empty())
+                {
+                    return {};
+                }
+
+                result.insert(result.end(), dnf.begin(), dnf.end());
+            }
+
+            return result;
+        }
+        else if (fn->name == "or")
+        {
+            const auto * args = typeid_cast<const ASTExpressionList *>(fn->children.front().get());
+
+            if (!args)
+            {
+                return {};
+            }
+
+            Disjunction result;
+
+            for (const auto & arg : args->children)
+            {
+                const auto dnf = analyzeFunction(typeid_cast<const ASTFunction *>(arg.get()), expr);
+
+                if (dnf.empty())
+                {
+                    return {};
+                }
+
+                result.insert(result.end(), dnf.begin(), dnf.end());
+            }
+
+            return result;
+        }
+        else if (fn->name == "and")
+        {
+            const auto * args = typeid_cast<const ASTExpressionList *>(fn->children.front().get());
+
+            if (!args)
+            {
+                return {};
+            }
+
+            Disjunction result;
+
+            for (const auto & arg : args->children)
+            {
+                const auto dnf = analyzeFunction(typeid_cast<const ASTFunction *>(arg.get()), expr);
+
+                if (dnf.empty())
+                {
+                    continue;
+                }
+
+                result = andDNF(result, dnf);
+            }
+
+            return result;
+        }
+
+        return {};
+    }
+}
+
+// TODO: distinguish always-false and failed evaluation results,
+//       assume failed if returned empty `Blocks` for now.
+Blocks evaluateConstantExpressionAsBlock(const ASTPtr & node, const ExpressionActionsPtr & target_expr)
+{
+    Blocks result;
+
+    // TODO: `node` may be always-false literal.
+
+    if (const auto fn = typeid_cast<const ASTFunction *>(node.get()))
+    {
+        const auto dnf = analyzeFunction(fn, target_expr);
+
+        if (dnf.empty())
+        {
+            return result;
+        }
+
+        auto hasRequiredColumns = [&target_expr](const Block & block) -> bool
+        {
+            for (const auto & name : target_expr->getRequiredColumns())
+            {
+                bool hasColumn = false;
+                for (const auto & column_name : block.getNames())
+                {
+                    if (column_name == name)
+                    {
+                        hasColumn = true;
+                        break;
+                    }
+                }
+
+                if (!hasColumn)
+                    return false;
+            }
+
+            return true;
+        };
+
+        for (const auto & conjunct : dnf)
+        {
+            Block block(conjunct);
+
+            // Block should contain all required columns from `target_expr`
+            if (!hasRequiredColumns(block))
+            {
+                return {};
+            }
+
+            target_expr->execute(block);
+
+            if (block.rows() == 1)
+            {
+                result.push_back(block);
+            }
+            else if (block.rows() == 0)
+            {
+                // filter out cases like "WHERE a = 1 AND a = 2"
+                continue;
+            }
+            else
+            {
+                // FIXME: shouldn't happen
+                return {};
+            }
+        }
+    }
+
+    return result;
+}
+
 }

dbms/src/Interpreters/evaluateConstantExpression.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <memory>
+#include <Core/Block.h>
 #include <Core/Field.h>
 #include <Parsers/IAST.h>
 #include <Parsers/IParser.h>
@@ -10,8 +11,10 @@ namespace DB
 {
 
 class Context;
+class ExpressionActions;
 class IDataType;
 
+using ExpressionActionsPtr = std::shared_ptr<ExpressionActions>;
 
 /** Evaluate constant expression and its type.
   * Used in rare cases - for elements of set for IN, for data to INSERT.
@@ -33,4 +36,7 @@ ASTPtr evaluateConstantExpressionAsLiteral(const ASTPtr & node, const Context &
   */
 ASTPtr evaluateConstantExpressionOrIdentifierAsLiteral(const ASTPtr & node, const Context & context);
 
+// FIXME: collapse returned blocks into a single block.
+Blocks evaluateConstantExpressionAsBlock(const ASTPtr & node, const ExpressionActionsPtr & target_expr);
+
 }

dbms/src/Storages/MergeTree/KeyCondition.cpp

@@ -313,7 +313,7 @@ bool KeyCondition::addCondition(const String & column, const Range & range)
     return true;
 }
 
-/** Computes value of constant expression and it data type.
+/** Computes value of constant expression and its data type.
   * Returns false, if expression isn't constant.
   */
 static bool getConstant(const ASTPtr & expr, Block & block_with_constants, Field & out_value, DataTypePtr & out_type)

dbms/src/Storages/MergeTree/KeyCondition.h

@@ -253,7 +253,7 @@ class KeyCondition
     /// Get the maximum number of the key element used in the condition.
     size_t getMaxKeyColumn() const;
 
-    /// Impose an additional condition: the value in the column column must be in the `range` range.
+    /// Impose an additional condition: the value in the column `column` must be in the range `range`.
     /// Returns whether there is such a column in the key.
     bool addCondition(const String & column, const Range & range);