Postgresql – Function in JOIN on DISTINCT values executes for each row instead of distinct parameters

I'm JOINing function execution using some of the values from table, expecting that it will execute only as many times as there are unique sets of parameters, but instead it executes for each row.

Function is set to STABLE, but I also attempted IMMUTABLE with no effect on execution time. I also did try to put it into SELECT clause, but as expected it executes for each row when done so.

Is there a different way to make it execute only as many times as there are distinct sets of parameters without querying table (or CTE) for the second time?

Schema

CREATE SCHEMA IF NOT EXISTS lkaminski;

CREATE TABLE lkaminski.table_1 (id serial, flag bool, fkey int4);

INSERT INTO lkaminski.table_1 (flag, fkey)
SELECT random() > 0.5, (random()*10)::int2+1
  FROM generate_series(1,500);

SELECT COUNT(DISTINCT CONCAT(flag, ':', fkey))
  FROM lkaminski.table_1;

CREATE OR REPLACE FUNCTION lkaminski.do_stuff(flag bool, fkey int4, OUT computed_value int4)
LANGUAGE plpgsql
STABLE
AS $body$
BEGIN
PERFORM pg_sleep(0.01);
SELECT COUNT(0)
  INTO computed_value
  /* Simplified, takes data from table, not immutable */
  ;
END;
$body$;

First attempt

EXPLAIN (ANALYZE, VERBOSE)
SELECT *, cnt
  FROM lkaminski.table_1
 CROSS JOIN lkaminski.do_stuff(flag, fkey) AS cnt;

QUERY PLAN

Nested Loop  (cost=0.25..76.25 rows=2200 width=17) (actual time=10.265..5084.277 rows=500 loops=1)
  Output: table_1.id, table_1.flag, table_1.fkey, cnt.computed_value, cnt.computed_value
  ->  Seq Scan on lkaminski.table_1  (cost=0.00..32.00 rows=2200 width=9) (actual time=0.047..1.027 rows=500 loops=1)
        Output: table_1.id, table_1.flag, table_1.fkey
  ->  Function Scan on lkaminski.do_stuff cnt  (cost=0.25..0.26 rows=1 width=4) (actual time=10.155..10.158 rows=1 loops=500)
        Output: cnt.computed_value
        Function Call: lkaminski.do_stuff(table_1.flag, table_1.fkey)
Planning time: 0.973 ms
Execution time: 5085.000 ms

Second attempt

EXPLAIN (ANALYZE, VERBOSE)
SELECT t1.*, cnt
  FROM lkaminski.table_1 AS t1
 CROSS JOIN LATERAL(SELECT fkey, lkaminski.do_stuff(s1.flag, s1.fkey) AS cnt
                      FROM (SELECT DISTINCT t1.flag, t1.fkey) AS s1
                   ) AS a;

QUERY PLAN

Nested Loop  (cost=0.00..648.00 rows=2200 width=13) (actual time=10.230..5072.033 rows=500 loops=1)
  Output: t1.id, t1.flag, t1.fkey, lkaminski.do_stuff((t1.flag), (t1.fkey))
  ->  Seq Scan on lkaminski.table_1 t1  (cost=0.00..32.00 rows=2200 width=9) (actual time=0.011..0.797 rows=500 loops=1)
        Output: t1.id, t1.flag, t1.fkey
  ->  Unique  (cost=0.00..0.01 rows=1 width=5) (actual time=0.004..0.010 rows=1 loops=500)
        Output: (t1.flag), (t1.fkey)
        ->  Result  (cost=0.00..0.01 rows=1 width=5) (actual time=0.001..0.003 rows=1 loops=500)
              Output: t1.flag, t1.fkey
Planning time: 0.184 ms
Execution time: 5072.785 ms

Third attempt, works, but makes second query on the same table/CTE, then JOIN and looks ugly

EXPLAIN (ANALYZE, VERBOSE)
WITH t1 AS (
  SELECT id, flag, fkey FROM lkaminski.table_1 --WHERE expensive conditions here
), c AS (
  SELECT fkey, lkaminski.do_stuff(flag, fkey) AS cnt
    FROM (SELECT DISTINCT flag, fkey
            FROM t1
         ) AS a
)
SELECT id, flag, fkey, cnt
  FROM t1
  JOIN c USING(fkey);

QUERY PLAN

Hash Join  (cost=208.00..411.50 rows=4400 width=13) (actual time=225.632..228.265 rows=1000 loops=1)
  Output: t1.id, t1.flag, t1.fkey, c.cnt
  Hash Cond: (t1.fkey = c.fkey)
  CTE t1
    ->  Seq Scan on lkaminski.table_1  (cost=0.00..32.00 rows=2200 width=9) (actual time=0.013..0.666 rows=500 loops=1)
          Output: table_1.id, table_1.flag, table_1.fkey
  CTE c
    ->  Subquery Scan on a  (cost=55.00..163.00 rows=400 width=8) (actual time=12.947..225.478 rows=22 loops=1)
          Output: a.fkey, lkaminski.do_stuff(a.flag, a.fkey)
          ->  HashAggregate  (cost=55.00..59.00 rows=400 width=5) (actual time=2.724..2.756 rows=22 loops=1)
                Output: t1_1.flag, t1_1.fkey
                Group Key: t1_1.flag, t1_1.fkey
                ->  CTE Scan on t1 t1_1  (cost=0.00..44.00 rows=2200 width=5) (actual time=0.002..2.025 rows=500 loops=1)
                      Output: t1_1.id, t1_1.flag, t1_1.fkey
  ->  CTE Scan on t1  (cost=0.00..44.00 rows=2200 width=9) (actual time=0.018..0.686 rows=500 loops=1)
        Output: t1.id, t1.flag, t1.fkey
  ->  Hash  (cost=8.00..8.00 rows=400 width=8) (actual time=225.600..225.600 rows=22 loops=1)
        Output: c.cnt, c.fkey
        Buckets: 1024  Batches: 1  Memory Usage: 9kB
        ->  CTE Scan on c  (cost=0.00..8.00 rows=400 width=8) (actual time=12.951..225.551 rows=22 loops=1)
              Output: c.cnt, c.fkey
Planning time: 0.291 ms
Execution time: 230.920 ms