SQL on DataFrame，多快好省

今天介绍一种可以让查询速度更快， 代码更简洁，省去记忆大量的 Pandas 命令的方法。

或许你听说过 pandasql（或者 sqldf)，这里说的不是它，因为它太慢了。

为了方便理解，做了三种方法的对比， 分别为直接对 DuckDB 的表（视图）进行 SQL 操作、Pandas API 以及把 Pandas DataFrame 当做表，通过 SQL 操作。

环境：Notebook

准备工作

下载数据

wget -q https://github.com/cwida/duckdb-data/releases/download/v1.0/lineitemsf1.snappy.parquet
wget -q https://github.com/cwida/duckdb-data/releases/download/v1.0/orders.parquet

安装包

!pip install pandasql ipython-sql duckdb-engine==0.1.8rc4  duckdb ipython-autotime

导入包及初始化

#自动计算运行时间
%load_ext autotime
#下面6行和ipython-sql有关
%load_ext sql
%config SqlMagic.autopandas=True
%config SqlMagic.feedback = False
%config SqlMagic.displaycon = False

%sql duckdb:///:memory:
%sql pragma threads=4
#导入必要的包
import pandas as pd
import numpy as np
from pandasql import sqldf
from duckdb import query
#方便需要定义的两个函数
sqlite = lambda q: sqldf(q, globals())
ducksql= lambda q:  query(q).df()

准备数据

DuckDB 以两种方式测试，第一种，通过创建视图的方式，直接访问 parquet 文件，不把数据放内存

%%sql
create view lineitem as select * from 'lineitemsf1.snappy.parquet';
create view orders as SELECT * FROM 'orders.parquet';
show tables;
#返回
 name
0 lineitem
1 orders

第二种方式就是把 parquet 加载到 DataFrame，然后 SQL on Pandas DataFrame

lineitem_df = %sql SELECT * FROM 'lineitemsf1.snappy.parquet'
orders_df = %sql SELECT * FROM 'orders.parquet'

对比测试

count

这一步 pandasql 基本就可以出局了，太慢了。

Limit

%sql select * from lineitem limit 5 # 60.9ms
ducksql("select * from lineitem_df limit 5") #23.5ms
sqlite("select * from lineitem_df limit 5") #3m 57s
lineitem_df.head(5) #13ms

这一轮， pandas 胜，pandasql 最慢，接下来放弃它了

选取部分列

lineitem_df[['l_shipdate','l_extendedprice','l_quantity','l_partkey']].head() #271ms
%sql select l_shipdate,l_extendedprice,l_quantity,l_partkey from lineitem limit 5 #18.8ms
ducksql( "select l_shipdate,l_extendedprice,l_quantity,l_partkey from lineitem_df limit 5") #8.1ms

pandas 最慢

排序(Sort/Order by)

lineitem_df.sort_values('l_shipdate',ascending=False).head(5) #28s
ducksql("select * from lineitem_df order by l_shipdate desc limit 5") #1.31s
%sql select * from lineitem order by l_shipdate desc limit 5

#1.14s

查询(Where)

%sql select * from lineitem where l_shipmode ='RAIL' limit 5 #38.5ms
ducksql("select * from lineitem_df where l_shipmode ='RAIL' limit 5") #15.3ms
lineitem_df.query("l_shipmode =='RAIL'").head()

#758ms

聚合（Groupby）

%sql select l_shipmode,l_shipinstruct, count(*) as cnt, sum(l_extendedprice*l_quantity) rev from lineitem group by 1,2 order by 3 desc limit 5 # 415ms

ducksql("select l_shipmode,l_shipinstruct, count(*) as cnt, sum(l_extendedprice*l_quantity) rev from lineitem_df group by 1,2 order by 3 desc limit 5") # 866ms

lineitem_df.assign(rev=lineitem_df["l_extendedprice"] * lineitem_df["l_quantity"]).groupby(['l_shipmode','l_shipinstruct']).agg({"rev": np.sum, "l_shipmode": np.size}).head()

#4.45s

差距有些大，而且代码很长，需要记忆很多函数写法了。

Join(连接)

这里只测试一个 inner join，其它的有兴趣可以自己来

%%sql
SELECT l_returnflag,
       l_linestatus,
       sum(l_extendedprice)
FROM  lineitem
JOIN  orders ON (l_orderkey=o_orderkey)
WHERE l_shipdate <= DATE '1998-09-02'
  AND o_orderstatus='O'
GROUP BY l_returnflag,
         l_linestatus

ducksql("""
SELECT l_returnflag,
       l_linestatus,
       sum(l_extendedprice)
FROM  lineitem_df
JOIN  orders_df ON (l_orderkey=o_orderkey)
WHERE l_shipdate <= DATE '1998-09-02'
  AND o_orderstatus='O'
GROUP BY l_returnflag,
         l_linestatus
""")

pd.merge(lineitem_df.query("l_shipdate <= '1998-09-02'")[['l_returnflag','l_orderkey','l_linestatus','l_extendedprice']],
         orders_df.query("o_orderstatus=='O'")[['o_orderkey']],how='inner', left_on='l_orderkey',right_on='o_orderkey')\
.groupby(['l_returnflag','l_linestatus']).agg(sum=('l_extendedprice','sum'))

Method	Time
View	518ms
SQL on DataFrame	854ms
Pandas	6.13s

Union（合并）

• union all（不去重）

pd.concat([lineitem_df.query("l_shipdate =='1992-01-02'  and l_shipmode =='FOB'"),
          lineitem_df.query("l_shipdate <='1992-01-03'  and l_shipmode =='FOB'")],axis=0)
#1.67s

%%sql
select * from lineitem where  l_shipdate ='1992-01-02' and l_shipmode ='FOB'
union all
select * from lineitem where  l_shipdate <='1992-01-03' and l_shipmode ='FOB'

#1.74s

ducksql("""
select * from lineitem_df where  l_shipdate ='1992-01-02' and l_shipmode ='FOB'
union all
select * from lineitem_df where  l_shipdate <='1992-01-03' and l_shipmode ='FOB'
""")
#2.31s

• union（去重）

pd.concat([lineitem_df.query("l_shipdate =='1992-01-02'  and l_shipmode =='FOB'"),
          lineitem_df.query("l_shipdate <='1992-01-03'  and l_shipmode =='FOB'")],axis=0).drop_duplicates()
# 1.16s

%%sql
select * from lineitem where  l_shipdate ='1992-01-02' and l_shipmode ='FOB'
union
select * from lineitem where  l_shipdate <='1992-01-03' and l_shipmode ='FOB'
#1.75s

ducksql("""
select * from lineitem_df where  l_shipdate ='1992-01-02' and l_shipmode ='FOB'
union
select * from lineitem_df where  l_shipdate <='1992-01-03' and l_shipmode ='FOB'
""")
# 2.51s

开窗函数(Window Function)

• row_number()

%%sql
SELECT * FROM (
  SELECT
    t.*,
    ROW_NUMBER() OVER(PARTITION BY l_shipdate ORDER BY l_extendedprice DESC) AS rn
  FROM lineitem t
)
WHERE rn =1
ORDER BY l_extendedprice, rn;
#14s

ducksql("""
SELECT * FROM (
  SELECT
    t.*,
    ROW_NUMBER() OVER(PARTITION BY l_shipdate ORDER BY l_extendedprice DESC) AS rn
  FROM lineitem_df t
)
WHERE rn =1
ORDER BY l_extendedprice, rn;
""")
#37.8s

(
    lineitem_df.assign(
        rn=lineitem_df.sort_values(["l_extendedprice"], ascending=False)
        .groupby(["l_shipdate"])
        .cumcount()
        + 1
    )
    .query("rn ==1")
    .sort_values(["l_shipdate", "rn"])
)
#24.6s

(
    lineitem_df.assign(
        rnk=lineitem_df.groupby(["l_shipdate"])["l_extendedprice"].rank(
            method="first", ascending=False
        )
    )
    .query("rnk == 1")
    .sort_values(["l_shipdate", "rnk"])
)
# 5.81s，这个快很多

• rank

%%sql
SELECT * FROM (
  SELECT
    t.*,
    RANK() OVER(PARTITION BY l_shipmode ORDER BY l_quantity) AS rnk
  FROM lineitem t
    where l_quantity<10
)
WHERE rnk < 3
ORDER BY l_shipmode, rnk;
#3.03s

ducksql("""
SELECT * FROM (
  SELECT
    t.*,
    RANK() OVER(PARTITION BY l_shipmode ORDER BY l_quantity) AS rnk
  FROM lineitem_df t
    where l_quantity<10
)
WHERE rnk < 3
ORDER BY l_shipmode, rnk;
""")
#4.17

(
    lineitem_df[lineitem_df["l_quantity"] < 10]
    .assign(rnk_min=lineitem_df.groupby(["l_shipmode"])["l_quantity"].rank(method="min"))
    .query("rnk_min < 3")
    .sort_values(["l_shipmode", "rnk_min"])
)
# 3.31

更新(UPDATE)与删除(DELETE)

SQL on DataFrame 不支持。

从上面的例子可以看出，使用 SQL on DataFrame， 省去了记忆复杂 Pandas 的函数、速度也比较保证（复杂的 SQL 反而快很多），另外如果数据比较大，还是建议直接放在 DuckDB 数据库里，不用担心丢失、节省内存，还保证速度，只在需要的时候才把数据取出来。