Is it possible to have stratified train-test split of a set based on two columns?
One option would be to feed an array of both variables to the stratify parameter which accepts multidimensional arrays too. Here's the description from the scikit documentation:

stratify array-like, default=None
If not None, data is split in a stratified fashion, using this as the class labels.


Here is an example:
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

# create dummy data with unbalanced feature value distribution
X = pd.DataFrame(np.concatenate((np.random.randint(0, 3, 500), np.random.randint(0, 10, 500)),axis=0).reshape((500, 2)), columns=[&quot;text&quot;, &quot;subreddit&quot;])
y = pd.DataFrame(np.random.randint(0,2, 500).reshape((500, 1)), columns=[&quot;label&quot;])

# split stratified to target variable and subreddit col
X_train, X_test, y_train, y_test = train_test_split(
    X, pd.concat([X[&quot;subreddit&quot;], y], axis=1), stratify=pd.concat([X[&quot;subreddit&quot;], y], axis=1))

# remove subreddit cols from target variable arrays
y_train = y_train.drop([&quot;subreddit&quot;], axis=1)
y_test = y_test.drop([&quot;subreddit&quot;], axis=1)

As you can see the split is stratified to subreddit too:
Train data shares for subreddits
X_train.groupby(&quot;subreddit&quot;).count()/len(X_train)

gives
text
subreddit   
0   0.232000
1   0.232000
2   0.213333
3   0.034667
4   0.037333
5   0.045333
6   0.056000
7   0.056000
8   0.048000
9   0.045333

Test data shares for subreddits
X_test.groupby(&quot;subreddit&quot;).count()/len(X_test)

gives
text
subreddit   
0   0.232
1   0.240
2   0.208
3   0.032
4   0.032
5   0.048
6   0.056
7   0.056
8   0.048
9   0.048


Naturally, this only works if you have sufficient data to stratify to subreddit and the target variable at the same time. Otherwise scikit learn will throw an exception.