Skip to content

Commit 1ebcf6e

Browse files
committed
Added test case to test EA and NUMPY dtypes
1 parent be09616 commit 1ebcf6e

File tree

1 file changed

+33
-27
lines changed

1 file changed

+33
-27
lines changed

pandas/tests/frame/test_arithmetic.py

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2207,30 +2207,36 @@ def test_df_fill_value_operations(op):
22072207
tm.assert_frame_equal(df_result, df_expected)
22082208

22092209

2210-
# ! Currently implementing
2211-
# @pytest.mark.parametrize("input_data, fill_val",
2212-
# [
2213-
# (np.arange(50).reshape(10, 5), 5), #Numpy
2214-
# (pd.array(np.random.choice([True, False], size=(10, 5)),
2215-
# dtype="boolean"), True),
2216-
# ]
2217-
# )
2218-
# def test_df_fill_value_dtype(input_data, fill_val):
2219-
# # GH 61581
2220-
# columns = list("ABCDE")
2221-
# df = DataFrame(input_data, columns=columns)
2222-
# for i in range(5):
2223-
# df.iat[i, i] = np.nan
2224-
# df.iat[i + 1, i] = np.nan
2225-
# df.iat[i + 4, i] = np.nan
2226-
2227-
# df_base = df.iloc[:, :-1]
2228-
# df_mult = df.iloc[:, -1]
2229-
# mask = df.isna().values
2230-
# mask = mask[:, :-1] & mask[:, [-1]]
2231-
2232-
# df_result = df_base.mul(df_mult, axis=0, fill_value=fill_val)
2233-
# df_expected = (df_base.fillna(fill_val).mul(df_mult.fillna(fill_val),
2234-
# axis=0)).mask(mask, np.nan)
2235-
2236-
# tm.assert_frame_equal(df_result, df_expected)
2210+
dt_params = [
2211+
(tm.ALL_INT_NUMPY_DTYPES, 5),
2212+
(tm.ALL_INT_EA_DTYPES, 5),
2213+
(tm.FLOAT_NUMPY_DTYPES, 4.9),
2214+
(tm.FLOAT_EA_DTYPES, 4.9),
2215+
]
2216+
2217+
dt_param_flat = [(dt, val) for lst, val in dt_params for dt in lst]
2218+
2219+
2220+
@pytest.mark.parametrize("data_type, fill_val", dt_param_flat)
2221+
def test_df_fill_value_dtype(data_type, fill_val):
2222+
# GH 61581
2223+
base_data = np.arange(50).reshape(10, 5)
2224+
df_data = pd.array(base_data, dtype=data_type)
2225+
columns = list("ABCDE")
2226+
df = DataFrame(df_data, columns=columns)
2227+
for i in range(5):
2228+
df.iat[i, i] = np.nan
2229+
df.iat[i + 1, i] = pd.NA
2230+
df.iat[i + 4, i] = pd.NA
2231+
2232+
df_base = df.iloc[:, :-1]
2233+
df_mult = df.iloc[:, -1]
2234+
mask = df.isna().values
2235+
mask = mask[:, :-1] & mask[:, [-1]]
2236+
2237+
df_result = df_base.mul(df_mult, axis=0, fill_value=fill_val)
2238+
df_expected = (df_base.fillna(fill_val).mul(df_mult.fillna(fill_val), axis=0)).mask(
2239+
mask, np.nan
2240+
)
2241+
2242+
tm.assert_frame_equal(df_result, df_expected)

0 commit comments

Comments
 (0)