Fixed bug when creating new column with missing values when setting a single string value

Author: MarcoGorelli

doc/source/whatsnew/v2.1.4.rst

@@ -13,7 +13,7 @@ including other versions of pandas.
 
 Fixed regressions
 ~~~~~~~~~~~~~~~~~
--
+- Fixed bug when creating new column with missing values when setting a single string value (:issue:`56204`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/core/dtypes/missing.py

@@ -624,7 +624,7 @@ def array_equals(left: ArrayLike, right: ArrayLike) -> bool:
         return array_equivalent(left, right, dtype_equal=True)
 
 
-def infer_fill_value(val):
+def infer_fill_value(val, index: Index = None):
     """
     infer the fill value for the nan/NaT from the provided
     scalar/ndarray/list-like if we are a NaT, return the correct dtyped
@@ -643,8 +643,13 @@ def infer_fill_value(val):
             return np.array("NaT", dtype=TD64NS_DTYPE)
         return np.array(np.nan, dtype=object)
     elif val.dtype.kind == "U":
-        return np.array(np.nan, dtype=object)
-        # return np.array(np.nan, dtype=val.dtype)
+        if get_option("future.infer_string"):
+            from pandas import Series
+
+            return Series(
+                [np.nan] * len(index), dtype="string[pyarrow_numpy]", index=index
+            )
+        return None
     return np.nan
 
 

pandas/core/indexing.py

@@ -1879,7 +1879,9 @@ def _setitem_with_indexer(self, indexer, value, name: str = "iloc"):
 
                         else:
                             # FIXME: GH#42099#issuecomment-864326014
-                            self.obj[key] = infer_fill_value(value)
+                            self.obj[key] = infer_fill_value(
+                                value, index=self.obj.index
+                            )
 
                         new_indexer = convert_from_missing_indexer_tuple(
                             indexer, self.obj.axes

pandas/tests/frame/indexing/test_indexing.py

@@ -1922,6 +1922,28 @@ def test_adding_new_conditional_column() -> None:
     tm.assert_frame_equal(df, expected)
 
 
+def test_adding_new_conditional_column_with_string() -> None:
+    # https://github.com/pandas-dev/pandas/issues/56204
+    df = DataFrame({"a": [1, 2], "b": [3, 4]})
+    df.loc[lambda x: x.a == 1, "c"] = "1"
+    expected = DataFrame({"a": [1, 2], "b": [3, 4], "c": ["1", None]}).astype(
+        {"a": "int64", "b": "int64", "c": "object"}
+    )
+    tm.assert_frame_equal(df, expected)
+
+
+def test_adding_new_conditional_column_with_infer_string() -> None:
+    # https://github.com/pandas-dev/pandas/issues/56204
+    pytest.importorskip("pyarrow")
+    df = DataFrame({"a": [1, 2], "b": [3, 4]})
+    with pd.option_context("future.infer_string", True):
+        df.loc[lambda x: x.a == 1, "c"] = "1"
+    expected = DataFrame({"a": [1, 2], "b": [3, 4], "c": ["1", None]}).astype(
+        {"a": "int64", "b": "int64", "c": "string[pyarrow_numpy]"}
+    )
+    tm.assert_frame_equal(df, expected)
+
+
 def test_add_new_column_infer_string():
     # GH#55366
     pytest.importorskip("pyarrow")