Pandas直接讀取sql腳本的方法

更新時間：2021年01月21日 11:01:28 作者：小小明（代碼實體）

這篇文章主要介紹了Pandas直接讀取sql腳本的方法,本文給大家介紹的非常詳細，對大家的學(xué)習(xí)或工作具有一定的參考借鑒價值，需要的朋友可以參考下

之前有群友反應(yīng)同事給了他一個幾百MB的sql腳本，導(dǎo)入數(shù)據(jù)庫再從數(shù)據(jù)庫讀取數(shù)據(jù)有點慢，想了解下有沒有可以直接讀取sql腳本到pandas的方法。

解析sql腳本文本文件替換成csv格式并加載

我考慮了一下sql腳本也就只是一個文本文件而已，而且只有幾百MB，現(xiàn)代的機器足以把它一次性全部加載到內(nèi)存中，使用python來處理也不會太慢。

我簡單研究了一下sql腳本的導(dǎo)出格式，并根據(jù)格式寫出了以下sql腳本的讀取方法。

注意：該讀取方法只針對SQLyog導(dǎo)出的mysql腳本測試，其他數(shù)據(jù)庫可能代碼需要根據(jù)實際情況微調(diào)。

讀取方法：

from io import StringIO
import pandas as pd
import re

def read_sql_script_all(sql_file_path, quotechar="'") -> (str, dict):
  insert_check = re.compile(r"insert +into +`?(\w+?)`?\(", re.I | re.A)
  with open(sql_file_path, encoding="utf-8") as f:
    sql_txt = f.read()
  end_pos = -1
  df_dict = {}
  while True:
    match_obj = insert_check.search(sql_txt, end_pos+1)
    if not match_obj:
      break
    table_name = match_obj.group(1)
    start_pos = match_obj.span()[1]+1
    end_pos = sql_txt.find(";", start_pos)
    tmp = re.sub(r"\)( values |,)\(", "\n", sql_txt[start_pos:end_pos])
    tmp = re.sub(r"[`()]", "", tmp)
    df = pd.read_csv(StringIO(tmp), quotechar=quotechar)
    dfs = df_dict.setdefault(table_name, [])
    dfs.append(df)
  for table_name, dfs in df_dict.items():
    df_dict[table_name] = pd.concat(dfs)
  return df_dict

參數(shù)：

sql_file_path：sql腳本的位置
quotechar：腳本中字符串是單引號還是雙引號，默認使用單引號解析

返回：

一個字典，鍵是表名，值是該表對應(yīng)的數(shù)據(jù)所組成的datafream對象

下面我測試讀取下面這個sql腳本：

其中的表名是index_test：

df_dict = read_sql_script_all("D:/tmp/test.sql")
df = df_dict['index_test']
df.head(10)

結(jié)果：

可以看到能順利的直接從sql腳本中讀取數(shù)據(jù)生成datafream。

當(dāng)然上面寫的方法是一次性讀取整個sql腳本的所有表，結(jié)果為一個字典（鍵為表名，值為datafream）。但大部分時候我們只需要讀取sql腳本的某一張表，我們可以改造一下上面的方法：

def read_sql_script_by_tablename(sql_file_path, table_name, quotechar="'") -> (str, dict):
  insert_check = re.compile(r"insert +into +`?(\w+?)`?\(", re.I | re.A)
  with open(sql_file_path, encoding="utf-8") as f:
    sql_txt = f.read()
  end_pos = -1
  dfs = []
  while True:
    match_obj = insert_check.search(sql_txt, end_pos+1)
    if not match_obj:
      break
    start_pos = match_obj.span()[1]+1
    end_pos = sql_txt.find(";", start_pos)
    if table_name != match_obj.group(1):
      continue
    tmp = re.sub(r"\)( values |,)\(", "\n", sql_txt[start_pos:end_pos])
    tmp = re.sub(r"[`()]", "", tmp)
    df = pd.read_csv(StringIO(tmp), quotechar=quotechar)
    dfs.append(df)
  return pd.concat(dfs)

參數(shù)：

sql_file_path：sql腳本的位置
table_name：被讀取的表名
quotechar：腳本中字符串是單引號還是雙引號，默認使用單引號解析

返回：

該表所對應(yīng)的datafream對象

讀取代碼：

df = read_sql_script_by_tablename("D:/tmp/test.sql", "index_test")
df.head()

結(jié)果：

將sql腳本轉(zhuǎn)換為sqlite格式并通過本地sql連接讀取

在寫完上面的方法后，我又想到另一種解決思路，就是將sql腳本轉(zhuǎn)換成sqlite語法的sql語句，然后直接加載。各種類型的數(shù)據(jù)庫的sql語句變化較大，下面的方法僅針對SQLyog導(dǎo)出的mysql腳本測試通過，如果是其他的數(shù)據(jù)庫，可能下面的方法仍然需要微調(diào)。最好是先自行將sql腳本轉(zhuǎn)換為sqlite語法的sql語句后，再使用我寫的方法加載。

加載sql腳本的方法：

from sqlalchemy import create_engine
import pandas as pd
import re


def load_sql2sqlite_conn(sqltxt_path):
  create_rule = re.compile("create +table [^;]+;", re.I)
  insert_rule = re.compile("insert +into [^;]+;", re.I)
  with open(sqltxt_path, encoding="utf-8") as f:
    sqltxt = f.read()
  engine = create_engine('sqlite:///:memory:')
  pos = -1
  while True:
    match_obj = create_rule.search(sqltxt, pos+1)
    if match_obj:
      pos = match_obj.span()[1]
      sql = match_obj.group(0).replace("AUTO_INCREMENT", "")
      sql = re.sub("\).+;", ");", sql)
      engine.execute(sql)
    match_obj = insert_rule.search(sqltxt, pos+1)
    if match_obj:
      pos = match_obj.span()[1]
      sql = match_obj.group(0)
      engine.execute(sql)
    else:
      break
  tablenames = [t[0] for t in engine.execute(
    "SELECT tbl_name FROM sqlite_master WHERE type='table';").fetchall()]
  return tablenames, engine.connect()

參數(shù)：

sql_file_path：sql腳本的位置

返回：

兩個元素的元組，第一個元素是表名列表，第二個元素是sqlite內(nèi)存虛擬連接

測試讀取：

tablenames, conn = load_sql2sqlite_conn("D:/tmp/test.sql")
tablename = tablenames[0]
print(tablename)
df = pd.read_sql(f"select * from {tablename};", conn)
df

結(jié)果：