Ibis Snowflake TPC-H Query 1

ibis_snowflake_tpc-h_1.py

# before running:
# 1. install Ibis and its Snowflake backend: https://ibis-project.org/backends/Snowflake/
# 2. create and activate a Snowflake trial account
# 3. set environment variables SNOWSQL_USER, SNOWSQL_PWD, SNOWSQL_ACCOUNT

import os
import ibis
from ibis import _

ibis.options.interactive = True

# connect to Snowflake
conn = ibis.snowflake.connect(
  user=os.getenv('SNOWSQL_USER'),
  password=os.getenv('SNOWSQL_PWD'),
  account=os.getenv('SNOWSQL_ACCOUNT'),
  database='snowflake_sample_data/tpch_sf1'
)

# assign the LINEITEM table to variable t (an Ibis table object)
t = conn.table('LINEITEM')

# use the Ibis dataframe API to run TPC-H query 1
t.filter(_.L_SHIPDATE.cast('date') <= ibis.date('1998-12-01') + ibis.interval(days=90)) \
 .mutate(DISCOUNT_PRICE = _.L_EXTENDEDPRICE * (1 - _.L_DISCOUNT)) \
 .mutate(CHARGE = _.DISCOUNT_PRICE * (1 + _.L_TAX)) \
 .group_by(['L_RETURNFLAG', 'L_LINESTATUS']) \
 .aggregate(
   sum_qty=_.L_QUANTITY.sum(),
   sum_base_price=_.L_EXTENDEDPRICE.sum(),
   sum_disc_price=_.DISCOUNT_PRICE.sum(),
   sum_charge=_.CHARGE.sum(),
   avg_qty=_.L_QUANTITY.mean(),
   avg_price=_.L_EXTENDEDPRICE.mean(),
   avg_disc=_.L_DISCOUNT.mean(),
   count_order=_.count()
 ) \
 .order_by(['L_RETURNFLAG', 'L_LINESTATUS'])