PySpark: подсчет пар слов и отдельных слов

Question

Я пытаюсь выполнить базовый c анализ текста с помощью PySpark.

Чтобы начать эту проблему, я загружаю текстовый файл, содержащий сонеты Шекспира, в RDD. Выводом должен быть RDD, содержащий кортежи, как показано ниже:

(word_pair, count_of_word_pair, word_1_count, word_2_count)

, где word_1 и word_2 - отдельные слова, составляющие пару слов.

Пока я могу правильно сгенерировать один СДР, который содержит отдельные уникальные слова и счетчик количества строк, на которых они встречаются, и другой СДР, содержащий word_pairs (порядок имеет значение) и подсчитывает:

distFile = sc.textFile("Shakespeare.txt")

    tokens = (
        distFile.map(simple_tokenize)
        .flatMap(set)
        .map(lambda x: (x, 1))
        .reduceByKey(lambda x, y: x + y)
    )

    token_pairs = (
        distFile.map(simple_tokenize)
        .flatMap(lambda x: itertools.permutations(x, 2))
        .filter(lambda x: x[0] != x[1])
        .map(lambda x: (x, 1))
        .reduceByKey(lambda x, y: x + y)

Примечание: simple_tokenize - это функция, которая устанавливает все символы в нижний регистр и разбивает предложение в каждой строке на список слов.

Пример выходных данных ниже: Первые 10 элементов токенов:

('shakespeare', 38),
  ('we', 3125),
  ('fairest', 42),
  ('increase', 28),
  ('creatures', 38),
  ('die', 461),
  ('rose', 68),
  ("beauty's", 30),
  ('thereby', 26),
  ('never', 996)]

Первые 10 элементов token_pairs:

[('the', 'sonnets', 2),
  ('sonnets', 'the', 2),
  ('by', 'william', 39),
  ('william', 'by', 39),
  ('from', 'desire', 4),
  ('fairest', 'creatures', 1),
  ('fairest', 'we', 1),
  ('fairest', 'increase', 1),
  ('creatures', 'fairest', 1),
  ('creatures', 'we', 4)],

Я застрял в том, что не понимаю как объединить эти разные значения в одном и том же СДР (согласно требуемому результату). Насколько я понимаю, я не могу выполнить какой-либо тип ключа l oop между СДР внутри функции карты (или что я вообще не могу этого сделать, пока не соберу один из них). Это правильно?

Как получить желаемый результат без необходимости сохранять что-либо как локальную переменную в Python?