ANSI SQL equivalent of pandas `factorize()`?

user_7212929 0
Viewed 135

So I have to create user IDs from user's emails and all data is present on BigQuery. In python, pandas, its an easy one-liner as:

all_data['user_id'] = all_data['email'].factorize()[0]

But I cannot figure out a way to do this in BigQuery SQL. I tried using RANK() function but its nowhere close. Currently I am trying to make use of Window functions with RANK() but it seems a bit of a stretch to have such approach for such a simple task. All data is on BigQuery already so any advice of doing it some other way even outside SQL would be good too.

A little context...

  1. pandas factorize() function assigns a unique ID based on column provided so if emails are like email1@example.com, email2@example.com, email1@example.com, email3@example.com, email1@example.com, email2@example.com, it would return: [0, 1, 0, 2, 0, 1] and so on.

  2. I have other columns in database as well so RANK() and ROW_NUMBER() seem not to help alone. I am trying to get around tho.

pandas google-bigquery python
2 Answers

Consider below two options

Note, I am using slightly modified data example - you will see why (I hope)

with `project.dataset.table` as (
  select '2021-01-01 00:01:00' sent , 'email4@example.com' recipient  union all 
  select '2021-01-01 00:02:00', 'email2@example.com' union all 
  select '2021-01-01 00:03:00', 'email4@example.com' union all 
  select '2021-01-01 00:04:00', 'email3@example.com' union all 
  select '2021-01-01 00:05:00', 'email4@example.com' union all 
  select '2021-01-01 00:06:00', 'email2@example.com'
)

Option 1:

In case if there is an order in which those emails should be set before assigning unique_id - for example by sent column. In this case consider below

#standardSQL
create temp function factorize(item string, list any type) as ((
  select unique_id from (
    select as struct recipient, row_number() over(order by min(sent)) - 1 unique_id
    from unnest(list)
    group by recipient
  ) 
  where recipient = item
));
select t.*, 
  factorize(recipient, array_agg(struct(recipient, sent)) over()) unique_id 
from `project.dataset.table` t

with output

enter image description here

Option 2:

In case if ordering does not matter much and you can just order alphabetically - consider below slightly [in my mind] simpler query with use of built-in range_bucket function

#standardSQL
create temp function factorize(item string, list any type) as (
  range_bucket(item, list) - 1 
);
with all_recipients as (
  select array_agg(recipient order by recipient) recipients from (
    select recipient
    from `project.dataset.table`
    group by recipient
  )
)
select t.*,
  factorize(recipient, recipients) unique_id
from `project.dataset.table` t, all_recipients

with output

enter image description here

Obviously in this case you can skip using udf and just simply use rabge_bucket within the final select (instead of within udf)

select t.*,
  range_bucket(recipient, recipients) - 1 unique_id
user_52219440

You can use the DENSE_RANK() window function for this purpose:

select dataset.*, DENSE_RANK() OVER (ORDER BY email)
from dataset
order by sent;

This would produce something like (using Mikhail Berlyant's example data as starting point):

SENT EMAIL DENSE_RANK
2021-01-01 00:01:00 email4@example.com 3
2021-01-01 00:02:00 email2@example.com 1
2021-01-01 00:03:00 email4@example.com 3
2021-01-01 00:04:00 email3@example.com 2
2021-01-01 00:05:00 email4@example.com 3
2021-01-01 00:06:00 email2@example.com 1
user_4668620
Related