From 61d40c3600ba6a1c6d6c0cf4919a3cdaebb3a31f Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Tue, 9 Mar 2021 22:00:38 +0300
Subject: [PATCH] Fix optimize_skip_unused_shards for zero shards case

v2: move check to the beginning of the StorageDistributed::read()
---
 src/Storages/StorageDistributed.cpp           | 19 +++++++++++++++++++
 ...e_skip_unused_shards_zero_shards.reference |  0
 ...ptimize_skip_unused_shards_zero_shards.sql |  2 ++
 .../queries/0_stateless/arcadia_skip_list.txt |  1 +
 4 files changed, 22 insertions(+)
 create mode 100644 tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.reference
 create mode 100644 tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.sql

diff --git a/src/Storages/StorageDistributed.cpp b/src/Storages/StorageDistributed.cpp
index 039cf63eca2..3a33e881611 100644
--- a/src/Storages/StorageDistributed.cpp
+++ b/src/Storages/StorageDistributed.cpp
@@ -48,6 +48,9 @@
 #include <Interpreters/getTableExpressions.h>
 #include <Functions/IFunction.h>
 
+#include <Processors/QueryPlan/ReadFromPreparedSource.h>
+#include <Processors/Sources/NullSource.h>
+
 #include <Core/Field.h>
 #include <Core/Settings.h>
 
@@ -83,6 +86,7 @@ namespace DB
 
 namespace ErrorCodes
 {
+    extern const int LOGICAL_ERROR;
     extern const int NOT_IMPLEMENTED;
     extern const int STORAGE_REQUIRES_PARAMETER;
     extern const int BAD_ARGUMENTS;
@@ -532,6 +536,17 @@ void StorageDistributed::read(
     Block header =
         InterpreterSelectQuery(query_info.query, context, SelectQueryOptions(processed_stage)).getSampleBlock();
 
+    /// Return directly (with correct header) if no shard to query.
+    if (query_info.cluster->getShardsInfo().empty())
+    {
+        Pipe pipe(std::make_shared<NullSource>(header));
+        auto read_from_pipe = std::make_unique<ReadFromPreparedSource>(std::move(pipe));
+        read_from_pipe->setStepDescription("Read from NullSource (Distributed)");
+        query_plan.addStep(std::move(read_from_pipe));
+
+        return;
+    }
+
     const Scalars & scalars = context.hasQueryContext() ? context.getQueryContext().getScalars() : Scalars{};
 
     bool has_virtual_shard_num_column = std::find(column_names.begin(), column_names.end(), "_shard_num") != column_names.end();
@@ -546,6 +561,10 @@ void StorageDistributed::read(
 
     ClusterProxy::executeQuery(query_plan, select_stream_factory, log,
         modified_query_ast, context, query_info);
+
+    /// This is a bug, it is possible only when there is no shards to query, and this is handled earlier.
+    if (!query_plan.isInitialized())
+        throw Exception("Pipeline is not initialized", ErrorCodes::LOGICAL_ERROR);
 }
 
 
diff --git a/tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.reference b/tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.reference
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.sql b/tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.sql
new file mode 100644
index 00000000000..b95d640ca1a
--- /dev/null
+++ b/tests/queries/0_stateless/01759_optimize_skip_unused_shards_zero_shards.sql
@@ -0,0 +1,2 @@
+create table dist_01756 (dummy UInt8) ENGINE = Distributed('test_cluster_two_shards', 'system', 'one', dummy);
+select ignore(1), * from dist_01756 where 0 settings optimize_skip_unused_shards=1, force_optimize_skip_unused_shards=1
diff --git a/tests/queries/0_stateless/arcadia_skip_list.txt b/tests/queries/0_stateless/arcadia_skip_list.txt
index 6926f16e027..a85ddba4752 100644
--- a/tests/queries/0_stateless/arcadia_skip_list.txt
+++ b/tests/queries/0_stateless/arcadia_skip_list.txt
@@ -218,3 +218,4 @@
 01682_cache_dictionary_complex_key
 01684_ssd_cache_dictionary_simple_key
 01685_ssd_cache_dictionary_complex_key
+01759_optimize_skip_unused_shards_zero_shards