fix(duckdb): free memtables based on operation lifetime (#10042)

ibis/backends/duckdb/__init__.py

@@ -7,6 +7,7 @@
 import os
 import urllib
 import warnings
+import weakref
 from operator import itemgetter
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -1605,6 +1606,12 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
             # only register if we haven't already done so
             self.con.register(name, op.data.to_pyarrow(op.schema))
 
+            # if we don't aggressively unregister tables duckdb will keep a
+            # reference to every memtable ever registered, even if there's no
+            # way for a user to access the operation anymore, resulting in a
+            # memory leak
+            weakref.finalize(op, self.con.unregister, name)
+
     def _register_udfs(self, expr: ir.Expr) -> None:
         con = self.con
 

ibis/backends/duckdb/tests/test_client.py

@@ -417,3 +417,12 @@ def test_read_csv_with_types(tmp_path, input, all_varchar):
     path.write_bytes(data)
     t = con.read_csv(path, all_varchar=all_varchar, **input)
     assert t.schema()["geom"].is_geospatial()
+
+
+def test_memtable_doesnt_leak(con, monkeypatch):
+    monkeypatch.setattr(ibis.options, "default_backend", con)
+    name = "memtable_doesnt_leak"
+    assert name not in con.list_tables()
+    df = ibis.memtable({"a": [1, 2, 3]}, name=name).execute()
+    assert name not in con.list_tables()
+    assert len(df) == 3