notebooks/data_frames.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PySpark DataFrames\n",
    "\n",
    "A short introduction to PySpark is provided in this notebook. It is merely a starting point for exploring core features such as PySpark DataFrames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, date\n",
    "\n",
    "import pandas as pd\n",
    "from pyspark.sql import SparkSession, Row\n",
    "from pyspark.sql.functions import lit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start session\n",
    "\n",
    "The main entry point for a PySpark application is provided by `SparkSession`. Before the unification that has been introduced in Spark 2.0, `SparkContext` used to be one of three different starting points. A session object can be initialized as it is shown in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start session\n",
    "spark = SparkSession.builder.getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataFrames\n",
    "\n",
    "A PySpark `DataFrame` can be created in different ways, for example through a `pd.DataFrame` or a list of rows and and an explicit schema. Note that PySpark uses **lazy evaluation** for transformations on DataFrames, whereas pandas is based on **eager execution**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create pandas DataFrame\n",
    "pandas_df = pd.DataFrame(\n",
    "    {\n",
    "        'a': [1, 2, 3],\n",
    "        'b': [2., 3., 4.],\n",
    "        'c': ['string1', 'string2', 'string3'],\n",
    "        'd': [date(2000, 1, 1), date(2000, 2, 1), date(2000, 3, 1)],\n",
    "        'e': [datetime(2000, 1, 1, 12, 0), datetime(2000, 1, 2, 12, 0), datetime(2000, 1, 3, 12, 0)]\n",
    "    }\n",
    ")\n",
    "\n",
    "# create DataFrame from pandas DataFrame\n",
    "spark_df = spark.createDataFrame(pandas_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create spark DataFrame from list of rows\n",
    "spark_df = spark.createDataFrame(\n",
    "    [\n",
    "        Row(a=1, b=2., c='string1', d=date(2000, 1, 1), e=datetime(2000, 1, 1, 12, 0)),\n",
    "        Row(a=2, b=3., c='string2', d=date(2000, 2, 1), e=datetime(2000, 1, 2, 12, 0)),\n",
    "        Row(a=4, b=5., c='string3', d=date(2000, 3, 1), e=datetime(2000, 1, 3, 12, 0))\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create spark DataFrame with an explicit schema\n",
    "spark_df = spark.createDataFrame(\n",
    "    [\n",
    "        (1, 2., 'string1', date(2000, 1, 1), datetime(2000, 1, 1, 12, 0)),\n",
    "        (2, 3., 'string2', date(2000, 2, 1), datetime(2000, 1, 2, 12, 0)),\n",
    "        (3, 4., 'string3', date(2000, 3, 1), datetime(2000, 1, 3, 12, 0))\n",
    "    ],\n",
    "    schema='a long, b double, c string, d date, e timestamp'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show summary\n",
    "spark_df.show()\n",
    "spark_df.printSchema()\n",
    "\n",
    "print(f'Number of rows: {spark_df.count()}')\n",
    "print(f'Number of columns: {len(spark_df.columns)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessing data\n",
    "\n",
    "The following section demonstrates how data in a DataFrame can be accessed. Very often one needs to select certain columns or investigate the values contained in a specific row."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select columns (note that spark DataFrames are lazily evaluated)\n",
    "one_col = spark_df.a\n",
    "\n",
    "two_cols = spark_df.select('a', 'b')\n",
    "two_cols = spark_df['a', 'b']\n",
    "\n",
    "print(one_col)\n",
    "print(two_cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# return the first rows\n",
    "list_of_first_rows = spark_df.head(2) # list_of_first_rows = spark_df.take(2)\n",
    "list_of_last_rows = spark_df.tail(2)\n",
    "\n",
    "print(list_of_first_rows)\n",
    "print(list_of_last_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect distributed data to the driver (note that this may cause an out-of-memory error)\n",
    "list_of_all_rows = spark_df.collect()\n",
    "\n",
    "for row in list_of_all_rows:\n",
    "    print(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get specific rows (with chosen columns)\n",
    "one_row = spark_df.select('a', 'b').collect()[0] # note that this collects to the driver node\n",
    "list_of_two_rows = spark_df.select('a', 'b').collect()[1:3]\n",
    "\n",
    "print(one_row)\n",
    "print(list_of_two_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop over columns\n",
    "for cname in spark_df.columns:\n",
    "    col = spark_df[cname]\n",
    "    print(col)\n",
    "\n",
    "# loop over rows (should be generally avoided for large datasets)\n",
    "for row in spark_df.select('a', 'b').collect(): # collects to driver node\n",
    "    print(row)\n",
    "\n",
    "for idx, row in spark_df.select('a', 'b').toPandas().iterrows(): # collects to driver node\n",
    "    print(row)\n",
    "\n",
    "for row in spark_df.select('a', 'b').toLocalIterator(): # iterates over rows\n",
    "    print(row)\n",
    "\n",
    "spark_df.select('a', 'b').rdd.foreach(lambda row: print(row)) # applies to each row"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter rows of DataFrame\n",
    "new_df = spark_df.filter(spark_df.a == 1) # based on condition\n",
    "new_df = spark_df.filter('a == 1') # based on SQL expression\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manipulating data\n",
    "\n",
    "Operations on DataFrames are either **transformations** (producing new DataFrames) or **actions** (triggering a computation in order to return a value). Examples of actions are `collect()`, `count()`, `show()` and `reduce()`. Transformations are for instance `filter()` and `map()`.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create DataFrame\n",
    "df = spark.createDataFrame(\n",
    "    [\n",
    "        ['red', 'banana', 1, 10],\n",
    "        ['blue', 'banana', 2, 20],\n",
    "        ['red', 'carrot', 3, 30],\n",
    "        ['blue', 'grape', 4, 40],\n",
    "        ['red', 'carrot', 5, 50],\n",
    "        ['black', 'carrot', 6, 60],\n",
    "        ['red', 'banana', 7, 70],\n",
    "        ['red', 'grape', 8, 80]\n",
    "    ],\n",
    "    schema=['color', 'fruit', 'v1', 'v2']\n",
    ")\n",
    "\n",
    "df.show()\n",
    "df.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# group by color and average\n",
    "color_avgs = df.groupby('color').avg()\n",
    "\n",
    "color_avgs.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rename column\n",
    "new_df = df.withColumnRenamed('v1', 'v3')\n",
    "new_df = new_df.withColumnRenamed('v2', 'v4')\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# drop column\n",
    "new_df = df.drop('v2')\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add columns (with constant values)\n",
    "new_df = df.withColumn('new_col1', lit('new_val'))\n",
    "new_df = new_df.withColumn('new_col2', lit(0))\n",
    "\n",
    "new_df.show()\n",
    "new_df.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add column (from existing one)\n",
    "new_df = df.withColumn('v3', df.v2 * 2)\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter rows of DataFrame\n",
    "new_df = df.filter(df.fruit == 'banana') # based on condition\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply row-wise function\n",
    "new_rdd = df.rdd.map(lambda row: (row.color, row.fruit, row.v1 * 2))\n",
    "new_df = new_rdd.toDF(('color', 'fruit', 'double_v1'))\n",
    "\n",
    "new_df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SQL queries\n",
    "\n",
    "DataFrames can also be accessed through SQL-like queries. Some simple examples of such queries are shown in the cells below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# register DataFrame as SQL table\n",
    "df.createOrReplaceTempView('table')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run SQL-style query\n",
    "spark.sql(\n",
    "    'SELECT * FROM table WHERE v2 >= 30'\n",
    ").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run SQL-style query\n",
    "spark.sql(\n",
    "    '''SELECT color, fruit, v2 FROM table\n",
    "    WHERE fruit in ('banana', 'grape')'''\n",
    ").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run SQL-style query\n",
    "spark.sql(\n",
    "    '''SELECT color, fruit, v1 FROM table\n",
    "    ORDER BY fruit'''\n",
    ").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run SQL-style query\n",
    "spark.sql(\n",
    "    '''SELECT fruit, COUNT(*) AS count FROM table\n",
    "    GROUP BY fruit'''\n",
    ").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run SQL-style query\n",
    "spark.sql(\n",
    "    'SELECT COUNT(*) AS count FROM table'\n",
    ").show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stop session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stop session\n",
    "spark.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}