# coding=utf-8
# Copyright 2021 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Largest Covid-19 epidemiological database with a powerful set of expansive covariates.

It includes open, publicly sourced, licensed data relating to demographics,
economy, epidemiology, geography, health, hospitalizations, mobility, government
response, weather, and more.
"""

import tensorflow as tf
import tensorflow_datasets.public_api as tfds

_DESCRIPTION = """
This repository attempts to assemble the largest Covid-19 epidemiological database
in addition to a powerful set of expansive covariates. It includes open, publicly sourced,
licensed data relating to demographics, economy, epidemiology, geography, health, hospitalizations,
mobility, government response, weather, and more.

This particular dataset corresponds to a join of all the different tables that are part of the repository.
Therefore, expect the resulting samples to be highly sparse.
"""

_CITATION = """
@article{Wahltinez2020,
  author = "O. Wahltinez and others",
  year = 2020,
  title = "COVID-19 Open-Data: curating a fine-grained, global-scale data repository for SARS-CoV-2",
  note = "Work in progress",
  url = {https://goo.gle/covid-19-open-data},
}
"""

_N_RECORDS = 11266423  # It should be an upper bound
_BATCH_SIZE = 10000


class Covid19(tfds.core.GeneratorBasedBuilder):
  """DatasetBuilder for covid19 dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
  }

  def _info(self) -> tfds.core.DatasetInfo:
    """Returns the dataset metadata."""
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            'location_key':
                tf.string,
            'date':
                tf.string,
            'place_id':
                tf.string,
            'wikidata_id':
                tf.string,
            'datacommons_id':
                tf.string,
            'country_code':
                tf.string,
            'country_name':
                tf.string,
            'iso_3166_1_alpha_2':
                tf.string,
            'iso_3166_1_alpha_3':
                tf.string,
            'aggregation_level':
                tf.float64,
            'new_confirmed':
                tf.float64,
            'new_deceased':
                tf.float64,
            'cumulative_confirmed':
                tf.float64,
            'cumulative_deceased':
                tf.float64,
            'cumulative_tested':
                tf.float64,
            'new_persons_vaccinated':
                tf.float64,
            'cumulative_persons_vaccinated':
                tf.float64,
            'new_persons_fully_vaccinated':
                tf.float64,
            'cumulative_persons_fully_vaccinated':
                tf.float64,
            'new_vaccine_doses_administered':
                tf.float64,
            'cumulative_vaccine_doses_administered':
                tf.float64,
            'population':
                tf.float64,
            'population_male':
                tf.float64,
            'population_female':
                tf.float64,
            'population_rural':
                tf.float64,
            'population_urban':
                tf.float64,
            'population_density':
                tf.float64,
            'human_development_index':
                tf.float64,
            'population_age_00_09':
                tf.float64,
            'population_age_10_19':
                tf.float64,
            'population_age_20_29':
                tf.float64,
            'population_age_30_39':
                tf.float64,
            'population_age_40_49':
                tf.float64,
            'population_age_50_59':
                tf.float64,
            'population_age_60_69':
                tf.float64,
            'population_age_70_79':
                tf.float64,
            'population_age_80_and_older':
                tf.float64,
            'gdp_usd':
                tf.float64,
            'gdp_per_capita_usd':
                tf.float64,
            'openstreetmap_id':
                tf.string,
            'latitude':
                tf.float64,
            'longitude':
                tf.float64,
            'area_sq_km':
                tf.float64,
            'smoking_prevalence':
                tf.float64,
            'diabetes_prevalence':
                tf.float64,
            'infant_mortality_rate':
                tf.float64,
            'nurses_per_1000':
                tf.float64,
            'physicians_per_1000':
                tf.float64,
            'health_expenditure_usd':
                tf.float64,
            'out_of_pocket_health_expenditure_usd':
                tf.float64,
            'school_closing':
                tf.float64,
            'workplace_closing':
                tf.float64,
            'cancel_public_events':
                tf.float64,
            'restrictions_on_gatherings':
                tf.float64,
            'public_transport_closing':
                tf.float64,
            'stay_at_home_requirements':
                tf.float64,
            'restrictions_on_internal_movement':
                tf.float64,
            'international_travel_controls':
                tf.float64,
            'income_support':
                tf.float64,
            'debt_relief':
                tf.float64,
            'fiscal_measures':
                tf.float64,
            'international_support':
                tf.float64,
            'public_information_campaigns':
                tf.float64,
            'testing_policy':
                tf.float64,
            'contact_tracing':
                tf.float64,
            'emergency_investment_in_healthcare':
                tf.float64,
            'investment_in_vaccines':
                tf.float64,
            'facial_coverings':
                tf.float64,
            'vaccination_policy':
                tf.float64,
            'stringency_index':
                tf.float64,
            'average_temperature_celsius':
                tf.float64,
            'minimum_temperature_celsius':
                tf.float64,
            'maximum_temperature_celsius':
                tf.float64,
            'rainfall_mm':
                tf.float64,
            'snowfall_mm':
                tf.float64,
            'dew_point':
                tf.float64,
            'relative_humidity':
                tf.float64,
            'new_tested':
                tf.float64,
            'population_largest_city':
                tf.float64,
            'population_clustered':
                tf.float64,
            'human_capital_index':
                tf.float64,
            'area_rural_sq_km':
                tf.float64,
            'area_urban_sq_km':
                tf.float64,
            'life_expectancy':
                tf.float64,
            'adult_male_mortality_rate':
                tf.float64,
            'adult_female_mortality_rate':
                tf.float64,
            'pollution_mortality_rate':
                tf.float64,
            'comorbidity_mortality_rate':
                tf.float64,
            'mobility_retail_and_recreation':
                tf.float64,
            'mobility_grocery_and_pharmacy':
                tf.float64,
            'mobility_parks':
                tf.float64,
            'mobility_transit_stations':
                tf.float64,
            'mobility_workplaces':
                tf.float64,
            'mobility_residential':
                tf.float64,
            'hospital_beds_per_1000':
                tf.float64,
            'subregion1_code':
                tf.string,
            'subregion1_name':
                tf.string,
            'new_recovered':
                tf.float64,
            'cumulative_recovered':
                tf.float64,
            'elevation_m':
                tf.float64,
            'new_hospitalized_patients':
                tf.float64,
            'cumulative_hospitalized_patients':
                tf.float64,
            'new_intensive_care_patients':
                tf.float64,
            'cumulative_intensive_care_patients':
                tf.float64,
            'new_confirmed_age_0':
                tf.float64,
            'new_confirmed_age_1':
                tf.float64,
            'new_confirmed_age_2':
                tf.float64,
            'new_confirmed_age_3':
                tf.float64,
            'new_confirmed_age_4':
                tf.float64,
            'new_confirmed_age_5':
                tf.float64,
            'new_confirmed_age_6':
                tf.float64,
            'new_confirmed_age_7':
                tf.float64,
            'new_confirmed_age_8':
                tf.float64,
            'new_confirmed_age_9':
                tf.float64,
            'cumulative_confirmed_age_0':
                tf.float64,
            'cumulative_confirmed_age_1':
                tf.float64,
            'cumulative_confirmed_age_2':
                tf.float64,
            'cumulative_confirmed_age_3':
                tf.float64,
            'cumulative_confirmed_age_4':
                tf.float64,
            'cumulative_confirmed_age_5':
                tf.float64,
            'cumulative_confirmed_age_6':
                tf.float64,
            'cumulative_confirmed_age_7':
                tf.float64,
            'cumulative_confirmed_age_8':
                tf.float64,
            'cumulative_confirmed_age_9':
                tf.float64,
            'new_deceased_age_0':
                tf.float64,
            'new_deceased_age_1':
                tf.float64,
            'new_deceased_age_2':
                tf.float64,
            'new_deceased_age_3':
                tf.float64,
            'new_deceased_age_4':
                tf.float64,
            'new_deceased_age_5':
                tf.float64,
            'new_deceased_age_6':
                tf.float64,
            'new_deceased_age_7':
                tf.float64,
            'new_deceased_age_8':
                tf.float64,
            'new_deceased_age_9':
                tf.float64,
            'cumulative_deceased_age_0':
                tf.float64,
            'cumulative_deceased_age_1':
                tf.float64,
            'cumulative_deceased_age_2':
                tf.float64,
            'cumulative_deceased_age_3':
                tf.float64,
            'cumulative_deceased_age_4':
                tf.float64,
            'cumulative_deceased_age_5':
                tf.float64,
            'cumulative_deceased_age_6':
                tf.float64,
            'cumulative_deceased_age_7':
                tf.float64,
            'cumulative_deceased_age_8':
                tf.float64,
            'cumulative_deceased_age_9':
                tf.float64,
            'new_tested_age_0':
                tf.float64,
            'new_tested_age_1':
                tf.float64,
            'new_tested_age_2':
                tf.float64,
            'new_tested_age_3':
                tf.float64,
            'new_tested_age_4':
                tf.float64,
            'new_tested_age_5':
                tf.float64,
            'new_tested_age_6':
                tf.float64,
            'new_tested_age_7':
                tf.float64,
            'new_tested_age_8':
                tf.float64,
            'new_tested_age_9':
                tf.float64,
            'cumulative_tested_age_0':
                tf.float64,
            'cumulative_tested_age_1':
                tf.float64,
            'cumulative_tested_age_2':
                tf.float64,
            'cumulative_tested_age_3':
                tf.float64,
            'cumulative_tested_age_4':
                tf.float64,
            'cumulative_tested_age_5':
                tf.float64,
            'cumulative_tested_age_6':
                tf.float64,
            'cumulative_tested_age_7':
                tf.float64,
            'cumulative_tested_age_8':
                tf.float64,
            'cumulative_tested_age_9':
                tf.float64,
            'new_hospitalized_patients_age_0':
                tf.float64,
            'new_hospitalized_patients_age_1':
                tf.float64,
            'new_hospitalized_patients_age_2':
                tf.float64,
            'new_hospitalized_patients_age_3':
                tf.float64,
            'new_hospitalized_patients_age_4':
                tf.float64,
            'new_hospitalized_patients_age_5':
                tf.float64,
            'new_hospitalized_patients_age_6':
                tf.float64,
            'new_hospitalized_patients_age_7':
                tf.float64,
            'new_hospitalized_patients_age_8':
                tf.float64,
            'new_hospitalized_patients_age_9':
                tf.float64,
            'cumulative_hospitalized_patients_age_0':
                tf.float64,
            'cumulative_hospitalized_patients_age_1':
                tf.float64,
            'cumulative_hospitalized_patients_age_2':
                tf.float64,
            'cumulative_hospitalized_patients_age_3':
                tf.float64,
            'cumulative_hospitalized_patients_age_4':
                tf.float64,
            'cumulative_hospitalized_patients_age_5':
                tf.float64,
            'cumulative_hospitalized_patients_age_6':
                tf.float64,
            'cumulative_hospitalized_patients_age_7':
                tf.float64,
            'cumulative_hospitalized_patients_age_8':
                tf.float64,
            'cumulative_hospitalized_patients_age_9':
                tf.float64,
            'new_intensive_care_patients_age_0':
                tf.float64,
            'new_intensive_care_patients_age_1':
                tf.float64,
            'new_intensive_care_patients_age_2':
                tf.float64,
            'new_intensive_care_patients_age_3':
                tf.float64,
            'new_intensive_care_patients_age_4':
                tf.float64,
            'new_intensive_care_patients_age_5':
                tf.float64,
            'new_intensive_care_patients_age_6':
                tf.float64,
            'new_intensive_care_patients_age_7':
                tf.float64,
            'new_intensive_care_patients_age_8':
                tf.float64,
            'new_intensive_care_patients_age_9':
                tf.float64,
            'cumulative_intensive_care_patients_age_0':
                tf.float64,
            'cumulative_intensive_care_patients_age_1':
                tf.float64,
            'cumulative_intensive_care_patients_age_2':
                tf.float64,
            'cumulative_intensive_care_patients_age_3':
                tf.float64,
            'cumulative_intensive_care_patients_age_4':
                tf.float64,
            'cumulative_intensive_care_patients_age_5':
                tf.float64,
            'cumulative_intensive_care_patients_age_6':
                tf.float64,
            'cumulative_intensive_care_patients_age_7':
                tf.float64,
            'cumulative_intensive_care_patients_age_8':
                tf.float64,
            'cumulative_intensive_care_patients_age_9':
                tf.float64,
            # Stored as string because values look like "30-39"
            'age_bin_0':
                tf.string,
            'age_bin_1':
                tf.string,
            'age_bin_2':
                tf.string,
            'age_bin_3':
                tf.string,
            'age_bin_4':
                tf.string,
            'age_bin_5':
                tf.string,
            'age_bin_6':
                tf.string,
            'age_bin_7':
                tf.string,
            'age_bin_8':
                tf.string,
            'age_bin_9':
                tf.string,
            'new_confirmed_male':
                tf.float64,
            'new_confirmed_female':
                tf.float64,
            'cumulative_confirmed_male':
                tf.float64,
            'cumulative_confirmed_female':
                tf.float64,
            'new_deceased_male':
                tf.float64,
            'new_deceased_female':
                tf.float64,
            'cumulative_deceased_male':
                tf.float64,
            'cumulative_deceased_female':
                tf.float64,
            'new_tested_male':
                tf.float64,
            'new_tested_female':
                tf.float64,
            'cumulative_tested_male':
                tf.float64,
            'cumulative_tested_female':
                tf.float64,
            'new_hospitalized_patients_male':
                tf.float64,
            'new_hospitalized_patients_female':
                tf.float64,
            'cumulative_hospitalized_patients_male':
                tf.float64,
            'cumulative_hospitalized_patients_female':
                tf.float64,
            'new_intensive_care_patients_male':
                tf.float64,
            'new_intensive_care_patients_female':
                tf.float64,
            'cumulative_intensive_care_patients_male':
                tf.float64,
            'cumulative_intensive_care_patients_female':
                tf.float64,
            'subregion2_code':
                tf.string,
            'subregion2_name':
                tf.string,
            'current_hospitalized_patients':
                tf.float64,
            'current_intensive_care_patients':
                tf.float64,
            'current_ventilator_patients':
                tf.float64,
            'search_trends_abdominal_obesity':
                tf.float64,
            'search_trends_abdominal_pain':
                tf.float64,
            'search_trends_acne':
                tf.float64,
            'search_trends_actinic_keratosis':
                tf.float64,
            'search_trends_acute_bronchitis':
                tf.float64,
            'search_trends_adrenal_crisis':
                tf.float64,
            'search_trends_ageusia':
                tf.float64,
            'search_trends_alcoholism':
                tf.float64,
            'search_trends_allergic_conjunctivitis':
                tf.float64,
            'search_trends_allergy':
                tf.float64,
            'search_trends_amblyopia':
                tf.float64,
            'search_trends_amenorrhea':
                tf.float64,
            'search_trends_amnesia':
                tf.float64,
            'search_trends_anal_fissure':
                tf.float64,
            'search_trends_anaphylaxis':
                tf.float64,
            'search_trends_anemia':
                tf.float64,
            'search_trends_angina_pectoris':
                tf.float64,
            'search_trends_angioedema':
                tf.float64,
            'search_trends_angular_cheilitis':
                tf.float64,
            'search_trends_anosmia':
                tf.float64,
            'search_trends_anxiety':
                tf.float64,
            'search_trends_aphasia':
                tf.float64,
            'search_trends_aphonia':
                tf.float64,
            'search_trends_apnea':
                tf.float64,
            'search_trends_arthralgia':
                tf.float64,
            'search_trends_arthritis':
                tf.float64,
            'search_trends_ascites':
                tf.float64,
            'search_trends_asperger_syndrome':
                tf.float64,
            'search_trends_asphyxia':
                tf.float64,
            'search_trends_asthma':
                tf.float64,
            'search_trends_astigmatism':
                tf.float64,
            'search_trends_ataxia':
                tf.float64,
            'search_trends_atheroma':
                tf.float64,
            'search_trends_attention_deficit_hyperactivity_disorder':
                tf.float64,
            'search_trends_auditory_hallucination':
                tf.float64,
            'search_trends_autoimmune_disease':
                tf.float64,
            'search_trends_avoidant_personality_disorder':
                tf.float64,
            'search_trends_back_pain':
                tf.float64,
            'search_trends_bacterial_vaginosis':
                tf.float64,
            'search_trends_balance_disorder':
                tf.float64,
            'search_trends_beaus_lines':
                tf.float64,
            'search_trends_bells_palsy':
                tf.float64,
            'search_trends_biliary_colic':
                tf.float64,
            'search_trends_binge_eating':
                tf.float64,
            'search_trends_bleeding':
                tf.float64,
            'search_trends_bleeding_on_probing':
                tf.float64,
            'search_trends_blepharospasm':
                tf.float64,
            'search_trends_bloating':
                tf.float64,
            'search_trends_blood_in_stool':
                tf.float64,
            'search_trends_blurred_vision':
                tf.float64,
            'search_trends_blushing':
                tf.float64,
            'search_trends_boil':
                tf.float64,
            'search_trends_bone_fracture':
                tf.float64,
            'search_trends_bone_tumor':
                tf.float64,
            'search_trends_bowel_obstruction':
                tf.float64,
            'search_trends_bradycardia':
                tf.float64,
            'search_trends_braxton_hicks_contractions':
                tf.float64,
            'search_trends_breakthrough_bleeding':
                tf.float64,
            'search_trends_breast_pain':
                tf.float64,
            'search_trends_bronchitis':
                tf.float64,
            'search_trends_bruise':
                tf.float64,
            'search_trends_bruxism':
                tf.float64,
            'search_trends_bunion':
                tf.float64,
            'search_trends_burn':
                tf.float64,
            'search_trends_burning_chest_pain':
                tf.float64,
            'search_trends_burning_mouth_syndrome':
                tf.float64,
            'search_trends_candidiasis':
                tf.float64,
            'search_trends_canker_sore':
                tf.float64,
            'search_trends_cardiac_arrest':
                tf.float64,
            'search_trends_carpal_tunnel_syndrome':
                tf.float64,
            'search_trends_cataplexy':
                tf.float64,
            'search_trends_cataract':
                tf.float64,
            'search_trends_chancre':
                tf.float64,
            'search_trends_cheilitis':
                tf.float64,
            'search_trends_chest_pain':
                tf.float64,
            'search_trends_chills':
                tf.float64,
            'search_trends_chorea':
                tf.float64,
            'search_trends_chronic_pain':
                tf.float64,
            'search_trends_cirrhosis':
                tf.float64,
            'search_trends_cleft_lip_and_cleft_palate':
                tf.float64,
            'search_trends_clouding_of_consciousness':
                tf.float64,
            'search_trends_cluster_headache':
                tf.float64,
            'search_trends_colitis':
                tf.float64,
            'search_trends_coma':
                tf.float64,
            'search_trends_common_cold':
                tf.float64,
            'search_trends_compulsive_behavior':
                tf.float64,
            'search_trends_compulsive_hoarding':
                tf.float64,
            'search_trends_confusion':
                tf.float64,
            'search_trends_congenital_heart_defect':
                tf.float64,
            'search_trends_conjunctivitis':
                tf.float64,
            'search_trends_constipation':
                tf.float64,
            'search_trends_convulsion':
                tf.float64,
            'search_trends_cough':
                tf.float64,
            'search_trends_crackles':
                tf.float64,
            'search_trends_cramp':
                tf.float64,
            'search_trends_crepitus':
                tf.float64,
            'search_trends_croup':
                tf.float64,
            'search_trends_cyanosis':
                tf.float64,
            'search_trends_dandruff':
                tf.float64,
            'search_trends_delayed_onset_muscle_soreness':
                tf.float64,
            'search_trends_dementia':
                tf.float64,
            'search_trends_dentin_hypersensitivity':
                tf.float64,
            'search_trends_depersonalization':
                tf.float64,
            'search_trends_depression':
                tf.float64,
            'search_trends_dermatitis':
                tf.float64,
            'search_trends_desquamation':
                tf.float64,
            'search_trends_developmental_disability':
                tf.float64,
            'search_trends_diabetes':
                tf.float64,
            'search_trends_diabetic_ketoacidosis':
                tf.float64,
            'search_trends_diarrhea':
                tf.float64,
            'search_trends_dizziness':
                tf.float64,
            'search_trends_dry_eye_syndrome':
                tf.float64,
            'search_trends_dysautonomia':
                tf.float64,
            'search_trends_dysgeusia':
                tf.float64,
            'search_trends_dysmenorrhea':
                tf.float64,
            'search_trends_dyspareunia':
                tf.float64,
            'search_trends_dysphagia':
                tf.float64,
            'search_trends_dysphoria':
                tf.float64,
            'search_trends_dystonia':
                tf.float64,
            'search_trends_dysuria':
                tf.float64,
            'search_trends_ear_pain':
                tf.float64,
            'search_trends_eczema':
                tf.float64,
            'search_trends_edema':
                tf.float64,
            'search_trends_encephalitis':
                tf.float64,
            'search_trends_encephalopathy':
                tf.float64,
            'search_trends_epidermoid_cyst':
                tf.float64,
            'search_trends_epilepsy':
                tf.float64,
            'search_trends_epiphora':
                tf.float64,
            'search_trends_erectile_dysfunction':
                tf.float64,
            'search_trends_erythema':
                tf.float64,
            'search_trends_erythema_chronicum_migrans':
                tf.float64,
            'search_trends_esophagitis':
                tf.float64,
            'search_trends_excessive_daytime_sleepiness':
                tf.float64,
            'search_trends_eye_pain':
                tf.float64,
            'search_trends_eye_strain':
                tf.float64,
            'search_trends_facial_nerve_paralysis':
                tf.float64,
            'search_trends_facial_swelling':
                tf.float64,
            'search_trends_fasciculation':
                tf.float64,
            'search_trends_fatigue':
                tf.float64,
            'search_trends_fatty_liver_disease':
                tf.float64,
            'search_trends_fecal_incontinence':
                tf.float64,
            'search_trends_fever':
                tf.float64,
            'search_trends_fibrillation':
                tf.float64,
            'search_trends_fibrocystic_breast_changes':
                tf.float64,
            'search_trends_fibromyalgia':
                tf.float64,
            'search_trends_flatulence':
                tf.float64,
            'search_trends_floater':
                tf.float64,
            'search_trends_focal_seizure':
                tf.float64,
            'search_trends_folate_deficiency':
                tf.float64,
            'search_trends_food_craving':
                tf.float64,
            'search_trends_food_intolerance':
                tf.float64,
            'search_trends_frequent_urination':
                tf.float64,
            'search_trends_gastroesophageal_reflux_disease':
                tf.float64,
            'search_trends_gastroparesis':
                tf.float64,
            'search_trends_generalized_anxiety_disorder':
                tf.float64,
            'search_trends_genital_wart':
                tf.float64,
            'search_trends_gingival_recession':
                tf.float64,
            'search_trends_gingivitis':
                tf.float64,
            'search_trends_globus_pharyngis':
                tf.float64,
            'search_trends_goitre':
                tf.float64,
            'search_trends_gout':
                tf.float64,
            'search_trends_grandiosity':
                tf.float64,
            'search_trends_granuloma':
                tf.float64,
            'search_trends_guilt':
                tf.float64,
            'search_trends_hair_loss':
                tf.float64,
            'search_trends_halitosis':
                tf.float64,
            'search_trends_hay_fever':
                tf.float64,
            'search_trends_headache':
                tf.float64,
            'search_trends_heart_arrhythmia':
                tf.float64,
            'search_trends_heart_murmur':
                tf.float64,
            'search_trends_heartburn':
                tf.float64,
            'search_trends_hematochezia':
                tf.float64,
            'search_trends_hematoma':
                tf.float64,
            'search_trends_hematuria':
                tf.float64,
            'search_trends_hemolysis':
                tf.float64,
            'search_trends_hemoptysis':
                tf.float64,
            'search_trends_hemorrhoids':
                tf.float64,
            'search_trends_hepatic_encephalopathy':
                tf.float64,
            'search_trends_hepatitis':
                tf.float64,
            'search_trends_hepatotoxicity':
                tf.float64,
            'search_trends_hiccup':
                tf.float64,
            'search_trends_hip_pain':
                tf.float64,
            'search_trends_hives':
                tf.float64,
            'search_trends_hot_flash':
                tf.float64,
            'search_trends_hydrocephalus':
                tf.float64,
            'search_trends_hypercalcaemia':
                tf.float64,
            'search_trends_hypercapnia':
                tf.float64,
            'search_trends_hypercholesterolemia':
                tf.float64,
            'search_trends_hyperemesis_gravidarum':
                tf.float64,
            'search_trends_hyperglycemia':
                tf.float64,
            'search_trends_hyperhidrosis':
                tf.float64,
            'search_trends_hyperkalemia':
                tf.float64,
            'search_trends_hyperlipidemia':
                tf.float64,
            'search_trends_hypermobility':
                tf.float64,
            'search_trends_hyperpigmentation':
                tf.float64,
            'search_trends_hypersomnia':
                tf.float64,
            'search_trends_hypertension':
                tf.float64,
            'search_trends_hyperthermia':
                tf.float64,
            'search_trends_hyperthyroidism':
                tf.float64,
            'search_trends_hypertriglyceridemia':
                tf.float64,
            'search_trends_hypertrophy':
                tf.float64,
            'search_trends_hyperventilation':
                tf.float64,
            'search_trends_hypocalcaemia':
                tf.float64,
            'search_trends_hypochondriasis':
                tf.float64,
            'search_trends_hypoglycemia':
                tf.float64,
            'search_trends_hypogonadism':
                tf.float64,
            'search_trends_hypokalemia':
                tf.float64,
            'search_trends_hypomania':
                tf.float64,
            'search_trends_hyponatremia':
                tf.float64,
            'search_trends_hypotension':
                tf.float64,
            'search_trends_hypothyroidism':
                tf.float64,
            'search_trends_hypoxemia':
                tf.float64,
            'search_trends_hypoxia':
                tf.float64,
            'search_trends_impetigo':
                tf.float64,
            'search_trends_implantation_bleeding':
                tf.float64,
            'search_trends_impulsivity':
                tf.float64,
            'search_trends_indigestion':
                tf.float64,
            'search_trends_infection':
                tf.float64,
            'search_trends_inflammation':
                tf.float64,
            'search_trends_inflammatory_bowel_disease':
                tf.float64,
            'search_trends_ingrown_hair':
                tf.float64,
            'search_trends_insomnia':
                tf.float64,
            'search_trends_insulin_resistance':
                tf.float64,
            'search_trends_intermenstrual_bleeding':
                tf.float64,
            'search_trends_intracranial_pressure':
                tf.float64,
            'search_trends_iron_deficiency':
                tf.float64,
            'search_trends_irregular_menstruation':
                tf.float64,
            'search_trends_itch':
                tf.float64,
            'search_trends_jaundice':
                tf.float64,
            'search_trends_kidney_failure':
                tf.float64,
            'search_trends_kidney_stone':
                tf.float64,
            'search_trends_knee_pain':
                tf.float64,
            'search_trends_kyphosis':
                tf.float64,
            'search_trends_lactose_intolerance':
                tf.float64,
            'search_trends_laryngitis':
                tf.float64,
            'search_trends_leg_cramps':
                tf.float64,
            'search_trends_lesion':
                tf.float64,
            'search_trends_leukorrhea':
                tf.float64,
            'search_trends_lightheadedness':
                tf.float64,
            'search_trends_low_back_pain':
                tf.float64,
            'search_trends_low_grade_fever':
                tf.float64,
            'search_trends_lymphedema':
                tf.float64,
            'search_trends_major_depressive_disorder':
                tf.float64,
            'search_trends_malabsorption':
                tf.float64,
            'search_trends_male_infertility':
                tf.float64,
            'search_trends_manic_disorder':
                tf.float64,
            'search_trends_melasma':
                tf.float64,
            'search_trends_melena':
                tf.float64,
            'search_trends_meningitis':
                tf.float64,
            'search_trends_menorrhagia':
                tf.float64,
            'search_trends_middle_back_pain':
                tf.float64,
            'search_trends_migraine':
                tf.float64,
            'search_trends_milium':
                tf.float64,
            'search_trends_mitral_insufficiency':
                tf.float64,
            'search_trends_mood_disorder':
                tf.float64,
            'search_trends_mood_swing':
                tf.float64,
            'search_trends_morning_sickness':
                tf.float64,
            'search_trends_motion_sickness':
                tf.float64,
            'search_trends_mouth_ulcer':
                tf.float64,
            'search_trends_muscle_atrophy':
                tf.float64,
            'search_trends_muscle_weakness':
                tf.float64,
            'search_trends_myalgia':
                tf.float64,
            'search_trends_mydriasis':
                tf.float64,
            'search_trends_myocardial_infarction':
                tf.float64,
            'search_trends_myoclonus':
                tf.float64,
            'search_trends_nasal_congestion':
                tf.float64,
            'search_trends_nasal_polyp':
                tf.float64,
            'search_trends_nausea':
                tf.float64,
            'search_trends_neck_mass':
                tf.float64,
            'search_trends_neck_pain':
                tf.float64,
            'search_trends_neonatal_jaundice':
                tf.float64,
            'search_trends_nerve_injury':
                tf.float64,
            'search_trends_neuralgia':
                tf.float64,
            'search_trends_neutropenia':
                tf.float64,
            'search_trends_night_sweats':
                tf.float64,
            'search_trends_night_terror':
                tf.float64,
            'search_trends_nocturnal_enuresis':
                tf.float64,
            'search_trends_nodule':
                tf.float64,
            'search_trends_nosebleed':
                tf.float64,
            'search_trends_nystagmus':
                tf.float64,
            'search_trends_obesity':
                tf.float64,
            'search_trends_onychorrhexis':
                tf.float64,
            'search_trends_oral_candidiasis':
                tf.float64,
            'search_trends_orthostatic_hypotension':
                tf.float64,
            'search_trends_osteopenia':
                tf.float64,
            'search_trends_osteophyte':
                tf.float64,
            'search_trends_osteoporosis':
                tf.float64,
            'search_trends_otitis':
                tf.float64,
            'search_trends_otitis_externa':
                tf.float64,
            'search_trends_otitis_media':
                tf.float64,
            'search_trends_pain':
                tf.float64,
            'search_trends_palpitations':
                tf.float64,
            'search_trends_pancreatitis':
                tf.float64,
            'search_trends_panic_attack':
                tf.float64,
            'search_trends_papule':
                tf.float64,
            'search_trends_paranoia':
                tf.float64,
            'search_trends_paresthesia':
                tf.float64,
            'search_trends_pelvic_inflammatory_disease':
                tf.float64,
            'search_trends_pericarditis':
                tf.float64,
            'search_trends_periodontal_disease':
                tf.float64,
            'search_trends_periorbital_puffiness':
                tf.float64,
            'search_trends_peripheral_neuropathy':
                tf.float64,
            'search_trends_perspiration':
                tf.float64,
            'search_trends_petechia':
                tf.float64,
            'search_trends_phlegm':
                tf.float64,
            'search_trends_photodermatitis':
                tf.float64,
            'search_trends_photophobia':
                tf.float64,
            'search_trends_photopsia':
                tf.float64,
            'search_trends_pleural_effusion':
                tf.float64,
            'search_trends_pleurisy':
                tf.float64,
            'search_trends_pneumonia':
                tf.float64,
            'search_trends_podalgia':
                tf.float64,
            'search_trends_polycythemia':
                tf.float64,
            'search_trends_polydipsia':
                tf.float64,
            'search_trends_polyneuropathy':
                tf.float64,
            'search_trends_polyuria':
                tf.float64,
            'search_trends_poor_posture':
                tf.float64,
            'search_trends_post_nasal_drip':
                tf.float64,
            'search_trends_postural_orthostatic_tachycardia_syndrome':
                tf.float64,
            'search_trends_prediabetes':
                tf.float64,
            'search_trends_proteinuria':
                tf.float64,
            'search_trends_pruritus_ani':
                tf.float64,
            'search_trends_psychosis':
                tf.float64,
            'search_trends_ptosis':
                tf.float64,
            'search_trends_pulmonary_edema':
                tf.float64,
            'search_trends_pulmonary_hypertension':
                tf.float64,
            'search_trends_purpura':
                tf.float64,
            'search_trends_pus':
                tf.float64,
            'search_trends_pyelonephritis':
                tf.float64,
            'search_trends_radiculopathy':
                tf.float64,
            'search_trends_rectal_pain':
                tf.float64,
            'search_trends_rectal_prolapse':
                tf.float64,
            'search_trends_red_eye':
                tf.float64,
            'search_trends_renal_colic':
                tf.float64,
            'search_trends_restless_legs_syndrome':
                tf.float64,
            'search_trends_rheum':
                tf.float64,
            'search_trends_rhinitis':
                tf.float64,
            'search_trends_rhinorrhea':
                tf.float64,
            'search_trends_rosacea':
                tf.float64,
            'search_trends_round_ligament_pain':
                tf.float64,
            'search_trends_rumination':
                tf.float64,
            'search_trends_scar':
                tf.float64,
            'search_trends_sciatica':
                tf.float64,
            'search_trends_scoliosis':
                tf.float64,
            'search_trends_seborrheic_dermatitis':
                tf.float64,
            'search_trends_self_harm':
                tf.float64,
            'search_trends_sensitivity_to_sound':
                tf.float64,
            'search_trends_sexual_dysfunction':
                tf.float64,
            'search_trends_shallow_breathing':
                tf.float64,
            'search_trends_sharp_pain':
                tf.float64,
            'search_trends_shivering':
                tf.float64,
            'search_trends_shortness_of_breath':
                tf.float64,
            'search_trends_shyness':
                tf.float64,
            'search_trends_sinusitis':
                tf.float64,
            'search_trends_skin_condition':
                tf.float64,
            'search_trends_skin_rash':
                tf.float64,
            'search_trends_skin_tag':
                tf.float64,
            'search_trends_skin_ulcer':
                tf.float64,
            'search_trends_sleep_apnea':
                tf.float64,
            'search_trends_sleep_deprivation':
                tf.float64,
            'search_trends_sleep_disorder':
                tf.float64,
            'search_trends_snoring':
                tf.float64,
            'search_trends_sore_throat':
                tf.float64,
            'search_trends_spasticity':
                tf.float64,
            'search_trends_splenomegaly':
                tf.float64,
            'search_trends_sputum':
                tf.float64,
            'search_trends_stomach_rumble':
                tf.float64,
            'search_trends_strabismus':
                tf.float64,
            'search_trends_stretch_marks':
                tf.float64,
            'search_trends_stridor':
                tf.float64,
            'search_trends_stroke':
                tf.float64,
            'search_trends_stuttering':
                tf.float64,
            'search_trends_subdural_hematoma':
                tf.float64,
            'search_trends_suicidal_ideation':
                tf.float64,
            'search_trends_swelling':
                tf.float64,
            'search_trends_swollen_feet':
                tf.float64,
            'search_trends_swollen_lymph_nodes':
                tf.float64,
            'search_trends_syncope':
                tf.float64,
            'search_trends_tachycardia':
                tf.float64,
            'search_trends_tachypnea':
                tf.float64,
            'search_trends_telangiectasia':
                tf.float64,
            'search_trends_tenderness':
                tf.float64,
            'search_trends_testicular_pain':
                tf.float64,
            'search_trends_throat_irritation':
                tf.float64,
            'search_trends_thrombocytopenia':
                tf.float64,
            'search_trends_thyroid_nodule':
                tf.float64,
            'search_trends_tic':
                tf.float64,
            'search_trends_tinnitus':
                tf.float64,
            'search_trends_tonsillitis':
                tf.float64,
            'search_trends_toothache':
                tf.float64,
            'search_trends_tremor':
                tf.float64,
            'search_trends_trichoptilosis':
                tf.float64,
            'search_trends_tumor':
                tf.float64,
            'search_trends_type_2_diabetes':
                tf.float64,
            'search_trends_unconsciousness':
                tf.float64,
            'search_trends_underweight':
                tf.float64,
            'search_trends_upper_respiratory_tract_infection':
                tf.float64,
            'search_trends_urethritis':
                tf.float64,
            'search_trends_urinary_incontinence':
                tf.float64,
            'search_trends_urinary_tract_infection':
                tf.float64,
            'search_trends_urinary_urgency':
                tf.float64,
            'search_trends_uterine_contraction':
                tf.float64,
            'search_trends_vaginal_bleeding':
                tf.float64,
            'search_trends_vaginal_discharge':
                tf.float64,
            'search_trends_vaginitis':
                tf.float64,
            'search_trends_varicose_veins':
                tf.float64,
            'search_trends_vasculitis':
                tf.float64,
            'search_trends_ventricular_fibrillation':
                tf.float64,
            'search_trends_ventricular_tachycardia':
                tf.float64,
            'search_trends_vertigo':
                tf.float64,
            'search_trends_viral_pneumonia':
                tf.float64,
            'search_trends_visual_acuity':
                tf.float64,
            'search_trends_vomiting':
                tf.float64,
            'search_trends_wart':
                tf.float64,
            'search_trends_water_retention':
                tf.float64,
            'search_trends_weakness':
                tf.float64,
            'search_trends_weight_gain':
                tf.float64,
            'search_trends_wheeze':
                tf.float64,
            'search_trends_xeroderma':
                tf.float64,
            'search_trends_xerostomia':
                tf.float64,
            'search_trends_yawn':
                tf.float64,
            'new_recovered_age_0':
                tf.float64,
            'new_recovered_age_1':
                tf.float64,
            'new_recovered_age_2':
                tf.float64,
            'new_recovered_age_3':
                tf.float64,
            'new_recovered_age_4':
                tf.float64,
            'new_recovered_age_5':
                tf.float64,
            'new_recovered_age_6':
                tf.float64,
            'new_recovered_age_7':
                tf.float64,
            'new_recovered_age_8':
                tf.float64,
            'new_recovered_age_9':
                tf.float64,
            'cumulative_recovered_age_0':
                tf.float64,
            'cumulative_recovered_age_1':
                tf.float64,
            'cumulative_recovered_age_2':
                tf.float64,
            'cumulative_recovered_age_3':
                tf.float64,
            'cumulative_recovered_age_4':
                tf.float64,
            'cumulative_recovered_age_5':
                tf.float64,
            'cumulative_recovered_age_6':
                tf.float64,
            'cumulative_recovered_age_7':
                tf.float64,
            'cumulative_recovered_age_8':
                tf.float64,
            'cumulative_recovered_age_9':
                tf.float64,
            'new_recovered_male':
                tf.float64,
            'new_recovered_female':
                tf.float64,
            'cumulative_recovered_male':
                tf.float64,
            'cumulative_recovered_female':
                tf.float64,
            'locality_code':
                tf.string,
            'locality_name':
                tf.string,
            'new_ventilator_patients':
                tf.float64,
            'cumulative_ventilator_patients':
                tf.float64,
            'new_persons_fully_vaccinated_pfizer':
                tf.float64,
            'cumulative_persons_fully_vaccinated_pfizer':
                tf.float64,
            'new_vaccine_doses_administered_pfizer':
                tf.float64,
            'cumulative_vaccine_doses_administered_pfizer':
                tf.float64,
            'new_persons_fully_vaccinated_moderna':
                tf.float64,
            'cumulative_persons_fully_vaccinated_moderna':
                tf.float64,
            'new_vaccine_doses_administered_moderna':
                tf.float64,
            'cumulative_vaccine_doses_administered_moderna':
                tf.float64,
            'new_persons_fully_vaccinated_janssen':
                tf.float64,
            'cumulative_persons_fully_vaccinated_janssen':
                tf.float64,
            'new_vaccine_doses_administered_janssen':
                tf.float64,
            'cumulative_vaccine_doses_administered_janssen':
                tf.float64,
        }),
        supervised_keys=None,
        homepage='https://github.com/GoogleCloudPlatform/covid-19-open-data',
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Returns SplitGenerators."""
    archive_path = dl_manager.download(
        'https://storage.googleapis.com/covid19-open-data/v3/aggregated.csv.gz?generation=1620814656792419'
    )
    return {tfds.Split.TRAIN: self._generate_examples(dl_manager, archive_path)}

  def _generate_examples(self, dl_manager: tfds.download.DownloadManager,
                         archive_path):
    """Yields examples."""
    pd = tfds.core.lazy_imports.pandas
    beam = tfds.core.lazy_imports.apache_beam

    def cast_according_to_column(feature_type, v):
      if feature_type == tf.string and isinstance(v, (float, int)):
        return str(v)
      return v

    file_handles = dl_manager.iter_archive(archive_path)
    _, file = next(file_handles)

    columns = pd.read_csv(file, nrows=1).columns

    def load_shard(index: int):
      # There is only one file so by using the for we guarantee that the file
      # will be closed.
      for _, file in dl_manager.iter_archive(archive_path):
        df = pd.read_csv(file, skiprows=index, nrows=_BATCH_SIZE)
        features = self.info.features
        result = []
        for i, row in df.iterrows():
          example = {
              k: cast_according_to_column(features[k].dtype, v)
              for k, v in zip(columns, row.values)
          }
          result.append((index + i, example))
        return result

    return (beam.Create(list(range(0, _N_RECORDS, _BATCH_SIZE)))
            | beam.FlatMap(load_shard))
