diff --git a/ibis_tpc/__init__.py b/ibis_tpc/__init__.py index 353e514..89025ac 100644 --- a/ibis_tpc/__init__.py +++ b/ibis_tpc/__init__.py @@ -1,2 +1 @@ -from .utils import * from .runners import * diff --git a/ibis_tpc/h01.py b/ibis_tpc/h01.py index f47fddc..40742c9 100644 --- a/ibis_tpc/h01.py +++ b/ibis_tpc/h01.py @@ -1,7 +1,7 @@ "Pricing Summary Report Query (Q1)" -from .utils import add_date +import ibis def tpc_h01(con, DELTA=90, DATE="1998-12-01"): @@ -18,8 +18,9 @@ def tpc_h01(con, DELTA=90, DATE="1998-12-01"): t = con.table("lineitem") - interval = add_date(DATE, dd=-1 * DELTA) - q = t.filter(t.l_shipdate <= interval) + q = t.filter( + t.l_shipdate <= ibis.date(DATE) - ibis.interval(days=DELTA) + ) discount_price = t.l_extendedprice * (1 - t.l_discount) charge = discount_price * (1 + t.l_tax) q = q.group_by(["l_returnflag", "l_linestatus"]) diff --git a/ibis_tpc/h03.py b/ibis_tpc/h03.py index 2299684..bf5599d 100644 --- a/ibis_tpc/h03.py +++ b/ibis_tpc/h03.py @@ -12,7 +12,11 @@ def tpc_h03(con, MKTSEGMENT="BUILDING", DATE="1995-03-15"): q = customer.join(orders, customer.c_custkey == orders.o_custkey) q = q.join(lineitem, lineitem.l_orderkey == orders.o_orderkey) q = q.filter( - [q.c_mktsegment == MKTSEGMENT, q.o_orderdate < DATE, q.l_shipdate > DATE] + [ + q.c_mktsegment == MKTSEGMENT, + q.o_orderdate < ibis.date(DATE), + q.l_shipdate > ibis.date(DATE) + ] ) qg = q.group_by([q.l_orderkey, q.o_orderdate, q.o_shippriority]) q = qg.aggregate(revenue=(q.l_extendedprice * (1 - q.l_discount)).sum()) diff --git a/ibis_tpc/h04.py b/ibis_tpc/h04.py index 2cfcd4f..aa1867d 100644 --- a/ibis_tpc/h04.py +++ b/ibis_tpc/h04.py @@ -1,6 +1,6 @@ "Order Priority Checking Query (Q4)" -from .utils import add_date +import ibis def tpc_h04(con, DATE="1993-07-01"): @@ -12,8 +12,9 @@ def tpc_h04(con, DATE="1993-07-01"): q = orders.filter( [ cond.any(), - orders.o_orderdate >= DATE, - orders.o_orderdate < add_date(DATE, dm=3), + orders.o_orderdate >= ibis.date(DATE), + orders.o_orderdate < ibis.date(DATE) + + ibis.interval(months=3), ] ) q = q.group_by([orders.o_orderpriority]) diff --git a/ibis_tpc/h05.py b/ibis_tpc/h05.py index 113f0e7..47f7528 100644 --- a/ibis_tpc/h05.py +++ b/ibis_tpc/h05.py @@ -2,9 +2,6 @@ import ibis -from .utils import add_date - - def tpc_h05(con, NAME="ASIA", DATE="1994-01-01"): customer = con.table("customer") orders = con.table("orders") @@ -25,7 +22,12 @@ def tpc_h05(con, NAME="ASIA", DATE="1994-01-01"): q = q.join(region, nation.n_regionkey == region.r_regionkey) q = q.filter( - [q.r_name == NAME, q.o_orderdate >= DATE, q.o_orderdate < add_date(DATE, dy=1)] + [ + q.r_name == NAME, + q.o_orderdate >= ibis.date(DATE), + q.o_orderdate < ibis.date(DATE) + + ibis.interval(years=1) + ] ) revexpr = q.l_extendedprice * (1 - q.l_discount) gq = q.group_by([q.n_name]) diff --git a/ibis_tpc/h06.py b/ibis_tpc/h06.py index 245482d..03f5f56 100644 --- a/ibis_tpc/h06.py +++ b/ibis_tpc/h06.py @@ -1,6 +1,6 @@ "Forecasting Revenue Change Query (Q6)" -from .utils import add_date +import ibis def tpc_h06(con, DATE="1994-01-01", DISCOUNT=0.06, QUANTITY=24): @@ -9,8 +9,9 @@ def tpc_h06(con, DATE="1994-01-01", DISCOUNT=0.06, QUANTITY=24): discount_max = round(DISCOUNT + 0.01, 2) q = q.filter( [ - q.l_shipdate >= DATE, - q.l_shipdate < add_date(DATE, dy=1), + q.l_shipdate >= ibis.date(DATE), + q.l_shipdate < ibis.date(DATE) + + ibis.interval(years=1), q.l_discount.between(discount_min, discount_max), q.l_quantity < QUANTITY, ] diff --git a/ibis_tpc/h07.py b/ibis_tpc/h07.py index ea04197..7f70502 100644 --- a/ibis_tpc/h07.py +++ b/ibis_tpc/h07.py @@ -1,6 +1,6 @@ "Volume Shipping Query (Q7)" -from .utils import add_date +import ibis def tpc_h07(con, NATION1="FRANCE", NATION2="GERMANY", DATE="1995-01-01"): @@ -33,7 +33,8 @@ def tpc_h07(con, NATION1="FRANCE", NATION2="GERMANY", DATE="1995-01-01"): [ ((q.cust_nation == NATION1) & (q.supp_nation == NATION2)) | ((q.cust_nation == NATION2) & (q.supp_nation == NATION1)), - q.l_shipdate.between(DATE, add_date(DATE, dy=2, dd=-1)), + q.l_shipdate.between(ibis.date(DATE), ibis.date(DATE) + + ibis.interval(years=2) - ibis.interval(days=1)), ] ) diff --git a/ibis_tpc/h08.py b/ibis_tpc/h08.py index 8f6eb80..82b0bab 100644 --- a/ibis_tpc/h08.py +++ b/ibis_tpc/h08.py @@ -2,8 +2,6 @@ import ibis -from .utils import add_date - def tpc_h08( con, @@ -42,7 +40,8 @@ def tpc_h08( q = q.filter( [ q.r_name == REGION, - q.o_orderdate.between(DATE, add_date(DATE, dy=2, dd=-1)), + q.o_orderdate.between(ibis.date(DATE), ibis.date(DATE) + + ibis.interval(years=2) - ibis.interval(days=1)), q.p_type == TYPE, ] ) diff --git a/ibis_tpc/h10.py b/ibis_tpc/h10.py index 4140c43..e84707f 100644 --- a/ibis_tpc/h10.py +++ b/ibis_tpc/h10.py @@ -1,7 +1,6 @@ "Returned Item Reporting Query (Q10)" import ibis -from .utils import add_date def tpc_h10(con, DATE="1993-10-01"): @@ -17,7 +16,9 @@ def tpc_h10(con, DATE="1993-10-01"): q = q.filter( [ - (q.o_orderdate >= DATE) & (q.o_orderdate < add_date(DATE, dm=3)), + (q.o_orderdate >= ibis.date(DATE)) + & (q.o_orderdate < ibis.date(DATE) + + ibis.interval(months=3)), q.l_returnflag == "R", ] ) diff --git a/ibis_tpc/h12.py b/ibis_tpc/h12.py index 90b1af1..dfb9664 100644 --- a/ibis_tpc/h12.py +++ b/ibis_tpc/h12.py @@ -1,4 +1,4 @@ -from .utils import add_date +import ibis def tpc_h12(con, SHIPMODE1="MAIL", SHIPMODE2="SHIP", DATE="1994-01-01"): @@ -18,8 +18,9 @@ def tpc_h12(con, SHIPMODE1="MAIL", SHIPMODE2="SHIP", DATE="1994-01-01"): q.l_shipmode.isin([SHIPMODE1, SHIPMODE2]), q.l_commitdate < q.l_receiptdate, q.l_shipdate < q.l_commitdate, - q.l_receiptdate >= DATE, - q.l_receiptdate < add_date(DATE, dy=1), + q.l_receiptdate >= ibis.date(DATE), + q.l_receiptdate < ibis.date(DATE) + + ibis.interval(years=1), ] ) diff --git a/ibis_tpc/h14.py b/ibis_tpc/h14.py index fdf1ddd..8eb1593 100644 --- a/ibis_tpc/h14.py +++ b/ibis_tpc/h14.py @@ -1,4 +1,4 @@ -from .utils import add_date +import ibis def tpc_h14(con, DATE="1995-09-01"): @@ -11,7 +11,12 @@ def tpc_h14(con, DATE="1995-09-01"): part = con.table("part") q = lineitem q = q.join(part, lineitem.l_partkey == part.p_partkey) - q = q.filter([q.l_shipdate >= DATE, q.l_shipdate < add_date(DATE, dm=1)]) + q = q.filter( + [ + q.l_shipdate >= ibis.date(DATE), + q.l_shipdate < ibis.date(DATE) + ibis.interval(months=1) + ] + ) revenue = q.l_extendedprice * (1 - q.l_discount) promo_revenue = q.p_type.like("PROMO%").ifelse(revenue, 0) diff --git a/ibis_tpc/h15.py b/ibis_tpc/h15.py index c837342..57c1a04 100644 --- a/ibis_tpc/h15.py +++ b/ibis_tpc/h15.py @@ -1,4 +1,4 @@ -from .utils import add_date +import ibis def tpc_h15(con, DATE="1996-01-01"): @@ -9,7 +9,11 @@ def tpc_h15(con, DATE="1996-01-01"): qrev = lineitem qrev = qrev.filter( - [lineitem.l_shipdate >= DATE, lineitem.l_shipdate < add_date(DATE, dm=3)] + [ + lineitem.l_shipdate >= ibis.date(DATE), + lineitem.l_shipdate < ibis.date(DATE) + + ibis.interval(months=3) + ] ) gqrev = qrev.group_by([lineitem.l_suppkey]) diff --git a/ibis_tpc/h20.py b/ibis_tpc/h20.py index e62a156..94540ac 100644 --- a/ibis_tpc/h20.py +++ b/ibis_tpc/h20.py @@ -1,4 +1,4 @@ -from .utils import add_date +import ibis def tpc_h20(con, COLOR="forest", DATE="1994-01-01", NATION="CANADA"): @@ -23,8 +23,9 @@ def tpc_h20(con, COLOR="forest", DATE="1994-01-01", NATION="CANADA"): [ lineitem.l_partkey == q2.ps_partkey, lineitem.l_suppkey == q2.ps_suppkey, - lineitem.l_shipdate >= DATE, - lineitem.l_shipdate < add_date(DATE, dy=1), + lineitem.l_shipdate >= ibis.date(DATE), + lineitem.l_shipdate < ibis.date(DATE) + + ibis.interval(years=1), ] ) diff --git a/ibis_tpc/utils.py b/ibis_tpc/utils.py deleted file mode 100644 index 2496557..0000000 --- a/ibis_tpc/utils.py +++ /dev/null @@ -1,16 +0,0 @@ -import datetime -import functools -from dateutil.relativedelta import relativedelta - - -@functools.singledispatch -def add_date(datestr, dy=0, dm=0, dd=0): - dt = datetime.date.fromisoformat(datestr) - dt += relativedelta(years=dy, months=dm, days=dd) - return dt.isoformat() - - -@add_date.register(datetime.date) -def _(dt, dy=0, dm=0, dd=0): - dt += relativedelta(years=dy, months=dm, days=dd) - return dt