Skip to content

Commit 8582104

Browse files
committed
Add more tests
1 parent d6ca465 commit 8582104

File tree

1 file changed

+13
-3
lines changed

1 file changed

+13
-3
lines changed

python/tests/test_dataframe.py

+13-3
Original file line numberDiff line numberDiff line change
@@ -1297,11 +1297,21 @@ def test_fill_nan(df):
12971297
df_with_nans = df.with_columns(
12981298
literal(float("nan")).alias("d").cast(pa.float64()),
12991299
literal(float("nan")).alias("e").cast(pa.float64()),
1300+
literal("abc").alias("f").cast(pa.string()), # non-numeric column
13001301
)
1301-
df_filled = df_with_nans.fill_nan(0, subset=["d", "e"])
1302+
df_filled = df_with_nans.fill_nan(0, subset=["d", "e", "f"])
13021303
result = df_filled.to_pydict()
1303-
assert result["d"] == [0, 0, 0]
1304-
assert result["e"] == [0, 0, 0]
1304+
assert result["d"] == [0, 0, 0] # succeeds
1305+
assert result["e"] == [0, 0, 0] # succeeds
1306+
assert result["f"] == ["abc", "abc", "abc"] # skipped because not numeric
1307+
1308+
# Test filling NaNs fails on non-numeric columns
1309+
df_with_mixed = df.with_columns(
1310+
literal(float("nan")).alias("d").cast(pa.float64()),
1311+
literal("abc").alias("e").cast(pa.string()),
1312+
)
1313+
with pytest.raises(ValueError, match="Column 'e' is not a numeric column"):
1314+
df_with_mixed.fill_nan(0, subset=["d", "e"])
13051315

13061316
# Test filling NaNs with subset of columns where all casts succeed
13071317
df_with_nans = df.with_columns(

0 commit comments

Comments
 (0)