feat: add astype(type, errors='null') to cast safely (#1122)

Author: TrevorBergeron

bigframes/core/compile/ibis_types.py

@@ -89,7 +89,7 @@
 
 
 def cast_ibis_value(
- value: ibis_types.Value, to_type: ibis_dtypes.DataType
+ value: ibis_types.Value, to_type: ibis_dtypes.DataType, safe: bool = False
 ) -> ibis_types.Value:
  """Perform compatible type casts of ibis values
 
@@ -176,7 +176,7 @@ def cast_ibis_value(
  value = ibis_value_to_canonical_type(value)
  if value.type() in good_casts:
  if to_type in good_casts[value.type()]:
- return value.cast(to_type)
+ return value.try_cast(to_type) if safe else value.cast(to_type)
  else:
  # this should never happen
  raise TypeError(
@@ -188,10 +188,16 @@ def cast_ibis_value(
  # BigQuery casts bools to lower case strings. Capitalize the result to match Pandas
  # TODO(bmil): remove this workaround after fixing Ibis
  if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.string:
- return cast(ibis_types.StringValue, value.cast(to_type)).capitalize()
+ if safe:
+ return cast(ibis_types.StringValue, value.try_cast(to_type)).capitalize()
+ else:
+ return cast(ibis_types.StringValue, value.cast(to_type)).capitalize()
 
  if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.float64:
- return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64)
+ if safe:
+ return value.try_cast(ibis_dtypes.int64).try_cast(ibis_dtypes.float64)
+ else:
+ return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64)
 
  if value.type() == ibis_dtypes.float64 and to_type == ibis_dtypes.bool:
  return value != ibis_types.literal(0)

bigframes/core/compile/scalar_op_compiler.py

@@ -947,7 +947,9 @@ def struct_field_op_impl(x: ibis_types.Value, op: ops.StructFieldOp):
  return result.cast(result.type()(nullable=True)).name(name)
 
 
-def numeric_to_datetime(x: ibis_types.Value, unit: str) -> ibis_types.TimestampValue:
+def numeric_to_datetime(
+ x: ibis_types.Value, unit: str, safe: bool = False
+) -> ibis_types.TimestampValue:
  if not isinstance(x, ibis_types.IntegerValue) and not isinstance(
  x, ibis_types.FloatingValue
  ):
@@ -956,7 +958,11 @@ def numeric_to_datetime(x: ibis_types.Value, unit: str) -> ibis_types.TimestampV
  if unit not in UNIT_TO_US_CONVERSION_FACTORS:
  raise ValueError(f"Cannot convert input with unit '{unit}'.")
  x_converted = x * UNIT_TO_US_CONVERSION_FACTORS[unit]
- x_converted = x_converted.cast(ibis_dtypes.int64)
+ x_converted = (
+ x_converted.try_cast(ibis_dtypes.int64)
+ if safe
+ else x_converted.cast(ibis_dtypes.int64)
+ )
 
  # Note: Due to an issue where casting directly to a timestamp
  # without a timezone does not work, we first cast to UTC. This
@@ -978,8 +984,11 @@ def astype_op_impl(x: ibis_types.Value, op: ops.AsTypeOp):
 
  # When casting DATETIME column into INT column, we need to convert the column into TIMESTAMP first.
  if to_type == ibis_dtypes.int64 and x.type() == ibis_dtypes.timestamp:
- x_converted = x.cast(ibis_dtypes.Timestamp(timezone="UTC"))
- return bigframes.core.compile.ibis_types.cast_ibis_value(x_converted, to_type)
+ utc_time_type = ibis_dtypes.Timestamp(timezone="UTC")
+ x_converted = x.try_cast(utc_time_type) if op.safe else x.cast(utc_time_type)
+ return bigframes.core.compile.ibis_types.cast_ibis_value(
+ x_converted, to_type, safe=op.safe
+ )
 
  if to_type == ibis_dtypes.int64 and x.type() == ibis_dtypes.time:
  # The conversion unit is set to "us" (microseconds) for consistency
@@ -991,15 +1000,20 @@ def astype_op_impl(x: ibis_types.Value, op: ops.AsTypeOp):
  # with pandas converting int64[pyarrow] to timestamp[us][pyarrow],
  # timestamp[us, tz=UTC][pyarrow], and time64[us][pyarrow].
  unit = "us"
- x_converted = numeric_to_datetime(x, unit)
+ x_converted = numeric_to_datetime(x, unit, safe=op.safe)
  if to_type == ibis_dtypes.timestamp:
- return x_converted.cast(ibis_dtypes.Timestamp())
+ return (
+ x_converted.try_cast(ibis_dtypes.Timestamp())
+ if op.safe
+ else x_converted.cast(ibis_dtypes.Timestamp())
+ )
  elif to_type == ibis_dtypes.Timestamp(timezone="UTC"):
  return x_converted
  elif to_type == ibis_dtypes.time:
  return x_converted.time()
 
- return bigframes.core.compile.ibis_types.cast_ibis_value(x, to_type)
+ # TODO: either inline this function, or push rest of this op into the function
+ return bigframes.core.compile.ibis_types.cast_ibis_value(x, to_type, safe=op.safe)
 
 
 @scalar_op_compiler.register_unary_op(ops.IsInOp, pass_op=True)

bigframes/core/indexes/base.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Hashable, Optional, Sequence, Union
+from typing import Hashable, Literal, Optional, Sequence, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.indexes.base as vendored_pandas_index
@@ -324,11 +324,17 @@ def sort_values(self, *, ascending: bool = True, na_position: str = "last"):
  def astype(
  self,
  dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
+ *,
+ errors: Literal["raise", "null"] = "raise",
  ) -> Index:
+ if errors not in ["raise", "null"]:
+ raise ValueError("Arg 'error' must be one of 'raise' or 'null'")
  if self.nlevels > 1:
  raise TypeError("Multiindex does not support 'astype'")
  return self._apply_unary_expr(
- ops.AsTypeOp(to_type=dtype).as_expr(ex.free_var("arg"))
+ ops.AsTypeOp(to_type=dtype, safe=(errors == "null")).as_expr(
+ ex.free_var("arg")
+ )
  )
 
  def all(self) -> bool:

bigframes/dataframe.py

@@ -365,8 +365,14 @@ def __iter__(self):
  def astype(
  self,
  dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
+ *,
+ errors: Literal["raise", "null"] = "raise",
  ) -> DataFrame:
- return self._apply_unary_op(ops.AsTypeOp(to_type=dtype))
+ if errors not in ["raise", "null"]:
+ raise ValueError("Arg 'error' must be one of 'raise' or 'null'")
+ return self._apply_unary_op(
+ ops.AsTypeOp(to_type=dtype, safe=(errors == "null"))
+ )
 
  def _to_sql_query(
  self, include_index: bool, enable_cache: bool = True

bigframes/operations/__init__.py

@@ -494,6 +494,7 @@ class AsTypeOp(UnaryOp):
  name: typing.ClassVar[str] = "astype"
  # TODO: Convert strings to dtype earlier
  to_type: dtypes.DtypeString | dtypes.Dtype
+ safe: bool = False
 
  def output_type(self, *input_types):
  # TODO: We should do this conversion earlier

bigframes/series.py

@@ -352,8 +352,14 @@ def __repr__(self) -> str:
  def astype(
  self,
  dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
+ *,
+ errors: Literal["raise", "null"] = "raise",
  ) -> Series:
- return self._apply_unary_op(bigframes.operations.AsTypeOp(to_type=dtype))
+ if errors not in ["raise", "null"]:
+ raise ValueError("Arg 'error' must be one of 'raise' or 'null'")
+ return self._apply_unary_op(
+ bigframes.operations.AsTypeOp(to_type=dtype, safe=(errors == "null"))
+ )
 
  def to_pandas(
  self,

tests/system/small/test_dataframe.py

@@ -3687,6 +3687,12 @@ def test_df_add_suffix(scalars_df_index, scalars_pandas_df_index, axis):
  )
 
 
+def test_df_astype_error_error(session):
+ input = pd.DataFrame(["hello", "world", "3.11", "4000"])
+ with pytest.raises(ValueError):
+ session.read_pandas(input).astype("Float64", errors="bad_value")
+
+
 def test_df_columns_filter_items(scalars_df_index, scalars_pandas_df_index):
  if pd.__version__.startswith("2.0") or pd.__version__.startswith("1."):
  pytest.skip("pandas filter items behavior different pre-2.1")

tests/system/small/test_index.py

@@ -123,6 +123,12 @@ def test_index_astype(scalars_df_index, scalars_pandas_df_index):
  pd.testing.assert_index_equal(bf_result, pd_result)
 
 
+def test_index_astype_error_error(session):
+ input = pd.Index(["hello", "world", "3.11", "4000"])
+ with pytest.raises(ValueError):
+ session.read_pandas(input).astype("Float64", errors="bad_value")
+
+
 def test_index_any(scalars_df_index, scalars_pandas_df_index):
  bf_result = scalars_df_index.set_index("int64_col").index.any()
  pd_result = scalars_pandas_df_index.set_index("int64_col").index.any()

tests/system/small/test_series.py

@@ -3092,6 +3092,7 @@ def foo(x):
  assert_series_equal(bf_result, pd_result, check_dtype=False)
 
 
+@pytest.mark.parametrize("errors", ["raise", "null"])
 @pytest.mark.parametrize(
  ("column", "to_type"),
  [
@@ -3107,6 +3108,7 @@ def foo(x):
  ("int64_col", "time64[us][pyarrow]"),
  ("bool_col", "Int64"),
  ("bool_col", "string[pyarrow]"),
+ ("bool_col", "Float64"),
  ("string_col", "binary[pyarrow]"),
  ("bytes_col", "string[pyarrow]"),
  # pandas actually doesn't let folks convert to/from naive timestamp and
@@ -3142,12 +3144,29 @@ def foo(x):
  ],
 )
 @skip_legacy_pandas
-def test_astype(scalars_df_index, scalars_pandas_df_index, column, to_type):
- bf_result = scalars_df_index[column].astype(to_type).to_pandas()
+def test_astype(scalars_df_index, scalars_pandas_df_index, column, to_type, errors):
+ bf_result = scalars_df_index[column].astype(to_type, errors=errors).to_pandas()
  pd_result = scalars_pandas_df_index[column].astype(to_type)
  pd.testing.assert_series_equal(bf_result, pd_result)
 
 
+def test_astype_safe(session):
+ input = pd.Series(["hello", "world", "3.11", "4000"])
+ exepcted = pd.Series(
+ [None, None, 3.11, 4000],
+ dtype="Float64",
+ index=pd.Index([0, 1, 2, 3], dtype="Int64"),
+ )
+ result = session.read_pandas(input).astype("Float64", errors="null").to_pandas()
+ pd.testing.assert_series_equal(result, exepcted)
+
+
+def test_series_astype_error_error(session):
+ input = pd.Series(["hello", "world", "3.11", "4000"])
+ with pytest.raises(ValueError):
+ session.read_pandas(input).astype("Float64", errors="bad_value")
+
+
 @skip_legacy_pandas
 def test_astype_numeric_to_int(scalars_df_index, scalars_pandas_df_index):
  column = "numeric_col"

third_party/bigframes_vendored/pandas/core/generic.py

@@ -180,6 +180,10 @@ def astype(self, dtype):
  ``pd.ArrowDtype(pa.time64("us"))``,
  ``pd.ArrowDtype(pa.timestamp("us"))``,
  ``pd.ArrowDtype(pa.timestamp("us", tz="UTC"))``.
+ errors ({'raise', 'null'}, default 'raise'):
+ Control raising of exceptions on invalid data for provided dtype.
+ If 'raise', allow exceptions to be raised if any value fails cast
+ If 'null', will assign null value if value fails cast
 
  Returns:
  bigframes.pandas.DataFrame: