{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "734557aa-90fb-468f-9ed0-0e6f295bb9eb",
   "metadata": {},
   "source": [
    "<img src=\"./images/DLI_Header.png\" width=400/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dbb572c-1291-4011-9ed4-120eb2ec7b29",
   "metadata": {},
   "source": [
    "# Fundamentals of Accelerated Data Science # "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "377ba9f0-0acc-4574-9a43-c475cfd52dd5",
   "metadata": {},
   "source": [
    "## 08 - Introduction to cuDF Polars  ##\n",
    "\n",
    "**Table of Contents**\n",
    "<br>\n",
    "This notebook briefly introduces Polars and covers the new GPU engine. This notebook covers the below sections: \n",
    "1. [Introduction to Polars](#Introduction-to-Polars)\n",
    "    * [Installation](#Installation)\n",
    "    * [Creating a DataFrame](#Creating-a-DataFrame)\n",
    "    * [Running Basic Operations](#Running-Basic-Operations)\n",
    "    * [Pandas Comparison](#Pandas-Comparison)\n",
    "    * [cuDF Pandas Comparison](#cuDF-Pandas-Comparison)\n",
    "2. [Basic Polars Operations](#Basic-Polars-Operations)\n",
    "    * [Polars Eager Execution API Reference](#Polars-Eager-Execution-API-Reference)\n",
    "    * [Exercise #1 - Load Data](#Exercise-#1---Load-Data)\n",
    "    * [Exercise #2 - Calculate Average Age of Population](#Exercise-#2---Calculate-Average-Age-of-Population)\n",
    "    * [Exercise #3 - Group By and Aggregation](#Exercise-#3---Group-By-and-Aggregation)\n",
    "    * [Exercise #4 - Gender Distribution](#Exercise-#4---Gender-Distribution)\n",
    "4. [Lazy Execution](#Lazy-Execution)\n",
    "    * [Polars Lazy Execution API Reference](#Polars-Lazy-Execution-API-Reference)\n",
    "    * [Execution Graph](#Execution-Graph)\n",
    "    * [Exercise #5 - Creating a Lazy Dataframe](#Exercise-#5---Creating-a-Lazy-Dataframe)\n",
    "    * [Exercise #6 - Query Creation](#Exercise-#6---Query-Creation)\n",
    "5. [cuDF Polars](#cuDF-Polars)\n",
    "    * [Accelerate Previous Code](#Accelerate-Previous-Code)\n",
    "    * [Verify Results Across Engines](#Verify-Results-Across-Engines)\n",
    "    * [Fallback](#Fallback)\n",
    "    * [Exercise #7 - Enable GPU Engine](#Exercise-#7---Enable-GPU-Engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7dd978a-a785-4b67-822e-048d56e231e2",
   "metadata": {},
   "source": [
    "## Introduction to Polars ##\n",
    "Polars is a data analysis and manipulation library that is designed for large data processing (10-100GB) on a single GPU and is known for its speed and memory efficiency. While Pandas makes use of eager execution, Polars additionally has the capability for lazy execution through the built-in query optimizer and makes use of zero-copy optimization techniques. Due to these improvements, Polars typically performs common operations 5-10x faster than Pandas, and requires 2-4 times less RAM. NVIDIA brings hardware acceleration to Polars through a new GPU engine named cuDF Polars, which is available as a pip install."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2187687-9a53-44f9-ba55-ba9e1081cac6",
   "metadata": {},
   "source": [
    "### Creating a DataFrame ###\n",
    "Now let's see how the syntax looks! We will create a dataframe to use within Polars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "be96aa2e-f1fc-4521-bf11-89505a2981ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Taken: 0.7802 seconds\n"
     ]
    }
   ],
   "source": [
    "import polars as pl\n",
    "import time\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "polars_df = pl.read_csv('./data/uk_pop.csv')\n",
    "\n",
    "polars_time = time.time() - start_time\n",
    "\n",
    "print(f\"Time Taken: {polars_time:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "90dedd84-0792-4a54-9380-4b3a020bccd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>age</th><th>sex</th><th>county</th><th>lat</th><th>long</th><th>name</th></tr><tr><td>i64</td><td>str</td><td>str</td><td>f64</td><td>f64</td><td>str</td></tr></thead><tbody><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.533644</td><td>-1.524401</td><td>&quot;FRANCIS&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.426256</td><td>-1.465314</td><td>&quot;EDWARD&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.5552</td><td>-1.496417</td><td>&quot;TEDDY&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.547906</td><td>-1.572341</td><td>&quot;ANGUS&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.477639</td><td>-1.605995</td><td>&quot;CHARLIE&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬─────┬────────────┬───────────┬───────────┬─────────┐\n",
       "│ age ┆ sex ┆ county     ┆ lat       ┆ long      ┆ name    │\n",
       "│ --- ┆ --- ┆ ---        ┆ ---       ┆ ---       ┆ ---     │\n",
       "│ i64 ┆ str ┆ str        ┆ f64       ┆ f64       ┆ str     │\n",
       "╞═════╪═════╪════════════╪═══════════╪═══════════╪═════════╡\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.533644 ┆ -1.524401 ┆ FRANCIS │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.426256 ┆ -1.465314 ┆ EDWARD  │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.5552   ┆ -1.496417 ┆ TEDDY   │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.547906 ┆ -1.572341 ┆ ANGUS   │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.477639 ┆ -1.605995 ┆ CHARLIE │\n",
       "└─────┴─────┴────────────┴───────────┴───────────┴─────────┘"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "polars_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12c66adc-5ac0-49a1-95fe-dc78fbb3950a",
   "metadata": {},
   "source": [
    "### Running Basic Operations ###\n",
    "That was simple- now let's try running a few operations on the dataset! We will be loading the dataset again for a fair comparison with Pandas later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7fc2280f-a89c-42dd-befc-39ffce2f01ec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 6)\n",
      "┌─────┬─────┬──────────────────────────┬───────────┬───────────┬───────┐\n",
      "│ age ┆ sex ┆ county                   ┆ lat       ┆ long      ┆ name  │\n",
      "│ --- ┆ --- ┆ ---                      ┆ ---       ┆ ---       ┆ ---   │\n",
      "│ i64 ┆ str ┆ str                      ┆ f64       ┆ f64       ┆ str   │\n",
      "╞═════╪═════╪══════════════════════════╪═══════════╪═══════════╪═══════╡\n",
      "│ 1   ┆ f   ┆ EAST RIDING OF YORKSHIRE ┆ 53.737344 ┆ -0.638535 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ SHEFFIELD                ┆ 53.35529  ┆ -1.669447 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ LINCOLNSHIRE             ┆ 53.164176 ┆ 0.015812  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ WORCESTERSHIRE           ┆ 52.258629 ┆ -2.31696  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ HERTFORDSHIRE            ┆ 51.731816 ┆ -0.377476 ┆ ZYRAH │\n",
      "└─────┴─────┴──────────────────────────┴───────────┴───────────┴───────┘\n",
      "Time Taken: 6.6280 seconds\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "\n",
    "#load data\n",
    "polars_df = pl.read_csv('./data/uk_pop.csv')\n",
    "\n",
    "# Filter for ages above 0\n",
    "filtered_df = polars_df.filter(pl.col('age') > 0.0)\n",
    "\n",
    "#Sort by name\n",
    "sorted_df = filtered_df.sort('name', descending=True)\n",
    "\n",
    "print(sorted_df.head())\n",
    "polars_time = time.time() - start_time\n",
    "print(f\"Time Taken: {polars_time:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf06893b-a772-4c15-a867-e4cf6aa0984b",
   "metadata": {},
   "source": [
    "### Pandas Comparison ###\n",
    "Let's see how long this would've taken in Pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99468ce0-ecb1-4da2-9e96-44c41c73605a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Taken: 148.7865 seconds\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import time\n",
    "start_time = time.time()\n",
    "pandas_df = pd.read_csv('./data/uk_pop.csv')\n",
    "\n",
    "filtered_df = pandas_df[pandas_df['age'] > 0.0]\n",
    "\n",
    "sorted_df = filtered_df.sort_values(by=['name'], ascending=False)\n",
    "\n",
    "pandas_time = time.time() - start_time\n",
    "print(f\"Time Taken: {pandas_time:.4f} seconds\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5f25cae-1ee8-462e-a487-68af7d9004ce",
   "metadata": {},
   "source": [
    "### cuDF Pandas Comparison ###\n",
    "Wow! That took quite some time to execute. Let's see if we can run it faster with cuDF Pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "23484488-4d35-4359-804d-80bb60463d2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Activate cuDF Pandas\n",
    "%load_ext cudf.pandas\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "714b14dd-0ebd-4814-a341-75cbf8e44fd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Taken for cuDF Pandas: 5.0425 seconds\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import time\n",
    "start_time = time.time()\n",
    "pandas_df = pd.read_csv('./data/uk_pop.csv')\n",
    "\n",
    "filtered_df = pandas_df[pandas_df['age'] > 0.0]\n",
    "\n",
    "sorted_df = filtered_df.sort_values(by=['name'], ascending=False)\n",
    "\n",
    "pandas_time = time.time() - start_time\n",
    "print(f\"Time Taken for cuDF Pandas: {pandas_time:.4f} seconds\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9a3692e-fe36-4bb2-bf1d-84ce5d4834f7",
   "metadata": {},
   "source": [
    "**Note**: Even with cuDF Pandas, we sometimes notice that the performance can be slower than Polars."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd499b82-9b82-472b-934c-5d7e34743458",
   "metadata": {},
   "source": [
    "## Basic Polars Operations ##\n",
    "Please refer to the following API reference guide to complete the exercises below.\n",
    "\n",
    "1. Load data\n",
    "2. Calculate average age of population\n",
    "3. Group By and Aggregation\n",
    "4. Gender Distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3972070-b02e-4870-90c2-bcae0c43f82e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Polars Eager Execution API Reference ###\n",
    "\n",
    "**DataFrame**\n",
    "\n",
    "The main data structure for eager execution in Polars.\n",
    "\n",
    "- `pl.DataFrame(data)`: Create a DataFrame from data\n",
    "- `pl.read_csv(file)`: Read CSV file into DataFrame\n",
    "- `pl.read_parquet(file)`: Read Parquet file into DataFrame\n",
    "\n",
    "**Key Methods**\n",
    "\n",
    "- `filter(mask)`: Filter rows based on a boolean mask\n",
    "- `select(columns)`: Select specific columns\n",
    "- `with_columns(expressions)`: Add or modify columns\n",
    "- `group_by(columns)`: Group by specified columns\n",
    "- `agg(aggregations)`: Perform aggregations on grouped data\n",
    "- `sort(columns)`: Sort the data by specified columns\n",
    "- `join(other, on)`: Join with another DataFrame\n",
    "\n",
    "**Expressions**\n",
    "\n",
    "Used to define operations on columns:\n",
    "\n",
    "- `pl.col(\"column\")`: Reference a column\n",
    "- `pl.lit(value)`: Create a literal value\n",
    "- `pl.when(predicate).then(value).otherwise(other)`: Conditional expression\n",
    "\n",
    "**Series Operations**\n",
    "\n",
    "- `series.sum()`: Calculate sum of series\n",
    "- `series.mean()`: Calculate mean of series\n",
    "- `series.max()`: Find maximum value in series\n",
    "- `series.min()`: Find minimum value in series\n",
    "- `series.sort()`: Sort series values\n",
    "\n",
    "**Data Types**\n",
    "\n",
    "- `pl.Int64`: 64-bit integer\n",
    "- `pl.Float64`: 64-bit float\n",
    "- `pl.Utf8`: String\n",
    "- `pl.Boolean`: Boolean\n",
    "- `pl.Date`: Date\n",
    "\n",
    "**Utilities**\n",
    "\n",
    "- `pl.concat([df1, df2])`: Concatenate DataFrames\n",
    "- `df.describe()`: Generate summary statistics\n",
    "- `df.to_csv(file)`: Write DataFrame to CSV\n",
    "- `df.to_parquet(file)`: Write DataFrame to Parquet\n",
    "\n",
    "The eager API executes operations immediately, providing direct access to results. It's suitable for interactive data exploration and smaller datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87e5c16-46df-4a5b-8454-64fc7bb761a4",
   "metadata": {},
   "source": [
    "### Exercise #1 - Load Data ###\n",
    "Load the csv file into a Dataframe using Polars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e138575b-8817-44e0-981c-128b07d6bf8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "my_df = pl.read_csv('./data/uk_pop.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b230dd-7c17-46fb-be96-1dd11dfbf90d",
   "metadata": {},
   "source": [
    "### Exercise #2 - Calculate Average Age of Population ###\n",
    "Now, filter for individuals aged 65 and above, and sort by ascending age."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d22e37d3-a0c0-4b05-891a-c74086a79e59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>age</th><th>sex</th><th>county</th><th>lat</th><th>long</th><th>name</th></tr><tr><td>i64</td><td>str</td><td>str</td><td>f64</td><td>f64</td><td>str</td></tr></thead><tbody><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.533644</td><td>-1.524401</td><td>&quot;FRANCIS&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.426256</td><td>-1.465314</td><td>&quot;EDWARD&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.5552</td><td>-1.496417</td><td>&quot;TEDDY&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.547906</td><td>-1.572341</td><td>&quot;ANGUS&quot;</td></tr><tr><td>0</td><td>&quot;m&quot;</td><td>&quot;DARLINGTON&quot;</td><td>54.477639</td><td>-1.605995</td><td>&quot;CHARLIE&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬─────┬────────────┬───────────┬───────────┬─────────┐\n",
       "│ age ┆ sex ┆ county     ┆ lat       ┆ long      ┆ name    │\n",
       "│ --- ┆ --- ┆ ---        ┆ ---       ┆ ---       ┆ ---     │\n",
       "│ i64 ┆ str ┆ str        ┆ f64       ┆ f64       ┆ str     │\n",
       "╞═════╪═════╪════════════╪═══════════╪═══════════╪═════════╡\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.533644 ┆ -1.524401 ┆ FRANCIS │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.426256 ┆ -1.465314 ┆ EDWARD  │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.5552   ┆ -1.496417 ┆ TEDDY   │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.547906 ┆ -1.572341 ┆ ANGUS   │\n",
       "│ 0   ┆ m   ┆ DARLINGTON ┆ 54.477639 ┆ -1.605995 ┆ CHARLIE │\n",
       "└─────┴─────┴────────────┴───────────┴───────────┴─────────┘"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_df = my_df.filter(pl.col('age') <= 65)\n",
    "my_df = my_df.sort('age', descending=False)\n",
    "my_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af1bb9b-91bc-4665-a371-699b66067cd2",
   "metadata": {},
   "source": [
    "### Exercise #3 - Group By and Aggregation ###\n",
    "Next, group by county and calculate the total population and average age."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "191a03a2-f420-4669-96ed-ddb290ce09f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>county</th><th>population</th><th>average age</th></tr><tr><td>str</td><td>u32</td><td>f64</td></tr></thead><tbody><tr><td>&quot;KENT&quot;</td><td>1270991</td><td>33.024753</td></tr><tr><td>&quot;ESSEX&quot;</td><td>1190726</td><td>33.357306</td></tr><tr><td>&quot;HAMPSHIRE&quot;</td><td>1096290</td><td>33.899235</td></tr><tr><td>&quot;BIRMINGHAM&quot;</td><td>1002376</td><td>29.690752</td></tr><tr><td>&quot;HERTFORDSHIRE&quot;</td><td>994042</td><td>32.660137</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌───────────────┬────────────┬─────────────┐\n",
       "│ county        ┆ population ┆ average age │\n",
       "│ ---           ┆ ---        ┆ ---         │\n",
       "│ str           ┆ u32        ┆ f64         │\n",
       "╞═══════════════╪════════════╪═════════════╡\n",
       "│ KENT          ┆ 1270991    ┆ 33.024753   │\n",
       "│ ESSEX         ┆ 1190726    ┆ 33.357306   │\n",
       "│ HAMPSHIRE     ┆ 1096290    ┆ 33.899235   │\n",
       "│ BIRMINGHAM    ┆ 1002376    ┆ 29.690752   │\n",
       "│ HERTFORDSHIRE ┆ 994042     ┆ 32.660137   │\n",
       "└───────────────┴────────────┴─────────────┘"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "by_county = my_df.group_by('county').agg([\n",
    "    pl.len().alias('population'),\n",
    "    pl.mean('age').alias('average age'),\n",
    "]).sort('population', descending=True)\n",
    "by_county.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6e2475-3823-43a7-b75b-460e971da230",
   "metadata": {},
   "source": [
    "### Exercise #4 - Gender Distribution ###\n",
    "Lastly, let's calculate the percentage of males to females in the sample data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1fa57226-f174-4967-a75a-b11d9df8765f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sex</th><th>len</th><th>percentage</th></tr><tr><td>str</td><td>u32</td><td>f64</td></tr></thead><tbody><tr><td>&quot;f&quot;</td><td>24076377</td><td>49.763676</td></tr><tr><td>&quot;m&quot;</td><td>24305051</td><td>50.236324</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 3)\n",
       "┌─────┬──────────┬────────────┐\n",
       "│ sex ┆ len      ┆ percentage │\n",
       "│ --- ┆ ---      ┆ ---        │\n",
       "│ str ┆ u32      ┆ f64        │\n",
       "╞═════╪══════════╪════════════╡\n",
       "│ f   ┆ 24076377 ┆ 49.763676  │\n",
       "│ m   ┆ 24305051 ┆ 50.236324  │\n",
       "└─────┴──────────┴────────────┘"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "by_sex = my_df\\\n",
    "    .group_by('sex')\\\n",
    "    .agg(pl.len())\\\n",
    "    .with_columns((pl.col('len') / my_df.shape[0] * 100).alias('percentage'))\n",
    "by_sex.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b44f79a-a542-492d-ae01-b1ee6f9203dc",
   "metadata": {},
   "source": [
    "## Lazy Execution ##\n",
    "Polars utilizes a technique called lazy execution to perform operations. Unlike eager execution, where operations are performed immediately, Polars defines and stores operations in a computational graph that isn't executed until explicitly required. This allows Polars to optimize the sequence of operations to minimize computation overhead and apply optimization techniques such as: applying filters early (predicate pushdown), selecting only necessary columns (projection pushdown), and executing operations in parallel. To make use of lazy execution in polars, a \"LazyFrame\" data structure is used.\n",
    "\n",
    "Now, lets run the same operations with lazy execution and visualize the graph!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a709c71-6fae-4eca-8236-853e3d97e12a",
   "metadata": {},
   "source": [
    "### Polars Lazy Execution API Reference ###\n",
    "\n",
    "**LazyFrame**\n",
    "\n",
    "The main entry point for lazy execution in Polars. Created from a DataFrame or data source.\n",
    "\n",
    "- `pl.LazyFrame(data)`: Create a LazyFrame from data.\n",
    "- `df.lazy()`: Convert a DataFrame to LazyFrame.\n",
    "\n",
    "**Key Methods**\n",
    "\n",
    "- `filter(predicate)`: Filter rows based on a condition.\n",
    "- `select(columns)`: Select specific columns.\n",
    "- `with_columns(expressions)`: Add or modify columns.\n",
    "- `group_by(columns)`: Group by specified columns.\n",
    "- `agg(aggregations)`: Perform aggregations on grouped data.\n",
    "- `sort(columns)`: Sort the data by specified columns.\n",
    "- `join(other, on)`: Join with another LazyFrame.\n",
    "- `collect()`: Execute the lazy query and return a DataFrame.\n",
    "\n",
    "**Expressions**\n",
    "\n",
    "Used to define operations on columns:\n",
    "\n",
    "- `pl.col(\"column\")`: Reference a column.\n",
    "- `pl.lit(value)`: Create a literal value.\n",
    "- `pl.when(predicate).then(value).otherwise(other)`: Define a conditional expression.\n",
    "\n",
    "**Execution**\n",
    "\n",
    "- `collect()`: Execute and return a DataFrame.\n",
    "- `fetch(n)`: Execute and return the first n rows.\n",
    "- `describe_plan()`: Show the query plan for optimization insights.\n",
    "- `explain()`: Explain the query execution process.\n",
    "\n",
    "**Optimization**\n",
    "\n",
    "- `cache()`: Cache intermediate results for faster access.\n",
    "- `optimize()`: Apply query optimizations to improve performance.\n",
    "\n",
    "The lazy API allows building complex queries that are optimized before execution, enabling better performance for large datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f183fa43-99e2-45dc-930c-44fee40d9135",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 6)\n",
      "┌─────┬─────┬──────────────────────────┬───────────┬───────────┬───────┐\n",
      "│ age ┆ sex ┆ county                   ┆ lat       ┆ long      ┆ name  │\n",
      "│ --- ┆ --- ┆ ---                      ┆ ---       ┆ ---       ┆ ---   │\n",
      "│ i64 ┆ str ┆ str                      ┆ f64       ┆ f64       ┆ str   │\n",
      "╞═════╪═════╪══════════════════════════╪═══════════╪═══════════╪═══════╡\n",
      "│ 1   ┆ f   ┆ EAST RIDING OF YORKSHIRE ┆ 53.737344 ┆ -0.638535 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ SHEFFIELD                ┆ 53.35529  ┆ -1.669447 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ LINCOLNSHIRE             ┆ 53.164176 ┆ 0.015812  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ WORCESTERSHIRE           ┆ 52.258629 ┆ -2.31696  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ HERTFORDSHIRE            ┆ 51.731816 ┆ -0.377476 ┆ ZYRAH │\n",
      "└─────┴─────┴──────────────────────────┴───────────┴───────────┴───────┘\n",
      "Time Taken: 6.8741 seconds\n"
     ]
    }
   ],
   "source": [
    "import polars as pl\n",
    "import time\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "# Create a lazy DataFrame\n",
    "lazy_df = pl.scan_csv('./data/uk_pop.csv')\n",
    "\n",
    "# Define the lazy operations\n",
    "lazy_result = (\n",
    "    lazy_df\n",
    "    .filter(pl.col('age') > 0.0)\n",
    "    .sort('name', descending=True)\n",
    ")\n",
    "\n",
    "# Execute the lazy query and collect the results\n",
    "result = lazy_result.collect()\n",
    "\n",
    "print(result.head())\n",
    "polars_time = time.time() - start_time\n",
    "print(f\"Time Taken: {polars_time:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ae1e01f-da85-453f-a708-37afca4753f6",
   "metadata": {},
   "source": [
    "### Execution Graph ###"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4affd1a7-80b2-4f9b-88d8-26a7008f6dc9",
   "metadata": {},
   "source": [
    "Let's see how the unoptimized execution graph looks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d43e9843-ef8e-48aa-b49c-482454d6e93f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"481pt\" height=\"190pt\" viewBox=\"0.00 0.00 481.00 190.00\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 186)\">\n",
       "<title>polars_query</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-186 477,-186 477,4 -4,4\"/>\n",
       "<!-- p1 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>p1</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"326,-182 147,-182 147,-146 326,-146 326,-182\"/>\n",
       "<text text-anchor=\"middle\" x=\"236.5\" y=\"-160.3\" font-family=\"Times,serif\" font-size=\"14.00\">SORT BY [col(&quot;name&quot;)]</text>\n",
       "</g>\n",
       "<!-- p2 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>p2</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"473,-110 0,-110 0,-74 473,-74 473,-110\"/>\n",
       "<text text-anchor=\"middle\" x=\"236.5\" y=\"-88.3\" font-family=\"Times,serif\" font-size=\"14.00\">FILTER BY [(col(&quot;age&quot;).cast(Unknown(Float))) &gt; (dyn float: 0.0)]</text>\n",
       "</g>\n",
       "<!-- p1&#45;&#45;p2 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>p1--p2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M236.5,-145.7C236.5,-134.85 236.5,-120.92 236.5,-110.1\"/>\n",
       "</g>\n",
       "<!-- p3 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>p3</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"350,-38 123,-38 123,0 350,0 350,-38\"/>\n",
       "<text text-anchor=\"middle\" x=\"236.5\" y=\"-22.8\" font-family=\"Times,serif\" font-size=\"14.00\">Csv SCAN [./data/uk_pop.csv]</text>\n",
       "<text text-anchor=\"middle\" x=\"236.5\" y=\"-7.8\" font-family=\"Times,serif\" font-size=\"14.00\">π */6;</text>\n",
       "</g>\n",
       "<!-- p2&#45;&#45;p3 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>p2--p3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M236.5,-73.81C236.5,-62.98 236.5,-49.01 236.5,-38.02\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show unoptimized Graph\n",
    "lazy_result.show_graph(optimized=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6f5883f3-6238-4f98-970b-f06adabfb50e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"414pt\" height=\"133pt\" viewBox=\"0.00 0.00 414.00 133.00\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 129)\">\n",
       "<title>polars_query</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-129 410,-129 410,4 -4,4\"/>\n",
       "<!-- p1 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>p1</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"292.5,-125 113.5,-125 113.5,-89 292.5,-89 292.5,-125\"/>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-103.3\" font-family=\"Times,serif\" font-size=\"14.00\">SORT BY [col(&quot;name&quot;)]</text>\n",
       "</g>\n",
       "<!-- p2 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>p2</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"406,-53 0,-53 0,0 406,0 406,-53\"/>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-37.8\" font-family=\"Times,serif\" font-size=\"14.00\">Csv SCAN [./data/uk_pop.csv]</text>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-22.8\" font-family=\"Times,serif\" font-size=\"14.00\">π */6;</text>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-7.8\" font-family=\"Times,serif\" font-size=\"14.00\">σ [(col(&quot;age&quot;).cast(Unknown(Float))) &gt; (dyn float: 0.0)]</text>\n",
       "</g>\n",
       "<!-- p1&#45;&#45;p2 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>p1--p2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M203,-88.97C203,-78.51 203,-64.9 203,-53.16\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show optimized Graph\n",
    "lazy_result.show_graph(optimized=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d0f6912-9b19-4473-80d3-4223271efe26",
   "metadata": {},
   "source": [
    "As we can see, during execution, Polars ran the age filter in parallel with reading the csv to save time! These type of optimizations is part of the reason why Polars is such a powerful Data Science tool."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5301bcfb-634e-4659-b413-e8279c9bc2ce",
   "metadata": {},
   "source": [
    "### Exercise #5 - Creating a Lazy Dataframe ###\n",
    "First, let's load the csv as a lazy dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "f096b49a-1395-4ddb-bf28-f6e04fb8b4f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "reg_df = pl.read_csv('./data/uk_pop.csv')\n",
    "#lazy_df = pl.scan_csv\n",
    "lazy_df = reg_df.lazy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3581e118-7902-4d04-b4b4-9887c6ff73ba",
   "metadata": {},
   "source": [
    "### Exercise #6 - Query Creation ###\n",
    "Now, let's create a query to find the 5 most common names for individuals under 30. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "1615956e-c6b5-485c-85d3-c0cf89cedbca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>name</th><th>len</th></tr><tr><td>str</td><td>u32</td></tr></thead><tbody><tr><td>&quot;OLIVER&quot;</td><td>218505</td></tr><tr><td>&quot;GEORGE&quot;</td><td>174261</td></tr><tr><td>&quot;HARRY&quot;</td><td>173862</td></tr><tr><td>&quot;OLIVIA&quot;</td><td>171424</td></tr><tr><td>&quot;AMELIA&quot;</td><td>163302</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌────────┬────────┐\n",
       "│ name   ┆ len    │\n",
       "│ ---    ┆ ---    │\n",
       "│ str    ┆ u32    │\n",
       "╞════════╪════════╡\n",
       "│ OLIVER ┆ 218505 │\n",
       "│ GEORGE ┆ 174261 │\n",
       "│ HARRY  ┆ 173862 │\n",
       "│ OLIVIA ┆ 171424 │\n",
       "│ AMELIA ┆ 163302 │\n",
       "└────────┴────────┘"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "most5 = lazy_df\\\n",
    "    .filter(pl.col('age') < 30)\\\n",
    "    .group_by('name')\\\n",
    "    .agg(pl.len())\\\n",
    "    .sort('len', descending=True)\\\n",
    "    .limit(5)\n",
    "most5.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cf331d-0e3d-407c-9990-c7c0873565da",
   "metadata": {},
   "source": [
    "## cuDF Polars ##\n",
    "cuDF Polars is built directly into the Polars Lazy API. The only requirement is to pass engine=\"gpu\" to the collect operation. Polars also allows defining an instance of the GPU engine for greater customization!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "9c6506b1-5c96-4b43-9d53-5826b5a8a44d",
   "metadata": {},
   "outputs": [],
   "source": [
    "lazy_df = pl.scan_csv('./data/uk_pop.csv').collect(engine=\"gpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "676d16d0-9da0-48a1-a9cc-dc60bde2b223",
   "metadata": {},
   "source": [
    "Now let's try defining our own engine object!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "66672414-a47d-48c4-9ad9-5329ae159765",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "import time\n",
    "\n",
    "gpu_engine = pl.GPUEngine(\n",
    "    device=0, # This is the default\n",
    "    raise_on_fail=True, # Fail loudly if we can't run on the GPU.\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "a939efd7-6c74-456b-baf0-8a2d9905956f",
   "metadata": {},
   "outputs": [],
   "source": [
    "lazy_df = pl.scan_csv('./data/uk_pop.csv').collect(engine=gpu_engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eafdfb6-92b0-4f0e-a2ea-90a2ee8a222c",
   "metadata": {},
   "source": [
    "Now that the GPU is warmed up, let's try accelerating the same code as before! Notice that we added an engine parameter to the collect call."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52995c50-bf64-44be-a22e-ba41c778f509",
   "metadata": {},
   "source": [
    "### Accelerate Previous Code ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69e5074d-d15b-471a-a9dd-e1f7a52013a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 6)\n",
      "┌─────┬─────┬──────────────────────────┬───────────┬───────────┬───────┐\n",
      "│ age ┆ sex ┆ county                   ┆ lat       ┆ long      ┆ name  │\n",
      "│ --- ┆ --- ┆ ---                      ┆ ---       ┆ ---       ┆ ---   │\n",
      "│ i64 ┆ str ┆ str                      ┆ f64       ┆ f64       ┆ str   │\n",
      "╞═════╪═════╪══════════════════════════╪═══════════╪═══════════╪═══════╡\n",
      "│ 1   ┆ f   ┆ EAST RIDING OF YORKSHIRE ┆ 53.737344 ┆ -0.638535 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ SHEFFIELD                ┆ 53.35529  ┆ -1.669447 ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ LINCOLNSHIRE             ┆ 53.164176 ┆ 0.015812  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ WORCESTERSHIRE           ┆ 52.258629 ┆ -2.31696  ┆ ZYRAH │\n",
      "│ 1   ┆ f   ┆ HERTFORDSHIRE            ┆ 51.731816 ┆ -0.377476 ┆ ZYRAH │\n",
      "└─────┴─────┴──────────────────────────┴───────────┴───────────┴───────┘\n",
      "Time Taken: 5.9919 seconds\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "\n",
    "# Create a lazy DataFrame\n",
    "lazy_df = pl.scan_csv('./data/uk_pop.csv')\n",
    "\n",
    "# Define the lazy operations\n",
    "lazy_result = (\n",
    "    lazy_df\n",
    "    .filter(pl.col('age') > 0.0)\n",
    "    .sort('name', descending=True)\n",
    ")\n",
    "\n",
    "# Switch to gpu_engine\n",
    "result = lazy_result.collect(engine=gpu_engine)\n",
    "\n",
    "print(result.head())\n",
    "polars_time = time.time() - start_time\n",
    "print(f\"Time Taken: {polars_time:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "025606e3-9d16-416b-9501-b4cf702fb317",
   "metadata": {},
   "source": [
    "### Verify Results Across Engines ###\n",
    "How do we know the results are the same with both the CPU and GPU engine? Luckily with Polars, we can execute the same query across both and compare results using the built in testing module! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "759ad96e-6bc1-4c79-bc50-ab89430fdb42",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The test frames are equal\n"
     ]
    }
   ],
   "source": [
    "from polars.testing import assert_frame_equal\n",
    "\n",
    "# Run on the CPU\n",
    "result_cpu = lazy_result.collect()\n",
    "\n",
    "# Run on the GPU\n",
    "result_gpu = lazy_result.collect(engine=\"gpu\")\n",
    "\n",
    "# assert both result are equal - Will error if not equal, return None otherwise\n",
    "if (assert_frame_equal(result_gpu, result_cpu) == None):\n",
    "    print(\"The test frames are equal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4325cd-0e2d-4cae-8f71-83967708c7b9",
   "metadata": {},
   "source": [
    "### Fallback ###\n",
    "What happens when an operation isn't supported? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "30401f1d-45ad-4365-a00f-e5e79e01412f",
   "metadata": {},
   "outputs": [
    {
     "ename": "ComputeError",
     "evalue": "'cuda' conversion failed: NotImplementedError: ('Query execution with GPU not possible: unsupported operations.\\nThe errors were:\\n- NotImplementedError: rolling mean', [NotImplementedError('rolling mean')])",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mComputeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[47], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m result \u001b[38;5;241m=\u001b[39m (\n\u001b[0;32m----> 2\u001b[0m     \u001b[43mlazy_df\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwith_columns\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpl\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcol\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mage\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrolling_mean\u001b[49m\u001b[43m(\u001b[49m\u001b[43mwindow_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m7\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43malias\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mage_rolling_mean\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfilter\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpl\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcol\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mage\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m  \u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mengine\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgpu_engine\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m )\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28mprint\u001b[39m(result[::\u001b[38;5;241m7\u001b[39m])\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/polars/lazyframe/frame.py:2029\u001b[0m, in \u001b[0;36mLazyFrame.collect\u001b[0;34m(self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, cluster_with_columns, collapse_joins, no_optimization, streaming, engine, background, _eager, **_kwargs)\u001b[0m\n\u001b[1;32m   2027\u001b[0m \u001b[38;5;66;03m# Only for testing purposes\u001b[39;00m\n\u001b[1;32m   2028\u001b[0m callback \u001b[38;5;241m=\u001b[39m _kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpost_opt_callback\u001b[39m\u001b[38;5;124m\"\u001b[39m, callback)\n\u001b[0;32m-> 2029\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m wrap_df(\u001b[43mldf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m)\u001b[49m)\n",
      "\u001b[0;31mComputeError\u001b[0m: 'cuda' conversion failed: NotImplementedError: ('Query execution with GPU not possible: unsupported operations.\\nThe errors were:\\n- NotImplementedError: rolling mean', [NotImplementedError('rolling mean')])"
     ]
    }
   ],
   "source": [
    "result = (\n",
    "    lazy_df\n",
    "    .with_columns(pl.col('age').rolling_mean(window_size=7).alias('age_rolling_mean'))\n",
    "    .filter(pl.col('age') > 0.0)  \n",
    "    .collect(engine=gpu_engine)\n",
    ")\n",
    "print(result[::7])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad531f3-ecfc-4214-b3c7-6506a3321e11",
   "metadata": {},
   "source": [
    "We intially constructed the GPU engine with raise_on_fail=True to ensure all operations ran on GPU. But as we can see, the rolling mean operation is not currently supported, which results in the query not executing. To enable fallback, we can simply change the raise_on_fail parameter to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b95c34-d3ea-494a-95b9-5d2fd8d49938",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_engine_with_fallback = pl.GPUEngine(\n",
    "    device=0, # This is the default\n",
    "    raise_on_fail=False, # Fallback to CPU if we can't run on the GPU (this is the default)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7e15a86-dad3-4b71-8ccb-c03b61a35e3c",
   "metadata": {},
   "source": [
    "Now let's try this query again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73c30f3c-4075-44aa-8715-648b49fec1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = (\n",
    "    lazy_df\n",
    "    .with_columns(pl.col('age').rolling_mean(window_size=7).alias('age_rolling_mean'))\n",
    "    .filter(pl.col('age') > 0.0)  \n",
    "    .collect(engine=gpu_engine_with_fallback)\n",
    ")\n",
    "print(result[::7])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19be5382-ba87-4bc3-8862-df73049598f7",
   "metadata": {},
   "source": [
    "### Exercise #7 - Enable GPU Engine ###\n",
    "The below code calculates the average latitude and longitude for each county. Let's try enabling the GPU Engine for this query!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94243bde-3a3a-4ce9-887c-13d0b0875750",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the lazy query with column pruning\n",
    "lazy_query = (\n",
    "    lazy_df\n",
    "    .select([\"county\", \"lat\", \"long\"])  # Column pruning: select only necessary columns\n",
    "    .group_by(\"county\")\n",
    "    .agg([\n",
    "        pl.col(\"lat\").mean().alias(\"avg_latitude\"),\n",
    "        pl.col(\"long\").mean().alias(\"avg_longitude\")\n",
    "    ])\n",
    "    .sort(\"county\")\n",
    ")\n",
    "\n",
    "# Execute the query\n",
    "result = lazy_query.collect()\n",
    "\n",
    "print(\"\\nAverage latitude and longitude for each county:\")\n",
    "print(result.head())  # Display first few rows"
   ]
  },
  {
   "cell_type": "raw",
   "id": "42acb118-cb9d-478f-9636-694a5ddcb071",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "\n",
    "# Create the lazy query with column pruning\n",
    "lazy_query = (\n",
    "    lazy_df\n",
    "    .select([\"county\", \"lat\", \"long\"])  # Column pruning: select only necessary columns\n",
    "    .group_by(\"county\")\n",
    "    .agg([\n",
    "        pl.col(\"lat\").mean().alias(\"avg_latitude\"),\n",
    "        pl.col(\"long\").mean().alias(\"avg_longitude\")\n",
    "    ])\n",
    "    .sort(\"county\")\n",
    ")\n",
    "\n",
    "# Execute the query\n",
    "result = lazy_query.collect(engine=\"gpu\")\n",
    "\n",
    "print(\"\\nAverage latitude and longitude for each county:\")\n",
    "print(result.head())  # Display first few rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b942cf78-b3cc-4fa1-8d62-b355e20dd2ae",
   "metadata": {},
   "source": [
    "Click ... for solution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf1158c3-8429-4637-a9c2-e8ca91a59965",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdecd196-8e0d-4b36-a428-4411bd480778",
   "metadata": {},
   "source": [
    "**Well Done!**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed99fbcd-6e89-401f-b0cc-0b4fe3f112ad",
   "metadata": {},
   "source": [
    "<img src=\"./images/DLI_Header.png\" width=400/>"
   ]
  }
 ],
 "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
