文件的输入输出

文件的输入输出#

Polars 的文件读写很快，CSV、Excel、JSON、Parquet、IPC 这些格式都能处理。读写本身并不难，难的是格式不干净时如何把参数设对：有没有表头、类型推断是否可靠、多余的空白算不算内容。这一章按格式分节，每节都从实际遇到的不规范数据出发，逐一确认有代表性的选项。

读文件时真正容易出错的，往往不是调用哪个函数，而是文件本身的约定：编码、分隔符、表头行数、空值写法、类型混用。这些细节只要有一处没对上，读进来的表就可能整列变字符串，后面的计算再正确也没有意义。

import polars as pl
from helper.jupyter import row, capture_except

CSV 文件#

CSV 的格式较为松散，读之前往往要先确定几件事：有没有表头、表头占几行、类型如何推断、多余的空白如何处理。

CSV 没有类型声明，也没有统一的空值记号。同一个文件里，数字可能带引号，空值可能写成空字符串、NA 或 -，列与列之间还可能混着制表符和空格。Polars 的读取参数就是用来逐项确认这些约定的；设错一项，结果可能从“能读”变成“读错”。

读取#

表头#

CSV 有时带表头，有时不带，表头也可能跨两行。下面用这份数据确认几种表头处理方式。

%%writefile data/csv_header.csv
A,B
a,b
0,1
2,3
4,5

Overwriting data/csv_header.csv

df1：默认把第 1 行当作表头读入。
df2：设 has_header=False，第 1 行也变成数据，列名用 new_columns 指定。
df3：用 skip_rows 跳过开头的 N 行。
df4：用 skip_rows_after_header 跳过表头之后的 N 行。
df5：先不带表头读取开头 2 行，组合出列名后再应用；表头跨两行时按这种方式合并。

fn = 'data/csv_header.csv'
df1 = pl.read_csv(fn)
df2 = pl.read_csv(fn, has_header=False, new_columns=['x', 'y'])
df3 = pl.read_csv(fn, skip_rows=1)
df4 = pl.read_csv(fn, skip_rows_after_header=1)

df_header = pl.read_csv(fn, n_rows=2, has_header=False)
columns = df_header.select(pl.all().str.join('-')).row(0)
df5 = pl.read_csv(fn, has_header=False, skip_rows=2, new_columns=columns)
row(df1, df2, df3, df4, df5)

shape: (4, 2)

A	B
str	str
"a"	"b"
"0"	"1"
"2"	"3"
"4"	"5"

shape: (5, 2)

x	y
str	str
"A"	"B"
"a"	"b"
"0"	"1"
"2"	"3"
"4"	"5"

shape: (3, 2)

a	b
i64	i64
0	1
2	3
4	5

shape: (3, 2)

A	B
i64	i64
0	1
2	3
4	5

shape: (3, 2)

A-a	B-b
i64	i64
0	1
2	3
4	5

infer_schema=True(默认)时，Polars 只检查开头 infer_schema_length 行来决定各列的类型，先推断类型，再按这个类型解析剩下的行。

推断的依据是样本行的字面值：整列都像整数，就推断为整数；出现小数点，就推断为浮点；出现非数字字符，就退回字符串。样本之外一旦出现不同类型的值，解析就会失败。这也是 CSV 读取最典型的错误来源：前 100 行很干净，第 101 行突然出现一个 "N/A"。

%%writefile data/csv_different_type.csv
A,B
0,1
2,3
4,5
a,5.5
10,20

Overwriting data/csv_different_type.csv

infer_schema_length 的默认值是 100，对大多数文件够用。下面的代码只取 2 行用于推断，于是第 4 行的 a 引发错误，恰好说明“推断范围不够”的后果。

%%capture_except
df = pl.read_csv('data/csv_different_type.csv', infer_schema_length=2)

ComputeError: could not parse `a` as dtype `i64` at column 'A' (column number 1)

The current offset in the file is 15 bytes.

You might want to try:
- increasing `infer_schema_length` (e.g. `infer_schema_length=10000`),
- increasing `infer_schema_files` (e.g. `infer_schema_files=50`),
- specifying correct dtype with the `schema_overrides` argument
- setting `ignore_errors` to `True`,
- adding `a` to the `null_values` list.

Original error: ```invalid primitive value found during CSV parsing```

错误信息中会列出几种解法，分别对应数据的干净程度。下面的代码把这些解法逐一试过。

df1：加大 infer_schema_length，A 读成 str、B 读成 f64。
df2：infer_schema_length=None 从所有行推断类型，再用 null_values=['a'] 把 a 当 NULL，于是 A 成为 i64。
df3：ignore_errors=True，把不符合推断类型的值变成 NULL。
df4：用 schema_overrides 明确指定类型，再用 ignore_errors=True，既节省内存，又不会因为个别异常值报错。

fn = 'data/csv_different_type.csv'
df1 = pl.read_csv(fn, infer_schema_length=1000)
df2 = pl.read_csv(fn, infer_schema_length=None, null_values=['a'])
df3 = pl.read_csv(fn, infer_schema_length=2, ignore_errors=True)
df4 = pl.read_csv(fn, schema_overrides={'A':pl.Int16, 'B':pl.Float32}, ignore_errors=True)
row(df1, df2, df3, df4)

shape: (5, 2)

A	B
str	f64
"0"	1.0
"2"	3.0
"4"	5.0
"a"	5.5
"10"	20.0

shape: (5, 2)

A	B
i64	f64
0	1.0
2	3.0
4	5.0
null	5.5
10	20.0

shape: (5, 2)

A	B
i64	i64
0	1
2	3
4	5
null	null
10	20

shape: (5, 2)

A	B
i16	f32
0	1.0
2	3.0
4	5.0
null	5.5
10	20.0

列值末尾带空白时，Polars 的标准引擎会原样保留，并把该列判定为 str，因为 "4 " 确实不是合法的整数。

标准引擎把空白视为值的一部分，不会自动裁剪；PyArrow 引擎则更宽松，解析数值时会忽略首尾空白。两种引擎没有绝对的好坏：标准引擎的行为更可预测，PyArrow 对脏数据更宽容。下面用这份数据对比它们的差别。

%%writefile data/csv_trailing_space.csv
str,int,float
日本語 ,4 ,5.67 
abc ,5 ,1.23 

Overwriting data/csv_trailing_space.csv

这份数据在标准引擎和 use_pyarrow=True 下结果不同。

df1：标准引擎下所有列都是 str，因为末尾的空白使该列无法判定为数值。
df2：设置 use_pyarrow=True 之后，int 和 float 两列被解释成数值，PyArrow 的解析器会自动去除空白。

选用哪个引擎取决于数据的干净程度：格式规范时用标准引擎，格式杂乱时 PyArrow 更宽容。

fn = 'data/csv_trailing_space.csv'
df1 = pl.read_csv(fn)
df2 = pl.read_csv(fn, use_pyarrow=True)
row(df1, df2)

shape: (2, 3)

str	int	float
str	str	str
"日本語 "	"4 "	"5.67 "
"abc "	"5 "	"1.23 "

shape: (2, 3)

str	int	float
str	i64	f64
"日本語 "	4	5.67
"abc "	5	1.23

要把标准引擎读出的字符串转成数值，需要自己写转换：先去空白，再按整数 → 小数 → 原字符串的顺序尝试。

写成函数可以集中逻辑，代价是必须逐列调用，因为每列的目标类型不同，无法用一个 return_type 概括。这种转换会失去向量化的优势，只适合列数不多、脏值较少的场景；如果整列都需要清洗，更合适的做法是在读取阶段用 PyArrow 或 schema_overrides 处理。

from polars.exceptions import InvalidOperationError


def try_cast_to_number(s: pl.Series) -> pl.Series:
    """去掉空白，按 Int64 → Float64 → 原字符串的顺序尝试转换。"""
    stripped = s.str.strip_chars()
    for dtype in (pl.Int64, pl.Float64):
        try:
            return stripped.cast(dtype)
        except InvalidOperationError:
            continue
    return stripped


# 各列的类型会变，所以先按列应用，再合成一个数据框
df_converted = pl.DataFrame({name: try_cast_to_number(df1[name]) for name in df1.columns})
row(df1, df_converted)

shape: (2, 3)

str	int	float
str	str	str
"日本語 "	"4 "	"5.67 "
"abc "	"5 "	"1.23 "

shape: (2, 3)

str	int	float
str	i64	f64
"日本語"	4	5.67
"abc"	5	1.23

读取多个文件#

实际数据常常按天、按对象拆成很多文件。Polars 的路径支持通配符，可以把这些文件一次读成一张表。

通配符的写法与 shell 类似：* 匹配单层目录，** 递归匹配多层。读取时 Polars 会把所有文件视为一个数据源，列名和类型不一致时尝试按上位类型统一，无法统一时会报错。下面的例子批量读取股票价格的 CSV，先把数据下载到 data/stock。

from helper.utils import download_folder_from_github, print_folder_structure
download_folder_from_github("https://github.com/jeroenjanssens/python-polars-the-definitive-guide/tree/main/data/stock", "data/stock")

data/stock exists, please delete and try again.

下载后的目录按股票代码划分子目录，每个子目录中存放按年份命名的 CSV 文件。

print_folder_structure("data/stock", max_file_count=3)

|-- asml
    |-- 1999.csv
    |-- 2000.csv
    |-- 2001.csv
    |-- ...
|-- nvda
    |-- 1999.csv
    |-- 2000.csv
    |-- 2001.csv
    |-- ...
|-- tsm
    |-- 1999.csv
    |-- 2000.csv
    |-- 2001.csv
    |-- ...

要一次读入所有 CSV，在路径里用 ** 递归匹配子目录即可。Polars 会把所有文件视为一个数据源，列名和类型不一致时按上位类型统一。

df_stock = pl.read_csv('data/stock/**/*.csv', try_parse_dates=True)
df_stock

shape: (18_476, 8)

symbol	date	open	high	low	close	adj close	volume
str	date	f64	f64	f64	f64	f64	i64
"ASML"	1999-01-04	11.765625	12.28125	11.765625	12.140625	7.522523	1801867
"ASML"	1999-01-05	11.859375	14.25	11.71875	13.96875	8.655257	8241600
"ASML"	1999-01-06	14.25	17.601563	14.203125	16.875	10.456018	16400267
"ASML"	1999-01-07	14.742188	17.8125	14.53125	16.851563	10.441495	17722133
"ASML"	1999-01-08	16.078125	16.289063	15.023438	15.796875	9.787995	10696000
…	…	…	…	…	…	…	…
"TSM"	2023-06-26	102.019997	103.040001	100.089996	100.110001	99.125954	8560000
"TSM"	2023-06-27	101.150002	102.790001	100.019997	102.080002	101.076591	9732000
"TSM"	2023-06-28	100.5	101.879997	100.220001	100.919998	99.927986	8160900
"TSM"	2023-06-29	101.339996	101.519997	100.019997	100.639999	99.650742	7383900
"TSM"	2023-06-30	101.400002	101.889999	100.410004	100.919998	99.927986	11701700

这份数据本身已包含股票名的列，不必再把文件名作为列。想知道某行来自哪个文件时，用 include_file_paths 指定列名即可。

Polars 2.0 的 read_csv() 也接受这个参数，但这里用延迟求值的 scan_csv()：它返回 LazyFrame，不会立刻读取数据，需要查看内容时才调用 collect()。

df_stock = pl.scan_csv('data/stock/**/*.csv', try_parse_dates=True, include_file_paths="path")
df_stock

NAIVE QUERY PLAN

run LazyFrame.show_graph() to see the optimized version

df_stock.collect()

shape: (18_476, 9)

symbol	date	open	high	low	close	adj close	volume	path
str	date	f64	f64	f64	f64	f64	i64	str
"ASML"	1999-01-04	11.765625	12.28125	11.765625	12.140625	7.522523	1801867	"data/stock/asml/1999.csv"
"ASML"	1999-01-05	11.859375	14.25	11.71875	13.96875	8.655257	8241600	"data/stock/asml/1999.csv"
"ASML"	1999-01-06	14.25	17.601563	14.203125	16.875	10.456018	16400267	"data/stock/asml/1999.csv"
"ASML"	1999-01-07	14.742188	17.8125	14.53125	16.851563	10.441495	17722133	"data/stock/asml/1999.csv"
"ASML"	1999-01-08	16.078125	16.289063	15.023438	15.796875	9.787995	10696000	"data/stock/asml/1999.csv"
…	…	…	…	…	…	…	…	…
"TSM"	2023-06-26	102.019997	103.040001	100.089996	100.110001	99.125954	8560000	"data/stock/tsm/2023.csv"
"TSM"	2023-06-27	101.150002	102.790001	100.019997	102.080002	101.076591	9732000	"data/stock/tsm/2023.csv"
"TSM"	2023-06-28	100.5	101.879997	100.220001	100.919998	99.927986	8160900	"data/stock/tsm/2023.csv"
"TSM"	2023-06-29	101.339996	101.519997	100.019997	100.639999	99.650742	7383900	"data/stock/tsm/2023.csv"
"TSM"	2023-06-30	101.400002	101.889999	100.410004	100.919998	99.927986	11701700	"data/stock/tsm/2023.csv"

非 CSV 的文本只要可以按行读成一列，也能用 Polars 的字符串操作解析。这比写 Python 循环快得多，因为解析本身也由 Polars 的引擎执行。

这类文本常见于日志、调试输出、配置转储：每行格式相似，但没有统一的分隔符。把它读成一列字符串，再用 str.split()、str.extract_groups() 拆开，就能在列式引擎里完成解析。下面的 variables_dump.txt 每行包含多个并列的“变量 = 值”转储。

%%writefile data/variables_dump.txt
A = 1, B = 2, C = 3.23
A = 11, B = 22, C = 33.23
B = 30, C = 20.0
D = 10, A = 3

Overwriting data/variables_dump.txt

读取时把 separator 指定为 NULL 字符 "\x00"：文本中不会出现这个字符，因此行不会被切分，每行整体成为一个值。

之后用下面的操作重新组织成表：

.str.split(",")：按 , 切分各行。
.explode("s")：将切分得到的列表展开为行。
.str.extract_groups(...)：用正则取出变量名(name)和值(value)。
.pivot(...)：将变量名转换为列，一行对应一个样本。
.cast(pl.Float64)：将值转换为浮点数。

最后找出小数部分全为 0 的列，转回 Int64，这样整数列不会显示成 3.0。

df = pl.read_csv(
    "data/variables_dump.txt", has_header=False, separator="\x00", new_columns=["s"]
)

df_vars = (
    df.select(pl.col("s").str.split(","))
    .with_row_index()
    .explode("s")  # 一个单元格一个变量
    .with_columns(
        pl.col("s").str.extract_groups(r"(?<name>\w)+\s*=\s*(?<value>.+)")
    )
    .unnest("s")  # 分成 name / value 两列
    .pivot(on="name", index="index")
    .with_columns(pl.exclude("index").cast(pl.Float64))
)

# 小数部分恒为 0 的列转回 Int64
int_cols = (
    df_vars.select((pl.col(pl.Float64) % 1 == 0).all())
    .transpose(include_header=True)
    .filter("column_0")
    .get_column("column")
    .to_list()
)

df_vars = df_vars.with_columns(pl.col(int_cols).cast(pl.Int64))
row(df, df_vars)

shape: (4, 1)

s
str
"A = 1, B = 2, C = 3.23"
"A = 11, B = 22, C = 33.23"
"B = 30, C = 20.0"
"D = 10, A = 3"

shape: (4, 5)

index	A	B	C	D
u32	i64	i64	f64	i64
0	1	2	3.23	null
1	11	22	33.23	null
2	null	30	20.0	null
3	3	null	null	10

写出#

用 DataFrame.write_csv() 可以把数据框直接输出为 CSV 文件，参数和读取端基本对称。

df1.write_csv('data/csv_output_utf8.csv')

write_csv() 可以指定 include_header、separator、line_terminator 等格式参数，但没有直接指定编码的选项。

旧系统需要 Shift-JIS 时，可以先写入 StringIO，再把内容转码保存。虽然多一步处理，但不影响性能。

import io

buf = io.StringIO()
df1.write_csv(buf)

with open('data/csv_output_shiftjis.csv', 'w', encoding='shift-jis') as f:
    f.write(buf.getvalue())

Excel 文件#

Excel 的读写使用 read_excel() 和 DataFrame.write_excel()。

读取支持 calamine、xlsx2csv、openpyxl 三个引擎，默认的 calamine 最快，因为它用 Rust 直接解析压缩包，不经过 Python 层。xlsx2csv 和 openpyxl 则作为备选，遇到 calamine 解析不了的文件时可以换用。写出使用 xlsxwriter 库。必要的库用下面的命令安装，缺少任何一个，导入时都会直接报错。

conda install fastexcel xlsxwriter

读取#

用 read_excel() 指定工作表读取：sheet_id 传入从 1 开始的工作表序号，或由序号组成的列表。

传入 0 则读取所有工作表，返回 {工作表名: DataFrame} 的字典。也可以用 sheet_names 直接指定工作表名，适合工作表顺序会变的情况。

工作表序号和名称各有用处：序号写起来短，但插入或删除工作表后会错位；名称更稳定，适合长期维护的表格。下面的例子先按序号读取前两个工作表，再读取全部工作表。

fn = 'data/xlsx_example.xlsx'
df1, df2 = pl.read_excel(fn, sheet_id=[1, 2]).values()
row(df1, df2)

shape: (3, 2)

A	B
i64	str
1	"aa"
3	"bb"
5	"cc"

shape: (3, 3)

x	y	z
f64	f64	f64
1.0	1.0	3.1
1.2	2.1	4.3
3.2	4.3	5.4

df1, df2, df3 = pl.read_excel(fn, sheet_id=0).values()
row(df1, df2, df3)

shape: (3, 2)

A	B
i64	str
1	"aa"
3	"bb"
5	"cc"

shape: (3, 3)

x	y	z
f64	f64	f64
1.0	1.0	3.1
1.2	2.1	4.3
3.2	4.3	5.4

shape: (4, 2)

A	B
str	str
"a"	"b"
"0"	"1"
"2"	"3"
"4"	"5"

Sheet3 的表头占两行，直接读取会把第 2 行当作数据，于是所有列都变成字符串，这是读取 Excel 时最容易出问题的地方。

修改 read_options 中传给引擎的设置即可解决：下面的代码先把开头 2 行作为数据读出并组合出列名，再跳过表头并应用这些列名。

df_header = pl.read_excel(fn, sheet_id=3, read_options=dict(n_rows=2, header_row=None))
columns = df_header.select(pl.all().str.join('-')).row(0)
df3 = pl.read_excel(fn, sheet_id=3, read_options=dict(skip_rows=2, column_names=columns))
df3

shape: (3, 2)

A-a	B-b
i64	i64
0	1
2	3
4	5

calamine 引擎能接受哪些参数，可以用 ExcelReader.load_sheet? 查看。

from fastexcel import ExcelReader

ExcelReader.load_sheet?

写出#

用 DataFrame.write_excel() 把数据写入指定的文件名和工作表名。

要把多个数据框写入同一文件的不同工作表，先创建 Workbook 对象，再对它多次调用 .write_excel()，这样只打开一次文件，不必反复读写同一个 xlsx。

import xlsxwriter

with xlsxwriter.Workbook("data/xlsx_example_output.xlsx") as wb:
    df1.write_excel(wb, "df1")
    df2.write_excel(wb, "df2")
    df3.write_excel(wb, "df3")    

JSON 文件#

JSON 可以直接表达嵌套结构，在 Web API 的响应和配置文件中十分常见。read_json() 可以直接接收 bytes，所以不必先保存为文件。

JSON 的灵活性来自嵌套：一个对象里可以放数组，数组里又可以放对象。Polars 把这种层次映射成 List 和 Struct 类型，读进来之后仍然是嵌套列，不会自动展平。下面的例子从一段字节串读取 JSON。

json_data1 = b'''
[
    {"color":"red","value":"#f00"},
    {"color":"green","value":"#0f0"},
    {"color":"blue","value":"#00f"}
]
'''

df = pl.read_json(json_data1)
df

shape: (3, 2)

color	value
str	str
"red"	"#f00"
"green"	"#0f0"
"blue"	"#00f"

每行存放一个 JSON 对象的格式是 NDJSON，适合流式写入和按行处理，用 read_ndjson() 读取。

ndjson_data = b'''
{"color":"red","value":"#f00"}
{"color":"green","value":"#0f0"}
{"color":"blue","value":"#00f"}
'''

df = pl.read_ndjson(ndjson_data)
df

shape: (3, 2)

color	value
str	str
"red"	"#f00"
"green"	"#0f0"
"blue"	"#00f"

把 JSON 转换为数据框只有两条规则。

JSON 按字典的列表处理，每个字典是一行，键是列。最外层是字典时，当作只有 1 个元素的列表。
嵌套的列表转换为 List 类型，嵌套的字典转换为 Struct 类型。

这两条规则决定了读入后的形状：顶层字典的每个键成为一列；值是列表时得到 List 列，值是字典时得到 Struct 列。读入下面的数据，会得到 3 行、列为 id、type、name、ppu、topping 的数据框，其中 topping 是“结构体的列表”(List[Struct])，嵌套结构被原样保留，没有被展平。

json_data2 = b'''
[
	{
		"id": "0001",
		"type": "donut",
		"name": "Cake",
		"ppu": 0.55,
		"topping":
			[
				{ "id": "5001", "type": "None" },
				{ "id": "5002", "type": "Glazed" },
				{ "id": "5005", "type": "Sugar" },
				{ "id": "5007", "type": "Powdered Sugar" },
				{ "id": "5006", "type": "Chocolate with Sprinkles" },
				{ "id": "5003", "type": "Chocolate" },
				{ "id": "5004", "type": "Maple" }
			]
	},
	{
		"id": "0002",
		"type": "donut",
		"name": "Raised",
		"ppu": 0.55,
		"topping":
			[
				{ "id": "5001", "type": "None" },
				{ "id": "5002", "type": "Glazed" },
				{ "id": "5005", "type": "Sugar" },
				{ "id": "5003", "type": "Chocolate" },
				{ "id": "5004", "type": "Maple" }
			]
	},
	{
		"id": "0003",
		"type": "donut",
		"name": "Old Fashioned",
		"ppu": 0.55,
		"topping":
			[
				{ "id": "5001", "type": "None" },
				{ "id": "5002", "type": "Glazed" },
				{ "id": "5003", "type": "Chocolate" },
				{ "id": "5004", "type": "Maple" }
			]
	}
]
'''
df = pl.read_json(json_data2)
df

shape: (3, 5)

id	type	name	ppu	topping
str	str	str	f64	list[struct[2]]
"0001"	"donut"	"Cake"	0.55	[{"5001","None"}, {"5002","Glazed"}, … {"5004","Maple"}]
"0002"	"donut"	"Raised"	0.55	[{"5001","None"}, {"5002","Glazed"}, … {"5004","Maple"}]
"0003"	"donut"	"Old Fashioned"	0.55	[{"5001","None"}, {"5002","Glazed"}, … {"5004","Maple"}]

topping 列的数据类型如下所示。

print(df.schema['topping'])

List(Struct({'id': String, 'type': String}))

嵌套结构不便处理时，用 explode() 把列表展开为行，用 struct.unnest() 把结构体展开为列，两步即可展平成普通表。

df.explode('topping').select(
    pl.exclude('topping'),
    pl.col('topping').struct.unnest().name.prefix('topping_')
)

shape: (16, 6)

id	type	name	ppu	topping_id	topping_type
str	str	str	f64	str	str
"0001"	"donut"	"Cake"	0.55	"5001"	"None"
"0001"	"donut"	"Cake"	0.55	"5002"	"Glazed"
"0001"	"donut"	"Cake"	0.55	"5005"	"Sugar"
"0001"	"donut"	"Cake"	0.55	"5007"	"Powdered Sugar"
"0001"	"donut"	"Cake"	0.55	"5006"	"Chocolate with Sprinkles"
…	…	…	…	…	…
"0002"	"donut"	"Raised"	0.55	"5004"	"Maple"
"0003"	"donut"	"Old Fashioned"	0.55	"5001"	"None"
"0003"	"donut"	"Old Fashioned"	0.55	"5002"	"Glazed"
"0003"	"donut"	"Old Fashioned"	0.55	"5003"	"Chocolate"
"0003"	"donut"	"Old Fashioned"	0.55	"5004"	"Maple"

也可以把原数据框拆成 3 张表再连接回原表：一张存储主体(甜甜圈)，一张存储配料(去重后的 topping)，一张存储两者的对应关系。

这是关系型建模的常规做法，数据量大、需要复用配料信息时比展平更节省空间。

df_topping = (
    df.select(
        pl.col('topping').explode().unique().struct.unnest().name.prefix(f'topping_')
    )
    .sort(by=f'topping_id')
)

df_donut = df.drop('topping')
df_donut_topping = (
    df.select(
        'id', 
        topping_id = pl.col('topping').list.eval(
            pl.element().struct.field('id')
        )
    )
    .explode(f'topping_id')
)

row(df_donut, df_donut_topping, df_topping)

shape: (3, 4)

id	type	name	ppu
str	str	str	f64
"0001"	"donut"	"Cake"	0.55
"0002"	"donut"	"Raised"	0.55
"0003"	"donut"	"Old Fashioned"	0.55

shape: (16, 2)

id	topping_id
str	str
"0001"	"5001"
"0001"	"5002"
"0001"	"5005"
"0001"	"5007"
"0001"	"5006"
…	…
"0002"	"5004"
"0003"	"5001"
"0003"	"5002"
"0003"	"5003"
"0003"	"5004"

shape: (7, 2)

topping_id	topping_type
str	str
"5001"	"None"
"5002"	"Glazed"
"5003"	"Chocolate"
"5004"	"Maple"
"5005"	"Sugar"
"5006"	"Chocolate with Sprinkles"
"5007"	"Powdered Sugar"

拆开的 3 张表用 join() 即可合并回一张表。

df_donut.join(df_donut_topping, on='id').join(df_topping, on='topping_id')

shape: (16, 6)

id	type	name	ppu	topping_id	topping_type
str	str	str	f64	str	str
"0001"	"donut"	"Cake"	0.55	"5001"	"None"
"0001"	"donut"	"Cake"	0.55	"5002"	"Glazed"
"0001"	"donut"	"Cake"	0.55	"5005"	"Sugar"
"0001"	"donut"	"Cake"	0.55	"5007"	"Powdered Sugar"
"0001"	"donut"	"Cake"	0.55	"5006"	"Chocolate with Sprinkles"
…	…	…	…	…	…
"0002"	"donut"	"Raised"	0.55	"5004"	"Maple"
"0003"	"donut"	"Old Fashioned"	0.55	"5001"	"None"
"0003"	"donut"	"Old Fashioned"	0.55	"5002"	"Glazed"
"0003"	"donut"	"Old Fashioned"	0.55	"5003"	"Chocolate"
"0003"	"donut"	"Old Fashioned"	0.55	"5004"	"Maple"

再看一个更复杂的例子：读取纽约市各区的多边形(GeoJSON)。

GeoJSON 是 JSON 的一种，几何数据层层嵌套，很适合用来练习嵌套结构的处理。它的顶层是一个 FeatureCollection 字典，features 字段是一个列表，每个元素包含 properties 和 geometry 两个结构体；geometry 里又有 coordinates 这样的多层列表。

df = pl.read_json("data/nyc-neighborhoods.geojson")
df

shape: (1, 3)

type	crs	features
str	struct[2]	list[struct[3]]
"FeatureCollection"	{"name",{"urn:ogc:def:crs:OGC:1.3:CRS84"}}	[{"Feature",{"Allerton","2","Bronx","http://nyc.pediacities.com/Resource/Neighborhood/Allerton"},{"Polygon",[[[-73.848597, 40.87167], [-73.845823, 40.870239], … [-73.848597, 40.87167]]]}}, {"Feature",{"Alley Pond Park","4","Queens","http://nyc.pediacities.com/Resource/Neighborhood/Alley_Pond_Park"},{"Polygon",[[[-73.743333, 40.738883], [-73.743714, 40.739403], … [-73.743333, 40.738883]]]}}, … {"Feature",{"Harlem","1","Manhattan",null},{"Polygon",[[[-73.934568, 40.828149], [-73.934424, 40.827043], … [-73.934568, 40.828149]]]}}]

这个文件最外层是一个字典，所以结果只有 1 行。地理信息都位于 features 列中，嵌套层次很深。

print(df.schema['features'])

List(Struct({'type': String, 'properties': Struct({'neighborhood': String, 'boroughCode': String, 'borough': String, 'X.id': String}), 'geometry': Struct({'type': String, 'coordinates': List(List(List(Float64)))})}))

为了看清深度嵌套的结构，本书提供了 dtype_pretty_format() 函数，把嵌套的数据类型按树状格式显示。

from helper.polars import dtype_pretty_format
print(dtype_pretty_format(df.schema['features']))

List(
  Struct(
    type: String
    properties:
      Struct(
        neighborhood: String
        boroughCode: String
        borough: String
        X.id: String
      )
    geometry:
      Struct(
        type: String
        coordinates:
          List(
            List(
              List(Float64)
            )
          )
      )
  )
)

features 是一个列表，元素是包含 properties 和 geometry 的结构体。

下面的代码把各字段展开为列，并让列表的每个元素成为一行。coordinates 外面还包裹着一层长度为 1 的列表，也一并拆开。

df_nb = (
    df
    .select(
        pl.col('features').list.explode().struct.unnest()
    )
    .select(
        pl.col('properties', 'geometry').struct.unnest()
    )
    .with_columns(
        pl.col('coordinates').list.get(0)
    )
)

df_nb

shape: (312, 6)

neighborhood	boroughCode	borough	X.id	type	coordinates
str	str	str	str	str	list[list[f64]]
"Allerton"	"2"	"Bronx"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-73.848597, 40.87167], [-73.845823, 40.870239], … [-73.848597, 40.87167]]
"Alley Pond Park"	"4"	"Queens"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-73.743333, 40.738883], [-73.743714, 40.739403], … [-73.743333, 40.738883]]
"Arden Heights"	"5"	"Staten Island"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-74.169827, 40.561078], [-74.169823, 40.561057], … [-74.169827, 40.561078]]
"Arlington"	"5"	"Staten Island"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-74.159748, 40.641417], [-74.159979, 40.641446], … [-74.159748, 40.641417]]
"Arrochar"	"5"	"Staten Island"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-74.06078, 40.593188], [-74.060792, 40.593178], … [-74.06078, 40.593188]]
…	…	…	…	…	…
"Woodside"	"4"	"Queens"	"http://nyc.pediacities.com/Res…	"Polygon"	[[-73.900866, 40.757674], [-73.90014, 40.75615], … [-73.900866, 40.757674]]
"East Williamsburg"	"3"	"Brooklyn"	null	"Polygon"	[[-73.9413, 40.7218], [-73.9406, 40.7186], … [-73.9413, 40.7218]]
"Hamilton Heights"	"1"	"Manhattan"	null	"Polygon"	[[-73.9505, 40.8349], [-73.9411, 40.8312], … [-73.9505, 40.8349]]
"Williamsburg"	"3"	"Brooklyn"	null	"Polygon"	[[-73.957572, 40.725097], [-73.952998, 40.722253], … [-73.957572, 40.725097]]
"Harlem"	"1"	"Manhattan"	null	"Polygon"	[[-73.934568, 40.828149], [-73.934424, 40.827043], … [-73.934568, 40.828149]]

要用 holoviews 的 Polygons 画图，需要把经度和纬度的列表分成两列。

list.eval() 是对列表每个元素求值的表达式，pl.element() 表示当前元素，list.get() 从嵌套列表中取值，一次拆开两层嵌套。

import holoviews as hv
hv.extension('bokeh', inline=False)

ⓘ

df_nb2 = df_nb.select(
    pl.col('neighborhood'),
    pl.col('borough'),
    pl.col('coordinates').list.eval(pl.element().list.get(0)).alias('lng'),
    pl.col('coordinates').list.eval(pl.element().list.get(1)).alias('lat'),
)

hv.Polygons(df_nb2.to_dicts(), kdims=['lng', 'lat'], vdims=['borough']).opts(cmap='tab10', data_aspect=1)

二进制文件#

二进制格式方面，Polars 支持 IPC(Arrow IPC)和 Parquet，两者都比 CSV 快得多，因为不必解析文本，类型信息也直接存储在文件中。

IPC：进程间传递数据的二进制格式，序列化和反序列化都快，大数据可以原样传递。
Parquet：列式存储格式，压缩率高，读取时只取需要的列，Spark 等分布式处理也把它当作标准。

选择哪一种，取决于数据要交给谁：同一套 Arrow 生态内部传递，IPC 最直接；需要长期存储、跨工具读取，Parquet 更通用。DataFrame 使用 read_*() / write_*()，LazyFrame 使用 scan_*() / sink_*()。

read_ipc() / read_parquet()：从文件读取为 DataFrame。
write_ipc() / write_parquet()：把 DataFrame 写入文件。
scan_ipc() / scan_parquet()：把文件读取为 LazyFrame。
sink_ipc() / sink_parquet()：把 LazyFrame 按块写出。

sink_*() 不会一次把数据全部放入内存，而是分块处理，所以内存装不下的规模也能处理。下面的代码把多个 CSV 转换为一个 IPC 文件。

df_stock = pl.scan_csv('data/stock/**/*.csv', try_parse_dates=True)
df_stock.sink_ipc('data/stock.arrow')

df_stock = pl.read_ipc('data/stock.arrow')
df_stock

shape: (18_476, 8)

symbol	date	open	high	low	close	adj close	volume
str	date	f64	f64	f64	f64	f64	i64
"ASML"	1999-01-04	11.765625	12.28125	11.765625	12.140625	7.522523	1801867
"ASML"	1999-01-05	11.859375	14.25	11.71875	13.96875	8.655257	8241600
"ASML"	1999-01-06	14.25	17.601563	14.203125	16.875	10.456018	16400267
"ASML"	1999-01-07	14.742188	17.8125	14.53125	16.851563	10.441495	17722133
"ASML"	1999-01-08	16.078125	16.289063	15.023438	15.796875	9.787995	10696000
…	…	…	…	…	…	…	…
"TSM"	2023-06-26	102.019997	103.040001	100.089996	100.110001	99.125954	8560000
"TSM"	2023-06-27	101.150002	102.790001	100.019997	102.080002	101.076591	9732000
"TSM"	2023-06-28	100.5	101.879997	100.220001	100.919998	99.927986	8160900
"TSM"	2023-06-29	101.339996	101.519997	100.019997	100.639999	99.650742	7383900
"TSM"	2023-06-30	101.400002	101.889999	100.410004	100.919998	99.927986	11701700

上面的函数是一个文件对应一个数据框，而使用 write_ipc_stream() / read_ipc_stream() 可以把多个数据框依次写入一个文件。

每个数据框是一条独立的记录流，读取时按写入顺序依次读取。下面的代码把两个数据框依次写入同一个文件。

df1 = pl.DataFrame({"a":[1, 2, 3], "b":["xx", "yyy", "zzzz"]})
df2 = pl.DataFrame({"x":[3.0, 4.0, 6.0, 7.0], "y":[1.0, 2.0, 3.0, 4.0], "z":[0.0, 0.1, 0.2, 0.3]})

with open('data/stream_test.arrow', 'wb') as f:
    df1.write_ipc_stream(f)
    df2.write_ipc_stream(f)

读取时按写出的顺序依次读取。

备注

在 Polars 2.0 中，连续读取仍需设置 use_pyarrow=True。使用默认设置读不到第 2 个及之后的数据框，这是目前的实现限制，不是用法问题。

with open("data/stream_test.arrow", "rb") as f:
    df1_r = pl.read_ipc_stream(f, use_pyarrow=True)
    df2_r = pl.read_ipc_stream(f, use_pyarrow=True)

print(df1.equals(df1_r), df2.equals(df2_r))

True True