示例合集

以下示例覆盖 Python 查询脚本的常见写法。所有脚本都以定义 main() 为入口，main() 的返回值即查询结果。

返回字面数据

def main():
    return [{"city": "北京", "value": 1}, {"city": "上海", "value": 2}]

使用参数

def main():
    return [{"result": args["k"]}]

从数据源查询

def main():
    df = query("SELECT id, name, amount FROM orders")
    return df

带占位符参数：

def main():
    return query("SELECT * FROM orders WHERE amount > ?", 100)

构造 DataFrame / Series

def main():
    return DataFrame([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}])

def main():
    return Series("value", [10, 20, 30])

过滤与派生列

def main():
    df = query("SELECT category, amount FROM sales")
    return (
        df.filter(pl.col("amount").is_not_null())
          .with_columns((pl.col("amount") * 1.1).round(2).alias("amount_with_tax"))
    )

分组聚合

def main():
    df = query("SELECT category, region, amount FROM sales")
    return df.group_by("category", "region").agg(
        pl.col("amount").sum().alias("total"),
        pl.col("amount").mean().alias("avg"),
        pl.col("amount").count().alias("count"),
    )

处理时间列

SQL 返回的时间列是字符串，需先转换：

def main():
    df = query("SELECT created_at, amount FROM orders")
    return (
        df.with_columns(pl.col("created_at").str.to_datetime())
          .group_by(pl.col("created_at").dt.truncate("1mo").alias("month"))
          .agg(pl.col("amount").sum().alias("monthly_total"))
    )

HTTP 请求（GET）

def main():
    res = fetch("https://api.example.com/items")
    if not res.ok:
        return [{"error": res.status}]
    return res.json()

HTTP 请求（POST）

def main():
    res = fetch(
        "https://api.example.com/query",
        method="POST",
        body={"page": 1, "size": 100},
        headers={"Authorization": "Bearer TOKEN"},
    )
    return res.json()

结合数据源与外部 API

def main():
    # 从数据源取出需要补全的 ID
    df = query("SELECT id FROM products WHERE detail IS NULL")
    rows = []
    for row in df.to_dicts():
        res = fetch(f"https://api.example.com/products/{row['id']}")
        if res.ok:
            rows.append(res.json())
    return rows

使用标准库

标准库在「逐行处理 / 构造请求 / 解析响应」等场景很有用。下例用 hashlib 对每个邮箱做脱敏，再返回处理后的行：

import hashlib

def main():
    df = query("SELECT id, email FROM users")
    rows = []
    for row in df.to_dicts():
        digest = hashlib.sha256(row["email"].encode()).hexdigest()
        rows.append({"id": row["id"], "email_hash": digest})
    return rows

备注

冻结的纯 Python 标准库（如 json、re、argparse）以及 math、datetime、struct、hashlib 等原生模块均可 import。

滚动窗口

def main():
    df = DataFrame([
        {"day": 1, "value": 10}, {"day": 2, "value": 20},
        {"day": 3, "value": 15}, {"day": 4, "value": 30},
        {"day": 5, "value": 25},
    ])
    return df.with_columns(
        pl.col("value").rolling_mean(3).alias("ma3"),
        pl.col("value").rolling_std(3, min_samples=2).alias("std3"),
    )

指数加权移动

def main():
    df = query("SELECT date, price FROM stock_prices ORDER BY date")
    return df.with_columns(
        pl.col("price").ewm_mean(span=7).alias("ewm_7"),
    )

条件表达式 when/then/otherwise

def main():
    df = query("SELECT product, amount FROM orders")
    return df.with_columns(
        pl.when(pl.col("amount") >= 1000).then(pl.lit("大额"))
        .when(pl.col("amount") >= 500).then(pl.lit("中额"))
        .otherwise(pl.lit("小额"))
        .alias("order_level")
    )

时间窗口分组 group_by_dynamic

def main():
    df = query("SELECT created_at, amount FROM orders")
    df = df.with_columns(pl.col("created_at").str.to_datetime())
    return df.group_by_dynamic("created_at", every="7d").agg(
        pl.col("amount").sum().alias("weekly_sum"),
        pl.col("amount").count().alias("order_count"),
    )

Join

def main():
    users = query("SELECT id, name FROM users")
    orders = query("SELECT user_id, amount FROM orders")
    return orders.join(users, left_on="user_id", right_on="id", how="left")

窗口函数 over

def main():
    df = query("SELECT category, product, amount FROM orders")
    return df.with_columns(
        (pl.col("amount") / pl.col("amount").sum().over("category") * 100)
        .alias("pct_of_category")
    )

累计求和与差值

def main():
    df = query("SELECT date, amount FROM daily_sales ORDER BY date")
    return df.with_columns(
        pl.col("amount").cum_sum().alias("cumulative"),
        pl.col("amount").diff(1).alias("day_change"),
    )

排序与选取

def main():
    df = query("SELECT * FROM products")
    return df.sort("price", descending=True).head(10)

值频次统计

def main():
    df = query("SELECT category FROM products")
    return df.get_column("category").value_counts(sort=True)

元素级映射

def main():
    df = query("SELECT name, score FROM students")
    return df.with_columns(
        pl.col("score")
        .map_elements(lambda x: "优秀" if x >= 90 else ("良好" if x >= 75 else "及格"))
        .alias("grade")
    )

去重与位移

def main():
    df = query("SELECT date, value FROM metrics ORDER BY date")
    return df.with_columns(
        pl.col("value").shift(1).alias("prev_value"),
        pl.col("value").forward_fill().alias("filled"),
    )