add partitioning test

dask_bigquery/tests/test_core.py

@@ -1,6 +1,7 @@
 import random
 
 import pandas as pd
+import pytest
 from distributed.utils_test import cluster_fixture # noqa: F401
 from distributed.utils_test import client, loop # noqa: F401
 from google.cloud import bigquery
@@ -38,6 +39,7 @@ def push_data():
  return df
 
 
+# test simple read
 def test_read_gbq(client):
  """Test simple read of data pushed to BigQuery using pandas-gbq"""
  try:
@@ -62,5 +64,37 @@ def test_read_gbq(client):
  assert ddf.npartitions == 2
 
  ddf_comp = ddf.set_index("idx").compute()
- # breakpoint()
  assert all(ddf_comp == df.set_index("idx"))
+
+
+# test partitioned data: this test requires a copy of the public dataset
+# bigquery-public-data.covid19_public_forecasts.county_14d into a the
+# project dask-bigquery
+
+
+@pytest.mark.parametrize(
+ "fields",
+ ([], ["county_name"], ["county_name", "county_fips_code"]),
+ ids=["no_fields", "missing_partition_field", "fields"],
+)
+def test_read_gbq_partitioning(fields, client):
+ partitions = ["Teton", "Loudoun"]
+ ddf = read_gbq(
+ project_id="dask-bigquery",
+ dataset_id="covid19_public_forecasts",
+ table_id="county_14d",
+ partition_field="county_name",
+ partitions=partitions,
+ fields=fields,
+ )
+
+ assert len(ddf) # check it's not empty
+ loaded = set(ddf.columns) | {ddf.index.name}
+
+ if fields:
+ assert loaded == set(fields) | {"county_name"}
+ else: # all columns loaded
+ assert loaded >= set(["county_name", "county_fips_code"])
+
+ assert ddf.npartitions == len(partitions)
+ assert list(ddf.divisions) == sorted(ddf.divisions)