Введение

При работе с большими коллекциями в MongoDB, размер которых превышал десятки миллионов записей, возникла необходимость формировать случайные выборки уникальных значений полей, принадлежащих документам этой коллекции.

Для такой операции, в MongoDB штатно предусмотрена функция $sample, которую можно использовать в составе pipeline при проведении агрегации данных. Есть еще один вариант, с использованием оператора $skip. Однако, как показала практика, выполнение выборки полей таким образом на большой коллекции и первым и втором способом может занимать весьма ощутимое время.

Чтобы сократить время выполнения таких выборок, потребовалось разработать собственный алгоритм, который на порядки увеличил скорость работы. Ниже приведен подход и вариант реализации данного алгоритма.

Исходные данные

В качестве исходных данных возьмем базу электронных писем, содержащую, в том числе, поле с адресом электронной почты и темой письма и количество элементов порядка 800 000.

Типовой документ, представляемый в БД Mongo будет выглядеть как стандартный JSON такого вида:

{
  "_id": {
    "$oid": "6461f8cbd2b81e62fcfo54ca"
  },
  "mail": "letter@bookvoed.ru",
  "subject": "Дорогу книголюбам",
}

Служебное поле "_id" в данном документе формируется при внесении документа в коллекцию автоматически. Важно, что это поле всегда индексируется по умолчанию и его значение вычисляется по строго установленным принципам, делающим его уникальным в пределах базы данных.

Согласно сведениям разработчика MongoDB, значение поля "_id" является BSON типом ObjectID и состоит из трех частей:

• 4-байтовая метка времени генерации записи, выраженная в секундах эпохи Unix.

• 5-байтное случайное число, генерируемое единожды для рабочего процесса, осуществляющего внесение документа в коллекцию.

• 3-байтный счетчик, инициализированный случайным образом.

  На рисунке 1 представлены варианты значений для тестовой коллекции.

Из рисунка видно, что при формировании коллекции был задействован единственный процесс, запись документов проводилась одним потоком (отметки времени являются регулярными) и счетчик инкрементируется на единицу для каждого документа.

В ходе экспериментов, было установлено, что при смене процесса меняется и идентификатор процесса, и стартовое значение счетчика. Таким образом, стабильно последовательным является только значение поля time.

Хотя, если коллекция заполняется документами нерегулярно - то в этом поле будут появляться разрывы. Для примера, на рисунке один приведена статистика для базы, емкостью около 100 000 000 записей (на Рисунке 2 поле time преобразовано в UNIX-время).

Из рисунка видно, что запись в таблицу проводилась в три этапа, в разное время, появились разрывы в псевдо-монотонной последовательности временных отметок документов.

Заметим, что при больших таблицах поле "счетчик" будет принимать все значения от 000000 до ffffff (см. Рисунок 2), что вполне логично.

С учетом сказанного, можно предположить, что для формирования случайной выборки можно работать с фрагментом поля _id, соответствующего времени внесения документа в коллекцию.

Так как этот фрагмент находится в старших разрядах значения поля "_id", индекс, созданный СУБД по умолчанию для этого поля будет корректно отрабатывать при выборе значений для этого фрагмента.

Алгоритм для регулярно заполняемых коллекций

Простой алгоритм случайного выбора нескольких oid может быть таким:

1. Найти минимальное (min_oid) и максимальное (max_oid) значение ObjectID и сгенерировать случайное значение (rand_oid) между ними.

2. Найти ближайший к сгенерированному значению реальный ObjectID из обрабатываемой коллекции.

3. Проверить, был ли уже ранее найден такой rand_oid, и если да - то сгенерировать новое случайное значение, после чего повторить процедуру выбора oid.

4. Повторить процедуру заданное число раз.

Алгоритм корректно отработает для коллекций, где записи вносятся регулярно, а не большими пачками с большими промежутками (как в примере на рисунке 2) и для некоторых случаев даже такого алгоритма будет уже достаточно. Тем не менее, если в псевдо-монотонной последовательности временных отметок будут разрывы, то случайное число на полном диапазоне будет часто попадать в эти разрывы и надо бы учесть такую особенность коллекции.

Что для этого можно сделать?

Алгоритм, учитывающий временные разрывы

Условно (масштаб не соблюден), можно представить коллекцию с разрывами как-то так (см. Рисунок 3):

Если мы попадаем в "молоко", то ближайшими идентификаторами реальных документов окажутся идентификаторы документов на краю диапазонов. При этом, процессы, которые эти документы разместили в базу, вероятнее всего будут разными. Будем использовать это как признак, для определения попадания в "молоко" таким образом:

1. Берем случайный rand_id и ищем левый (оператор MongoDB $lt) и правый ($gt) по отношению к нему ObjectID реального документа.

2. Проверяем, совпадают ли у найденных документов идентификаторы процесса. Если нет - то мы нашли граничные документы диапазонов. В таком случае повторно генерируем новый случайный rand_id в диапазоне, который лежит между краем (началом или концом) общего диапазона и найденным граничным документом.

3. Повторяем проверку и сужение диапазона генерации случайного rand_id пока не попадем в "зеленку", о чем будет свидетельствовать совпадение идентификаторов процесса у "левого" и "правого" по отношению к rand_id реального ObjectID.

4. Включаем найденный реальный ObjectID (либо "левый2, либо "правый", либо оба) в формируемый набор.

Реализация и проверка

Проведем первичный эксперимент по формированию выборок почтовых адресов и оценим скорость их выполнения.

Из специфических модулей, которые потребуются для тестирования, нам потребуется штатный модуль pymongo, который можно скачать из открытого репозитория.

Для оценки времени исполнения удобно использовать декоратор, дополняющий результат выполнения произвольной функции Python расчетом времени ее отработки.

def timer(func):
    def _wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        runtime = time.perf_counter() - start
        print(f"{func.__name__} took {runtime:.4f} secs")
        return result
    return _wrapper

Прежде всего составим функцию, которая будет находить случайный идентификатор ObjectID:

def get_random_objid(min_oid, max_oid):
    """
    Функция генерирует случайный ObjectID в диапазоне
    """
    f = str(min_oid)[:8] # берем первые 8 символов, соответствующие временной отметке записи
    l = str(max_oid)[:8]
    fn = int(f, 16)
    ln = int(l, 16)
    random.seed()
    r = random.randrange(fn, ln)
    rand_oid = ObjectId("".join([hex(r)[2:],"0000000000000000"]))
    return rand_oid

Замечу, что для генерации имеет смысл задавать только разряды ObjectID соответствующие отметке времени, то есть первым 4 байтам идентификатора.

Для использования функции надо будет задать минимальный и максимальный идентификатор соответствующей временным границам обрабатываемой коллекции, которые можно вычислить следующим образом:

# MongoClient(used_connection)["tmp"]
min_oid = list(coll.find().sort("_id", 1).limit(1))[0]["_id"]
max_oid = list(coll.find().sort("_id", -1).limit(1))[0]["_id"]

Сортировка с параметром "-1" упорядочивает коллекцию от большего к меньшему значению идентификатора.

Ниже приведена функция, которая находит идентификаторы реальных документов в коллекции, расположенных слева и справа от случайно найденного идентификатора, а также функция, которая разбивает идентификатор на логические составляющие - "Время", "Процесс", "Счетчик".

def trim_objid(obj_id):
    s = str(obj_id)
    return {"id":obj_id, "time":s[:8],"proc":s[8:18], "counter":s[18:]}
  
def get_near_oid(coll, reference_oid):
    left = list(coll.aggregate([{"$sort": {"_id": -1}},
                                {"$match": {"_id": {"$lte": reference_oid}}}, 
                                {"$limit":1}]))
    if not left:
        left = [coll.find_one()]
    right = list(coll.aggregate([{"$match": {"_id": {"$gte": reference_oid}}}, 
                                 {"$limit":1}]))
    if not right:
        right = list(coll.aggregate([{"$sort": {"_id": -1}}, {"$limit":1}]))

    return {"left":trim_objid(left[0]["_id"]), "right":trim_objid(right[0]["_id"])}

В итоге отработки, функция поиска ближайших идентификаторов get_near_oid выдаст примерно такой результат:

{'left': {'id': ObjectId('6461f8cdd2b81e62fcf16ffd'), 
          'time': '6461f8cd', 
          'proc': 'd2b81e62fc', 
          'counter': 'f16ffd'}, 
 'right': {'id': ObjectId('6461f8ced2b81e62fcf16ffe'), 
           'time': '6461f8ce', 
           'proc': 'd2b81e62fc', 
           'counter': 'f16ffe'}}

На данном примере видно, что у "правого" и "левого" идентификатора процесса ("proc") совпадают, то есть мы попали в "зеленку", содержательную часть коллекции, а не в "молоко", разрыв, и можно спокойно брать полученный идентификатор для последующего включения в набор.

Функция, которая возвращает "корректные" идентификаторы из содержательных блоков коллекции будет выглядеть так:

def get_inline_oids(coll):
    min_oid = list(coll.find().sort("_id", 1).limit(1))[0]["_id"]
    max_oid = list(coll.find().sort("_id", -1).limit(1))[0]["_id"]
    rand_oid = get_random_objid(min_oid, max_oid)

    # ищем идентификаторы документов, которые находятся слева и справа от референсного
    gap = get_near_oid(coll, rand_oid)

    while gap["left"]["proc"] != gap["right"]["proc"]:
        rand_oid = get_random_objid(min_oid, gap["left"]["id"])
        gap = get_near_oid(coll, rand_oid)
        if gap["left"]["proc"] != gap["right"]["proc"]:
            rand_oid = get_random_objid(gap["right"]["id"], max_oid)
            gap = get_near_oid(coll, rand_oid)
    print(gap)
    return {gap["left"]["id"], gap["left"]["id"]}

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

Здесь же приведем стандартную функцию выборки случайных значений для последующего сравнения производительности.

Ну и проведем тест на выборках различной длины: 5, 10, 50.

@timer
def get_sample_oids(coll, count=1):
    oids = set()
    while len(oids) < count:
        oids = oids.union(get_inline_oids(coll))
    if len(oids) != count:
        oids.pop()
    return oids

@timer
def get_samples_standard(coll, count=1):
    res = list(coll.aggregate([{"$sample": {"size": count}},{"$project":{"_id":1}}]))

    return {value["_id"] for value in res}

for i in [5,10,50]:
    print("Collection size: {}, sample size: {}".format(coll.estimated_document_count(), i))
    get_sample_oids(coll,i)
    get_samples_standard(coll,i)

В результате отработки, получаем следующие результаты:

Collection size: 102316457, sample size: 5
get_sample_oids took 0.0247 secs
get_samples_standard took 2.9753 secs
Collection size: 102316457, sample size: 10
get_sample_oids took 0.0495 secs
get_samples_standard took 3.2230 secs
Collection size: 102316457, sample size: 50
get_sample_oids took 1.5310 secs
get_samples_standard took 3.2344 secs

То есть для небольших выборок, скорость отличается на два порядка, что весьма и весьма неплохо!

Понятно, что данный код не является идеально оптимизированным, что заметно на больших выборках. Здесь приходится тратить время на "избегание" разрывов и при попадании в разрыв повторно производить поиск нового идентификатора.

Если сделать модернизацию алгоритма, в которой единожды найденные разрывы будут запоминаться и не участвовать в дальнейшем для выбора случайных значений идентификатора - то скорость не будет падать и на больших выборках тоже.

Также, следует отметить, что вся эта канитель со случайными выборками может вообще перестать быть проблемой, если сразу принять за правило при внесении документа в коллекцию базы MongoDB формировать идентификатор не по умолчанию, а свой собственный, уникальный, инкрементальный.

Но... Всякое бывает в жизни и имеет право на жизнь. К разработанному алгоритму это тоже относится.

:)