examples/basics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "529f4422-5c3a-4bd6-abe0-a15edfc62abb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "import polars_ds as pds\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a314316",
   "metadata": {},
   "source": [
    "# This notebook illustrates the basic usage of this package\n",
    "\n",
    "You need to create an environment with this package installed to run this notebook. (usually latest version)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aef5c69-fff3-4779-9b58-f939d725f0b0",
   "metadata": {},
   "source": [
    "# Num Extensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "430fec01-5d0b-422f-b099-c86037512b6d",
   "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, 11)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th><th>time_idx</th><th>dummy</th><th>a</th><th>b</th><th>x1</th><th>x2</th><th>y</th><th>actual</th><th>predicted</th><th>dummy_groups</th></tr><tr><td>f64</td><td>i64</td><td>str</td><td>f64</td><td>f64</td><td>i64</td><td>i64</td><td>i64</td><td>i32</td><td>f64</td><td>str</td></tr></thead><tbody><tr><td>0.0</td><td>0</td><td>&quot;a&quot;</td><td>0.634299</td><td>0.495647</td><td>0</td><td>10000</td><td>-10000</td><td>1</td><td>0.198915</td><td>&quot;a&quot;</td></tr><tr><td>0.841471</td><td>1</td><td>&quot;a&quot;</td><td>0.835205</td><td>0.169188</td><td>1</td><td>10001</td><td>-9999</td><td>0</td><td>0.701006</td><td>&quot;a&quot;</td></tr><tr><td>0.909297</td><td>2</td><td>&quot;a&quot;</td><td>0.492115</td><td>0.738706</td><td>2</td><td>10002</td><td>-9998</td><td>1</td><td>0.562431</td><td>&quot;a&quot;</td></tr><tr><td>0.14112</td><td>3</td><td>&quot;a&quot;</td><td>0.903647</td><td>0.284837</td><td>3</td><td>10003</td><td>-9997</td><td>1</td><td>0.811163</td><td>&quot;a&quot;</td></tr><tr><td>-0.756802</td><td>4</td><td>&quot;a&quot;</td><td>0.65933</td><td>0.870768</td><td>4</td><td>10004</td><td>-9996</td><td>1</td><td>0.238711</td><td>&quot;a&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 11)\n",
       "┌───────────┬──────────┬───────┬──────────┬───┬────────┬────────┬───────────┬──────────────┐\n",
       "│ f         ┆ time_idx ┆ dummy ┆ a        ┆ … ┆ y      ┆ actual ┆ predicted ┆ dummy_groups │\n",
       "│ ---       ┆ ---      ┆ ---   ┆ ---      ┆   ┆ ---    ┆ ---    ┆ ---       ┆ ---          │\n",
       "│ f64       ┆ i64      ┆ str   ┆ f64      ┆   ┆ i64    ┆ i32    ┆ f64       ┆ str          │\n",
       "╞═══════════╪══════════╪═══════╪══════════╪═══╪════════╪════════╪═══════════╪══════════════╡\n",
       "│ 0.0       ┆ 0        ┆ a     ┆ 0.634299 ┆ … ┆ -10000 ┆ 1      ┆ 0.198915  ┆ a            │\n",
       "│ 0.841471  ┆ 1        ┆ a     ┆ 0.835205 ┆ … ┆ -9999  ┆ 0      ┆ 0.701006  ┆ a            │\n",
       "│ 0.909297  ┆ 2        ┆ a     ┆ 0.492115 ┆ … ┆ -9998  ┆ 1      ┆ 0.562431  ┆ a            │\n",
       "│ 0.14112   ┆ 3        ┆ a     ┆ 0.903647 ┆ … ┆ -9997  ┆ 1      ┆ 0.811163  ┆ a            │\n",
       "│ -0.756802 ┆ 4        ┆ a     ┆ 0.65933  ┆ … ┆ -9996  ┆ 1      ┆ 0.238711  ┆ a            │\n",
       "└───────────┴──────────┴───────┴──────────┴───┴────────┴────────┴───────────┴──────────────┘"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 10_000\n",
    "df = pl.DataFrame({\n",
    "    \"f\": np.sin(list(range(size)))\n",
    "    , \"time_idx\": range(size)\n",
    "    , \"dummy\": [\"a\"] * (size // 2) + [\"b\"] * (size // 2)\n",
    "    , \"a\": np.random.random(size = size)\n",
    "    , \"b\": np.random.random(size = size)\n",
    "    , \"x1\" : range(size)\n",
    "    , \"x2\" : range(size, size + size)\n",
    "    , \"y\": range(-size, 0)\n",
    "    , \"actual\": np.round(np.random.random(size=size)).astype(np.int32)\n",
    "    , \"predicted\": np.random.random(size=size)\n",
    "    , \"dummy_groups\":[\"a\"] * (size//2) + [\"b\"] * (size//2) \n",
    "})\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6f98453-34cd-4afc-b35d-db58fa60a69a",
   "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: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>x1</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────┐\n",
       "│ x1  │\n",
       "│ --- │\n",
       "│ f64 │\n",
       "╞═════╡\n",
       "│ 0.0 │\n",
       "└─────┘"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Column-wise Jaccard Similarity. Result should be 0 as they are distinct\n",
    "df.select(\n",
    "    pds.query_jaccard_col(\"x1\", pl.col(\"x2\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "416d5346-e75b-4769-a953-e898d6a4d84c",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th></tr><tr><td>array[f64, 2]</td></tr></thead><tbody><tr><td>[1.939505, 0.0]</td></tr><tr><td>[1.939506, 0.000209]</td></tr><tr><td>[1.939508, 0.000418]</td></tr><tr><td>[1.939512, 0.000627]</td></tr><tr><td>[1.939518, 0.000835]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────┐\n",
       "│ f                    │\n",
       "│ ---                  │\n",
       "│ array[f64, 2]        │\n",
       "╞══════════════════════╡\n",
       "│ [1.939505, 0.0]      │\n",
       "│ [1.939506, 0.000209] │\n",
       "│ [1.939508, 0.000418] │\n",
       "│ [1.939512, 0.000627] │\n",
       "│ [1.939518, 0.000835] │\n",
       "└──────────────────────┘"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. First is real part, second is complex part\n",
    "# By default, this behaves the same as np's rfft, which returns a non-redundant \n",
    "# compact representation of fft output.\n",
    "df.select(\n",
    "    pds.rfft(\"f\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "71c76353",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. But return the full length\n",
    "df.select(\n",
    "    pds.rfft(\"f\", return_full=True)\n",
    ").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cd6662d2",
   "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>f</th><th>a</th><th>b</th></tr><tr><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.0</td><td>-0.634299</td><td>-0.495647</td></tr><tr><td>-0.841471</td><td>-0.835205</td><td>-0.169188</td></tr><tr><td>-0.909297</td><td>-0.492115</td><td>-0.738706</td></tr><tr><td>-0.14112</td><td>-0.903647</td><td>-0.284837</td></tr><tr><td>0.756802</td><td>-0.025031</td><td>-0.375121</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌───────────┬───────────┬───────────┐\n",
       "│ f         ┆ a         ┆ b         │\n",
       "│ ---       ┆ ---       ┆ ---       │\n",
       "│ f64       ┆ f64       ┆ f64       │\n",
       "╞═══════════╪═══════════╪═══════════╡\n",
       "│ 0.0       ┆ -0.634299 ┆ -0.495647 │\n",
       "│ -0.841471 ┆ -0.835205 ┆ -0.169188 │\n",
       "│ -0.909297 ┆ -0.492115 ┆ -0.738706 │\n",
       "│ -0.14112  ┆ -0.903647 ┆ -0.284837 │\n",
       "│ 0.756802  ┆ -0.025031 ┆ -0.375121 │\n",
       "└───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convolution (by FFT). \n",
    "# Modes: `same`, `left` (left-aligned same), `right` (right-aligned same), `valid` or `full`\n",
    "# Currently slower than SciPy but provides parallelism because of Polars\n",
    "df.select(\n",
    "    pds.convolve(\"f\", [-1, 0, 0, 0, 1], mode = \"full\"),\n",
    "    pds.convolve(\"a\", [-1, 0, 0, 0, 1], mode = \"full\"),\n",
    "    pds.convolve(\"b\", [-1, 0, 0, 0, 1], mode = \"full\"),\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ed47b643-6bcc-43f6-9a25-82168c33e7fd",
   "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: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[2.0, -1.0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────┐\n",
       "│ y           │\n",
       "│ ---         │\n",
       "│ list[f64]   │\n",
       "╞═════════════╡\n",
       "│ [2.0, -1.0] │\n",
       "└─────────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Least Square (Linear Regression)\n",
    "df.select(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e6da23d",
   "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, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>idx</th><th>coeff</th><th>std_err</th><th>t</th><th>p&gt;|t|</th></tr><tr><td>u16</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0</td><td>2.0</td><td>2.3854e-16</td><td>8.3842e15</td><td>0.0</td></tr><tr><td>1</td><td>-1.0</td><td>9.0158e-17</td><td>-1.1092e16</td><td>0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 5)\n",
       "┌─────┬───────┬────────────┬────────────┬───────┐\n",
       "│ idx ┆ coeff ┆ std_err    ┆ t          ┆ p>|t| │\n",
       "│ --- ┆ ---   ┆ ---        ┆ ---        ┆ ---   │\n",
       "│ u16 ┆ f64   ┆ f64        ┆ f64        ┆ f64   │\n",
       "╞═════╪═══════╪════════════╪════════════╪═══════╡\n",
       "│ 0   ┆ 2.0   ┆ 2.3854e-16 ┆ 8.3842e15  ┆ 0.0   │\n",
       "│ 1   ┆ -1.0  ┆ 9.0158e-17 ┆ -1.1092e16 ┆ 0.0   │\n",
       "└─────┴───────┴────────────┴────────────┴───────┘"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.query_lstsq_report(\n",
    "        # str | pl.Expr\n",
    "        \"x1\", \"x2\",\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    ).alias(\"report\")\n",
    ").unnest(\"report\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4c34e7b2",
   "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: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[2.0, -1.0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────┐\n",
       "│ y           │\n",
       "│ ---         │\n",
       "│ list[f64]   │\n",
       "╞═════════════╡\n",
       "│ [2.0, -1.0] │\n",
       "└─────────────┘"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.lazy().select(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = \"y\", # We can either put pl.col(\"y\") here or just the string \"y\"\n",
    "        add_bias=False\n",
    "    )\n",
    ").collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "16511624-fc7f-45fc-b28e-ad9a91c1bfe9",
   "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>dummy</th><th>coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌───────┬─────────────┐\n",
       "│ dummy ┆ coeffs      │\n",
       "│ ---   ┆ ---         │\n",
       "│ str   ┆ list[f64]   │\n",
       "╞═══════╪═════════════╡\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "└───────┴─────────────┘"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    \"dummy\",\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    ).over(pl.col(\"dummy\"))\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f550c7c6",
   "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, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>x1</th><th>x2</th><th>y</th><th>pred</th><th>resid</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0</td><td>10000</td><td>-10000</td><td>-10000.0</td><td>5.8208e-11</td></tr><tr><td>1</td><td>10001</td><td>-9999</td><td>-9999.0</td><td>5.8208e-11</td></tr><tr><td>2</td><td>10002</td><td>-9998</td><td>-9998.0</td><td>5.8208e-11</td></tr><tr><td>3</td><td>10003</td><td>-9997</td><td>-9997.0</td><td>5.8208e-11</td></tr><tr><td>4</td><td>10004</td><td>-9996</td><td>-9996.0</td><td>5.8208e-11</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌─────┬───────┬────────┬──────────┬────────────┐\n",
       "│ x1  ┆ x2    ┆ y      ┆ pred     ┆ resid      │\n",
       "│ --- ┆ ---   ┆ ---    ┆ ---      ┆ ---        │\n",
       "│ i64 ┆ i64   ┆ i64    ┆ f64      ┆ f64        │\n",
       "╞═════╪═══════╪════════╪══════════╪════════════╡\n",
       "│ 0   ┆ 10000 ┆ -10000 ┆ -10000.0 ┆ 5.8208e-11 │\n",
       "│ 1   ┆ 10001 ┆ -9999  ┆ -9999.0  ┆ 5.8208e-11 │\n",
       "│ 2   ┆ 10002 ┆ -9998  ┆ -9998.0  ┆ 5.8208e-11 │\n",
       "│ 3   ┆ 10003 ┆ -9997  ┆ -9997.0  ┆ 5.8208e-11 │\n",
       "│ 4   ┆ 10004 ┆ -9996  ┆ -9996.0  ┆ 5.8208e-11 │\n",
       "└─────┴───────┴────────┴──────────┴────────────┘"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# If you want prediction and residue instead of coefficients\n",
    "df.select(\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    \"y\",\n",
    "    pds.query_lstsq(\n",
    "        \"x1\", pl.col(\"x2\"),\n",
    "        target = \"y\",\n",
    "        add_bias=False, \n",
    "        return_pred=True\n",
    "    ).alias(\"prediction\")\n",
    ").unnest(\"prediction\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0e9fb061-340d-423d-9107-772387006ff2",
   "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, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;b&quot;</td><td>[2.0, -1.0]</td></tr><tr><td>&quot;a&quot;</td><td>[2.0, -1.0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬─────────────┐\n",
       "│ dummy ┆ coeffs      │\n",
       "│ ---   ┆ ---         │\n",
       "│ str   ┆ list[f64]   │\n",
       "╞═══════╪═════════════╡\n",
       "│ b     ┆ [2.0, -1.0] │\n",
       "│ a     ┆ [2.0, -1.0] │\n",
       "└───────┴─────────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy\").agg(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "765ff27b",
   "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: (10, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>time_idx</th><th>coefficients</th></tr><tr><td>i64</td><td>list[f64]</td></tr></thead><tbody><tr><td>30</td><td>[2.0, -1.0]</td></tr><tr><td>31</td><td>[2.0, -1.0]</td></tr><tr><td>32</td><td>[2.0, -1.0]</td></tr><tr><td>33</td><td>[2.0, -1.0]</td></tr><tr><td>34</td><td>[2.0, -1.0]</td></tr><tr><td>35</td><td>[2.0, -1.0]</td></tr><tr><td>36</td><td>[2.0, -1.0]</td></tr><tr><td>37</td><td>[2.0, -1.0]</td></tr><tr><td>38</td><td>[2.0, -1.0]</td></tr><tr><td>39</td><td>[2.0, -1.0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10, 2)\n",
       "┌──────────┬──────────────┐\n",
       "│ time_idx ┆ coefficients │\n",
       "│ ---      ┆ ---          │\n",
       "│ i64      ┆ list[f64]    │\n",
       "╞══════════╪══════════════╡\n",
       "│ 30       ┆ [2.0, -1.0]  │\n",
       "│ 31       ┆ [2.0, -1.0]  │\n",
       "│ 32       ┆ [2.0, -1.0]  │\n",
       "│ 33       ┆ [2.0, -1.0]  │\n",
       "│ 34       ┆ [2.0, -1.0]  │\n",
       "│ 35       ┆ [2.0, -1.0]  │\n",
       "│ 36       ┆ [2.0, -1.0]  │\n",
       "│ 37       ┆ [2.0, -1.0]  │\n",
       "│ 38       ┆ [2.0, -1.0]  │\n",
       "│ 39       ┆ [2.0, -1.0]  │\n",
       "└──────────┴──────────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Rolling regression, kind of slow rn\n",
    "df.lazy().rolling(\n",
    "    index_column = pl.col(\"time_idx\").set_sorted(),\n",
    "    period = \"30i\",\n",
    "    # offset = \"-1i\"\n",
    ").agg(\n",
    "    pds.query_lstsq(pl.col(\"x1\"), pl.col(\"x2\"), target = pl.col(\"y\"), add_bias=False).alias(\"coefficients\")\n",
    ").slice(offset = 30).select(\n",
    "    \"time_idx\",\n",
    "    \"coefficients\",\n",
    ").collect().head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d8fda8ca-57e7-4e02-a3f0-283ecce66a59",
   "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: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>-0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌──────┐\n",
       "│ y    │\n",
       "│ ---  │\n",
       "│ f64  │\n",
       "╞══════╡\n",
       "│ -0.0 │\n",
       "└──────┘"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Conditional Entropy, should be 0 because x1 is an ID\n",
    "df.select(\n",
    "    pds.query_cond_entropy(\"y\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "81def1cf",
   "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: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[288675.133151, 28.891478, 28.621473]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────────────────────────┐\n",
       "│ a                               │\n",
       "│ ---                             │\n",
       "│ list[f64]                       │\n",
       "╞═════════════════════════════════╡\n",
       "│ [288675.133151, 28.891478, 28.… │\n",
       "└─────────────────────────────────┘"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Only want singular values (principal values?)\n",
    "df.select(\n",
    "    pds.query_singular_values(\"a\", \"b\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cc497383",
   "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, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>singular_value</th><th>weight_vector</th></tr><tr><td>f64</td><td>list[f64]</td></tr></thead><tbody><tr><td>28.891478</td><td>[0.728048, 0.685527]</td></tr><tr><td>28.62156</td><td>[-0.685527, 0.728048]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌────────────────┬───────────────────────┐\n",
       "│ singular_value ┆ weight_vector         │\n",
       "│ ---            ┆ ---                   │\n",
       "│ f64            ┆ list[f64]             │\n",
       "╞════════════════╪═══════════════════════╡\n",
       "│ 28.891478      ┆ [0.728048, 0.685527]  │\n",
       "│ 28.62156       ┆ [-0.685527, 0.728048] │\n",
       "└────────────────┴───────────────────────┘"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Singular values + The principal components\n",
    "df.select(\n",
    "    pds.query_pca(\"a\", \"b\")\n",
    ").unnest(\"a\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2e036f",
   "metadata": {},
   "source": [
    "# ML Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "85d0d094-3c4c-4230-a589-1027c5690162",
   "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, 8)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy_groups</th><th>l2</th><th>log loss</th><th>precision</th><th>recall</th><th>f</th><th>average_precision</th><th>roc_auc</th></tr><tr><td>str</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>&quot;b&quot;</td><td>0.337956</td><td>1.010538</td><td>0.482977</td><td>0.497148</td><td>0.24498</td><td>0.482876</td><td>0.491038</td></tr><tr><td>&quot;a&quot;</td><td>0.328728</td><td>0.991122</td><td>0.518827</td><td>0.515354</td><td>0.258542</td><td>0.510587</td><td>0.50804</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 8)\n",
       "┌──────────────┬──────────┬──────────┬───────────┬──────────┬──────────┬────────────────┬──────────┐\n",
       "│ dummy_groups ┆ l2       ┆ log loss ┆ precision ┆ recall   ┆ f        ┆ average_precis ┆ roc_auc  │\n",
       "│ ---          ┆ ---      ┆ ---      ┆ ---       ┆ ---      ┆ ---      ┆ ion            ┆ ---      │\n",
       "│ str          ┆ f64      ┆ f64      ┆ f64       ┆ f64      ┆ f64      ┆ ---            ┆ f64      │\n",
       "│              ┆          ┆          ┆           ┆          ┆          ┆ f64            ┆          │\n",
       "╞══════════════╪══════════╪══════════╪═══════════╪══════════╪══════════╪════════════════╪══════════╡\n",
       "│ b            ┆ 0.337956 ┆ 1.010538 ┆ 0.482977  ┆ 0.497148 ┆ 0.24498  ┆ 0.482876       ┆ 0.491038 │\n",
       "│ a            ┆ 0.328728 ┆ 0.991122 ┆ 0.518827  ┆ 0.515354 ┆ 0.258542 ┆ 0.510587       ┆ 0.50804  │\n",
       "└──────────────┴──────────┴──────────┴───────────┴──────────┴──────────┴────────────────┴──────────┘"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy_groups\").agg(\n",
    "    pds.query_l2(\"actual\", \"predicted\").alias(\"l2\"),\n",
    "    pds.query_log_loss(\"actual\", \"predicted\").alias(\"log loss\"),\n",
    "    pds.query_binary_metrics(actual=\"actual\", pred=\"predicted\").alias(\"combo\")\n",
    ").unnest(\"combo\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8d7c6e3-0f1d-45f0-9fdb-cdb303b98556",
   "metadata": {},
   "source": [
    "# Str Extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "54ad36f9-264e-4a49-bf36-936639440edf",
   "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>sen</th><th>word</th></tr><tr><td>str</td><td>str</td></tr></thead><tbody><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────────────────────────────┬───────┐\n",
       "│ sen                             ┆ word  │\n",
       "│ ---                             ┆ ---   │\n",
       "│ str                             ┆ str   │\n",
       "╞═════════════════════════════════╪═══════╡\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "│ Hello, world! I'm going to chu… ┆ word  │\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "│ Hello, world! I'm going to chu… ┆ word  │\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "└─────────────────────────────────┴───────┘"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 100_000\n",
    "df2 = pl.DataFrame({\n",
    "    \"sen\":[\"Hello, world! I'm going to church.\"] * size,\n",
    "    \"word\":[\"words\", \"word\"] * (size //2)\n",
    "})\n",
    "df2.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ee123a7e-7f9b-4f48-a5d5-6354799201ab",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;going&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;world&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;to&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ going  │\n",
       "│ church │\n",
       "│ world  │\n",
       "│ hello  │\n",
       "│ to     │\n",
       "└────────┘"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Tokenize\n",
    "df2.select(\n",
    "    pds.str_tokenize(pl.col(\"sen\").str.to_lowercase()).explode().unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f33017e3-17df-498b-93d9-1d656a344388",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;world&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;go&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ world  │\n",
       "│ church │\n",
       "│        │\n",
       "│ hello  │\n",
       "│ go     │\n",
       "└────────┘"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select(\n",
    "    pds.str_tokenize(pl.col(\"sen\").str.to_lowercase(), stem=True).explode().unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "69237c02-5f9f-4e92-b68d-6ac43aad1a79",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ u32  │\n",
       "╞══════╡\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "└──────┘"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select(\n",
    "    pds.str_leven(\"word\", pl.lit(\"world\"))\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2eba320c",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ u32  │\n",
       "╞══════╡\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "└──────┘"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Damerau-Levenshtein\n",
    "df2.select(\n",
    "    pds.str_d_leven(\"word\", pl.lit(\"world\"))\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "795396dc",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.6</td></tr><tr><td>0.8</td></tr><tr><td>0.6</td></tr><tr><td>0.8</td></tr><tr><td>0.6</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ f64  │\n",
       "╞══════╡\n",
       "│ 0.6  │\n",
       "│ 0.8  │\n",
       "│ 0.6  │\n",
       "│ 0.8  │\n",
       "│ 0.6  │\n",
       "└──────┘"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select( # column \"word\" vs. the word \"world\"\n",
    "    pds.str_leven(\"word\", pl.lit(\"world\"), return_sim = True)\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2dad7633-67fa-47f3-b86a-9f4cd097a650",
   "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>sen</th><th>word</th></tr><tr><td>str</td><td>str</td></tr></thead><tbody><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────────────────────────────┬──────┐\n",
       "│ sen                             ┆ word │\n",
       "│ ---                             ┆ ---  │\n",
       "│ str                             ┆ str  │\n",
       "╞═════════════════════════════════╪══════╡\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "└─────────────────────────────────┴──────┘"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.filter(\n",
    "    # This is way faster than computing ditance and then doing a filter\n",
    "    pds.filter_by_levenshtein(pl.col(\"word\"), pl.lit(\"world\"), 1) # <= 1. \n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "dc9477c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame({\n",
    "    \"word\":[\"apple\", \"banana\", \"pineapple\", \"asasasas\", \"sasasass\"],\n",
    "    \"other_data\": [1,2,3,4,5]\n",
    "})\n",
    "gibberish = [\"asasasa\", \"sasaaasss\", \"asdasadadfa\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c50591e0",
   "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, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th><th>other_data</th></tr><tr><td>str</td><td>i64</td></tr></thead><tbody><tr><td>&quot;asasasas&quot;</td><td>4</td></tr><tr><td>&quot;sasasass&quot;</td><td>5</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌──────────┬────────────┐\n",
       "│ word     ┆ other_data │\n",
       "│ ---      ┆ ---        │\n",
       "│ str      ┆ i64        │\n",
       "╞══════════╪════════════╡\n",
       "│ asasasas ┆ 4          │\n",
       "│ sasasass ┆ 5          │\n",
       "└──────────┴────────────┘"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    pds.similar_to_vocab(\n",
    "        pl.col(\"word\"),\n",
    "        vocab = gibberish,\n",
    "        threshold = 0.5,\n",
    "        metric = \"lv\", # Levenshtein similarity. Other options: dleven, osa, jw\n",
    "        strategy = \"any\" # True if the word is similar to any word in vocab. Other options: \"all\", \"avg\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "7ece3794",
   "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>asasasa</th><th>sasaaasss</th><th>asdasadadfa</th><th>LCS based Fuzz match - apples</th><th>Optimal String Alignment - apples</th><th>Jaro-Winkler - apples</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.142857</td><td>0.111111</td><td>0.090909</td><td>0.833333</td><td>0.833333</td><td>0.966667</td></tr><tr><td>0.428571</td><td>0.333333</td><td>0.272727</td><td>0.166667</td><td>0.0</td><td>0.444444</td></tr><tr><td>0.111111</td><td>0.111111</td><td>0.090909</td><td>0.555556</td><td>0.444444</td><td>0.5</td></tr><tr><td>0.875</td><td>0.666667</td><td>0.545455</td><td>0.25</td><td>0.25</td><td>0.527778</td></tr><tr><td>0.75</td><td>0.777778</td><td>0.454545</td><td>0.25</td><td>0.25</td><td>0.527778</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌──────────┬───────────┬─────────────┬────────────────┬──────────────────┬─────────────────────────┐\n",
       "│ asasasa  ┆ sasaaasss ┆ asdasadadfa ┆ LCS based Fuzz ┆ Optimal String   ┆ Jaro-Winkler - apples   │\n",
       "│ ---      ┆ ---       ┆ ---         ┆ match - apples ┆ Alignment - app… ┆ ---                     │\n",
       "│ f64      ┆ f64       ┆ f64         ┆ ---            ┆ ---              ┆ f64                     │\n",
       "│          ┆           ┆             ┆ f64            ┆ f64              ┆                         │\n",
       "╞══════════╪═══════════╪═════════════╪════════════════╪══════════════════╪═════════════════════════╡\n",
       "│ 0.142857 ┆ 0.111111  ┆ 0.090909    ┆ 0.833333       ┆ 0.833333         ┆ 0.966667                │\n",
       "│ 0.428571 ┆ 0.333333  ┆ 0.272727    ┆ 0.166667       ┆ 0.0              ┆ 0.444444                │\n",
       "│ 0.111111 ┆ 0.111111  ┆ 0.090909    ┆ 0.555556       ┆ 0.444444         ┆ 0.5                     │\n",
       "│ 0.875    ┆ 0.666667  ┆ 0.545455    ┆ 0.25           ┆ 0.25             ┆ 0.527778                │\n",
       "│ 0.75     ┆ 0.777778  ┆ 0.454545    ┆ 0.25           ┆ 0.25             ┆ 0.527778                │\n",
       "└──────────┴───────────┴─────────────┴────────────────┴──────────────────┴─────────────────────────┘"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.str_leven(\"word\", pl.lit(\"asasasa\"), return_sim=True).alias(\"asasasa\"),\n",
    "    pds.str_leven(\"word\", pl.lit(\"sasaaasss\"), return_sim=True).alias(\"sasaaasss\"),\n",
    "    pds.str_leven(\"word\", pl.lit(\"asdasadadfa\"), return_sim=True).alias(\"asdasadadfa\"),\n",
    "    pds.str_fuzz(\"word\", pl.lit(\"apples\")).alias(\"LCS based Fuzz match - apples\"),\n",
    "    pds.str_osa(\"word\", pl.lit(\"apples\"), return_sim=True).alias(\"Optimal String Alignment - apples\"),\n",
    "    pds.str_jw(\"word\", pl.lit(\"apples\")).alias(\"Jaro-Winkler - apples\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8841f2a1",
   "metadata": {},
   "source": [
    "# Stats Extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2c6171b0",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>null</td></tr><tr><td>null</td></tr><tr><td>-0.343264</td></tr><tr><td>0.215017</td></tr><tr><td>-0.632978</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ a         │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ null      │\n",
       "│ null      │\n",
       "│ -0.343264 │\n",
       "│ 0.215017  │\n",
       "│ -0.632978 │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "df = pl.DataFrame({\n",
    "    \"a\": [None, None] + list(np.random.normal(size = 998))\n",
    "})\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2f6e7445",
   "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>a</th><th>random</th></tr><tr><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>null</td><td>null</td></tr><tr><td>null</td><td>null</td></tr><tr><td>-0.343264</td><td>-0.655657</td></tr><tr><td>0.215017</td><td>-0.666812</td></tr><tr><td>-0.632978</td><td>0.646823</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌───────────┬───────────┐\n",
       "│ a         ┆ random    │\n",
       "│ ---       ┆ ---       │\n",
       "│ f64       ┆ f64       │\n",
       "╞═══════════╪═══════════╡\n",
       "│ null      ┆ null      │\n",
       "│ null      ┆ null      │\n",
       "│ -0.343264 ┆ -0.655657 │\n",
       "│ 0.215017  ┆ -0.666812 │\n",
       "│ -0.632978 ┆ 0.646823  │\n",
       "└───────────┴───────────┘"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random numbers, respecting null positions in reference column (pl.col(\"a\"))\n",
    "df.with_columns(\n",
    "    pl.col(\"a\").stats.rand_normal(mean = 0.5, std = 1., respect_null=True).alias(\"random\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "11e13f55",
   "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>a</th><th>random_str</th></tr><tr><td>f64</td><td>str</td></tr></thead><tbody><tr><td>null</td><td>null</td></tr><tr><td>null</td><td>null</td></tr><tr><td>-0.343264</td><td>&quot;OD2g&quot;</td></tr><tr><td>0.215017</td><td>&quot;PZw&quot;</td></tr><tr><td>-0.632978</td><td>&quot;FoQU&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌───────────┬────────────┐\n",
       "│ a         ┆ random_str │\n",
       "│ ---       ┆ ---        │\n",
       "│ f64       ┆ str        │\n",
       "╞═══════════╪════════════╡\n",
       "│ null      ┆ null       │\n",
       "│ null      ┆ null       │\n",
       "│ -0.343264 ┆ OD2g       │\n",
       "│ 0.215017  ┆ PZw        │\n",
       "│ -0.632978 ┆ FoQU       │\n",
       "└───────────┴────────────┘"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random string\n",
    "df.with_columns(\n",
    "    pl.col(\"a\").stats.rand_str(min_size = 1, max_size = 5, respect_null=True).alias(\"random_str\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "43c37394",
   "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>a</th><th>random_str</th></tr><tr><td>f64</td><td>str</td></tr></thead><tbody><tr><td>null</td><td>null</td></tr><tr><td>null</td><td>null</td></tr><tr><td>-0.343264</td><td>&quot;OIqE9&quot;</td></tr><tr><td>0.215017</td><td>&quot;tQhcG&quot;</td></tr><tr><td>-0.632978</td><td>&quot;bRnZL&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌───────────┬────────────┐\n",
       "│ a         ┆ random_str │\n",
       "│ ---       ┆ ---        │\n",
       "│ f64       ┆ str        │\n",
       "╞═══════════╪════════════╡\n",
       "│ null      ┆ null       │\n",
       "│ null      ┆ null       │\n",
       "│ -0.343264 ┆ OIqE9      │\n",
       "│ 0.215017  ┆ tQhcG      │\n",
       "│ -0.632978 ┆ bRnZL      │\n",
       "└───────────┴────────────┘"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate fixed size random string, while respecting column a's nulls\n",
    "df.with_columns(\n",
    "    pl.col(\"a\").stats.rand_str(min_size = 5, max_size = 5, respect_null=True).alias(\"random_str\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "3d0c06a4",
   "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, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>test1</th><th>test2</th><th>test1_perturbed</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>null</td><td>0.239046</td><td>null</td><td>0.239152</td></tr><tr><td>null</td><td>0.63256</td><td>null</td><td>0.632709</td></tr><tr><td>-0.343264</td><td>-0.90806</td><td>2.663809</td><td>-0.908275</td></tr><tr><td>0.215017</td><td>-1.601119</td><td>1.172176</td><td>-1.601176</td></tr><tr><td>-0.632978</td><td>0.602753</td><td>2.513366</td><td>0.602348</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌───────────┬───────────┬──────────┬─────────────────┐\n",
       "│ a         ┆ test1     ┆ test2    ┆ test1_perturbed │\n",
       "│ ---       ┆ ---       ┆ ---      ┆ ---             │\n",
       "│ f64       ┆ f64       ┆ f64      ┆ f64             │\n",
       "╞═══════════╪═══════════╪══════════╪═════════════════╡\n",
       "│ null      ┆ 0.239046  ┆ null     ┆ 0.239152        │\n",
       "│ null      ┆ 0.63256   ┆ null     ┆ 0.632709        │\n",
       "│ -0.343264 ┆ -0.90806  ┆ 2.663809 ┆ -0.908275       │\n",
       "│ 0.215017  ┆ -1.601119 ┆ 1.172176 ┆ -1.601176       │\n",
       "│ -0.632978 ┆ 0.602753  ┆ 2.513366 ┆ 0.602348        │\n",
       "└───────────┴───────────┴──────────┴─────────────────┘"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    # Sample from a normal distribution, using reference column \"a\" 's mean and std\n",
    "    pl.col(\"a\").stats.rand_normal().alias(\"test1\") \n",
    "    # Sample from uniform distribution, with low = 0 and high = \"a\"'s max, and respect the nulls in \"a\"\n",
    "    , pl.col(\"a\").stats.rand_uniform(low = 0., high = None, respect_null=True).alias(\"test2\")\n",
    ").with_columns(\n",
    "    # Add a random pertubation to test1\n",
    "    pds.perturb(\"test1\", epsilon=0.001).alias(\"test1_perturbed\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "67dc6583",
   "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, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>[0, 1)</th><th>Normal</th><th>Int from [0, 10)</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>i32</td></tr></thead><tbody><tr><td>null</td><td>0.583019</td><td>-0.509022</td><td>6</td></tr><tr><td>null</td><td>0.678392</td><td>-0.222092</td><td>2</td></tr><tr><td>-0.343264</td><td>0.356277</td><td>0.427585</td><td>9</td></tr><tr><td>0.215017</td><td>0.740707</td><td>-0.624831</td><td>4</td></tr><tr><td>-0.632978</td><td>0.784875</td><td>0.281009</td><td>5</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌───────────┬──────────┬───────────┬──────────────────┐\n",
       "│ a         ┆ [0, 1)   ┆ Normal    ┆ Int from [0, 10) │\n",
       "│ ---       ┆ ---      ┆ ---       ┆ ---              │\n",
       "│ f64       ┆ f64      ┆ f64       ┆ i32              │\n",
       "╞═══════════╪══════════╪═══════════╪══════════════════╡\n",
       "│ null      ┆ 0.583019 ┆ -0.509022 ┆ 6                │\n",
       "│ null      ┆ 0.678392 ┆ -0.222092 ┆ 2                │\n",
       "│ -0.343264 ┆ 0.356277 ┆ 0.427585  ┆ 9                │\n",
       "│ 0.215017  ┆ 0.740707 ┆ -0.624831 ┆ 4                │\n",
       "│ -0.632978 ┆ 0.784875 ┆ 0.281009  ┆ 5                │\n",
       "└───────────┴──────────┴───────────┴──────────────────┘"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# New in v0.3.5\n",
    "# This way, we don't have a reference column, so we cannot respect nulls, but is more convenient to use.\n",
    "df.with_columns(\n",
    "    pds.random().alias(\"[0, 1)\"),\n",
    "    pds.random_normal(pl.col(\"a\").mean(), pl.col(\"a\").std()).alias(\"Normal\"),\n",
    "    pds.random_int(0, 10).alias(\"Int from [0, 10)\"),\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "7b63f636",
   "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: (1, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>t-tests: statistics</th><th>t-tests: pvalue</th><th>normality_test: statistics</th><th>normality_test: pvalue</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>-1.080921</td><td>0.279913</td><td>1.221052</td><td>0.543065</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 4)\n",
       "┌─────────────────────┬─────────────────┬────────────────────────────┬────────────────────────┐\n",
       "│ t-tests: statistics ┆ t-tests: pvalue ┆ normality_test: statistics ┆ normality_test: pvalue │\n",
       "│ ---                 ┆ ---             ┆ ---                        ┆ ---                    │\n",
       "│ f64                 ┆ f64             ┆ f64                        ┆ f64                    │\n",
       "╞═════════════════════╪═════════════════╪════════════════════════════╪════════════════════════╡\n",
       "│ -1.080921           ┆ 0.279913        ┆ 1.221052                   ┆ 0.543065               │\n",
       "└─────────────────────┴─────────────────┴────────────────────────────┴────────────────────────┘"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate 2 random sample, both normally distributed\n",
    "# Run Welch's t test on them, p value should be big since they have equal mean\n",
    "# Run a normality test. Again, p value should be big since they are normally distributed \n",
    "\n",
    "df.with_columns(\n",
    "    pds.random_normal(0.5, 1.0).alias(\"test1\"),\n",
    "    pds.random_normal(0.5, 2.0).alias(\"test2\"),\n",
    ").select(\n",
    "    pds.query_ttest_ind(\"test1\", \"test2\", equal_var=False).alias(\"t-test\"),\n",
    "    pds.normal_test(\"test1\").alias(\"normality_test\")\n",
    ").select(\n",
    "    pl.col(\"t-test\").struct.field(\"statistic\").alias(\"t-tests: statistics\")\n",
    "    , pl.col(\"t-test\").struct.field(\"pvalue\").alias(\"t-tests: pvalue\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"statistic\").alias(\"normality_test: statistics\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"pvalue\").alias(\"normality_test: pvalue\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b46a72a5",
   "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, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>market_id</th><th>var1</th><th>var2</th><th>category_1</th><th>category_2</th></tr><tr><td>i64</td><td>f64</td><td>f64</td><td>i32</td><td>i32</td></tr></thead><tbody><tr><td>0</td><td>0.789191</td><td>0.657872</td><td>4</td><td>1</td></tr><tr><td>1</td><td>0.072238</td><td>0.205825</td><td>3</td><td>5</td></tr><tr><td>2</td><td>0.96792</td><td>0.868846</td><td>1</td><td>1</td></tr><tr><td>0</td><td>0.44313</td><td>0.609293</td><td>2</td><td>4</td></tr><tr><td>1</td><td>0.126652</td><td>0.335964</td><td>1</td><td>6</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌───────────┬──────────┬──────────┬────────────┬────────────┐\n",
       "│ market_id ┆ var1     ┆ var2     ┆ category_1 ┆ category_2 │\n",
       "│ ---       ┆ ---      ┆ ---      ┆ ---        ┆ ---        │\n",
       "│ i64       ┆ f64      ┆ f64      ┆ i32        ┆ i32        │\n",
       "╞═══════════╪══════════╪══════════╪════════════╪════════════╡\n",
       "│ 0         ┆ 0.789191 ┆ 0.657872 ┆ 4          ┆ 1          │\n",
       "│ 1         ┆ 0.072238 ┆ 0.205825 ┆ 3          ┆ 5          │\n",
       "│ 2         ┆ 0.96792  ┆ 0.868846 ┆ 1          ┆ 1          │\n",
       "│ 0         ┆ 0.44313  ┆ 0.609293 ┆ 2          ┆ 4          │\n",
       "│ 1         ┆ 0.126652 ┆ 0.335964 ┆ 1          ┆ 6          │\n",
       "└───────────┴──────────┴──────────┴────────────┴────────────┘"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 5_000\n",
    "df = pl.DataFrame({\n",
    "    \"market_id\": range(size),\n",
    "}).with_columns(\n",
    "    pl.col(\"market_id\").mod(3),\n",
    "    var1 = pds.random(),\n",
    "    var2 = pds.random(),\n",
    "    category_1 = pds.random_int(0, 5),\n",
    "    category_2 = pds.random_int(0, 10),\n",
    ")\n",
    "\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "adc4f66f",
   "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: (1, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>t-test</th><th>chi2-test</th><th>f-test</th></tr><tr><td>struct[2]</td><td>struct[2]</td><td>struct[2]</td></tr></thead><tbody><tr><td>{-1.340382,0.180152}</td><td>{38.269766,0.366801}</td><td>{1.836938,0.118797}</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 3)\n",
       "┌──────────────────────┬──────────────────────┬─────────────────────┐\n",
       "│ t-test               ┆ chi2-test            ┆ f-test              │\n",
       "│ ---                  ┆ ---                  ┆ ---                 │\n",
       "│ struct[2]            ┆ struct[2]            ┆ struct[2]           │\n",
       "╞══════════════════════╪══════════════════════╪═════════════════════╡\n",
       "│ {-1.340382,0.180152} ┆ {38.269766,0.366801} ┆ {1.836938,0.118797} │\n",
       "└──────────────────────┴──────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In dataframe statistical tests!\n",
    "df.select(\n",
    "    pds.query_ttest_ind(\"var1\", \"var2\", equal_var=True).alias(\"t-test\"),\n",
    "    pds.query_chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "    pds.query_f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "65dbb6bd",
   "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: (3, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>market_id</th><th>t-test</th><th>chi2-test</th><th>f-test</th></tr><tr><td>i64</td><td>struct[2]</td><td>struct[2]</td><td>struct[2]</td></tr></thead><tbody><tr><td>0</td><td>{-1.397243,0.162434}</td><td>{33.053173,0.609504}</td><td>{1.077703,0.365947}</td></tr><tr><td>1</td><td>{-0.953147,0.340585}</td><td>{32.40336,0.6404}</td><td>{0.637505,0.635777}</td></tr><tr><td>2</td><td>{0.047554,0.962075}</td><td>{27.160966,0.855728}</td><td>{1.245919,0.289437}</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (3, 4)\n",
       "┌───────────┬──────────────────────┬──────────────────────┬─────────────────────┐\n",
       "│ market_id ┆ t-test               ┆ chi2-test            ┆ f-test              │\n",
       "│ ---       ┆ ---                  ┆ ---                  ┆ ---                 │\n",
       "│ i64       ┆ struct[2]            ┆ struct[2]            ┆ struct[2]           │\n",
       "╞═══════════╪══════════════════════╪══════════════════════╪═════════════════════╡\n",
       "│ 0         ┆ {-1.397243,0.162434} ┆ {33.053173,0.609504} ┆ {1.077703,0.365947} │\n",
       "│ 1         ┆ {-0.953147,0.340585} ┆ {32.40336,0.6404}    ┆ {0.637505,0.635777} │\n",
       "│ 2         ┆ {0.047554,0.962075}  ┆ {27.160966,0.855728} ┆ {1.245919,0.289437} │\n",
       "└───────────┴──────────────────────┴──────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Can also be done in group by context\n",
    "df.group_by(\"market_id\").agg(\n",
    "    pds.query_ttest_ind(\"var1\", \"var2\", equal_var=False).alias(\"t-test\"),\n",
    "    pds.query_chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "    pds.query_f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "843d54c5",
   "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: (9, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>first_digit_cnt</th><th>first_digit_distribution</th></tr><tr><td>u32</td><td>f64</td></tr></thead><tbody><tr><td>568</td><td>0.1136</td></tr><tr><td>539</td><td>0.1078</td></tr><tr><td>561</td><td>0.1122</td></tr><tr><td>574</td><td>0.1148</td></tr><tr><td>552</td><td>0.1104</td></tr><tr><td>558</td><td>0.1116</td></tr><tr><td>569</td><td>0.1138</td></tr><tr><td>515</td><td>0.103</td></tr><tr><td>564</td><td>0.1128</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (9, 2)\n",
       "┌─────────────────┬──────────────────────────┐\n",
       "│ first_digit_cnt ┆ first_digit_distribution │\n",
       "│ ---             ┆ ---                      │\n",
       "│ u32             ┆ f64                      │\n",
       "╞═════════════════╪══════════════════════════╡\n",
       "│ 568             ┆ 0.1136                   │\n",
       "│ 539             ┆ 0.1078                   │\n",
       "│ 561             ┆ 0.1122                   │\n",
       "│ 574             ┆ 0.1148                   │\n",
       "│ 552             ┆ 0.1104                   │\n",
       "│ 558             ┆ 0.1116                   │\n",
       "│ 569             ┆ 0.1138                   │\n",
       "│ 515             ┆ 0.103                    │\n",
       "│ 564             ┆ 0.1128                   │\n",
       "└─────────────────┴──────────────────────────┘"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Benford's law\n",
    "df.select(\n",
    "    first_digit_cnt = pds.query_first_digit_cnt(pl.col(\"var1\")).explode()\n",
    ").with_columns(\n",
    "    # This doesn't follow benford's law because it is random data\n",
    "    first_digit_distribution = pl.col(\"first_digit_cnt\") / pl.col(\"first_digit_cnt\").sum()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b232e4d7",
   "metadata": {},
   "source": [
    "# Nearest Neighbors Related Tasks\n",
    "\n",
    "These queries can be very slow when data/dimension gets huge, even when processed in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "87aff1ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars_ds as pds\n",
    "size = 2000\n",
    "df = pl.DataFrame({\n",
    "    \"id\": range(size), \n",
    "}).with_columns(\n",
    "    pds.random().alias(\"var1\"),\n",
    "    pds.random().alias(\"var2\"),\n",
    "    pds.random().alias(\"var3\"),\n",
    "    pds.random().alias(\"r\"),\n",
    "    (pds.random() * 10).alias(\"rh\"),\n",
    "    pl.col(\"id\").cast(pl.UInt32)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "2fae4b5e",
   "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, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>nb_l_inf_cnt</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>0.549494</td><td>0.078403</td><td>0.437876</td><td>0.804484</td><td>1.048376</td><td>14</td></tr><tr><td>1</td><td>0.355507</td><td>0.852117</td><td>0.871496</td><td>0.546271</td><td>5.079249</td><td>15</td></tr><tr><td>2</td><td>0.069114</td><td>0.415601</td><td>0.355308</td><td>0.307507</td><td>3.284447</td><td>14</td></tr><tr><td>3</td><td>0.349286</td><td>0.229361</td><td>0.395056</td><td>0.132254</td><td>1.135489</td><td>21</td></tr><tr><td>4</td><td>0.607897</td><td>0.74279</td><td>0.063657</td><td>0.925873</td><td>8.774421</td><td>6</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ nb_l_inf_cnt │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---          │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ u32          │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════════╡\n",
       "│ 0   ┆ 0.549494 ┆ 0.078403 ┆ 0.437876 ┆ 0.804484 ┆ 1.048376 ┆ 14           │\n",
       "│ 1   ┆ 0.355507 ┆ 0.852117 ┆ 0.871496 ┆ 0.546271 ┆ 5.079249 ┆ 15           │\n",
       "│ 2   ┆ 0.069114 ┆ 0.415601 ┆ 0.355308 ┆ 0.307507 ┆ 3.284447 ┆ 14           │\n",
       "│ 3   ┆ 0.349286 ┆ 0.229361 ┆ 0.395056 ┆ 0.132254 ┆ 1.135489 ┆ 21           │\n",
       "│ 4   ┆ 0.607897 ┆ 0.74279  ┆ 0.063657 ┆ 0.925873 ┆ 8.774421 ┆ 6            │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────┘"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get neighbor count. The point itself is always considered a neighbor to itself.\n",
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        0.1, # radius \n",
    "        pl.col(\"var1\"), \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        dist = \"inf\", # L Infinity distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l_inf_cnt\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "69ad83d1",
   "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, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>nb_l1_r_cnt</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>0.549494</td><td>0.078403</td><td>0.437876</td><td>0.804484</td><td>1.048376</td><td>794</td></tr><tr><td>1</td><td>0.355507</td><td>0.852117</td><td>0.871496</td><td>0.546271</td><td>5.079249</td><td>222</td></tr><tr><td>2</td><td>0.069114</td><td>0.415601</td><td>0.355308</td><td>0.307507</td><td>3.284447</td><td>62</td></tr><tr><td>3</td><td>0.349286</td><td>0.229361</td><td>0.395056</td><td>0.132254</td><td>1.135489</td><td>6</td></tr><tr><td>4</td><td>0.607897</td><td>0.74279</td><td>0.063657</td><td>0.925873</td><td>8.774421</td><td>846</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬─────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ nb_l1_r_cnt │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---         │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ u32         │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪═════════════╡\n",
       "│ 0   ┆ 0.549494 ┆ 0.078403 ┆ 0.437876 ┆ 0.804484 ┆ 1.048376 ┆ 794         │\n",
       "│ 1   ┆ 0.355507 ┆ 0.852117 ┆ 0.871496 ┆ 0.546271 ┆ 5.079249 ┆ 222         │\n",
       "│ 2   ┆ 0.069114 ┆ 0.415601 ┆ 0.355308 ┆ 0.307507 ┆ 3.284447 ┆ 62          │\n",
       "│ 3   ┆ 0.349286 ┆ 0.229361 ┆ 0.395056 ┆ 0.132254 ┆ 1.135489 ┆ 6           │\n",
       "│ 4   ┆ 0.607897 ┆ 0.74279  ┆ 0.063657 ┆ 0.925873 ┆ 8.774421 ┆ 846         │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────────┘"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        pl.col(\"r\"), # radius be an expression too\n",
    "        \"var1\", \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        dist = \"l1\", # L 1 distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l1_r_cnt\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ce1a2c7a",
   "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, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>best friends</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>list[u32]</td></tr></thead><tbody><tr><td>0</td><td>0.549494</td><td>0.078403</td><td>0.437876</td><td>0.804484</td><td>1.048376</td><td>[0, 1877, … 607]</td></tr><tr><td>1</td><td>0.355507</td><td>0.852117</td><td>0.871496</td><td>0.546271</td><td>5.079249</td><td>[1, 1331, … 766]</td></tr><tr><td>2</td><td>0.069114</td><td>0.415601</td><td>0.355308</td><td>0.307507</td><td>3.284447</td><td>[2, 603, … 972]</td></tr><tr><td>3</td><td>0.349286</td><td>0.229361</td><td>0.395056</td><td>0.132254</td><td>1.135489</td><td>[3, 112, … 1543]</td></tr><tr><td>4</td><td>0.607897</td><td>0.74279</td><td>0.063657</td><td>0.925873</td><td>8.774421</td><td>[4, 571, … 1612]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ best friends     │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---              │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ list[u32]        │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════════════╡\n",
       "│ 0   ┆ 0.549494 ┆ 0.078403 ┆ 0.437876 ┆ 0.804484 ┆ 1.048376 ┆ [0, 1877, … 607] │\n",
       "│ 1   ┆ 0.355507 ┆ 0.852117 ┆ 0.871496 ┆ 0.546271 ┆ 5.079249 ┆ [1, 1331, … 766] │\n",
       "│ 2   ┆ 0.069114 ┆ 0.415601 ┆ 0.355308 ┆ 0.307507 ┆ 3.284447 ┆ [2, 603, … 972]  │\n",
       "│ 3   ┆ 0.349286 ┆ 0.229361 ┆ 0.395056 ┆ 0.132254 ┆ 1.135489 ┆ [3, 112, … 1543] │\n",
       "│ 4   ┆ 0.607897 ┆ 0.74279  ┆ 0.063657 ┆ 0.925873 ┆ 8.774421 ┆ [4, 571, … 1612] │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┘"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get ids of the k nearest neighbors. \n",
    "# The point itself is always considered a neighbor to itself, so k + 1 elements will be returned.\n",
    "df.with_columns(\n",
    "    pds.query_knn_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        index = \"id\",  # pl.col(\"id\"), str | pl.Expr\n",
    "        k = 3, \n",
    "        dist = \"l2\", # squared l2\n",
    "        parallel = True\n",
    "    ).alias(\"best friends\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "67a769f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 3)\n",
      "┌─────┬───────────────────┬────────────────────┐\n",
      "│ id  ┆ best friends      ┆ best friends count │\n",
      "│ --- ┆ ---               ┆ ---                │\n",
      "│ u32 ┆ list[u32]         ┆ u32                │\n",
      "╞═════╪═══════════════════╪════════════════════╡\n",
      "│ 0   ┆ [0, 1877, … 1574] ┆ 213                │\n",
      "│ 1   ┆ [1, 1331, … 437]  ┆ 151                │\n",
      "│ 2   ┆ [2, 603, … 262]   ┆ 169                │\n",
      "│ 3   ┆ [3, 112, … 1259]  ┆ 254                │\n",
      "│ 4   ┆ [4, 571, … 474]   ┆ 167                │\n",
      "└─────┴───────────────────┴────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Get all neighbors within radius r\n",
    "# The point itself is always considered a neighbor to itself.\n",
    "print(df.select(\n",
    "    pl.col(\"id\"),\n",
    "    pds.query_radius_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        index = pl.col(\"id\"),\n",
    "        r = 0.1, \n",
    "        dist = \"l2\", # actually this is squared l2\n",
    "        parallel = True\n",
    "    ).alias(\"best friends\"),\n",
    ").with_columns( # -1 to remove the point itself\n",
    "    (pl.col(\"best friends\").list.len() - 1).alias(\"best friends count\")\n",
    ").head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "06cd7fb2",
   "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, 8)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>idx</th><th>dist</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>list[u32]</td><td>list[f64]</td></tr></thead><tbody><tr><td>0</td><td>0.549494</td><td>0.078403</td><td>0.437876</td><td>0.804484</td><td>1.048376</td><td>[0, 1877, … 607]</td><td>[0.0, 0.000535, … 0.004595]</td></tr><tr><td>1</td><td>0.355507</td><td>0.852117</td><td>0.871496</td><td>0.546271</td><td>5.079249</td><td>[1, 1331, … 766]</td><td>[0.0, 0.001853, … 0.004976]</td></tr><tr><td>2</td><td>0.069114</td><td>0.415601</td><td>0.355308</td><td>0.307507</td><td>3.284447</td><td>[2, 603, … 972]</td><td>[0.0, 0.000692, … 0.006261]</td></tr><tr><td>3</td><td>0.349286</td><td>0.229361</td><td>0.395056</td><td>0.132254</td><td>1.135489</td><td>[3, 112, … 1543]</td><td>[0.0, 0.000113, … 0.001303]</td></tr><tr><td>4</td><td>0.607897</td><td>0.74279</td><td>0.063657</td><td>0.925873</td><td>8.774421</td><td>[4, 571, … 1612]</td><td>[0.0, 0.005093, … 0.010396]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 8)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────────────┬──────────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ idx              ┆ dist             │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---              ┆ ---              │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ list[u32]        ┆ list[f64]        │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════════════╪══════════════════╡\n",
       "│ 0   ┆ 0.549494 ┆ 0.078403 ┆ 0.437876 ┆ 0.804484 ┆ 1.048376 ┆ [0, 1877, … 607] ┆ [0.0, 0.000535,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.004595]      │\n",
       "│ 1   ┆ 0.355507 ┆ 0.852117 ┆ 0.871496 ┆ 0.546271 ┆ 5.079249 ┆ [1, 1331, … 766] ┆ [0.0, 0.001853,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.004976]      │\n",
       "│ 2   ┆ 0.069114 ┆ 0.415601 ┆ 0.355308 ┆ 0.307507 ┆ 3.284447 ┆ [2, 603, … 972]  ┆ [0.0, 0.000692,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.006261]      │\n",
       "│ 3   ┆ 0.349286 ┆ 0.229361 ┆ 0.395056 ┆ 0.132254 ┆ 1.135489 ┆ [3, 112, … 1543] ┆ [0.0, 0.000113,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.001303]      │\n",
       "│ 4   ┆ 0.607897 ┆ 0.74279  ┆ 0.063657 ┆ 0.925873 ┆ 8.774421 ┆ [4, 571, … 1612] ┆ [0.0, 0.005093,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.010396]      │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┴──────────────────┘"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get ids of the k nearest neighbors and distances\n",
    "# The point itself is always considered a neighbor to itself, so k + 1 elements will be returned.\n",
    "df.with_columns(\n",
    "    pds.query_knn_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        index = pl.col(\"id\"),\n",
    "        k = 3, \n",
    "        dist = \"l2\", # actually this is squared l2\n",
    "        parallel = True,\n",
    "        return_dist = True\n",
    "    ).alias(\"best_friends_w_dist\")\n",
    ").unnest(\"best_friends_w_dist\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "b5c69ae4",
   "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>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0</td><td>0.549494</td><td>0.078403</td><td>0.437876</td><td>0.804484</td><td>1.048376</td></tr><tr><td>3</td><td>0.349286</td><td>0.229361</td><td>0.395056</td><td>0.132254</td><td>1.135489</td></tr><tr><td>12</td><td>0.661558</td><td>0.734605</td><td>0.259757</td><td>0.253638</td><td>9.009864</td></tr><tr><td>14</td><td>0.146869</td><td>0.395578</td><td>0.465677</td><td>0.853553</td><td>3.554655</td></tr><tr><td>15</td><td>0.155285</td><td>0.497808</td><td>0.654169</td><td>0.707903</td><td>6.872508</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 0   ┆ 0.549494 ┆ 0.078403 ┆ 0.437876 ┆ 0.804484 ┆ 1.048376 │\n",
       "│ 3   ┆ 0.349286 ┆ 0.229361 ┆ 0.395056 ┆ 0.132254 ┆ 1.135489 │\n",
       "│ 12  ┆ 0.661558 ┆ 0.734605 ┆ 0.259757 ┆ 0.253638 ┆ 9.009864 │\n",
       "│ 14  ┆ 0.146869 ┆ 0.395578 ┆ 0.465677 ┆ 0.853553 ┆ 3.554655 │\n",
       "│ 15  ┆ 0.155285 ┆ 0.497808 ┆ 0.654169 ┆ 0.707903 ┆ 6.872508 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Filter to only points near the given point\n",
    "df.filter(\n",
    "    pds.query_within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        pt = [0.5, 0.5, 0.5],\n",
    "        r = 0.2,\n",
    "        dist = \"l2\" # actually this is squared l2, so this is asking for squared l2 <= 0.2\n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "4ab9e8f8",
   "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>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>10</td><td>0.563265</td><td>0.511338</td><td>0.969821</td><td>0.498458</td><td>6.901672</td></tr><tr><td>58</td><td>0.453018</td><td>0.430267</td><td>0.796741</td><td>0.337218</td><td>0.84054</td></tr><tr><td>122</td><td>0.499126</td><td>0.559724</td><td>0.898422</td><td>0.244714</td><td>7.379702</td></tr><tr><td>148</td><td>0.557274</td><td>0.494576</td><td>0.159659</td><td>0.875239</td><td>0.788142</td></tr><tr><td>159</td><td>0.514593</td><td>0.414222</td><td>0.158342</td><td>0.092265</td><td>6.599698</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 10  ┆ 0.563265 ┆ 0.511338 ┆ 0.969821 ┆ 0.498458 ┆ 6.901672 │\n",
       "│ 58  ┆ 0.453018 ┆ 0.430267 ┆ 0.796741 ┆ 0.337218 ┆ 0.84054  │\n",
       "│ 122 ┆ 0.499126 ┆ 0.559724 ┆ 0.898422 ┆ 0.244714 ┆ 7.379702 │\n",
       "│ 148 ┆ 0.557274 ┆ 0.494576 ┆ 0.159659 ┆ 0.875239 ┆ 0.788142 │\n",
       "│ 159 ┆ 0.514593 ┆ 0.414222 ┆ 0.158342 ┆ 0.092265 ┆ 6.599698 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Haversine distance is available when dimension is 2\n",
    "df.filter(\n",
    "    pds.query_within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), # Columns used as the coordinates in n-d space\n",
    "        pt = [0.5, 0.5],\n",
    "        r = 10, # in km\n",
    "        dist = \"h\" \n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "7d3f5ae4",
   "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>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>122</td><td>0.499126</td><td>0.559724</td><td>0.898422</td><td>0.244714</td><td>7.379702</td></tr><tr><td>185</td><td>0.485795</td><td>0.486589</td><td>0.046022</td><td>0.706548</td><td>2.1732</td></tr><tr><td>194</td><td>0.48362</td><td>0.474265</td><td>0.741543</td><td>0.090014</td><td>7.930565</td></tr><tr><td>256</td><td>0.474407</td><td>0.531546</td><td>0.189203</td><td>0.775875</td><td>6.445429</td></tr><tr><td>362</td><td>0.55823</td><td>0.468139</td><td>0.521009</td><td>0.922885</td><td>8.851264</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 122 ┆ 0.499126 ┆ 0.559724 ┆ 0.898422 ┆ 0.244714 ┆ 7.379702 │\n",
       "│ 185 ┆ 0.485795 ┆ 0.486589 ┆ 0.046022 ┆ 0.706548 ┆ 2.1732   │\n",
       "│ 194 ┆ 0.48362  ┆ 0.474265 ┆ 0.741543 ┆ 0.090014 ┆ 7.930565 │\n",
       "│ 256 ┆ 0.474407 ┆ 0.531546 ┆ 0.189203 ┆ 0.775875 ┆ 6.445429 │\n",
       "│ 362 ┆ 0.55823  ┆ 0.468139 ┆ 0.521009 ┆ 0.922885 ┆ 8.851264 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    pds.query_within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), \n",
    "        pt = [0.5, 0.5],\n",
    "        # radius can also be an existing column in the dataframe.\n",
    "        r = pl.col(\"rh\"), \n",
    "        dist = \"h\" \n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "f14627bc",
   "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>id</th><th>friends</th><th>count</th></tr><tr><td>u64</td><td>list[u32]</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>[0, 308, … 599]</td><td>123</td></tr><tr><td>1</td><td>[1, 966, … 1470]</td><td>115</td></tr><tr><td>2</td><td>[2, 1987, … 75]</td><td>97</td></tr><tr><td>3</td><td>[3, 112, … 203]</td><td>126</td></tr><tr><td>4</td><td>[4, 678, … 922]</td><td>109</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌─────┬──────────────────┬───────┐\n",
       "│ id  ┆ friends          ┆ count │\n",
       "│ --- ┆ ---              ┆ ---   │\n",
       "│ u64 ┆ list[u32]        ┆ u32   │\n",
       "╞═════╪══════════════════╪═══════╡\n",
       "│ 0   ┆ [0, 308, … 599]  ┆ 123   │\n",
       "│ 1   ┆ [1, 966, … 1470] ┆ 115   │\n",
       "│ 2   ┆ [2, 1987, … 75]  ┆ 97    │\n",
       "│ 3   ┆ [3, 112, … 203]  ┆ 126   │\n",
       "│ 4   ┆ [4, 678, … 922]  ┆ 109   │\n",
       "└─────┴──────────────────┴───────┘"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "friends = df.select(\n",
    "    pl.col(\"id\").cast(pl.UInt64),\n",
    "    pds.query_radius_ptwise(\n",
    "        # Columns used as the coordinates in n-d space\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), \n",
    "        index=pl.col(\"id\"),\n",
    "        r = 0.02, \n",
    "        dist = \"l2\",\n",
    "    ).alias(\"friends\")\n",
    ").with_columns(\n",
    "    pl.col(\"friends\").list.len().alias(\"count\")\n",
    ")\n",
    "friends.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea3e4400",
   "metadata": {},
   "source": [
    "# Simple Graph Queries\n",
    "\n",
    "There is limited functionality in the Graph module currently. E.g. Only constant cost per edge.\n",
    "\n",
    "Graph queries are very expensive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "fba0149e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# friends.select(\n",
    "#     pl.col(\"friends\").graph.eigen_centrality() # .arg_max()\n",
    "# ).head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "9d2b70d7",
   "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>id</th><th>friends</th></tr><tr><td>u32</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>0</td></tr><tr><td>0</td><td>308</td></tr><tr><td>0</td><td>1877</td></tr><tr><td>0</td><td>937</td></tr><tr><td>0</td><td>1635</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────┬─────────┐\n",
       "│ id  ┆ friends │\n",
       "│ --- ┆ ---     │\n",
       "│ u32 ┆ u32     │\n",
       "╞═════╪═════════╡\n",
       "│ 0   ┆ 0       │\n",
       "│ 0   ┆ 308     │\n",
       "│ 0   ┆ 1877    │\n",
       "│ 0   ┆ 937     │\n",
       "│ 0   ┆ 1635    │\n",
       "└─────┴─────────┘"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Turn friends to a table suitable for graph analytics\n",
    "df_graph = friends.select(\n",
    "    pl.col(\"id\"),\n",
    "    pl.col(\"friends\"),\n",
    ").explode(pl.col(\"friends\")).with_columns(\n",
    "    pl.col(\"id\").cast(pl.UInt32),\n",
    "    pl.col(\"friends\").cast(pl.UInt32),\n",
    ")\n",
    "df_graph.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "8ff9aa74",
   "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_000, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>path</th></tr><tr><td>u32</td><td>list[u32]</td></tr></thead><tbody><tr><td>0</td><td>[599, 3]</td></tr><tr><td>1</td><td>[275, 1521, … 3]</td></tr><tr><td>2</td><td>[1187, 1057, 3]</td></tr><tr><td>3</td><td>[]</td></tr><tr><td>4</td><td>[1804, 1573, … 3]</td></tr><tr><td>&hellip;</td><td>&hellip;</td></tr><tr><td>1995</td><td>[3]</td></tr><tr><td>1996</td><td>[170, 1157, … 3]</td></tr><tr><td>1997</td><td>[437, 1612, … 3]</td></tr><tr><td>1998</td><td>[305, 3]</td></tr><tr><td>1999</td><td>[1065, 737, 3]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2_000, 2)\n",
       "┌──────┬───────────────────┐\n",
       "│ id   ┆ path              │\n",
       "│ ---  ┆ ---               │\n",
       "│ u32  ┆ list[u32]         │\n",
       "╞══════╪═══════════════════╡\n",
       "│ 0    ┆ [599, 3]          │\n",
       "│ 1    ┆ [275, 1521, … 3]  │\n",
       "│ 2    ┆ [1187, 1057, 3]   │\n",
       "│ 3    ┆ []                │\n",
       "│ 4    ┆ [1804, 1573, … 3] │\n",
       "│ …    ┆ …                 │\n",
       "│ 1995 ┆ [3]               │\n",
       "│ 1996 ┆ [170, 1157, … 3]  │\n",
       "│ 1997 ┆ [437, 1612, … 3]  │\n",
       "│ 1998 ┆ [305, 3]          │\n",
       "│ 1999 ┆ [1065, 737, 3]    │\n",
       "└──────┴───────────────────┘"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_graph.select(\n",
    "    # Shortest path to the node with id = 3\n",
    "    # Node and link can be str | pl.Expr\n",
    "    pds.query_shortest_path(node = \"id\", link = pl.col(\"friends\"), target = 3, cost = None, parallel=True).alias(\"shortest_path\")\n",
    ").unnest(\"shortest_path\").sort(\"id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "76d43fd5",
   "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_000, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>reachable</th><th>steps</th></tr><tr><td>u32</td><td>bool</td><td>u32</td></tr></thead><tbody><tr><td>1492</td><td>true</td><td>4</td></tr><tr><td>542</td><td>true</td><td>2</td></tr><tr><td>1382</td><td>true</td><td>1</td></tr><tr><td>1954</td><td>true</td><td>3</td></tr><tr><td>824</td><td>true</td><td>5</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>1149</td><td>true</td><td>2</td></tr><tr><td>1656</td><td>true</td><td>7</td></tr><tr><td>1881</td><td>true</td><td>4</td></tr><tr><td>789</td><td>true</td><td>5</td></tr><tr><td>1910</td><td>true</td><td>5</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2_000, 3)\n",
       "┌──────┬───────────┬───────┐\n",
       "│ id   ┆ reachable ┆ steps │\n",
       "│ ---  ┆ ---       ┆ ---   │\n",
       "│ u32  ┆ bool      ┆ u32   │\n",
       "╞══════╪═══════════╪═══════╡\n",
       "│ 1492 ┆ true      ┆ 4     │\n",
       "│ 542  ┆ true      ┆ 2     │\n",
       "│ 1382 ┆ true      ┆ 1     │\n",
       "│ 1954 ┆ true      ┆ 3     │\n",
       "│ 824  ┆ true      ┆ 5     │\n",
       "│ …    ┆ …         ┆ …     │\n",
       "│ 1149 ┆ true      ┆ 2     │\n",
       "│ 1656 ┆ true      ┆ 7     │\n",
       "│ 1881 ┆ true      ┆ 4     │\n",
       "│ 789  ┆ true      ┆ 5     │\n",
       "│ 1910 ┆ true      ┆ 5     │\n",
       "└──────┴───────────┴───────┘"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_graph.select(\n",
    "    # Almost every node can reach node 3, and the number is the number steps to reach it\n",
    "    # This is a way faster way to filter results if you don't need the actual path\n",
    "    pds.query_node_reachable(\"id\", \"friends\", target = 3).alias(\"reach\")\n",
    ").unnest(\"reach\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "4af6e6b7",
   "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: (11, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>connections</th><th>close-ness</th></tr><tr><td>u32</td><td>u32</td><td>f64</td></tr></thead><tbody><tr><td>0</td><td>1</td><td>0.4</td></tr><tr><td>0</td><td>2</td><td>0.3</td></tr><tr><td>0</td><td>3</td><td>0.2</td></tr><tr><td>0</td><td>4</td><td>0.1</td></tr><tr><td>1</td><td>2</td><td>0.1</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>2</td><td>4</td><td>0.5</td></tr><tr><td>3</td><td>0</td><td>0.1</td></tr><tr><td>3</td><td>1</td><td>0.1</td></tr><tr><td>3</td><td>2</td><td>0.1</td></tr><tr><td>4</td><td>1</td><td>0.1</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (11, 3)\n",
       "┌─────┬─────────────┬────────────┐\n",
       "│ id  ┆ connections ┆ close-ness │\n",
       "│ --- ┆ ---         ┆ ---        │\n",
       "│ u32 ┆ u32         ┆ f64        │\n",
       "╞═════╪═════════════╪════════════╡\n",
       "│ 0   ┆ 1           ┆ 0.4        │\n",
       "│ 0   ┆ 2           ┆ 0.3        │\n",
       "│ 0   ┆ 3           ┆ 0.2        │\n",
       "│ 0   ┆ 4           ┆ 0.1        │\n",
       "│ 1   ┆ 2           ┆ 0.1        │\n",
       "│ …   ┆ …           ┆ …          │\n",
       "│ 2   ┆ 4           ┆ 0.5        │\n",
       "│ 3   ┆ 0           ┆ 0.1        │\n",
       "│ 3   ┆ 1           ┆ 0.1        │\n",
       "│ 3   ┆ 2           ┆ 0.1        │\n",
       "│ 4   ┆ 1           ┆ 0.1        │\n",
       "└─────┴─────────────┴────────────┘"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relationships = pl.DataFrame({\n",
    "    \"id\": range(5),\n",
    "    \"connections\":[[1,2,3,4], [2,3], [4], [0,1,2], [1]],\n",
    "    # Small values means closer\n",
    "    \"close-ness\":[[0.4, 0.3, 0.2, 0.1], [0.1, 1.0], [0.5], [0.1, 0.1, 0.1], [0.1]]\n",
    "}).with_columns(\n",
    "    pl.col(\"id\").cast(pl.UInt32),\n",
    "    pl.col(\"connections\").list.eval(pl.element().cast(pl.UInt32))\n",
    ").explode(\n",
    "    pl.col(\"connections\"), pl.col(\"close-ness\")\n",
    ")\n",
    "\n",
    "relationships.head(50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "239f4e85",
   "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>id</th><th>path</th><th>cost</th></tr><tr><td>u32</td><td>list[u32]</td><td>f64</td></tr></thead><tbody><tr><td>2</td><td>[4, 1]</td><td>0.6</td></tr><tr><td>3</td><td>[1]</td><td>0.1</td></tr><tr><td>0</td><td>[4, 1]</td><td>0.2</td></tr><tr><td>4</td><td>[1]</td><td>0.1</td></tr><tr><td>1</td><td>[]</td><td>0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌─────┬───────────┬──────┐\n",
       "│ id  ┆ path      ┆ cost │\n",
       "│ --- ┆ ---       ┆ ---  │\n",
       "│ u32 ┆ list[u32] ┆ f64  │\n",
       "╞═════╪═══════════╪══════╡\n",
       "│ 2   ┆ [4, 1]    ┆ 0.6  │\n",
       "│ 3   ┆ [1]       ┆ 0.1  │\n",
       "│ 0   ┆ [4, 1]    ┆ 0.2  │\n",
       "│ 4   ┆ [1]       ┆ 0.1  │\n",
       "│ 1   ┆ []        ┆ 0.0  │\n",
       "└─────┴───────────┴──────┘"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# To go to node at id = 1, node 0 would rather go to 4 first and then 1.\n",
    "relationships.select(\n",
    "    pds.query_shortest_path(\"id\", \"connections\", target = 1, cost = \"close-ness\").alias(\"path\")\n",
    ").unnest(\"path\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "df458bdd",
   "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>node</th><th>deg</th></tr><tr><td>u32</td><td>u32</td></tr></thead><tbody><tr><td>2</td><td>1</td></tr><tr><td>3</td><td>3</td></tr><tr><td>0</td><td>4</td></tr><tr><td>4</td><td>1</td></tr><tr><td>1</td><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌──────┬─────┐\n",
       "│ node ┆ deg │\n",
       "│ ---  ┆ --- │\n",
       "│ u32  ┆ u32 │\n",
       "╞══════╪═════╡\n",
       "│ 2    ┆ 1   │\n",
       "│ 3    ┆ 3   │\n",
       "│ 0    ┆ 4   │\n",
       "│ 4    ┆ 1   │\n",
       "│ 1    ┆ 2   │\n",
       "└──────┴─────┘"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In and out deg\n",
    "relationships.select(\n",
    "    pds.query_node_deg(\"id\", \"connections\", directed=True).alias(\"deg\")\n",
    ").unnest(\"deg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f8eabf1",
   "metadata": {},
   "source": [
    "# String Nearest Neighbors\n",
    "\n",
    "This might be very slow for very large vocab / column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "f2c9ed80",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame({\n",
    "    \"a\":[\"AAAAA\", \"ABCABC\", \"AAAADDD\", \"ADSDSDS\", \"WORD\"],\n",
    "    \"b\":[\"AAAAT\", \"ABCACD\", \"ADSSD\", \"APPLES\", \"WORLD\"] \n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "062fd694",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;AAAAT&quot;</td></tr><tr><td>&quot;ABCACD&quot;</td></tr><tr><td>&quot;AAAAT&quot;</td></tr><tr><td>&quot;ADSSD&quot;</td></tr><tr><td>&quot;WORLD&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ str                      │\n",
       "╞══════════════════════════╡\n",
       "│ AAAAT                    │\n",
       "│ ABCACD                   │\n",
       "│ AAAAT                    │\n",
       "│ ADSSD                    │\n",
       "│ WORLD                    │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use Levenshtein to find the nearest neighbor in vocab to word in column a\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 1, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "29437b8b",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>list[str]</td></tr></thead><tbody><tr><td>[&quot;AAAAT&quot;, &quot;ADSSD&quot;]</td></tr><tr><td>[&quot;ABCACD&quot;, &quot;AAAAT&quot;]</td></tr><tr><td>[&quot;AAAAT&quot;, &quot;ABCACD&quot;]</td></tr><tr><td>[&quot;ADSSD&quot;, &quot;APPLES&quot;]</td></tr><tr><td>[&quot;WORLD&quot;, &quot;ADSSD&quot;]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ list[str]                │\n",
       "╞══════════════════════════╡\n",
       "│ [\"AAAAT\", \"ADSSD\"]       │\n",
       "│ [\"ABCACD\", \"AAAAT\"]      │\n",
       "│ [\"AAAAT\", \"ABCACD\"]      │\n",
       "│ [\"ADSSD\", \"APPLES\"]      │\n",
       "│ [\"WORLD\", \"ADSSD\"]       │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use Levenshtein to find 2 nearest neighbors\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 2, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "9150bece",
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>list[str]</td></tr></thead><tbody><tr><td>[&quot;AAAAT&quot;, &quot;ADSSD&quot;]</td></tr><tr><td>[&quot;ABCACD&quot;]</td></tr><tr><td>[]</td></tr><tr><td>[]</td></tr><tr><td>[]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ list[str]                │\n",
       "╞══════════════════════════╡\n",
       "│ [\"AAAAT\", \"ADSSD\"]       │\n",
       "│ [\"ABCACD\"]               │\n",
       "│ []                       │\n",
       "│ []                       │\n",
       "│ []                       │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Currently only Levenshtein and hamming are implemented for this\n",
    "# Empty means nothing in vocab can be compared in the hamming sense with the corresponding word in a\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 2, \n",
    "        threshold = 4,\n",
    "        metric = \"hamming\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "b35776a1",
   "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>a</th><th>similar_words_from_vocab</th></tr><tr><td>str</td><td>list[str]</td></tr></thead><tbody><tr><td>&quot;AAAAA&quot;</td><td>[&quot;AAAAA&quot;, &quot;ZIV&quot;, &quot;WORLD&quot;]</td></tr><tr><td>&quot;ABCABC&quot;</td><td>[&quot;ABCDEFG&quot;, &quot;AAAAA&quot;, &quot;ZIV&quot;]</td></tr><tr><td>&quot;AAAADDD&quot;</td><td>[&quot;AAAAA&quot;, &quot;WORLD&quot;, &quot;ABCDEFG&quot;]</td></tr><tr><td>&quot;ADSDSDS&quot;</td><td>[&quot;ABCDEFG&quot;, &quot;WORLD&quot;, &quot;AAAAA&quot;]</td></tr><tr><td>&quot;WORD&quot;</td><td>[&quot;WORLD&quot;, &quot;ZIV&quot;, &quot;TQQQ&quot;]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────┬───────────────────────────────┐\n",
       "│ a       ┆ similar_words_from_vocab      │\n",
       "│ ---     ┆ ---                           │\n",
       "│ str     ┆ list[str]                     │\n",
       "╞═════════╪═══════════════════════════════╡\n",
       "│ AAAAA   ┆ [\"AAAAA\", \"ZIV\", \"WORLD\"]     │\n",
       "│ ABCABC  ┆ [\"ABCDEFG\", \"AAAAA\", \"ZIV\"]   │\n",
       "│ AAAADDD ┆ [\"AAAAA\", \"WORLD\", \"ABCDEFG\"] │\n",
       "│ ADSDSDS ┆ [\"ABCDEFG\", \"WORLD\", \"AAAAA\"] │\n",
       "│ WORD    ┆ [\"WORLD\", \"ZIV\", \"TQQQ\"]      │\n",
       "└─────────┴───────────────────────────────┘"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# You may provide a vocab like this\n",
    "df.select(\n",
    "    pl.col(\"a\"),\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = [\"WORLD\", \"AAAAA\", \"ABCDEFG\", \"ZIV\", \"TQQQ\"],\n",
    "        k = 3, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a42a771",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}