costflow 0.1.2

Cost estimation & scaling analysis for pandas pipelines

CostFlow

CostFlow measures how your pandas pipelines scale. It runs a pipeline across multiple data sizes, traces expensive DataFrame ops (merge, groupby, sort, pivot, etc.), and fits simple complexity models so you can spot bottlenecks before the code hits production.

Why use it

• See how wall time changes as rows grow and get quick projections for larger sizes.
• Identify the pandas operations consuming most of the traced time.
• Estimate memory needs (when psutil is installed) to avoid surprises on shared machines.

Install

pip install .
# for memory tracking
pip install ".[mem]"

Quickstart

import pandas as pd
from costflow import analyze
from costflow.report import pretty_print

def make_df(n: int) -> pd.DataFrame:
    import numpy as np
    rng = np.random.default_rng(0)
    return pd.DataFrame(
        {
            "customer_id": rng.integers(0, 100, size=n),
            "region": rng.choice(["NA", "EU", "APAC", "LATAM"], size=n),
            "quantity": rng.integers(1, 10, size=n),
            "price": rng.normal(50, 10, size=n).clip(5),
            "discount": rng.uniform(0, 0.3, size=n),
        }
    )

def pipeline(df):
    df["revenue"] = df["quantity"] * df["price"] * (1 - df["discount"])
    out = (
        df.groupby(["customer_id", "region"])
          .agg(total_revenue=("revenue", "sum"))
          .reset_index()
          .sort_values("total_revenue", ascending=False)
    )
    return out

report = analyze(
    pipeline_fn=pipeline,
    make_df=make_df,
    sizes=[1_000, 5_000, 10_000],
    trace_ops=True,
)
pretty_print(report)

# If you already have real data, skip writing make_df and use analyze_with_df:
base_df = pd.read_parquet("your_data.parquet")  # or read_csv/DB pull

sizes = [
    len(base_df) // 2,              # smaller downsample
    len(base_df),                   # full dataset
    min(len(base_df) * 2, 200_000)  # scaled up (capped)
]

report = analyze_with_df(
    pipeline_fn=pipeline,
    base_df=base_df,
    sizes=sizes,
    trace_ops=True,
    warmup=False,
)
pretty_print(report)

# Multiple DataFrames: use analyze_with_dfs for convenience
def pipeline_multi(df_left, df_right):
    return (
        df_left.merge(df_right, on="key")
               .groupby("key")
               .agg(total=("value", "sum"))
               .reset_index()
    )

report = analyze_with_dfs(
    pipeline_fn=pipeline_multi,
    base_dfs=[base_df, other_df],
    sizes=[10_000, 50_000, 100_000],
    trace_ops=True,
)
pretty_print(report)

### Plotting
```python
from costflow.report import plot_report

plot_report(report)

See `examples/quickstart.py` for a fuller walk-through with multiple pipelines.
More scenarios:
- `examples/wide_table.py` for wide numeric tables with many columns.
- `examples/string_heavy.py` for string-heavy joins/pivots.

### Interpreting the report
- **Runs**: shows wall time and (when available) peak RSS for each dataset size.
- **Time model**: picks the candidate scaling model with the lowest RMSE; check `r2` to see how well it fits.
- **Memory model**: only populated when `psutil` is installed (`pip install ".[mem]"`). Without it, the memory fit is `N/A` and projections are omitted.
- **Dominant ops**: fraction of traced time spent in the most expensive pandas operations. Tracing is best-effort: common methods are wrapped explicitly, and other DataFrame/GroupBy methods are generically wrapped, but deep pandas internals (C/numba/numexpr) are not captured.
- **Projections**: extrapolations to common target sizes using the chosen model and the column count from your largest run. Treat these as coarse estimates—validate them with a real run when possible and prefer models with a reasonable `r2`.

### Tips for better signal
- Use at least two or three increasing sizes so the fit can distinguish between linear vs. super-linear growth.
- Set `trace_ops=True` to find bottlenecks, then re-run with `trace_ops=False` if you want to remove tracing overhead from timings.
- Keep pipelines functional: return a DataFrame/Series; avoid mutating global state.

### CLI
Analyze a pipeline from the command line (defaults: `pipeline`/`make_df` function names):
```bash
costflow --pipeline examples/quickstart.py:simple_pipeline --make-df examples/quickstart.py:make_df --sizes 1000 5000 10000

• Add --json to emit machine-readable output.
• Use --no-trace-ops to remove tracing overhead when you just want wall-clock measurements.

Limitations and guidance

• Tracing is shallow: pandas operations implemented in C/numexpr/numba are not visible, so dominant-op fractions reflect Python-level pandas calls.
• Tracing wraps most DataFrame and Series methods; common top-level pandas functions (pd.concat, pd.merge) are also traced. Other top-level functions still won't appear unless called as methods.
• Projections assume the chosen model holds; verify with a real run near your target size before making infra decisions.
• Memory projections require psutil. Without it, memory fit is N/A and only wall-time projections are shown.

Development

pip install ".[dev]"
pytest

License

MIT (see LICENSE)