Add Dask-SQL example notebook

index.rst

@@ -22,6 +22,7 @@ You can run these examples in a live session here: |Binder|
    delayed
    futures
    machine-learning
+   sql
    xarray
 
 .. toctree::

sql.ipynb

@@ -0,0 +1,387 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Operating on Dask Dataframes with SQL\n",
+    "\n",
+    "[Dask-SQL](https://dask-sql.readthedocs.io/en/stable/) is an open source project and Python package leveraging [Apache Calcite](https://calcite.apache.org/) to provide a SQL frontend for [Dask](https://dask.org/) dataframe operations, allowing SQL users to take advantage of Dask's distributed capabilities without requiring an extensive knowledge of the dataframe API."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "! pip install dask-sql"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set up a Dask cluster\n",
+    "\n",
+    "Setting up a Dask [Cluster](https://docs.dask.org/en/latest/deploying.html) is optional, but can dramatically expand our options for distributed computation by giving us access to Dask workers on GPUs, remote machines, common cloud providers, and more).\n",
+    "Additionally, connecting our cluster to a Dask [Client](https://distributed.dask.org/en/stable/client.html) will give us access to a dashboard, which can be used to monitor the progress of active computations and diagnose issues.\n",
+    "\n",
+    "For this notebook, we will create a local cluster and connect it to a client.\n",
+    "Once the client has been created, a link will appear to its associated dashboard, which can be viewed throughout the following computations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dask.distributed import Client\n",
+    "\n",
+    "client = Client(n_workers=2, threads_per_worker=2, memory_limit='1GB')\n",
+    "client"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create a context\n",
+    "\n",
+    "A `dask_sql.Context` is the Python equivalent to a SQL database, serving as an interface to register all tables and functions used in SQL queries, as well as execute the queries themselves.\n",
+    "In typical usage, a single `Context` is created and used for the duration of a Python script or notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dask_sql import Context\n",
+    "c = Context()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load and register data\n",
+    "\n",
+    "Once a `Context` has been created, there are a variety of ways to register tables in it.\n",
+    "The simplest way to do this is through the `create_table` method, which accepts a variety of input types which Dask-SQL then uses to infer the table creation method.\n",
+    "Supported input types include:\n",
+    "\n",
+    "- Dask / [Pandas](https://pandas.pydata.org/)-like dataframes\n",
+    "- String locations of local or remote datasets\n",
+    "- [Apache Hive](https://github.com/apache/hive) tables served through [PyHive](https://github.com/dropbox/PyHive) or [SQLAlchemy](https://www.sqlalchemy.org/)\n",
+    "\n",
+    "Input type can also be specified explicitly by providing a `format`.\n",
+    "When being registered, tables can optionally be persisted into memory by passing `persist=True`, which can greatly speed up repeated queries on the same table at the cost of loading the entire table into memory.\n",
+    "For more information, see [Data Loading and Input](https://dask-sql.readthedocs.io/en/latest/pages/data_input.html)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from dask.datasets import timeseries\n",
+    "\n",
+    "# register and persist a dask table\n",
+    "ddf = timeseries()\n",
+    "c.create_table(\"dask\", ddf, persist=True)\n",
+    "\n",
+    "# register a pandas table (implicitly converted to a dask table)\n",
+    "df = pd.DataFrame({\"a\": [1, 2, 3]})\n",
+    "c.create_table(\"pandas\", df)\n",
+    "\n",
+    "# register a table from local storage; kwargs are passed on to the underlying table creation method\n",
+    "c.create_table(\n",
+    "    \"local\",\n",
+    "    \"surveys/data/2021-user-survey-results.csv.gz\",\n",
+    "    format=\"csv\",\n",
+    "    parse_dates=['Timestamp'],\n",
+    "    blocksize=None\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Tables can also be registered through SQL `CREATE TABLE WITH` or `CREATE TABLE AS` statements, using the `sql` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# replace our table from local storage\n",
+    "c.sql(\"\"\"\n",
+    "    CREATE OR REPLACE TABLE\n",
+    "        \"local\"\n",
+    "    WITH (\n",
+    "        location = 'surveys/data/2021-user-survey-results.csv.gz',\n",
+    "        format = 'csv',\n",
+    "        parse_dates = ARRAY [ 'Timestamp' ]\n",
+    "    )\n",
+    "\"\"\")\n",
+    "\n",
+    "# create a new table from a SQL query\n",
+    "c.sql(\"\"\"\n",
+    "    CREATE TABLE filtered AS (\n",
+    "        SELECT id, name FROM dask WHERE name = 'Zelda'\n",
+    "    )\n",
+    "\"\"\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Dask-SQL currently offers experimental GPU support, powered by the [RAPIDS](https://rapids.ai/) suite of open source GPU data science libraries.\n",
+    "Input support is currently limited to Dask / Pandas-like dataframes and data in local/remote storage, and though most queries run without issue, users should expect some bugs or undefined behavior.\n",
+    "To register a table and mark it for use on GPUs, `gpu=True` can be passed to a standard `create_table` call, or its equivalent `CREATE TABLE WITH` query (note that this requires [cuDF and Dask-cuDF](https://github.com/rapidsai/cudf)).\n",
+    "\n",
+    "```python\n",
+    "# register a dask table for use on GPUs (not possible in this binder)\n",
+    "c.create_table(\"gpu_dask\", ddf, gpu=True)\n",
+    "\n",
+    "# load in a table from disk using GPU-accelerated IO operations\n",
+    "c.sql(\"\"\"\n",
+    "    CREATE TABLE\n",
+    "        \"gpu_local\"\n",
+    "    WITH (\n",
+    "        location = 'surveys/data/2021-user-survey-results.csv.gz',\n",
+    "        format = 'csv',\n",
+    "        parse_dates = ARRAY [ 'Timestamp' ],\n",
+    "        gpu = True\n",
+    "    )\n",
+    "\"\"\")\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Query the data\n",
+    "\n",
+    "When the `sql` method is called, Dask-SQL hands the query off to Apache Calcite to convert into a relational algebra - essentially a list of SQL tasks that must be executed in order to get a result.\n",
+    "The relational algebra of any query can be viewed directly using the `explain` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(c.explain(\"SELECT AVG(x) FROM dask\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "From here, this relational algebra is then converted into a Dask computation graph, which ultimately returns (or in the case of `CREATE TABLE` statements, implicitly assigns) a Dask dataframe."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "c.sql(\"SELECT AVG(x) FROM dask\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Dask dataframes are lazy, meaning that at the time of their creation, none of their dependent tasks have been executed yet.\n",
+    "To actually execute these tasks and get a result, we must call `compute`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "c.sql(\"SELECT AVG(x) FROM dask\").compute()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Looking at the dashboard, we can see that executing this query has triggered some Dask computations.\n",
+    "\n",
+    "Because the return value of a query is a Dask dataframe, it is also possible to do follow-up operations on it using Dask's dataframe API.\n",
+    "This can be useful if we want to perform some complex operations on a dataframe that are not possible through Dask, then follow up with some simpler operations that can easily be expressed through the dataframe API."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# perform a multi-column sort that isn't possible in Dask\n",
+    "res = c.sql(\"\"\"\n",
+    "    SELECT * FROM dask ORDER BY name ASC, id DESC, x ASC\n",
+    "\"\"\")\n",
+    "\n",
+    "# now do some follow groupby aggregations\n",
+    "res.groupby(\"name\").agg({\"x\": \"sum\", \"y\": \"mean\"}).compute()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Custom functions and aggregations\n",
+    "\n",
+    "When standard SQL functionality is insufficient, it is possible to register custom functions for use in queries.\n",
+    "These functions can be classified as one of the following:\n",
+    "\n",
+    "- Column-wise functions\n",
+    "- Row-wise functions\n",
+    "- Aggregations\n",
+    "\n",
+    "### Column-wise functions\n",
+    "\n",
+    "Column-wise functions can take columns or literal values as input and return a column of an identical length.\n",
+    "Column-wise functions can be registered in a `Context` using the `register_function` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "def f(x):\n",
+    "    return x ** 2\n",
+    "\n",
+    "c.register_function(f, \"f\", [(\"x\", np.float64)], np.float64)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Function registration requires the following inputs:\n",
+    "\n",
+    "- A callable function\n",
+    "- A name for the function to be referred to in queries\n",
+    "- A list of tuples, representing the input variables and their respective types, which can be either Pandas or [NumPy](https://numpy.org/) types\n",
+    "- A type for the output column\n",
+    "\n",
+    "Once a function has been registered, it can be called like any other standard SQL function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "c.sql(\"SELECT F(x) FROM dask\").compute()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Row-wise functions\n",
+    "\n",
+    "In some cases, it may be easier to write a custom function that processes a dict-like `row` object - otherwise known as a row-wise function.\n",
+    "These functions can also be registered using `register_function` by passing `row_udf=True`, and used in the same manner as a column-wise function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def g(row):\n",
+    "    if row[\"x\"] > row[\"y\"]:\n",
+    "        return row[\"x\"] - row[\"y\"]\n",
+    "    return row[\"y\"] - row[\"x\"]\n",
+    "\n",
+    "c.register_function(g, \"g\", [(\"x\", np.float64), (\"y\", np.float64)], np.float64, row_udf=True)\n",
+    "\n",
+    "c.sql(\"SELECT G(x, y) FROM dask\").compute()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that unlike column-wise functions, which are called directly using specified columns and literals as input, row-wise functions are called using `apply`, which can have unpredictable performance depending on the underlying dataframe library (e.g. Pandas, cuDF) and the function itself.\n",
+    "\n",
+    "### Aggregations\n",
+    "\n",
+    "Aggregations take a single column as input and return a single value - thus, they can only be used to reduce the results of a `GROUP BY` query.\n",
+    "Aggregations can be registered using the `register_aggregation` method, which is functionally similar to `register_function` but takes a Dask [Aggregation](https://docs.dask.org/en/latest/dataframe-groupby.html#aggregate) as input instead of a callable function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import dask.dataframe as dd\n",
+    "\n",
+    "my_sum = dd.Aggregation(\"my_sum\", lambda x: x.sum(), lambda x: x.sum())\n",
+    "\n",
+    "c.register_aggregation(my_sum, \"my_sum\", [(\"x\", np.float64)], np.float64)\n",
+    "\n",
+    "c.sql(\"SELECT MY_SUM(x) FROM dask\").compute()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}