Batch get_stat_all calls if more than QUERY_BATCH_SIZE places. (#155)

Also replace mixer reference with REST API reference.

Author: tjann

datacommons/examples/stat_vars.py

@@ -126,6 +126,37 @@ def call_str(pvs):
         dc.get_stat_all(["badPlaceId", "country/FRA"],
                         ["Median_Age_Person", "Count_Person"]))
 
+    print('\nSTRESS TEST FOR GET_STAT_ALL')
+    try:
+        dc.get_stat_all(
+            dc.get_places_in(['country/USA'], 'County')['country/USA'], [
+                'Count_Person_Unemployed', 'Median_Age_Person', 'Count_Death',
+                'Count_Death_CertainInfectiousParasiticDiseases',
+                'Count_UnemploymentInsuranceClaim_StateUnemploymentInsurance',
+                'Count_Worker_NAICSAccommodationFoodServices',
+                'Count_Household_With0AvailableVehicles',
+                'Count_Person_WhiteAlone', 'Count_Person_WorkedFullTime',
+                'Count_Person_Employed', 'Count_Person_EnrolledInSchool',
+                'Count_Person_EnrolledInGrade9ToGrade12',
+                'Count_Person_ResidesInGroupQuarters',
+                'Count_Person_NowMarried',
+                'RetailDrugDistribution_DrugDistribution_DMethamphetamine',
+                'Count_Household_With2Person',
+                'Count_Person_EducationalAttainment2NdGrade',
+                'Count_Person_EducationalAttainmentGedOrAlternativeCredential',
+                'Count_Person_EducationalAttainmentKindergarten',
+                'Count_Person_EducationalAttainmentRegularHighSchoolDiploma',
+                'Count_Death_DiseasesOfBloodAndBloodFormingOrgansAndImmuneDisorders',
+                'Count_Household_NoComputer',
+                'Median_Income_Household_HouseholderRaceHispanicOrLatino,',
+                'Count_HousingUnit_RenterOccupied',
+                'Count_Worker_NAICSInformation'
+            ])
+    except ValueError:
+        print('Stress test for get_stat_all FAILED!')
+    else:
+        print('Stress test for get_stat_all succeeded.')
+
 
 if __name__ == '__main__':
     main()

datacommons/stat_vars.py

@@ -203,17 +203,46 @@ def get_stat_all(places, stat_vars):
       }
     """
     url = utils._API_ROOT + utils._API_ENDPOINTS['get_stat_all']
-    req_json = {'stat_vars': stat_vars, 'places': places}
-
-    # Send the request
-    res_json = utils._send_request(url, req_json=req_json, use_payload=False)
-
-    if 'placeData' not in res_json:
-        raise ValueError('No data in response.')
-
-    # Unnest the REST response for keys that have single-element values.
-    place_statvar_series = collections.defaultdict(dict)
-    for place_dcid, place in res_json['placeData'].items():
-        for stat_var_dcid, stat_var in place['statVarData'].items():
-            place_statvar_series[place_dcid][stat_var_dcid] = stat_var
-    return dict(place_statvar_series)
+    places = list(places)
+    # Get number of batches via an arithmetic ceiling trick:
+    # 11//10 rounds down to 1.
+    # -11//10 rounds down to -2.
+    # We can divide with, then remove the negative to get the ceiling.
+    batches = -(-len(places) // utils._QUERY_BATCH_SIZE)
+    res = {}
+    for i in range(batches):
+        req_json = {
+            'stat_vars':
+                stat_vars,
+            'places':
+                places[i * utils._QUERY_BATCH_SIZE:(i + 1) *
+                       utils._QUERY_BATCH_SIZE]
+        }
+        # Send the request
+        res_json = utils._send_request(url,
+                                       req_json=req_json,
+                                       use_payload=False)
+        if 'placeData' not in res_json:
+            # The REST API spec will always return a dictionary under
+            # placeData, even if no places exist or have no
+            # data. If no Places are provided, REST will return an
+            # error, which will have been caught and passed on in
+            # _send_request.
+            raise ValueError("Unexpected response from REST stat/all API.")
+
+        # Unnest the REST response for keys that have single-element values.
+        place_statvar_series = collections.defaultdict(dict)
+        for place_dcid, place in res_json['placeData'].items():
+            stat_var_data = place.get('statVarData')
+            if not stat_var_data:
+                # The REST API spec will always return a dictionary under
+                # statVarData, even if no StatVars exist or have no
+                # data. If no StatVars are provided, REST will return an
+                # error, which will have been caught and passed on in
+                # _send_request.
+                raise ValueError("Unexpected response from REST stat/all API.")
+            for stat_var_dcid, stat_var in stat_var_data.items():
+                place_statvar_series[place_dcid][stat_var_dcid] = stat_var
+        res.update(dict(place_statvar_series))
+
+    return res

datacommons/utils.py

@@ -111,12 +111,12 @@ def _send_request(req_url, req_json={}, compress=False, post=True, use_payload=T
     res = six.moves.urllib.request.urlopen(req)
   except six.moves.urllib.error.HTTPError as e:
     raise ValueError(
-        'Response error: An HTTP {} code was returned by the mixer. Printing '
-        'response\n\n{}'.format(e.code, e.read()))
+        'Response error: An HTTP {} code was returned by the REST API. '
+        'Printing response\n\n{}'.format(e.code, e.read()))
   if isinstance(res, six.moves.urllib.error.HTTPError):
       raise ValueError(
-          'Response error: An HTTP {} code was returned by the mixer. Printing '
-          'response\n\n{}'.format(res.code, res.msg))
+          'Response error: An HTTP {} code was returned by the REST API. '
+          'Printing response\n\n{}'.format(res.code, res.msg))
   # Get the JSON
   res_json = json.loads(res.read())
   if not use_payload: