Sampling column value for each group from a Bigquery table after group by where the value is not present in the column values of current group

Viewed 105

I have a table in the form

QUERY | DOC | DOC_ID
---------------------
q1    | d1  |  1
---------------------
q1    | d2  |  2
---------------------
q2    | d3  |  3
---------------------
q3    | d4  |  4
---------------------

The goal is to group the table by (query, doc) pairs and for each group, select K docs (random) not present in the group and assign them to a new column. For example one possible output with K=1 is

QUERY | DOC | DOC_ID | RND_DOC
------------------------------
q1    | d1  |  1     | d3
------------------------------
q1    | d2  |  2     | d3
------------------------------
q2    | d3  |  3     | d4
------------------------------
q3    | d4  |  4     | d1
------------------------------

The RND_DOC selected is random from rows not in the group. For K>1, we would create new rows to accommodate the new samples. I started doing a cross join with the table and then sampling, but as suspected for big tables that computation is not feasible. What is a good way to go about it considering the table in Bigquery is quite big?

1 Answers

Consider below (with K=2)

select * except(ALL_DOC)
from (
  select *, array_agg(DOC) over() ALL_DOC
  from `project.dataset.table`
), unnest(array(
   select x
   from unnest(ALL_DOC) x
   where x != DOC
   order by rand()
   limit 2
 )) RND_DOC

if applied to sample data in your question - output is

enter image description here

Related