Pandas直接讀取sql腳本的方法

Posted on by

之前有群友反應同事給瞭他一個幾百MB的sql腳本，導入數據庫再從數據庫讀取數據有點慢，想瞭解下有沒有可以直接讀取sql腳本到pandas的方法。

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

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

我簡單研究瞭一下sql腳本的導出格式，並根據格式寫出瞭以下sql腳本的讀取方法。

註意：該讀取方法隻針對SQLyog導出的mysql腳本測試，其他數據庫可能代碼需要根據實際情況微調。

讀取方法：

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

參數：

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

返回：

一個字典，鍵是表名，值是該表對應的數據所組成的datafream對象

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

image-20210119212414433

其中的表名是index_test：

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

結果：

image-20210119212546592

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

當然上面寫的方法是一次性讀取整個sql腳本的所有表，結果為一個字典（鍵為表名，值為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)

參數：

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

返回：

該表所對應的datafream對象

讀取代碼：

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

結果：

image-20210119222939611

將sql腳本轉換為sqlite格式並通過本地sql連接讀取

在寫完上面的方法後，我又想到另一種解決思路，就是將sql腳本轉換成sqlite語法的sql語句，然後直接加載。各種類型的數據庫的sql語句變化較大，下面的方法僅針對SQLyog導出的mysql腳本測試通過，如果是其他的數據庫，可能下面的方法仍然需要微調。最好是先自行將sql腳本轉換為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()

參數：

sql_file_path：sql腳本的位置

返回：

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

測試讀取：

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

結果：

image-20210119234403297

到此這篇關於Pandas直接讀取sql腳本的文章就介紹到這瞭,更多相關Pandas讀取sql腳本內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！

推薦閱讀：