adds test for query results to_dataframe() with nested schema

googleapis · Nov 17, 2017 · f93f51b · f93f51b
1 parent 0cdc967
commit f93f51b
Showing 1 changed file with 49 additions and 0 deletions.
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
@@ -1437,6 +1437,55 @@ def test_create_table_rows_fetch_nested_schema(self):
             e_favtime = datetime.datetime(*parts[0:6])
             self.assertEqual(found[7], e_favtime)
 
+    def _fetch_dataframe(self, query):
+        return Config.CLIENT.query(query).result().to_dataframe()
+
+    def test_nested_table_to_dataframe(self):
+        SF = bigquery.SchemaField
+        schema = [
+            SF('string_col', 'STRING', mode='NULLABLE'),
+            SF('record_col', 'RECORD', mode='NULLABLE', fields=[
+                SF('nested_string', 'STRING', mode='NULLABLE'),
+                SF('nested_repeated', 'INTEGER', mode='REPEATED'),
+                SF('nested_record', 'RECORD', mode='NULLABLE', fields=[
+                    SF('nested_nested_string', 'STRING', mode='NULLABLE'),
+                ]),
+            ]),
+        ]
+        record = {
+            'nested_string': 'another string value',
+            'nested_repeated': [0, 1, 2],
+            'nested_record': {'nested_nested_string': 'some deep insight'},
+        }
+        to_insert = [
+            ('Some value', record)
+        ]
+        table_id = 'test_table'
+        dataset = self.temp_dataset(_make_dataset_id('nested_df'))
+        table_arg = Table(dataset.table(table_id), schema=schema)
+        table = retry_403(Config.CLIENT.create_table)(table_arg)
+        self.to_delete.insert(0, table)
+        Config.CLIENT.create_rows(table, to_insert)
+        QUERY = 'SELECT * from `{}.{}.{}`'.format(
+            Config.CLIENT.project, dataset.dataset_id, table_id)
+
+        retry = RetryResult(_has_rows, max_tries=8)
+        df = retry(self._fetch_dataframe)(QUERY)
+
+        self.assertIsInstance(df, pandas.DataFrame)
+        self.assertEqual(len(df), 1)  # verify the number of rows
+        exp_columns = ['string_col', 'record_col']
+        self.assertEqual(list(df), exp_columns)  # verify the column names
+        row = df.iloc[0]
+        # verify the row content
+        self.assertEqual(row['string_col'], 'Some value')
+        self.assertEqual(row['record_col'], record)
+        # verify that nested data can be accessed with indices/keys
+        self.assertEqual(row['record_col']['nested_repeated'][0], 0)
+        self.assertEqual(
+            row['record_col']['nested_record']['nested_nested_string'],
+            'some deep insight')
+
     def temp_dataset(self, dataset_id):
         dataset = retry_403(Config.CLIENT.create_dataset)(
             Dataset(Config.CLIENT.dataset(dataset_id)))