spark_cache_example2.py

sales = spark.read.option("header", True).csv("sales_train_evaluation.csv")

# select d_1~d_100 and turn into long format
cols = ["d_" + str(i) for i in range(1, 100)]
sales = sales \
  .selectExpr("id", "item_id", "dept_id", "cat_id", "store_id", "state_id",
              "stack({}, {}) as (d, amount)".format(len(cols), ', '.join(("'{}', {}".format(i, i) for i in cols)))) \
  .cache()

# group by state_id
groupby_state = sales \
  .groupBy("state_id") \
  .agg(F.sum("amount").alias('amt_tot')) \
  .orderBy(F.col("amt_tot").desc())
groupby_state.show()

# group by store_id
groupby_store = sales \
  .groupBy("store_id") \
  .agg(F.sum("amount").alias('amt_tot')) \
  .orderBy(F.col("amt_tot").desc())
groupby_store.show()