def split_by_length_threshold(
    df, test_set_size):
  """Splits off a test set based on text length.

  Finds a length threshold such that all texts longer than this threshold make
  up a test set that is at max `test_set_size` in size. Note that the exact
  size cannot be guaranteed.

  Args:
    df: DataFrame with 'text_len' column containing the length of every text.
    test_set_size: Number of elements the test set must contain.

  Returns:
    1. Length threshold.
    2. Lengths that exist in the dataset and that are larger than the threshold.
    3. Binary mask in the size of the original data, where True indicates that
      the example must be part of the test set.

  Raises:
    ValueError: In case no threshold can be found such that `test_set_size`
      elements end up in the test set.
  """
  current_count = 0

  # Start from the longest texts.
  for i in range(max(df['text_len']), 0, -1):
    current_count += len(df.loc[df['text_len'] == i, :])

    if current_count > test_set_size:
      return (i,
              set(df['text_len']) & set(range(i + 1,
                                              max(df['text_len']) + 1)),
              df.loc[:, 'text_len'] > i)
  raise ValueError(
      'No length threshold found to create a test set with size {}.'.format(
          test_set_size))