Polars 数据处理¶

# 安装Polars（只需一行）
pip install polars  # 基础版

# 安装带所有功能的版本
pip install 'polars[all]'  # 包含Excel、数据库、时区等额外功能

# 验证
python -c "import polars as pl; print(pl.__version__)"

import polars as pl  # 约定简写为pl（类似pandas的pd）

# 从字典创建
df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie", "Diana"],  # 姓名
    "age": [25, 30, 35, 28],                        # 年龄
    "city": ["北京", "上海", "广州", "深圳"],         # 城市
    "salary": [15000, 20000, 25000, 18000],          # 薪资
})

print(df)          # 打印表格（自带漂亮格式）
print(df.schema)   # 查看数据类型
print(df.shape)    # (行数, 列数)
print(df.columns)  # 列名列表

# 读取CSV
df = pl.read_csv("data.csv")  # 急切模式：立即读取全部数据

# 懒读取CSV（推荐，大文件更高效）
lf = pl.scan_csv("data.csv")  # 不立即读取，返回LazyFrame

# 读取Parquet（最推荐的格式）
df = pl.read_parquet("data.parquet")
lf = pl.scan_parquet("data.parquet")  # 懒读取

# 读取JSON
df = pl.read_json("data.json")

# 读取Excel（需要安装openpyxl）
df = pl.read_excel("data.xlsx", sheet_name="Sheet1")

# ===== 选择列 =====
df.select("name", "age")  # 选择指定列
df.select(pl.col("name"), pl.col("age") + 1)  # 用表达式选择

# ===== 过滤行 =====
df.filter(pl.col("age") > 28)  # 年龄大于28
df.filter(
    (pl.col("age") > 25) & (pl.col("city") == "上海")  # 多条件
)

# ===== 添加新列 =====
df = df.with_columns(
    (pl.col("salary") * 12).alias("annual_salary"),  # 年薪 = 月薪 * 12
    pl.when(pl.col("age") < 30)  # 条件判断
        .then(pl.lit("年轻"))
        .otherwise(pl.lit("资深"))
        .alias("age_group"),  # 年龄分组
)

# ===== 排序 =====
df.sort("salary", descending=True)  # 按薪资降序
df.sort(["city", "age"])  # 多列排序

# ===== 分组聚合 =====
df.group_by("city").agg(
    pl.col("salary").mean().alias("avg_salary"),   # 平均薪资
    pl.col("name").count().alias("person_count"),  # 人数
    pl.col("age").max().alias("max_age"),           # 最大年龄
)

# ===== 去重 =====
df.unique(subset=["email"])  # 按email去重
df.n_unique()  # 每列唯一值数量

# 懒执行的优势：Polars会优化整个查询计划再执行，比急切模式更快

result = (
    pl.scan_csv("big_data.csv")  # 1. 懒读取（不加载数据）
    .filter(pl.col("status") == "active")  # 2. 过滤（还没执行）
    .group_by("department")  # 3. 分组（还没执行）
    .agg(
        pl.col("salary").mean().alias("avg_salary"),
        pl.col("id").count().alias("count"),
    )  # 4. 聚合（还没执行）
    .sort("avg_salary", descending=True)  # 5. 排序（还没执行）
    .collect()  # 6. 触发执行！Polars优化后一次性计算
)

# 查看查询计划（理解Polars怎么优化你的查询）
lf = pl.scan_csv("data.csv").filter(pl.col("age") > 25)
print(lf.explain())  # 打印优化后的执行计划

# 在每个部门内排名
df = df.with_columns(
    pl.col("salary")
        .rank(descending=True)  # 降序排名
        .over("department")     # 在department组内执行
        .alias("salary_rank"),  # 部门内薪资排名
)

# 移动平均
df = df.with_columns(
    pl.col("value")
        .rolling_mean(window_size=7)  # 7天移动平均
        .alias("ma_7"),
)

# 内连接
result = df_users.join(
    df_orders,
    on="user_id",    # 连接键
    how="inner",     # inner/left/right/full/cross/semi/anti
)

# 左连接（保留左表所有行）
result = df_users.join(df_orders, on="user_id", how="left")

# Anti Join（找出没有订单的用户）
no_orders = df_users.join(df_orders, on="user_id", how="anti")

df = df.with_columns(
    pl.col("name").str.to_uppercase().alias("name_upper"),  # 转大写
    pl.col("email").str.contains("@gmail").alias("is_gmail"),  # 包含判断
    pl.col("text").str.replace_all(r"\s+", " ").alias("clean_text"),  # 正则替换
    pl.col("date_str").str.to_date("%Y-%m-%d").alias("date"),  # 字符串转日期
)

# 创建时间序列
df = pl.DataFrame({
    "date": pl.date_range(  # 生成日期范围
        start=pl.date(2026, 1, 1),
        end=pl.date(2026, 12, 31),
        interval="1d",  # 每天
        eager=True,
    ),
    "value": range(365),
})

# 按月聚合
monthly = df.group_by_dynamic(
    "date",
    every="1mo",  # 每月
).agg(
    pl.col("value").sum().alias("monthly_total"),
    pl.col("value").mean().alias("monthly_avg"),
)

# 当数据超出内存时，用流式模式
result = (
    pl.scan_csv("huge_file.csv")  # 懒读取
    .filter(pl.col("amount") > 100)
    .group_by("category")
    .agg(pl.col("amount").sum())
    .collect(engine="streaming")  # 流式执行，分块处理不爆内存
)

# pandas → Polars
import pandas as pd
pandas_df = pd.DataFrame({"a": [1, 2, 3]})
polars_df = pl.from_pandas(pandas_df)  # pandas转Polars

# Polars → pandas
pandas_df = polars_df.to_pandas()  # Polars转pandas

# ===== 读写 =====
pl.read_csv("f.csv")           # 读CSV
pl.scan_csv("f.csv")           # 懒读CSV
pl.read_parquet("f.parquet")   # 读Parquet
df.write_csv("out.csv")        # 写CSV
df.write_parquet("out.parquet") # 写Parquet

# ===== 基本操作 =====
df.select("col1", "col2")     # 选列
df.filter(pl.col("x") > 10)   # 过滤
df.with_columns(...)           # 添加/修改列
df.sort("col", descending=True) # 排序
df.group_by("col").agg(...)    # 分组聚合
df.join(df2, on="key")         # 连接
df.unique(subset=["col"])      # 去重
df.head(10)                    # 前10行
df.sample(100)                 # 随机抽样100行

# ===== 常用表达式 =====
pl.col("name")                # 引用列
pl.lit(42)                     # 字面量
pl.when(...).then(...).otherwise(...)  # 条件
pl.col("x").alias("new_name")  # 重命名
pl.col("x").cast(pl.Float64)   # 类型转换
pl.col("x").is_null()          # 判断空值
pl.col("x").fill_null(0)       # 填充空值