Thinking ahead to other stratification strategies, that use columns other than the target y_column_name, to guarantee you have true out of sample rows. e.g. user_id may be used to train with multiple instances from a particular user, but you want to ensure there is no data leakage where the model is just memorizing user_ids rather than the more abstract, so user_id should be excluded as a feature, but used for stratification.

I think that could work as an additional stratified_column_name parameter to train. In that case though, the sampling wouldn't be stratified_random. So we'd need to add a different stratified type, or we could just call this stratified, and if you don't specify a column, it's random by the y_column_name. If you do specify stratification column(s), then those column(s) get removed from features, and strictly used for stratification.

This can happen in a follow up PR, I'm just commenting so we get a forward looking name on the sampling strategy.

Good ideas/points. I changed it to just stratified. For this PR, there will be no ability to specify the columns to stratify by, stratified only uses y_column_name. I think that should be a non-breaking change to add an optional parameter in the future that would change it from y_column_name to something else.

I still need to do some more testing to make sure this migration works as intended.

@montanalow , are there integration tests that assert migrations work? I didn't see any...

There aren't integration tests. I typically run the test/test.sql to populate a database with the previous version, then alter extension pgml update. You'll want to alter type add value... here rather than dropping the enum, as that would only work on an empty database.