improvement: leverage new aggregate loading optimization

Author: zachdaniel

lib/data_layer.ex

@@ -3619,14 +3619,11 @@ defmodule AshPostgres.DataLayer do
   end
 
   @impl true
-  def add_aggregates(query, aggregates, resource) do
-    AshSql.Aggregate.add_aggregates(
-      query,
-      aggregates,
-      resource,
-      true,
-      query.__ash_bindings__.root_binding
-    )
+  def add_aggregates(query, aggregates, _resource) do
+    {:ok,
+     Map.update!(query, :__ash_bindings__, fn bindings ->
+       Map.put(bindings, :load_aggregates, aggregates)
+     end)}
   end
 
   @impl true

mix.lock

@@ -1,6 +1,6 @@
 %{
   "ash": {:hex, :ash, "3.6.2", "90d1c8296be777b90caabf51b99323d6618a0b92594dfab92b02bdf848ac38bf", [:mix], [{:decimal, "~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:ecto, "~> 3.7", [hex: :ecto, repo: "hexpm", optional: false]}, {:ets, "~> 0.8", [hex: :ets, repo: "hexpm", optional: false]}, {:igniter, ">= 0.6.29 and < 1.0.0-0", [hex: :igniter, repo: "hexpm", optional: true]}, {:jason, ">= 1.0.0", [hex: :jason, repo: "hexpm", optional: false]}, {:picosat_elixir, "~> 0.2", [hex: :picosat_elixir, repo: "hexpm", optional: true]}, {:plug, ">= 0.0.0", [hex: :plug, repo: "hexpm", optional: true]}, {:reactor, "~> 0.11", [hex: :reactor, repo: "hexpm", optional: false]}, {:simple_sat, ">= 0.1.1 and < 1.0.0-0", [hex: :simple_sat, repo: "hexpm", optional: true]}, {:spark, ">= 2.3.3 and < 3.0.0-0", [hex: :spark, repo: "hexpm", optional: false]}, {:splode, ">= 0.2.6 and < 1.0.0-0", [hex: :splode, repo: "hexpm", optional: false]}, {:stream_data, "~> 1.0", [hex: :stream_data, repo: "hexpm", optional: false]}, {:telemetry, "~> 1.1", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "3546b5798dd24576cc451f6e03f3d6e3bb62666c0921bfe8aae700c599d9c38d"},
-  "ash_sql": {:hex, :ash_sql, "0.3.2", "e2d65dac1c813cbd2569a750bf1c063109778e840052e44535ced294d7638a19", [:mix], [{:ash, ">= 3.5.43 and < 4.0.0-0", [hex: :ash, repo: "hexpm", optional: false]}, {:ecto, "~> 3.9", [hex: :ecto, repo: "hexpm", optional: false]}, {:ecto_sql, "~> 3.9", [hex: :ecto_sql, repo: "hexpm", optional: false]}], "hexpm", "1f6e5d827c0eb55fc5a07f58eb97f9bb3e6b290d83df75883f422537b98c9c68"},
+  "ash_sql": {:hex, :ash_sql, "0.3.4", "c8c0446fbd6d3e6920f793b971c83ba3d14f96095036366d313b72656400509d", [:mix], [{:ash, ">= 3.5.43 and < 4.0.0-0", [hex: :ash, repo: "hexpm", optional: false]}, {:ecto, "~> 3.9", [hex: :ecto, repo: "hexpm", optional: false]}, {:ecto_sql, "~> 3.9", [hex: :ecto_sql, repo: "hexpm", optional: false]}], "hexpm", "4edb1fb707048a41f7944274b1a0b571aa3c9117b8a7a12809ca023b6f955cb4"},
   "benchee": {:hex, :benchee, "1.4.0", "9f1f96a30ac80bab94faad644b39a9031d5632e517416a8ab0a6b0ac4df124ce", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}, {:statistex, "~> 1.0", [hex: :statistex, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "299cd10dd8ce51c9ea3ddb74bb150f93d25e968f93e4c1fa31698a8e4fa5d715"},
   "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
   "credo": {:hex, :credo, "1.7.12", "9e3c20463de4b5f3f23721527fcaf16722ec815e70ff6c60b86412c695d426c1", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "8493d45c656c5427d9c729235b99d498bd133421f3e0a683e5c1b561471291e5"},
@@ -17,7 +17,7 @@
   "ets": {:hex, :ets, "0.9.0", "79c6a6c205436780486f72d84230c6cba2f8a9920456750ddd1e47389107d5fd", [:mix], [], "hexpm", "2861fdfb04bcaeff370f1a5904eec864f0a56dcfebe5921ea9aadf2a481c822b"},
   "ex_check": {:hex, :ex_check, "0.16.0", "07615bef493c5b8d12d5119de3914274277299c6483989e52b0f6b8358a26b5f", [:mix], [], "hexpm", "4d809b72a18d405514dda4809257d8e665ae7cf37a7aee3be6b74a34dec310f5"},
   "ex_doc": {:hex, :ex_doc, "0.38.4", "ab48dff7a8af84226bf23baddcdda329f467255d924380a0cf0cee97bb9a9ede", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "f7b62346408a83911c2580154e35613eb314e0278aeea72ed7fedef9c1f165b2"},
-  "file_system": {:hex, :file_system, "1.1.0", "08d232062284546c6c34426997dd7ef6ec9f8bbd090eb91780283c9016840e8f", [:mix], [], "hexpm", "bfcf81244f416871f2a2e15c1b515287faa5db9c6bcf290222206d120b3d43f6"},
+  "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"},
   "finch": {:hex, :finch, "0.20.0", "5330aefb6b010f424dcbbc4615d914e9e3deae40095e73ab0c1bb0968933cadf", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.6.2 or ~> 1.7", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "2658131a74d051aabfcba936093c903b8e89da9a1b63e430bee62045fa9b2ee2"},
   "git_cli": {:hex, :git_cli, "0.3.0", "a5422f9b95c99483385b976f5d43f7e8233283a47cda13533d7c16131cb14df5", [:mix], [], "hexpm", "78cb952f4c86a41f4d3511f1d3ecb28edb268e3a7df278de2faa1bd4672eaf9b"},
   "git_ops": {:hex, :git_ops, "2.9.0", "b74f6040084f523055b720cc7ef718da47f2cbe726a5f30c2871118635cb91c1", [:mix], [{:git_cli, "~> 0.2", [hex: :git_cli, repo: "hexpm", optional: false]}, {:igniter, ">= 0.5.27 and < 1.0.0-0", [hex: :igniter, repo: "hexpm", optional: true]}, {:nimble_parsec, "~> 1.0", [hex: :nimble_parsec, repo: "hexpm", optional: false]}, {:req, "~> 0.5", [hex: :req, repo: "hexpm", optional: false]}], "hexpm", "7fdf84be3490e5692c5dc1f8a1084eed47a221c1063e41938c73312f0bfea259"},
@@ -42,7 +42,7 @@
   "req": {:hex, :req, "0.5.15", "662020efb6ea60b9f0e0fac9be88cd7558b53fe51155a2d9899de594f9906ba9", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "a6513a35fad65467893ced9785457e91693352c70b58bbc045b47e5eb2ef0c53"},
   "rewrite": {:hex, :rewrite, "1.2.0", "80220eb14010e175b67c939397e1a8cdaa2c32db6e2e0a9d5e23e45c0414ce21", [:mix], [{:glob_ex, "~> 0.1", [hex: :glob_ex, repo: "hexpm", optional: false]}, {:sourceror, "~> 1.0", [hex: :sourceror, repo: "hexpm", optional: false]}, {:text_diff, "~> 0.1", [hex: :text_diff, repo: "hexpm", optional: false]}], "hexpm", "a1cd702bbb9d51613ab21091f04a386d750fc6f4516b81900df082d78b2d8c50"},
   "simple_sat": {:hex, :simple_sat, "0.1.3", "f650fc3c184a5fe741868b5ac56dc77fdbb428468f6dbf1978e14d0334497578", [:mix], [], "hexpm", "a54305066a356b7194dc81db2a89232bacdc0b3edaef68ed9aba28dcbc34887b"},
-  "sobelow": {:hex, :sobelow, "0.14.0", "dd82aae8f72503f924fe9dd97ffe4ca694d2f17ec463dcfd365987c9752af6ee", [:mix], [{:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "7ecf91e298acfd9b24f5d761f19e8f6e6ac585b9387fb6301023f1f2cd5eed5f"},
+  "sobelow": {:hex, :sobelow, "0.14.1", "2f81e8632f15574cba2402bcddff5497b413c01e6f094bc0ab94e83c2f74db81", [:mix], [{:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "8fac9a2bd90fdc4b15d6fca6e1608efb7f7c600fa75800813b794ee9364c87f2"},
   "sourceror": {:hex, :sourceror, "1.10.0", "38397dedbbc286966ec48c7af13e228b171332be1ad731974438c77791945ce9", [:mix], [], "hexpm", "29dbdfc92e04569c9d8e6efdc422fc1d815f4bd0055dc7c51b8800fb75c4b3f1"},
   "spark": {:hex, :spark, "2.3.5", "f30d30ecc3b4ab9b932d9aada66af7677fc1f297a2c349b0bcec3eafb9f996e8", [:mix], [{:igniter, ">= 0.3.64 and < 1.0.0-0", [hex: :igniter, repo: "hexpm", optional: true]}, {:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: true]}, {:sourceror, "~> 1.2", [hex: :sourceror, repo: "hexpm", optional: true]}], "hexpm", "0e9d339704d5d148f77f2b2fef3bcfc873a9e9bb4224fcf289c545d65827202f"},
   "spitfire": {:hex, :spitfire, "0.2.1", "29e154873f05444669c7453d3d931820822cbca5170e88f0f8faa1de74a79b47", [:mix], [], "hexpm", "6eeed75054a38341b2e1814d41bb0a250564092358de2669fdb57ff88141d91b"},

test/aggregate_test.exs

@@ -1734,4 +1734,130 @@ defmodule AshSql.AggregateTest do
 
     assert loaded_post.count_of_comments == 1
   end
+
+  test "aggregate with sort and limit is accurate" do
+    # Setup: Create an author with multiple posts
+    author =
+      Author
+      |> Ash.Changeset.for_create(:create, %{first_name: "John", last_name: "Doe"})
+      |> Ash.create!()
+
+    # Create posts with different titles to test sorting
+    post1 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "A First Post"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    post2 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "Z Last Post"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    post3 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "M Middle Post"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    # Add comments to posts
+    Comment
+    |> Ash.Changeset.for_create(:create, %{title: "Comment 1"})
+    |> Ash.Changeset.manage_relationship(:post, post1, type: :append_and_remove)
+    |> Ash.create!()
+
+    Comment
+    |> Ash.Changeset.for_create(:create, %{title: "Comment 2"})
+    |> Ash.Changeset.manage_relationship(:post, post1, type: :append_and_remove)
+    |> Ash.create!()
+
+    Comment
+    |> Ash.Changeset.for_create(:create, %{title: "Comment 3"})
+    |> Ash.Changeset.manage_relationship(:post, post2, type: :append_and_remove)
+    |> Ash.create!()
+
+    Comment
+    |> Ash.Changeset.for_create(:create, %{title: "Comment 4"})
+    |> Ash.Changeset.manage_relationship(:post, post3, type: :append_and_remove)
+    |> Ash.create!()
+
+    # Query with aggregate, sort, and limit
+    # This should ideally use a subquery to apply sort/limit before loading the aggregate
+    results =
+      Post
+      |> Ash.Query.load(:count_of_comments)
+      |> Ash.Query.sort(:title)
+      |> Ash.Query.limit(2)
+      |> Ash.read!()
+
+    # Verify we got the right posts (sorted by title, limited to 2)
+    assert length(results) == 2
+    assert Enum.at(results, 0).title == "A First Post"
+    assert Enum.at(results, 1).title == "M Middle Post"
+
+    # Verify the aggregates are correct
+    assert Enum.at(results, 0).count_of_comments == 2
+    assert Enum.at(results, 1).count_of_comments == 1
+  end
+
+  test "aggregate with sort by aggregate value and limit is accurate" do
+    # This tests sorting by the aggregate itself, not by another field
+    author =
+      Author
+      |> Ash.Changeset.for_create(:create, %{first_name: "Jane", last_name: "Smith"})
+      |> Ash.create!()
+
+    # Create posts with different numbers of comments
+    post1 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "Post with 3 comments"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    post2 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "Post with 1 comment"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    post3 =
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "Post with 2 comments"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.create!()
+
+    # Add varying numbers of comments
+    for _i <- 1..3 do
+      Comment
+      |> Ash.Changeset.for_create(:create, %{title: "Comment on post 1"})
+      |> Ash.Changeset.manage_relationship(:post, post1, type: :append_and_remove)
+      |> Ash.create!()
+    end
+
+    Comment
+    |> Ash.Changeset.for_create(:create, %{title: "Comment on post 2"})
+    |> Ash.Changeset.manage_relationship(:post, post2, type: :append_and_remove)
+    |> Ash.create!()
+
+    for _i <- 1..2 do
+      Comment
+      |> Ash.Changeset.for_create(:create, %{title: "Comment on post 3"})
+      |> Ash.Changeset.manage_relationship(:post, post3, type: :append_and_remove)
+      |> Ash.create!()
+    end
+
+    # Query sorting by the aggregate value itself
+    results =
+      Post
+      |> Ash.Query.load(:count_of_comments)
+      |> Ash.Query.sort(count_of_comments: :desc)
+      |> Ash.Query.limit(2)
+      |> Ash.read!()
+
+    # Should get the posts with most comments first
+    assert length(results) == 2
+    assert Enum.at(results, 0).count_of_comments == 3
+    assert Enum.at(results, 1).count_of_comments == 2
+  end
 end