一、介绍

Polars 是一个用于操作结构化数据的高性能 DataFrame 库,由于 Polars 是从0开始用Rust编写,紧密与机器结合。其矢量化和列式处理可在现代处理器上实现缓存一致性算法和高性能。如果您经常使用 pandas,那么用起 Polars 会感觉很轻松,可以说是平替 Pandas 最有潜质的包。

Polars 在独立的 TPCH 基准测试中与其他几个解决方案进行了基准测试。该基准测试旨在复制实践中使用的数据整理操作。由于其并行执行引擎、高效算法以及 SIMD(单指令、多数据)矢量化的使用,Polars 轻松胜过其他解决方案。与pandas相比,它可以实现30倍以上的性能提升

Polars 的目标是提供一个闪电般快速的DataFrame库:

  • 利用机器上所有可用的内核。
  • 优化查询以减少不必要的工作/内存分配。
  • 处理比可用 RAM 大得多的数据集。
  • 拥有一致且可预测的 API。
  • 具有严格的架构(在运行查询之前应该知道数据类型)。

User guide: https://pola-rs.github.io/polars/user-guide/ API reference: https://pola-rs.github.io/polars/py-polars/html/reference/io.html


打开命令行, 执行 polars 安装命令

pip3 install 'polars[all]'

二、数据读写

Polars 读写数据支持

  • 常见的数据文件,如 csv、xlsx、json、parquet ;
  • 云存储,如 S3、Azure Blob, BigQuery;
  • 数据库,如postgres、mysql

咱们主要分享常见的代码操作

2.1 DataFrame

import polars as pl
import polars.selectors as cs
from datetime import datetime

df = pl.DataFrame(
    {
        "idx": [1, 2, 3, 4],
        "name": ["张三", "李四", "王五", "赵六"],
        "birthday": [
            datetime(2009, 5, 1),
            datetime(2005, 10, 15),
            datetime(2000, 12, 31),
            datetime(1995, 6, 15),
        ],
        "gender": ["男", "男", "男", "女"],
        "bio": ["好好学习,天天向上", 
                "泰难了", 
                "学习有毛用", 
                "躺平ing"],
    }
)

#存入csv、excel、json、parquet
df.write_csv("data.csv")
df.write_excel("data.xlsx")
df.write_json("data.json")
df.write_parquet("data.parquet")


df

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆    birthday	       | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       |
│ 4 ┆"赵六" ┆ 1995-06-15 00:00:00 ┆ "女"  │"躺平ing"         |
└──────────┴─────────────────────┴───────┘─────────────────┴

2.2 csv、excel

  • df.write_csv 存入csv
  • pl.read_csv 读取csv
  • df.write_excel 存入xlsx文件
  • pl.read_excel 读取xlsx
df_csv = pl.read_csv('data.csv')
df_xlsx = pl.read_excel('data.xlsx')

df_csv

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆       birthday	     | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆        str          ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆  "2009-05-01T00:…   ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆  "2005-10-15T00:…	  ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆  "2000-12-31T00:…   ┆ "男"  │"学习有毛用"       |
│ 4 ┆"赵六" ┆  "1995-06-15T00:…	  ┆ "女"  │"躺平ing"         |
└──────────┴─────────────────────┴───────┘─────────────────┴

注意哦, 此时的 date 字段数据类型是 str


2.3 json/parquet

  • df.write_json
  • pl.read_json
  • df.write_parquet
  • pl.read_parquet
df_json = pl.read_json("data.json")
df_parquet = pl.read_parquet("data.parquet")

df_json

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆    birthday	       | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       |
│ 4 ┆"赵六" ┆ 1995-06-15 00:00:00 ┆ "女"  │"躺平ing"         |
└──────────┴─────────────────────┴───────┘─────────────────┴

注意, 使用 df.write_json 或 df.write_parquet 将数据存入 json、parquet, 都可以保留 date 字段的 datetime 类型。而 csv、xlsx 只会将date字段存储为 str 类型。



三、常用表达式

Expressions是Polars的核心功能, expressions 既可以解决简单的查询,又可以轻松扩展到复杂的查询。下面是 polars 的基本表达式

  • pl.col 列选择器
  • df.select 结合pl.col, 返回dataframe
  • selector selector选择器
  • df.filter 结合pl.col, 返回dataframe
  • df.with_columns 结合pl.col, 返回dataframe
  • df.grouby 结合pl.col, 返回dataframe

3.1 pl.col

选择某一(多)个字段(列)

pl.col('birthday')

Run

col("birthday")

pl.col('name', 'birthday')

Run

col(["name", "birthday"])

3.2 df.select

选择 namebirthday 两个字段, 实现该功能有多种写法

#df[['name', 'birthday']]

#df.select(
#    pl.col("name"), 
#    pl.col("birthday"), 
#)


#df.select(["name", "birthday"])


df.select(
    pl.col("name", "birthday")
)

Run

shape: (4, 2)
┌──────┬─────────────────────┬
│ name ┆    birthday	       | 
│------┆ ------------------  ┆ 
│  str ┆    datetime[μs]     ┆ 
╞══════╪═════════════════════╪
│"张三" ┆ 2009-05-01 00:00:00 ┆
│"李四" ┆ 2005-10-15 00:00:00 ┆
│"王五" ┆ 2000-12-31 00:00:00 ┆
│"赵六" ┆ 1995-06-15 00:00:00 ┆
└─────────────────────────────

polars 即使选择一个字段, 返回的也是dataframe

#df[['name']]

#df.select(["name"])

df.select("name")

Run

shape: (4, 1)
┌──────┬
│ 姓名  ┆
│------┆ 
│  str ┆ 
╞══════╪
│"张三" ┆
│"李四" ┆
│"王五" ┆
│"赵六" ┆
└───────

3.3 df.with_columns

与 df.select 功能类似,但是df.with_columns可以在选择字段的同时,保留之前的字段

df.with_columns(
    pl.col('name')
)

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆       birthday	     | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆        str          ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆  "2009-05-01T00:…   ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆  "2005-10-15T00:…	  ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆  "2000-12-31T00:…   ┆ "男"  │"学习有毛用"       |
│ 4 ┆"赵六" ┆  "1995-06-15T00:…	  ┆ "女"  │"躺平ing"         |
└──────────┴─────────────────────┴───────┘─────────────────┴

df.with_columns(
    pl.col('name').alias('姓名')
)

Run

shape: (4, 6)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐───────┬
│idx| name ┆    birthday	       | gender┆	    bio        ┆  姓名  ┆ 
│---┆------┆-------------------- ┆  ---  ┆ --------------  │-------┆
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆  str  ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡═══════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|"张三"  ┆
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          |"李四"  ┆
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       |"王五"  ┆
│ 4 ┆"赵六" ┆ 1995-06-15 00:00:00 ┆ "女"  │"躺平ing"         |"赵六"  ┆ 
└──────────┴─────────────────────┴───────┘─────────────────┴───────┴

3.4 df.filter

筛选出生日是 00 后的记录

df.filter(
  pl.col('birthday') > datetime(2000, 1, 1)
)

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆    birthday	       | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       |
└──────────┴─────────────────────┴───────┘─────────────────┴

3.5 df.groupby

性别gender 进行分组功能

#for gender, gender_df in df.groupby('gender'):
for gender, gender_df in df.groupby(pl.col('gender')):
    print(gender, len(gender_df), type(gender_df))

Run

男 3 <class 'polars.dataframe.frame.DataFrame'>
女 1 <class 'polars.dataframe.frame.DataFrame'>

分别计算男女学生的bio的文本长度的均值

for gender, gender_df in df.groupby(pl.col('gender')):
    print(gender,  gender_df['bio'].apply(lambda t: len(t)).mean())

Run

男 5.666666666666667
女 5.0

df.groupby('gender').agg(
    pl.count(),
    pl.col('bio').str.len_chars().mean().alias('mean_len')
)

Run

shape: (2, 3)
┌──────┬───────┬───────────┬
│gender| count ┆ mean_len	 |
│------┆ ----- ┆-----------┆
│  str ┆  u32	 ┆   f64     ┆
╞══════╪═══════╪═══════════╡
│ "女" ┆  1    ┆    5.0    ┆
│ "男" ┆  3    ┆  5.666667 ┆
└──────┴───────┴───────────┘



四、选择器

  • cs.integer、cs.string、cs.numeric 、cs.datetime()、cs.temporal() 按照数据格式筛选字段
  • cs.contains 、cs.matches 使用正则表达式筛选字段

4.1 按数据格式筛选

筛选出字段数据类型为字符和数字的字段,返回dataframe

import polars.selectors as cs

df.select(
    cs.integer(), cs.string()
)

Run

shape: (4, 4)
┌───┬──────┬───────┬─────────────────┐
│idx| name ┆ gender┆	    bio        ┆
│---┆ ---  ┆  ---  ┆ --------------  │
│i64┆ str  ┆  str  ┆      str        ┆
╞═══╪══════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆ "男"  │"好好学习,天天向上"|
│ 2 ┆"李四" ┆ "男"  │"泰难了"          |
│ 3 ┆"王五" ┆ "男"  │"学习有毛用"       |
│ 4 ┆"赵六" ┆ "女"  │"躺平ing"         |
└──────────┴───────┘─────────────────┴

筛选出 datetime 格式的字段,返回 dataframe

#df.select(cs.temporal())

df.select(
    cs.datetime()
)

Run

shape: (4, 1)
┌───────────────────┬
│ birthday	        |
│-------------------┆
│ datetime[μs]      ┆
╞═══════════════════╪
│2009-05-01 00:00:00┆
│2005-10-15 00:00:00┆
│2000-12-31 00:00:00┆
│1995-06-15 00:00:00┆
└───────────────────┴

4.2 cs.contains/ cs.matches

筛选出含 r 字段,返回dataframe

#筛选出字段名含 r 的字段
df.select(
    cs.contains('r')
)

Run

shape: (4, 2)
┌───────────────────┬───────┬
│ birthday	        | gender┆
│-------------------┆  ---  ┆
│ datetime[μs]      ┆  str  ┆
╞═══════════════════╪═══════╡
│2009-05-01 00:00:00┆ "男"  │
│2005-10-15 00:00:00┆ "男"  │
│2000-12-31 00:00:00┆ "男"  │
│1995-06-15 00:00:00┆ "女"  │
└───────────────────┴───────┘

筛选出含 naio 的字段,返回dataframe

df.select(
    cs.matches('na|io')
)

Run

shape: (4, 2)
┌─────┬───────────────────┐
│name ┆       bio         ┆
│ --- ┆ ---------------   ┆
│ str ┆  str              ┆
╞═════╪═══════════════════╡
│"张三"┆ "好好学习,天天向上" |
│"李四"┆ "泰难了"           |
│"王五"┆ "学习有毛用"        |
│"赵六"┆ "躺平ing"          |
└─────┴───────────────────┴



五、逻辑条件

pl.when(condition).then(result1).otherwise(result2)

当满足condition时, 值为result1; 反之,则result2

df.with_columns(
    pl.when(pl.col('birthday')>datetime(2000, 1, 1))
    .then(True)
    .otherwise(False)
    .alias('00后')
)

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐───────┬
│idx| name ┆    birthday	       | gender┆	    bio        ┆ 00后  ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │  ---- ┆
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆  str  ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡═══════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"| true  |
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          | true  |
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       | true  |
│ 4 ┆"赵六" ┆ 1995-06-15 00:00:00 ┆ "女"  │"躺平ing"         | false |
└──────────┴─────────────────────┴───────┘─────────────────┴───────┴



六、字符串操作

  • pl.col().str.len_chars() 字符长度
  • pl.col().str.contains(pat) 是否含某字符(符合pat模式)
  • pl.col().str.extract(pat) 提取出符合模式的文本
  • pl.col().str.replace(old_pat, new_pat) 把old_pat替换为new_pat

6.1 str.len_chars()

计算 bio 的文字长度,计算结果存储到 lenth 字段中

df.select(
    pl.col('bio'),
    pl.col('bio').str.len_chars().alias('lenth')
)

Run

shape: (4, 2)
┌─────────────────┐───────┬
│	      bio       ┆ lenth ┆
│ --------------  │  ---- ┆
│      str        ┆  u32  ┆
╞═════════════════╡═══════╡
│ "好好学习,天天向上"|   9  |
│ "泰难了"          |   3  |
│ "学习有毛用"       |   5  |
│ "躺平ing"         |   5  |
└──────────────────┴───────┴

6.2 str.contains()

从 bio 中筛选出含 学习 字眼的记录

df.filter(
  pl.col('bio').str.contains("学习")
)

Run

shape: (4, 5)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐
│idx| name ┆    birthday	       | gender┆	    bio        ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       |
└──────────┴─────────────────────┴───────┘─────────────────┴

6.3 str.extract()

根据负面词典 '躺平|难|毛' 选出负面词, 结果存储到字段 neg

df.with_columns(
    pl.col('bio').str.extract_all('躺平|难|毛').alias('neg')
)

Run

shape: (4, 6)
┌───┬──────┬─────────────────────┬───────┬─────────────────┐───────┬
│idx| name ┆    birthday	       | gender┆	    bio        ┆  neg  ┆
│---┆ ---  ┆    -------------    ┆  ---  ┆ --------------  │  ---  ┆
│i64┆ str  ┆    datetime[μs]     ┆  str  ┆      str        ┆  str  ┆
╞═══╪══════╪═════════════════════╪═══════╡═════════════════╡═══════╡
│ 1 ┆"张三" ┆ 2009-05-01 00:00:00 ┆ "男"  │"好好学习,天天向上"|   []  |
│ 2 ┆"李四" ┆ 2005-10-15 00:00:00 ┆ "男"  │"泰难了"          | ["难"]|
│ 3 ┆"王五" ┆ 2000-12-31 00:00:00 ┆ "男"  │"学习有毛用"       | ["毛"]|
│ 4 ┆"赵六" ┆ 1995-06-15 00:00:00 ┆ "女"  │"躺平ing"         |["躺平"]|
└──────────┴─────────────────────┴───────┘─────────────────┴───────┴



广而告之