refactor(backends): clean up memtable resources

ibis/backends/__init__.py

@@ -1116,13 +1116,20 @@ def _register_in_memory_tables(self, expr: ir.Expr) -> None:
         for memtable in expr.op().find(ops.InMemoryTable):
             if not self._table_exists(memtable.name):
                 self._register_in_memory_table(memtable)
+                self._register_memtable_finalizer(memtable)
 
     def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         if self.supports_in_memory_tables:
             raise NotImplementedError(
                 f"{self.name} must implement `_register_in_memory_table` to support in-memory tables"
             )
 
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable) -> None:
+        if self.supports_in_memory_tables:
+            raise NotImplementedError(
+                f"{self.name} must implement `_register_memtable_finalizer` to support in-memory tables"
+            )
+
     def _run_pre_execute_hooks(self, expr: ir.Expr) -> None:
         """Backend-specific hooks to run before an expression is executed."""
         self._register_udfs(expr)

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
@@ -158,12 +159,9 @@ def create_table(
             properties.append(sge.TemporaryProperty())
             catalog = "temp"
 
-        temp_memtable_view = None
-
         if obj is not None:
             if not isinstance(obj, ir.Expr):
                 table = ibis.memtable(obj)
-                temp_memtable_view = table.op().name
             else:
                 table = obj
 
@@ -242,9 +240,6 @@ def create_table(
                         ).sql(self.name)
                     )
 
-        if temp_memtable_view is not None:
-            self.con.unregister(temp_memtable_view)
-
         return self.table(name, database=(catalog, database))
 
     def table(
@@ -1601,6 +1596,11 @@ def _table_exists(self, name: str) -> bool:
     def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         self.con.register(op.name, op.data.to_pyarrow(op.schema))
 
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable):
+        # we can't use drop_table, because self.con.register creates a view, so
+        # use the corresponding unregister method
+        weakref.finalize(op, self.con.unregister, op.name)
+
     def _register_udfs(self, expr: ir.Expr) -> None:
         con = self.con
 

ibis/backends/exasol/__init__.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-import atexit
 import contextlib
 import datetime
 import re
+import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote_plus
 
@@ -293,12 +293,14 @@ def process_item(item: Any):
         with self._safe_raw_sql(create_stmt_sql):
             if not df.empty:
                 self.con.ext.insert_multi(name, rows)
-        atexit.register(self._clean_up_tmp_table, ident)
 
-    def _clean_up_tmp_table(self, ident: sge.Identifier) -> None:
-        with self._safe_raw_sql(
-            sge.Drop(kind="TABLE", this=ident, exists=True, cascade=True)
-        ):
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable):
+        weakref.finalize(op, self._clean_up_tmp_table, op.name)
+
+    def _clean_up_tmp_table(self, name: str) -> None:
+        ident = sg.to_identifier(name, quoted=self.compiler.quoted)
+        sql = sge.Drop(kind="TABLE", this=ident, exists=True, cascade=True)
+        with self._safe_raw_sql(sql):
             pass
 
     def create_table(
@@ -349,11 +351,9 @@ def create_table(
 
         quoted = self.compiler.quoted
 
-        temp_memtable_view = None
         if obj is not None:
             if not isinstance(obj, ir.Expr):
                 table = ibis.memtable(obj)
-                temp_memtable_view = table.op().name
             else:
                 table = obj
 
@@ -382,31 +382,27 @@ def create_table(
         else:
             temp_name = name
 
-        table = sg.table(temp_name, catalog=database, quoted=quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        table_expr = sg.table(temp_name, catalog=database, quoted=quoted)
+        target = sge.Schema(this=table_expr, expressions=column_defs)
 
         create_stmt = sge.Create(kind="TABLE", this=target)
 
         this = sg.table(name, catalog=database, quoted=quoted)
         with self._safe_raw_sql(create_stmt):
             if query is not None:
                 self.con.execute(
-                    sge.Insert(this=table, expression=query).sql(self.name)
+                    sge.Insert(this=table_expr, expression=query).sql(self.name)
                 )
 
             if overwrite:
                 self.con.execute(
                     sge.Drop(kind="TABLE", this=this, exists=True).sql(self.name)
                 )
                 self.con.execute(
-                    f"RENAME TABLE {table.sql(self.name)} TO {this.sql(self.name)}"
+                    f"RENAME TABLE {table_expr.sql(self.name)} TO {this.sql(self.name)}"
                 )
 
         if schema is None:
-            # Clean up temporary memtable if we've created one
-            # for in-memory reads
-            if temp_memtable_view is not None:
-                self.drop_table(temp_memtable_view)
             return self.table(name, database=database)
 
         # preserve the input schema if it was provided

ibis/backends/mssql/__init__.py

@@ -625,11 +625,9 @@ def create_table(
             properties.append(sge.TemporaryProperty())
             catalog, db = None, None
 
-        temp_memtable_view = None
         if obj is not None:
             if not isinstance(obj, ir.Expr):
                 table = ibis.memtable(obj)
-                temp_memtable_view = table.op().name
             else:
                 table = obj
 
@@ -657,19 +655,22 @@ def create_table(
         else:
             temp_name = name
 
-        table = sg.table(
-            "#" * temp + temp_name, catalog=catalog, db=db, quoted=self.compiler.quoted
-        )
+        quoted = self.compiler.quoted
         raw_table = sg.table(temp_name, catalog=catalog, db=db, quoted=False)
-        target = sge.Schema(this=table, expressions=column_defs)
+        target = sge.Schema(
+            this=sg.table(
+                "#" * temp + temp_name, catalog=catalog, db=db, quoted=quoted
+            ),
+            expressions=column_defs,
+        )
 
         create_stmt = sge.Create(
             kind="TABLE",
             this=target,
             properties=sge.Properties(expressions=properties),
         )
 
-        this = sg.table(name, catalog=catalog, db=db, quoted=self.compiler.quoted)
+        this = sg.table(name, catalog=catalog, db=db, quoted=quoted)
         raw_this = sg.table(name, catalog=catalog, db=db, quoted=False)
         with self._safe_ddl(create_stmt) as cur:
             if query is not None:
@@ -702,10 +703,6 @@ def create_table(
             db = "dbo"
 
         if schema is None:
-            # Clean up temporary memtable if we've created one
-            # for in-memory reads
-            if temp_memtable_view is not None:
-                self.drop_table(temp_memtable_view)
             return self.table(name, database=(catalog, db))
 
         # preserve the input schema if it was provided

ibis/backends/mysql/__init__.py

@@ -403,11 +403,9 @@ def create_table(
         if temp:
             properties.append(sge.TemporaryProperty())
 
-        temp_memtable_view = None
         if obj is not None:
             if not isinstance(obj, ir.Expr):
                 table = ibis.memtable(obj)
-                temp_memtable_view = table.op().name
             else:
                 table = obj
 
@@ -435,8 +433,8 @@ def create_table(
         else:
             temp_name = name
 
-        table = sg.table(temp_name, catalog=database, quoted=self.compiler.quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        table_expr = sg.table(temp_name, catalog=database, quoted=self.compiler.quoted)
+        target = sge.Schema(this=table_expr, expressions=column_defs)
 
         create_stmt = sge.Create(
             kind="TABLE",
@@ -447,23 +445,20 @@ def create_table(
         this = sg.table(name, catalog=database, quoted=self.compiler.quoted)
         with self._safe_raw_sql(create_stmt) as cur:
             if query is not None:
-                insert_stmt = sge.Insert(this=table, expression=query).sql(self.name)
+                insert_stmt = sge.Insert(this=table_expr, expression=query).sql(
+                    self.name
+                )
                 cur.execute(insert_stmt)
 
             if overwrite:
                 cur.execute(
                     sge.Drop(kind="TABLE", this=this, exists=True).sql(self.name)
                 )
                 cur.execute(
-                    f"ALTER TABLE IF EXISTS {table.sql(self.name)} RENAME TO {this.sql(self.name)}"
+                    f"ALTER TABLE IF EXISTS {table_expr.sql(self.name)} RENAME TO {this.sql(self.name)}"
                 )
 
         if schema is None:
-            # Clean up temporary memtable if we've created one
-            # for in-memory reads
-            if temp_memtable_view is not None:
-                self.drop_table(temp_memtable_view)
-
             return self.table(name, database=database)
 
         # preserve the input schema if it was provided

ibis/backends/oracle/__init__.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-import atexit
 import contextlib
 import re
 import warnings
+import weakref
 from functools import cached_property
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any
@@ -419,11 +419,9 @@ def create_table(
         if temp:
             properties.append(sge.TemporaryProperty())
 
-        temp_memtable_view = None
         if obj is not None:
             if not isinstance(obj, ir.Expr):
                 table = ibis.memtable(obj)
-                temp_memtable_view = table.op().name
             else:
                 table = obj
 
@@ -478,10 +476,6 @@ def create_table(
                 )
 
         if schema is None:
-            # Clean up temporary memtable if we've created one
-            # for in-memory reads
-            if temp_memtable_view is not None:
-                self.drop_table(temp_memtable_view)
             return self.table(name, database=database)
 
         # preserve the input schema if it was provided
@@ -551,7 +545,8 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
                     insert_stmt, list(data.iloc[start:end].itertuples(index=False))
                 )
 
-        atexit.register(self._clean_up_tmp_table, name)
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable):
+        weakref.finalize(op, self._clean_up_tmp_table, op.name)
 
     def _get_schema_using_query(self, query: str) -> sch.Schema:
         name = util.gen_name("oracle_metadata")
@@ -632,6 +627,13 @@ def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
         return OraclePandasData.convert_table(df, schema)
 
     def _clean_up_tmp_table(self, name: str) -> None:
+        dialect = self.dialect
+
+        ident = sg.to_identifier(name, quoted=self.compiler.quoted)
+
+        truncate = sge.TruncateTable(expressions=[ident]).sql(dialect)
+        drop = sge.Drop(kind="TABLE", this=ident).sql(dialect)
+
         with self.begin() as bind:
             # global temporary tables cannot be dropped without first truncating them
             #
@@ -640,9 +642,9 @@ def _clean_up_tmp_table(self, name: str) -> None:
             # ignore DatabaseError exceptions because the table may not exist
             # because it's already been deleted
             with contextlib.suppress(oracledb.DatabaseError):
-                bind.execute(f'TRUNCATE TABLE "{name}"')
+                bind.execute(truncate)
             with contextlib.suppress(oracledb.DatabaseError):
-                bind.execute(f'DROP TABLE "{name}"')
+                bind.execute(drop)
 
-    def _drop_cached_table(self, name):
+    def _drop_cached_table(self, name: str) -> None:
         self._clean_up_tmp_table(name)

ibis/backends/pandas/__init__.py

@@ -331,6 +331,9 @@ def execute(self, query, params=None, limit="default", **kwargs):
     def _create_cached_table(self, name, expr):
         return self.create_table(name, expr.execute())
 
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable):
+        """No-op, let Python handle clean up."""
+
 
 @lazy_singledispatch
 def _convert_object(obj: Any, _conn):

ibis/backends/polars/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import weakref
 from collections.abc import Iterable, Mapping
 from functools import lru_cache
 from pathlib import Path
@@ -75,9 +76,14 @@ def table(self, name: str) -> ir.Table:
         schema = sch.infer(self._tables[name])
         return ops.DatabaseTable(name, schema, self).to_expr()
 
+    def _table_exists(self, name: str) -> bool:
+        return name in self._tables
+
     def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
-        if (name := op.name) not in self._tables:
-            self._add_table(name, op.data.to_polars(op.schema).lazy())
+        self._add_table(op.name, op.data.to_polars(op.schema).lazy())
+
+    def _register_memtable_finalizer(self, op: ops.InMemoryTable) -> None:
+        weakref.finalize(op, self.drop_table, op.name, force=True)
 
     @deprecated(
         as_of="9.1",

ibis/backends/postgres/__init__.py

@@ -701,8 +701,8 @@ def create_table(
         else:
             temp_name = name
 
-        table = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        table_expr = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
+        target = sge.Schema(this=table_expr, expressions=column_defs)
 
         create_stmt = sge.Create(
             kind="TABLE",
@@ -713,15 +713,17 @@ def create_table(
         this = sg.table(name, catalog=database, quoted=self.compiler.quoted)
         with self._safe_raw_sql(create_stmt) as cur:
             if query is not None:
-                insert_stmt = sge.Insert(this=table, expression=query).sql(self.dialect)
+                insert_stmt = sge.Insert(this=table_expr, expression=query).sql(
+                    self.dialect
+                )
                 cur.execute(insert_stmt)
 
             if overwrite:
                 cur.execute(
                     sge.Drop(kind="TABLE", this=this, exists=True).sql(self.dialect)
                 )
                 cur.execute(
-                    f"ALTER TABLE IF EXISTS {table.sql(self.dialect)} RENAME TO {this.sql(self.dialect)}"
+                    f"ALTER TABLE IF EXISTS {table_expr.sql(self.dialect)} RENAME TO {this.sql(self.dialect)}"
                 )
 
         if schema is None: