Splits stratified dataset into train and test using SQL

stratified_train_test_split.sql

/* Stratified splits dataset into training and test 

It guarantees that each group has the minimum size to be split.
*/

with ssize as (
	select
		group
	from  to_split_table
	group by group
	having count(*) >= {{ MINIMUM GROUP SIZE }}) -- {{ MINIMUM GROUP SIZE }} = 1 / {{ TEST_THRESHOLD }}
select
	id,
	ts.group,
	case
		when
    	cast(row_number() over (partition by ts.group order by rand()) as double) / cast(count() over (partition by ts.group) as double)
		< {{ TEST_THRESHOLD }} then 'test'
		else 'train'
	end as splitting
from  to_split_table ts
join ssize
on ts.group = ssize.group