Memory is not reclaimed with unregister() when inside a transaction

[jprafael]

What happens?

Repeatedly inserting pandas DataFrames via connection.register() inside a single open transaction causes process RSS to grow linearly with each chunk, even after unregister(), del, and gc.collect().

Using SET python_scan_all_frames=true and referencing a local variable instead of explicit use of register()/unregister() avoids the problem.

To Reproduce

The following unmodified script shows the memory growing with each chunk until the process hits >10GB.
Change the USE_REGISTER flag to False to avoid the use of register() which bypasses the problem.

#!/usr/bin/env python

import gc
import os
import tempfile
import shutil
from pathlib import Path

import duckdb
import numpy as np
import pandas as pd
import psutil


CHUNKS = 100
MAX_RSS_GB = 10.0

# ~1 GiB/chunk:
# 8 int64 columns * 16,000,000 rows ~= 1.0 GiB, plus pandas overhead.
ROWS_PER_CHUNK = 16_000_000
COLS = 8

USE_REGISTER = True


def rss_gb() -> float:
    return psutil.Process(os.getpid()).memory_info().rss / (1024 ** 3)


def make_chunk(chunk_idx: int) -> pd.DataFrame:
    base = np.random.randint(
        0,
        1_000_000,
        size=(ROWS_PER_CHUNK, COLS),
        dtype=np.int64,
    )
    return pd.DataFrame(
        base,
        columns=[f"c{i}" for i in range(COLS)],
    )


tmp = Path(tempfile.mkdtemp(prefix="duckdb-register-rss-repro-"))

try:
    con = duckdb.connect(str(tmp / "db.duckdb"))
    
    if not USE_REGISTER:
        con.execute("SET python_scan_all_frames=false;")
    
    con.execute(
        """
        CREATE TABLE t (
            c0 BIGINT,
            c1 BIGINT,
            c2 BIGINT,
            c3 BIGINT,
            c4 BIGINT,
            c5 BIGINT,
            c6 BIGINT,
            c7 BIGINT
        );
        """
    )

    con.begin()

    peak = rss_gb()
    print(f"duckdb={duckdb.__version__} start_rss={peak:.2f}GB", flush=True)

    for i in range(CHUNKS):
        chunk = make_chunk(i)

        if USE_REGISTER:
            con.register("input", chunk)
            con.execute("INSERT INTO t SELECT * FROM input;")
            con.unregister("input")
        else:
            con.execute("INSERT INTO t SELECT * FROM chunk;")


        del chunk
        gc.collect()

        current = rss_gb()
        peak = max(peak, current)
        print(
            f"chunk={i + 1}/{CHUNKS} rss={current:.2f}GB peak={peak:.2f}GB",
            flush=True,
        )

        assert current < MAX_RSS_GB, (
            f"RSS exceeded {MAX_RSS_GB}GB after chunk {i + 1}: "
            f"rss={current:.2f}GB peak={peak:.2f}GB"
        )

    con.rollback()
    print(f"done peak={peak:.2f}GB", flush=True)

finally:
    shutil.rmtree(tmp, ignore_errors=True)

OS:

Ubuntu 24.04LTS

DuckDB Package Version:

1.5.2

Python Version:

3.14

Full Name:

João Pedro Maia Rafael

Affiliation:

Upper Delta, Unipessoal LDA

What is the latest build you tested with? If possible, we recommend testing with the latest nightly build.

I have tested with a nightly build

Did you include all relevant data sets for reproducing the issue?

Not applicable - the reproduction does not require a data set

Did you include all code required to reproduce the issue?

• Yes, I have

Did you include all relevant configuration to reproduce the issue?

• Yes, I have

[jprafael]

The underlying issue is that when the object is register()ed, to make it available for the query planner a view is created with the same name. This view is added to the catalog, and is DROPed when unregister() is invoked. The issue is that when inside a transaction the DROP VIEW doesn't clear the view object from the catalog because it is required to maintain transactionality. That view holds the reference to the data object which prevents its garbage collection.

In contrast, the approach when register() is bypassed doesn't suffer from this. It is PythonReplacementScan that reads the variable from python's variable stack directly.

I can see some alternative solutions for this problem but would love some guidance on which option is more preferable:

1. We keep the CREATE VIEW, but make it so that the view instead stores only a weakref. If the view is re-read after the python object is deleted then an duckdb runtime error is raised.
2. We remove the CREATE VIEW, and make register()/unregister() behave just like an explicit version of PythonReplacementScan. In this format, the two functions simply modify the ClientContext to maintain the a map of name -> python object (outside of transaction semantics). Then PythonReplacementScan is modified to first check this map of names instead of checking python locals.

I'm more inclined on the second option (my mental model of register/unregister was simply an explicit version of python_scan_all_frames), but maybe there are some use cases that require the VIEWs to be created. I didn't find any unit test that makes such assumption explicit if it exists.

[evertlammerts]

Thanks for the report! I agree that this is not behavior you might expect and we should update our docs. But I do think we should keep register() the way it is. It creates a view, and the mental model should be more or less the same as when you would create and drop a (materialized) view repeatedly in a transaction. This would result in accumulating memory in the same way, as long as the transaction is not committed or rolled back. The only real difference is that a Python object lives outside of DuckDB's managed heap so it can't spill etc. But otherwise the semantics line up, and I think they are correct.

For this sort of bulk loading it is indeed better to use a replacement scan rather than register.

[jprafael]

The gripe I have with that is I'd rather not have duckdb magically have access to all python objects in scope and would prefer the option to have explicit/whitelisted access only. Is there a way we can keep register() with its current semantics but also have the option to not create the view? e.g. register(.., create_view=False) or register_unsafe(...)