Query stops parallelizing when adding LIMIT or ROW_NUMBER

Question

I have 20 tables with identical structure, and I need to aggregate the results.

I'm using UNION ALL to collect the results, in the following structure:

SELECT something, SUM(total) AS overall_total
FROM (
  SELECT something, SUM(something_else) AS total
  FROM table1
  GROUP BY something

  UNION ALL
  SELECT something, SUM(something_else) AS total
  FROM table2
  GROUP BY something

  UNION ALL
  ...

  ) AS X
GROUP BY something
ORDER BY overall_total DESC
LIMIT 10

In this case the DB uses Parallel Append with multiple workers, which is very fast.

However, each table returns a large amount of data so aggregation is slow.

To improve performance, I want to limit the number of rows returned from each table to the top 100 results only, since the overall will always rely on those.

I tried using either LIMIT or ROW_NUMBER() OVER() to retrieve only the top rows from each table:

SELECT something, SUM(something_else) AS total
FROM table2
GROUP BY something
ORDER BY total DESC
LIMIT 100

SELECT something, SUM(something_else) AS total, ROW_NUMBER() OVER() as r
FROM table2
WHERE r <= 100
GROUP BY something
ORDER BY total DESC

But now the plan never uses Parallel Append, so data fetch is very slow and overall performance is worse. Can't get the parallel plan even using flags like:

set force_parallel_mode = on;
set max_parallel_workers_per_gather = 10;
set parallel_setup_cost = 0;
set parallel_tuple_cost = 0;

Reading PostgreSQL documentation, parallelization can't be done when using unsafe operations. But it didn't mention LIMIT and ROW_NUMBER as being unsafe.

Answer 1

Why?

I can reproduce the issue. As soon as I add a LIMIT clause to one of the SELECT terms of the UNION ALL, the parallel plan goes away, even with set force_parallel_mode = on;.

Maybe it's connected to this comment in the source code:

      * We can't push sub-select containing LIMIT/OFFSET to workers as
      * there is no guarantee that the row order will be fully
      * deterministic, and applying LIMIT/OFFSET will lead to
      * inconsistent results at the top-level.  (In some cases, where
      * the result is ordered, we could relax this restriction.  But it
      * doesn't currently seem worth expending extra effort to do so.)

But I am not entirely sure. This applies to "subquery in FROM". But it would seem it shouldn't keep Postgres from scheduling parallel workers for whole subqueries in UNION ALL in a Parallel Append node. Maybe just a shortcoming in current Postgres 15?

Test case

SELECT something, sum(total) AS overall_total
FROM  (
   (  -- !
   SELECT something, sum(something_else) AS total
   FROM   table1
   GROUP  BY 1
   ORDER  BY 2 DESC NULLS LAST
   LIMIT  100
   )  -- !

   UNION ALL
   (
   SELECT something, SUM(something_else)
   FROM   table2
   GROUP  BY 1
   ORDER  BY 2 DESC NULLS LAST, 1
   LIMIT  100
   )

   --  ... more?
   ) sub
GROUP  BY 1
ORDER  BY 2 DESC NULLS LAST, 1
LIMIT  10;

See:

• Combine multiple SELECT statements

You may or may not need / want NULLS LAST:

• Sort NULL values to the end of a table

I made the sort order deterministic by adding something as second ORDER BY item. Else, repeated calls might get you different results.

Workaround

We can work around the limitation by encapsulating each root query in a function with a PARALLEL SAFE label:

CREATE FUNCTION public.f_top_sums(_tbl regclass)
  RETURNS TABLE (
     something int  -- adjust type to your case!
   , total bigint   -- adjust type to your case!
  )
  LANGUAGE plpgsql PARALLEL SAFE ROWS 100 COST 1000000 AS
$func$
BEGIN
   RETURN QUERY EXECUTE format(
      '
      SELECT something, sum(something_else) AS total
      FROM   %s
      GROUP  BY 1
      ORDER  BY 2 DESC NULLS LAST, 1
      LIMIT  100
      '
    , _tbl);
END   
$func$;

I chose a generic function with dynamic SQL and EXECUTE, passing table names safely. If you are unfamiliar, here is more about the basics:

• Table name as a PostgreSQL function parameter

I also declared ROWS 100 (because we know that) and, more importantly, COST 1000000 to make the query planner consider a parallel plan. You may want to fiddle with the COST setting to influence the query plan ...

You can also write one simple SQL function for each table instead.

Either way, this now allows parallel plans:

SELECT something, SUM(total) AS overall_total
FROM  (
   SELECT * FROM public.f_top_sums('public.table1')
   UNION ALL
   SELECT * FROM public.f_top_sums('public.table2')
   --  more ...
   ) sub
GROUP  BY 1
ORDER  BY 2 DESC NULLS LAST
LIMIT  10;

---

Aside, your second attempt with the window function row_number() is illegal syntax. You cannot reference the output column r in the WHERE clause. You would need a subquery to make it work. But a single subquery level would suffice, as you can wrap the window function around the aggregate function:

SELECT *
FROM (
   SELECT something, sum(something_else) AS total
        , row_number() OVER (ORDER BY sum(something_else) DESC NULLS LAST) AS r
   FROM   table2
   GROUP  BY 1
   ) sub
WHERE  r <= 100

That said, the first approach with LIMIT version is simpler and faster.