diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 6fd6fc23c2..ce9c22132b 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -229,6 +229,10 @@ def assign_constant( value: typing.Any, dtype: typing.Optional[bigframes.dtypes.Dtype], ) -> ArrayValue: + if pandas.isna(value): + # Need to assign a data type when value is NaN. + dtype = dtype or bigframes.dtypes.DEFAULT_DTYPE + if destination_id in self.column_ids: # Mutate case exprs = [ ( diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 355849538e..645914285e 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -605,17 +605,24 @@ def test_assign_new_column_w_loc(scalars_dfs): pd.testing.assert_frame_equal(bf_result, pd_result) -def test_assign_new_column_w_setitem(scalars_dfs): +@pytest.mark.parametrize( + ("scalar",), + [ + (2.1,), + (None,), + ], +) +def test_assign_new_column_w_setitem(scalars_dfs, scalar): scalars_df, scalars_pandas_df = scalars_dfs bf_df = scalars_df.copy() pd_df = scalars_pandas_df.copy() - bf_df["new_col"] = 2 - pd_df["new_col"] = 2 + bf_df["new_col"] = scalar + pd_df["new_col"] = scalar bf_result = bf_df.to_pandas() pd_result = pd_df - # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. - pd_result["new_col"] = pd_result["new_col"].astype("Int64") + # Convert default pandas dtypes `float64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Float64") pd.testing.assert_frame_equal(bf_result, pd_result)