Backport PR #54778 on branch 2.1.x (REGR: Index.union loses python string dtype) (#54813)

Backport PR #54778: REGR: Index.union loses python string dtype

Co-authored-by: Patrick Hoefler <[email protected]>

Author: meeseeksmachine

pandas/core/indexes/base.py

@@ -5204,7 +5204,7 @@ def _from_join_target(self, result: np.ndarray) -> ArrayLike:
         """
         if isinstance(self.values, BaseMaskedArray):
             return type(self.values)(result, np.zeros(result.shape, dtype=np.bool_))
-        elif isinstance(self.values, ArrowExtensionArray):
+        elif isinstance(self.values, (ArrowExtensionArray, StringArray)):
             return type(self.values)._from_sequence(result)
         return result
 

pandas/tests/indexes/test_setops.py

@@ -899,3 +899,10 @@ def test_union_ea_dtypes(self, any_numeric_ea_and_arrow_dtype):
         result = idx.union(idx2)
         expected = Index([1, 2, 3, 4, 5], dtype=any_numeric_ea_and_arrow_dtype)
         tm.assert_index_equal(result, expected)
+
+    def test_union_string_array(self, any_string_dtype):
+        idx1 = Index(["a"], dtype=any_string_dtype)
+        idx2 = Index(["b"], dtype=any_string_dtype)
+        result = idx1.union(idx2)
+        expected = Index(["a", "b"], dtype=any_string_dtype)
+        tm.assert_index_equal(result, expected)