fix: assign NaN scalar to column error. (#513)

* fix: assign NaN scalar to column error. * Update test. * remove import.

Author: Genesis929

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 = [
  (

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)