ENH: Arrow backed string array - implement factorize() method without casting to objects

pandas/core/arrays/string_arrow.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 from distutils.version import LooseVersion
-from typing import TYPE_CHECKING, Any, Sequence, Type, Union
+from typing import TYPE_CHECKING, Any, Sequence, Tuple, Type, Union
 
 import numpy as np
 
 from pandas._libs import lib, missing as libmissing
+from pandas.util._decorators import doc
 from pandas.util._validators import validate_fillna_kwargs
 
 from pandas.core.dtypes.base import ExtensionDtype
@@ -15,10 +16,12 @@
 from pandas.api.types import (
  is_array_like,
  is_bool_dtype,
+ is_int64_dtype,
  is_integer,
  is_integer_dtype,
  is_scalar,
 )
+from pandas.core.algorithms import factorize
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays.base import ExtensionArray
 from pandas.core.indexers import check_array_indexer, validate_indices
@@ -252,9 +255,20 @@ def __len__(self) -> int:
  """
  return len(self._data)
 
- @classmethod
- def _from_factorized(cls, values, original):
- return cls._from_sequence(values)
+ @doc(ExtensionArray.factorize)
+ def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]:
+ if self._data.num_chunks == 1:
+ encoded = self._data.chunk(0).dictionary_encode()
+ indices = encoded.indices.to_pandas()
+ if indices.dtype.kind == "f":
+ indices[np.isnan(indices)] = na_sentinel
+ indices = indices.astype(int)
+ if not is_int64_dtype(indices):
+ indices = indices.astype(np.int64)
+ return indices.values, type(self)(encoded.dictionary)
+ else:
+ np_array = self._data.to_pandas().values
+ return factorize(np_array, na_sentinel=na_sentinel)
 
  @classmethod
  def _concat_same_type(cls, to_concat) -> ArrowStringArray: