用 Python 定义 Schema 并生成 Parquet 文件

原来用 Java 和 Python 实现过 Avro 转换成 Parquet 格式,所以 Schema 都是在 Avro 中定义的。这里要尝试的是如何定义 Parquet 的 Schema, 然后据此填充数据并生成 Parquet 文件。

本文将演示两个例子,一个是没有层级的两个字段,另一个是含于嵌套级别的字段,将要使用到的 Python 模块有 pandas 和 pyarrow

简单字段定义

定义 Schema 并生成 Parquet 文件

 1import pandas as pd
 2import pyarrow as pa
 3import pyarrow.parquet as pq
 4
 5# 定义 Schema
 6schema = pa.schema([
 7    ('id', pa.int32()),
 8    ('email', pa.string())
 9])
10
11# 准备数据
12ids = pa.array([1, 2], type = pa.int32())
13emails = pa.array(['first@example.com', 'second@example.com'], pa.string())
14
15# 生成 Parquet 数据
16batch = pa.RecordBatch.from_arrays(
17    [ids, emails],
18    schema = schema
19)
20table = pa.Table.from_batches([batch])
21
22# 写 Parquet 文件 plain.parquet
23pq.write_table(table, 'plain.parquet')

验证 Parquet 数据文件

我们可以用工具 parquet-tools  来查看 plain.parquet 文件的数据和 Schema
$ parquet-tools schema plain.parquet
message schema {
    optional int32 id;
    optional binary email (STRING);
} $ parquet-tools cat --json plain.parquet
{"id":1,"email":"first@example.com"}
{"id":2,"email":"second@example.com"}
没问题,与我们期望的一致。也可以用 pyarrow 代码来获取其中的 Schema 和数据
1schema = pq.read_schema('plain.parquet')
2print(schema)
3
4df = pd.read_parquet('plain.parquet')
5print(df.to_json())

输出为
1id: int32
2  -- field metadata --
3  PARQUET:field_id: '1'
4email: string
5  -- field metadata --
6  PARQUET:field_id: '2'
7{"id":{"0":1,"1":2},"email":{"0":"first@example.com","1":"second@example.com"}}

含嵌套字段定义

下面的 Schema 定义加入一个嵌套对象,在 address 下分 email_address 和 post_address,Schema 定义及生成 Parquet 文件的代码如下

 1import pandas as pd
 2import pyarrow as pa
 3import pyarrow.parquet as pq
 4
 5# 内部字段
 6address_fields = [
 7    ('email_address', pa.string()),
 8    ('post_address', pa.string()),
 9]
10
11# 定义 Parquet Schema,address 嵌套了 address_fields
12schema = pa.schema(j)
13
14# 准备数据
15ids = pa.array([1, 2], type = pa.int32())
16addresses = pa.array(
17    [('first@example.com', 'city1'), ('second@example.com', 'city2')],
18    pa.struct(address_fields)
19)
20
21# 生成 Parquet 数据
22batch = pa.RecordBatch.from_arrays(
23    [ids, addresses],
24    schema = schema
25)
26table = pa.Table.from_batches([batch])
27
28# 写 Parquet 数据到文件
29pq.write_table(table, 'nested.parquet')

验证 Parquet 数据文件

同样用 parquet-tools 来查看下 nested.parquet 文件
$ parquet-tools schema nested.parquet
message schema {
    optional int32 id;
    optional group address {
        optional binary email_address (STRING);
        optional binary post_address (STRING);
    }
} $ parquet-tools cat --json nested.parquet
{"id":1,"address":{"email_address":"first@example.com","post_address":"city1"}}
{"id":2,"address":{"email_address":"second@example.com","post_address":"city2"}}
parquet-tools 看到的 Schama 并没有 struct 的字样,但体现了它 address 与下级属性的嵌套关系。

用 pyarrow 代码来读取 nested.parquet 文件的 Schema 和数据是什么样子
1schema = pq.read_schema("nested.parquet")
2print(schema)
3
4df = pd.read_parquet('nested.parquet')
5print(df.to_json())

 1id: int32
 2  -- field metadata --
 3  PARQUET:field_id: '1'
 4address: struct<email_address: string, post_address: string>
 5  child 0, email_address: string
 6    -- field metadata --
 7    PARQUET:field_id: '3'
 8  child 1, post_address: string
 9    -- field metadata --
10    PARQUET:field_id: '4'
11  -- field metadata --
12  PARQUET:field_id: '2'
13{"id":{"0":1,"1":2},"address":{"0":{"email_address":"first@example.com","post_address":"city1"},"1":{"email_address":"second@example.com","post_address":"city2"}}}

数据当然是一样的,有略微不同的是显示的 Schema 中, address 标识为 struct<email_address: string, post_address: string>, 明确的表明它是一个 struct 类型,而不是只展示嵌套层次。

最后留下一个问题,前面我们定义 Parquet Schema 都是在 Python 代码中完成了,Parquet 是否也能像 Avro 一样用外部文件来定义 Schema,  然后编译给 Python 用?

链接:

  1. Data Types and In-Memory Data Model
  2. pyarrow.struct
  3. How to write to a Parquet file in Python
永久链接 https://yanbin.blog/python-define-schema-generate-parquet-file/, 来自 隔叶黄莺 Yanbin's Blog
[版权声明] 本文采用 署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0) 进行许可。