File size: 1,597 Bytes
de46078 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | id: sql_having_to_dataframe
category: sql_translation
difficulty: medium
probes: >
HAVING filters after aggregation, WHERE filters before. The query has both,
with different predicates, so swapping them changes the answer. Tests whether
the model understands the clause order rather than pattern-matching filter().
tags: [sql, having, where, groupby]
prompt: |
Translate this SQL into the PySpark DataFrame API against `txns`
(do not use spark.sql or createOrReplaceTempView):
SELECT store, SUM(amount) AS total
FROM txns
WHERE status = 'ok'
GROUP BY store
HAVING COUNT(*) >= 2
ORDER BY total DESC
Return columns: store, total.
fixtures:
- name: txns
schema: txn_id INT, store STRING, status STRING, amount DOUBLE
rows:
- [1, "s1", "ok", 100.0]
- [2, "s1", "ok", 50.0]
- [3, "s1", "void", 999.0]
- [4, "s2", "ok", 400.0]
- [5, "s2", "void", 1.0]
- [6, "s3", "ok", 10.0]
- [7, "s3", "ok", 20.0]
- [8, "s3", "ok", 30.0]
solution: |
from pyspark.sql import functions as F
def solve(spark, txns):
# WHERE before groupBy, HAVING as a filter on the aggregated frame.
# s2 has only one 'ok' row, so HAVING COUNT(*) >= 2 removes it.
return (
txns
.filter(F.col("status") == "ok")
.groupBy("store")
.agg(F.sum("amount").alias("total"), F.count(F.lit(1)).alias("_n"))
.filter(F.col("_n") >= 2)
.drop("_n")
.orderBy(F.col("total").desc())
)
compare:
mode: ordered_rows
float_tolerance: 1.0e-6
|