Memory usage for merge_insert

[ellchow]

Hi,

I am trying out merge_insert for large data and memory usage seems to be very high. I have a simple script here that should reproduce this and also provides other points for comparison. The incoming data source is a generator of fake data, so it should require limited memory. For upsert and insert-if-not-exists, it seems to load the entirety of the incoming data source in memory. Other operations like creating a table from scratch and merge for adding a new column seem to behave as expected though.

Is my usage correct /is this expected? Is there a reason why merge_insert requires all data to be loaded in memory while `merge does not?

import os
import math
import lancedb
import lancedb.common
import lancedb.db
import pyarrow as pa
import time

if __name__ == "__main__":
    db_url = "test-db"
    table_name = "test-table-{}".format(int(time.time()))
    n = 10_000 # number of rows
    bsize = 5_000_000 # binary data size
    operation = 1 # choose operation to test below

    ######################
    b = os.urandom(bsize)

    def source(n, with_bin=None, batch_size=10):
        if with_bin is not None:
            full_bin_batch = pa.array([b[:with_bin]] * batch_size, pa.large_binary())

        for i in range(int(math.ceil(n / batch_size))):
            ids = range(i * batch_size, min(n, (i + 1) * batch_size))
            res = {"id": ids}

            if with_bin is not None:
                res["bin"] = full_bin_batch.slice(0, len(ids))

            t = pa.Table.from_pydict(res)

            yield t

    ##################
    if operation == 1:
        print("iterate data")
        for _ in source(n, with_bin=bsize):
            pass

    ##################
    if operation == 2:
        print("create table")
        db = lancedb.connect(db_url)
        t = db.create_table(table_name, source(n, with_bin=bsize))

    ##################
    if operation == 3:
        print("add column")
        db = lancedb.connect(db_url)
        t = db.create_table(table_name, source(n, with_bin=None))
        t.merge(source(n, with_bin=bsize), "id")

    ##################
    if operation == 4:
        print("upsert")
        db = lancedb.connect(db_url)
        t = db.create_table(table_name, source(n, with_bin=0))
        t.merge_insert("id"). \
            when_matched_update_all(). \
            when_not_matched_insert_all(). \
            execute(source(n, with_bin=bsize))
    ##################
    if operation == 5:
        print("insert if not exists")
        db = lancedb.connect(db_url)
        t = db.create_table(table_name, source(n, with_bin=0))
        t.merge_insert("id"). \
            when_not_matched_insert_all(). \
            execute(source(n, with_bin=bsize))

[wjones127]

merge_insert implementation currently takes a lot of memory. We have plans to eventually optimize this, including with:

• non-indexed merge_insert is loading too many columns #3481
• merge_insert is loading too many columns during indexed merge insert #3480