Prompting for pandas: Getting Correct groupby Logic Instead of Plausible-Looking Code
Generated pandas aggregations often run clean and return the wrong number; here is how to prompt and verify so the logic is actually right.
The dangerous thing about asking a model for a pandas aggregation is that the code almost always runs. It imports cleanly, it produces a DataFrame with sensible column names, and the numbers look like numbers. Then three weeks later someone notices the monthly revenue total is off by four percent, and it turns out the groupby silently dropped rows with a null category, or averaged an already-averaged rate, or double-counted after a one-to-many join.
Correct-looking pandas and correct pandas are different things, and the gap is exactly where groupby lives. If you are going to generate aggregation code with an LLM, the skill is not writing a clever prompt once. It is structuring the request so the model cannot skip the decisions that actually determine correctness, and then verifying the output against something you trust.
Where generated groupby code goes wrong
Almost every silent error I see falls into one of a few buckets, and they are worth naming because they tell you what your prompt has to pin down.
- Nulls in the grouping key. By default
groupbydrops rows where the key isNaN. If ten percent of your orders have a nullregion, the regional breakdown quietly omits them and the parts no longer sum to the whole. - Categorical explosion or collapse. Grouping on a categorical dtype with
observed=Falseproduces a row for every category combination that could exist, including ones with zero data. Get this wrong and you either invent empty groups or, in a cross-tab, generate a cartesian product that blows up memory. - Averaging an average. Taking
mean()of a per-row rate or percentage gives you an unweighted average that is almost never what the business means. A conversion rate across regions has to be summed-then-divided, not meaned. - Double counting after a merge. If the DataFrame is the result of a join that fanned out one order into many line items, summing
order_totalcounts each order once per line. The aggregation is correct; the grain is wrong.
None of these throw an error. That is the whole problem.
Prompt at the grain, not the syntax
The most useful shift is to stop asking for code and start describing the shape of the answer. The model is good at pandas syntax. What it cannot infer is the grain of your data or the grain you want out. So state both.
A weak prompt: "Write pandas to get total revenue by region and month." A prompt that gets correct code:
DataFrame `df`, one row per ORDER LINE ITEM. Columns and dtypes:
order_id: str (an order has 1..n line items)
region: category (may be null for ~8% of rows)
order_ts: datetime64
line_revenue: float (revenue for THIS line item)
I want: total revenue by region and calendar month.
Grain of output: one row per (region, month).
Requirements:
- Include rows where region is null; label them "Unknown".
- line_revenue is already per-line, so summing is correct.
- Do not create empty (region, month) combinations.
- Return region and month as columns, not index.
Before the code, state your assumptions about grain and null handling in two bullets.Three things do the work here. Giving dtypes and the row grain tells the model whether it is looking at orders or line items, which is the single most common source of double counting. Naming the output grain forces a decision about what one output row means. And asking for assumptions before the code turns the model's implicit choices into text you can check in five seconds, before you ever run anything.
Demand the unglamorous flags
Certain pandas defaults are wrong often enough that I ask for them explicitly every time. Put these directly in the prompt as requirements so the model does not fall back to defaults:
dropna=Falseongroupbywhen you want null keys represented, paired with afillnato a real label so the null group is visible rather than a silentNaNrow.observed=Truewhenever a grouping key is categorical, so you get groups that exist in the data rather than the full categorical product.as_index=Falseor an explicitreset_index()so the result is a flat DataFrame, which is what almost every downstream step actually wants.- Named aggregation via
.agg(total=('line_revenue','sum'))rather than positional aggregation, so the output columns are self-documenting and you cannot mix up which column got which function.
For weighted metrics, do not accept a mean(). Ask the model to aggregate the numerator and denominator separately and divide after grouping. "Conversion rate by channel" should sum conversions and sum sessions per channel, then divide, and the prompt should say so.
Verification is not optional
Even a well-specified prompt can produce wrong code, so the last step is reconciliation. Two cheap checks catch the overwhelming majority of errors.
Reconcile the total. The sum of your aggregated metric should equal the sum of the source column, within floating-point tolerance. If result['total'].sum() does not match df['line_revenue'].sum(), you have either dropped rows (a null key) or double counted (a grain problem). This one line is the highest-value assertion you can add.
assert abs(result['total'].sum() - df['line_revenue'].sum()) < 0.01, \
"aggregate does not reconcile to source total"Check the row count of the output. If you expect one row per region-month and you get more, a categorical exploded or a key was not what you thought. If you get fewer, groups were dropped. I often ask the model to print the output row count and the number of distinct group keys as a built-in tripwire.
For anything that will feed a report, I keep a tiny reference implementation, sometimes in SQL against the same warehouse table, and compare totals. Two independent computations agreeing is far stronger evidence than one that looks right.
The honest trade-off
Generating aggregation code with a model saves real time on the syntax and the boilerplate, especially for multi-key groupings with named aggregations you would otherwise have to look up. It does not save you the analytical thinking, and pretending it does is how the four-percent error ships. The model does not know your data's grain, does not know which columns are already rates, and does not know that your region column has nulls that matter. You do, or you can find out in one value_counts().
So the workflow that actually holds up is narrow. Describe the data and the desired grain precisely, force the model to declare its assumptions, require the flags that override bad defaults, and reconcile the total before you trust a single number. The model writes the pandas. You still own the logic. Treat a generated aggregation exactly as you would a pull request from a fast but junior colleague: read the assumptions, run the checks, and never merge on the fact that it ran.
A note on shelf life. AI products change fast. This guide deliberately focuses on the parts that stay true — how to judge a tool, what the trade-offs are — rather than ranking products that will have changed by the time you read it. Prices and feature claims should always be checked against the provider before you rely on them.