test: replace add_date() calls in TPC-H queries with Ibis function calls

ibis_tpc/__init__.py

@@ -1,2 +1 @@
-from .utils import *
 from .runners import *

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"])

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())

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])

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])

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,
         ]

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)),
         ]
     )
 

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,
         ]
     )

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",
         ]
     )

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),
         ]
     )
 

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)

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])

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),
         ]
     )