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 additionalstratified_column_name parameter to train. In that case though, the sampling wouldn't bestratified_random. So we'd need to add a different stratified type, or we could just call thisstratified, 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 juststratified. 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, thenalter extension pgml update. You'll want toalter type add value... here rather than dropping the enum, as that would only work on an empty database.