Pandas直接讀取sql腳本的方法

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

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

之前有群友反應(yīng)同事給了他一個(gè)幾百M(fèi)B的sql腳本，導(dǎo)入數(shù)據(jù)庫(kù)再?gòu)臄?shù)據(jù)庫(kù)讀取數(shù)據(jù)有點(diǎn)慢，想了解下有沒(méi)有可以直接讀取sql腳本到pandas的方法。

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

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

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

注意：該讀取方法只針對(duì)SQLyog導(dǎo)出的mysql腳本測(cè)試，其他數(shù)據(jù)庫(kù)可能代碼需要根據(jù)實(shí)際情況微調(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：腳本中字符串是單引號(hào)還是雙引號(hào)，默認(rèn)使用單引號(hào)解析

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

下面我測(cè)試讀取下面這個(gè)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)然上面寫(xiě)的方法是一次性讀取整個(gè)sql腳本的所有表，結(jié)果為一個(gè)字典（鍵為表名，值為datafream）。但大部分時(shí)候我們只需要讀取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：腳本中字符串是單引號(hào)還是雙引號(hào)，默認(rèn)使用單引號(hào)解析

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

讀取代碼：

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

結(jié)果：

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

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

加載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腳本的位置

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

測(cè)試讀?。?/p>

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é)果：