feat: distributed correctness — pipeline splitter, aggregate fixes

Major distributed correctness improvements:

PipelineSplitter: classifies ops into worker-safe vs coordinator-only.
- filter, mutate, select, discard, rename, drop_nil → push to workers
- sort_by, head, distinct → push to workers AND re-apply on coordinator
- slice, pivot_wider, pivot_longer → coordinator-only (positional/schema)
- summarise with AVG → rewritten to SUM+COUNT on workers, divided on coordinator

Merger aggregate fixes:
- MIN() re-aggregated with MIN() (was SUM — WRONG)
- MAX() re-aggregated with MAX() (was SUM — WRONG)
- SUM/COUNT still use SUM (correct)
- AVG handled via pipeline rewrite (SUM/COUNT split)

Coordinator now:
- Splits pipeline via PipelineSplitter before fan-out
- Applies AVG rewrites after merge (SUM/COUNT → division)
- Applies coordinator-only ops (slice, pivot) post-merge

14 correctness tests: MIN/MAX re-aggregation, grouped MIN/MAX,
AVG rewrite (global + grouped), mixed aggregates (SUM+COUNT+AVG+MIN+MAX),
slice on coordinator, pivot on coordinator, pipeline splitter unit tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: cigrainger

lib/dux/remote/coordinator.ex

@@ -19,7 +19,7 @@ defmodule Dux.Remote.Coordinator do
   The result is a `%Dux{}` struct with the merged data.
   """
 
-  alias Dux.Remote.{Merger, Partitioner, Worker}
+  alias Dux.Remote.{Merger, Partitioner, PipelineSplitter, Worker}
 
   @doc """
   Execute a `%Dux{}` pipeline across distributed workers.
@@ -44,8 +44,14 @@ defmodule Dux.Remote.Coordinator do
       raise ArgumentError, "no workers available for distributed execution"
     end
 
-    # Partition the pipeline across workers
-    assignments = Partitioner.assign(pipeline, workers, strategy: strategy)
+    # Split pipeline: worker ops push down, coordinator ops apply post-merge
+    %{worker_ops: worker_ops, coordinator_ops: coord_ops, agg_rewrites: rewrites} =
+      PipelineSplitter.split(pipeline.ops)
+
+    worker_pipeline = %{pipeline | ops: worker_ops}
+
+    # Partition the worker pipeline across workers
+    assignments = Partitioner.assign(worker_pipeline, workers, strategy: strategy)
 
     # Fan out: each worker executes its partition
     results = fan_out(assignments, timeout)
@@ -59,7 +65,13 @@ defmodule Dux.Remote.Coordinator do
     end
 
     # Merge partial results on coordinator
-    Merger.merge_to_dux(successes, pipeline)
+    merged = Merger.merge_to_dux(successes, worker_pipeline)
+
+    # Apply AVG rewrites if any
+    merged = apply_avg_rewrites(merged, rewrites)
+
+    # Apply coordinator-only ops (slice, pivot, etc.)
+    apply_coordinator_ops(merged, coord_ops)
   end
 
   @doc """
@@ -106,4 +118,53 @@ defmodule Dux.Remote.Coordinator do
       {Enum.map(ok, fn {:ok, ipc} -> ipc end), err}
     end)
   end
+
+  # Apply AVG rewrites: compute __sum / __count for each rewritten AVG column
+  defp apply_avg_rewrites(dux, rewrites) when map_size(rewrites) == 0, do: dux
+
+  defp apply_avg_rewrites(dux, rewrites) do
+    # For each AVG rewrite, add a mutate computing sum/count,
+    # then discard the intermediate columns
+    avg_exprs =
+      Enum.map(rewrites, fn {name, {:avg, sum_col, count_col, _inner}} ->
+        {name, "\"#{esc(sum_col)}\" / \"#{esc(count_col)}\""}
+      end)
+
+    intermediate_cols =
+      Enum.flat_map(rewrites, fn {_name, {:avg, sum_col, count_col, _}} ->
+        [String.to_atom(sum_col), String.to_atom(count_col)]
+      end)
+
+    dux
+    |> Dux.mutate_with(avg_exprs)
+    |> Dux.discard(intermediate_cols)
+  end
+
+  # Apply coordinator-only ops to the merged result
+  defp apply_coordinator_ops(dux, []), do: dux
+
+  defp apply_coordinator_ops(dux, [op | rest]) do
+    updated = apply_single_op(dux, op)
+    apply_coordinator_ops(updated, rest)
+  end
+
+  defp apply_single_op(dux, {:slice, offset, length}), do: Dux.slice(dux, offset, length)
+  defp apply_single_op(dux, {:head, n}), do: Dux.head(dux, n)
+  defp apply_single_op(dux, {:sort_by, spec}), do: %{dux | ops: dux.ops ++ [{:sort_by, spec}]}
+  defp apply_single_op(dux, {:distinct, cols}), do: %{dux | ops: dux.ops ++ [{:distinct, cols}]}
+
+  defp apply_single_op(dux, {:pivot_wider, names_col, values_col, agg}) do
+    Dux.pivot_wider(dux, names_col, values_col, agg: agg)
+  end
+
+  defp apply_single_op(dux, {:pivot_longer, cols, names_to, values_to}) do
+    Dux.pivot_longer(dux, cols, names_to: names_to, values_to: values_to)
+  end
+
+  defp apply_single_op(dux, op) do
+    # Unknown op — append directly
+    %{dux | ops: dux.ops ++ [op]}
+  end
+
+  defp esc(name), do: String.replace(name, ~s("), ~s(""))
 end

lib/dux/remote/merger.ex

@@ -121,13 +121,12 @@ defmodule Dux.Remote.Merger do
         "SELECT DISTINCT * FROM (#{union_sql}) __merged"
 
       {:re_aggregate, groups, aggs} ->
-        # Re-aggregate partial results
+        # Re-aggregate partial results with correct functions
         group_cols = Enum.map_join(groups, ", ", &qi/1)
 
         agg_cols =
-          Enum.map_join(aggs, ", ", fn {name, _expr} ->
-            quoted = qi(name)
-            "SUM(#{quoted}) AS #{quoted}"
+          Enum.map_join(aggs, ", ", fn {name, expr} ->
+            re_aggregate_expr(name, expr)
           end)
 
         select = if groups == [], do: agg_cols, else: "#{group_cols}, #{agg_cols}"
@@ -179,6 +178,28 @@ defmodule Dux.Remote.Merger do
     end)
   end
 
+  # Determine the correct re-aggregation function based on the original expression.
+  # SUM → SUM, COUNT → SUM, MIN → MIN, MAX → MAX
+  # AVG columns should have been rewritten by PipelineSplitter before reaching here.
+  defp re_aggregate_expr(name, expr) when is_binary(expr) do
+    upper = String.upcase(expr)
+    quoted = qi(name)
+
+    cond do
+      String.contains?(upper, "MIN(") -> "MIN(#{quoted}) AS #{quoted}"
+      String.contains?(upper, "MAX(") -> "MAX(#{quoted}) AS #{quoted}"
+      String.contains?(upper, "SUM(") -> "SUM(#{quoted}) AS #{quoted}"
+      String.contains?(upper, "COUNT(") -> "SUM(#{quoted}) AS #{quoted}"
+      # Default: SUM (safe for additive aggregates)
+      true -> "SUM(#{quoted}) AS #{quoted}"
+    end
+  end
+
+  defp re_aggregate_expr(name, _expr) do
+    quoted = qi(name)
+    "SUM(#{quoted}) AS #{quoted}"
+  end
+
   # Quote identifier — escape double quotes to prevent SQL injection
   defp qi(name) do
     escaped = String.replace(name, ~s("), ~s(""))

lib/dux/remote/pipeline_splitter.ex

@@ -0,0 +1,168 @@
+defmodule Dux.Remote.PipelineSplitter do
+  @moduledoc false
+
+  # Splits a pipeline's ops into worker-safe ops (pushed to workers)
+  # and coordinator ops (applied after merge).
+  #
+  # Safe to push: filter, mutate, select, discard, rename, drop_nil,
+  #   group_by, summarise (with rewrite), sort_by, head, distinct
+  #
+  # Coordinator-only: slice, pivot_wider, pivot_longer (after merge),
+  #   join (requires special handling)
+  #
+  # The splitter also rewrites AVG aggregations: workers compute
+  # SUM + COUNT, coordinator divides to get the correct average.
+
+  @doc """
+  Split pipeline ops into worker and coordinator portions.
+
+  Returns `%{worker_ops: [...], coordinator_ops: [...], agg_rewrites: %{}}`.
+
+  `agg_rewrites` maps output column names to their rewrite info
+  (e.g., AVG rewritten to SUM/COUNT pair).
+  """
+  def split(ops) do
+    {worker_ops, coordinator_ops, rewrites} = do_split(ops, [], [], %{})
+
+    %{
+      worker_ops: Enum.reverse(worker_ops),
+      coordinator_ops: Enum.reverse(coordinator_ops),
+      agg_rewrites: rewrites
+    }
+  end
+
+  # Walk ops left-to-right. Push safe ops to workers, keep others for coordinator.
+
+  # Filter, mutate, select, discard, rename, drop_nil — always safe
+  defp do_split([{:filter, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:mutate, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:select, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:discard, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:rename, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:drop_nil, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  # Group by — push to workers (sets state for summarise)
+  defp do_split([{:group_by, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  defp do_split([{:ungroup} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  # Summarise — push to workers, but rewrite AVG and track for re-aggregation
+  defp do_split([{:summarise, aggs} | rest], worker, coord, rewrites) do
+    {worker_aggs, new_rewrites} = rewrite_aggregates(aggs)
+    do_split(rest, [{:summarise, worker_aggs} | worker], coord, Map.merge(rewrites, new_rewrites))
+  end
+
+  # Sort, head, distinct — push to workers AND add to coordinator for re-merge
+  defp do_split([{:sort_by, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], [op | coord], rewrites)
+  end
+
+  defp do_split([{:head, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], [op | coord], rewrites)
+  end
+
+  defp do_split([{:distinct, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], [op | coord], rewrites)
+  end
+
+  # Slice — coordinator only (OFFSET is positional, not safe per-partition)
+  defp do_split([{:slice, _, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, worker, [op | coord], rewrites)
+  end
+
+  # Pivot — coordinator only (schema may differ across partitions)
+  defp do_split([{:pivot_wider, _, _, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, worker, [op | coord], rewrites)
+  end
+
+  defp do_split([{:pivot_longer, _, _, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, worker, [op | coord], rewrites)
+  end
+
+  # Join — push to workers (right side will be broadcast/replicated by the partitioner)
+  defp do_split([{:join, _, _, _, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  # Concat — push to workers
+  defp do_split([{:concat_rows, _} = op | rest], worker, coord, rewrites) do
+    do_split(rest, [op | worker], coord, rewrites)
+  end
+
+  # Unknown op — keep on coordinator for safety
+  defp do_split([op | rest], worker, coord, rewrites) do
+    do_split(rest, worker, [op | coord], rewrites)
+  end
+
+  defp do_split([], worker, coord, rewrites) do
+    {worker, coord, rewrites}
+  end
+
+  # ---------------------------------------------------------------------------
+  # Aggregate rewrites
+  # ---------------------------------------------------------------------------
+
+  # Rewrite aggregates that can't be naively re-aggregated with SUM.
+  # AVG(x) → workers compute __sum_name = SUM(x), __count_name = COUNT(x)
+  #           coordinator computes name = __sum_name / __count_name
+  defp rewrite_aggregates(aggs) do
+    {worker_aggs, rewrites} =
+      Enum.reduce(aggs, {[], %{}}, fn {name, expr}, {acc_aggs, acc_rewrites} ->
+        if avg_expr?(expr) do
+          inner = extract_inner_expr(expr, "AVG")
+          sum_name = "__avg_sum_#{name}"
+          count_name = "__avg_count_#{name}"
+
+          new_aggs = [
+            {sum_name, rewrite_func(expr, "AVG", "SUM")},
+            {count_name, rewrite_func(expr, "AVG", "COUNT")} | acc_aggs
+          ]
+
+          new_rewrites = Map.put(acc_rewrites, name, {:avg, sum_name, count_name, inner})
+          {new_aggs, new_rewrites}
+        else
+          {[{name, expr} | acc_aggs], acc_rewrites}
+        end
+      end)
+
+    {Enum.reverse(worker_aggs), rewrites}
+  end
+
+  defp avg_expr?(expr) when is_binary(expr), do: String.contains?(String.upcase(expr), "AVG(")
+  defp avg_expr?(_), do: false
+
+  defp extract_inner_expr(expr, func) do
+    # Extract the inner expression from FUNC(inner)
+    regex = ~r/#{func}\((.+)\)/i
+
+    case Regex.run(regex, expr) do
+      [_, inner] -> inner
+      _ -> expr
+    end
+  end
+
+  defp rewrite_func(expr, from, to) do
+    String.replace(expr, ~r/#{from}\(/i, "#{to}(")
+  end
+end