fix: retain direct ordering for compatible distinct pagination

lib/query.ex

@@ -328,7 +328,7 @@ defmodule AshSql.Query do
       {:ok, query} ->
         query =
           if query.__ash_bindings__[:__order__?] && query.windows[:order] do
-            if query.distinct do
+            if query.distinct && !distinct_matches_sort?(query) do
               {calculations_require_rewrite, aggregates_require_rewrite, query} =
                 rewrite_nested_selects(query)
 
@@ -431,6 +431,42 @@ defmodule AshSql.Query do
     end
   end
 
+  # Ecto prepends DISTINCT ON fields to ORDER BY, merging their common prefix.
+  # If either ordering is a prefix of the other, that preserves the requested sort.
+  # Keep it directly usable by the planner instead of sorting computed row numbers
+  # outside a subquery, which can force every row to be read before LIMIT.
+  defp distinct_matches_sort?(
+         %{
+           distinct: %Ecto.Query.ByExpr{expr: distinct, params: [], subqueries: []}
+         } = query
+       )
+       when is_list(distinct) do
+    case query.windows[:order] do
+      %Ecto.Query.ByExpr{expr: [order_by: order], params: [], subqueries: []} ->
+        normalize = fn expr ->
+          Macro.prewalk(expr, fn
+            {:as, _, [name]} = binding ->
+              case Map.fetch(query.aliases, name) do
+                {:ok, index} -> {:&, [], [index]}
+                :error -> binding
+              end
+
+            other ->
+              other
+          end)
+        end
+
+        distinct = normalize.(distinct)
+        order = normalize.(order)
+        List.starts_with?(distinct, order) || List.starts_with?(order, distinct)
+
+      _ ->
+        false
+    end
+  end
+
+  defp distinct_matches_sort?(_), do: false
+
   defp set_lateral_join_prefix(ash_query, query) do
     if Ash.Resource.Info.multitenancy_strategy(ash_query.resource) == :context do
       tenant = query.prefix || ash_query.to_tenant

test/query_test.exs

@@ -5,6 +5,8 @@
 defmodule AshSql.QueryTest do
   use ExUnit.Case, async: false
 
+  import Ecto.Query
+
   defmodule RecordSharedContext do
     use Ash.Resource.Preparation
 
@@ -123,4 +125,141 @@ defmodule AshSql.QueryTest do
 
     assert_receive {:shared_context, ^shared}
   end
+
+  describe "distinct pagination ordering" do
+    test "keeps compatible ordering and pagination in the same query" do
+      query =
+        from(post in "posts",
+          as: ^0,
+          left_join: comment in "comments",
+          on: comment.post_id == post.id,
+          where: post.author_id == ^"author",
+          distinct: [desc: post.author_id, asc: post.id, asc: post.id],
+          select: %{id: post.id, author_id: post.author_id},
+          limit: ^51,
+          offset: ^50
+        )
+        |> sorted_query(author_id: :desc, id: :asc)
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+
+      assert result.from == query.from
+      assert result.distinct == query.distinct
+      assert result.joins == query.joins
+      assert result.wheres == query.wheres
+      assert result.select == query.select
+      assert result.limit == query.limit
+      assert result.offset == query.offset
+      assert result.windows == []
+      assert hd(result.order_bys).expr == query.windows[:order].expr[:order_by]
+    end
+
+    test "keeps tie breakers after a shorter distinct ordering" do
+      query =
+        from(post in "posts", as: ^0, distinct: [desc: post.author_id], select: post.id)
+        |> sorted_query(author_id: :desc, id: :asc)
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+      assert result.from == query.from
+      assert hd(result.order_bys).expr == query.windows[:order].expr[:order_by]
+      assert result.windows == []
+    end
+
+    test "preserves explicit null ordering" do
+      query =
+        from(post in "posts",
+          as: ^0,
+          distinct: [desc_nulls_last: post.author_id, asc: post.id],
+          select: post.id
+        )
+        |> sorted_query(author_id: :desc_nils_last, id: :asc)
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+      assert result.from == query.from
+      assert hd(result.order_bys).expr == query.windows[:order].expr[:order_by]
+    end
+
+    test "recognizes ordering through a nonzero named root binding" do
+      query =
+        from(post in "posts", as: ^500, distinct: [asc: post.id], select: post.id)
+        |> AshSql.Bindings.default_bindings(Post, __MODULE__, %{
+          data_layer: %{start_bindings_at: 500}
+        })
+        |> Map.update!(:__ash_bindings__, &Map.put(&1, :sort, id: :asc))
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+      assert result.from == query.from
+      assert result.windows == []
+    end
+
+    test "retains the wrapper when distinct would change the requested ordering" do
+      for distinct <- [[asc: :id], [asc: :author_id], [desc_nulls_first: :author_id]] do
+        query =
+          from(post in "posts", as: ^0, distinct: ^distinct, select: %{id: post.id}, limit: ^51)
+          |> sorted_query(author_id: :desc, id: :asc)
+
+        assert {:ok, result} = AshSql.Query.return_query(query, Post)
+        assert %Ecto.SubQuery{query: inner} = result.from.source
+        assert inner.distinct == query.distinct
+        assert inner.limit == nil
+        assert result.limit == query.limit
+        assert [asc: {{:., _, [_, :__order__]}, _, []}] = hd(result.order_bys).expr
+      end
+    end
+
+    test "retains the wrapper for bound sort expressions" do
+      query =
+        from(post in "posts",
+          as: ^0,
+          distinct: [asc: fragment("coalesce(?, ?)", post.author_id, ^"first")],
+          select: %{id: post.id},
+          windows: [
+            order: [order_by: [asc: fragment("coalesce(?, ?)", post.author_id, ^"second")]]
+          ]
+        )
+        |> AshSql.Bindings.default_bindings(Post, __MODULE__)
+        |> Map.update!(
+          :__ash_bindings__,
+          &Map.merge(&1, %{sort_applied?: true, __order__?: true})
+        )
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+      assert %Ecto.SubQuery{query: inner} = result.from.source
+      assert inner.distinct.params == query.distinct.params
+      assert inner.windows[:order].params == query.windows[:order].params
+    end
+
+    test "keeps independent subqueries in the distinct and display orders" do
+      highest = from(comment in "comments", select: max(comment.score))
+      lowest = from(comment in "comments", select: min(comment.score))
+
+      query =
+        from(post in "posts",
+          as: ^0,
+          distinct: [asc: subquery(highest)],
+          select: %{id: post.id},
+          windows: [order: [order_by: [asc: subquery(lowest)]]]
+        )
+        |> AshSql.Bindings.default_bindings(Post, __MODULE__)
+        |> Map.update!(
+          :__ash_bindings__,
+          &Map.merge(&1, %{sort_applied?: true, __order__?: true})
+        )
+
+      assert {:ok, result} = AshSql.Query.return_query(query, Post)
+      assert %Ecto.SubQuery{query: inner} = result.from.source
+      assert inner.distinct.subqueries == query.distinct.subqueries
+      assert inner.windows[:order].subqueries == query.windows[:order].subqueries
+    end
+  end
+
+  defp sorted_query(query, sort) do
+    query =
+      query
+      |> AshSql.Bindings.default_bindings(Post, __MODULE__)
+      |> Map.update!(:__ash_bindings__, &Map.put(&1, :sort, sort))
+
+    {:ok, query} = AshSql.Sort.apply_sort(query, sort, Post)
+    query
+  end
 end