ColaboratoryでBigQueryのクエリ結果をDataFrameにしてqueryメソッドを実行したときのエラー対応

概要

Google Colaboratory(以下 Colab)ではBigQueryのクエリ結果を直にPandas DataFrameに変換できる。ちょっとした調査でColabを使う事が多いのだが、過去に作成した調査用Colabノートが軒並エラーが出て動かなくなったので対応内容をメモしておく。

具体的にはPandas DataFrameの query メソッドがそのままでは動作しなくなったのでdtypeを変換する処理を挟んで回避した。

再現コード

%%bigquery df_sample --project xxx

select
  `DATE` as pub_date,
  title,
  citations
from `bigquery-public-data.breathe.arxiv`
where title like "%COVID-19%"

df_sample.query("citations > 0")

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-691abbefba30> in <cell line: 1>()
----> 1 df_sample.query("citations > 0")

9 frames
/usr/local/lib/python3.10/dist-packages/numexpr/necompiler.py in getType(a)
    698     if kind == 'U':
    699         raise ValueError('NumExpr 2 does not support Unicode as a dtype.')
--> 700     raise ValueError("unknown type %s" % a.dtype.name)
    701 
    702 

ValueError: unknown type object

これは困った df_sample[df_sample['citations'] > 0] と書けば動くが面倒くさい。BigQuery上で DATE 型の列も同様にエラーになる。

workaround

dtypesを確認すると見慣れない型 dbdate, Int64 が登場してくる

df_sample.info()

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 27 entries, 0 to 26
Data columns (total 3 columns):
 #   Column     Non-Null Count  Dtype 
---  ------     --------------  ----- 
 0   pub_date   27 non-null     dbdate
 1   title      27 non-null     object
 2   citations  27 non-null     Int64 
dtypes: Int64(1), dbdate(1), object(1)
memory usage: 803.0+ bytes

なのでいつもの(?) np.int64 などに変換してやれば動く

import numpy as np
import pandas as pd

def update_dtypes(df: pd.DataFrame) -> None:
    for column in df.columns:
        if df[column].dtype == "dbdate":
            df[column] = df[column].astype(dtype=np.datetime64)
            print(f'Replace dtype of {column}, dbdate to np.datetime64')
        if df[column].dtype == "Int64":
            df[column] = df[column].astype(dtype=np.int64)
            print(f"Replace dtype of {column}, Int64 to np.int64")
	    
update_dtypes(df_sample)

おそらく比較に使う変数を db-dtypes パッケージにある型にしてあげるのが丁寧なソシューションだと思うけど面倒なのでこれで。
https://googleapis.dev/python/db-dtypes/latest/usage.html

再現ノート

スクリーンショット