title: How to test/train split by column value rather than by row in pyspark
body: <p>I would like to generate a train and test set for machine learning.  Let's say I have a dataframe with the following columns:</p>
  <pre><code>account_id | session_id | feature_1 | feature_2 | label
  </code></pre>
  <p>In this dataset, each row will have a unique session_id, but an account_id can show up multiple times.  However, I want my train and test sets to have mutually exclusive account_ids.  (Almost seems like the opposite of stratified sampling).</p>
  <p>In pandas, this is simple enough.  I have something like the following:</p>
  <pre><code>def train_test_split(df, split_col, feature_cols, label_col, test_fraction=0.2):
      &quot;&quot;&quot;
      While sklearn train_test_split splits by each row in the dataset,
      this function will split by a specific column. In that way, we can 
      separate account_id such that train and test sets will have mutually
      exclusive accounts, to minimize cross-talk between train and test sets.
      &quot;&quot;&quot;
      split_values = df[split_col].drop_duplicates()
      test_values = split_values.sample(frac=test_fraction, random_state=42)
      
      df_test = df[df[split_col].isin(test_values)]
      df_train = df[~df[split_col].isin(test_values)]
  
      return df_test, df_train
  </code></pre>
  <p>Now, my dataset is large enough that it cannot fit into memory, and I have to switch over from pandas to doing all of this in pyspark.  How can I split a train and test set to have mutually exclusive account_ids in pyspark, without fitting everything into memory?</p>