R/tpch.R

In duckplyr: A 'DuckDB'-Backed Version of 'dplyr'

select_opt <- select
# select_opt <- function(x, ...) x

TPCH_NA_MATCHES <- "never"

#' @autoglobal
tpch_01 <- function() {
  lineitem |>
    select_opt(l_shipdate, l_returnflag, l_linestatus, l_quantity, l_extendedprice, l_discount, l_tax) |>
    filter(l_shipdate <= !!as.Date("1998-09-02")) |>
    select_opt(l_returnflag, l_linestatus, l_quantity, l_extendedprice, l_discount, l_tax) |>
    summarise(
      sum_qty = sum(l_quantity),
      sum_base_price = sum(l_extendedprice),
      sum_disc_price = sum(l_extendedprice * (1 - l_discount)),
      sum_charge = sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)),
      avg_qty = mean(l_quantity),
      avg_price = mean(l_extendedprice),
      avg_disc = mean(l_discount),
      count_order = n(),
      .by = c(l_returnflag, l_linestatus)
    ) |>
    arrange(l_returnflag, l_linestatus)
}

#' @autoglobal
tpch_02 <- function() {
  ps <- partsupp |> select_opt(ps_partkey, ps_suppkey, ps_supplycost)

  p <- part |>
    select_opt(p_partkey, p_type, p_size, p_mfgr) |>
    filter(p_size == 15, grepl("BRASS$", p_type)) |>
    select_opt(p_partkey, p_mfgr)

  psp <- inner_join(na_matches = TPCH_NA_MATCHES, p, ps, by = c("p_partkey" = "ps_partkey"))

  sp <- supplier |>
    select_opt(
      s_suppkey, s_nationkey, s_acctbal, s_name,
      s_address, s_phone, s_comment
    )

  psps <- inner_join(psp, sp,
    by = c("ps_suppkey" = "s_suppkey")
  ) |>
    select_opt(
      p_partkey, ps_supplycost, p_mfgr, s_nationkey,
      s_acctbal, s_name, s_address, s_phone, s_comment
    )

  nr <- inner_join(
    nation,
    region |> filter(r_name == "EUROPE"),
    by = c("n_regionkey" = "r_regionkey")
  ) |>
    select_opt(n_nationkey, n_name)

  pspsnr <- inner_join(psps, nr, by = c("s_nationkey" = "n_nationkey")) |>
    select_opt(
      p_partkey, ps_supplycost, p_mfgr, n_name, s_acctbal,
      s_name, s_address, s_phone, s_comment
    )

  aggr <- pspsnr |>
    summarise(min_ps_supplycost = min(ps_supplycost), .by = p_partkey)

  sj <- inner_join(pspsnr, aggr,
    by = c(
      "p_partkey" = "p_partkey",
      "ps_supplycost" = "min_ps_supplycost"
    )
  )

  res <- sj |>
    select(
      s_acctbal, s_name, n_name, p_partkey, p_mfgr,
      s_address, s_phone, s_comment
    ) |>
    arrange(desc(s_acctbal), n_name, s_name, p_partkey) |>
    head(100)

  res
}

#' @autoglobal
tpch_03 <- function() {
  oc <- inner_join(
    orders |>
      select_opt(o_orderkey, o_custkey, o_orderdate, o_shippriority) |>
      filter(o_orderdate < !!as.Date("1995-03-15")),
    customer |>
      select_opt(c_custkey, c_mktsegment) |>
      filter(c_mktsegment == "BUILDING"),
    by = c("o_custkey" = "c_custkey")
  ) |>
    select_opt(o_orderkey, o_orderdate, o_shippriority)

  loc <- inner_join(
    lineitem |>
      select_opt(l_orderkey, l_shipdate, l_extendedprice, l_discount) |>
      filter(l_shipdate > !!as.Date("1995-03-15")) |>
      select_opt(l_orderkey, l_extendedprice, l_discount),
    oc,
    by = c("l_orderkey" = "o_orderkey")
  )

  aggr <- loc |>
    mutate(volume = l_extendedprice * (1 - l_discount)) |>
    summarise(revenue = sum(volume), .by = c(l_orderkey, o_orderdate, o_shippriority)) |>
    select(l_orderkey, revenue, o_orderdate, o_shippriority) |>
    arrange(desc(revenue), o_orderdate) |>
    head(10)
  aggr
}

#' @autoglobal
tpch_04 <- function() {
  l <- lineitem |>
    select_opt(l_orderkey, l_commitdate, l_receiptdate) |>
    filter(l_commitdate < l_receiptdate) |>
    select(l_orderkey)

  o <- orders |>
    select_opt(o_orderkey, o_orderdate, o_orderpriority) |>
    filter(o_orderdate >= !!as.Date("1993-07-01"), o_orderdate < !!as.Date("1993-10-01")) |>
    select(o_orderkey, o_orderpriority)

  # distinct after join, tested and indeed faster
  lo <- inner_join(na_matches = TPCH_NA_MATCHES, l, o, by = c("l_orderkey" = "o_orderkey")) |>
    distinct() |>
    select_opt(o_orderpriority)

  aggr <- lo |>
    summarise(order_count = n(), .by = o_orderpriority) |>
    arrange(o_orderpriority)

  aggr
}

#' @autoglobal
tpch_05 <- function() {
  nr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    nation |>
      select_opt(n_nationkey, n_regionkey, n_name),
    region |>
      select_opt(r_regionkey, r_name) |>
      filter(r_name == "ASIA"),
    by = c("n_regionkey" = "r_regionkey")
  ) |>
    select_opt(n_nationkey, n_name)

  snr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    supplier |>
      select_opt(s_suppkey, s_nationkey),
    nr,
    by = c("s_nationkey" = "n_nationkey")
  ) |>
    select_opt(s_suppkey, s_nationkey, n_name)

  lsnr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    lineitem |> select_opt(l_suppkey, l_orderkey, l_extendedprice, l_discount),
    snr,
    by = c("l_suppkey" = "s_suppkey")
  )

  o <- orders |>
    select_opt(o_orderdate, o_orderkey, o_custkey) |>
    filter(o_orderdate >= !!as.Date("1994-01-01"), o_orderdate < !!as.Date("1995-01-01")) |>
    select_opt(o_orderkey, o_custkey)

  oc <- inner_join(
    na_matches = TPCH_NA_MATCHES, o, customer |> select_opt(c_custkey, c_nationkey),
    by = c("o_custkey" = "c_custkey")
  ) |>
    select_opt(o_orderkey, c_nationkey)

  lsnroc <- inner_join(
    na_matches = TPCH_NA_MATCHES, lsnr, oc,
    by = c("l_orderkey" = "o_orderkey", "s_nationkey" = "c_nationkey")
  ) |>
    select_opt(l_extendedprice, l_discount, n_name)

  aggr <- lsnroc |>
    mutate(volume = l_extendedprice * (1 - l_discount)) |>
    summarise(revenue = sum(volume), .by = n_name) |>
    arrange(desc(revenue))

  aggr
}

#' @autoglobal
tpch_06 <- function() {
  lineitem |>
    select_opt(l_shipdate, l_extendedprice, l_discount, l_quantity) |>
    filter(
      l_shipdate >= !!as.Date("1994-01-01"),
      l_shipdate < !!as.Date("1995-01-01"),
      l_discount >= 0.05,
      l_discount <= 0.07,
      l_quantity < 24
    ) |>
    select_opt(l_extendedprice, l_discount) |>
    summarise(revenue = sum(l_extendedprice * l_discount))
}

#' @autoglobal
tpch_07 <- function() {
  sn <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    supplier |>
      select_opt(s_nationkey, s_suppkey),
    nation |>
      select(n1_nationkey = n_nationkey, n1_name = n_name) |>
      # filter(n1_name %in% c("FRANCE", "GERMANY")),  TODO
      filter(n1_name == "FRANCE" | n1_name == "GERMANY"),
    by = c("s_nationkey" = "n1_nationkey")
  ) |>
    select_opt(s_suppkey, n1_name)

  cn <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    customer |>
      select_opt(c_custkey, c_nationkey),
    nation |>
      select(n2_nationkey = n_nationkey, n2_name = n_name) |>
      # filter(n2_name %in% c("FRANCE", "GERMANY")),
      filter(n2_name == "FRANCE" | n2_name == "GERMANY"),
    by = c("c_nationkey" = "n2_nationkey")
  ) |>
    select_opt(c_custkey, n2_name)

  cno <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    orders |>
      select_opt(o_custkey, o_orderkey),
    cn,
    by = c("o_custkey" = "c_custkey")
  ) |>
    select_opt(o_orderkey, n2_name)

  cnol <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    lineitem |>
      select_opt(l_orderkey, l_suppkey, l_shipdate, l_extendedprice, l_discount) |>
      filter(l_shipdate >= !!as.Date("1995-01-01"), l_shipdate <= !!as.Date("1996-12-31")),
    cno,
    by = c("l_orderkey" = "o_orderkey")
  ) |>
    select_opt(l_suppkey, l_shipdate, l_extendedprice, l_discount, n2_name)

  all <- inner_join(na_matches = TPCH_NA_MATCHES, cnol, sn, by = c("l_suppkey" = "s_suppkey"))

  aggr <- all |>
    filter((n1_name == "FRANCE" & n2_name == "GERMANY") |
      (n1_name == "GERMANY" & n2_name == "FRANCE")) |>
    mutate(
      supp_nation = n1_name,
      cust_nation = n2_name,
      l_year = as.integer(strftime(l_shipdate, "%Y")),
      volume = l_extendedprice * (1 - l_discount)
    ) |>
    select_opt(supp_nation, cust_nation, l_year, volume) |>
    summarise(revenue = sum(volume), .by = c(supp_nation, cust_nation, l_year)) |>
    arrange(supp_nation, cust_nation, l_year)

  aggr
}

#' @autoglobal
tpch_08 <- function() {
  nr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    nation |>
      select(n1_nationkey = n_nationkey, n1_regionkey = n_regionkey),
    region |>
      select_opt(r_regionkey, r_name) |>
      filter(r_name == "AMERICA") |>
      select_opt(r_regionkey),
    by = c("n1_regionkey" = "r_regionkey")
  ) |>
    select_opt(n1_nationkey)

  cnr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    customer |>
      select_opt(c_custkey, c_nationkey),
    nr,
    by = c("c_nationkey" = "n1_nationkey")
  ) |>
    select_opt(c_custkey)

  ocnr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    orders |>
      select_opt(o_orderkey, o_custkey, o_orderdate) |>
      filter(o_orderdate >= !!as.Date("1995-01-01"), o_orderdate <= !!as.Date("1996-12-31")),
    cnr,
    by = c("o_custkey" = "c_custkey")
  ) |>
    select_opt(o_orderkey, o_orderdate)

  locnr <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    lineitem |>
      select_opt(l_orderkey, l_partkey, l_suppkey, l_extendedprice, l_discount),
    ocnr,
    by = c("l_orderkey" = "o_orderkey")
  ) |>
    select_opt(l_partkey, l_suppkey, l_extendedprice, l_discount, o_orderdate)

  locnrp <- inner_join(
    na_matches = TPCH_NA_MATCHES, locnr,
    part |>
      select_opt(p_partkey, p_type) |>
      filter(p_type == "ECONOMY ANODIZED STEEL") |>
      select_opt(p_partkey),
    by = c("l_partkey" = "p_partkey")
  ) |>
    select_opt(l_suppkey, l_extendedprice, l_discount, o_orderdate)

  locnrps <- inner_join(
    na_matches = TPCH_NA_MATCHES, locnrp,
    supplier |>
      select_opt(s_suppkey, s_nationkey),
    by = c("l_suppkey" = "s_suppkey")
  ) |>
    select_opt(l_extendedprice, l_discount, o_orderdate, s_nationkey)

  all <- inner_join(
    na_matches = TPCH_NA_MATCHES, locnrps,
    nation |>
      select(n2_nationkey = n_nationkey, n2_name = n_name),
    by = c("s_nationkey" = "n2_nationkey")
  ) |>
    select_opt(l_extendedprice, l_discount, o_orderdate, n2_name)

  aggr <- all |>
    mutate(
      o_year =  as.integer(strftime(o_orderdate, "%Y")),
      volume = l_extendedprice * (1 - l_discount),
      nation = n2_name
    ) |>
    select_opt(o_year, volume, nation) |>
    summarise(
      mkt_share = sum(if_else(nation == "BRAZIL", volume, 0)) / sum(volume),
      .by = o_year
    ) |>
    arrange(o_year)
  aggr
}

#' @autoglobal
tpch_09 <- function() {
  p <- part |>
    select_opt(p_name, p_partkey) |>
    filter(grepl("green", p_name)) |>
    select_opt(p_partkey)

  psp <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    partsupp |>
      select_opt(ps_suppkey, ps_partkey, ps_supplycost),
    p,
    by = c("ps_partkey" = "p_partkey")
  )

  sn <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    supplier |>
      select_opt(s_suppkey, s_nationkey),
    nation |>
      select_opt(n_nationkey, n_name),
    by = c("s_nationkey" = "n_nationkey")
  ) |>
    select_opt(s_suppkey, n_name)

  pspsn <- inner_join(na_matches = TPCH_NA_MATCHES, psp, sn, by = c("ps_suppkey" = "s_suppkey"))

  lpspsn <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    lineitem |>
      select_opt(l_suppkey, l_partkey, l_orderkey, l_extendedprice, l_discount, l_quantity),
    pspsn,
    by = c("l_suppkey" = "ps_suppkey", "l_partkey" = "ps_partkey")
  ) |>
    select_opt(l_orderkey, l_extendedprice, l_discount, l_quantity, ps_supplycost, n_name)

  all <- inner_join(
    na_matches = TPCH_NA_MATCHES,
    orders |>
      select_opt(o_orderkey, o_orderdate),
    lpspsn,
    by = c("o_orderkey" = "l_orderkey")
  ) |>
    select_opt(l_extendedprice, l_discount, l_quantity, ps_supplycost, n_name, o_orderdate)

  aggr <- all |>
    mutate(
      nation = n_name,
      o_year = as.integer(strftime(o_orderdate, "%Y")),
      amount = l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity
    ) |>
    select_opt(nation, o_year, amount) |>
    summarise(sum_profit = sum(amount), .by = c(nation, o_year)) |>
    arrange(nation, desc(o_year))

  aggr
}

#' @autoglobal
tpch_10 <- function() {
  l <- lineitem |>
    select_opt(l_orderkey, l_returnflag, l_extendedprice, l_discount) |>
    filter(l_returnflag == "R") |>
    select_opt(l_orderkey, l_extendedprice, l_discount)

  o <- orders |>
    select_opt(o_orderkey, o_custkey, o_orderdate) |>
    filter(o_orderdate >= !!as.Date("1993-10-01"), o_orderdate < !!as.Date("1994-01-01")) |>
    select_opt(o_orderkey, o_custkey)

  lo <- inner_join(
    na_matches = TPCH_NA_MATCHES, l, o,
    by = c("l_orderkey" = "o_orderkey")
  ) |>
    select_opt(l_extendedprice, l_discount, o_custkey)
  # first aggregate, then join with customer/nation,
  # otherwise we need to aggr over lots of cols

  lo_aggr <- lo |>
    mutate(volume = l_extendedprice * (1 - l_discount)) |>
    summarise(revenue = sum(volume), .by = o_custkey)

  c <- customer |>
    select_opt(c_custkey, c_nationkey, c_name, c_acctbal, c_phone, c_address, c_comment)

  loc <- inner_join(na_matches = TPCH_NA_MATCHES, c, lo_aggr, by = c("c_custkey" = "o_custkey"))

  locn <- inner_join(
    na_matches = TPCH_NA_MATCHES, loc, nation |> select_opt(n_nationkey, n_name),
    by = c("c_nationkey" = "n_nationkey")
  )

  res <- locn |>
    select(
      c_custkey, c_name, revenue, c_acctbal, n_name,
      c_address, c_phone, c_comment
    ) |>
    arrange(desc(revenue)) |>
    head(20)

  res
}