fix(duckdb): workaround an ownership bug at the interaction of duckdb…

…, pandas and pyarrow

ibis/backends/base/sql/__init__.py

@@ -248,10 +248,10 @@ def execute(
 
         return result
 
-    def _register_in_memory_table(self, table_op):
-        raise NotImplementedError
+    def _register_in_memory_table(self, _: ops.InMemoryTable) -> None:
+        raise NotImplementedError(self.name)
 
-    def _register_in_memory_tables(self, expr):
+    def _register_in_memory_tables(self, expr: ir.Expr) -> None:
         if self.compiler.cheap_in_memory_tables:
             for memtable in an.find_memtables(expr.op()):
                 self._register_in_memory_table(memtable)

ibis/backends/clickhouse/__init__.py

@@ -92,8 +92,8 @@ def __init__(self, *args, external_tables=None, **kwargs):
         super().__init__(*args, **kwargs)
         self._external_tables = external_tables or {}
 
-    def _register_in_memory_table(self, table_op):
-        self._external_tables[table_op.name] = table_op.data.to_frame()
+    def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
+        self._external_tables[op.name] = op.data.to_frame()
 
     def _log(self, sql: str) -> None:
         """Log the SQL, usually to the standard output.

ibis/backends/duckdb/__init__.py

@@ -20,6 +20,8 @@
     import pandas as pd
     import pyarrow as pa
 
+    import ibis.expr.operations as ops
+
 import ibis.expr.schema as sch
 import ibis.expr.types as ir
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
@@ -568,10 +570,13 @@ def _metadata(self, query: str) -> Iterator[tuple[str, dt.DataType]]:
                 ibis_type = parse(type)
                 yield name, ibis_type.copy(nullable=null.lower() == "yes")
 
-    def _register_in_memory_table(self, table_op):
-        df = table_op.data.to_frame()
+    def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
+        # in theory we could use pandas dataframes, but when using dataframes
+        # with pyarrow datatypes later reads of this data segfault
+        name = op.name
+        table = op.data.to_pyarrow()
         with self.begin() as con:
-            con.connection.register(table_op.name, df)
+            con.connection.register(name, table)
 
     def _get_sqla_table(
         self, name: str, schema: str | None = None, **kwargs: Any

ibis/backends/duckdb/tests/test_register.py

@@ -2,7 +2,6 @@
 import tempfile
 from pathlib import Path
 
-import duckdb
 import pandas as pd
 import pytest
 import sqlalchemy as sa
@@ -140,10 +139,6 @@ def test_memtable_with_nullable_pyarrow_string():
     assert len(res) == len(data)
 
 
-@pytest.mark.xfail(
-    raises=duckdb.NotImplementedException,
-    reason="DuckDB only supports the `string[pyarrow]` pandas dtype",
-)
 def test_memtable_with_nullable_pyarrow_not_string():
     pytest.importorskip("pyarrow")
 

ibis/backends/pandas/client.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -45,6 +46,9 @@
     },
 )
 
+if TYPE_CHECKING:
+    import pyarrow as pa
+
 
 @dt.dtype.register(DatetimeTZDtype)
 def from_pandas_tzdtype(value):
@@ -196,20 +200,25 @@ def try_json(x):
 class DataFrameProxy(Immutable, util.ToFrame):
     __slots__ = ('_df', '_hash')
 
-    def __init__(self, df):
+    def __init__(self, df: pd.DataFrame) -> None:
         object.__setattr__(self, "_df", df)
         object.__setattr__(self, "_hash", hash((type(df), id(df))))
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return self._hash
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         df_repr = util.indent(repr(self._df), spaces=2)
         return f"{self.__class__.__name__}:\n{df_repr}"
 
-    def to_frame(self):
+    def to_frame(self) -> pd.DataFrame:
         return self._df
 
+    def to_pyarrow(self) -> pa.Table:
+        import pyarrow as pa
+
+        return pa.Table.from_pandas(self._df)
+
 
 class PandasInMemoryTable(ops.InMemoryTable):
     data = rlz.instance_of(DataFrameProxy)

ibis/backends/pyspark/__init__.py

@@ -468,9 +468,8 @@ def create_table(
 
         return self.raw_sql(statement.compile())
 
-    def _register_in_memory_table(self, table_op):
-        spark_df = self.compile(table_op.to_expr())
-        spark_df.createOrReplaceTempView(table_op.name)
+    def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
+        self.compile(op.to_expr()).createOrReplaceTempView(op.name)
 
     def create_view(
         self,

ibis/util.py

@@ -32,6 +32,7 @@
     from pathlib import Path
 
     import pandas as pd
+    import pyarrow as pa
 
     import ibis.expr.operations as ops
 
@@ -506,12 +507,19 @@ def experimental(func):
 
 
 class ToFrame(abc.ABC):
-    """Interface for in-memory objects that can be converted to a DataFrame."""
+    """Interface for in-memory objects that can be converted to an in-memory structure.
+
+    Supports pandas DataFrames and PyArrow Tables.
+    """
 
     __slots__ = ()
 
     @abc.abstractmethod
-    def to_frame(self) -> pd.DataFrame:
+    def to_frame(self) -> pd.DataFrame:  # pragma: no cover
+        ...
+
+    @abc.abstractmethod
+    def to_pyarrow(self) -> pa.Table:  # pragma: no cover
         ...