Implementation of over, partition by and window for PostgreSQL

Author: Anber

lib/ecto/adapters/postgres/connection.ex

@@ -109,7 +109,7 @@ if Code.ensure_loaded?(Postgrex) do
     end
 
     alias Ecto.Query
-    alias Ecto.Query.{BooleanExpr, JoinExpr, QueryExpr}
+    alias Ecto.Query.{BooleanExpr, JoinExpr, QueryExpr, PartitionByExpr}
 
     def all(query) do
       sources = create_names(query)
@@ -121,12 +121,13 @@ if Code.ensure_loaded?(Postgrex) do
       where    = where(query, sources)
       group_by = group_by(query, sources)
       having   = having(query, sources)
+      window   = window(query, sources)
       order_by = order_by(query, order_by_distinct, sources)
       limit    = limit(query, sources)
       offset   = offset(query, sources)
       lock     = lock(query.lock)
 
-      [select, from, join, where, group_by, having, order_by, limit, offset | lock]
+      [select, from, join, where, group_by, having, window, order_by, limit, offset | lock]
     end
 
     def update_all(%{from: %{source: source}} = query, prefix \\ nil) do
@@ -380,6 +381,22 @@ if Code.ensure_loaded?(Postgrex) do
        end)]
     end
 
+    defp window(%Query{windows: []}, _sources), do: []
+    defp window(%Query{windows: windows} = query, sources) do
+      [" WINDOW " |
+       intersperse_map(windows, ", ", fn
+         {name, definition} ->
+           [Atom.to_string(name), " AS ", partition_by(definition, sources, query)] end)]
+    end
+
+    defp partition_by(%PartitionByExpr{fields: fields, order_bys: order_bys}, sources, %Query{} = query) do
+      partition_query = %Query{ query | order_bys: order_bys }
+      ["(PARTITION BY ",
+        intersperse_map(fields, ", ", &expr(&1, sources, query)),
+        order_by(partition_query, [], sources),
+        ?)]
+    end
+
     defp order_by(%Query{order_bys: []}, _distinct, _sources), do: []
     defp order_by(%Query{order_bys: order_bys} = query, distinct, sources) do
       order_bys = Enum.flat_map(order_bys, & &1.expr)
@@ -509,6 +526,25 @@ if Code.ensure_loaded?(Postgrex) do
       [aggregate, " FILTER (WHERE ", expr(filter, sources, query), ?)]
     end
 
+    defp expr({:over, _, [agg, %PartitionByExpr{} = window]}, sources, query) do
+      aggregate = expr(agg, sources, query)
+      [aggregate, " OVER ", partition_by(window, sources, query)]
+    end
+
+    defp expr({:over, _, [agg, nil]}, sources, query) do
+      aggregate = expr(agg, sources, query)
+      [aggregate, " OVER ()"]
+    end
+
+    defp expr({:over, _, [agg, name]}, sources, %Ecto.Query{ windows: windows } = query) when is_atom(name) do
+      if Keyword.has_key?(windows, name) do
+        aggregate = expr(agg, sources, query)
+        [aggregate, " OVER ", Atom.to_string(name)]
+      else
+        error!(query, "window #{name} is undefined")
+      end
+    end
+
     defp expr({:{}, _, elems}, sources, query) do
       [?(, intersperse_map(elems, ?,, &expr(&1, sources, query)), ?)]
     end

lib/ecto/query.ex

@@ -354,7 +354,7 @@ defmodule Ecto.Query do
 
   defstruct [prefix: nil, sources: nil, from: nil, joins: [], aliases: %{}, wheres: [], select: nil,
              order_bys: [], limit: nil, offset: nil, group_bys: [], updates: [],
-             havings: [], preloads: [], assocs: [], distinct: nil, lock: nil]
+             havings: [], preloads: [], assocs: [], distinct: nil, lock: nil, windows: []]
 
   defmodule FromExpr do
     @moduledoc false
@@ -366,6 +366,11 @@ defmodule Ecto.Query do
     defstruct [:fun, :binding, :file, :line]
   end
 
+  defmodule PartitionByExpr do
+    @moduledoc false
+    defstruct [:fields, :file, :line, order_bys: []]
+  end
+
   defmodule QueryExpr do
     @moduledoc false
     defstruct [:expr, :file, :line, params: []]
@@ -398,8 +403,8 @@ defmodule Ecto.Query do
   @opaque dynamic :: %DynamicExpr{}
 
   alias Ecto.Query.Builder
-  alias Ecto.Query.Builder.{Distinct, Dynamic, Filter, From, GroupBy, Join,
-                            LimitOffset, Lock, OrderBy, Preload, Select, Update}
+  alias Ecto.Query.Builder.{Distinct, Dynamic, Filter, From, GroupBy, Join, PartitionBy,
+                            Window, LimitOffset, Lock, OrderBy, Preload, Select, Update}
 
   @doc """
   Builds a dynamic query expression.
@@ -464,6 +469,18 @@ defmodule Ecto.Query do
     Dynamic.build(binding, expr, __CALLER__)
   end
 
+  defmacro partition_by(binding \\ [], expr, opts \\ []) do
+    partition = PartitionBy.build(binding, expr, __CALLER__)
+    case Keyword.get(opts, :order_by) do
+      nil -> partition
+      fields -> OrderBy.build(partition, binding, fields, __CALLER__)
+    end
+  end
+
+  defmacro window(query, binding \\ [], expr) do
+    Window.build(query, binding, expr, __CALLER__)
+  end
+
   @doc """
   Converts a query into a subquery.
 
@@ -648,7 +665,7 @@ defmodule Ecto.Query do
 
   @from_join_opts [:as, :prefix, :hints]
   @no_binds [:lock]
-  @binds [:where, :or_where, :select, :distinct, :order_by, :group_by] ++
+  @binds [:where, :or_where, :select, :distinct, :order_by, :group_by, :window] ++
            [:having, :or_having, :limit, :offset, :preload, :update, :select_merge]
 
   defp from([{type, expr}|t], env, count_bind, quoted, binds) when type in @binds do

lib/ecto/query/api.ex

@@ -176,6 +176,11 @@ defmodule Ecto.Query.API do
   """
   def coalesce(value, expr), do: doc! [value, expr]
 
+  ## FIXME doc
+  @doc """
+  """
+  def over(value, window), do: doc! [value, window]
+
   @doc """
   Applies the given expression as a FILTER clause against an
   aggregate. This is currently only supported by Postgres.

lib/ecto/query/builder.ex

@@ -266,6 +266,12 @@ defmodule Ecto.Query.Builder do
     {expr, params_acc}
   end
 
+  def escape({:row_number, _, [arg]}, type, params_acc, vars, env) do
+    {arg, params_acc} = escape(arg, type, params_acc, vars, env)
+    expr = {:{}, [], [:row_number, [], [arg]]}
+    {expr, params_acc}
+  end
+
   def escape({:filter, _, [aggregate]}, type, params_acc, vars, env) do
     escape(aggregate, type, params_acc, vars, env)
   end
@@ -282,6 +288,32 @@ defmodule Ecto.Query.Builder do
     {{:{}, [], [:coalesce, [], [left, right]]}, params_acc}
   end
 
+  def escape({:over, _, [aggregate]}, type, params_acc, vars, env) do
+    {aggregate, params_acc} = escape(aggregate, type, params_acc, vars, env)
+    {{:{}, [], [:over, [], [aggregate, nil]]}, params_acc}
+  end
+
+  def escape({:over, _, [aggregate, {:^, _, [window]}]}, type, params_acc, vars, env) do
+    {aggregate, params_acc} = escape(aggregate, type, params_acc, vars, env)
+    {{:{}, [], [:over, [], [aggregate, window]]}, params_acc}
+  end
+
+  def escape({:over, _, [aggregate, {:partition_by, meta, fields}]}, type, params_acc, vars, env) do
+    {aggregate, params_acc} = escape(aggregate, type, params_acc, vars, env)
+    window = {:partition_by, meta, [{:vars, vars} | fields]}
+    {{:{}, [], [:over, [], [aggregate, window]]}, params_acc}
+  end
+
+  def escape({:over, _, [aggregate, {window_name, _, nil}]}, type, params_acc, vars, env) when is_atom(window_name) do
+    {aggregate, params_acc} = escape(aggregate, type, params_acc, vars, env)
+    {{:{}, [], [:over, [], [aggregate, window_name]]}, params_acc}
+  end
+
+  def escape({:over, _, [aggregate, window_name]}, type, params_acc, vars, env) when is_atom(window_name) do
+    {aggregate, params_acc} = escape(aggregate, type, params_acc, vars, env)
+    {{:{}, [], [:over, [], [aggregate, window_name]]}, params_acc}
+  end
+
   def escape({:=, _, _} = expr, _type, _params_acc, _vars, _env) do
     error! "`#{Macro.to_string(expr)}` is not a valid query expression. " <>
             "The match operator is not supported: `=`. " <>

lib/ecto/query/builder/order_by.ex

@@ -138,9 +138,18 @@ defmodule Ecto.Query.Builder.OrderBy do
                            params: unquote(params),
                            file: unquote(env.file),
                            line: unquote(env.line)}
+
     Builder.apply_query(query, __MODULE__, [order_by], env)
   end
 
+  @doc """
+  The callback applied by `build/4` to build the partition.
+  """
+  @spec apply(Ecto.Query.PartitionByExpr.t, term) :: Ecto.Query.PartitionByExpr.t
+  def apply(%Ecto.Query.PartitionByExpr{order_bys: order_bys} = partition, expr) do
+    %{partition | order_bys: order_bys ++ [expr]}
+  end
+
   @doc """
   The callback applied by `build/4` to build the query.
   """

lib/ecto/query/builder/partition_by.ex

@@ -0,0 +1,35 @@
+import Kernel, except: [apply: 2]
+
+defmodule Ecto.Query.Builder.PartitionBy do
+  @moduledoc false
+
+  alias Ecto.Query.Builder
+
+  @doc false
+  @spec build(Tuple.t, [Macro.t], Macro.Env.t) :: Macro.t
+  def build({:vars, vars}, fields, env) when is_list(fields) do
+    {fields, {_, :acc}} = Builder.escape(fields, :any, {%{}, :acc}, vars, env)
+
+    quote do
+      %Ecto.Query.PartitionByExpr{
+        fields: unquote(fields),
+        file: unquote(env.file),
+        line: unquote(env.line)}
+    end
+  end
+
+  @doc false
+  @spec build(Tuple.t, Macro.t, Macro.Env.t) :: Macro.t
+  def build({:vars, _} = vars, field, env) do
+    build(vars, [field], env)
+  end
+
+  @doc """
+  Builds a partition by expression.
+  """
+  @spec build([Macro.t], Macro.t, Macro.Env.t) :: Macro.t
+  def build(binding, expr, env) do
+    {_, vars} = Builder.escape_binding(quote(do: query), binding, env)
+    build({:vars, vars}, expr, env)
+  end
+end

lib/ecto/query/builder/window.ex

@@ -0,0 +1,63 @@
+import Kernel, except: [apply: 3]
+
+defmodule Ecto.Query.Builder.Window do
+  @moduledoc false
+
+  alias Ecto.Query.Builder
+
+  @doc """
+  """
+  @spec escape(Macro.t, Keyword.t, Macro.Env.t) :: {Atom.t, Macro.t}
+  def escape({name, _, [{:as, _, [expr]}]}, binding, env) when is_atom(name) do
+    escape(name, expr, binding, env)
+  end
+
+  @spec escape(Atom.t, Macro.t, Keyword.t, Macro.Env.t) :: {String.t, Macro.t}
+  def escape(name, {:^, _, [expr]}, _binding, _env) do
+    {name, expr}
+  end
+
+  def escape(name, {:partition_by, meta, [fields, opts]}, binding, _env) when is_list(opts) do
+    expr = {:partition_by, meta, [binding, fields, opts]}
+    {name, expr}
+  end
+
+  def escape(name, {:partition_by, meta, [fields]}, binding, _env) when is_list(fields) do
+    expr = {:partition_by, meta, [binding, fields]}
+    {name, expr}
+  end
+
+  def escape(name, {:partition_by, meta, [field]}, binding, _env) do
+    expr = {:partition_by, meta, [binding, [field]]}
+    {name, expr}
+  end
+
+  @doc """
+  Builds a quoted expression.
+
+  The quoted expression should evaluate to a query at runtime.
+  If possible, it does all calculations at compile time to avoid
+  runtime work.
+  """
+  @spec build(Macro.t, [Macro.t], Macro.t, Macro.Env.t) :: Macro.t
+  def build(query, binding, expr, env) do
+    {name, partition} = escape(expr, binding, env)
+
+    Builder.apply_query(query, __MODULE__, [partition, name], env)
+  end
+
+  @doc """
+  The callback applied by `build/4` to build the query.
+  """
+  @spec apply(Ecto.Queryable.t, Macro.t, term) :: Ecto.Query.t
+  def apply(%Ecto.Query{windows: windows} = query, expr, name) do
+    if Keyword.has_key?(windows, name) do
+      Builder.error! "window with name #{name} is already defined"
+    else
+      %{query | windows: windows ++ [{name, expr}]}
+    end
+  end
+  def apply(query, expr, name) do
+    apply(Ecto.Queryable.to_query(query), expr, name)
+  end
+end

lib/ecto/query/inspect.ex

@@ -1,7 +1,7 @@
 import Inspect.Algebra
 import Kernel, except: [to_string: 1]
 
-alias Ecto.Query.{DynamicExpr, JoinExpr, QueryExpr}
+alias Ecto.Query.{DynamicExpr, JoinExpr, QueryExpr, PartitionByExpr}
 
 defimpl Inspect, for: Ecto.Query.DynamicExpr do
   def inspect(%DynamicExpr{binding: binding} = dynamic, opts) do
@@ -52,6 +52,7 @@ defimpl Inspect, for: Ecto.Query do
     joins     = joins(query.joins, names)
     preloads  = preloads(query.preloads)
     assocs    = assocs(query.assocs, names)
+    windows   = windows(query.windows, names)
 
     wheres    = bool_exprs(%{and: :where, or: :or_where}, query.wheres, names)
     group_bys = kw_exprs(:group_by, query.group_bys, names)
@@ -65,7 +66,7 @@ defimpl Inspect, for: Ecto.Query do
     select    = kw_expr(:select, query.select, names)
     distinct  = kw_expr(:distinct, query.distinct, names)
 
-    Enum.concat [from, joins, wheres, group_bys, havings, order_bys,
+    Enum.concat [from, joins, wheres, group_bys, havings, windows, order_bys,
                  limit, offset, lock, distinct, updates, select, preloads, assocs]
   end
 
@@ -121,6 +122,30 @@ defimpl Inspect, for: Ecto.Query do
     end
   end
 
+  defp windows(windows, names) do
+    Enum.map(windows, &window(&1, names))
+  end
+
+  defp window({name, definition}, names) do
+    {:window, "(#{name} as " <> partition_by(definition, names) <> ")"}
+  end
+
+  defp window({name, %{ order_bys: order_bys } = definition}, names) do
+    {:window, string} = window({name, %{ definition | order_bys: [] }}, names)
+    [order_by: order_string] = kw_exprs(:order_by, order_bys, names)
+    {:window, string <> " order_by " <> order_string}
+  end
+
+  defp partition_by(%{ order_bys: order_bys } = definition, names) do
+    partition = "partition_by " <> expr(definition, names)
+    case order_bys do
+      [] -> partition
+      fields ->
+        [order_by: order_string] = kw_exprs(:order_by, fields, names)
+        partition <> ", order_by: " <> order_string
+    end
+  end
+
   defp bool_exprs(keys, exprs, names) do
     Enum.map exprs, fn %{expr: expr, op: op} = part ->
       {Map.fetch!(keys, op), expr(expr, names, part)}
@@ -141,6 +166,10 @@ defimpl Inspect, for: Ecto.Query do
     kw_inspect(:as, as) ++ kw_inspect(:prefix, prefix)
   end
 
+  defp expr(%PartitionByExpr{fields: fields} = part, names) do
+    Macro.to_string(fields, &expr_to_string(&1, &2, names, part))
+  end
+
   defp expr(%{expr: expr} = part, names) do
     expr(expr, names, part)
   end

lib/ecto/query/planner.ex

@@ -4,7 +4,7 @@ defmodule Ecto.Query.Planner do
 
   alias Ecto.Query.{BooleanExpr, DynamicExpr, JoinExpr, QueryExpr, SelectExpr}
 
-  if map_size(%Ecto.Query{}) != 18 do
+  if map_size(%Ecto.Query{}) != 19 do
     raise "Ecto.Query match out of date in builder"
   end
 
@@ -1022,6 +1022,11 @@ defmodule Ecto.Query.Planner do
     {type, [expr | fields], from}
   end
 
+  defp collect_fields({:over, _, [call, _]} = expr, fields, from, query, take) do
+    {type, _, _} = collect_fields(call, fields, from, query, take)
+    {type, [expr | fields], from}
+  end
+
   defp collect_fields({{:., _, [{:&, _, [ix]}, field]}, _, []} = expr,
                       fields, from, %{select: select} = query, _take) do
     type = source_type!(:select, query, select, ix, field)