[Python] Extension array data type should default to the storage type if to_pandas_dtype is not implemented

[AlenkaF]

Describe the bug, including details regarding any error messages, version, and platform.

When working on the extension type for tensors in PyArrow I came across a behaviour of the conversion to pandas that could be improved.

Creating an extension array (fixed shape tensor in this case) and converting it to pandas works well

>>> arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
>>> storage = pa.array(arr, pa.list_(pa.int32(), 4))
>>> tensor = pa.ExtensionArray.from_storage(tensor_type, storage)
>>> tensor.to_pandas()
0            [1, 2, 3, 4]
1        [10, 20, 30, 40]
2    [100, 200, 300, 400]
dtype: object

But creating a table with an extension array and then converting it to pandas fails:

>>> data = [
...     pa.array([1, 2, 3]),
...     pa.array(['foo', 'bar', None]),
...     pa.array([True, None, True]),
...     tensor
... ]
>>> my_schema = pa.schema([('f0', pa.int8()),
...                        ('f1', pa.string()),
...                        ('f2', pa.bool_()),
...                        ('tensors_int', tensor_type)])
>>> table = pa.Table.from_arrays(data, schema=my_schema)
>>> table.to_pandas()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "pyarrow/array.pxi", line 830, in pyarrow.lib._PandasConvertible.to_pandas
    return self._to_pandas(options, categories=categories,
  File "pyarrow/table.pxi", line 4004, in pyarrow.lib.Table._to_pandas
    mgr = table_to_blockmanager(
  File "/Users/alenkafrim/repos/arrow/python/pyarrow/pandas_compat.py", line 820, in table_to_blockmanager
    blocks = _table_to_blocks(options, table, categories, ext_columns_dtypes)
  File "/Users/alenkafrim/repos/arrow/python/pyarrow/pandas_compat.py", line 1171, in _table_to_blocks
    return [_reconstruct_block(item, columns, extension_columns)
  File "/Users/alenkafrim/repos/arrow/python/pyarrow/pandas_compat.py", line 1171, in <listcomp>
    return [_reconstruct_block(item, columns, extension_columns)
  File "/Users/alenkafrim/repos/arrow/python/pyarrow/pandas_compat.py", line 776, in _reconstruct_block
    pandas_dtype = extension_columns[name]
KeyError: 'tensors_int'

The issue is due to the extension array in this example not having to_pandas_dtype method implemented. In this case ext_columns does not get populated in _get_extension_dtypes method with the name of the column with an extension type:

arrow/python/pyarrow/pandas_compat.py

Lines 870 to 879 in 0368e41

	# infer from extension type in the schema
	for field in table.schema:
	typ = field.type
	if isinstance(typ, pa.BaseExtensionType):
	try:
	pandas_dtype = typ.to_pandas_dtype()
	except NotImplementedError:
	pass
	else:
	ext_columns[field.name] = pandas_dtype

It would be good if it would, in case to_pandas_dtype method is not implemented, convert the storage array

arrow/python/pyarrow/pandas_compat.py

Line 776 in 0368e41

pandas_dtype = extension_columns[name]

similar to

arrow/python/pyarrow/array.pxi

Lines 2888 to 2889 in 925cbd8

	# otherwise convert the storage array with the base implementation
	return Array._to_pandas(self.storage, options, **kwargs)

Component(s)

Python