from google.cloud import bigquery

client = bigquery.Client()

query = """
WITH a AS (
  SELECT x AS i
  FROM UNNEST(GENERATE_ARRAY(1, 10000)) AS x
),
b AS (
  SELECT y AS j
  FROM UNNEST(GENERATE_ARRAY(1, 10000)) AS y
),
pairs AS (
  SELECT
    i,
    j,
    SIN(i * j) AS s,
    COS(i + j) AS c
  FROM a
  CROSS JOIN b
),
heavy AS (
  SELECT
    i,
    j,
    -- intentionally nasty CPU-heavy expression
    EXP(
      SIN(s) * COS(c)
      + LOG(ABS(s) + 1.000001)
      + LOG(ABS(c) + 1.000001)
    )
    + EXP(
      SIN(i) * COS(j)
      + LOG(i + 1.000001)
      + LOG(j + 1.000001)
    ) AS heavy_val
  FROM pairs
)
SELECT
  COUNT(*) AS row_count,
  SUM(heavy_val) AS total_heavy_val
FROM heavy;
"""

rows = client.query_and_wait(query)

for row in rows:
    print(f"row_count={row['row_count']}, total_heavy_val={row['total_heavy_val']}")