Python 数据处理：Pandas 2.0 与 PyArrow 集成深度解析

import pandas as pd

# Pandas 2.0+ 启用 Arrow -backed DataFrame
pd.set_option("dataframe.default_dtype", "pyarrow")

# 或在读龅时指定
df = pd.read_parquet("data.parquet", dtype_backend="pyarrow")

# 或使用 pyarrow-backed CSV 读取
df = pd.read_csv("data.csv", dtype_backend="pyarrow")

# 查看内存布局
df.memory_usage(deep=True)
# vs NumPy-backed: 内存占用通常减少 30-50%

# NumPy 行式存储 (每行连续)
# [row0: a, b, c, d] [row1: e, f, g, h] [row2: i, j, k, l]
# 访问一行高效，访问一列需要跳步访问

# Arrow 列式存储 (每列连续)
# [col0: a, e, i] [col1: b, f, j] [col2: c, g, k] [col3: d, h, l]
# 访问一列高效，支持 SIMD 批量处理

# Arrow 还包含 metadata（null bitmap, offset array）
# null_count: [0, 1, 0] 表示第2行有 NULL
# offsets: [0, 3, 6, 9] 表示变长字符串的边界

# 读取 Parquet（自动使用 PyArrow）
df = pd.read_parquet("sales.parquet")
# df.dtypes 会显示 pyarrow 类型

# 写入 Parquet
df.to_parquet("output.parquet", engine="pyarrow", compression="zstd")

# 读取远程 Parquet (S3)
import pyarrow.parquet as pq
dataset = pq.ParquetDataset("s3://bucket/data/")
df = dataset.read().to_pandas()

# 分片读取大文件
import pyarrow.parquet as pq
pf = pq.ParquetFile("large.parquet")
for batch in pf.iter_batches(batch_size=100_000):
    process(batch.to_pandas())

# Pandas 2.0 默认启用 CoW
df = pd.DataFrame({"a": [1, 2, 3]})
df_v2 = df  # 不复制，共享 Arrow 数组

# 写入操作触发复制
df_v2.iloc[0, 0] = 99  # 此时才复制

# 如果只是读取，不会复制
df_readonly = df
print(df_readonly.head())  # 无复制开销

# 强制复制（当确认需要独立副本时）
df_copy = df.copy()  # 或 df.copy()
# 使用浅拷贝
df_shallow = df.__copy__()

# Pandas 2.0: 链式操作只在最终写入时复制
result = (
    df.query("category == 'electronics'")
    .groupby("region")
    .agg("revenue": "sum")
    .sort_values("revenue", ascending=False)
)

# 旧版本 Pandas 1.x 行为：
# step1 = df.query()     -> 复制 1
# step2 = step1.groupby() -> 复制 2
# step3 = step2.agg()    -> 复制 3
# result = step3.sort()  -> 复制 4
# 内存峰值 = 4x df size

# Pandas 2.0 CoW 行为：
# 所有步骤共享 Arrow 数组引用
# 仅在 sort_values 写入时复制一次
# 内存峰值 = ~1.5x df size

# 传统方式：整数 + NaN = float（类型丢失）
s = pd.Series([1, 2, 3, np.nan])
s.dtype  # dtype('float64') - 整数信息丢失

# Pandas 2.0 nullable: 保留整数类型
s_nullable = pd.Series([1, 2, 3, None], dtype="Int64")
s_nullable.dtype  # Int64（带大写 I，与 Python int 区分）

# 支持的 nullable 类型
pd.Series([1, 2, None], dtype="Int8")   # Int8/16/32/64
pd.Series([1.0, 2.0, None], dtype="Float64")  # Float32/64
pd.Series([True, False, None], dtype="boolean") # boolean（带小写）
pd.Series(["a", "b", None], dtype="string")      # string 类型

# Arrow 类型 <-> Pandas dtype 映射
pyarrow.int8()   <-> pd.Int8Dtype()
pyarrow.int16()  <-> pd.Int16Dtype()
pyarrow.int32()  <-> pd.Int32Dtype()
pyarrow.int64()  <-> pd.Int64Dtype()
pyarrow.float()  <-> pd.Float32Dtype()
pyarrow.double() <-> pd.Float64Dtype()
pyarrow.bool_() <-> pd.BooleanDtype()
pyarrow.string() <-> pd.StringDtype()

# 验证类型一致性
import pyarrow as pa
arr = pa.array([1, 2, None], type=pa.int64())
s = pd.Series(arr.to_pandas())
print(s.dtype)  # Int64

# 内存占用对比 (1M 行字符串)
# object dtype: ~120 MB（Python object + pointer  overhead）
# StringDtype (PyArrow): ~45 MB（连续 Arrow 字符串）
# StringDtype (NumPy): ~80 MB（NumPy 字符串数组）

import sys
s_object = pd.Series(["hello"]*1_000_000, dtype="object")
s_string = pd.Series(["hello"]*1_000_000, dtype="string")

sys.getsizeof(s_object)    # ~80 MB
sys.getsizeof(s_string)    # ~8 MB（字符串数据在 Arrow 堆外）
s_string.memory_usage()    # ~45 MB（包含 metadata）

# Arrow StringArray 内部结构
# offsets: [0, 5, 10, 16, 22, 28]
# data: "hellohelloworldpythopytho"
#         ^^^^^  ^^^^^ ^^^^^ ^^^^^ ^^^^^
#         str0   str1   str2   str3   str4

# null bitmap (如果有 NULL)
# nulls: [0, 1, 0, 0, 1] -> str1 和 str4 是 NULL

import pyarrow as pa
arr = pa.array(["hello", "world", "python"])
print(arr.type)        # string
print(arr.buffers())   # [null_bitmap, offset, data]

# PyArrow-backed 字符串操作（自动使用向量化实现）
df = pd.DataFrame({"email": ["alice@example.com", "bob@company.org"]})
df["domain"] = df["email"].str.split("@").str[1]
df["is_company"] = df["domain"].str.endswith(".org")

# 使用 .str accessor 的性能
# str.contains / str.replace / str.split 等操作
# 在 PyArrow 后端下使用 pyarrow.compute 引擎
# 性能比纯 Python 实现快 5-10x

# 直接使用 pyarrow.compute (极限优化)
import pyarrow.compute as pc
domains = pc.split_pattern(df["email"].array, "@")
# domains 是 Arrow ChunkedArray，直接操作，无需转换

# Pandas 2.0 支持正则表达式的向量化操作
s = pd.Series(["user123@example.com", "test@domain.org", "invalid"])

# 提取邮箱用户名部分
s.str.extract(r"^([^@]+)@")

# 判断是否是有效邮箱
s.str.match(r"^[a-zA-Z0-9_.]+@[a-zA-Z0-9]+\.[a-zA-Z]{2,}$")

# 批量替换
s.str.replace(r"[0-9]+", "#", regex=True)

# 使用 re2 正则引擎（可选）
# pip install google-re2 后自动启用，内存更安全

import pandas as pd
import time
import tracemalloc

# 启用 PyArrow 后端
pd.set_option("dataframe.default_dtype", "pyarrow")

# 内存基准
tracemalloc.start()
df = pd.read_csv("large_file.csv", dtype_backend="pyarrow")
current, peak = tracemalloc.get_traced_memory()
print(f"Current: {current/1024/1024:.1f}MB, Peak: {peak/1024/1024:.1f}MB")
tracemalloc.stop()

# 操作性能基准
start = time.perf_counter()
result = df.groupby("category").agg("value": "sum")
elapsed = time.perf_counter() - start
print(f"GroupBy: {elapsed*1000:.2f}ms")