Allow PREWHERE for Merge with different DEFAULT expression for column

It can be not that special to have Merge engine w/o default expression,
while MergeTree will have it. So improve PREWHERE for such kinds, and
also allow PREWHERE for DEFAULT and MATERIALIZED (since both are
physical columns on disk and this make sense).

Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>

src/Storages/StorageMerge.cpp

@@ -44,6 +44,27 @@
 #include <algorithm>
 
 
+namespace
+{
+
+using namespace DB;
+bool columnIsPhysical(ColumnDefaultKind kind)
+{
+    return kind == ColumnDefaultKind::Default || kind == ColumnDefaultKind::Materialized;
+}
+bool columnDefaultKindHasSameType(ColumnDefaultKind lhs, ColumnDefaultKind rhs)
+{
+    if (lhs == rhs)
+        return true;
+
+    if (columnIsPhysical(lhs) == columnIsPhysical(rhs))
+        return true;
+
+    return false;
+}
+
+}
+
 namespace DB
 {
 
@@ -172,11 +193,13 @@ std::optional<NameSet> StorageMerge::supportedPrewhereColumns() const
 
     NameSet supported_columns;
 
-    std::unordered_map<std::string, std::pair<const IDataType *, std::optional<ColumnDefault>>> column_type_default;
+    std::unordered_map<std::string, std::pair<const IDataType *, ColumnDefaultKind>> column_info;
     for (const auto & name_type : columns.getAll())
     {
-        column_type_default.emplace(name_type.name, std::make_pair(
-            name_type.type.get(), columns.getDefault(name_type.name)));
+        const auto & column_default = columns.getDefault(name_type.name).value_or(ColumnDefault{});
+        column_info.emplace(name_type.name, std::make_pair(
+            name_type.type.get(),
+            column_default.kind));
         supported_columns.emplace(name_type.name);
     }
 
@@ -191,11 +214,10 @@ std::optional<NameSet> StorageMerge::supportedPrewhereColumns() const
         const auto & table_columns = table_metadata_ptr->getColumns();
         for (const auto & column : table_columns.getAll())
         {
-            const auto & root_type_default = column_type_default[column.name];
-            const IDataType * root_type = root_type_default.first;
-            const std::optional<ColumnDefault> & src_default = root_type_default.second;
+            const auto & column_default = table_columns.getDefault(column.name).value_or(ColumnDefault{});
+            const auto & [root_type, src_default_kind] = column_info[column.name];
             if ((root_type && !root_type->equals(*column.type)) ||
-                src_default != table_columns.getDefault(column.name))
+                !columnDefaultKindHasSameType(src_default_kind, column_default.kind))
             {
                 supported_columns.erase(column.name);
             }

tests/queries/0_stateless/02575_merge_prewhere_default_expression.reference

@@ -0,0 +1,8 @@
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a, f;
+OK	1
+OK	2
+SELECT * FROM m PREWHERE f = 1 ORDER BY a, f;
+OK	1
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=0;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=1;

tests/queries/0_stateless/02575_merge_prewhere_default_expression.sql

@@ -0,0 +1,38 @@
+-- Allow PREWHERE when Merge() and MergeTree has different DEFAULT expression
+
+DROP TABLE IF EXISTS m;
+DROP TABLE IF EXISTS t1;
+DROP TABLE IF EXISTS t2;
+
+CREATE TABLE m
+(
+    `a` String,
+    `f` UInt8 DEFAULT 0
+)
+ENGINE = Merge(currentDatabase(), '^(t1|t2)$');
+
+CREATE TABLE t1
+(
+    a String,
+    f UInt8 DEFAULT 1
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t1 (a) VALUES ('OK');
+
+CREATE TABLE t2
+(
+    a String,
+    f UInt8 DEFAULT 2
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t2 (a) VALUES ('OK');
+
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a, f;
+SELECT * FROM m PREWHERE f = 1 ORDER BY a, f;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=0;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=1;

tests/queries/0_stateless/02575_merge_prewhere_different_default_kind.sql

@@ -1,3 +1,5 @@
+-- Prohibit PREWHERE when Merge and MergeTree has different default type of the column
+
 DROP TABLE IF EXISTS m;
 DROP TABLE IF EXISTS t1;
 DROP TABLE IF EXISTS t2;

tests/queries/0_stateless/02575_merge_prewhere_ephemeral.reference

@@ -0,0 +1,11 @@
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a;
+OK
+OK
+SELECT * FROM m PREWHERE f = 1 ORDER BY a; -- { serverError ILLEGAL_PREWHERE }
+SELECT * FROM m WHERE a = 'OK' SETTINGS optimize_move_to_prewhere=0;
+OK
+OK
+SELECT * FROM m WHERE a = 'OK' SETTINGS optimize_move_to_prewhere=1;
+OK
+OK

tests/queries/0_stateless/02575_merge_prewhere_ephemeral.sql

@@ -0,0 +1,38 @@
+-- You cannot query EPHEMERAL
+
+DROP TABLE IF EXISTS m;
+DROP TABLE IF EXISTS t1;
+DROP TABLE IF EXISTS t2;
+
+CREATE TABLE m
+(
+    `a` String,
+    `f` UInt8 EPHEMERAL 0
+)
+ENGINE = Merge(currentDatabase(), '^(t1|t2)$');
+
+CREATE TABLE t1
+(
+    a String,
+    f UInt8 DEFAULT 1
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t1 (a) VALUES ('OK');
+
+CREATE TABLE t2
+(
+    a String,
+    f UInt8 DEFAULT 2
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t2 (a) VALUES ('OK');
+
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a;
+SELECT * FROM m PREWHERE f = 1 ORDER BY a; -- { serverError ILLEGAL_PREWHERE }
+SELECT * FROM m WHERE a = 'OK' SETTINGS optimize_move_to_prewhere=0;
+SELECT * FROM m WHERE a = 'OK' SETTINGS optimize_move_to_prewhere=1;

tests/queries/0_stateless/02575_merge_prewhere_materialized.reference

@@ -0,0 +1,8 @@
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a, f;
+OK	1
+OK	2
+SELECT * FROM m PREWHERE f = 1 ORDER BY a, f;
+OK	1
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=0;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=1;

tests/queries/0_stateless/02575_merge_prewhere_materialized.sql

@@ -0,0 +1,38 @@
+-- Allow PREWHERE when Merge has DEFAULT and MergeTree has MATERIALIZED
+
+DROP TABLE IF EXISTS m;
+DROP TABLE IF EXISTS t1;
+DROP TABLE IF EXISTS t2;
+
+CREATE TABLE m
+(
+    `a` String,
+    `f` UInt8 DEFAULT 0
+)
+ENGINE = Merge(currentDatabase(), '^(t1|t2)$');
+
+CREATE TABLE t1
+(
+    a String,
+    f UInt8 MATERIALIZED 1
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t1 (a) VALUES ('OK');
+
+CREATE TABLE t2
+(
+    a String,
+    f UInt8 DEFAULT 2
+)
+ENGINE = MergeTree
+ORDER BY tuple()
+SETTINGS index_granularity = 8192;
+INSERT INTO t2 (a) VALUES ('OK');
+
+-- { echoOn }
+SELECT * FROM m PREWHERE a = 'OK' ORDER BY a, f;
+SELECT * FROM m PREWHERE f = 1 ORDER BY a, f;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=0;
+SELECT * FROM m WHERE f = 0 SETTINGS optimize_move_to_prewhere=1;