title: Sorting train_test_split data by numpy array
body: <p>I want to split the following numpy arrays for training and testing: <code>X</code>, <code>y</code> and <code>qid</code></p>
  <ul>
  <li><code>X</code> is a set of featurized documents - <strong>shape</strong>: (140, 105)</li>
  <li><code>qid</code> is a set of query identifiers for each document - <strong>shape</strong>: (140,)</li>
  <li><code>y</code> is a set of labels for each (<code>X</code>, <code>qid</code>) pair - <strong>shape</strong>: (140,)</li>
  </ul>
  <p>At the moment, what I do for splitting is:</p>
  <pre><code># Split documents, labels, and query_ids into training (70%) and testing (30%)
      X_tr, X_tst, y_tr, y_tst, qid_tr, qid_tst= train_test_split(X, y, qid, test_size=0.3, random_state=1, shuffle=True, stratify=qid)
  </code></pre>
  <p>The problem is that after splitting, I need the returning numpy arrays to be sorted by <code>qid</code>. That is, all the documents with the <strong>same</strong> <code>qid</code> need to be together (one after another) as a block (both in training and testing).</p>
  <p><strong>Example</strong></p>
  <p>Correct split:</p>
  <pre><code>X              qid           y       
  ------------------------------
  document 1     0             0
  document 5     0             1
  document 4     1             1
  document 6     1             0
  document 9     2             1
  </code></pre>
  <p>Incorrect split:</p>
  <pre><code>X              qid           y       
  ------------------------------
  document 1     0             0
  document 4     1             1
  document 9     2             1
  document 5     0             1
  document 6     1             0
  </code></pre>
  <p>Is there any way to make this possible?</p>