Memory leak on writes/merges

[echai58]

Environment

Binding: python

---

Bug

What happened:
We're noticing constantly rising memory in our processes that write to deltalake. I wrote a minimal reproduction that loops and writes to deltalake, and the memory usage seems to indicate a memory leak.

What you expected to happen:
Memory to be reclaimed after writes.

How to reproduce it:
This is my script I tested with:

import datetime
from deltalake import write_deltalake, DeltaTable
import numpy as np
import pandas as pd 
import pyarrow as pa

schema = pa.schema(
    [
        ("date", pa.date32()),
        ("c1", pa.string()),
        ("c2", pa.int32()),
        ("v1", pa.float64()),
        ("v2", pa.float64()),
        ("v3", pa.int32()),
    ]
)

def write_table(dt: DeltaTable, date: datetime.date):
    num_rows = 150
    data = pd.DataFrame(
        {
            "date": np.repeat(date, num_rows),
            "c1": [f"str_{i}" for i in np.random.randint(0, 100, num_rows)],
            "c2": np.random.randint(1, 100, num_rows),
            "v1": np.random.rand(num_rows) * 100,
            "v2": np.random.rand(num_rows) * 100,
            "v3": np.random.randint(0, 100, num_rows),
        }
    )
    dt.merge(
        pa.Table.from_pandas(data, schema=schema),
        "s.date = t.date and s.c1 = t.c1 and s.c2 = t.c2",
        source_alias="s",
        target_alias="t",
    ).when_matched_update_all().when_not_matched_insert_all().execute()
    if dt.version() % 10 == 0:
        dt.create_checkpoint()

def main():
    DeltaTable.create("test", schema)

    for date in pd.date_range("2022-01-01", periods=25):
        for _ in range(50):
            write_table(DeltaTable("test"), date.date())

if __name__ == "__main__":
    main()

More details:
Here's the memray graph:

I also tested this with just write_deltalake(mode="append"), and the issue seems to also persist:

I saw #2068 and tried setting that env var, and got the following (doesn't seem to help):

[ion-elgreco]

Can you do the following, change your script to sleep first for 30 secs, then create the table and then sleep again for 30 secs, then start writing in a loop those 50 times?

I think the slow increase in resident size might just be because at every write you update the table state at the end of the commit since it includes more info now.

[echai58]

here's the mem graph with the 30 second sleeps:

[echai58]

@ion-elgreco

This is a graph for 250 merges:

and for 250 appends:

The slow increase in the 250 appends I think would correspond to the increased table state metadata. But the merge memory use rises much faster, so maybe it is something particular with merges?

[ion-elgreco]

Merge operations probably holds more info, so this looks normal to me

[echai58]

@ion-elgreco The _last_checkpoint file says:
{"size":360,"size_in_bytes":75770,"version":100}

for the last checkpoint after running a script with 1000 merges resulting in the following memory increase:

The amount the memory increased seems much larger than the metadata

[ion-elgreco]

@echai58 the checkpoint is compressed and also would never translate 1:1 from disk to mem afaik

[echai58]

@ion-elgreco profiling a script that just instantiates the same delta table gives the following:

~13 mb , which is still much less than the >100mb seen from the merge script

[malachibazar]

Has there been any progress on this? I'm experiencing the same issue with merges.

[rtyler]

I have a theory that this might be related to some of the performance issues that @roeap and I were hunting after this week. He's got some fixes in mind after which we can look into this specific issue some more

[aldder]

Hello, any update for this topic?

[kapteinkaald]

I think we have the same issue here.

[ion-elgreco]

There are a couple features in main that have improved the memory usage:

• chore(python): use custom global alloc #3176 <-- improves memory allocation (more stable across)
• feat: streamed execution in MERGE #3145 <-- consumes less memory by not collecting everything in memory using a LazyTableProvider
• refactor: logical plans in writer #3141 <-- potentially better optimized plan execution
• feat: use LazyTableProvider by default for write_to_deltalake for memory efficiency #3196 (WIP) <-- same approach as feat: streamed execution in MERGE #3145.

@rtyler is doing some testing at the moment on the memory perf. The numbers are promising imho, even if it's a micro benchmark for normal appends. Before these PRs in main it was around 500-700MB memory usage. Now it's at a stable 500MB, with the WIP PR #3196 using the LazyTableProvider it's at 30-50MB ^^