tpch examples: rewrite queries idiomatically and embed reference SQL (#1504)

* tpch examples: add reference SQL to each query, fix Q20

- Append the canonical TPC-H reference SQL (from benchmarks/tpch/queries/)
  to each q01..q22 module docstring so readers can compare the DataFrame
  translation against the SQL at a glance.
- Fix Q20: `df = df.filter(col("ps_availqty") > lit(0.5) * col("total_sold"))`
  was missing the assignment so the filter was dropped from the pipeline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* tpch examples: rewrite non-idiomatic queries in idiomatic DataFrame form

Rewrite the seven TPC-H example queries that did not demonstrate the
idiomatic DataFrame pattern. The remaining queries (Q02/Q11/Q15/Q17/Q22,
which use window functions in place of correlated subqueries) already are
idiomatic and are left unchanged.

- Q04: replace `.aggregate([col("l_orderkey")], [])` with
  `.select("l_orderkey").distinct()`, which is the natural way to express
  "reduce to one row per order" on a DataFrame.
- Q07: remove the CASE-as-filter on `n_name` and use
  `F.in_list(col("n_name"), [nation_1, nation_2])` instead. Drops a
  comment block that admitted the filter form was simpler.
- Q08: rewrite the switched CASE `F.case(...).when(lit(False), ...)` as a
  searched `F.when(col(...).is_not_null(), ...).otherwise(...)`. That
  mirrors the reference SQL's `case when ... then ... else 0 end` shape.
- Q12: replace `array_position(make_array(...), col)` with
  `F.in_list(col("l_shipmode"), [...])`. Same semantics, without routing
  through array construction / array search.
- Q19: remove the pyarrow UDF that re-implemented a disjunctive predicate
  in Python. Build the same predicate in DataFusion by OR-combining one
  `in_list` + range-filter expression per brand. Keeps the per-brand
  constants in the existing `items_of_interest` dict.
- Q20: use `F.starts_with` instead of an explicit substring slice. Replace
  the inner-join + `select(...).distinct()` tail with a semi join against
  a precomputed set of excess-quantity suppliers so the supplier columns
  are preserved without deduplication after the fact.
- Q21: replace the `array_agg` / `array_length` / `array_element` pipeline
  with two semi joins. One semi join keeps orders with more than one
  distinct supplier (stand-in for the reference SQL's `exists` subquery),
  the other keeps orders with exactly one late supplier (stand-in for the
  `not exists` subquery).

All 22 answer-file comparisons and 22 plan-comparison diagnostics still
pass (`pytest examples/tpch/_tests.py`: 44 passed).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* tpch examples: align reference SQL constants with DataFrame queries

The reference SQL embedded in each q01..q22 module docstring was carried
over verbatim from ``benchmarks/tpch/queries/`` and uses a different set
of TPC-H substitution parameters than the DataFrame examples
(answer-file-validated at scale factor 1). Update each reference SQL to
use the substitution parameters the DataFrame uses, so both expressions
describe the same query and would produce the same results against the
same data.

Constants aligned:

- Q01: ``90 days`` cutoff (DataFrame ``DAYS_BEFORE_FINAL = 90``).
- Q02: ``p_size = 15``, ``p_type like '%BRASS'``, ``r_name = 'EUROPE'``.
- Q04: base date ``1993-07-01`` (``3 month`` interval preserved per the
  "quarter of a year" wording).
- Q05: ``r_name = 'ASIA'``.
- Q06: ``l_discount between 0.06 - 0.01 and 0.06 + 0.01``.
- Q07: nations ``'FRANCE'`` / ``'GERMANY'``.
- Q08: ``r_name = 'AMERICA'``, ``p_type = 'ECONOMY ANODIZED STEEL'``,
  inner-case ``nation = 'BRAZIL'``.
- Q09: ``p_name like '%green%'``.
- Q10: base date ``1993-10-01`` (``3 month`` interval preserved).
- Q11: ``n_name = 'GERMANY'``.
- Q12: ship modes ``('MAIL', 'SHIP')``, base date ``1994-01-01``.
- Q13: ``o_comment not like '%special%requests%'``.
- Q14: base date ``1995-09-01``.
- Q15: base date ``1996-01-01``.
- Q16: ``p_brand <> 'Brand#45'``, ``p_type not like 'MEDIUM POLISHED%'``,
  sizes ``(49, 14, 23, 45, 19, 3, 36, 9)``.
- Q17: ``p_brand = 'Brand#23'``, ``p_container = 'MED BOX'``.
- Q18: ``sum(l_quantity) > 300``.
- Q19: brands ``Brand#12`` / ``Brand#23`` / ``Brand#34`` with the matching
  minimum quantities (1, 10, 20).
- Q20: ``p_name like 'forest%'``, base date ``1994-01-01``,
  ``n_name = 'CANADA'``.
- Q21: ``n_name = 'SAUDI ARABIA'``.
- Q22: country codes ``('13', '31', '23', '29', '30', '18', '17')``.

Interval units (month / year) are preserved where the problem-statement
text reads "given quarter", "given year", "given month". Q01 keeps the
literal "days" unit because the TPC-H problem statement itself describes
the cutoff in days.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* tpch examples: apply SKILL.md idioms across all 22 queries

Sweep every q01..q22 example for idiomatic DataFrame style as described in
the repo-root SKILL.md:

- ``col("x") == "s"`` in place of ``col("x") == lit("s")`` on comparison
  right-hand sides (auto-wrap applies).
- Plain-name strings in ``select``/``aggregate``/``sort`` group/sort key
  lists when the key is a bare column.
- Drop redundant ``how="inner"`` and single-element ``left_on``/``right_on``
  list wrapping on equi-joins.
- Collapse chained ``.filter(a).filter(b)`` runs into ``.filter(a, b)``
  and chained ``.with_column`` runs into ``.with_columns(a=..., b=...)``.
- ``df.sort_by(...)`` or plain-name ``df.sort(...)`` when no null-placement
  override is needed.
- ``F.count_star()`` in place of ``F.count(col("x"))`` whenever the SQL
  reads ``count(*)``.
- ``F.starts_with(col, lit(prefix))`` and ``~F.starts_with(...)`` in place
  of substring-prefix equality/inequality tricks.
- ``F.in_list(col, [lit(...)])`` in place of ``~F.array_position(...).
  is_null()`` and in place of disjunctions of equality comparisons.
- Searched ``F.when(cond, x).otherwise(y)`` in place of switched
  ``F.case(bool_expr).when(lit(True/False), x).end()`` forms.
- Semi-joins as the DataFrame form of ``EXISTS`` (Q04); anti-joins as
  ``NOT EXISTS`` (Q22 was already using this idiom).
- Whole-frame window aggregates as the DataFrame stand-in for a SQL
  scalar subquery (Q11/Q15/Q17/Q22).

Individual query fixes of note:

- Q16 — add the secondary sort keys (``p_brand``, ``p_type``, ``p_size``)
  that the TPC-H spec requires but the original DataFrame omitted.
- Q22 — drop a stray ``df.show()`` mid-pipeline; replace the 0-based
  substring slice with ``F.left(col("c_phone"), lit(2))``.
- Q14 — rewrite the promo/non-promo factor split as a searched CASE inside
  ``F.sum(...)`` so the DataFrame expression matches the reference SQL
  shape exactly.

All 22 answer-file comparisons still pass at scale factor 1.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* tpch examples: more idiomatic aggregate FILTER, string funcs, date handling

Additional sweep of the TPC-H DataFrame examples informed by comparing
against a fresh set of SKILL.md-only generations under
``examples/tpch/agentic_queries/``:

- Q02: ``F.ends_with(col("p_type"), lit(TYPE_OF_INTEREST))`` in place of
  ``F.strpos(col, lit) > 0``. The reference SQL is ``p_type like '%BRASS'``,
  which is an ends_with check, not contains. ``F.strpos > 0`` returned the
  correct rows on TPC-H data by coincidence but is semantically wrong.
- Q09: ``F.contains(col("p_name"), lit(part_color))`` in place of
  ``F.strpos(col, lit) > 0``. The SQL is ``p_name like '%green%'``.
- Q08, Q12, Q14: use the ``filter`` keyword on ``F.sum`` / ``F.count`` —
  the DataFrame form of SQL ``sum(...) FILTER (WHERE ...)`` — instead of
  wrapping the aggregate input in ``F.when(cond, x).otherwise(0)``. Q08
  also reorganises to inner-join the supplier's nation onto the regional
  sales, which removes the previous left-join + ``F.when(is_not_null, ...)``
  dance.
- Q15: compute the grand maximum revenue as a separate scalar aggregate
  and ``join_on(...)`` on equality, instead of the whole-frame window
  ``F.max`` + filter shape. Simpler plan, same result.
- Q16: ``F.regexp_like(col, pattern)`` in place of
  ``F.regexp_match(col, pattern).is_not_null()``.
- Q04, Q05, Q06, Q07, Q08, Q10, Q12, Q14, Q15, Q20: store both the start
  and the end of the date window as plain ``datetime.date`` objects and
  compare with ``lit(end_date)``, instead of carrying the start date +
  ``pa.month_day_nano_interval`` and adding them at query-build time.
  Drops unused ``pyarrow`` imports from the files that no longer need
  Arrow scalars.

All 22 answer-file comparisons still pass at scale factor 1.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

examples/tpch/q01_pricing_summary_report.py

@@ -27,6 +27,30 @@
 
 The above problem statement text is copyrighted by the Transaction Processing Performance Council
 as part of their TPC Benchmark H Specification revision 2.18.0.
+
+Reference SQL (from TPC-H specification, used by the benchmark suite)::
+
+    select
+        l_returnflag,
+        l_linestatus,
+        sum(l_quantity) as sum_qty,
+        sum(l_extendedprice) as sum_base_price,
+        sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+        sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+        avg(l_quantity) as avg_qty,
+        avg(l_extendedprice) as avg_price,
+        avg(l_discount) as avg_disc,
+        count(*) as count_order
+    from
+        lineitem
+    where
+        l_shipdate <= date '1998-12-01' - interval '90 days'
+    group by
+        l_returnflag,
+        l_linestatus
+    order by
+        l_returnflag,
+        l_linestatus;
 """
 
 import pyarrow as pa
@@ -58,31 +82,25 @@
 
 # Aggregate the results
 
+disc_price = col("l_extendedprice") * (lit(1) - col("l_discount"))
+
 df = df.aggregate(
-    [col("l_returnflag"), col("l_linestatus")],
+    ["l_returnflag", "l_linestatus"],
     [
         F.sum(col("l_quantity")).alias("sum_qty"),
         F.sum(col("l_extendedprice")).alias("sum_base_price"),
-        F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias(
-            "sum_disc_price"
-        ),
-        F.sum(
-            col("l_extendedprice")
-            * (lit(1) - col("l_discount"))
-            * (lit(1) + col("l_tax"))
-        ).alias("sum_charge"),
+        F.sum(disc_price).alias("sum_disc_price"),
+        F.sum(disc_price * (lit(1) + col("l_tax"))).alias("sum_charge"),
         F.avg(col("l_quantity")).alias("avg_qty"),
         F.avg(col("l_extendedprice")).alias("avg_price"),
         F.avg(col("l_discount")).alias("avg_disc"),
-        F.count(col("l_returnflag")).alias(
-            "count_order"
-        ),  # Counting any column should return same result
+        F.count_star().alias("count_order"),
     ],
 )
 
 # Sort per the expected result
 
-df = df.sort(col("l_returnflag").sort(), col("l_linestatus").sort())
+df = df.sort_by("l_returnflag", "l_linestatus")
 
 # Note: There appears to be a discrepancy between what is returned here and what is in the generated
 # answers file for the case of return flag N and line status O, but I did not investigate further.

examples/tpch/q02_minimum_cost_supplier.py

@@ -27,6 +27,52 @@
 
 The above problem statement text is copyrighted by the Transaction Processing Performance Council
 as part of their TPC Benchmark H Specification revision 2.18.0.
+
+Reference SQL (from TPC-H specification, used by the benchmark suite)::
+
+    select
+        s_acctbal,
+        s_name,
+        n_name,
+        p_partkey,
+        p_mfgr,
+        s_address,
+        s_phone,
+        s_comment
+    from
+        part,
+        supplier,
+        partsupp,
+        nation,
+        region
+    where
+        p_partkey = ps_partkey
+        and s_suppkey = ps_suppkey
+        and p_size = 15
+        and p_type like '%BRASS'
+        and s_nationkey = n_nationkey
+        and n_regionkey = r_regionkey
+        and r_name = 'EUROPE'
+        and ps_supplycost = (
+                select
+                        min(ps_supplycost)
+                from
+                        partsupp,
+                        supplier,
+                        nation,
+                        region
+                where
+                        p_partkey = ps_partkey
+                        and s_suppkey = ps_suppkey
+                        and s_nationkey = n_nationkey
+                        and n_regionkey = r_regionkey
+                        and r_name = 'EUROPE'
+        )
+    order by
+        s_acctbal desc,
+        n_name,
+        s_name,
+        p_partkey limit 100;
 """
 
 import datafusion
@@ -67,35 +113,30 @@
     "r_regionkey", "r_name"
 )
 
-# Filter down parts. Part names contain the type of interest, so we can use strpos to find where
-# in the p_type column the word is. `strpos` will return 0 if not found, otherwise the position
-# in the string where it is located.
+# Filter down parts. The reference SQL uses ``p_type like '%BRASS'`` which
+# is an ``ends_with`` check; use the dedicated string function rather than
+# a manual substring match.
 
 df_part = df_part.filter(
-    F.strpos(col("p_type"), lit(TYPE_OF_INTEREST)) > lit(0)
-).filter(col("p_size") == lit(SIZE_OF_INTEREST))
+    F.ends_with(col("p_type"), lit(TYPE_OF_INTEREST)),
+    col("p_size") == SIZE_OF_INTEREST,
+)
 
 # Filter regions down to the one of interest
 
-df_region = df_region.filter(col("r_name") == lit(REGION_OF_INTEREST))
+df_region = df_region.filter(col("r_name") == REGION_OF_INTEREST)
 
 # Now that we have the region, find suppliers in that region. Suppliers are tied to their nation
 # and nations are tied to the region.
 
-df_nation = df_nation.join(
-    df_region, left_on=["n_regionkey"], right_on=["r_regionkey"], how="inner"
-)
-df_supplier = df_supplier.join(
-    df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner"
-)
+df_nation = df_nation.join(df_region, left_on="n_regionkey", right_on="r_regionkey")
+df_supplier = df_supplier.join(df_nation, left_on="s_nationkey", right_on="n_nationkey")
 
 # Now that we know who the potential suppliers are for the part, we can limit out part
 # supplies table down. We can further join down to the specific parts we've identified
 # as matching the request
 
-df = df_partsupp.join(
-    df_supplier, left_on=["ps_suppkey"], right_on=["s_suppkey"], how="inner"
-)
+df = df_partsupp.join(df_supplier, left_on="ps_suppkey", right_on="s_suppkey")
 
 # Locate the minimum cost across all suppliers. There are multiple ways you could do this,
 # but one way is to create a window function across all suppliers, find the minimum, and
@@ -112,9 +153,9 @@
     ),
 )
 
-df = df.filter(col("min_cost") == col("ps_supplycost"))
-
-df = df.join(df_part, left_on=["ps_partkey"], right_on=["p_partkey"], how="inner")
+df = df.filter(col("min_cost") == col("ps_supplycost")).join(
+    df_part, left_on="ps_partkey", right_on="p_partkey"
+)
 
 # From the problem statement, these are the values we wish to output
 
@@ -132,12 +173,10 @@
 # Sort and display 100 entries
 df = df.sort(
     col("s_acctbal").sort(ascending=False),
-    col("n_name").sort(),
-    col("s_name").sort(),
-    col("p_partkey").sort(),
-)
-
-df = df.limit(100)
+    "n_name",
+    "s_name",
+    "p_partkey",
+).limit(100)
 
 # Show results
 

examples/tpch/q03_shipping_priority.py

@@ -25,6 +25,31 @@
 
 The above problem statement text is copyrighted by the Transaction Processing Performance Council
 as part of their TPC Benchmark H Specification revision 2.18.0.
+
+Reference SQL (from TPC-H specification, used by the benchmark suite)::
+
+    select
+        l_orderkey,
+        sum(l_extendedprice * (1 - l_discount)) as revenue,
+        o_orderdate,
+        o_shippriority
+    from
+        customer,
+        orders,
+        lineitem
+    where
+        c_mktsegment = 'BUILDING'
+        and c_custkey = o_custkey
+        and l_orderkey = o_orderkey
+        and o_orderdate < date '1995-03-15'
+        and l_shipdate > date '1995-03-15'
+    group by
+        l_orderkey,
+        o_orderdate,
+        o_shippriority
+    order by
+        revenue desc,
+        o_orderdate limit 10;
 """
 
 from datafusion import SessionContext, col, lit
@@ -50,38 +75,34 @@
 
 # Limit dataframes to the rows of interest
 
-df_customer = df_customer.filter(col("c_mktsegment") == lit(SEGMENT_OF_INTEREST))
+df_customer = df_customer.filter(col("c_mktsegment") == SEGMENT_OF_INTEREST)
 df_orders = df_orders.filter(col("o_orderdate") < lit(DATE_OF_INTEREST))
 df_lineitem = df_lineitem.filter(col("l_shipdate") > lit(DATE_OF_INTEREST))
 
 # Join all 3 dataframes
 
-df = df_customer.join(
-    df_orders, left_on=["c_custkey"], right_on=["o_custkey"], how="inner"
-).join(df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner")
+df = df_customer.join(df_orders, left_on="c_custkey", right_on="o_custkey").join(
+    df_lineitem, left_on="o_orderkey", right_on="l_orderkey"
+)
 
 # Compute the revenue
 
 df = df.aggregate(
-    [col("l_orderkey")],
+    ["l_orderkey"],
     [
         F.first_value(col("o_orderdate")).alias("o_orderdate"),
         F.first_value(col("o_shippriority")).alias("o_shippriority"),
         F.sum(col("l_extendedprice") * (lit(1.0) - col("l_discount"))).alias("revenue"),
     ],
 )
 
-# Sort by priority
-
-df = df.sort(col("revenue").sort(ascending=False), col("o_orderdate").sort())
-
-# Only return 10 results
+# Sort by priority, take 10, and project in the order expected by the spec.
 
-df = df.limit(10)
-
-# Change the order that the columns are reported in just to match the spec
-
-df = df.select("l_orderkey", "revenue", "o_orderdate", "o_shippriority")
+df = (
+    df.sort(col("revenue").sort(ascending=False), "o_orderdate")
+    .limit(10)
+    .select("l_orderkey", "revenue", "o_orderdate", "o_shippriority")
+)
 
 # Show result
 

examples/tpch/q04_order_priority_checking.py

@@ -24,18 +24,40 @@
 
 The above problem statement text is copyrighted by the Transaction Processing Performance Council
 as part of their TPC Benchmark H Specification revision 2.18.0.
+
+Reference SQL (from TPC-H specification, used by the benchmark suite)::
+
+    select
+        o_orderpriority,
+        count(*) as order_count
+    from
+        orders
+    where
+        o_orderdate >= date '1993-07-01'
+        and o_orderdate < date '1993-07-01' + interval '3' month
+        and exists (
+                select
+                        *
+                from
+                        lineitem
+                where
+                        l_orderkey = o_orderkey
+                        and l_commitdate < l_receiptdate
+        )
+    group by
+        o_orderpriority
+    order by
+        o_orderpriority;
 """
 
-from datetime import datetime
+from datetime import date
 
-import pyarrow as pa
 from datafusion import SessionContext, col, lit
 from datafusion import functions as F
 from util import get_data_path
 
-# Ideally we could put 3 months into the interval. See note below.
-INTERVAL_DAYS = 92
-DATE_OF_INTEREST = "1993-07-01"
+QUARTER_START = date(1993, 7, 1)
+QUARTER_END = date(1993, 10, 1)
 
 # Load the dataframes we need
 
@@ -48,36 +70,23 @@
     "l_orderkey", "l_commitdate", "l_receiptdate"
 )
 
-# Create a date object from the string
-date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date()
-
-interval = pa.scalar((0, INTERVAL_DAYS, 0), type=pa.month_day_nano_interval())
-
-# Limit results to cases where commitment date before receipt date
-# Aggregate the results so we only get one row to join with the order table.
-# Alternately, and likely more idiomatic is instead of `.aggregate` you could
-# do `.select("l_orderkey").distinct()`. The goal here is to show
-# multiple examples of how to use Data Fusion.
-df_lineitem = df_lineitem.filter(col("l_commitdate") < col("l_receiptdate")).aggregate(
-    [col("l_orderkey")], []
+# Keep only orders in the quarter of interest, then restrict to those that
+# have at least one late lineitem via a semi join (the DataFrame form of
+# ``EXISTS`` from the reference SQL).
+df_orders = df_orders.filter(
+    col("o_orderdate") >= lit(QUARTER_START),
+    col("o_orderdate") < lit(QUARTER_END),
 )
 
-# Limit orders to date range of interest
-df_orders = df_orders.filter(col("o_orderdate") >= lit(date)).filter(
-    col("o_orderdate") < lit(date) + lit(interval)
-)
+late_lineitems = df_lineitem.filter(col("l_commitdate") < col("l_receiptdate"))
 
-# Perform the join to find only orders for which there are lineitems outside of expected range
 df = df_orders.join(
-    df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner"
+    late_lineitems, left_on="o_orderkey", right_on="l_orderkey", how="semi"
 )
 
-# Based on priority, find the number of entries
-df = df.aggregate(
-    [col("o_orderpriority")], [F.count(col("o_orderpriority")).alias("order_count")]
+# Count the number of orders in each priority group and sort.
+df = df.aggregate(["o_orderpriority"], [F.count_star().alias("order_count")]).sort_by(
+    "o_orderpriority"
 )
 
-# Sort the results
-df = df.sort(col("o_orderpriority").sort())
-
 df.show()