diff --git a/google/cloud/bigquery_storage_v1/reader.py b/google/cloud/bigquery_storage_v1/reader.py index 09b0d99c..3587ff9e 100644 --- a/google/cloud/bigquery_storage_v1/reader.py +++ b/google/cloud/bigquery_storage_v1/reader.py @@ -649,7 +649,7 @@ def to_dataframe(self, message, dtypes=None): def _parse_arrow_message(self, message): self._parse_arrow_schema() - return pyarrow.read_record_batch( + return pyarrow.ipc.read_record_batch( pyarrow.py_buffer(message.arrow_record_batch.serialized_record_batch), self._schema, ) @@ -658,7 +658,7 @@ def _parse_arrow_schema(self): if self._schema: return - self._schema = pyarrow.read_schema( + self._schema = pyarrow.ipc.read_schema( pyarrow.py_buffer(self._read_session.arrow_schema.serialized_schema) ) self._column_names = [field.name for field in self._schema] diff --git a/google/cloud/bigquery_storage_v1beta1/reader.py b/google/cloud/bigquery_storage_v1beta1/reader.py index 7e07392d..60942f29 100644 --- a/google/cloud/bigquery_storage_v1beta1/reader.py +++ b/google/cloud/bigquery_storage_v1beta1/reader.py @@ -607,7 +607,7 @@ def to_dataframe(self, message, dtypes=None): def _parse_arrow_message(self, message): self._parse_arrow_schema() - return pyarrow.read_record_batch( + return pyarrow.ipc.read_record_batch( pyarrow.py_buffer(message.arrow_record_batch.serialized_record_batch), self._schema, ) @@ -616,7 +616,7 @@ def _parse_arrow_schema(self): if self._schema: return - self._schema = pyarrow.read_schema( + self._schema = pyarrow.ipc.read_schema( pyarrow.py_buffer(self._read_session.arrow_schema.serialized_schema) ) self._column_names = [field.name for field in self._schema] diff --git a/tests/system/v1/test_reader_dataframe_v1.py b/tests/system/v1/test_reader_dataframe_v1.py index 2162ba48..ec3e983c 100644 --- a/tests/system/v1/test_reader_dataframe_v1.py +++ b/tests/system/v1/test_reader_dataframe_v1.py @@ -46,12 +46,12 @@ def test_read_v1(client, project_id): assert tbl.num_columns == 4 schema = tbl.schema - # Use field_by_name because the order doesn't currently match that of - # selected_fields. - assert pyarrow.types.is_int64(schema.field_by_name("station_id").type) - assert pyarrow.types.is_float64(schema.field_by_name("latitude").type) - assert pyarrow.types.is_float64(schema.field_by_name("longitude").type) - assert pyarrow.types.is_string(schema.field_by_name("name").type) + # Use field with a name specifier as there may be ordering differences + # when selected_fields is used + assert pyarrow.types.is_int64(schema.field("station_id").type) + assert pyarrow.types.is_float64(schema.field("latitude").type) + assert pyarrow.types.is_float64(schema.field("longitude").type) + assert pyarrow.types.is_string(schema.field("name").type) @pytest.mark.parametrize( diff --git a/tests/system/v1beta1/test_reader_dataframe_v1beta1.py b/tests/system/v1beta1/test_reader_dataframe_v1beta1.py index 07dcab38..20143f0f 100644 --- a/tests/system/v1beta1/test_reader_dataframe_v1beta1.py +++ b/tests/system/v1beta1/test_reader_dataframe_v1beta1.py @@ -48,12 +48,12 @@ def test_read_rows_to_arrow(client, project_id): assert tbl.num_columns == 4 schema = tbl.schema - # Use field_by_name because the order doesn't currently match that of - # selected_fields. - assert pyarrow.types.is_int64(schema.field_by_name("station_id").type) - assert pyarrow.types.is_float64(schema.field_by_name("latitude").type) - assert pyarrow.types.is_float64(schema.field_by_name("longitude").type) - assert pyarrow.types.is_string(schema.field_by_name("name").type) + # Use field with a name specifier as there may be ordering differences + # when selected_fields is used + assert pyarrow.types.is_int64(schema.field("station_id").type) + assert pyarrow.types.is_float64(schema.field("latitude").type) + assert pyarrow.types.is_float64(schema.field("longitude").type) + assert pyarrow.types.is_string(schema.field("name").type) @pytest.mark.parametrize(