feat: support RANGE in queries Part 2: Arrow

google/cloud/bigquery/_helpers.py

@@ -66,6 +66,8 @@
 _UNIVERSE_DOMAIN_ENV = "GOOGLE_CLOUD_UNIVERSE_DOMAIN"
 """Environment variable for setting universe domain."""
 
+_SUPPORTED_RANGE_ELEMENTS = {"TIMESTAMP", "DATETIME", "DATE"}
+
 
 def _get_client_universe(
  client_options: Optional[Union[client_options_lib.ClientOptions, dict]]
@@ -310,17 +312,13 @@ def _json_from_json(value, field):
 
 
 def _range_element_from_json(value, field):
- """Coerce 'value' to a range element value, if set or not nullable."""
+ """Coerce 'value' to a range element value."""
  if value == "UNBOUNDED":
  return None
- elif field.element_type == "DATE":
- return _date_from_json(value, None)
- elif field.element_type == "DATETIME":
- return _datetime_from_json(value, None)
- elif field.element_type == "TIMESTAMP":
- return _timestamp_from_json(value, None)
+ if field.element_type in _SUPPORTED_RANGE_ELEMENTS:
+ return _CELLDATA_FROM_JSON[field.element_type](value, field.element_type)
  else:
- raise ValueError(f"Unsupported range field type: {value}")
+ raise ValueError(f"Unsupported range element type: {field.element_type}")
 
 
 def _range_from_json(value, field):
@@ -344,7 +342,7 @@ def _range_from_json(value, field):
  end = _range_element_from_json(end, field.range_element_type)
  return {"start": start, "end": end}
  else:
- raise ValueError(f"Unknown range format: {value}")
+ raise ValueError(f"Unknown format for range value: {value}")
  else:
  return None
 

google/cloud/bigquery/_pandas_helpers.py

@@ -142,6 +142,17 @@ def bq_to_arrow_struct_data_type(field):
  return pyarrow.struct(arrow_fields)
 
 
+def bq_to_arrow_range_data_type(field):
+ if field is None:
+ raise ValueError(
+ "Range element type cannot be None, must be one of "
+ "DATE, DATETIME, or TIMESTAMP"
+ )
+ element_type = field.element_type.upper()
+ arrow_element_type = _pyarrow_helpers.bq_to_arrow_scalars(element_type)()
+ return pyarrow.struct([("start", arrow_element_type), ("end", arrow_element_type)])
+
+
 def bq_to_arrow_data_type(field):
  """Return the Arrow data type, corresponding to a given BigQuery column.
 
@@ -160,6 +171,9 @@ def bq_to_arrow_data_type(field):
  if field_type_upper in schema._STRUCT_TYPES:
  return bq_to_arrow_struct_data_type(field)
 
+ if field_type_upper == "RANGE":
+ return bq_to_arrow_range_data_type(field.range_element_type)
+
  data_type_constructor = _pyarrow_helpers.bq_to_arrow_scalars(field_type_upper)
  if data_type_constructor is None:
  return None
@@ -220,6 +234,9 @@ def default_types_mapper(
  datetime_dtype: Union[Any, None] = None,
  time_dtype: Union[Any, None] = None,
  timestamp_dtype: Union[Any, None] = None,
+ range_date_dtype: Union[Any, None] = None,
+ range_datetime_dtype: Union[Any, None] = None,
+ range_timestamp_dtype: Union[Any, None] = None,
 ):
  """Create a mapping from pyarrow types to pandas types.
 
@@ -274,6 +291,22 @@ def types_mapper(arrow_data_type):
  elif time_dtype is not None and pyarrow.types.is_time(arrow_data_type):
  return time_dtype
 
+ elif pyarrow.types.is_struct(arrow_data_type):
+ if range_datetime_dtype is not None and arrow_data_type.equals(
+ range_datetime_dtype.pyarrow_dtype
+ ):
+ return range_datetime_dtype
+
+ elif range_date_dtype is not None and arrow_data_type.equals(
+ range_date_dtype.pyarrow_dtype
+ ):
+ return range_date_dtype
+
+ elif range_timestamp_dtype is not None and arrow_data_type.equals(
+ range_timestamp_dtype.pyarrow_dtype
+ ):
+ return range_timestamp_dtype
+
  return types_mapper
 
 

google/cloud/bigquery/dbapi/_helpers.py

@@ -277,12 +277,14 @@ def complex_query_parameter(
  param = query.ArrayQueryParameter(
  name,
  sub_type,
- value
- if isinstance(sub_type, query.ScalarQueryParameterType)
- else [
- complex_query_parameter(None, v, sub_type._complex__src, base)
- for v in value
- ],
+ (
+ value
+ if isinstance(sub_type, query.ScalarQueryParameterType)
+ else [
+ complex_query_parameter(None, v, sub_type._complex__src, base)
+ for v in value
+ ]
+ ),
  )
  elif type_type == STRUCT:
  if not isinstance(value, collections_abc.Mapping):

google/cloud/bigquery/enums.py

@@ -99,6 +99,15 @@ class DefaultPandasDTypes(enum.Enum):
  TIME_DTYPE = object()
  """Specifies default time dtype"""
 
+ RANGE_DATE_DTYPE = object()
+ """Specifies default range date dtype"""
+
+ RANGE_DATETIME_DTYPE = object()
+ """Specifies default range datetime dtype"""
+
+ RANGE_TIMESTAMP_DTYPE = object()
+ """Specifies default range timestamp dtype"""
+
 
 class DestinationFormat(object):
  """The exported file format. The default value is :attr:`CSV`.

google/cloud/bigquery/job/query.py

@@ -1784,6 +1784,13 @@ def to_dataframe(
  datetime_dtype: Union[Any, None] = None,
  time_dtype: Union[Any, None] = DefaultPandasDTypes.TIME_DTYPE,
  timestamp_dtype: Union[Any, None] = None,
+ range_date_dtype: Union[Any, None] = DefaultPandasDTypes.RANGE_DATE_DTYPE,
+ range_datetime_dtype: Union[
+ Any, None
+ ] = DefaultPandasDTypes.RANGE_DATETIME_DTYPE,
+ range_timestamp_dtype: Union[
+ Any, None
+ ] = DefaultPandasDTypes.RANGE_TIMESTAMP_DTYPE,
  ) -> "pandas.DataFrame":
  """Return a pandas DataFrame from a QueryJob
 
@@ -1919,6 +1926,63 @@ def to_dataframe(
 
  .. versionadded:: 3.10.0
 
+ range_date_dtype (Optional[pandas.Series.dtype, None]):
+ If set, indicate a pandas ExtensionDtype, such as:
+
+ .. code-block:: python
+
+ pandas.ArrowDtype(pyarrow.struct(
+ [("start", pyarrow.date32()), ("end", pyarrow.date32())]
+ ))
+
+ to convert BigQuery RANGE<DATE> type, instead of relying on
+ the default ``object``. If you explicitly set the value to
+ ``None``, the data type will be ``object``. BigQuery Range type
+ can be found at:
+ https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#range_type
+
+ .. versionadded:: 3.21.0
+
+ range_datetime_dtype (Optional[pandas.Series.dtype, None]):
+ If set, indicate a pandas ExtensionDtype, such as:
+
+ .. code-block:: python
+
+ pandas.ArrowDtype(pyarrow.struct(
+ [
+ ("start", pyarrow.timestamp("us")),
+ ("end", pyarrow.timestamp("us")),
+ ]
+ ))
+
+ to convert BigQuery RANGE<DATETIME> type, instead of relying on
+ the default ``object``. If you explicitly set the value to
+ ``None``, the data type will be ``object``. BigQuery Range type
+ can be found at:
+ https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#range_type
+
+ .. versionadded:: 3.21.0
+
+ range_timestamp_dtype (Optional[pandas.Series.dtype, None]):
+ If set, indicate a pandas ExtensionDtype, such as:
+
+ .. code-block:: python
+
+ pandas.ArrowDtype(pyarrow.struct(
+ [
+ ("start", pyarrow.timestamp("us", tz="UTC")),
+ ("end", pyarrow.timestamp("us", tz="UTC")),
+ ]
+ ))
+
+ to convert BigQuery RANGE<TIMESTAMP> type, instead of relying
+ on the default ``object``. If you explicitly set the value to
+ ``None``, the data type will be ``object``. BigQuery Range type
+ can be found at:
+ https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#range_type
+
+ .. versionadded:: 3.21.0
+
  Returns:
  pandas.DataFrame:
  A :class:`~pandas.DataFrame` populated with row data
@@ -1949,6 +2013,9 @@ def to_dataframe(
  datetime_dtype=datetime_dtype,
  time_dtype=time_dtype,
  timestamp_dtype=timestamp_dtype,
+ range_date_dtype=range_date_dtype,
+ range_datetime_dtype=range_datetime_dtype,
+ range_timestamp_dtype=range_timestamp_dtype,
  )
 
  # If changing the signature of this method, make sure to apply the same

google/cloud/bigquery/query.py

@@ -24,14 +24,13 @@
 from google.cloud.bigquery._helpers import _rows_from_json
 from google.cloud.bigquery._helpers import _QUERY_PARAMS_FROM_JSON
 from google.cloud.bigquery._helpers import _SCALAR_VALUE_TO_JSON_PARAM
+from google.cloud.bigquery._helpers import _SUPPORTED_RANGE_ELEMENTS
 
 
 _SCALAR_VALUE_TYPE = Optional[
  Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]
 ]
 
-_RANGE_ELEMENT_TYPE_STR = {"TIMESTAMP", "DATETIME", "DATE"}
-
 
 class ConnectionProperty:
  """A connection-level property to customize query behavior.
@@ -388,14 +387,14 @@ def _parse_range_element_type(self, type_):
  google.cloud.bigquery.query.ScalarQueryParameterType: Instance
  """
  if isinstance(type_, str):
- if type_ not in _RANGE_ELEMENT_TYPE_STR:
+ if type_ not in _SUPPORTED_RANGE_ELEMENTS:
  raise ValueError(
  "If given as a string, range element type must be one of "
  "'TIMESTAMP', 'DATE', or 'DATETIME'."
  )
  return ScalarQueryParameterType(type_)
  elif isinstance(type_, ScalarQueryParameterType):
- if type_._type not in _RANGE_ELEMENT_TYPE_STR:
+ if type_._type not in _SUPPORTED_RANGE_ELEMENTS:
  raise ValueError(
  "If given as a ScalarQueryParameter object, range element "
  "type must be one of 'TIMESTAMP', 'DATE', or 'DATETIME' "
@@ -960,14 +959,14 @@ class RangeQueryParameter(_AbstractQueryParameter):
  @classmethod
  def _parse_range_element_type(self, range_element_type):
  if isinstance(range_element_type, str):
- if range_element_type not in _RANGE_ELEMENT_TYPE_STR:
+ if range_element_type not in _SUPPORTED_RANGE_ELEMENTS:
  raise ValueError(
  "If given as a string, range_element_type must be one of "
  f"'TIMESTAMP', 'DATE', or 'DATETIME'. Got {range_element_type}."
  )
  return RangeQueryParameterType(range_element_type)
  elif isinstance(range_element_type, RangeQueryParameterType):
- if range_element_type.type_._type not in _RANGE_ELEMENT_TYPE_STR:
+ if range_element_type.type_._type not in _SUPPORTED_RANGE_ELEMENTS:
  raise ValueError(
  "If given as a RangeQueryParameterType object, "
  "range_element_type must be one of 'TIMESTAMP', 'DATE', "