Рандомизация очень больших массивов данных

Рассмотрим проблему рандомизации набора данных, который настолько велик, что даже не помещается в памяти. В данной статье описывается, как это можно легко и (относительно) быстро сделать на языке Python.

В наши дни нередко встречаются наборы данных, размер которых измеряется гигабайтами и даже терабайтами. Такой объем данных может оказать огромную помощь в процессе обучения для создания надежных моделей машинного обучения. Но как рандомизировать такие большие массивы данных?

Представьте себе, что у вас есть очень большой набор данных, содержащий по одному элементу на строку в файле. Детали данных не имеют значения для наших целей. Это могут быть строки в файле CSV (значения, разделенные запятыми) или TSV (значения, разделенные табуляцией), или каждая строка может быть объектом JSON, или списком значений X, Y, Z точки в большом облаке точек. Все, что нам нужно, – это чтобы набор данных был отформатирован по одному элементу в строке.

Для файлов, содержащих небольшие наборы данных, можно произвести рандомизацию файла (так называемую “перетасовку”) в памяти с помощью простой функции Python, например, такой:

import random

def shuffle_in_memory(filename_in, filename_out):
    # Shuffle a file, line-by-line
    with open(filename_in) as fp:
        lines = fp.readlines()
    # Randomize them in place:
    random.shuffle(lines)
    # Write the new order out:
    with open(filename_out, "w") as fp:
        fp.writelines(lines)

Функция shuffle_in_memory() принимает на вход имя файла и имя выходного файла, перемешивает строки в памяти с помощью встроенной функции random.shuffle() и записывает рандомизированные данные наружу. Как следует из названия, эта функция требует, чтобы в память были загружены сразу все строки файла.

Чтобы проверить работу этой функции, создадим несколько тестовых файлов. Функция make_file() принимает количество строк, которое вы хотите видеть в тестовом файле. Функция создаст файл и вернет его имя.

import os

def make_file(lines):
    filename = "test-%s.txt" % lines
    print("Making test file '%s'..." % filename)

    with open(filename, "w") as fp:
        for i in range(lines):
            fp.write(f"Line {i}\n")

    print("Done!")
    return filename

Например, создание файла с именем “test-1000.txt”, содержащего 100 строк, выглядит следующим образом:

filename_in = make_file(1000)

После выполнения этой функции в текущем каталоге должен появиться файл с именем “test-1000.txt”, содержащий 1000 строк текста, вида:

Line 0
Line 1
Line 2
Line 3
Line 4
Line 5
Line 6
Line 7
Line 8
Line 9
...

Для тестирования функции shuffle_in_memory() назовем выходной файл, сохраним строку в переменной filename_out и вызовем функцию:

filename_out = "test-randomized-1000.txt"
shuffle_in_memory(filename_in, filename_out)

Теперь в вашем каталоге должен появиться второй файл с именем “test-randomized-1000.txt”. Он должен быть точно такого же размера, как и “test-1000.txt”, с точно такими же строками, но расположенными в случайном порядке:

Line 110
Line 592
Line 887
Line 366
Line 52
Line 22
Line 891
Line 83
Line 931
Line 408
...

Хорошо, теперь главный вопрос: что делать, если у нас очень большой файл? Давайте сделаем файл среднего размера, скажем, с 10 млн. строк. (Для большинства компьютеров это все еще достаточно мало для рандомизации в памяти, но достаточно большой размер, чтобы с ним можно было потренироваться). Как и раньше, мы создаем входной файл вызовом make_file():

filename_in_big = make_file(10_000_000)

Это займет несколько секунд. После этого в Вашем каталоге должен появиться файл с именем “test-10000000.txt”. Он должен начинаться, как и раньше, но содержать 10 миллионов строк. Размер файла составляет около 128 Мбайт.

Как его рандомизировать? Если мы не хотим использовать всю оперативную память или у нас ее недостаточно, то вместо нее можно использовать жесткий диск. Приведем рекурсивный алгоритм, основанный на аналогичной задаче – сортировке. Следующая функция shuffle() основана на алгоритме сортировки слиянием.

Сначала она проверяет, достаточно ли мал файл для того, чтобы его можно было перемешать в памяти (базовый случай в языке рекурсивных функций). Параметр memory_limit задается в байтах. Если размер файла меньше memory_limit, то он будет перемешан в памяти. Если он слишком велик, то он случайным образом разбивается на несколько файлов меньшего размера, и каждый из них рекурсивно перемешивается. Наконец, содержимое меньших файлов объединяется вместе.

Вот функция:

import tempfile

def shuffle(filename_in, filename_out, memory_limit, file_split_count, 
            depth=0, debug=False):
    if os.path.getsize(filename_in) < memory_limit:
        if debug: print(" " * depth, f"Level {depth + 1}",
            "Shuffle in memory...")
        shuffle_in_memory(filename_in, filename_out)
    else:
        if debug: print(
            " " * depth, f"Level {depth + 1}",
            f"{os.path.getsize(filename_in)} is too big;",
            f"Split into {file_split_count} files..."
        )
        # Split the big file into smaller files
        temp_files = [tempfile.NamedTemporaryFile('w+', delete=False)
                      for i in range(file_split_count)]
        for line in open(filename_in):
            random_index = random.randint(0, len(temp_files) - 1)
            temp_files[random_index].write(line)

        # Now we shuffle each smaller file
        for temp_file in temp_files:
            temp_file.close()
            shuffle(temp_file.name, temp_file.name, memory_limit, 
                    file_split_count, depth+1, debug)

        # And merge back in place of the original
        if debug: print(" " * depth, f"Level {depth + 1}", 
            "Merge files...")
        merge_files(temp_files, filename_out)

Если бы речь шла об алгоритме сортировки, то мы бы объединили файлы в обратном порядке, чтобы создать отсортированный порядок. Однако для перетасовки мы не заботимся о том, чтобы объединить их в определенном порядке, поскольку хотим, чтобы они были рандомизированы. Поэтому функция merge_files() выглядит следующим образом:

def merge_files(temp_files, filename_out):
    with open(filename_out, "w") as fp_out:
        for temp_file in temp_files:
            with open(temp_file.name) as fp:
                line = fp.readline()
                while line:
                    fp_out.write(line)
                    line = fp.readline()

Обратите внимание, что мы стараемся не считывать в память сразу все строки файлов. Проверим это, задав ограничение на перетасовку памяти равным размеру файла. Поскольку размер файла не меньше 128 888 890, он будет разбит на несколько файлов меньшего размера. Для данного примера разобьем большой файл на 2, каждый из которых будет достаточно мал для тасования в памяти:

filename_out_big = "test-randomized-10000000.txt"
shuffle(filename_in_big, filename_out_big, 128_888_890, 2, debug=True)

Этот вызов приводит к следующему результату:

 Level 1 128888890 is too big; Split into 2 files...
  Level 2 Shuffle in memory...
  Level 2 Shuffle in memory...
 Level 1 Merge files...

А содержимое результирующего файла “test-randomized-10000000.txt” должно содержать 10 миллионов строк, причем все они должны быть рандомизированы. Более эффективным тестом было бы уменьшение объема памяти до размера, значительно меньшего, чем размер рандомизируемого файла, и разбиение слишком больших файлов более чем на 2. Допустим, мы хотим использовать только около 1 МБ оперативной памяти и разбить файлы на 20 меньших файлов:

shuffle(filename_in_big, filename_out_big, 1_000_000, 20, debug=True)

Данный пример будет использовать не более 1 Мбайт оперативной памяти и рекурсивно декомпозировать подфайлы, превышающие этот размер, по 20 за раз.

Этот алгоритм будет работать с файлами любого размера (главное, чтобы хватило места на диске!). Чем больше памяти вы можете выделить для функции shuffle_in_memory(), тем быстрее она будет работать. Если количество мелких файлов слишком велико, то вы будете тратить слишком много времени на открытие и закрытие файлов. Вы можете попробовать разные числа для memory_limit, но у меня получалось от 20 до 200. Чем больше начальный файл, тем больше подфайлов вам, вероятно, потребуется.

Существуют и другие алгоритмы, которые также можно использовать. Я возлагал большие надежды на запись всех строк в базу данных SQLite и выборку их в случайном порядке, но это оказалось не быстрее, чем приведенный выше код.

import sqlite3

def shuffle_sql(filename_in, filename_out, memory_limit, depth=0, debug=False):
    if os.path.getsize(filename_in) < memory_limit:
        if debug: print(" " * depth, f"Level {depth + 1}",
            "Shuffle in memory...")
        shuffle_in_memory(filename_in, filename_out)
    else:
        if debug: print(
            " " * depth, f"Level {depth + 1}",
            f"{os.path.getsize(filename_in)} is too big;",
            f"Writing to SQLite database..."
        )
        temp_db = tempfile.NamedTemporaryFile(delete=False)
        connection = sqlite3.connect(temp_db.name)
        cursor = connection.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS lines (
                line TEXT
            );
        """)
        with open(filename_in) as fp:
            line = fp.readline()
            while line:
                cursor.execute("INSERT INTO lines (line) VALUES (?);", [line])
                line = fp.readline()
            connection.commit()
        with open(filename_out, "w") as fp:
          for line in cursor.execute("""
              SELECT line FROM lines ORDER BY random();
              """):
              fp.write(line[0])

shuffle_sql(filename_in_big, filename_out_big, 1_000_000, debug=True)

Можете ли вы победить рекурсивный алгоритм перетасовки, оставаясь чисто на Python?