Fix `count(:all)` to correctly work `distinct` with custom SELECT list

Currently `count(:all)` with `distinct` doesn't work correctly because
SELECT list is always replaced to `*` or primary key in that case even
if having custom SELECT list.

And also, PostgreSQL has a limitation that ORDER BY expressions must
appear in select list for SELECT DISTINCT.

Therefore, we should not replace custom SELECT list when using
`count(:all)` with `distinct`.

Closes #31277.

activerecord/CHANGELOG.md

+*   Fix `count(:all)` to correctly work `distinct` with custom SELECT list.
+
+    *Ryuta Kamizono*
+
 *   Fix conflicts `counter_cache` with `touch: true` by optimistic locking.
 
     ```

activerecord/lib/active_record/relation/calculations.rb

@@ -195,7 +195,7 @@ def perform_calculation(operation, column_name)
         if operation == "count"
           column_name ||= select_for_count
           if column_name == :all
-            if distinct && (group_values.any? || !(has_limit_or_offset? && order_values.any?))
+            if distinct && (group_values.any? || select_values.empty? && order_values.empty?)
               column_name = primary_key
             end
           elsif column_name =~ /\s*DISTINCT[\s(]+/i
@@ -227,7 +227,7 @@ def operation_over_aggregate_column(column, operation, distinct)
       def execute_simple_calculation(operation, column_name, distinct) #:nodoc:
         column_alias = column_name
 
-        if operation == "count" && has_limit_or_offset?
+        if operation == "count" && (column_name == :all && distinct || has_limit_or_offset?)
           # Shortcut when limit is zero.
           return 0 if limit_value == 0
 
@@ -369,14 +369,12 @@ def select_for_count
       end
 
       def build_count_subquery(relation, column_name, distinct)
-        relation.select_values = [
-          if column_name == :all
-            distinct ? table[Arel.star] : Arel.sql(FinderMethods::ONE_AS_ONE)
-          else
-            column_alias = Arel.sql("count_column")
-            aggregate_column(column_name).as(column_alias)
-          end
-        ]
+        if column_name == :all
+          relation.select_values = [ Arel.sql(FinderMethods::ONE_AS_ONE) ] unless distinct
+        else
+          column_alias = Arel.sql("count_column")
+          relation.select_values = [ aggregate_column(column_name).as(column_alias) ]
+        end
 
         subquery = relation.arel.as(Arel.sql("subquery_for_count"))
         select_value = operation_over_aggregate_column(column_alias || Arel.star, "count", false)

activerecord/test/cases/calculations_test.rb

@@ -241,6 +241,12 @@ def test_apply_distinct_in_count
     end
   end
 
+  def test_distinct_count_all_with_custom_select_and_order
+    accounts = Account.distinct.select("credit_limit % 10").order(Arel.sql("credit_limit % 10"))
+    assert_queries(1) { assert_equal 3, accounts.count(:all) }
+    assert_queries(1) { assert_equal 3, accounts.load.size }
+  end
+
   def test_distinct_count_with_order_and_limit
     assert_equal 4, Account.distinct.order(:firm_id).limit(4).count
   end

activerecord/test/cases/relations_test.rb

@@ -1121,6 +1121,12 @@ def test_count_with_distinct
     assert_equal 11, posts.distinct(false).select(:comments_count).count
   end
 
+  def test_size_with_distinct
+    posts = Post.distinct.select(:author_id, :comments_count)
+    assert_queries(1) { assert_equal 8, posts.size }
+    assert_queries(1) { assert_equal 8, posts.load.size }
+  end
+
   def test_update_all_with_scope
     tag = Tag.first
     Post.tagged_with(tag.id).update_all title: "rofl"