Skip to main content

Examples

The following examples cover common patterns in Python query scripts. All scripts define main() as the entry point; its return value is the query result.

Returning Literal Data​

def main():
    return [{"city": "Beijing", "value": 1}, {"city": "Shanghai", "value": 2}]

Using Parameters​

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

Querying from Data Sources​

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

With placeholder parameters:

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

Constructing DataFrame / Series​

def main():
    return DataFrame([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}])
def main():
    return Series("value", [10, 20, 30])

Filter & Derived Columns​

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"))
    )

Group-by Aggregation​

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"),
    )

Handling Timestamp Columns​

SQL timestamp columns are returned as strings — convert them first:

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 Request (GET)​

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

HTTP Request (POST)​

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

Combining Data Sources with External APIs​

def main():
    # Fetch IDs that need enrichment from the data source
    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

Using the Standard Library​

The standard library is useful for row-by-row processing, building requests, and parsing responses. This example uses hashlib to anonymize email addresses:

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
note

Frozen pure-Python standard libraries (such as json, re, argparse) and native modules like math, datetime, struct, hashlib can all be imported.

Rolling Windows​

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"),
    )

Exponentially Weighted Moving​

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"),
    )

Conditional Expressions 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("Large"))
        .when(pl.col("amount") >= 500).then(pl.lit("Medium"))
        .otherwise(pl.lit("Small"))
        .alias("order_level")
    )

Time-window Grouping 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")

Window Function 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")
    )

Cumulative Sum & Difference​

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"),
    )

Sort & Select​

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

Value Counts​

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

Element-wise Mapping​

def main():
    df = query("SELECT name, score FROM students")
    return df.with_columns(
        pl.col("score")
        .map_elements(lambda x: "Excellent" if x >= 90 else ("Good" if x >= 75 else "Pass"))
        .alias("grade")
    )

Dedup & Shift​

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"),
    )