diff --git a/bigquery/google/cloud/bigquery/_pandas_helpers.py b/bigquery/google/cloud/bigquery/_pandas_helpers.py index 5ac0505e91ae2..5cc69e434b04e 100644 --- a/bigquery/google/cloud/bigquery/_pandas_helpers.py +++ b/bigquery/google/cloud/bigquery/_pandas_helpers.py @@ -47,7 +47,6 @@ "please install google-cloud-bigquery-storage to use bqstorage features." ) -STRUCT_TYPES = ("RECORD", "STRUCT") _PROGRESS_INTERVAL = 0.2 # Maximum time between download status checks, in seconds. @@ -126,7 +125,7 @@ def bq_to_arrow_data_type(field): return pyarrow.list_(inner_type) return None - if field.field_type.upper() in STRUCT_TYPES: + if field.field_type.upper() in schema._STRUCT_TYPES: return bq_to_arrow_struct_data_type(field) data_type_constructor = BQ_TO_ARROW_SCALARS.get(field.field_type.upper()) @@ -168,7 +167,7 @@ def bq_to_arrow_array(series, bq_field): arrow_type = bq_to_arrow_data_type(bq_field) if bq_field.mode.upper() == "REPEATED": return pyarrow.ListArray.from_pandas(series, type=arrow_type) - if bq_field.field_type.upper() in STRUCT_TYPES: + if bq_field.field_type.upper() in schema._STRUCT_TYPES: return pyarrow.StructArray.from_pandas(series, type=arrow_type) return pyarrow.array(series, type=arrow_type) diff --git a/bigquery/google/cloud/bigquery/schema.py b/bigquery/google/cloud/bigquery/schema.py index 61bc0bcedfd64..e0673d85baf69 100644 --- a/bigquery/google/cloud/bigquery/schema.py +++ b/bigquery/google/cloud/bigquery/schema.py @@ -17,6 +17,8 @@ from google.cloud.bigquery_v2 import types +_STRUCT_TYPES = ("RECORD", "STRUCT") + # SQL types reference: # https://cloud.google.com/bigquery/data-types#legacy_sql_data_types # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types @@ -150,7 +152,7 @@ def to_api_repr(self): # If this is a RECORD type, then sub-fields are also included, # add this to the serialized representation. - if self.field_type.upper() == "RECORD": + if self.field_type.upper() in _STRUCT_TYPES: answer["fields"] = [f.to_api_repr() for f in self.fields] # Done; return the serialized dictionary. diff --git a/bigquery/tests/unit/test_schema.py b/bigquery/tests/unit/test_schema.py index 682e458958528..862d8a823e620 100644 --- a/bigquery/tests/unit/test_schema.py +++ b/bigquery/tests/unit/test_schema.py @@ -71,25 +71,26 @@ def test_to_api_repr(self): ) def test_to_api_repr_with_subfield(self): - subfield = self._make_one("bar", "INTEGER", "NULLABLE") - field = self._make_one("foo", "RECORD", "REQUIRED", fields=(subfield,)) - self.assertEqual( - field.to_api_repr(), - { - "fields": [ - { - "mode": "NULLABLE", - "name": "bar", - "type": "INTEGER", - "description": None, - } - ], - "mode": "REQUIRED", - "name": "foo", - "type": "RECORD", - "description": None, - }, - ) + for record_type in ("RECORD", "STRUCT"): + subfield = self._make_one("bar", "INTEGER", "NULLABLE") + field = self._make_one("foo", record_type, "REQUIRED", fields=(subfield,)) + self.assertEqual( + field.to_api_repr(), + { + "fields": [ + { + "mode": "NULLABLE", + "name": "bar", + "type": "INTEGER", + "description": None, + } + ], + "mode": "REQUIRED", + "name": "foo", + "type": record_type, + "description": None, + }, + ) def test_from_api_repr(self): field = self._get_target_class().from_api_repr(