使用 Schema

Arrow 在读取数据或将 Python 对象转换为 Arrow 对象时，会自动推断最合适的数据类型。

但是，您可能希望手动告诉 Arrow 使用哪些数据类型，例如，为了确保与数据库和数据仓库系统的互操作性。本章包含处理模式的食谱。

设置 Arrow 数组的数据类型

如果您有一个现有的数组，并且想要更改其数据类型，可以通过 cast 函数来完成。

import pyarrow as pa

arr = pa.array([1, 2, 3, 4, 5])
print(arr.type)

int64

arr = arr.cast(pa.int8())
print(arr.type)

int8

您也可以通过在创建数组时提供类型来创建请求类型的数组。

import pyarrow as pa

arr = pa.array([1, 2, 3, 4, 5], type=pa.int8())
print(arr.type)

int8

设置表的 Schema

表包含多个列，每个列都有自己的名称和类型。类型和名称的并集定义了 Schema。

Arrow 中的 Schema 可以使用 pyarrow.schema() 定义。

import pyarrow as pa

schema = pa.schema([
    ("col1", pa.int8()),
    ("col2", pa.string()),
    ("col3", pa.float64())
])

然后可以在创建表时将 Schema 提供给表。

table = pa.table([
    [1, 2, 3, 4, 5],
    ["a", "b", "c", "d", "e"],
    [1.0, 2.0, 3.0, 4.0, 5.0]
], schema=schema)

print(table)

pyarrow.Table
col1: int8
col2: string
col3: double
----
col1: [[1,2,3,4,5]]
col2: [["a","b","c","d","e"]]
col3: [[1,2,3,4,5]]

与数组一样，只要表兼容，就可以将表转换为不同的 Schema。

schema_int32 = pa.schema([
    ("col1", pa.int32()),
    ("col2", pa.string()),
    ("col3", pa.float64())
])

table = table.cast(schema_int32)

print(table)

pyarrow.Table
col1: int32
col2: string
col3: double
----
col1: [[1,2,3,4,5]]
col2: [["a","b","c","d","e"]]
col3: [[1,2,3,4,5]]

合并多个 Schema

当您有多个独立的数据组需要合并时，可能需要统一它们的 Schema，以创建一个适用于所有数据源的超集。

import pyarrow as pa

first_schema = pa.schema([
    ("country", pa.string()),
    ("population", pa.int32())
])

second_schema = pa.schema([
    ("country_code", pa.string()),
    ("language", pa.string())
])

unify_schemas() 可用于将多个 Schema 合并为一个。

union_schema = pa.unify_schemas([first_schema, second_schema])

print(union_schema)

country: string
population: int32
country_code: string
language: string

如果合并的 Schema 具有重叠的列，只要冲突的列保留相同的类型（country_code），它们仍然可以合并。

third_schema = pa.schema([
    ("country_code", pa.string()),
    ("lat", pa.float32()),
    ("long", pa.float32()),
])

union_schema =  pa.unify_schemas([first_schema, second_schema, third_schema])

print(union_schema)

country: string
population: int32
country_code: string
language: string
lat: float
long: float

如果合并的字段在合并的 Schema 中具有不同的类型，则尝试合并 Schema 将失败。例如，如果 country_code 是数字而不是字符串，我们将无法统一 Schema，因为在 second_schema 中它已被声明为 pa.string()。

third_schema = pa.schema([
    ("country_code", pa.int32()),
    ("lat", pa.float32()),
    ("long", pa.float32()),
])

try:
    union_schema =  pa.unify_schemas([first_schema, second_schema, third_schema])
except (pa.ArrowInvalid, pa.ArrowTypeError) as e:
    print(e)

Unable to merge: Field country_code has incompatible types: string vs int32