Allow moving conditions with subcolumns to prewhere

Author: Avogar

src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp

@@ -59,6 +59,21 @@ static Int64 findMinPosition(const NameSet & condition_table_columns, const Name
     return min_position;
 }
 
+static NameSet getTableColumns(const StorageMetadataPtr & metadata_snapshot, const Names & queried_columns)
+{
+    const auto & columns_description = metadata_snapshot->getColumns();
+    NameSet table_columns = collections::map<std::unordered_set>(columns_description.getAllPhysical(), [](const NameAndTypePair & col) { return col.name; });
+
+    /// Add also requested subcolumns to known table columns.
+    for (const auto & column : queried_columns)
+    {
+        if (columns_description.hasSubcolumn(column))
+            table_columns.insert(column);
+    }
+
+    return table_columns;
+}
+
 MergeTreeWhereOptimizer::MergeTreeWhereOptimizer(
     std::unordered_map<std::string, UInt64> column_sizes_,
     const StorageMetadataPtr & metadata_snapshot,
@@ -67,8 +82,7 @@ MergeTreeWhereOptimizer::MergeTreeWhereOptimizer(
     const std::optional<NameSet> & supported_columns_,
     LoggerPtr log_)
     : estimator(estimator_)
-    , table_columns{collections::map<std::unordered_set>(
-        metadata_snapshot->getColumns().getAllPhysical(), [](const NameAndTypePair & col) { return col.name; })}
+    , table_columns(getTableColumns(metadata_snapshot, queried_columns_))
     , queried_columns{queried_columns_}
     , supported_columns{supported_columns_}
     , sorting_key_names{NameSet(

tests/performance/prewhere_subcolumns.xml

@@ -0,0 +1,15 @@
+<test>
+    <settings>
+        <max_threads>1</max_threads>
+        <use_variant_as_common_type>1</use_variant_as_common_type>
+    </settings>
+
+    <create_query>create table test (x UInt32, n Nullable(UInt32), s String, t Tuple(a UInt32, b String), json JSON) engine=MergeTree order by tuple()</create_query>
+    <fill_query>insert into test select number, number = 0 ? NULL : number, repeat('a', 1000), tuple(number, randomString(1000)), toJSONString(map('a', number::Int32, 'b', repeat('a', 1000))) from numbers(1000000)</fill_query>
+
+    <query>select * from test where t.a = 42 format Null</query>
+    <query>select * from test where n.null format Null</query>
+    <query>select * from test where json.a = 42 format Null</query>
+
+    <drop_query>drop table if exists test</drop_query>
+</test>

tests/queries/0_stateless/03445_subcolumns_prewhere_pushdown.reference

@@ -0,0 +1,42 @@
+Expression ((Project names + Projection))
+  Expression
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}
+Expression ((Project names + Projection))
+  Filter ((WHERE + Change column names to column identifiers))
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}
+Expression ((Project names + Projection))
+  Expression
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}
+Expression ((Project names + Projection))
+  Filter ((WHERE + Change column names to column identifiers))
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}
+Expression ((Project names + Projection))
+  Expression
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}
+Expression ((Project names + Projection))
+  Filter ((WHERE + Change column names to column identifiers))
+    ReadFromMergeTree (default.test)
+0	\N	(0,0)	{"a":"0","b":"0"}
+1	\N	(1,1)	{"a":"1","b":"1"}
+2	\N	(2,2)	{"a":"2","b":"2"}
+3	\N	(3,3)	{"a":"3","b":"3"}

tests/queries/0_stateless/03445_subcolumns_prewhere_pushdown.sql

@@ -0,0 +1,20 @@
+drop table if exists test;
+create table test (x UInt64, n Nullable(UInt32), t Tuple(a UInt32, b UInt32), json JSON) engine=MergeTree order by tuple() settings index_granularity=4;
+insert into test select number, number < 4 ? NULL : number, tuple(number, number), toJSONString(map('a', number, 'b', number)) from numbers(12);
+explain select * from test where n.null settings optimize_move_to_prewhere=1;
+select * from test where n.null settings optimize_move_to_prewhere=1;
+explain select * from test where n.null settings optimize_move_to_prewhere=0;
+select * from test where n.null settings optimize_move_to_prewhere=0;
+
+explain select * from test where t.a < 4 settings optimize_move_to_prewhere=1;
+select * from test where t.a < 4 settings optimize_move_to_prewhere=1;
+explain select * from test where t.a < 4 settings optimize_move_to_prewhere=0;
+select * from test where t.a < 4 settings optimize_move_to_prewhere=0;
+
+explain select * from test where json.a::Int64 < 4 settings optimize_move_to_prewhere=1;
+select * from test where json.a::Int64 < 4 settings optimize_move_to_prewhere=1;
+explain select * from test where json.a::Int64 < 4 settings optimize_move_to_prewhere=0;
+select * from test where json.a::Int64 < 4 settings optimize_move_to_prewhere=0;
+
+drop table test;
+