Fix optimize_skip_unused_shards with virtual columns of the underlying table in WHERE

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -957,7 +957,7 @@ ExpressionAnalysisResult::ExpressionAnalysisResult(
     const ASTSelectQuery & query = *query_analyzer.getSelectQuery();
     const Context & context = query_analyzer.context;
     const Settings & settings = context.getSettingsRef();
-    const StoragePtr & storage = query_analyzer.storage();
+    const ConstStoragePtr & storage = query_analyzer.storage();
 
     bool finalized = false;
     size_t where_step_num = 0;

dbms/src/Interpreters/ExpressionAnalyzer.h

@@ -122,7 +122,7 @@ class ExpressionAnalyzer : protected ExpressionAnalyzerData, private boost::nonc
 
     SyntaxAnalyzerResultPtr syntax;
 
-    const StoragePtr & storage() const { return syntax->storage; } /// The main table in FROM clause, if exists.
+    const ConstStoragePtr & storage() const { return syntax->storage; } /// The main table in FROM clause, if exists.
     const AnalyzedJoin & analyzedJoin() const { return *syntax->analyzed_join; }
     const NamesAndTypesList & sourceColumns() const { return syntax->required_source_columns; }
     const std::vector<const ASTFunction *> & aggregates() const { return syntax->aggregates; }

dbms/src/Interpreters/SyntaxAnalyzer.cpp

@@ -867,7 +867,7 @@ SyntaxAnalyzerResultPtr SyntaxAnalyzer::analyzeSelect(
     return std::make_shared<const SyntaxAnalyzerResult>(result);
 }
 
-SyntaxAnalyzerResultPtr SyntaxAnalyzer::analyze(ASTPtr & query, const NamesAndTypesList & source_columns, StoragePtr storage) const
+SyntaxAnalyzerResultPtr SyntaxAnalyzer::analyze(ASTPtr & query, const NamesAndTypesList & source_columns, ConstStoragePtr storage) const
 {
     if (query->as<ASTSelectQuery>())
         throw Exception("Not select analyze for select asts.", ErrorCodes::LOGICAL_ERROR);

dbms/src/Interpreters/SyntaxAnalyzer.h

@@ -19,7 +19,7 @@ using Scalars = std::map<String, Block>;
 
 struct SyntaxAnalyzerResult
 {
-    StoragePtr storage;
+    ConstStoragePtr storage;
     std::shared_ptr<AnalyzedJoin> analyzed_join;
 
     NamesAndTypesList source_columns;
@@ -51,7 +51,7 @@ struct SyntaxAnalyzerResult
 
     bool maybe_optimize_trivial_count = false;
 
-    SyntaxAnalyzerResult(const NamesAndTypesList & source_columns_, StoragePtr storage_ = {}, bool add_virtuals = true)
+    SyntaxAnalyzerResult(const NamesAndTypesList & source_columns_, ConstStoragePtr storage_ = {}, bool add_virtuals = true)
         : storage(storage_)
         , source_columns(source_columns_)
     {
@@ -86,7 +86,7 @@ class SyntaxAnalyzer
     {}
 
     /// Analyze and rewrite not select query
-    SyntaxAnalyzerResultPtr analyze(ASTPtr & query, const NamesAndTypesList & source_columns_, StoragePtr storage = {}) const;
+    SyntaxAnalyzerResultPtr analyze(ASTPtr & query, const NamesAndTypesList & source_columns_, ConstStoragePtr storage = {}) const;
 
     /// Analyze and rewrite select query
     SyntaxAnalyzerResultPtr analyzeSelect(

dbms/src/Storages/IStorage_fwd.h

@@ -10,6 +10,7 @@ namespace DB
 
 class IStorage;
 
+using ConstStoragePtr = std::shared_ptr<const IStorage>;
 using StoragePtr = std::shared_ptr<IStorage>;
 using Tables = std::map<String, StoragePtr>;
 

dbms/src/Storages/StorageDistributed.cpp

@@ -227,9 +227,9 @@ class ReplacingConstantExpressionsMatcher
     }
 };
 
-void replaceConstantExpressions(ASTPtr & node, const Context & context, const NamesAndTypesList & columns)
+void replaceConstantExpressions(ASTPtr & node, const Context & context, const NamesAndTypesList & columns, ConstStoragePtr storage)
 {
-    auto syntax_result = SyntaxAnalyzer(context).analyze(node, columns);
+    auto syntax_result = SyntaxAnalyzer(context).analyze(node, columns, storage);
     Block block_with_constants = KeyCondition::getBlockWithConstants(node, syntax_result, context);
 
     InDepthNodeVisitor<ReplacingConstantExpressionsMatcher, true> visitor(block_with_constants);
@@ -662,7 +662,7 @@ ClusterPtr StorageDistributed::skipUnusedShards(ClusterPtr cluster, const Select
         condition_ast = select.prewhere() ? select.prewhere()->clone() : select.where()->clone();
     }
 
-    replaceConstantExpressions(condition_ast, context, getColumns().getAllPhysical() /** TODO: sharding_key_column_name */);
+    replaceConstantExpressions(condition_ast, context, getColumns().getAll(), shared_from_this());
     const auto blocks = evaluateExpressionOverConstantCondition(condition_ast, sharding_key_expr);
 
     // Can't get definite answer if we can skip any shards

dbms/tests/queries/0_stateless/01072_optimize_skip_unused_shards_const_expr_eval.sql

@@ -45,5 +45,12 @@ select * from dist_01072 where key=toInt32(value); -- { serverError 507; }
 select * from dist_01072 where key=value settings force_optimize_skip_unused_shards=0;
 select * from dist_01072 where key=toInt32(value) settings force_optimize_skip_unused_shards=0;
 
+-- check virtual columns
+drop table data_01072;
+drop table dist_01072;
+create table data_01072 (key Int) Engine=MergeTree() ORDER BY key;
+create table dist_01072 (key Int) Engine=Distributed(test_cluster_two_shards, currentDatabase(), data_01072, key);
+select * from dist_01072 where key=0 and _part='0' settings force_optimize_skip_unused_shards=2;
+
 drop table data_01072;
 drop table dist_01072;