pyspark_window_function_examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## create Spark context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark import SQLContext,SparkConf,SparkContext,HiveContext\n",
    "from pyspark.sql.window import Window\n",
    "from pyspark.sql.types import DateType\n",
    "import pyspark.sql.functions as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_conf = SparkConf(loadDefaults=True)\n",
    "sc = SparkContext(conf=my_conf)\n",
    "sql_context = HiveContext(sc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## generate PySpark dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "shopping_data = \\\n",
    "[('Alex','2018-10-10','Paint',80),('Alex','2018-04-02','Ladder',20),('Alex','2018-06-22','Stool',20),\\\n",
    "('Alex','2018-12-09','Vacuum',40),('Alex','2018-07-12','Bucket',5),('Alex','2018-02-18','Gloves',5),\\\n",
    "('Alex','2018-03-03','Brushes',30),('Alex','2018-09-26','Sandpaper',10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = sql_context.createDataFrame(shopping_data, ['name','date','product','price']).withColumn('date',F.col('date').cast(DateType()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- name: string (nullable = true)\n",
      " |-- date: date (nullable = true)\n",
      " |-- product: string (nullable = true)\n",
      " |-- price: long (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+\n",
      "|name|      date|  product|price|\n",
      "+----+----------+---------+-----+\n",
      "|Alex|2018-10-10|    Paint|   80|\n",
      "|Alex|2018-04-02|   Ladder|   20|\n",
      "|Alex|2018-06-22|    Stool|   20|\n",
      "|Alex|2018-12-09|   Vacuum|   40|\n",
      "|Alex|2018-07-12|   Bucket|    5|\n",
      "|Alex|2018-02-18|   Gloves|    5|\n",
      "|Alex|2018-03-03|  Brushes|   30|\n",
      "|Alex|2018-09-26|Sandpaper|   10|\n",
      "+----+----------+---------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## set the base Window Partition "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = Window.partitionBy('name')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rank , Dense Rank, Percentile of purchases by price"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+----------+\n",
      "|name|      date|  product|price|price_rank|\n",
      "+----+----------+---------+-----+----------+\n",
      "|Alex|2018-10-10|    Paint|   80|         1|\n",
      "|Alex|2018-12-09|   Vacuum|   40|         2|\n",
      "|Alex|2018-03-03|  Brushes|   30|         3|\n",
      "|Alex|2018-04-02|   Ladder|   20|         4|\n",
      "|Alex|2018-06-22|    Stool|   20|         4|\n",
      "|Alex|2018-09-26|Sandpaper|   10|         5|\n",
      "|Alex|2018-07-12|   Bucket|    5|         6|\n",
      "|Alex|2018-02-18|   Gloves|    5|         6|\n",
      "+----+----------+---------+-----+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('price_rank',F.dense_rank().over(w0.orderBy(F.col('price').desc()))).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+----------+\n",
      "|name|      date|  product|price|price_rank|\n",
      "+----+----------+---------+-----+----------+\n",
      "|Alex|2018-07-12|   Bucket|    5|         1|\n",
      "|Alex|2018-02-18|   Gloves|    5|         1|\n",
      "|Alex|2018-09-26|Sandpaper|   10|         3|\n",
      "|Alex|2018-04-02|   Ladder|   20|         4|\n",
      "|Alex|2018-06-22|    Stool|   20|         4|\n",
      "|Alex|2018-03-03|  Brushes|   30|         6|\n",
      "|Alex|2018-12-09|   Vacuum|   40|         7|\n",
      "|Alex|2018-10-10|    Paint|   80|         8|\n",
      "+----+----------+---------+-----+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('price_rank',F.rank().over(w0.orderBy(F.col('price').asc()))).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+------------+\n",
      "|name|      date|  product|price|price_bucket|\n",
      "+----+----------+---------+-----+------------+\n",
      "|Alex|2018-10-10|    Paint|   80|           1|\n",
      "|Alex|2018-12-09|   Vacuum|   40|           1|\n",
      "|Alex|2018-03-03|  Brushes|   30|           2|\n",
      "|Alex|2018-04-02|   Ladder|   20|           2|\n",
      "|Alex|2018-06-22|    Stool|   20|           3|\n",
      "|Alex|2018-09-26|Sandpaper|   10|           3|\n",
      "|Alex|2018-07-12|   Bucket|    5|           4|\n",
      "|Alex|2018-02-18|   Gloves|    5|           4|\n",
      "+----+----------+---------+-----+------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('price_bucket',F.ntile(4).over(w0.orderBy(F.col('price').desc()))).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+-------------------+\n",
      "|name|      date|  product|price|     price_rel_rank|\n",
      "+----+----------+---------+-----+-------------------+\n",
      "|Alex|2018-10-10|    Paint|   80|                0.0|\n",
      "|Alex|2018-12-09|   Vacuum|   40|0.14285714285714285|\n",
      "|Alex|2018-03-03|  Brushes|   30| 0.2857142857142857|\n",
      "|Alex|2018-04-02|   Ladder|   20|0.42857142857142855|\n",
      "|Alex|2018-06-22|    Stool|   20|0.42857142857142855|\n",
      "|Alex|2018-09-26|Sandpaper|   10| 0.7142857142857143|\n",
      "|Alex|2018-07-12|   Bucket|    5| 0.8571428571428571|\n",
      "|Alex|2018-02-18|   Gloves|    5| 0.8571428571428571|\n",
      "+----+----------+---------+-----+-------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('price_rel_rank',F.percent_rank().over(w0.orderBy(F.col('price').desc()))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moving Average and Running Sum of price"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+-----------+----------------+-----------+------------+-------------+\n",
      "|name|      date|  product|price|avg_to_date|accumulating_sum|max_to_date|max_of_last2|items_to_date|\n",
      "+----+----------+---------+-----+-----------+----------------+-----------+------------+-------------+\n",
      "|Alex|2018-02-18|   Gloves|    5|        5.0|               5|          5|           5|            1|\n",
      "|Alex|2018-03-03|  Brushes|   30|       17.5|              35|         30|          30|            2|\n",
      "|Alex|2018-04-02|   Ladder|   20|      18.33|              55|         30|          30|            3|\n",
      "|Alex|2018-06-22|    Stool|   20|      18.75|              75|         30|          20|            4|\n",
      "|Alex|2018-07-12|   Bucket|    5|       16.0|              80|         30|          20|            5|\n",
      "|Alex|2018-09-26|Sandpaper|   10|       15.0|              90|         30|          10|            6|\n",
      "|Alex|2018-10-10|    Paint|   80|      24.29|             170|         80|          80|            7|\n",
      "|Alex|2018-12-09|   Vacuum|   40|      26.25|             210|         80|          80|            8|\n",
      "+----+----------+---------+-----+-----------+----------------+-----------+------------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('avg_to_date',     F.round(F.avg('price').over(w0.orderBy(F.col('date'))),2))\\\n",
    "  .withColumn('accumulating_sum',F.sum('price').over(w0.orderBy(F.col('date'))))\\\n",
    "  .withColumn('max_to_date',     F.max('price').over(w0.orderBy(F.col('date'))))\\\n",
    "  .withColumn('max_of_last2',    F.max('price').over(w0.orderBy(F.col('date')).rowsBetween(-1,Window.currentRow)))\\\n",
    "  .withColumn('items_to_date',   F.count('*').over(w0.orderBy(F.col('date'))))\\\n",
    "  .show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time Lag and Lead for purchases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+-----------------------+-------------------------+\n",
      "|name|      date|  product|price|days_from_last_purchase|days_before_next_purchase|\n",
      "+----+----------+---------+-----+-----------------------+-------------------------+\n",
      "|Alex|2018-02-18|   Gloves|    5|                   null|                       13|\n",
      "|Alex|2018-03-03|  Brushes|   30|                     13|                       30|\n",
      "|Alex|2018-04-02|   Ladder|   20|                     30|                       81|\n",
      "|Alex|2018-06-22|    Stool|   20|                     81|                       20|\n",
      "|Alex|2018-07-12|   Bucket|    5|                     20|                       76|\n",
      "|Alex|2018-09-26|Sandpaper|   10|                     76|                       14|\n",
      "|Alex|2018-10-10|    Paint|   80|                     14|                       60|\n",
      "|Alex|2018-12-09|   Vacuum|   40|                     60|                     null|\n",
      "+----+----------+---------+-----+-----------------------+-------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('days_from_last_purchase', F.datediff('date',F.lag('date',1).over(w0.orderBy(F.col('date')))))\\\n",
    "  .withColumn('days_before_next_purchase', F.datediff(F.lead('date',1).over(w0.orderBy(F.col('date'))),'date'))\\\n",
    "  .show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df.withColumn('days_before_next_purchase', F.datediff(F.lead('date',1).over(w0.orderBy(F.col('date'))),'date')).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## List and Sets  of purchases by member"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+----------------+-----------------------+\n",
      "|name|date      |product  |price|items_by_price  |all_prices             |\n",
      "+----+----------+---------+-----+----------------+-----------------------+\n",
      "|Alex|2018-07-12|Bucket   |5    |[Bucket, Gloves]|[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-02-18|Gloves   |5    |[Bucket, Gloves]|[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-09-26|Sandpaper|10   |[Sandpaper]     |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-10-10|Paint    |80   |[Paint, Paint]  |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-10-11|Paint    |80   |[Paint, Paint]  |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-03-03|Brushes  |30   |[Brushes]       |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-04-02|Ladder   |20   |[Ladder, Stool] |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-06-22|Stool    |20   |[Ladder, Stool] |[30, 5, 20, 10, 40, 80]|\n",
      "|Alex|2018-12-09|Vacuum   |40   |[Vacuum]        |[30, 5, 20, 10, 40, 80]|\n",
      "+----+----------+---------+-----+----------------+-----------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "newRow = sql_context.createDataFrame([('Alex','2018-10-11','Paint',80)])\n",
    "df2 = df.union(newRow)\n",
    "\n",
    "df2.withColumn('items_by_price', F.collect_list('product').over(w0.partitionBy('price')))\\\n",
    "   .withColumn('all_prices',     F.collect_set('price').over(w0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+-----+----------------+\n",
      "|name|Price|           items|\n",
      "+----+-----+----------------+\n",
      "|Alex|    5|[Bucket, Gloves]|\n",
      "|Alex|   10|     [Sandpaper]|\n",
      "|Alex|   80|         [Paint]|\n",
      "|Alex|   30|       [Brushes]|\n",
      "|Alex|   20| [Ladder, Stool]|\n",
      "|Alex|   40|        [Vacuum]|\n",
      "+----+-----+----------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df2.withColumn('items', F.collect_set('product').over(w0.partitionBy('price')))\\\n",
    "   .select('name','Price','items')\\\n",
    "   .distinct()\\\n",
    "   .show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other tricks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### multiple  order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "w1 = Window.partitionBy('name').orderBy(F.col('price').desc(),F.col('date').asc())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+-------------------+\n",
      "|name|      date|  product|price|price_and_date_rank|\n",
      "+----+----------+---------+-----+-------------------+\n",
      "|Alex|2018-10-10|    Paint|   80|                  1|\n",
      "|Alex|2018-12-09|   Vacuum|   40|                  2|\n",
      "|Alex|2018-03-03|  Brushes|   30|                  3|\n",
      "|Alex|2018-04-02|   Ladder|   20|                  4|\n",
      "|Alex|2018-06-22|    Stool|   20|                  5|\n",
      "|Alex|2018-09-26|Sandpaper|   10|                  6|\n",
      "|Alex|2018-02-18|   Gloves|    5|                  7|\n",
      "|Alex|2018-07-12|   Bucket|    5|                  8|\n",
      "+----+----------+---------+-----+-------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.withColumn('price_and_date_rank',F.rank().over(w1)).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### moving average"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+----------+----------------+------------------+\n",
      "|name|      date|  product|price| unix_time|30day_moving_avg| 30day_moving_stdv|\n",
      "+----+----------+---------+-----+----------+----------------+------------------+\n",
      "|Alex|2018-02-18|   Gloves|    5|1518933600|             5.0| 17.67766952966369|\n",
      "|Alex|2018-03-03|  Brushes|   30|1520056800|            17.5| 17.67766952966369|\n",
      "|Alex|2018-04-02|   Ladder|   20|1522645200|            25.0|7.0710678118654755|\n",
      "|Alex|2018-06-22|    Stool|   20|1529643600|            20.0|               NaN|\n",
      "|Alex|2018-07-12|   Bucket|    5|1531371600|            12.5|10.606601717798213|\n",
      "|Alex|2018-09-26|Sandpaper|   10|1537938000|            10.0| 49.49747468305833|\n",
      "|Alex|2018-10-10|    Paint|   80|1539147600|            45.0| 49.49747468305833|\n",
      "|Alex|2018-12-09|   Vacuum|   40|1544335200|            40.0|               NaN|\n",
      "+----+----------+---------+-----+----------+----------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "days = lambda i: i * 86400 # 86400 seconds in a day  \n",
    "\n",
    "df.withColumn('unix_time',F.col('date').cast('timestamp').cast('long'))\\\n",
    "  .withColumn('30day_moving_avg', F.avg('price').over(w0.orderBy(F.col('unix_time')).rangeBetween(-days(30),0)))\\\n",
    "  .withColumn('30day_moving_stdv',F.stddev('price').over(w0.orderBy(F.col('unix_time')).rangeBetween(-days(30),days(15))))\\\n",
    "  .show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### multiple windows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------+---------+-----+----------+---------+\n",
      "|name|      date|  product|price|price_rank|date_rank|\n",
      "+----+----------+---------+-----+----------+---------+\n",
      "|Alex|2018-02-18|   Gloves|    5|         1|        1|\n",
      "|Alex|2018-03-03|  Brushes|   30|         4|        2|\n",
      "|Alex|2018-04-02|   Ladder|   20|         3|        3|\n",
      "|Alex|2018-06-22|    Stool|   20|         3|        4|\n",
      "|Alex|2018-07-12|   Bucket|    5|         1|        5|\n",
      "|Alex|2018-09-26|Sandpaper|   10|         2|        6|\n",
      "|Alex|2018-10-10|    Paint|   80|         6|        7|\n",
      "|Alex|2018-12-09|   Vacuum|   40|         5|        8|\n",
      "+----+----------+---------+-----+----------+---------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "w2 = Window.partitionBy('name').orderBy(F.col('price'))\n",
    "w3 = Window.partitionBy('name').orderBy(F.col('date'))\n",
    "\n",
    "df.withColumn('price_rank',F.dense_rank().over(w2))\\\n",
    "  .withColumn('date_rank',F.row_number().over(w3))\\\n",
    "  .show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}