WIP frame addition with fill_value test

Author: eicchen

pandas/tests/frame/test_arithmetic.py

@@ -626,9 +626,29 @@ def test_arith_flex_frame_corner(self, float_frame):
         expected = float_frame.sort_index() * np.nan
         tm.assert_frame_equal(result, expected)
 
-        res = float_frame.add(float_frame.iloc[0], fill_value=3)
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_arith_flex_frame_fill_value_corner(self, float_frame, axis):
+        rng = np.random.default_rng(60)
+        mask = rng.random(float_frame.shape) < 0.2
+        left = float_frame.mask(mask)
+        right = left.iloc[0].value
 
-        res = float_frame.add(float_frame.iloc[0], axis="index", fill_value=3)
+        result = left.add(right, axis=axis, fill_value=3)
+
+        if axis == 0:  # axis = index, vertical
+            expected = result * np.nan
+
+        else:  # axis = columns, horizontal
+            right_df = DataFrame([right] * result.shape[0], index=result.index)
+            left_filled = left.fillna(3)
+            right_filled = right_df.fillna(3)
+            expected = right_filled + left_filled
+            expected = expected.mask(expected == 6, pd.NA)
+
+        print(result)
+        print(expected)
+
+        tm.assert_frame_equal(result, expected)
 
     @pytest.mark.parametrize("op", ["add", "sub", "mul", "mod"])
     def test_arith_flex_series_ops(self, simple_frame, op):
@@ -2190,3 +2210,61 @@ def test_mixed_col_index_dtype(string_dtype_no_object):
     expected.columns = expected.columns.astype(string_dtype_no_object)
 
     tm.assert_frame_equal(result, expected)
+
+
+dt_params = [
+    (tm.ALL_INT_NUMPY_DTYPES[0], 5),
+    (tm.ALL_INT_EA_DTYPES[0], 5),
+    (tm.FLOAT_NUMPY_DTYPES[0], 4.9),
+    (tm.FLOAT_EA_DTYPES[0], 4.9),
+]
+
+axes = [0, 1]
+
+
+@pytest.mark.parametrize(
+    "data_type,fill_val, axis",
+    [(dt, val, axis) for axis in axes for dt, val in dt_params],
+)
+def test_df_fill_value_dtype(data_type, fill_val, axis):
+    # GH 61581
+    base_data = np.arange(25).reshape(5, 5)
+    mult_list = [1, np.nan, 5, np.nan, 3]
+    np_int_flag = 0
+
+    try:
+        mult_data = pd.array(mult_list, dtype=data_type)
+    except ValueError as e:
+        # Numpy int type cannot represent NaN, it will end up here
+        if "cannot convert float NaN to integer" in str(e):
+            mult_data = np.asarray(mult_list)
+            np_int_flag = 1
+
+    columns = list("ABCDE")
+    df = DataFrame(base_data, columns=columns)
+
+    for i in range(df.shape[0]):
+        try:
+            df.iat[i, i] = np.nan
+            df.iat[i + 1, i] = pd.NA
+            df.iat[i + 3, i] = pd.NA
+        except IndexError:
+            pass
+
+    mult_mat = np.broadcast_to(mult_data, df.shape)
+    if axis == 0:
+        mask = np.isnan(mult_mat).T
+    else:
+        mask = np.isnan(mult_mat)
+    mask = df.isna().values & mask
+
+    df_result = df.mul(mult_data, axis=axis, fill_value=fill_val)
+    if np_int_flag == 1:
+        mult_np = np.nan_to_num(mult_data, nan=fill_val)
+        df_expected = (df.fillna(fill_val).mul(mult_np, axis=axis)).mask(mask, np.nan)
+    else:
+        df_expected = (
+            df.fillna(fill_val).mul(mult_data.fillna(fill_val), axis=axis)
+        ).mask(mask, np.nan)
+
+    tm.assert_frame_equal(df_result, df_expected)