Select random row from a PostgreSQL table with weighted row probabilities

Question

Example input:

SELECT * FROM test;
 id | percent   
----+----------
  1 | 50 
  2 | 35   
  3 | 15   
(3 rows)

How would you write such query, that

Answer 1

Branko's accepted solution is great (thanks!). However, I'd like to contribute an alternative that is just as performant (according to my tests), and perhaps easier to visualize.

Let's recap. The original question can perhaps be generalized as follows:

Given an map of ids and relative weights, create a query that returns a random id in the map, but with a probability proportional to its relative weight.

Note the emphasis on relative weights, not percent. As Branko points out in his answer, using relative weights will work for anything, including percents.

Now, consider some test data, which we'll put in a temporary table:

CREATE TEMP TABLE test AS
SELECT * FROM (VALUES
    (1, 25),
    (2, 10),
    (3, 10),
    (4, 05)
) AS test(id, weight);

Note that I'm using a more complicated example than that in the original question, in that it does not conveniently add up to 100, and in that the same weight (20) is used more than once (for ids 2 and 3), which is important to consider, as you'll see later.

The first thing we have to do is turn the weights into probabilities from 0 to 1, which is nothing more than a simple normalization (weight / sum(weights)):

WITH p AS ( -- probability
    SELECT *,
        weight::NUMERIC / sum(weight) OVER () AS probability
    FROM test
),
cp AS ( -- cumulative probability
    SELECT *,
        sum(p.probability) OVER (
            ORDER BY probability DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS cumprobability
    FROM p
)
SELECT
    cp.id,
    cp.weight,
    cp.probability,
    cp.cumprobability - cp.probability AS startprobability,
    cp.cumprobability AS endprobability
FROM cp
;

This will result in the following output:

 id | weight | probability | startprobability | endprobability
----+--------+-------------+------------------+----------------
  1 |     25 |         0.5 |              0.0 |            0.5
  2 |     10 |         0.2 |              0.5 |            0.7
  3 |     10 |         0.2 |              0.7 |            0.9
  4 |      5 |         0.1 |              0.9 |            1.0

The query above is admittedly doing more work than strictly necessary for our needs, but I find it helpful to visualize the relative probabilities this way, and it does make the final step of choosing the id trivial:

SELECT id FROM (queryabove)
WHERE random() BETWEEN startprobability AND endprobability;

Now, let's put it all together with a test that ensures the query is returning data with the expected distribution. We'll use generate_series() to generate a random number a million times:

WITH p AS ( -- probability
    SELECT *,
        weight::NUMERIC / sum(weight) OVER () AS probability
    FROM test
),
cp AS ( -- cumulative probability
    SELECT *,
        sum(p.probability) OVER (
            ORDER BY probability DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS cumprobability
    FROM p
),
fp AS ( -- final probability
    SELECT
        cp.id,
        cp.weight,
        cp.probability,
        cp.cumprobability - cp.probability AS startprobability,
        cp.cumprobability AS endprobability
    FROM cp
)
SELECT *
FROM fp
CROSS JOIN (SELECT random() FROM generate_series(1, 1000000)) AS random(val)
WHERE random.val BETWEEN fp.startprobability AND fp.endprobability
;

This will result in output similar to the following:

 id | count  
----+--------
 1  | 499679 
 3  | 200652 
 2  | 199334 
 4  | 100335 

Which, as you can see, tracks the expected distribution perfectly.

Performance

The query above is quite performant. Even in my average machine, with PostgreSQL running in a WSL1 instance (the horror!), execution is relatively fast:

     count | time (ms)
-----------+----------
     1,000 |         7
    10,000 |        25
   100,000 |       210
 1,000,000 |      1950 

Adaptation to generate test data

I often use a variation of the query above when generating test data for unit/integration tests. The idea is to generate random data that approximates a probability distribution that tracks reality.

In that situation I find it useful to compute the start and end distributions once and storing the results in a table:

CREATE TEMP TABLE test AS
WITH test(id, weight) AS (VALUES
    (1, 25),
    (2, 10),
    (3, 10),
    (4, 05)
),
p AS ( -- probability
    SELECT *, (weight::NUMERIC / sum(weight) OVER ()) AS probability
    FROM test
),
cp AS ( -- cumulative probability
    SELECT *,
        sum(p.probability) OVER (
            ORDER BY probability DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) cumprobability
    FROM p
)
SELECT
    cp.id,
    cp.weight,
    cp.probability,
    cp.cumprobability - cp.probability AS startprobability,
    cp.cumprobability AS endprobability
FROM cp
;

I can then use these precomputed probabilities repeatedly, which results in extra performance and simpler use.

I can even wrap it all in a function that I can call any time I want to get a random id:

CREATE OR REPLACE FUNCTION getrandomid(p_random FLOAT8 = random())
RETURNS INT AS
$$
    SELECT id
    FROM test
    WHERE p_random BETWEEN startprobability AND endprobability
    ;
$$
LANGUAGE SQL STABLE STRICT

Window function frames

It's worth noting that the technique above is using a window function with a non-standard frame ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW. This is necessary to deal with the fact that some weights might be repeated, which is why I chose test data with repeated weights in the first place!

Answer 2

This should do the trick:

WITH CTE AS (
    SELECT random() * (SELECT SUM(percent) FROM YOUR_TABLE) R
)
SELECT *
FROM (
    SELECT id, SUM(percent) OVER (ORDER BY id) S, R
    FROM YOUR_TABLE CROSS JOIN CTE
) Q
WHERE S >= R
ORDER BY id
LIMIT 1;

The sub-query Q gives the following result:

1  50
2  85
3  100

We then simply generate a random number in range [0, 100) and pick the first row that is at or beyond that number (the WHERE clause). We use common table expression (WITH) to ensure the random number is calculated only once.

BTW, the SELECT SUM(percent) FROM YOUR_TABLE allows you to have any weights in percent - they don't strictly need to be percentages (i.e. add-up to 100).

[SQL Fiddle]

Answer 3

Your proposed query appears to work; see this SQLFiddle demo. It creates the wrong distribution though; see below.

To prevent PostgreSQL from optimising the subquery I've wrapped it in a VOLATILE SQL function. PostgreSQL has no way to know that you intend the subquery to run once for every row of the outer query, so if you don't force it to volatile it'll just execute it once. Another possibility - though one that the query planner might optimize out in future - is to make it appear to be a correlated subquery, like this hack that uses an always-true where clause, like this: http://sqlfiddle.com/#!12/3039b/9

At a guess (before you updated to explain why it didn't work) your testing methodology was at fault, or you're using this as a subquery in an outer query where PostgreSQL is noticing it isn't a correlated subquery and executing it just once, like in this example. .

UPDATE: The distribution produced isn't what you're expecting. The issue here is that you're skewing the distribution by taking multiple samples of random(); you need a single sample.

This query produces the correct distribution (SQLFiddle):

WITH random_weight(rw) AS (SELECT random() * (SELECT sum(percent) FROM test))
 SELECT id
FROM (                   
  SELECT 
    id,
    sum(percent) OVER (ORDER BY id),
    coalesce(sum(prev_percent) OVER (ORDER BY id),0) FROM (
      SELECT 
        id,
        percent,
        lag(percent) OVER () AS prev_percent
      FROM test
    ) x
) weighted_ids(id, weight_upper, weight_lower)
CROSS JOIN random_weight
WHERE rw BETWEEN weight_lower AND weight_upper;

Performance is, needless to say, horrible. It's using two nested sets of windows. What I'm doing is:

• Creating (id, percent, previous_percent) then using that to create two running sums of weights that are used as range brackets; then
• Taking a random value, scaling it to the range of weights, and then picking a value that has weights within the target bracket

Answer 4

Based on Branko Dimitrijevic's answer, I wrote this query, which may or may not be faster by using the sum total of percent using tiered windowing functions (not unlike a ROLLUP).

WITH random AS (SELECT random() AS random)
SELECT id FROM (
    SELECT id, percent,
    SUM(percent) OVER (ORDER BY id) AS rank,
    SUM(percent) OVER () * random AS roll
    FROM test CROSS JOIN random
) t WHERE roll <= rank LIMIT 1

If the ordering isn't important, SUM(percent) OVER (ROWS UNBOUNDED PRECEDING) AS rank, may be preferable because it avoids having to sort the data first.

I also tried Mechanic Wei's answer (as described in this paper, apparently), which seems very promising in terms of performance, but after some testing, the distribution appear to be off :

SELECT id
FROM test
ORDER BY random() ^ (1.0/percent)
LIMIT 1

Answer 5

ORDER BY random() ^ (1.0 / p)

from the algorithm described by Efraimidis and Spirakis.

Answer 6

Here is something for you to play with:

select t1.id as id1
  , case when t2.id is null then 0 else t2.id end as id2
  , t1.percent as percent1
  , case when t2.percent is null then 0 else t2.percent end as percent2 
from "Test1" t1 
  left outer join "Test1" t2 on t1.id = t2.id + 1
where random() * 100 between t1.percent and 
  case when t2.percent is null then 0 else t2.percent end;

Essentially perform a left outer join so that you have two columns to apply a between clause.

Note that it will only work if you get your table ordered in the right way.