@@ -2720,6 +2720,70 @@ def test_reload_w_alternate_client(self):
2720
2720
self .assertEqual (req ['path' ], PATH )
2721
2721
self ._verifyResourceProperties (job , RESOURCE )
2722
2722
2723
+ def test_to_dataframe (self ):
2724
+ import pandas as pd
2725
+
2726
+ begun_resource = self ._makeResource ()
2727
+ query_resource = {
2728
+ 'jobComplete' : True ,
2729
+ 'jobReference' : {
2730
+ 'projectId' : self .PROJECT ,
2731
+ 'jobId' : self .JOB_ID ,
2732
+ },
2733
+ 'schema' : {
2734
+ 'fields' : [
2735
+ {'name' : 'name' , 'type' : 'STRING' , 'mode' : 'NULLABLE' },
2736
+ {'name' : 'age' , 'type' : 'INTEGER' , 'mode' : 'NULLABLE' },
2737
+ ],
2738
+ },
2739
+ 'rows' : [
2740
+ {'f' : [{'v' : 'Phred Phlyntstone' }, {'v' : '32' }]},
2741
+ {'f' : [{'v' : 'Bharney Rhubble' }, {'v' : '33' }]},
2742
+ {'f' : [{'v' : 'Wylma Phlyntstone' }, {'v' : '29' }]},
2743
+ {'f' : [{'v' : 'Bhettye Rhubble' }, {'v' : '27' }]},
2744
+ ],
2745
+ }
2746
+ done_resource = copy .deepcopy (begun_resource )
2747
+ done_resource ['status' ] = {'state' : 'DONE' }
2748
+ connection = _Connection (
2749
+ begun_resource , query_resource , done_resource , query_resource )
2750
+ client = _make_client (project = self .PROJECT , connection = connection )
2751
+ job = self ._make_one (self .JOB_ID , self .QUERY , client )
2752
+ df = job .to_dataframe ()
2753
+
2754
+ self .assertIsInstance (df , pd .DataFrame )
2755
+ self .assertEqual (len (df ), 4 )
2756
+ self .assertEqual (list (df ), ['name' , 'age' ])
2757
+
2758
+ def test_to_dataframe_w_empty_results (self ):
2759
+ import pandas as pd
2760
+
2761
+ begun_resource = self ._makeResource ()
2762
+ query_resource = {
2763
+ 'jobComplete' : True ,
2764
+ 'jobReference' : {
2765
+ 'projectId' : self .PROJECT ,
2766
+ 'jobId' : self .JOB_ID ,
2767
+ },
2768
+ 'schema' : {
2769
+ 'fields' : [
2770
+ {'name' : 'name' , 'type' : 'STRING' , 'mode' : 'NULLABLE' },
2771
+ {'name' : 'age' , 'type' : 'INTEGER' , 'mode' : 'NULLABLE' },
2772
+ ],
2773
+ },
2774
+ }
2775
+ done_resource = copy .deepcopy (begun_resource )
2776
+ done_resource ['status' ] = {'state' : 'DONE' }
2777
+ connection = _Connection (
2778
+ begun_resource , query_resource , done_resource , query_resource )
2779
+ client = _make_client (project = self .PROJECT , connection = connection )
2780
+ job = self ._make_one (self .JOB_ID , self .QUERY , client )
2781
+ df = job .to_dataframe ()
2782
+
2783
+ self .assertIsInstance (df , pd .DataFrame )
2784
+ self .assertEqual (len (df ), 0 )
2785
+ self .assertEqual (list (df ), ['name' , 'age' ])
2786
+
2723
2787
2724
2788
class TestQueryPlanEntryStep (unittest .TestCase , _Base ):
2725
2789
KIND = 'KIND'
0 commit comments