From 7c76ce808e1d11dcb399927174f289a8c590bfc4 Mon Sep 17 00:00:00 2001 From: Jechol Lee Date: Sat, 7 Nov 2026 06:39:23 +0000 Subject: [PATCH] perf: narrow implicit from_many subquery selections --- lib/aggregate/grouped/query.ex | 7 +- lib/aggregate/lateral/query.ex | 8 +- lib/join.ex | 9 ++ lib/query/projection.ex | 235 +++++++++++++++++++++++++++++++++ test/projection_test.exs | 94 +++++++++++++ 5 files changed, 350 insertions(+), 3 deletions(-) create mode 100644 lib/query/projection.ex create mode 100644 test/projection_test.exs diff --git a/lib/aggregate/grouped/query.ex b/lib/aggregate/grouped/query.ex index 90fab94..c3dbbdc 100644 --- a/lib/aggregate/grouped/query.ex +++ b/lib/aggregate/grouped/query.ex @@ -62,7 +62,11 @@ defmodule AshSql.Aggregate.Grouped.Query do repo = AshSql.dynamic_repo(resource, implementation, query) - value = repo.one(query, AshSql.repo_opts(repo, implementation, nil, nil, resource)) + value = + repo.one( + AshSql.Query.Projection.prune(query), + AshSql.repo_opts(repo, implementation, nil, nil, resource) + ) {:ok, maybe_default_value(value, aggregate)} end @@ -78,6 +82,7 @@ defmodule AshSql.Aggregate.Grouped.Query do result = query + |> AshSql.Query.Projection.prune() |> repo.one(AshSql.repo_opts(repo, implementation, nil, nil, resource)) |> Map.get(aggregate.name) diff --git a/lib/aggregate/lateral/query.ex b/lib/aggregate/lateral/query.ex index bdc067d..89b1500 100644 --- a/lib/aggregate/lateral/query.ex +++ b/lib/aggregate/lateral/query.ex @@ -101,7 +101,11 @@ defmodule AshSql.Aggregate.Lateral.Query do _ -> repo = AshSql.dynamic_repo(resource, implementation, query) - repo.one(group_query, AshSql.repo_opts(repo, implementation, nil, nil, resource)) + + repo.one( + AshSql.Query.Projection.prune(group_query), + AshSql.repo_opts(repo, implementation, nil, nil, resource) + ) end {:ok, add_single_aggs(result, resource, query, cant_group, implementation)} @@ -244,7 +248,7 @@ defmodule AshSql.Aggregate.Lateral.Query do Map.merge( result || %{}, repo.one( - query, + AshSql.Query.Projection.prune(query), AshSql.repo_opts(repo, query.__ash_bindings__.sql_behaviour, nil, nil, resource) ) ) diff --git a/lib/join.ex b/lib/join.ex index fbfbeda..00ad30f 100644 --- a/lib/join.ex +++ b/lib/join.ex @@ -694,6 +694,13 @@ defmodule AshSql.Join do ) end + query = + if opts[:project_relationship?] do + put_in(query.__ash_bindings__[:relationship_projection?], true) + else + query + end + {:ok, query} end end @@ -1540,6 +1547,8 @@ defmodule AshSql.Join do refs_at_path: path, require_lateral?: require_lateral?, filter_subquery?: true, + project_relationship?: + Map.get(relationship, :from_many?) && is_nil(relationship.manual), join_binding: if(!Map.get(relationship, :from_many?), do: initial_ash_bindings.current), on_subquery: fn subquery -> diff --git a/lib/query/projection.ex b/lib/query/projection.ex new file mode 100644 index 0000000..c0ec6aa --- /dev/null +++ b/lib/query/projection.ex @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: 2024 ash_sql contributors +# +# SPDX-License-Identifier: MIT + +defmodule AshSql.Query.Projection do + @moduledoc false + import Ecto.Query, only: [from: 2] + + @clauses [ + :select, + :wheres, + :havings, + :group_bys, + :order_bys, + :distinct, + :windows, + :limit, + :offset + ] + + # Run at the execution boundary: later calculations can reuse an existing join. + def prune(query) do + if candidate?(query) do + {query, fields} = collect_query(query, [], %{}) + rewrite(query, fields) + else + query + end + end + + defp candidate?(%Ecto.SubQuery{query: query}), do: candidate?(query) + + defp candidate?(%{from: %{source: source}, joins: joins} = query) do + Map.get(query, :__ash_bindings__, %{})[:relationship_projection?] || + candidate?(source) || Enum.any?(joins, &candidate?(&1.source)) || + Enum.any?(@clauses, &candidate?(Map.get(query, &1))) + end + + defp candidate?(%{subqueries: subqueries}), do: Enum.any?(subqueries, &candidate?/1) + defp candidate?(list) when is_list(list), do: Enum.any?(list, &candidate?/1) + defp candidate?({_name, value}), do: candidate?(value) + defp candidate?(_), do: false + + defp collect_query(query, parents, fields) do + id = make_ref() + scope = [{id, query.aliases} | parents] + query = Map.put(query, :__projection_scope__, id) + + {query, fields} = map_children(query, fields, &collect_query(&1, scope, &2)) + + expressions = + [query.from && query.from.source | Enum.map(@clauses, &Map.get(query, &1))] ++ + Enum.flat_map(query.joins, &[&1.source, &1.on]) + + fields = Enum.reduce(expressions, fields, &collect_expr(&1, scope, &2)) + + fields = + if query.assocs != [] || query.preloads != [] || query.combinations != [] || query.with_ctes do + Enum.reduce(scope, fields, fn {scope_id, aliases}, fields -> + Enum.reduce(Map.values(aliases), fields, &Map.put(&2, {scope_id, &1}, :all)) + end) + else + fields + end + + {query, fields} + end + + defp collect_expr({{:., _, [binding, field]}, _, []}, scope, fields) when is_atom(field), + do: reference(binding, field, scope, fields) + + defp collect_expr({kind, _, [_]} = binding, scope, fields) when kind in [:&, :as, :parent_as], + do: reference(binding, :all, scope, fields) + + defp collect_expr({:fragment, _, args}, scope, fields) do + # The CASE fragment emitted by AshSql has no column references in its SQL + # text. Other fragments may name columns without using a bound expression. + if Enum.filter(args, &match?({:raw, _}, &1)) == + [raw: "(CASE WHEN ", raw: " THEN ", raw: " ELSE ", raw: " END)"] do + collect_expr(args, scope, fields) + else + Enum.reduce(scope, fields, fn {id, aliases}, fields -> + Enum.reduce(Map.values(aliases), fields, &Map.put(&2, {id, &1}, :all)) + end) + end + end + + # Subqueries have already been visited with their own binding scope. Do not walk + # parameter values or Ash context/actor data while looking for SQL references. + defp collect_expr(%{expr: expr}, scope, fields), do: collect_expr(expr, scope, fields) + + defp collect_expr(tuple, scope, fields) when is_tuple(tuple), + do: tuple |> Tuple.to_list() |> collect_expr(scope, fields) + + defp collect_expr(list, scope, fields) when is_list(list), + do: Enum.reduce(list, fields, &collect_expr(&1, scope, &2)) + + defp collect_expr(_, _, fields), do: fields + + defp reference({:&, _, [index]}, field, [{id, _} | _], fields), + do: add_field(fields, {id, index}, field) + + defp reference({:as, _, [name]}, field, [{id, aliases} | _], fields), + do: named_field(fields, id, aliases, name, field) + + defp reference({:parent_as, _, [name]}, field, [_ | parents], fields) do + # Interpolated join queries can be flattened by Ecto. Retaining matches in + # every enclosing scope is conservative when a binding name is shadowed. + Enum.reduce(parents, fields, fn {id, aliases}, fields -> + named_field(fields, id, aliases, name, field) + end) + end + + defp reference(_, _, _, fields), do: fields + + defp named_field(fields, id, aliases, name, field) do + case Map.fetch(aliases, name) do + {:ok, index} -> add_field(fields, {id, index}, field) + :error -> fields + end + end + + defp add_field(fields, key, field) do + Map.update(fields, key, if(field == :all, do: :all, else: MapSet.new([field])), fn + :all -> :all + _ when field == :all -> :all + previous -> MapSet.put(previous, field) + end) + end + + defp rewrite(query, fields) do + joins = + Enum.with_index(query.joins, 1) + |> Enum.map(fn {join, index} -> + needed = Map.get(fields, {query.__projection_scope__, index}, MapSet.new()) + + source = + if is_map(join.source) && + Map.get(join.source, :__ash_bindings__, %{})[:relationship_projection?] do + restrict(join.source, needed, fields) + else + join.source + end + + %{join | source: source} + end) + + {query, _} = + map_children(%{query | joins: joins}, nil, fn child, _ -> {rewrite(child, fields), nil} end) + + Map.delete(query, :__projection_scope__) + end + + defp restrict(query, :all, _fields), do: query + + defp restrict(query, needed, fields) do + if is_nil(query.select) && is_nil(query.distinct) && query.combinations == [] && + query.group_bys == [] && query.havings == [] && is_nil(query.with_ctes) do + case {query.from.source, needed} do + {%Ecto.SubQuery{query: inner} = subquery, needed} -> + local = Map.get(fields, {query.__projection_scope__, 0}, MapSet.new()) + needed = if local == :all, do: :all, else: MapSet.union(needed, local) + + %{ + query + | from: %{query.from | source: %{subquery | query: restrict(inner, needed, fields)}} + } + + {{table, schema}, needed} + when is_binary(table) and is_atom(schema) and not is_nil(schema) -> + attributes = schema.__schema__(:fields) + selected = Enum.filter(attributes, &MapSet.member?(needed, &1)) + + if selected != [] && length(selected) == MapSet.size(needed) && + length(selected) < length(attributes) do + from(row in query, select: struct(row, ^selected)) + else + query + end + + _ -> + query + end + else + query + end + end + + defp map_children(query, acc, mapper) do + {from, acc} = + if query.from do + {source, acc} = map_source(query.from.source, acc, mapper) + {%{query.from | source: source}, acc} + else + {nil, acc} + end + + {joins, acc} = + Enum.map_reduce(query.joins, acc, fn join, acc -> + {source, acc} = map_source(join.source, acc, mapper) + {on, acc} = map_clause(join.on, acc, mapper) + {%{join | source: source, on: on}, acc} + end) + + query = %{query | from: from, joins: joins} + + Enum.reduce(@clauses, {query, acc}, fn key, {query, acc} -> + {value, acc} = map_clause(Map.get(query, key), acc, mapper) + {Map.put(query, key, value), acc} + end) + end + + defp map_source(%Ecto.SubQuery{query: query} = subquery, acc, mapper) do + {query, acc} = mapper.(query, acc) + {%{subquery | query: query}, acc} + end + + defp map_source(%{from: _, joins: _} = query, acc, mapper), do: mapper.(query, acc) + defp map_source(other, acc, _mapper), do: {other, acc} + + defp map_clause(%{subqueries: subqueries} = clause, acc, mapper) do + {subqueries, acc} = Enum.map_reduce(subqueries, acc, &map_source(&1, &2, mapper)) + {%{clause | subqueries: subqueries}, acc} + end + + defp map_clause(list, acc, mapper) when is_list(list), + do: Enum.map_reduce(list, acc, &map_clause(&1, &2, mapper)) + + defp map_clause({name, value}, acc, mapper) do + {value, acc} = map_clause(value, acc, mapper) + {{name, value}, acc} + end + + defp map_clause(other, acc, _mapper), do: {other, acc} +end diff --git a/test/projection_test.exs b/test/projection_test.exs new file mode 100644 index 0000000..816c1a6 --- /dev/null +++ b/test/projection_test.exs @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: 2024 ash_sql contributors +# +# SPDX-License-Identifier: MIT + +defmodule AshSql.Query.ProjectionTest do + use ExUnit.Case, async: true + import Ecto.Query + alias AshSql.Query.Projection + + defmodule Comment do + use Ecto.Schema + + schema "comments" do + field(:post_id, :integer) + field(:title, :string) + field(:likes, :integer) + end + end + + defp related_comment(inner_fun \\ & &1) do + inner = from(c in Comment, as: :comment, where: c.post_id == parent_as(:post).id, limit: 1) + + from(c in subquery(inner_fun.(inner)), as: :comment) + |> Map.put(:__ash_bindings__, %{relationship_projection?: true}) + end + + defp posts(source \\ related_comment()) do + from(p in "posts", + as: :post, + left_lateral_join: c in ^source, + as: :related_comment, + on: true, + select: %{id: p.id, title: c.title} + ) + end + + defp selected_fields(query) do + query.joins + |> hd() + |> Map.fetch!(:source) + |> Map.fetch!(:from) + |> Map.fetch!(:source) + |> Map.fetch!(:query) + |> Map.fetch!(:select) + end + + test "retains references from later filters and sorts on a reused join" do + result = + posts() |> where([_p, c], c.likes > 0) |> order_by([_p, c], c.id) |> Projection.prune() + + assert selected_fields(result).take == %{0 => {:struct, [:id, :title, :likes]}} + refute Map.has_key?(result, :__projection_scope__) + end + + test "retains fields referenced only by a correlated subquery" do + ratings = from(r in "ratings", where: r.score > parent_as(:related_comment).likes, select: 1) + + result = + posts() |> select_merge(%{has_rating: exists(subquery(ratings))}) |> Projection.prune() + + assert selected_fields(result).take == %{0 => {:struct, [:title, :likes]}} + end + + test "retains a wrapper filter's fields below the relationship limit" do + source = related_comment() |> where([c], c.likes > 0) + result = source |> posts() |> Projection.prune() + assert selected_fields(result).take == %{0 => {:struct, [:title, :likes]}} + assert hd(result.joins).source.wheres == source.wheres + end + + test "keeps the implicit selection for a whole-row fragment" do + result = + posts() + |> select_merge([_p, c], %{comment: fragment("row_to_json(?)", c)}) + |> Projection.prune() + + assert selected_fields(result) == nil + end + + test "keeps fields that opaque SQL may reference without a binding" do + result = posts() |> where(fragment("likes > 0")) |> Projection.prune() + assert selected_fields(result) == nil + end + + test "does not change a DISTINCT subquery's implicit selection" do + result = related_comment(&distinct(&1, true)) |> posts() |> Projection.prune() + assert selected_fields(result) == nil + end + + test "does not replace an explicit selection" do + result = related_comment(&select(&1, [c], %{title: c.title, likes: c.likes})) |> posts() + assert selected_fields(Projection.prune(result)) == selected_fields(result) + end +end