feat: Support python type as astype arg (#1316)

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -707,7 +707,7 @@ def split(
  # Create an ordering col and convert to string
  block, ordering_col = block.promote_offsets()
  block, string_ordering_col = block.apply_unary_op(
- ordering_col, ops.AsTypeOp(to_type="string[pyarrow]")
+ ordering_col, ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE)
  )
 
  # Apply hash method to sum col and order by it.
@@ -1479,7 +1479,9 @@ def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block:
  expr, new_col = expr.project_to_id(
  expression=ops.add_op.as_expr(
  ex.const(prefix),
- ops.AsTypeOp(to_type="string").as_expr(index_col),
+ ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE).as_expr(
+ index_col
+ ),
  ),
  )
  new_index_cols.append(new_col)
@@ -1502,7 +1504,9 @@ def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block:
  for index_col in self._index_columns:
  expr, new_col = expr.project_to_id(
  expression=ops.add_op.as_expr(
- ops.AsTypeOp(to_type="string").as_expr(index_col),
+ ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE).as_expr(
+ index_col
+ ),
  ex.const(suffix),
  ),
  )

bigframes/core/compile/ibis_types.py

@@ -13,9 +13,8 @@
 # limitations under the License.
 from __future__ import annotations
 
-import textwrap
 import typing
-from typing import Any, cast, Dict, Iterable, Optional, Tuple, Union
+from typing import cast, Dict, Iterable, Optional, Tuple, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.ibis
@@ -28,7 +27,6 @@
 import db_dtypes # type: ignore
 import geopandas as gpd # type: ignore
 import google.cloud.bigquery as bigquery
-import numpy as np
 import pandas as pd
 import pyarrow as pa
 
@@ -228,9 +226,7 @@ def ibis_value_to_canonical_type(value: ibis_types.Value) -> ibis_types.Value:
 
 
 def bigframes_dtype_to_ibis_dtype(
- bigframes_dtype: Union[
- bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype, np.dtype[Any]
- ]
+ bigframes_dtype: bigframes.dtypes.Dtype,
 ) -> ibis_dtypes.DataType:
  """Converts a BigQuery DataFrames supported dtype to an Ibis dtype.
 
@@ -244,36 +240,14 @@ def bigframes_dtype_to_ibis_dtype(
  Raises:
  ValueError: If passed a dtype not supported by BigQuery DataFrames.
  """
- if str(bigframes_dtype) in bigframes.dtypes.BIGFRAMES_STRING_TO_BIGFRAMES:
- bigframes_dtype = bigframes.dtypes.BIGFRAMES_STRING_TO_BIGFRAMES[
- cast(bigframes.dtypes.DtypeString, str(bigframes_dtype))
- ]
-
  if bigframes_dtype in BIGFRAMES_TO_IBIS.keys():
  return BIGFRAMES_TO_IBIS[bigframes_dtype]
 
  elif isinstance(bigframes_dtype, pd.ArrowDtype) and bigframes_dtype.pyarrow_dtype:
  return _arrow_dtype_to_ibis_dtype(bigframes_dtype.pyarrow_dtype)
 
  else:
- raise ValueError(
- textwrap.dedent(
- f"""
- Unexpected data type {bigframes_dtype}. The following
- str dtypes are supppted: 'boolean','Float64','Int64',
- 'int64[pyarrow]','string','string[pyarrow]',
- 'timestamp[us, tz=UTC][pyarrow]','timestamp[us][pyarrow]',
- 'date32[day][pyarrow]','time64[us][pyarrow]'.
- The following pandas.ExtensionDtype are supported:
- pandas.BooleanDtype(), pandas.Float64Dtype(),
- pandas.Int64Dtype(), pandas.StringDtype(storage="pyarrow"),
- pd.ArrowDtype(pa.date32()), pd.ArrowDtype(pa.time64("us")),
- pd.ArrowDtype(pa.timestamp("us")),
- pd.ArrowDtype(pa.timestamp("us", tz="UTC")).
- {constants.FEEDBACK_LINK}
- """
- )
- )
+ raise ValueError(f"Datatype has no ibis type mapping: {bigframes_dtype}")
 
 
 def ibis_dtype_to_bigframes_dtype(

bigframes/core/indexes/base.py

@@ -78,7 +78,8 @@ def __new__(
  if name is not None:
  index.name = name
  if dtype is not None:
- index = index.astype(dtype)
+ bf_dtype = bigframes.dtypes.bigframes_type(dtype)
+ index = index.astype(bf_dtype)
  block = index._block
  elif isinstance(data, pandas.Index):
  pd_df = pandas.DataFrame(index=data)
@@ -310,14 +311,15 @@ def sort_values(self, *, ascending: bool = True, na_position: str = "last"):
 
  def astype(
  self,
- dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
+ dtype,
  *,
  errors: Literal["raise", "null"] = "raise",
  ) -> Index:
  if errors not in ["raise", "null"]:
  raise ValueError("Argument 'errors' must be one of 'raise' or 'null'")
  if self.nlevels > 1:
  raise TypeError("Multiindex does not support 'astype'")
+ dtype = bigframes.dtypes.bigframes_type(dtype)
  return self._apply_unary_expr(
  ops.AsTypeOp(to_type=dtype, safe=(errors == "null")).as_expr(
  ex.free_var("arg")

bigframes/core/local_data.py

@@ -59,6 +59,12 @@ def arrow_type_replacements(type: pa.DataType) -> pa.DataType:
  if pa.types.is_time64(type):
  # This is potentially lossy, but BigFrames doesn't support ns
  return pa.time64("us")
+ if pa.types.is_decimal128(type):
+ return pa.decimal128(38, 9)
+ if pa.types.is_decimal256(type):
+ return pa.decimal256(76, 38)
+ if pa.types.is_dictionary(type):
+ return arrow_type_replacements(type.value_type)
  if pa.types.is_large_string(type):
  # simple string type can handle the largest strings needed
  return pa.string()

bigframes/dataframe.py

@@ -180,7 +180,8 @@ def __init__(
  if columns:
  block = block.select_columns(list(columns)) # type:ignore
  if dtype:
- block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=dtype))
+ bf_dtype = bigframes.dtypes.bigframes_type(dtype)
+ block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=bf_dtype))
  self._block = block
 
  else:
@@ -368,6 +369,7 @@ def astype(
  dtype: Union[
  bigframes.dtypes.DtypeString,
  bigframes.dtypes.Dtype,
+ type,
  dict[str, Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype]],
  ],
  *,
@@ -378,23 +380,15 @@ def astype(
 
  safe_cast = errors == "null"
 
- # Type strings check
- if dtype in bigframes.dtypes.DTYPE_STRINGS:
- return self._apply_unary_op(ops.AsTypeOp(dtype, safe_cast))
-
- # Type instances check
- if type(dtype) in bigframes.dtypes.DTYPES:
- return self._apply_unary_op(ops.AsTypeOp(dtype, safe_cast))
-
  if isinstance(dtype, dict):
  result = self.copy()
  for col, to_type in dtype.items():
  result[col] = result[col].astype(to_type)
  return result
 
- raise TypeError(
- f"Invalid type {type(dtype)} for dtype input. {constants.FEEDBACK_LINK}"
- )
+ dtype = bigframes.dtypes.bigframes_type(dtype)
+
+ return self._apply_unary_op(ops.AsTypeOp(dtype, safe_cast))
 
  def _to_sql_query(
  self, include_index: bool, enable_cache: bool = True

bigframes/dtypes.py

@@ -17,6 +17,7 @@
 from dataclasses import dataclass
 import datetime
 import decimal
+import textwrap
 import typing
 from typing import Any, Dict, List, Literal, Union
 
@@ -422,7 +423,7 @@ def arrow_dtype_to_bigframes_dtype(arrow_dtype: pa.DataType) -> Dtype:
  return DEFAULT_DTYPE
 
  # No other types matched.
- raise ValueError(
+ raise TypeError(
  f"Unexpected Arrow data type {arrow_dtype}. {constants.FEEDBACK_LINK}"
  )
 
@@ -447,7 +448,7 @@ def bigframes_dtype_to_arrow_dtype(
  if pa.types.is_struct(bigframes_dtype.pyarrow_dtype):
  return bigframes_dtype.pyarrow_dtype
  else:
- raise ValueError(
+ raise TypeError(
  f"No arrow conversion for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
  )
 
@@ -474,7 +475,7 @@ def bigframes_dtype_to_literal(
  if isinstance(bigframes_dtype, gpd.array.GeometryDtype):
  return shapely.Point((0, 0))
 
- raise ValueError(
+ raise TypeError(
  f"No literal conversion for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
  )
 
@@ -515,11 +516,91 @@ def arrow_type_to_literal(
  if pa.types.is_time(arrow_type):
  return datetime.time(1, 1, 1)
 
- raise ValueError(
+ raise TypeError(
  f"No literal conversion for {arrow_type}. {constants.FEEDBACK_LINK}"
  )
 
 
+def bigframes_type(dtype) -> Dtype:
+ """Convert type object to canoncial bigframes dtype."""
+ if _is_bigframes_dtype(dtype):
+ return dtype
+ elif isinstance(dtype, str):
+ return _dtype_from_string(dtype)
+ elif isinstance(dtype, type):
+ return _infer_dtype_from_python_type(dtype)
+ elif isinstance(dtype, pa.DataType):
+ return arrow_dtype_to_bigframes_dtype(dtype)
+ else:
+ raise TypeError(
+ f"Cannot infer supported datatype for: {dtype}. {constants.FEEDBACK_LINK}"
+ )
+
+
+def _is_bigframes_dtype(dtype) -> bool:
+ """True iff dtyps is a canonical bigframes dtype"""
+ # have to be quite strict, as pyarrow dtypes equal their string form, and we don't consider that a canonical form.
+ if (type(dtype), dtype) in set(
+ (type(item.dtype), item.dtype) for item in SIMPLE_TYPES
+ ):
+ return True
+ if isinstance(dtype, pd.ArrowDtype):
+ try:
+ _ = arrow_dtype_to_bigframes_dtype(dtype.pyarrow_dtype)
+ return True
+ except TypeError:
+ return False
+ return False
+
+
+def _infer_dtype_from_python_type(type: type) -> Dtype:
+ if issubclass(type, (bool, np.bool_)):
+ return BOOL_DTYPE
+ if issubclass(type, (int, np.integer)):
+ return INT_DTYPE
+ if issubclass(type, (float, np.floating)):
+ return FLOAT_DTYPE
+ if issubclass(type, decimal.Decimal):
+ return NUMERIC_DTYPE
+ if issubclass(type, (str, np.str_)):
+ return STRING_DTYPE
+ if issubclass(type, (bytes, np.bytes_)):
+ return BYTES_DTYPE
+ if issubclass(type, datetime.date):
+ return DATE_DTYPE
+ if issubclass(type, datetime.time):
+ return TIME_DTYPE
+ else:
+ raise TypeError(
+ f"No matching datatype for python type: {type}. {constants.FEEDBACK_LINK}"
+ )
+
+
+def _dtype_from_string(dtype_string: str) -> typing.Optional[Dtype]:
+ if str(dtype_string) in BIGFRAMES_STRING_TO_BIGFRAMES:
+ return BIGFRAMES_STRING_TO_BIGFRAMES[
+ typing.cast(DtypeString, str(dtype_string))
+ ]
+ raise TypeError(
+ textwrap.dedent(
+ f"""
+ Unexpected data type string {dtype_string}. The following
+ dtypes are supppted: 'boolean','Float64','Int64',
+ 'int64[pyarrow]','string','string[pyarrow]',
+ 'timestamp[us, tz=UTC][pyarrow]','timestamp[us][pyarrow]',
+ 'date32[day][pyarrow]','time64[us][pyarrow]'.
+ The following pandas.ExtensionDtype are supported:
+ pandas.BooleanDtype(), pandas.Float64Dtype(),
+ pandas.Int64Dtype(), pandas.StringDtype(storage="pyarrow"),
+ pd.ArrowDtype(pa.date32()), pd.ArrowDtype(pa.time64("us")),
+ pd.ArrowDtype(pa.timestamp("us")),
+ pd.ArrowDtype(pa.timestamp("us", tz="UTC")).
+ {constants.FEEDBACK_LINK}
+ """
+ )
+ )
+
+
 def infer_literal_type(literal) -> typing.Optional[Dtype]:
  # Maybe also normalize literal to canonical python representation to remove this burden from compilers?
  if pd.api.types.is_list_like(literal):
@@ -539,30 +620,17 @@ def infer_literal_type(literal) -> typing.Optional[Dtype]:
  return pd.ArrowDtype(pa.struct(fields))
  if pd.isna(literal):
  return None # Null value without a definite type
- if isinstance(literal, (bool, np.bool_)):
- return BOOL_DTYPE
- if isinstance(literal, (int, np.integer)):
- return INT_DTYPE
- if isinstance(literal, (float, np.floating)):
- return FLOAT_DTYPE
- if isinstance(literal, decimal.Decimal):
- return NUMERIC_DTYPE
- if isinstance(literal, (str, np.str_)):
- return STRING_DTYPE
- if isinstance(literal, (bytes, np.bytes_)):
- return BYTES_DTYPE
  # 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 isinstance(literal, datetime.date):
- return DATE_DTYPE
- if isinstance(literal, datetime.time):
- return TIME_DTYPE
+ from_python_type = _infer_dtype_from_python_type(type(literal))
+ if from_python_type is not None:
+ return from_python_type
  else:
- raise ValueError(f"Unable to infer type for value: {literal}")
+ raise TypeError(f"Unable to infer type for value: {literal}")
 
 
 def infer_literal_arrow_type(literal) -> typing.Optional[pa.DataType]:
@@ -602,7 +670,7 @@ def convert_schema_field(
  return field.name, pd.ArrowDtype(pa_type)
  return field.name, _TK_TO_BIGFRAMES[field.field_type]
  else:
- raise ValueError(f"Cannot handle type: {field.field_type}")
+ raise TypeError(f"Cannot handle type: {field.field_type}")
 
 
 def convert_to_schema_field(
@@ -636,7 +704,7 @@ def convert_to_schema_field(
  if bigframes_dtype.pyarrow_dtype == pa.duration("us"):
  # Timedeltas are represented as integers in microseconds.
  return google.cloud.bigquery.SchemaField(name, "INTEGER")
- raise ValueError(
+ raise TypeError(
  f"No arrow conversion for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
  )
 

bigframes/operations/base.py

@@ -87,7 +87,8 @@ def __init__(
  if name is not None:
  data.name = name
  if dtype is not None:
- data = data.astype(dtype)
+ bf_dtype = bigframes.dtypes.bigframes_type(dtype)
+ data = data.astype(bf_dtype)
  else: # local dict-like data
  data = read_pandas_func(pd.Series(data, name=name, dtype=dtype)) # type: ignore
  data_block = data._block