diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index dca99af81204..6e9dc46a9a3d 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -290,27 +290,34 @@ def array(object obj, type=None, mask=None, size=None, from_pandas=None, indices = _codes_to_indices( values.codes, mask, index_type, memory_pool) - try: - dictionary = array( - values.categories.values, type=value_type, - memory_pool=memory_pool) - except TypeError: - # TODO when removing the deprecation warning, this whole - # try/except can be removed (to bubble the TypeError of - # the first array(..) call) - if value_type is not None: - warnings.warn( - "The dtype of the 'categories' of the passed " - "categorical values ({0}) does not match the " - "specified type ({1}). For now ignoring the specified " - "type, but in the future this mismatch will raise a " - "TypeError".format( - values.categories.dtype, value_type), - FutureWarning, stacklevel=2) + + if hasattr(values.categories.values, '__arrow_array__'): + dictionary = _handle_arrow_array_protocol( + values.categories.values, type, mask, size) + if isinstance(dictionary, ChunkedArray): + dictionary = dictionary.combine_chunks() + else: + try: dictionary = array( - values.categories.values, memory_pool=memory_pool) - else: - raise + values.categories.values, type=value_type, + memory_pool=memory_pool) + except TypeError: + # TODO when removing the deprecation warning, this whole + # try/except can be removed (to bubble the TypeError of + # the first array(..) call) + if value_type is not None: + warnings.warn( + "The dtype of the 'categories' of the passed " + "categorical values ({0}) does not match the " + "specified type ({1}). For now ignoring the specified " + "type, but in the future this mismatch will raise a " + "TypeError".format( + values.categories.dtype, value_type), + FutureWarning, stacklevel=2) + dictionary = array( + values.categories.values, memory_pool=memory_pool) + else: + raise return DictionaryArray.from_arrays( indices, dictionary, ordered=values.ordered, safe=safe) diff --git a/python/pyarrow/tests/test_pandas.py b/python/pyarrow/tests/test_pandas.py index 4d0ddf875474..e3d3d044f242 100644 --- a/python/pyarrow/tests/test_pandas.py +++ b/python/pyarrow/tests/test_pandas.py @@ -1728,6 +1728,25 @@ def test_table_str_to_categorical_with_na(self): table.to_pandas(strings_to_categorical=True, zero_copy_only=True) + def test_categorical_with_arrow_backed_ea_to_array(self): + if Version(pd.__version__) < Version("1.5.0"): + pytest.skip("ArrowDtype missing") + + ser = pd.Series([1, 2, 1], dtype=pd.ArrowDtype(pa.int64())).astype("category") + result = pa.array(ser) + expected = pa.DictionaryArray.from_arrays( + pa.array([0, 1, 0], type=pa.int8()), pa.array([1, 2])) + assert result.equals(expected) + + ser = pd.Series( + ["a", "b", "a"], + dtype=pd.ArrowDtype(pa.string()), + ).astype("category") + result = pa.array(ser) + expected = pa.DictionaryArray.from_arrays( + pa.array([0, 1, 0], type=pa.int8()), pa.array(["a", "b"])) + assert result.equals(expected) + # Regression test for ARROW-2101 def test_array_of_bytes_to_strings(self): converted = pa.array(np.array([b'x'], dtype=object), pa.string())