add test case

fix order key check

src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp

@@ -37,7 +37,8 @@ MergeTreeWhereOptimizer::MergeTreeWhereOptimizer(
     : table_columns{ext::map<std::unordered_set>(
         metadata_snapshot->getColumns().getAllPhysical(), [](const NameAndTypePair & col) { return col.name; })}
     , queried_columns{queried_columns_}
-    , primary_key_columns{metadata_snapshot->getPrimaryKey().column_names}
+    , sorting_key_names{NameSet(
+          metadata_snapshot->getSortingKey().column_names.begin(), metadata_snapshot->getSortingKey().column_names.end())}
     , block_with_constants{KeyCondition::getBlockWithConstants(query_info.query, query_info.syntax_analyzer_result, context)}
     , log{log_}
     , column_sizes{std::move(column_sizes_)}
@@ -301,9 +302,9 @@ bool MergeTreeWhereOptimizer::isPrimaryKeyAtom(const ASTPtr & ast) const
 }
 
 
-bool MergeTreeWhereOptimizer::isPrimaryKey(const String & column_name) const
+bool MergeTreeWhereOptimizer::isSortingKey(const String & column_name) const
 {
-    return std::find(primary_key_columns.begin(), primary_key_columns.end(), column_name) != primary_key_columns.end();
+    return sorting_key_names.count(column_name);
 }
 
 
@@ -344,7 +345,7 @@ bool MergeTreeWhereOptimizer::cannotBeMoved(const ASTPtr & ptr, bool is_final) c
         /// disallow moving result of ARRAY JOIN to PREWHERE
         if (array_joined_names.count(*opt_name) ||
             array_joined_names.count(Nested::extractTableName(*opt_name)) ||
-            (is_final && !isPrimaryKey(*opt_name)))
+            (is_final && !isSortingKey(*opt_name)))
             return true;
     }
 

src/Storages/MergeTree/MergeTreeWhereOptimizer.h

@@ -85,7 +85,7 @@ private:
 
     bool isPrimaryKeyAtom(const ASTPtr & ast) const;
 
-    bool isPrimaryKey(const String & column_name) const;
+    bool isSortingKey(const String & column_name) const;
 
     bool isConstant(const ASTPtr & expr) const;
 
@@ -106,7 +106,7 @@ private:
     String first_primary_key_column;
     const StringSet table_columns;
     const Names queried_columns;
-    const Names primary_key_columns;
+    const NameSet sorting_key_names;
     const Block block_with_constants;
     Poco::Logger * log;
     std::unordered_map<std::string, UInt64> column_sizes;

tests/queries/0_stateless/01737_move_order_key_to_prewhere_select_final.reference

@@ -0,0 +1,35 @@
+SELECT
+    x,
+    y,
+    z
+FROM prewhere_move_select_final
+PREWHERE y > 100
+SELECT
+    x,
+    y,
+    z
+FROM prewhere_move_select_final
+FINAL
+PREWHERE y > 100
+SELECT
+    x,
+    y,
+    z
+FROM prewhere_move_select_final
+FINAL
+WHERE z > 400
+SELECT
+    x,
+    y,
+    z
+FROM prewhere_move_select_final
+FINAL
+WHERE z > 400
+SELECT
+    x,
+    y,
+    z
+FROM prewhere_move_select_final
+FINAL
+PREWHERE y > 100
+WHERE (y > 100) AND (z > 400)

tests/queries/0_stateless/01737_move_order_to_prewhere_select_final.sql

@@ -0,0 +1,15 @@
+DROP TABLE IF EXISTS prewhere_move_select_final;
+CREATE TABLE prewhere_move_select_final (x Int, y Int, z Int) ENGINE = ReplacingMergeTree() ORDER BY (x, y);
+INSERT INTO prewhere_move_select_final SELECT number, number * 2, number * 3 FROM numbers(1000);
+
+-- order key can be pushed down with final
+EXPLAIN SYNTAX SELECT * FROM prewhere_move_select_final WHERE y > 100;
+EXPLAIN SYNTAX SELECT * FROM prewhere_move_select_final FINAL WHERE y > 100;
+
+-- can not be pushed down
+EXPLAIN SYNTAX SELECT * FROM prewhere_move_select_final FINAL WHERE z > 400;
+
+-- only y can be pushed down
+EXPLAIN SYNTAX SELECT * FROM prewhere_move_select_final FINAL WHERE y > 100 and z > 400;
+
+DROP TABLE prewhere_move_select_final;