fix: support setitem with NaTType

Author: chelsea-lin

bigframes/core/array_value.py

@@ -324,10 +324,6 @@ def create_constant(
         value: typing.Any,
         dtype: typing.Optional[bigframes.dtypes.Dtype],
     ) -> Tuple[ArrayValue, str]:
-        if pandas.isna(value):
-            # Need to assign a data type when value is NaN.
-            dtype = dtype or bigframes.dtypes.DEFAULT_DTYPE
-
         return self.project_to_id(ex.const(value, dtype))
 
     def select_columns(

bigframes/dtypes.py

@@ -674,14 +674,14 @@ def infer_literal_type(literal) -> typing.Optional[Dtype]:
                 pa.field(key, field_type, nullable=(not pa.types.is_list(field_type)))
             )
         return pd.ArrowDtype(pa.struct(fields))
-    if pd.isna(literal):
-        return None  # Null value without a definite type
     # Make sure to check datetime before date as datetimes are also dates
     if isinstance(literal, (datetime.datetime, pd.Timestamp)):
         if literal.tzinfo is not None:
             return TIMESTAMP_DTYPE
         else:
             return DATETIME_DTYPE
+    if pd.isna(literal):
+        return None  # Null value without a definite type
     from_python_type = _infer_dtype_from_python_type(type(literal))
     if from_python_type is not None:
         return from_python_type

tests/system/small/test_dataframe.py

@@ -3764,6 +3764,12 @@ def test_at_no_duplicate(scalars_df_index, scalars_pandas_df_index):
     assert bf_result == pd_result
 
 
+def test_setitem_w_timestamp_none():
+    b_df = bpd.DataFrame({'rowindex': [1, 2, 3]})
+    b_df['temp_timestamp'] = pd.Timestamp(ts_input=None, unit="us", tz="utc")
+    assert b_df['temp_timestamp'].dtype == "timestamp[us][pyarrow]"
+
+
 def test_loc_setitem_bool_series_scalar_new_col(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     bf_df = scalars_df.copy()
@@ -5714,4 +5720,4 @@ def test_agg_with_dict_containing_non_existing_col_raise_key_error(scalars_dfs):
     }
 
     with pytest.raises(KeyError):
-        bf_df.agg(agg_funcs)
+        bf_df.agg(agg_funcs)