diff --git a/src/Interpreters/TreeRewriter.cpp b/src/Interpreters/TreeRewriter.cpp
index 1025e64e2f4..1ad22efc3fc 100644
--- a/src/Interpreters/TreeRewriter.cpp
+++ b/src/Interpreters/TreeRewriter.cpp
@@ -562,7 +562,7 @@ void TreeRewriterResult::collectUsedColumns(const ASTPtr & query, bool is_select
             /// If we have no information about columns sizes, choose a column of minimum size of its data type.
             required.insert(ExpressionActions::getSmallestColumn(source_columns));
     }
-    else if (is_select && metadata_snapshot)
+    else if (is_select && metadata_snapshot && !columns_context.has_array_join)
     {
         const auto & partition_desc = metadata_snapshot->getPartitionKey();
         if (partition_desc.expression)
diff --git a/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.reference b/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.reference
index 9db37bb5f81..b8b8fae2830 100644
--- a/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.reference
+++ b/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.reference
@@ -7,4 +7,6 @@
 0
 2
 0
+4
+6
 3
diff --git a/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.sql b/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.sql
index 47ad0128130..ecf0b791a49 100644
--- a/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.sql
+++ b/tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.sql
@@ -2,7 +2,7 @@ drop table if exists test1;
 drop table if exists test_tuple;
 drop table if exists test_two_args;
 
-create table test1(p DateTime, k int) engine MergeTree partition by toDate(p) order by k;
+create table test1(p DateTime, k int) engine MergeTree partition by toDate(p) order by k settings index_granularity = 1;
 insert into test1 values ('2020-09-01 00:01:02', 1), ('2020-09-01 20:01:03', 2), ('2020-09-02 00:01:03', 3);
 
 set max_rows_to_read = 1;
@@ -22,7 +22,7 @@ select count() from test1 where toDate(p) > '2020-09-01';
 -- non-optimized
 select count() from test1 where toDate(p) >= '2020-09-01' and p <= '2020-09-01 00:00:00';
 
-create table test_tuple(p DateTime, i int, j int) engine MergeTree partition by (toDate(p), i) order by j;
+create table test_tuple(p DateTime, i int, j int) engine MergeTree partition by (toDate(p), i) order by j settings index_granularity = 1;
 
 insert into test_tuple values ('2020-09-01 00:01:02', 1, 2), ('2020-09-01 00:01:03', 2, 3), ('2020-09-02 00:01:03', 3, 4);
 
@@ -34,8 +34,14 @@ select count() from test_tuple where toDate(p) > '2020-09-01' and i = 1;
 select count() from test_tuple where i > 1;
 -- optimized
 select count() from test_tuple where i < 1;
+-- non-optimized
+select count() from test_tuple array join [p,p] as c where toDate(p) = '2020-09-01'; -- { serverError 158; }
+select count() from test_tuple array join [1,2] as c where toDate(p) = '2020-09-01' settings max_rows_to_read = 4;
+-- non-optimized
+select count() from test_tuple array join [1,2,3] as c where toDate(p) = '2020-09-01'; -- { serverError 158; }
+select count() from test_tuple array join [1,2,3] as c where toDate(p) = '2020-09-01' settings max_rows_to_read = 6;
 
-create table test_two_args(i int, j int, k int) engine MergeTree partition by i + j order by k;
+create table test_two_args(i int, j int, k int) engine MergeTree partition by i + j order by k settings index_granularity = 1;
 
 insert into test_two_args values (1, 2, 3), (2, 1, 3), (0, 3, 4);