diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 262fa9dde7..2c5df89665 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -754,7 +754,7 @@ def bf_type_from_type_kind( def is_dtype(scalar: typing.Any, dtype: Dtype) -> bool: """Captures whether a scalar can be losslessly represented by a dtype.""" - if scalar is None: + if pd.isna(scalar): return True if pd.api.types.is_bool_dtype(dtype): return pd.api.types.is_bool(scalar) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 710e1481be..c391370805 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -629,6 +629,18 @@ def test_series_replace_list_scalar(scalars_dfs): ) +def test_series_replace_nans_with_pd_na(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = scalars_df[col_name].replace({pd.NA: "UNKNOWN"}).to_pandas() + pd_result = scalars_pandas_df[col_name].replace({pd.NA: "UNKNOWN"}) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + @pytest.mark.parametrize( ("replacement_dict",), ( diff --git a/third_party/bigframes_vendored/ibis/expr/types/core.py b/third_party/bigframes_vendored/ibis/expr/types/core.py index 9685e4ddca..5704dc993a 100644 --- a/third_party/bigframes_vendored/ibis/expr/types/core.py +++ b/third_party/bigframes_vendored/ibis/expr/types/core.py @@ -19,6 +19,7 @@ import bigframes_vendored.ibis.expr.operations as ops from bigframes_vendored.ibis.expr.types.pretty import to_rich from bigframes_vendored.ibis.util import experimental +import pandas as pd from public import public from rich.console import Console from rich.jupyter import JupyterMixin @@ -34,7 +35,6 @@ EdgeAttributeGetter, NodeAttributeGetter, ) - import pandas as pd import polars as pl import pyarrow as pa import torch @@ -744,9 +744,9 @@ def _binop(op_class: type[ops.Binary], left: ir.Value, right: ir.Value) -> ir.Va def _is_null_literal(value: Any) -> bool: """Detect whether `value` will be treated by ibis as a null literal.""" - if value is None: - return True if isinstance(value, Expr): op = value.op() return isinstance(op, ops.Literal) and op.value is None + if pd.isna(value): + return True return False