如何在Colab中使用SQL

如今，編碼測試在數(shù)據(jù)科學面試過程中幾乎是標準的。

作為一名數(shù)據(jù)科學招聘經理，我發(fā)現(xiàn)一個20-30分鐘的現(xiàn)場編碼測試和一些準備好的任務能夠有效地識別那些能夠勝任職位的候選人。

Google Colab［https://colab.research.google.com/notebooks/intro.ipynb？utm_source=scs-index］是各種離線和實時數(shù)據(jù)科學編碼的優(yōu)秀工具，因為它熟悉的筆記本環(huán)境和并且方便共享。但是Colab幾乎只限于Python（還有一些黑客攻擊）。

根據(jù)我個人的經驗，SQL是成為一名成功的數(shù)據(jù)科學家的關鍵技能。深度學習是很好的，但是如果你不能編寫SQL查詢，你可能永遠都會被你的貓對狗分類器所困擾。

因此，我開始尋找進行實時SQL編碼測試的方法，瞧！我偶然發(fā)現(xiàn)了這種在Colab（或任何Python環(huán)境）中運行SQL的簡單方法，實際上只有2個helper函數(shù)。

在本文中，我將介紹兩個Python函數(shù)，我們可以使用它們在Google Colab中設置和運行SQL，具體如下：

在Google Colab中創(chuàng)建一個數(shù)據(jù)庫并將數(shù)據(jù)幀上傳到該數(shù)據(jù)庫中的一個表中

在GoogleColab中對該數(shù)據(jù)庫和表編寫和執(zhí)行SQL查詢。

所有代碼都可以在我的Github帳戶上找到。你也可以直接在谷歌Colab上打開代碼， 單擊此鏈接：

https://colab.research.google.com/github/stephenleo/medium-python-hacks/blob/main/02_sql_on_colab/main.ipynb

SQLite

失敗是成功之母

在這次搜索之前，我只隱約聽說過SQLite，但從未需要深入研究它。

SQLite…實現(xiàn)了一個小型、快速、自包含、高可靠性、功能齊全的SQL數(shù)據(jù)庫引擎。

SQLite是世界上使用最多的數(shù)據(jù)庫引擎。SQLite內置于所有手機和大多數(shù)計算機中，并捆綁在人們每天使用的應用程序中。

聽起來很有希望！SQLite的工作原理是創(chuàng)建一個本地的.db文件，我們可以像普通的SQL數(shù)據(jù)庫一樣連接到該文件。然后我們可以在這個.db文件上創(chuàng)建表、上傳數(shù)據(jù)和查詢數(shù)據(jù)。

另外，用于創(chuàng)建SQLite并與之交互的Python模塊（sqlite3）也是Python標準庫的一部分。所以它在Colab上是開箱即用的。

首先，讓我們導入sqlite3和pandas模塊，它們是我們任務的唯一需求

importsqlite3
importpandasaspd

函數(shù)1：將Dataframe轉換為SQL DB表

我們的第一個任務是創(chuàng)建一個數(shù)據(jù)庫，并將一個數(shù)據(jù)幀上傳到該數(shù)據(jù)庫中的一個表中。

下面的代碼將其作為一個函數(shù)實現(xiàn)，我將逐步介紹。該函數(shù)將輸入數(shù)據(jù)幀、表名稱和數(shù)據(jù)庫名稱（.db）作為輸入，并運行以下步驟。

1. 設置一些日志以跟蹤函數(shù)的執(zhí)行

2. 查找數(shù)據(jù)框中的所有列。這是必要的，因為我們需要在創(chuàng)建表和將數(shù)據(jù)上載到表時提供此信息。

3. 連接到.db文件（如果存在）。如果文件不存在，請在本地計算機上創(chuàng)建新文件。

4. 在.db文件中創(chuàng)建一個表，我們在上一步中連接到該表（或剛剛創(chuàng)建）

5. 將 input_df 中的數(shù)據(jù)行上傳到我們在上一步中剛剛創(chuàng)建的表中

6. 提交更改并關閉與數(shù)據(jù)庫的連接

defpd_to_sqlDB(input_df:pd.DataFrame,
table_name:str,
db_name:str='default.db')->None:

'''
取一個數(shù)據(jù)幀'input_df'并將其上傳到'table_name'SQLITE表
Args:
input_df(pd.DataFrame):包含要上傳到SQLITE的數(shù)據(jù)的數(shù)據(jù)幀
table_name(str):要上傳的SQLITE表的名稱
db_name (str, optional):創(chuàng)建表的SQLITE數(shù)據(jù)庫的名稱。
默認為“default.db”
'''

#步驟1:設置本地日志
importlogging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s%(levelname)s:%(message)s',
datefmt='%Y-%m-%d%H:%M:%S')

#步驟2:在數(shù)據(jù)幀中查找列
cols=input_df.columns
cols_string=','.join(cols)
val_wildcard_string=','.join(['?']*len(cols))

#步驟3:如果DB文件存在，連接到它，否則創(chuàng)建一個新文件
con=sqlite3.connect(db_name)
cur=con.cursor()
logging.info(f'SQLDB{db_name}created')

#步驟4:創(chuàng)建表
sql_string=f"""CREATETABLE{table_name}({cols_string});"""
cur.execute(sql_string)
logging.info(f'SQLTable{table_name}createdwith{len(cols)}columns')

#步驟5:上傳數(shù)據(jù)幀
rows_to_upload=input_df.to_dict(orient='split')['data']
sql_string=f"""INSERTINTO{table_name}({cols_string})VALUES({val_wildcard_string});"""
cur.executemany(sql_string,rows_to_upload)
logging.info(f'{len(rows_to_upload)}rowsuploadedto{table_name}')

#步驟6:提交更改并關閉連接
con.commit()
con.close()

函數(shù)2：對數(shù)據(jù)幀的SQL查詢

我們的第二個任務是在第一個任務中創(chuàng)建的數(shù)據(jù)庫和表上編寫和執(zhí)行SQL查詢。

下面的代碼將其作為另一個函數(shù)實現(xiàn)，我也將逐步介紹。該函數(shù)將sql_query_string和db_name（.db file）作為輸入，運行以下步驟，并返回一個dataframe作為輸出。

1. 連接到.DB文件中的SQL DB

2. 在SQL查詢字符串中執(zhí)行SQL查詢

3. 運行SQL查詢后獲取結果數(shù)據(jù)和輸出的列名

4. 關閉與數(shù)據(jù)庫的連接

5. 將結果作為數(shù)據(jù)幀返回

defsql_query_to_pd(sql_query_string:str,db_name:str='default.db')->pd.DataFrame:
'''執(zhí)行一個SQL查詢，并以數(shù)據(jù)幀的形式返回結果
Args:
sql_query_string(str):要執(zhí)行的SQL查詢字符串
db_name(str,optional)::要執(zhí)行查詢的SQLITE數(shù)據(jù)庫的名稱
默認為“default.db”
Returns:
pd.DataFrame:數(shù)據(jù)幀中SQL查詢的結果
'''
#步驟1:連接SQL數(shù)據(jù)庫
con=sqlite3.connect(db_name)

#步驟2:執(zhí)行SQL查詢
cursor=con.execute(sql_query_string)

#步驟3:獲取數(shù)據(jù)和列名
result_data=cursor.fetchall()
cols=[description[0]fordescriptionincursor.description]

#步驟4:關閉連接
con.close()

#步驟5:返回數(shù)據(jù)幀
returnpd.DataFrame(result_data,columns=cols)

差不多就是這樣！我們現(xiàn)在有兩個函數(shù)，可以粘貼到任何Colab筆記本中，以解決我們的兩個任務。讓我們通過一個例子來看看如何使用這兩個函數(shù)

使用這兩個函數(shù)在Colab中運行SQL

由于新冠病毒基本上是我們這些天談論的唯一話題，我從Kaggle下載了一個新冠病毒疫苗樣本數(shù)據(jù)集，以測試我們的兩個函數(shù)函數(shù)的使用非常簡單，如下所示

1. 將下載的csv文件加載到數(shù)據(jù)幀中

2. 使用我們的第一個函數(shù)pd_to_sqlDB函數(shù)將我們在上一步中加載的數(shù)據(jù)幀上載到名為default.DB的數(shù)據(jù)庫中名為 country_vaccination 的表中

3. 在名為 sql_query_string 的字符串變量中編寫SQL查詢。正如你在下面的代碼片段中所看到的，我們可以編寫任何可以想到的復雜SQL查詢。

4. 使用第二個函數(shù)sql_query_to_pd在default.db上執(zhí)行上一步的sql查詢，并將結果保存在result_df數(shù)據(jù)框中

#步驟1:讀取csv文件到一個數(shù)據(jù)幀
#數(shù)據(jù)集來自https://www.kaggle.com/gpreda/covid-world-vaccination-progress
input_df=pd.read_csv('country_vaccinations.csv')

#步驟2:上傳數(shù)據(jù)幀到SQL表
pd_to_sqlDB(input_df,
table_name='country_vaccinations',
db_name='default.db')

#步驟3:在字符串變量中寫入SQL查詢
sql_query_string="""
SELECTcountry,SUM(daily_vaccinations)astotal_vaccinated
FROMcountry_vaccinations
WHEREdaily_vaccinationsISNOTNULL
GROUPBYcountry
ORDERBYtotal_vaccinatedDESC
"""

#步驟4:執(zhí)行SQL查詢
result_df=sql_query_to_pd(sql_query_string,db_name='default.db')
result_df

從結果中，我們可以看到（在撰寫本文時），中國的接種人數(shù)居世界首位，其次是印度和美國。我們通過在GoogleColab中對DB執(zhí)行SQL查詢得到了這些結果！

給你。現(xiàn)在，你可以在下次面試中展示你卓越的Python和SQL知識了！

另外，節(jié)省一些時間，直接使用我在本文開頭鏈接的GoogleColab筆記本，它包含上述函數(shù)和示例查詢。