fix: load calculations from sorts properly

lib/ash/query/aggregate.ex

@@ -466,14 +466,18 @@ defmodule Ash.Query.Aggregate do
         end)
 
       sort_calculations =
-        Enum.flat_map(query.sort, fn {field, _} ->
-          case Map.fetch(query.calculations, field) do
-            :error ->
-              []
+        Enum.flat_map(query.sort, fn
+          {%Ash.Query.Calculation{} = calc, _} ->
+            [calc]
 
-            {:ok, calc} ->
-              [calc]
-          end
+          {field, _} ->
+            case Map.fetch(query.calculations, field) do
+              :error ->
+                []
+
+              {:ok, calc} ->
+                [calc]
+            end
         end)
 
       sort_calc_aggregates =

lib/ash/query/query.ex

@@ -1554,12 +1554,39 @@ defmodule Ash.Query do
         end)
         |> validate_sort()
 
-      Enum.reduce(query_with_sort.sort || [], query_with_sort, fn {key, _value}, query ->
-        if Ash.Resource.Info.aggregate(query.resource, key) do
-          Ash.Query.load(query, key)
-        else
-          query
-        end
+      Enum.reduce(query_with_sort.sort || [], query_with_sort, fn
+        {%Ash.Query.Calculation{name: name, module: module, opts: opts} = calculation, _},
+        query ->
+          {resource_load, resource_select} =
+            if resource_calculation = Ash.Resource.Info.calculation(query.resource, name) do
+              {resource_calculation.load, resource_calculation.select}
+            else
+              {[], []}
+            end
+
+          fields_to_select =
+            resource_select
+            |> Kernel.||([])
+            |> Enum.concat(module.select(query, opts, calculation.context) || [])
+
+          calculation = %{calculation | load: name, select: fields_to_select}
+
+          query =
+            query
+            |> module.load(
+              opts,
+              calculation.context
+              |> Map.put(:context, query.context)
+            )
+
+          Ash.Query.load(query, resource_load)
+
+        {key, _value}, query ->
+          if Ash.Resource.Info.aggregate(query.resource, key) do
+            Ash.Query.load(query, key)
+          else
+            query
+          end
       end)
     else
       add_error(query, :sort, "Data layer does not support sorting")