refactor: optimize step output storage and aggregation (#569)

# Optimize Step Output Storage for Improved Performance

This PR changes how step outputs are stored and accessed in PGFlow to improve performance and reduce redundant data processing:

1. Store aggregated outputs directly in the `step_states` table instead of recalculating them from individual task outputs
2. Add a new `output` column to `step_states` with a constraint ensuring it's only populated for completed steps
3. Implement data backfill to populate the `output` field for existing completed steps:
   - For single steps: store the task output directly
   - For map steps: aggregate task outputs into an array

4. Update functions to use the pre-aggregated outputs:
   - `maybe_complete_run`: Use step_states.output directly instead of recalculating
   - `start_tasks`: Read dependency outputs from step_states instead of aggregating from tasks

This approach eliminates expensive output aggregation operations that were previously performed repeatedly, making the system more efficient.

Author: jumski

.changeset/step-output-storage.md

@@ -0,0 +1,7 @@
+---
+'@pgflow/core': minor
+---
+
+Performance: Store step outputs atomically for 2x faster downstream task startup
+
+Step outputs are now stored in step_states.output when steps complete, eliminating expensive aggregation queries. Benchmarks show 2.17x improvement for Map->Map chains. Includes data migration to backfill existing completed steps.

pkgs/core/schemas/0100_function_maybe_complete_run.sql

@@ -15,51 +15,22 @@ begin
   SET
     status = 'completed',
     completed_at = now(),
-    -- Only compute expensive aggregation when actually completing the run
+    -- Gather outputs from leaf steps (already stored in step_states.output by writers)
     output = (
-      -- ---------- Gather outputs from leaf steps ----------
       -- Leaf steps = steps with no dependents
-      -- For map steps: aggregate all task outputs into array
-      -- For single steps: use the single task output
       SELECT jsonb_object_agg(
-        step_slug,
-        CASE
-          WHEN step_type = 'map' THEN aggregated_output
-          ELSE single_output
-        END
+        leaf_state.step_slug,
+        leaf_state.output  -- Already aggregated by writers
       )
-      FROM (
-        SELECT DISTINCT
-          leaf_state.step_slug,
-          leaf_step.step_type,
-          -- For map steps: aggregate all task outputs
-          CASE WHEN leaf_step.step_type = 'map' THEN
-            (SELECT COALESCE(jsonb_agg(leaf_task.output ORDER BY leaf_task.task_index), '[]'::jsonb)
-             FROM pgflow.step_tasks leaf_task
-             WHERE leaf_task.run_id = leaf_state.run_id
-               AND leaf_task.step_slug = leaf_state.step_slug
-               AND leaf_task.status = 'completed')
-          END as aggregated_output,
-          -- For single steps: get the single output
-          CASE WHEN leaf_step.step_type = 'single' THEN
-            (SELECT leaf_task.output
-             FROM pgflow.step_tasks leaf_task
-             WHERE leaf_task.run_id = leaf_state.run_id
-               AND leaf_task.step_slug = leaf_state.step_slug
-               AND leaf_task.status = 'completed'
-             LIMIT 1)
-          END as single_output
-        FROM pgflow.step_states leaf_state
-        JOIN pgflow.steps leaf_step ON leaf_step.flow_slug = leaf_state.flow_slug AND leaf_step.step_slug = leaf_state.step_slug
-        WHERE leaf_state.run_id = maybe_complete_run.run_id
-          AND leaf_state.status = 'completed'
-          AND NOT EXISTS (
-            SELECT 1
-            FROM pgflow.deps dep
-            WHERE dep.flow_slug = leaf_state.flow_slug
-              AND dep.dep_slug = leaf_state.step_slug
-          )
-      ) leaf_outputs
+      FROM pgflow.step_states leaf_state
+      WHERE leaf_state.run_id = maybe_complete_run.run_id
+        AND leaf_state.status = 'completed'
+        AND NOT EXISTS (
+          SELECT 1
+          FROM pgflow.deps dep
+          WHERE dep.flow_slug = leaf_state.flow_slug
+            AND dep.dep_slug = leaf_state.step_slug
+        )
     )
   WHERE pgflow.runs.run_id = maybe_complete_run.run_id
     AND pgflow.runs.remaining_steps = 0

pkgs/core/schemas/0120_function_start_tasks.sql

@@ -47,28 +47,13 @@ as $$
       st.run_id,
       st.step_slug,
       dep.dep_slug,
-      -- Aggregate map outputs or use single output
-      CASE
-        WHEN dep_step.step_type = 'map' THEN
-          -- Aggregate all task outputs ordered by task_index
-          -- Use COALESCE to return empty array if no tasks
-          (SELECT COALESCE(jsonb_agg(dt.output ORDER BY dt.task_index), '[]'::jsonb)
-           FROM pgflow.step_tasks dt
-           WHERE dt.run_id = st.run_id
-             AND dt.step_slug = dep.dep_slug
-             AND dt.status = 'completed')
-        ELSE
-          -- Single step: use the single task output
-          dep_task.output
-      END as dep_output
+      -- Read output directly from step_states (already aggregated by writers)
+      dep_state.output as dep_output
     from tasks st
     join pgflow.deps dep on dep.flow_slug = st.flow_slug and dep.step_slug = st.step_slug
-    join pgflow.steps dep_step on dep_step.flow_slug = dep.flow_slug and dep_step.step_slug = dep.dep_slug
-    left join pgflow.step_tasks dep_task on
-      dep_task.run_id = st.run_id and
-      dep_task.step_slug = dep.dep_slug and
-      dep_task.status = 'completed'
-      and dep_step.step_type = 'single'  -- Only join for single steps
+    join pgflow.step_states dep_state on
+      dep_state.run_id = st.run_id and
+      dep_state.step_slug = dep.dep_slug
   ),
   deps_outputs as (
     select

pkgs/core/scripts/benchmarks/README.md

@@ -0,0 +1,129 @@
+# step_output_storage Benchmark
+
+Measures performance improvements from storing step outputs in `step_states.output` instead of aggregating from `step_tasks` on every read.
+
+## Quick Start
+
+```bash
+# From pkgs/core directory
+pnpm nx supabase:reset core
+pnpm nx supabase:status core  # Note the port
+
+PGPASSWORD=postgres psql -h 127.0.0.1 -p PORT -U postgres -d postgres \
+  -f scripts/benchmarks/step_output_storage.sql
+```
+
+## What Changed
+
+### OLD Code (main branch)
+
+In `start_tasks`, when a task needs its dependency outputs:
+
+```sql
+-- For each dependent task, aggregate all completed task outputs
+CASE WHEN dep_step.step_type = 'map' THEN
+  (SELECT jsonb_agg(dt.output ORDER BY dt.task_index)
+   FROM pgflow.step_tasks dt
+   WHERE dt.run_id = ... AND dt.status = 'completed')
+ELSE ...
+```
+
+**Problem**: If a map step has 500 completed tasks, and 500 downstream tasks need to start, each downstream task triggers an aggregation query over 500 rows = **250,000 row scans**.
+
+### NEW Code (step_output_storage branch)
+
+Outputs are stored in `step_states.output` when a step completes:
+
+```sql
+-- Just read the pre-stored output
+dep_state.output as dep_output
+```
+
+**Improvement**: 500 downstream tasks each read 1 column = **500 column reads**.
+
+## Benchmark Tests
+
+### Test 1: `complete_task_final`
+
+**What it measures**: Time to complete the last task of a map step.
+
+**Setup**: Start a flow with N-element array, complete N-1 tasks, then time the final `complete_task()`.
+
+**Why it matters**: The final task triggers output storage. In NEW code, this aggregates once and stores. In OLD code, this just completes the task (aggregation happens on read).
+
+### Test 2: `start_tasks_read_N`
+
+**What it measures**: Time to start a single downstream task that reads N-element dependency output.
+
+**Setup**:
+1. Map step `producer` with N tasks - all completed
+2. Single step `consumer` depends on `producer`
+3. Time `start_tasks()` for the consumer task
+
+**Flow structure**:
+```
+[producer: map(N)] --> [consumer: single]
+```
+
+**Expected difference**:
+- OLD: Aggregates N task outputs on every read
+- NEW: Reads from `step_states.output` (O(1))
+
+### Test 3: `start_tasks_batch_NxN`
+
+**What it measures**: Time to start N downstream tasks, each reading N-element dependency output.
+
+**Setup**:
+1. Map step `producer` with N tasks - all completed
+2. Map step `consumer` depends on `producer`, has N tasks
+3. Time `start_tasks()` for all N consumer tasks at once
+
+**Flow structure**:
+```
+[producer: map(N)] --> [consumer: map(N)]
+```
+
+**Expected difference**:
+- OLD: N tasks * aggregate(N outputs) = O(N^2) row scans
+- NEW: N tasks * read(1 column) = O(N) column reads
+
+**This is the key benchmark** - should show the largest improvement.
+
+## Configuration
+
+Edit line 19 to change array size:
+
+```sql
+\set ARRAY_SIZE 500  -- Default: ~3 min runtime
+\set ARRAY_SIZE 100  -- Quick test: ~30 sec
+\set ARRAY_SIZE 1000 -- Thorough: ~15+ min
+```
+
+## Comparing Branches
+
+```bash
+# 1. Run on step_output_storage branch
+git checkout step_output_storage
+pnpm nx supabase:reset core
+psql ... -f scripts/benchmarks/step_output_storage.sql | tee results_new.txt
+
+# 2. Run on main branch
+git checkout main
+pnpm nx supabase:reset core
+psql ... -f scripts/benchmarks/step_output_storage.sql | tee results_old.txt
+
+# 3. Compare start_tasks_batch_NxN results
+grep "start_tasks_batch" results_*.txt
+```
+
+## Expected Results
+
+| Test | OLD (main) | NEW (step_output_storage) | Improvement |
+|------|------------|---------------------------|-------------|
+| `complete_task_final` | ~same | ~same | minimal |
+| `start_tasks_read_N` | slower | faster | 30-50% |
+| `start_tasks_batch_NxN` | much slower | faster | 50-80% |
+
+The batch test improvement grows with N because:
+- OLD scales as O(N^2)
+- NEW scales as O(N)