Implement pd.StringDtype handling

python/interpret-core/interpret/utils/_clean_x.py

@@ -1036,6 +1036,24 @@ def _process_pandas_column(X_col, is_initial, feature_type, min_unique_continuou
             X_col.codes,
             None,
         )
+    elif isinstance(dt, pd.StringDtype):
+        # this handles pd.StringDtype both the numpy and arrow versions
+        if X_col.hasnans:
+            # if hasnans is true then there is definetly a real missing value in there and not just a mask
+            return _process_ndarray(
+                X_col.dropna().values.astype(np.str_, copy=False),
+                X_col.notna().values,
+                is_initial,
+                feature_type,
+                min_unique_continuous,
+            )
+        return _process_ndarray(
+            X_col.values.astype(np.str_, copy=False),
+            None,
+            is_initial,
+            feature_type,
+            min_unique_continuous,
+        )
     elif issubclass(tt, _intbool_types):
         # this handles Int8Dtype to Int64Dtype, UInt8Dtype to UInt64Dtype, and BooleanDtype
         if X_col.hasnans:
@@ -1058,8 +1076,6 @@ def _process_pandas_column(X_col, is_initial, feature_type, min_unique_continuou
         )
 
     # TODO: implement pd.SparseDtype
-    # TODO: implement pd.StringDtype both the numpy and arrow versions
-    # https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.StringDtype.html#pandas.StringDtype
     msg = f"{type(dt)} not supported"
     _log.error(msg)
     raise TypeError(msg)

python/interpret-core/tests/utils/test_clean_x.py

@@ -1231,6 +1231,10 @@ def test_unify_columns_pandas_missings_BooleanDtype():
     check_pandas_missings(pd.BooleanDtype(), False, True)
 
 
+def test_unify_columns_pandas_missings_StringDtype():
+    check_pandas_missings(pd.StringDtype(), "abc", "def")
+
+
 def test_unify_columns_pandas_missings_str():
     check_pandas_missings(np.object_, "abc", "def")