Heavy & increasing memory usage when looping write_parquet

[SanilK2108]

Describe the bug

We’re building a pipeline that reads a large dataset, runs memory-intensive SQL/window transforms partition-by-partition, and writes the results to partitioned Parquet. While iterating partitions (and even in a much simpler write-only loop), the pod/process RSS steadily increases over time, despite deleting DataFrames, clearing Daft caches, and forcing GC / Arrow pool release.

Daft version: Tried with 0.4.12 and 0.6.9
Python: 3.10 (also repro on 3.11)
OS: Linux (Kubernetes pod) and macOS (local)
Storage target: local filesystem (Hive-style partitions)

To Reproduce

import asyncio, gc
import daft
import pyarrow
from daft.runners.runner import get_or_create_runner

num_rows = 5000000
batch_size = 1000 
num_batches = (num_rows + batch_size - 1) // batch_size
partition_count = 1000
output_path = "/tmp/daft_out"

with daft.execution_config_ctx(
    parquet_target_filesize=100_000,
    default_morsel_size=1_000,  # also repros without this
    native_parquet_writer=True,   # also repros without this
):
    for batch_num in range(num_batches):
        start_idx = batch_num * batch_size
        end_idx = min(start_idx + batch_size, num_rows)

        # build a small batch with a partition_id column
        ids = [f"entity_{i}" for i in range(start_idx, end_idx)]
        partition_ids = [i % partition_count for i in range(start_idx, end_idx)]
        values = [i * 10 for i in range(start_idx, end_idx)]
        timestamps = [1000 + i * 10 for i in range(start_idx, end_idx)]

        batch_df = daft.from_pydict({
            "id": ids,
            "partition_id": partition_ids,
            "value": values,
            "timestamp": timestamps,
        }).with_column("id", daft.col("id").cast(daft.DataType.string())
        ).with_column("partition_id", daft.col("partition_id").cast(daft.DataType.int64())
        ).with_column("value", daft.col("value").cast(daft.DataType.int64())
        ).with_column("timestamp", daft.col("timestamp").cast(daft.DataType.int64()))

        # Partitioned write
        batch_df.write_parquet(
            root_dir=output_path,
            write_mode="append",
            partition_cols=["partition_id"],
        )

        # cleanup attempts
        get_or_create_runner()._part_set_cache.clear()
        del batch_df
        pyarrow.default_memory_pool().release_unused()
        gc.collect()

        # small pause
        asyncio.run(asyncio.sleep(0.05))

        if (batch_num + 1) % 50 == 0:
            print("pyarrow total_allocated_bytes:", pyarrow.total_allocated_bytes())

Process RSS increases steadily with each batch (pod memory graph attached in our environment)

Expected behavior

We expected that once the call to write_parquet is done, all the memory being used by it should go out of scope and be cleared.

Component(s)

Parquet

Additional context

We are trying to read massive tables (tens of millions of rows), apply complex SQL transformations, including window functions, deduplication, and other memory-intensive operations, and then write the results efficiently back to partitioned Parquet outputs for downstream consumption. To manage scale, we partition the dataset and process each partition iteratively, aiming to keep memory usage stable and predictable within each loop.

This is the rough description of what we are doing -

for partition_id in range(partition_count):
    # read one partition
    pattern = f"{staging_path}/partition_id={partition_id}/**/*.parquet"

    try:
        df = daft.read_parquet(pattern, infer_schema=False)  # schema provided in real code
    except Exception:
        continue

    # window: keep latest per id by timestamp desc
    transformed = (
        df.with_column(
            "row_num",
            row_number().over(Window().partition_by("id").order_by(col("timestamp"), desc=True)),
        )
        .filter("row_num = 1")
        .exclude("row_num")
    )

    # write partitioned parquet (same partition_id column)
    transformed.write_parquet(
        root_dir=output_path,
        write_mode="append",
        partition_cols=["partition_id"],
    )

    # cleanup
    del df, transformed

Are we using write_parquet properly for this pattern (iterative writes to Hive-style partition folders)? Any flags or APIs we should use instead (e.g., different writer, different execution config, flushing, write mode, partition handling)?

If not, what approach should we follow for our usecase?

[colin-ho]

Hey @SanilK2108 , thanks for the detailed issue and repro instructions.

I'm running it now with latest daft on my macbook and unfortunately can't repro it, memory hovers consistently around 160mb. I also removed all the manual cleanup attempts and execution config and memory hovers around 190mb.

I'm gonna try running on linux next. In the meantime, one thing that might be useful is using https://github.com/bloomberg/memray, which is what i use to profile memory. It can create flamegraphs of which functions allocate the most memory.

This is a screenshot of a live run:

[SanilK2108]

Hi @colin-ho , Thanks for replying.

I reran the script I sent you with and without memray live tracking and I am seeing the same graph on grafana of pod memory usage -

This is the exact metric I am looking at -

sum(container_memory_working_set_bytes{namespace="memory-leak-test"}) by (pod,container)

However, on memray live tracking, I don't see any increase in memory usage, and memory consistently stays under ~100MB

I had also tried using psutil to get RSS memory of the current process locally, and found similar consistent memory usage around 100-200MB.

I tried looking at sum by(pod,container)(container_memory_rss{namespace="memory-leak-test"}) on my pod, and that is also under 200MB

Is there any reason we would not see this memory growth in RSS, but in container_memory_working_set_bytes?

[colin-ho]

Are there any other metrics you are able to see? e.g. container_memory_cache or container_fs_writes_bytes_total

A possible theory is that during the writes, the kernel has not started flushing the dirty pages used for the writes, and so those count towards working_set_memory .

[SanilK2108]

Hi @colin-ho, adding more details on container_memory_cache, container_fs_writes_bytes_total -

container_memory_rss and container_fs_writes_bytes_total are relatively small.

container_memory_cache does grow with the same step-pattern we are seeing for container_memory_working_set_bytes, but seems to be growing at half the speed of container_memory_working_set_bytes

I had a few questions based on these -

Do you think this is a Daft issue or a Kernel / OS level issue? I had some doubts on Daft configuration we could experiment with

• Is is possible its linked to concurrency / target file sizes / other related parameters? Could those cause similar behaviour and any specific configuration parameters you think I should take a look into?

• Daft does spawn child processes for some operations (and as far as I understand, memory usage in these processes would not count towards RSS), but I couldn't find any evidence of this happening in write_parquet, could you confirm whether I should look into this?

• Is there anything else we should be looking into related to either Daft / Kernel level configuration that you know of?

Thanks for taking the time and helping us out here. Really appreciate it!

[colin-ho]

Is is possible its linked to concurrency / target file sizes / other related parameters? Could those cause similar behaviour and any specific configuration parameters you think I should take a look into?

What i think is happening is that the page cache is being populated by the writes, which is why memory_cache and working_set_bytes are increasing. Lowering parameters like parquet_target_file_size or parquet_target_row_group_size will make daft write more frequently, but when those pages get flushed to disk is up to the kernel, not Daft.

Daft does spawn child processes for some operations (and as far as I understand, memory usage in these processes would not count towards RSS), but I couldn't find any evidence of this happening in write_parquet, could you confirm whether I should look into this?

Only in UDFs. Your query doesn't have any UDFs so there should be no other processes.

Is there anything else we should be looking into related to either Daft / Kernel level configuration that you know of?

parquet_target_row_group_size is another configuration you can control for parquet writes.