Repository: ibis-project/ibis

Issue:
feat: memtable() should support nullable pandas types
### What happened?

nullable dtypes like pd.Int64Dtype(). Probably fails with other ones too, but I haven't tested.

```python
import pandas as pd
import ibis

df = pd.DataFrame(
    {
        "x": pd.Series([1, 2, None], dtype=pd.Int64Dtype()),
    }
)
ibis.memtable(df)
```

### What version of ibis are you using?

4.0.0

### What backend(s) are you using, if any?

duckdb

### Relevant log output

```sh
stacktrace
---------------------------------------------------------------------------
IbisTypeError                             Traceback (most recent call last)
Cell In[14], line 9
      2 import ibis
      4 df = pd.DataFrame(
      5     {
      6         "x": pd.Series([1, 2, 3, None], dtype=pd.Int64Dtype()),
      7     }
      8 )
----> 9 ibis.memtable(df)

File ~/.pyenv/versions/3.10.4/lib/python3.10/functools.py:889, in singledispatch.<locals>.wrapper(*args, **kw)
    885 if not args:
    886     raise TypeError(f'{funcname} requires at least '
    887                     '1 positional argument')
--> 889 return dispatch(args[0].__class__)(*args, **kw)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/expr/api.py:409, in memtable(data, columns, schema, name)
    403     newcols = getattr(
    404         schema,
    405         "names",
    406         (f"col{i:d}" for i in range(len(cols))),
    407     )
    408     df = df.rename(columns=dict(zip(cols, newcols)))
--> 409 return _memtable_from_dataframe(df, name=name, schema=schema)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/expr/api.py:425, in _memtable_from_dataframe(df, name, schema)
    415 def _memtable_from_dataframe(
    416     df: pd.DataFrame,
    417     *,
    418     name: str | None = None,
    419     schema: SupportsSchema | None = None,
    420 ) -> Table:
    421     from ibis.backends.pandas.client import DataFrameProxy, PandasInMemoryTable
    423     op = PandasInMemoryTable(
    424         name=name if name is not None else next(_gen_memtable_name),
--> 425         schema=sch.infer(df) if schema is None else schema,
    426         data=DataFrameProxy(df),
    427     )
    428     return op.to_expr()

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/multipledispatch/dispatcher.py:278, in Dispatcher.__call__(self, *args, **kwargs)
    276     self._cache[types] = func
    277 try:
--> 278     return func(*args, **kwargs)
    280 except MDNotImplementedError:
    281     funcs = self.dispatch_iter(*types)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/backends/pandas/client.py:81, in infer_pandas_schema(df, schema)
     79         ibis_dtype = dt.dtype(schema[column_name])
     80     else:
---> 81         ibis_dtype = dt.infer(df[column_name]).value_type
     83     pairs.append((column_name, ibis_dtype))
     85 return sch.schema(pairs)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/common/dispatch.py:88, in lazy_singledispatch.<locals>.call(arg, *args, **kwargs)
     86 @functools.wraps(func)
     87 def call(arg, *args, **kwargs):
---> 88     return dispatch(type(arg))(arg, *args, **kwargs)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/expr/datatypes/value.py:233, in infer_pandas_series(value)
    231     value_dtype = _infer_object_array_dtype(value)
    232 else:
--> 233     value_dtype = dt.dtype(value.dtype)
    235 return dt.Array(value_dtype)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/multipledispatch/dispatcher.py:278, in Dispatcher.__call__(self, *args, **kwargs)
    276     self._cache[types] = func
    277 try:
--> 278     return func(*args, **kwargs)
    280 except MDNotImplementedError:
    281     funcs = self.dispatch_iter(*types)

File ~/Library/Application Support/hatch/env/virtual/noatak-UM6-FHel/noatak/lib/python3.10/site-packages/ibis/expr/datatypes/core.py:29, in dtype_from_object(value, **kwargs)
     27 @dtype.register(object)
     28 def dtype_from_object(value, **kwargs) -> DataType:
---> 29     raise IbisTypeError(f'Value {value!r} is not a valid datatype')

IbisTypeError: Value Int64Dtype() is not a valid datatype
```


### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
I believe we can support inferring these types (but not yet returning them from execute) without having to implement #5316.

Resolving patch:
diff --git a/ibis/backends/pandas/client.py b/ibis/backends/pandas/client.py
--- a/ibis/backends/pandas/client.py
+++ b/ibis/backends/pandas/client.py
@@ -56,9 +56,16 @@ def from_pandas_categorical(_):
     return dt.Category()
 
 
-@dt.dtype.register(pd.core.arrays.string_.StringDtype)
-def from_pandas_string(_):
-    return dt.String()
+@dt.dtype.register(pd.core.dtypes.base.ExtensionDtype)
+def from_pandas_extension_dtype(t):
+    return getattr(dt, t.__class__.__name__.replace("Dtype", "").lower())
+
+
+@dt.dtype.register(pd.core.arrays.arrow.dtype.ArrowDtype)
+def from_pandas_arrow_extension_dtype(t):
+    import ibis.backends.pyarrow.datatypes as _  # noqa: F401
+
+    return dt.dtype(t.pyarrow_dtype)
 
 
 @sch.schema.register(pd.Series)


Test patch:
diff --git a/ibis/backends/duckdb/tests/test_register.py b/ibis/backends/duckdb/tests/test_register.py
--- a/ibis/backends/duckdb/tests/test_register.py
+++ b/ibis/backends/duckdb/tests/test_register.py
@@ -2,6 +2,7 @@
 import tempfile
 from pathlib import Path
 
+import duckdb
 import pandas as pd
 import pytest
 import sqlalchemy as sa
@@ -109,3 +110,56 @@ def test_register_sqlite(data_directory):
     path = data_directory / "ibis_testing.db"
     ft = con.register(f"sqlite://{path}", "functional_alltypes")
     assert ft.count().execute()
+
+
+def test_memtable_with_nullable_dtypes():
+    data = pd.DataFrame(
+        {
+            "a": pd.Series(["a", None, "c"], dtype="string"),
+            "b": pd.Series([None, 1, 2], dtype="Int8"),
+            "c": pd.Series([0, None, 2], dtype="Int16"),
+            "d": pd.Series([0, 1, None], dtype="Int32"),
+            "e": pd.Series([None, None, -1], dtype="Int64"),
+            "f": pd.Series([None, 1, 2], dtype="UInt8"),
+            "g": pd.Series([0, None, 2], dtype="UInt16"),
+            "h": pd.Series([0, 1, None], dtype="UInt32"),
+            "i": pd.Series([None, None, 42], dtype="UInt64"),
+            "j": pd.Series([None, False, True], dtype="boolean"),
+        }
+    )
+    expr = ibis.memtable(data)
+    res = expr.execute()
+    assert len(res) == len(data)
+
+
+def test_memtable_with_nullable_pyarrow_string():
+    pytest.importorskip("pyarrow")
+    data = pd.DataFrame({"a": pd.Series(["a", None, "c"], dtype="string[pyarrow]")})
+    expr = ibis.memtable(data)
+    res = expr.execute()
+    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")
+
+    data = pd.DataFrame(
+        {
+            "b": pd.Series([None, 1, 2], dtype="int8[pyarrow]"),
+            "c": pd.Series([0, None, 2], dtype="int16[pyarrow]"),
+            "d": pd.Series([0, 1, None], dtype="int32[pyarrow]"),
+            "e": pd.Series([None, None, -1], dtype="int64[pyarrow]"),
+            "f": pd.Series([None, 1, 2], dtype="uint8[pyarrow]"),
+            "g": pd.Series([0, None, 2], dtype="uint16[pyarrow]"),
+            "h": pd.Series([0, 1, None], dtype="uint32[pyarrow]"),
+            "i": pd.Series([None, None, 42], dtype="uint64[pyarrow]"),
+            "j": pd.Series([None, False, True], dtype="boolean[pyarrow]"),
+        }
+    )
+    expr = ibis.memtable(data)
+    res = expr.execute()
+    assert len(res) == len(data)
diff --git a/ibis/backends/pandas/tests/test_client.py b/ibis/backends/pandas/tests/test_client.py
--- a/ibis/backends/pandas/tests/test_client.py
+++ b/ibis/backends/pandas/tests/test_client.py
@@ -5,6 +5,7 @@
 from pytest import param
 
 import ibis
+import ibis.expr.datatypes as dt
 from ibis.backends.pandas.client import PandasTable
 
 
@@ -98,3 +99,61 @@ def test_datetime64_infer(client, unit):
     expr = ibis.literal(value, type='timestamp')
     result = client.execute(expr)
     assert result == pd.Timestamp(value).to_pydatetime()
+
+
+@pytest.mark.parametrize(
+    ("ext_dtype", "expected"),
+    [
+        (pd.StringDtype(), dt.string),
+        (pd.Int8Dtype(), dt.int8),
+        (pd.Int16Dtype(), dt.int16),
+        (pd.Int32Dtype(), dt.int32),
+        (pd.Int64Dtype(), dt.int64),
+        (pd.UInt8Dtype(), dt.uint8),
+        (pd.UInt16Dtype(), dt.uint16),
+        (pd.UInt32Dtype(), dt.uint32),
+        (pd.UInt64Dtype(), dt.uint64),
+        (pd.BooleanDtype(), dt.boolean),
+    ],
+    ids=str,
+)
+def test_infer_nullable_dtypes(ext_dtype, expected):
+    assert dt.dtype(ext_dtype) == expected
+
+
+@pytest.mark.parametrize(
+    ("arrow_dtype", "expected"),
+    [
+        ("string", dt.string),
+        ("int8", dt.int8),
+        ("int16", dt.int16),
+        ("int32", dt.int32),
+        ("int64", dt.int64),
+        ("uint8", dt.uint8),
+        ("uint16", dt.uint16),
+        ("uint32", dt.uint32),
+        ("uint64", dt.uint64),
+        param(
+            "list<item: string>",
+            dt.Array(dt.string),
+            marks=pytest.mark.xfail(
+                reason="list repr in dtype Series argument doesn't work",
+                raises=TypeError,
+            ),
+            id="list_string",
+        ),
+    ],
+    ids=str,
+)
+def test_infer_pandas_arrow_dtype(arrow_dtype, expected):
+    pytest.importorskip("pyarrow")
+    ser = pd.Series([], dtype=f"{arrow_dtype}[pyarrow]")
+    dtype = ser.dtype
+    assert dt.dtype(dtype) == expected
+
+
+def test_infer_pandas_arrow_list_dtype():
+    pa = pytest.importorskip("pyarrow")
+    ser = pd.Series([], dtype=pd.ArrowDtype(pa.list_(pa.string())))
+    dtype = ser.dtype
+    assert dt.dtype(dtype) == dt.Array(dt.string)