Python中如何操作Parquet文件？pyarrow使用指南

import pyarrow.parquet as pq
import pyarrow as pa
import pandas as pd

# 假设我们有一个DataFrame
data = {'col1': [1, 2, 3, 4],
        'col2': ['A', 'B', 'C', 'D'],
        'col3': [True, False, True, False]}
df = pd.DataFrame(data)

# 将Pandas DataFrame写入Parquet文件
# 这里可以指定压缩方式，比如'snappy'，'gzip'，'brotli'等
table = pa.Table.from_pandas(df)
pq.write_table(table, 'example.parquet', compression='snappy')
print("文件已写入: example.parquet")

# 从Parquet文件读取数据
read_table = pq.read_table('example.parquet')
print("\n从Parquet文件读取的数据 (pyarrow.Table):")
print(read_table)

# 如果想转回Pandas DataFrame
read_df = read_table.to_pandas()
print("\n转换回Pandas DataFrame:")
print(read_df)

# 读取特定列
# 有时候我们只关心文件中的几列数据，而不是全部加载
partial_table = pq.read_table('example.parquet', columns=['col1', 'col3'])
print("\n只读取特定列 (col1, col3):")
print(partial_table)

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

# 准备一个DataFrame，包含一些常见的数据类型
df_complex = pd.DataFrame({
    'int_col': [1, 2, None, 4],
    'float_col': [1.1, 2.2, 3.3, None],
    'str_col': ['apple', 'banana', 'cherry', 'date'],
    'date_col': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04']),
    'bool_col': [True, False, True, False]
})

# 写入Parquet，注意这里可以控制Pandas索引是否写入
# 默认情况下，Pandas的索引会作为一列写入Parquet，这在某些场景下可能不是你想要的
table_from_df = pa.Table.from_pandas(df_complex, preserve_index=False)
pq.write_table(table_from_df, 'complex_data.parquet')
print("复杂数据DataFrame已写入: complex_data.parquet")

# 从Parquet读取，并转换回Pandas DataFrame
# 注意：pyarrow在处理缺失值和数据类型推断上可能与Pandas默认行为略有不同
# 比如Pandas的NaN在pyarrow中会映射为null，datetime with timezone可能需要特殊处理
read_back_df = pq.read_table('complex_data.parquet').to_pandas()
print("\n从Parquet读取并转换回Pandas DataFrame:")
print(read_back_df.info())

# 假设有一个非常大的Parquet文件
# 我们可以模拟一个大的Parquet文件
large_df = pd.DataFrame({
    'id': range(1_000_000),
    'value': [float(i) / 100 for i in range(1_000_000)]
})
large_table = pa.Table.from_pandas(large_df)
pq.write_table(large_table, 'large_example.parquet', row_group_size=100000) # 设置行组大小，方便分块读取

# 分块读取大型Parquet文件
parquet_file = pq.ParquetFile('large_example.parquet')
print(f"\n大型文件共有 {parquet_file.num_row_groups} 个行组。")

# 迭代读取每个行组，并转换为Pandas DataFrame
# 这种方式可以有效控制内存使用
for i, batch in enumerate(parquet_file.iter_batches(batch_size=50000)):
    batch_df = batch.to_pandas()
    print(f"处理第 {i+1} 批数据，大小: {len(batch_df)} 行")
    # 在这里可以对 batch_df 进行处理，比如写入数据库或聚合
    if i >= 1: # 只演示前两批
        break

# 假设我们有一个包含日期和数值的Parquet文件
# 模拟数据，包含多个行组
data_rows = []
for i in range(100000):
    data_rows.append({
        'date': pd.Timestamp(f'2023-01-01') + pd.Timedelta(days=i),
        'value': i % 100,
        'category': f'cat_{i % 5}'
    })
large_df_filtered = pd.DataFrame(data_rows)
large_table_filtered = pa.Table.from_pandas(large_df_filtered)
pq.write_table(large_table_filtered, 'filtered_data.parquet', row_group_size=10000)

# 使用谓词下推和列裁剪
# 只读取 'value' 和 'category' 列，并且只读取 value > 90 的行
# 注意：pyarrow的过滤器表达式语法
filtered_table = pq.read_table(
    'filtered_data.parquet',
    columns=['value', 'category'],
    filters=[('value', '>', 90)]
)
print(f"\n使用谓词下推和列裁剪读取的数据量: {len(filtered_table)} 行")
print(filtered_table.to_pandas().head())

import pyarrow.dataset as ds
import os

# 模拟一个分区数据集
# 创建一些目录和文件
os.makedirs('my_partitioned_data/year=2023/month=01', exist_ok=True)
os.makedirs('my_partitioned_data/year=2023/month=02', exist_ok=True)

df_p1 = pd.DataFrame({'col_a': [1, 2], 'col_b': ['x', 'y']})
pq.write_table(pa.Table.from_pandas(df_p1), 'my_partitioned_data/year=2023/month=01/part1.parquet')

df_p2 = pd.DataFrame({'col_a': [3, 4], 'col_b': ['z', 'w']})
pq.write_table(pa.Table.from_pandas(df_p2), 'my_partitioned_data/year=2023/month=02/part2.parquet')

# 使用pyarrow.dataset读取分区数据集
dataset = ds.dataset('my_partitioned_data', format='parquet', partitioning='hive')

# 可以像查询一个表一样查询整个数据集
full_dataset_table = dataset.to_table()
print("\n读取整个分区数据集:")
print(full_dataset_table.to_pandas())

# 也可以对分区进行过滤
# 只读取 month=01 的数据
filtered_dataset_table = dataset.to_table(filter=(ds.field("month") == "01"))
print("\n过滤分区 (month='01') 后读取的数据:")
print(filtered_dataset_table.to_pandas())