from google.cloud import bigquery
from timeit import default_timer as timer
from datetime import timedelta

def show_query(query):
    client = bigquery.Client()
    query_job = client.query(query, bigquery.job.QueryJobConfig(use_query_cache=False))
    df = query_job.result().to_dataframe()
    print("Compute: {} slotms Bytes: {:.1f} MB".format(query_job.slot_millis, query_job.total_bytes_processed/(1024*1024)))
    return df

"""
Explanation: Speeding up BigQuery queries with BI Engine
To speed up small queries in BigQuery, simply turn on BI Engine.
The Client API remains exactly the same.
Accompanies https://medium.com/@lakshmanok/speeding-up-small-queries-in-bigquery-with-bi-engine-4ac8420a2ef0
Queries
End of explanation
"""


COST_BY_STATE="""
    SELECT 
       nppes_provider_state, SUM(total_drug_cost)/SUM(total_claim_count) AS avg_cost
    FROM `bigquery-public-data.medicare.part_d_prescriber_2014`
    GROUP BY 1
    ORDER BY 2 DESC
    LIMIT 10
"""
show_query(COST_BY_STATE)

"""
Explanation: This query finds the average prescription claim by state.
End of explanation
"""


TIPS_BY_TYPE="""
SELECT payment_type, AVG(tips) AS avg_tip
FROM bigquery-public-data.chicago_taxi_trips.taxi_trips
GROUP BY 1
ORDER BY 2 DESC
"""
show_query(TIPS_BY_TYPE)

"""
Explanation: This query finds the average tip based on payment type.
End of explanation
"""


AIR_QUALITY="""
            SELECT
               site_num,
               ANY_VALUE(state_name) AS state,
               AVG(aqi) as air_quality_index,
            FROM `bigquery-public-data.epa_historical_air_quality.pm10_daily_summary`
            GROUP BY site_num
            ORDER BY air_quality_index DESC
            LIMIT 10
"""
show_query(AIR_QUALITY)

"""
Explanation: This query uses a discrete number for grouping and will be inherently faster than the names queries (which use strings).
This query finds the sites with the worst (on average) air quality
End of explanation
"""


from google.cloud import bigquery
from timeit import default_timer as timer
from datetime import timedelta

# Construct a BigQuery client object.
client = bigquery.Client()

def run_query(query, n=5):
    tot_slotmillis, tot_timeelapsed = 0, timedelta(0)
    for iter in range(n):
        query_job = client.query(query, bigquery.job.QueryJobConfig(use_query_cache=False))
        df = query_job.result().to_dataframe()
        tot_timeelapsed += (query_job.ended - query_job.started)
        tot_slotmillis += query_job.slot_millis
    print("Job stat: slot_mills={} server_time={}".format(tot_slotmillis/n, tot_timeelapsed/n))

"""
Explanation: Without BI Engine
Time it. Note that I am measuring the time taken on the server
using query_job.started and query_job.ended.
This takes out variability due to the time it takes to send the query
over the network to the BigQuery API.
End of explanation
"""


run_query(COST_BY_STATE)

run_query(TIPS_BY_TYPE)

run_query(AIR_QUALITY)

"""
Explanation: Here, I'm running the query without BI Engine turned on.
End of explanation
"""


run_query(COST_BY_STATE)

run_query(TIPS_BY_TYPE)

run_query(AIR_QUALITY)

"""
Explanation: The slot milliseconds is a proxy for the cost if you have a reservation -- it measures how much your BigQuery slots are getting used.
The server_time is the time taken to process the request (we don't measure the network roundtrip time because it's going to be the
same whether or not you use BI Engine).
With BI Engine
Then, I went to the web console and turned on a 10 GB BI Engine reservation (monthly cost: $300).
Note: It seems to take about 3 minutes for the memory to become available, so this is something
you should consider doing for a few hours at least, not on a per-query basis.
End of explanation
"""


