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, 13)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th><th>time_idx</th><th>dummy</th><th>actual</th><th>predicted</th><th>dummy_groups</th><th>x1</th><th>x2</th><th>x3</th><th>a</th><th>b</th><th>y</th><th>y2</th></tr><tr><td>f64</td><td>i64</td><td>str</td><td>i32</td><td>f64</td><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>0.0</td><td>0</td><td>&quot;a&quot;</td><td>0</td><td>0.870124</td><td>&quot;a&quot;</td><td>0.499676</td><td>0.116938</td><td>0.587946</td><td>0.131978</td><td>0.324241</td><td>-0.771857</td><td>0.058872</td></tr><tr><td>0.841471</td><td>1</td><td>&quot;a&quot;</td><td>0</td><td>0.626794</td><td>&quot;a&quot;</td><td>0.791879</td><td>0.094017</td><td>0.374373</td><td>0.489786</td><td>0.504696</td><td>-0.414556</td><td>0.107824</td></tr><tr><td>0.909297</td><td>2</td><td>&quot;a&quot;</td><td>0</td><td>0.362253</td><td>&quot;a&quot;</td><td>0.526741</td><td>0.544468</td><td>0.649537</td><td>0.824565</td><td>0.414665</td><td>-0.73191</td><td>0.248602</td></tr><tr><td>0.14112</td><td>3</td><td>&quot;a&quot;</td><td>1</td><td>0.236681</td><td>&quot;a&quot;</td><td>0.196969</td><td>0.264493</td><td>0.736293</td><td>0.42734</td><td>0.734808</td><td>-0.995451</td><td>0.071098</td></tr><tr><td>-0.756802</td><td>4</td><td>&quot;a&quot;</td><td>0</td><td>0.147215</td><td>&quot;a&quot;</td><td>0.786499</td><td>0.964315</td><td>0.764042</td><td>0.788627</td><td>0.2992</td><td>-0.738782</td><td>0.459861</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 13)\n",
       "┌───────────┬──────────┬───────┬────────┬───┬──────────┬──────────┬───────────┬──────────┐\n",
       "│ f         ┆ time_idx ┆ dummy ┆ actual ┆ … ┆ a        ┆ b        ┆ y         ┆ y2       │\n",
       "│ ---       ┆ ---      ┆ ---   ┆ ---    ┆   ┆ ---      ┆ ---      ┆ ---       ┆ ---      │\n",
       "│ f64       ┆ i64      ┆ str   ┆ i32    ┆   ┆ f64      ┆ f64      ┆ f64       ┆ f64      │\n",
       "╞═══════════╪══════════╪═══════╪════════╪═══╪══════════╪══════════╪═══════════╪══════════╡\n",
       "│ 0.0       ┆ 0        ┆ a     ┆ 0      ┆ … ┆ 0.131978 ┆ 0.324241 ┆ -0.771857 ┆ 0.058872 │\n",
       "│ 0.841471  ┆ 1        ┆ a     ┆ 0      ┆ … ┆ 0.489786 ┆ 0.504696 ┆ -0.414556 ┆ 0.107824 │\n",
       "│ 0.909297  ┆ 2        ┆ a     ┆ 0      ┆ … ┆ 0.824565 ┆ 0.414665 ┆ -0.73191  ┆ 0.248602 │\n",
       "│ 0.14112   ┆ 3        ┆ a     ┆ 1      ┆ … ┆ 0.42734  ┆ 0.734808 ┆ -0.995451 ┆ 0.071098 │\n",
       "│ -0.756802 ┆ 4        ┆ a     ┆ 0      ┆ … ┆ 0.788627 ┆ 0.2992   ┆ -0.738782 ┆ 0.459861 │\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",
    "    , \"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",
    "}).with_columns(\n",
    "    pds.random(0., 1.).alias(\"x1\")\n",
    "    , pds.random(0., 1.).alias(\"x2\")\n",
    "    , pds.random(0., 1.).alias(\"x3\")\n",
    "    , pds.random(0., 1.).alias(\"a\")\n",
    "    , pds.random(0., 1.).alias(\"b\")\n",
    ").with_columns(\n",
    "    y = pl.col(\"x1\") * 0.15 + pl.col(\"x2\") * 0.3 - pl.col(\"x3\") * 1.5 + pds.random() * 0.0001,\n",
    "    y2 = pl.col(\"x1\") * 0.13 + pl.col(\"x2\") * 0.45 - pl.col(\"x3\") * 0.1 + pds.random() * 0.0001\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.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>1.3944e-15</td><td>-0.131978</td><td>-0.324241</td></tr><tr><td>-0.841471</td><td>-0.489786</td><td>-0.504696</td></tr><tr><td>-0.909297</td><td>-0.824565</td><td>-0.414665</td></tr><tr><td>-0.14112</td><td>-0.42734</td><td>-0.734808</td></tr><tr><td>0.756802</td><td>-0.656649</td><td>0.025041</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌────────────┬───────────┬───────────┐\n",
       "│ f          ┆ a         ┆ b         │\n",
       "│ ---        ┆ ---       ┆ ---       │\n",
       "│ f64        ┆ f64       ┆ f64       │\n",
       "╞════════════╪═══════════╪═══════════╡\n",
       "│ 1.3944e-15 ┆ -0.131978 ┆ -0.324241 │\n",
       "│ -0.841471  ┆ -0.489786 ┆ -0.504696 │\n",
       "│ -0.909297  ┆ -0.824565 ┆ -0.414665 │\n",
       "│ -0.14112   ┆ -0.42734  ┆ -0.734808 │\n",
       "│ 0.756802   ┆ -0.656649 ┆ 0.025041  │\n",
       "└────────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Multiple Convolutions at once\n",
    "# Modes: `same`, `left` (left-aligned same), `right` (right-aligned same), `valid` or `full`\n",
    "# Method: `fft`, `direct`\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\", method = \"fft\"), # column f with the kernel given here\n",
    "    pds.convolve(\"a\", [-1, 0, 0, 0, 1], mode = \"full\", method = \"direct\"),\n",
    "    pds.convolve(\"b\", [-1, 0, 0, 0, 1], mode = \"full\", method = \"direct\"),\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>coeffs</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.497998, -0.333503]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌────────────────────────┐\n",
       "│ coeffs                 │\n",
       "│ ---                    │\n",
       "│ list[f64]              │\n",
       "╞════════════════════════╡\n",
       "│ [-0.497998, -0.333503] │\n",
       "└────────────────────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Linear Regression\n",
    "df.select(\n",
    "    pds.lin_reg(\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": "e94324b7",
   "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, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>target_0</th><th>target_1</th></tr><tr><td>list[f64]</td><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.497998, -0.333503]</td><td>[0.08684, 0.407807]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 2)\n",
       "┌────────────────────────┬─────────────────────┐\n",
       "│ target_0               ┆ target_1            │\n",
       "│ ---                    ┆ ---                 │\n",
       "│ list[f64]              ┆ list[f64]           │\n",
       "╞════════════════════════╪═════════════════════╡\n",
       "│ [-0.497998, -0.333503] ┆ [0.08684, 0.407807] │\n",
       "└────────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Linear Regression, multi-target\n",
    "df.select(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = [pl.col(\"y\"), pl.col(\"y2\")],\n",
    "        add_bias=False\n",
    "    )\n",
    ").unnest(\"coeffs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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: (4, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>features</th><th>beta</th><th>std_err</th><th>t</th><th>p&gt;|t|</th><th>0.025</th><th>0.975</th></tr><tr><td>str</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;ln(x1+1)&quot;</td><td>0.219173</td><td>0.001673</td><td>131.006017</td><td>0.0</td><td>0.215893</td><td>0.222452</td></tr><tr><td>&quot;exp(x2)&quot;</td><td>0.175514</td><td>0.000678</td><td>258.836417</td><td>0.0</td><td>0.174185</td><td>0.176843</td></tr><tr><td>&quot;sin(x3)&quot;</td><td>-1.741122</td><td>0.001334</td><td>-1304.782909</td><td>0.0</td><td>-1.743738</td><td>-1.738507</td></tr><tr><td>&quot;__bias__&quot;</td><td>-0.110494</td><td>0.001485</td><td>-74.408752</td><td>0.0</td><td>-0.113404</td><td>-0.107583</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (4, 7)\n",
       "┌──────────┬───────────┬──────────┬──────────────┬───────┬───────────┬───────────┐\n",
       "│ features ┆ beta      ┆ std_err  ┆ t            ┆ p>|t| ┆ 0.025     ┆ 0.975     │\n",
       "│ ---      ┆ ---       ┆ ---      ┆ ---          ┆ ---   ┆ ---       ┆ ---       │\n",
       "│ str      ┆ f64       ┆ f64      ┆ f64          ┆ f64   ┆ f64       ┆ f64       │\n",
       "╞══════════╪═══════════╪══════════╪══════════════╪═══════╪═══════════╪═══════════╡\n",
       "│ ln(x1+1) ┆ 0.219173  ┆ 0.001673 ┆ 131.006017   ┆ 0.0   ┆ 0.215893  ┆ 0.222452  │\n",
       "│ exp(x2)  ┆ 0.175514  ┆ 0.000678 ┆ 258.836417   ┆ 0.0   ┆ 0.174185  ┆ 0.176843  │\n",
       "│ sin(x3)  ┆ -1.741122 ┆ 0.001334 ┆ -1304.782909 ┆ 0.0   ┆ -1.743738 ┆ -1.738507 │\n",
       "│ __bias__ ┆ -0.110494 ┆ 0.001485 ┆ -74.408752   ┆ 0.0   ┆ -0.113404 ┆ -0.107583 │\n",
       "└──────────┴───────────┴──────────┴──────────────┴───────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.lin_reg_report(\n",
    "        # formulaic input is also available for lstsq related queries, \n",
    "        # or you can always use polars expressions, e.g. pl.col('x1') + 1, pl.col('x2').exp(), pl.col('x3').sin()\n",
    "        \"ln(x1+1)\", \"exp(x2)\", \"sin(x3)\",\n",
    "        target = \"y\",\n",
    "        add_bias = True\n",
    "    ).alias(\"report\")\n",
    ").unnest(\"report\")"
   ]
  },
  {
   "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: (10_000, 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>[-0.491841, -0.342778]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.491841, -0.342778]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.491841, -0.342778]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.491841, -0.342778]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.491841, -0.342778]</td></tr><tr><td>&hellip;</td><td>&hellip;</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10_000, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ …     ┆ …                      │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    \"dummy\",\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    ).over(pl.col(\"dummy\"))\n",
    ")"
   ]
  },
  {
   "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>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.499676</td><td>0.116938</td><td>-0.771857</td><td>-0.287837</td><td>-0.48402</td></tr><tr><td>0.791879</td><td>0.094017</td><td>-0.414556</td><td>-0.425709</td><td>0.011154</td></tr><tr><td>0.526741</td><td>0.544468</td><td>-0.73191</td><td>-0.443897</td><td>-0.288013</td></tr><tr><td>0.196969</td><td>0.264493</td><td>-0.995451</td><td>-0.186299</td><td>-0.809152</td></tr><tr><td>0.786499</td><td>0.964315</td><td>-0.738782</td><td>-0.713276</td><td>-0.025506</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌──────────┬──────────┬───────────┬───────────┬───────────┐\n",
       "│ x1       ┆ x2       ┆ y         ┆ pred      ┆ resid     │\n",
       "│ ---      ┆ ---      ┆ ---       ┆ ---       ┆ ---       │\n",
       "│ f64      ┆ f64      ┆ f64       ┆ f64       ┆ f64       │\n",
       "╞══════════╪══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 0.499676 ┆ 0.116938 ┆ -0.771857 ┆ -0.287837 ┆ -0.48402  │\n",
       "│ 0.791879 ┆ 0.094017 ┆ -0.414556 ┆ -0.425709 ┆ 0.011154  │\n",
       "│ 0.526741 ┆ 0.544468 ┆ -0.73191  ┆ -0.443897 ┆ -0.288013 │\n",
       "│ 0.196969 ┆ 0.264493 ┆ -0.995451 ┆ -0.186299 ┆ -0.809152 │\n",
       "│ 0.786499 ┆ 0.964315 ┆ -0.738782 ┆ -0.713276 ┆ -0.025506 │\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.lin_reg(\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;a&quot;</td><td>[-0.491841, -0.342778]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.504246, -0.324211]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.491841, -0.342778] │\n",
       "│ b     ┆ [-0.504246, -0.324211] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy\").agg(\n",
    "    pds.lin_reg(\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": "955f2db2",
   "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;a&quot;</td><td>[-0.329471, -0.157246]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.338692, -0.144217]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.329471, -0.157246] │\n",
       "│ b     ┆ [-0.338692, -0.144217] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lasso\n",
    "df.group_by(\"dummy\").agg(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        l1_reg = 0.1,\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1bdae8e5",
   "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>lasso_r2</th></tr><tr><td>str</td><td>f64</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>-0.539511</td></tr><tr><td>&quot;b&quot;</td><td>-0.52702</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬───────────┐\n",
       "│ dummy ┆ lasso_r2  │\n",
       "│ ---   ┆ ---       │\n",
       "│ str   ┆ f64       │\n",
       "╞═══════╪═══════════╡\n",
       "│ a     ┆ -0.539511 │\n",
       "│ b     ┆ -0.52702  │\n",
       "└───────┴───────────┘"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# R2 metric of lasso regressions on each group\n",
    "df.group_by(\"dummy\").agg(\n",
    "    pds.query_r2(\n",
    "        actual = pl.col(\"y\"),\n",
    "        pred = pds.lin_reg(\n",
    "            pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "            target = pl.col(\"y\"),\n",
    "            l1_reg = 0.1,\n",
    "            return_pred = True,\n",
    "            add_bias=False\n",
    "        ).struct.field(\"pred\")\n",
    "    ).alias(\"lasso_r2\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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_000, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th><th>x1</th><th>x2</th><th>coeffs</th><th>pred</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>list[f64]</td><td>f64</td></tr></thead><tbody><tr><td>-0.771857</td><td>0.499676</td><td>0.116938</td><td>null</td><td>null</td></tr><tr><td>-0.414556</td><td>0.791879</td><td>0.094017</td><td>null</td><td>null</td></tr><tr><td>-0.73191</td><td>0.526741</td><td>0.544468</td><td>null</td><td>null</td></tr><tr><td>-0.995451</td><td>0.196969</td><td>0.264493</td><td>null</td><td>null</td></tr><tr><td>-0.738782</td><td>0.786499</td><td>0.964315</td><td>[-0.714089, -0.473885]</td><td>-1.018605</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>0.126856</td><td>0.126946</td><td>0.964879</td><td>[-1.072946, 0.029]</td><td>-0.108224</td></tr><tr><td>-1.313661</td><td>0.899251</td><td>0.051631</td><td>[-1.268071, 0.358157]</td><td>-1.121822</td></tr><tr><td>-0.550832</td><td>0.911417</td><td>0.220975</td><td>[-1.06123, 0.355866]</td><td>-0.888586</td></tr><tr><td>-1.193538</td><td>0.333315</td><td>0.54282</td><td>[-1.067394, 0.007257]</td><td>-0.351839</td></tr><tr><td>-1.364033</td><td>0.138382</td><td>0.379413</td><td>[-1.131432, -0.386638]</td><td>-0.303265</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10_000, 5)\n",
       "┌───────────┬──────────┬──────────┬────────────────────────┬───────────┐\n",
       "│ y         ┆ x1       ┆ x2       ┆ coeffs                 ┆ pred      │\n",
       "│ ---       ┆ ---      ┆ ---      ┆ ---                    ┆ ---       │\n",
       "│ f64       ┆ f64      ┆ f64      ┆ list[f64]              ┆ f64       │\n",
       "╞═══════════╪══════════╪══════════╪════════════════════════╪═══════════╡\n",
       "│ -0.771857 ┆ 0.499676 ┆ 0.116938 ┆ null                   ┆ null      │\n",
       "│ -0.414556 ┆ 0.791879 ┆ 0.094017 ┆ null                   ┆ null      │\n",
       "│ -0.73191  ┆ 0.526741 ┆ 0.544468 ┆ null                   ┆ null      │\n",
       "│ -0.995451 ┆ 0.196969 ┆ 0.264493 ┆ null                   ┆ null      │\n",
       "│ -0.738782 ┆ 0.786499 ┆ 0.964315 ┆ [-0.714089, -0.473885] ┆ -1.018605 │\n",
       "│ …         ┆ …        ┆ …        ┆ …                      ┆ …         │\n",
       "│ 0.126856  ┆ 0.126946 ┆ 0.964879 ┆ [-1.072946, 0.029]     ┆ -0.108224 │\n",
       "│ -1.313661 ┆ 0.899251 ┆ 0.051631 ┆ [-1.268071, 0.358157]  ┆ -1.121822 │\n",
       "│ -0.550832 ┆ 0.911417 ┆ 0.220975 ┆ [-1.06123, 0.355866]   ┆ -0.888586 │\n",
       "│ -1.193538 ┆ 0.333315 ┆ 0.54282  ┆ [-1.067394, 0.007257]  ┆ -0.351839 │\n",
       "│ -1.364033 ┆ 0.138382 ┆ 0.379413 ┆ [-1.131432, -0.386638] ┆ -0.303265 │\n",
       "└───────────┴──────────┴──────────┴────────────────────────┴───────────┘"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Rolling regression\n",
    "df.select(\n",
    "    \"y\",\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    pds.rolling_lin_reg(\n",
    "        \"x1\", \"x2\",\n",
    "        target = \"y\",\n",
    "        window_size = 5,\n",
    "        null_policy = \"zero\"\n",
    "    ).alias(\"result\")\n",
    ").unnest(\"result\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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": 16,
     "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": 17,
   "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>[29.123673, 28.887074, 28.317706]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────────────────────────┐\n",
       "│ a                               │\n",
       "│ ---                             │\n",
       "│ list[f64]                       │\n",
       "╞═════════════════════════════════╡\n",
       "│ [29.123673, 28.887074, 28.3177… │\n",
       "└─────────────────────────────────┘"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Only want singular values (principal values?)\n",
    "df.select(\n",
    "    pds.singular_values(\"a\", \"b\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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>29.0503</td><td>[0.835278, -0.549827]</td></tr><tr><td>28.458703</td><td>[0.549827, 0.835278]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌────────────────┬───────────────────────┐\n",
       "│ singular_value ┆ weight_vector         │\n",
       "│ ---            ┆ ---                   │\n",
       "│ f64            ┆ list[f64]             │\n",
       "╞════════════════╪═══════════════════════╡\n",
       "│ 29.0503        ┆ [0.835278, -0.549827] │\n",
       "│ 28.458703      ┆ [0.549827, 0.835278]  │\n",
       "└────────────────┴───────────────────────┘"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Singular values + The principal components\n",
    "df.select(\n",
    "    pds.pca(\"a\", \"b\")\n",
    ").unnest(\"a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e047d40a",
   "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>pc1</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>-0.210686</td></tr><tr><td>-0.011035</td></tr><tr><td>0.318099</td></tr><tr><td>-0.189717</td></tr><tr><td>0.351567</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ pc1       │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ -0.210686 │\n",
       "│ -0.011035 │\n",
       "│ 0.318099  │\n",
       "│ -0.189717 │\n",
       "│ 0.351567  │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# PC1\n",
    "df.select(\n",
    "    pds.principal_components(\"a\", \"b\", k =1).alias(\"principal_components\")\n",
    ").unnest(\"principal_components\").head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2e036f",
   "metadata": {},
   "source": [
    "# ML Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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.334322</td><td>1.001752</td><td>0.490312</td><td>0.506743</td><td>0.498392</td><td>0.483651</td><td>0.496841</td></tr><tr><td>&quot;a&quot;</td><td>0.33018</td><td>0.994606</td><td>0.508013</td><td>0.500988</td><td>0.504476</td><td>0.513971</td><td>0.508563</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.334322 ┆ 1.001752 ┆ 0.490312  ┆ 0.506743 ┆ 0.498392 ┆ 0.483651       ┆ 0.496841 │\n",
       "│ a            ┆ 0.33018  ┆ 0.994606 ┆ 0.508013  ┆ 0.500988 ┆ 0.504476 ┆ 0.513971       ┆ 0.508563 │\n",
       "└──────────────┴──────────┴──────────┴───────────┴──────────┴──────────┴────────────────┴──────────┘"
      ]
     },
     "execution_count": 20,
     "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": 21,
   "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": 21,
     "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": 22,
   "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;world&quot;</td></tr><tr><td>&quot;going&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;to&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ world  │\n",
       "│ going  │\n",
       "│ church │\n",
       "│ to     │\n",
       "│ hello  │\n",
       "└────────┘"
      ]
     },
     "execution_count": 22,
     "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": 23,
   "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;&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;world&quot;</td></tr><tr><td>&quot;go&quot;</td></tr><tr><td>&quot;church&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│        │\n",
       "│ hello  │\n",
       "│ world  │\n",
       "│ go     │\n",
       "│ church │\n",
       "└────────┘"
      ]
     },
     "execution_count": 23,
     "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": 24,
   "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": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select(\n",
    "    pds.str_leven(\"word\", pl.lit(\"world\"))\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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": 25,
     "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": 26,
   "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": 26,
     "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": 27,
   "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": 27,
     "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": 28,
   "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": 29,
   "id": "cf0c0e72",
   "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>word</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;banana&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌────────┐\n",
       "│ word   │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ banana │\n",
       "└────────┘"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    # Nearest string\n",
    "    pds.str_nearest(\"word\", word = \"banana\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "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": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    # Filters to words that are similar to any word in vocab\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": 31,
   "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": 31,
     "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": 32,
   "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.620102</td></tr><tr><td>0.354738</td></tr><tr><td>-1.146126</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ a         │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ null      │\n",
       "│ null      │\n",
       "│ 0.620102  │\n",
       "│ 0.354738  │\n",
       "│ -1.146126 │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 32,
     "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": 33,
   "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, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>random_normal</th><th>random_normal_that_respects_null_of_a</th></tr><tr><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>null</td><td>1.48806</td><td>null</td></tr><tr><td>null</td><td>1.578787</td><td>null</td></tr><tr><td>0.620102</td><td>1.197807</td><td>1.4971</td></tr><tr><td>0.354738</td><td>1.315951</td><td>1.310793</td></tr><tr><td>-1.146126</td><td>0.020791</td><td>-0.163486</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌───────────┬───────────────┬─────────────────────────────────┐\n",
       "│ a         ┆ random_normal ┆ random_normal_that_respects_nu… │\n",
       "│ ---       ┆ ---           ┆ ---                             │\n",
       "│ f64       ┆ f64           ┆ f64                             │\n",
       "╞═══════════╪═══════════════╪═════════════════════════════════╡\n",
       "│ null      ┆ 1.48806       ┆ null                            │\n",
       "│ null      ┆ 1.578787      ┆ null                            │\n",
       "│ 0.620102  ┆ 1.197807      ┆ 1.4971                          │\n",
       "│ 0.354738  ┆ 1.315951      ┆ 1.310793                        │\n",
       "│ -1.146126 ┆ 0.020791      ┆ -0.163486                       │\n",
       "└───────────┴───────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random numbers, respecting null positions in reference column (pl.col(\"a\"))\n",
    "df.with_columns(\n",
    "    pds.random_normal(mean = 0.5, std = 1.0).alias(\"random_normal\"),\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_normal(mean = 0.5, std = 1.0).alias(\"random_normal\")\n",
    "    ).alias(\"random_normal_that_respects_null_of_a\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "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, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>random_str</th><th>random_str_that_respects_null_of_a</th></tr><tr><td>f64</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>null</td><td>&quot;Tyl&quot;</td><td>null</td></tr><tr><td>null</td><td>&quot;9o&quot;</td><td>null</td></tr><tr><td>0.620102</td><td>&quot;j8&quot;</td><td>&quot;tIAdD&quot;</td></tr><tr><td>0.354738</td><td>&quot;aqGdl&quot;</td><td>&quot;yqXA&quot;</td></tr><tr><td>-1.146126</td><td>&quot;9y&quot;</td><td>&quot;BaOP&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌───────────┬────────────┬─────────────────────────────────┐\n",
       "│ a         ┆ random_str ┆ random_str_that_respects_null_… │\n",
       "│ ---       ┆ ---        ┆ ---                             │\n",
       "│ f64       ┆ str        ┆ str                             │\n",
       "╞═══════════╪════════════╪═════════════════════════════════╡\n",
       "│ null      ┆ Tyl        ┆ null                            │\n",
       "│ null      ┆ 9o         ┆ null                            │\n",
       "│ 0.620102  ┆ j8         ┆ tIAdD                           │\n",
       "│ 0.354738  ┆ aqGdl      ┆ yqXA                            │\n",
       "│ -1.146126 ┆ 9y         ┆ BaOP                            │\n",
       "└───────────┴────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random string\n",
    "df.with_columns(\n",
    "    pds.random_str(min_size = 1, max_size = 5).alias(\"random_str\"),\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_str(min_size = 1, max_size = 5)\n",
    "    ).alias(\"random_str_that_respects_null_of_a\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "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.620102</td><td>&quot;Grcv7&quot;</td></tr><tr><td>0.354738</td><td>&quot;FnKt9&quot;</td></tr><tr><td>-1.146126</td><td>&quot;CifSX&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.620102  ┆ Grcv7      │\n",
       "│ 0.354738  ┆ FnKt9      │\n",
       "│ -1.146126 ┆ CifSX      │\n",
       "└───────────┴────────────┘"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate fixed size random string, while respecting column a's nulls\n",
    "df.with_columns(\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_str(min_size = 5, max_size = 5)\n",
    "    ).alias(\"random_str\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "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>literal</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>-1.354802</td><td>null</td><td>-1.355085</td></tr><tr><td>null</td><td>0.395681</td><td>null</td><td>0.39548</td></tr><tr><td>0.620102</td><td>-1.468948</td><td>1.766922</td><td>-1.469441</td></tr><tr><td>0.354738</td><td>-0.834784</td><td>0.233956</td><td>-0.835179</td></tr><tr><td>-1.146126</td><td>0.201755</td><td>0.710931</td><td>0.201758</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌───────────┬───────────┬──────────┬─────────────────┐\n",
       "│ a         ┆ test1     ┆ literal  ┆ test1_perturbed │\n",
       "│ ---       ┆ ---       ┆ ---      ┆ ---             │\n",
       "│ f64       ┆ f64       ┆ f64      ┆ f64             │\n",
       "╞═══════════╪═══════════╪══════════╪═════════════════╡\n",
       "│ null      ┆ -1.354802 ┆ null     ┆ -1.355085       │\n",
       "│ null      ┆ 0.395681  ┆ null     ┆ 0.39548         │\n",
       "│ 0.620102  ┆ -1.468948 ┆ 1.766922 ┆ -1.469441       │\n",
       "│ 0.354738  ┆ -0.834784 ┆ 0.233956 ┆ -0.835179       │\n",
       "│ -1.146126 ┆ 0.201755  ┆ 0.710931 ┆ 0.201758        │\n",
       "└───────────┴───────────┴──────────┴─────────────────┘"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    # Sample from a normal distribution, using reference column \"a\" 's mean and std\n",
    "    pds.random_normal(pl.col(\"a\").mean(), pl.col(\"a\").std()).alias(\"test1\") \n",
    "    # Sample from uniform distribution, with low = 0 and high = \"a\"'s max, and respect the nulls in \"a\"\n",
    "    , pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random(lower = 0., upper = pl.col(\"a\").max()).alias(\"test2\")\n",
    "    )\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": 37,
   "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.537011</td><td>-0.593783</td><td>2</td></tr><tr><td>null</td><td>0.556977</td><td>-1.17522</td><td>7</td></tr><tr><td>0.620102</td><td>0.74351</td><td>0.977437</td><td>9</td></tr><tr><td>0.354738</td><td>0.722634</td><td>-1.324135</td><td>7</td></tr><tr><td>-1.146126</td><td>0.192882</td><td>-0.144544</td><td>8</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.537011 ┆ -0.593783 ┆ 2                │\n",
       "│ null      ┆ 0.556977 ┆ -1.17522  ┆ 7                │\n",
       "│ 0.620102  ┆ 0.74351  ┆ 0.977437  ┆ 9                │\n",
       "│ 0.354738  ┆ 0.722634 ┆ -1.324135 ┆ 7                │\n",
       "│ -1.146126 ┆ 0.192882 ┆ -0.144544 ┆ 8                │\n",
       "└───────────┴──────────┴───────────┴──────────────────┘"
      ]
     },
     "execution_count": 37,
     "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": 38,
   "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>0.806691</td><td>0.419973</td><td>1.536018</td><td>0.463936</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",
       "│ 0.806691            ┆ 0.419973        ┆ 1.536018                   ┆ 0.463936               │\n",
       "└─────────────────────┴─────────────────┴────────────────────────────┴────────────────────────┘"
      ]
     },
     "execution_count": 38,
     "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.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": 39,
   "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.760853</td><td>0.166417</td><td>2</td><td>8</td></tr><tr><td>1</td><td>0.149607</td><td>0.886886</td><td>2</td><td>1</td></tr><tr><td>2</td><td>0.378087</td><td>0.709648</td><td>2</td><td>1</td></tr><tr><td>0</td><td>0.012671</td><td>0.785445</td><td>3</td><td>3</td></tr><tr><td>1</td><td>0.451061</td><td>0.8895</td><td>3</td><td>1</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.760853 ┆ 0.166417 ┆ 2          ┆ 8          │\n",
       "│ 1         ┆ 0.149607 ┆ 0.886886 ┆ 2          ┆ 1          │\n",
       "│ 2         ┆ 0.378087 ┆ 0.709648 ┆ 2          ┆ 1          │\n",
       "│ 0         ┆ 0.012671 ┆ 0.785445 ┆ 3          ┆ 3          │\n",
       "│ 1         ┆ 0.451061 ┆ 0.8895   ┆ 3          ┆ 1          │\n",
       "└───────────┴──────────┴──────────┴────────────┴────────────┘"
      ]
     },
     "execution_count": 39,
     "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": 40,
   "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>{-0.236185,0.813294}</td><td>{44.257796,0.162397}</td><td>{0.284087,0.888447}</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",
       "│ {-0.236185,0.813294} ┆ {44.257796,0.162397} ┆ {0.284087,0.888447} │\n",
       "└──────────────────────┴──────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In dataframe statistical tests!\n",
    "df.select(\n",
    "    pds.ttest_ind(\"var1\", \"var2\", equal_var=True).alias(\"t-test\"),\n",
    "    pds.chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "    pds.f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "65dbb6bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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         ┆ {-0.035991,0.971292} ┆ {49.261984,0.069371} ┆ {1.242167,0.290986} │\n",
      "│ 1         ┆ {-1.297004,0.19472}  ┆ {52.860716,0.034598} ┆ {0.558786,0.69263}  │\n",
      "│ 2         ┆ {0.910669,0.362536}  ┆ {37.555952,0.397746} ┆ {1.407349,0.229106} │\n",
      "└───────────┴──────────────────────┴──────────────────────┴─────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Can also be done in group by context\n",
    "print(\n",
    "    df.group_by(\"market_id\").agg(\n",
    "        pds.ttest_ind(\"var1\", \"var2\", equal_var=False).alias(\"t-test\"),\n",
    "        pds.chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "        pds.f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "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>569</td><td>0.1138</td></tr><tr><td>585</td><td>0.117</td></tr><tr><td>555</td><td>0.111</td></tr><tr><td>533</td><td>0.1066</td></tr><tr><td>545</td><td>0.109</td></tr><tr><td>578</td><td>0.1156</td></tr><tr><td>550</td><td>0.11</td></tr><tr><td>530</td><td>0.106</td></tr><tr><td>555</td><td>0.111</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (9, 2)\n",
       "┌─────────────────┬──────────────────────────┐\n",
       "│ first_digit_cnt ┆ first_digit_distribution │\n",
       "│ ---             ┆ ---                      │\n",
       "│ u32             ┆ f64                      │\n",
       "╞═════════════════╪══════════════════════════╡\n",
       "│ 569             ┆ 0.1138                   │\n",
       "│ 585             ┆ 0.117                    │\n",
       "│ 555             ┆ 0.111                    │\n",
       "│ 533             ┆ 0.1066                   │\n",
       "│ 545             ┆ 0.109                    │\n",
       "│ 578             ┆ 0.1156                   │\n",
       "│ 550             ┆ 0.11                     │\n",
       "│ 530             ┆ 0.106                    │\n",
       "│ 555             ┆ 0.111                    │\n",
       "└─────────────────┴──────────────────────────┘"
      ]
     },
     "execution_count": 42,
     "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": 43,
   "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": 44,
   "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.720898</td><td>0.991827</td><td>0.347226</td><td>0.467622</td><td>2.156275</td><td>10</td></tr><tr><td>1</td><td>0.104165</td><td>0.519944</td><td>0.157892</td><td>0.895848</td><td>5.869907</td><td>16</td></tr><tr><td>2</td><td>0.823754</td><td>0.340088</td><td>0.019865</td><td>0.395699</td><td>2.7325</td><td>10</td></tr><tr><td>3</td><td>0.875717</td><td>0.577812</td><td>0.859156</td><td>0.361623</td><td>0.193279</td><td>18</td></tr><tr><td>4</td><td>0.766846</td><td>0.074721</td><td>0.249287</td><td>0.666755</td><td>9.584758</td><td>20</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.720898 ┆ 0.991827 ┆ 0.347226 ┆ 0.467622 ┆ 2.156275 ┆ 10           │\n",
       "│ 1   ┆ 0.104165 ┆ 0.519944 ┆ 0.157892 ┆ 0.895848 ┆ 5.869907 ┆ 16           │\n",
       "│ 2   ┆ 0.823754 ┆ 0.340088 ┆ 0.019865 ┆ 0.395699 ┆ 2.7325   ┆ 10           │\n",
       "│ 3   ┆ 0.875717 ┆ 0.577812 ┆ 0.859156 ┆ 0.361623 ┆ 0.193279 ┆ 18           │\n",
       "│ 4   ┆ 0.766846 ┆ 0.074721 ┆ 0.249287 ┆ 0.666755 ┆ 9.584758 ┆ 20           │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────┘"
      ]
     },
     "execution_count": 44,
     "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",
    "        pl.col(\"var1\"), \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        r = 0.1, # radius \n",
    "        dist = \"inf\", # L Infinity distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l_inf_cnt\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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.720898</td><td>0.991827</td><td>0.347226</td><td>0.467622</td><td>2.156275</td><td>137</td></tr><tr><td>1</td><td>0.104165</td><td>0.519944</td><td>0.157892</td><td>0.895848</td><td>5.869907</td><td>760</td></tr><tr><td>2</td><td>0.823754</td><td>0.340088</td><td>0.019865</td><td>0.395699</td><td>2.7325</td><td>88</td></tr><tr><td>3</td><td>0.875717</td><td>0.577812</td><td>0.859156</td><td>0.361623</td><td>0.193279</td><td>110</td></tr><tr><td>4</td><td>0.766846</td><td>0.074721</td><td>0.249287</td><td>0.666755</td><td>9.584758</td><td>363</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.720898 ┆ 0.991827 ┆ 0.347226 ┆ 0.467622 ┆ 2.156275 ┆ 137         │\n",
       "│ 1   ┆ 0.104165 ┆ 0.519944 ┆ 0.157892 ┆ 0.895848 ┆ 5.869907 ┆ 760         │\n",
       "│ 2   ┆ 0.823754 ┆ 0.340088 ┆ 0.019865 ┆ 0.395699 ┆ 2.7325   ┆ 88          │\n",
       "│ 3   ┆ 0.875717 ┆ 0.577812 ┆ 0.859156 ┆ 0.361623 ┆ 0.193279 ┆ 110         │\n",
       "│ 4   ┆ 0.766846 ┆ 0.074721 ┆ 0.249287 ┆ 0.666755 ┆ 9.584758 ┆ 363         │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────────┘"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        \"var1\", \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        r = pl.col(\"r\"), # radius be an expression too\n",
    "        dist = \"l1\", # L 1 distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l1_r_cnt\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "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.720898</td><td>0.991827</td><td>0.347226</td><td>0.467622</td><td>2.156275</td><td>[0, 466, … 1831]</td></tr><tr><td>1</td><td>0.104165</td><td>0.519944</td><td>0.157892</td><td>0.895848</td><td>5.869907</td><td>[1, 1982, … 1117]</td></tr><tr><td>2</td><td>0.823754</td><td>0.340088</td><td>0.019865</td><td>0.395699</td><td>2.7325</td><td>[2, 665, … 612]</td></tr><tr><td>3</td><td>0.875717</td><td>0.577812</td><td>0.859156</td><td>0.361623</td><td>0.193279</td><td>[3, 377, … 1429]</td></tr><tr><td>4</td><td>0.766846</td><td>0.074721</td><td>0.249287</td><td>0.666755</td><td>9.584758</td><td>[4, 1358, … 19]</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.720898 ┆ 0.991827 ┆ 0.347226 ┆ 0.467622 ┆ 2.156275 ┆ [0, 466, … 1831]  │\n",
       "│ 1   ┆ 0.104165 ┆ 0.519944 ┆ 0.157892 ┆ 0.895848 ┆ 5.869907 ┆ [1, 1982, … 1117] │\n",
       "│ 2   ┆ 0.823754 ┆ 0.340088 ┆ 0.019865 ┆ 0.395699 ┆ 2.7325   ┆ [2, 665, … 612]   │\n",
       "│ 3   ┆ 0.875717 ┆ 0.577812 ┆ 0.859156 ┆ 0.361623 ┆ 0.193279 ┆ [3, 377, … 1429]  │\n",
       "│ 4   ┆ 0.766846 ┆ 0.074721 ┆ 0.249287 ┆ 0.666755 ┆ 9.584758 ┆ [4, 1358, … 19]   │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴───────────────────┘"
      ]
     },
     "execution_count": 46,
     "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": 47,
   "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, 466, … 229]   ┆ 5                  │\n",
      "│ 1   ┆ [1, 1982, … 579]  ┆ 6                  │\n",
      "│ 2   ┆ [2, 665, … 1288]  ┆ 7                  │\n",
      "│ 3   ┆ [3, 377, … 86]    ┆ 9                  │\n",
      "│ 4   ┆ [4, 1358, … 1104] ┆ 15                 │\n",
      "└─────┴───────────────────┴────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Get all neighbors within radius r, call them best friends\n",
    "print(\n",
    "\n",
    "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 3d 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()\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "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.720898</td><td>0.991827</td><td>0.347226</td><td>0.467622</td><td>2.156275</td><td>[0, 466, … 1831]</td><td>[0.0, 0.059742, … 0.087801]</td></tr><tr><td>1</td><td>0.104165</td><td>0.519944</td><td>0.157892</td><td>0.895848</td><td>5.869907</td><td>[1, 1982, … 1117]</td><td>[0.0, 0.047833, … 0.070508]</td></tr><tr><td>2</td><td>0.823754</td><td>0.340088</td><td>0.019865</td><td>0.395699</td><td>2.7325</td><td>[2, 665, … 612]</td><td>[0.0, 0.042865, … 0.061837]</td></tr><tr><td>3</td><td>0.875717</td><td>0.577812</td><td>0.859156</td><td>0.361623</td><td>0.193279</td><td>[3, 377, … 1429]</td><td>[0.0, 0.070834, … 0.079681]</td></tr><tr><td>4</td><td>0.766846</td><td>0.074721</td><td>0.249287</td><td>0.666755</td><td>9.584758</td><td>[4, 1358, … 19]</td><td>[0.0, 0.018153, … 0.04832]</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.720898 ┆ 0.991827 ┆ 0.347226 ┆ 0.467622 ┆ 2.156275 ┆ [0, 466, … 1831] ┆ [0.0, 0.059742,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.087801]      │\n",
       "│ 1   ┆ 0.104165 ┆ 0.519944 ┆ 0.157892 ┆ 0.895848 ┆ 5.869907 ┆ [1, 1982, …      ┆ [0.0, 0.047833,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆ 1117]            ┆ … 0.070508]      │\n",
       "│ 2   ┆ 0.823754 ┆ 0.340088 ┆ 0.019865 ┆ 0.395699 ┆ 2.7325   ┆ [2, 665, … 612]  ┆ [0.0, 0.042865,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.061837]      │\n",
       "│ 3   ┆ 0.875717 ┆ 0.577812 ┆ 0.859156 ┆ 0.361623 ┆ 0.193279 ┆ [3, 377, … 1429] ┆ [0.0, 0.070834,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.079681]      │\n",
       "│ 4   ┆ 0.766846 ┆ 0.074721 ┆ 0.249287 ┆ 0.666755 ┆ 9.584758 ┆ [4, 1358, … 19]  ┆ [0.0, 0.018153,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.04832]       │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┴──────────────────┘"
      ]
     },
     "execution_count": 48,
     "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": 49,
   "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>8</td><td>0.441137</td><td>0.19478</td><td>0.183594</td><td>0.398282</td><td>9.18641</td></tr><tr><td>12</td><td>0.734951</td><td>0.649102</td><td>0.500223</td><td>0.883317</td><td>9.904633</td></tr><tr><td>15</td><td>0.108074</td><td>0.469015</td><td>0.479899</td><td>0.918489</td><td>9.804659</td></tr><tr><td>16</td><td>0.485801</td><td>0.530277</td><td>0.590574</td><td>0.706186</td><td>6.602236</td></tr><tr><td>20</td><td>0.360227</td><td>0.713703</td><td>0.231448</td><td>0.138263</td><td>3.325791</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",
       "│ 8   ┆ 0.441137 ┆ 0.19478  ┆ 0.183594 ┆ 0.398282 ┆ 9.18641  │\n",
       "│ 12  ┆ 0.734951 ┆ 0.649102 ┆ 0.500223 ┆ 0.883317 ┆ 9.904633 │\n",
       "│ 15  ┆ 0.108074 ┆ 0.469015 ┆ 0.479899 ┆ 0.918489 ┆ 9.804659 │\n",
       "│ 16  ┆ 0.485801 ┆ 0.530277 ┆ 0.590574 ┆ 0.706186 ┆ 6.602236 │\n",
       "│ 20  ┆ 0.360227 ┆ 0.713703 ┆ 0.231448 ┆ 0.138263 ┆ 3.325791 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Filter to only points near the given point\n",
    "df.filter(\n",
    "    pds.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": 50,
   "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>16</td><td>0.485801</td><td>0.530277</td><td>0.590574</td><td>0.706186</td><td>6.602236</td></tr><tr><td>172</td><td>0.486064</td><td>0.434635</td><td>0.266546</td><td>0.299513</td><td>3.402964</td></tr><tr><td>178</td><td>0.501763</td><td>0.522313</td><td>0.934925</td><td>0.709687</td><td>1.89558</td></tr><tr><td>184</td><td>0.543118</td><td>0.529606</td><td>0.545816</td><td>0.0628</td><td>3.587051</td></tr><tr><td>189</td><td>0.494074</td><td>0.571409</td><td>0.414045</td><td>0.330102</td><td>6.852999</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",
       "│ 16  ┆ 0.485801 ┆ 0.530277 ┆ 0.590574 ┆ 0.706186 ┆ 6.602236 │\n",
       "│ 172 ┆ 0.486064 ┆ 0.434635 ┆ 0.266546 ┆ 0.299513 ┆ 3.402964 │\n",
       "│ 178 ┆ 0.501763 ┆ 0.522313 ┆ 0.934925 ┆ 0.709687 ┆ 1.89558  │\n",
       "│ 184 ┆ 0.543118 ┆ 0.529606 ┆ 0.545816 ┆ 0.0628   ┆ 3.587051 │\n",
       "│ 189 ┆ 0.494074 ┆ 0.571409 ┆ 0.414045 ┆ 0.330102 ┆ 6.852999 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Haversine distance is available when dimension is 2\n",
    "df.filter(\n",
    "    pds.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": 51,
   "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>16</td><td>0.485801</td><td>0.530277</td><td>0.590574</td><td>0.706186</td><td>6.602236</td></tr><tr><td>303</td><td>0.502767</td><td>0.51603</td><td>0.978096</td><td>0.738166</td><td>9.051765</td></tr><tr><td>418</td><td>0.508616</td><td>0.476217</td><td>0.221188</td><td>0.421526</td><td>7.903868</td></tr><tr><td>467</td><td>0.489657</td><td>0.483175</td><td>0.523553</td><td>0.093809</td><td>7.757587</td></tr><tr><td>481</td><td>0.516688</td><td>0.550573</td><td>0.866126</td><td>0.848545</td><td>9.566443</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",
       "│ 16  ┆ 0.485801 ┆ 0.530277 ┆ 0.590574 ┆ 0.706186 ┆ 6.602236 │\n",
       "│ 303 ┆ 0.502767 ┆ 0.51603  ┆ 0.978096 ┆ 0.738166 ┆ 9.051765 │\n",
       "│ 418 ┆ 0.508616 ┆ 0.476217 ┆ 0.221188 ┆ 0.421526 ┆ 7.903868 │\n",
       "│ 467 ┆ 0.489657 ┆ 0.483175 ┆ 0.523553 ┆ 0.093809 ┆ 7.757587 │\n",
       "│ 481 ┆ 0.516688 ┆ 0.550573 ┆ 0.866126 ┆ 0.848545 ┆ 9.566443 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    pds.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": 52,
   "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, 1421, 1777]</td><td>3</td></tr><tr><td>1</td><td>[1, 17, … 842]</td><td>4</td></tr><tr><td>2</td><td>[2, 1857, … 665]</td><td>4</td></tr><tr><td>3</td><td>[3, 1868, 1468]</td><td>3</td></tr><tr><td>4</td><td>[4, 982, … 774]</td><td>8</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, 1421, 1777]  ┆ 3     │\n",
       "│ 1   ┆ [1, 17, … 842]   ┆ 4     │\n",
       "│ 2   ┆ [2, 1857, … 665] ┆ 4     │\n",
       "│ 3   ┆ [3, 1868, 1468]  ┆ 3     │\n",
       "│ 4   ┆ [4, 982, … 774]  ┆ 8     │\n",
       "└─────┴──────────────────┴───────┘"
      ]
     },
     "execution_count": 52,
     "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": "98264a05",
   "metadata": {},
   "source": [
    "# Using PDS Expressions On pl.Series, NumPy arrays, or pd.Series, etc.\n",
    "\n",
    "The output by default is always a Polars Series. The user gets to choose whether to turn it into NumPy, Pandas, or other data structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "0a42a771",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_23937/3354819425.py:3: UserWarning: The compatibility layer is considered experimental.\n",
      "  from polars_ds.compat import compat as pds2\n"
     ]
    },
    {
     "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>actual</th><th>predicted</th><th>0-2</th><th>0-9</th><th>s1</th><th>s2</th></tr><tr><td>f64</td><td>f64</td><td>i32</td><td>i32</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>0.0</td><td>0.192482</td><td>2</td><td>2</td><td>&quot;Mk&quot;</td><td>&quot;DR&quot;</td></tr><tr><td>1.0</td><td>0.010681</td><td>2</td><td>8</td><td>&quot;en&quot;</td><td>&quot;9&quot;</td></tr><tr><td>1.0</td><td>0.096006</td><td>0</td><td>6</td><td>&quot;R&quot;</td><td>&quot;Q&quot;</td></tr><tr><td>1.0</td><td>0.238899</td><td>0</td><td>2</td><td>&quot;m&quot;</td><td>&quot;t&quot;</td></tr><tr><td>0.0</td><td>0.451355</td><td>1</td><td>8</td><td>&quot;x8&quot;</td><td>&quot;a&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌────────┬───────────┬─────┬─────┬─────┬─────┐\n",
       "│ actual ┆ predicted ┆ 0-2 ┆ 0-9 ┆ s1  ┆ s2  │\n",
       "│ ---    ┆ ---       ┆ --- ┆ --- ┆ --- ┆ --- │\n",
       "│ f64    ┆ f64       ┆ i32 ┆ i32 ┆ str ┆ str │\n",
       "╞════════╪═══════════╪═════╪═════╪═════╪═════╡\n",
       "│ 0.0    ┆ 0.192482  ┆ 2   ┆ 2   ┆ Mk  ┆ DR  │\n",
       "│ 1.0    ┆ 0.010681  ┆ 2   ┆ 8   ┆ en  ┆ 9   │\n",
       "│ 1.0    ┆ 0.096006  ┆ 0   ┆ 6   ┆ R   ┆ Q   │\n",
       "│ 1.0    ┆ 0.238899  ┆ 0   ┆ 2   ┆ m   ┆ t   │\n",
       "│ 0.0    ┆ 0.451355  ┆ 1   ┆ 8   ┆ x8  ┆ a   │\n",
       "└────────┴───────────┴─────┴─────┴─────┴─────┘"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from polars_ds.compat import compat as pds2\n",
    "\n",
    "df = pds.frame(size=100_000).select(\n",
    "    pds.random(0.0, 1.0).round().alias(\"actual\"),\n",
    "    pds.random(0.0, 1.0).alias(\"predicted\"),\n",
    "    pds.random_int(0, 3).alias(\"0-2\"),\n",
    "    pds.random_int(0, 10).alias(\"0-9\"),\n",
    "    pds.random_str(min_size=1, max_size=2).alias(\"s1\"),\n",
    "    pds.random_str(min_size=1, max_size=2).alias(\"s2\"),\n",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "e6574a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pd = df.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "70830d04",
   "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,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>jaccard_col</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.3</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'jaccard_col' [f64]\n",
       "[\n",
       "\t0.3\n",
       "]"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Pandas Series\n",
    "pds2.jaccard_col(df_pd[\"0-2\"], df_pd[\"0-9\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "471cbbbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.497157  ┆ 0.496392 ┆ 0.496774 ┆ 0.498906          ┆ 0.496817 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n",
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.497157  ┆ 0.496392 ┆ 0.496774 ┆ 0.498906          ┆ 0.496817 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n",
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.497157  ┆ 0.496392 ┆ 0.496774 ┆ 0.498906          ┆ 0.496817 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n"
     ]
    }
   ],
   "source": [
    "# Polars Series\n",
    "print(pds2.query_binary_metrics(df[\"actual\"], df[\"predicted\"]).struct.unnest())\n",
    "# Pyarrow Series\n",
    "print(pds2.query_binary_metrics(df[\"actual\"].to_arrow(), df[\"predicted\"].to_arrow()).struct.unnest())\n",
    "# NumPy\n",
    "print(pds2.query_binary_metrics(df[\"actual\"].to_numpy(), df[\"predicted\"].to_numpy()).struct.unnest())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "9d15fc63",
   "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>&lt;=</th><th>baseline_pct</th><th>actual_pct</th><th>psi_bin</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.198333</td><td>0.2</td><td>0.216</td><td>0.001231</td></tr><tr><td>0.412338</td><td>0.2</td><td>0.218</td><td>0.001551</td></tr><tr><td>0.597273</td><td>0.2</td><td>0.174</td><td>0.003621</td></tr><tr><td>0.784626</td><td>0.2</td><td>0.184</td><td>0.001334</td></tr><tr><td>inf</td><td>0.2</td><td>0.208</td><td>0.000314</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌──────────┬──────────────┬────────────┬──────────┐\n",
       "│ <=       ┆ baseline_pct ┆ actual_pct ┆ psi_bin  │\n",
       "│ ---      ┆ ---          ┆ ---        ┆ ---      │\n",
       "│ f64      ┆ f64          ┆ f64        ┆ f64      │\n",
       "╞══════════╪══════════════╪════════════╪══════════╡\n",
       "│ 0.198333 ┆ 0.2          ┆ 0.216      ┆ 0.001231 │\n",
       "│ 0.412338 ┆ 0.2          ┆ 0.218      ┆ 0.001551 │\n",
       "│ 0.597273 ┆ 0.2          ┆ 0.174      ┆ 0.003621 │\n",
       "│ 0.784626 ┆ 0.2          ┆ 0.184      ┆ 0.001334 │\n",
       "│ inf      ┆ 0.2          ┆ 0.208      ┆ 0.000314 │\n",
       "└──────────┴──────────────┴────────────┴──────────┘"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NumPy Arrays\n",
    "pds2.psi(\n",
    "    np.random.random(size = 1000), \n",
    "    np.random.random(size = 1000), \n",
    "    n_bins = 5, \n",
    "    return_report = True\n",
    ").struct.unnest()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "8c5c3fee",
   "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,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>cid_ce</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>13.47064</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'cid_ce' [f64]\n",
       "[\n",
       "\t13.47064\n",
       "]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.query_cid_ce(\n",
    "    np.random.random(size = 1000),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "b0c14067",
   "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,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>c3_stats</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.141005</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'c3_stats' [f64]\n",
       "[\n",
       "\t0.141005\n",
       "]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.query_c3_stats(\n",
    "    pl.Series(values=np.random.random(size = 1000)), \n",
    "    lag = 3\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "f7cc8f20",
   "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,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>str_leven</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>1</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10,)\n",
       "Series: 'str_leven' [u32]\n",
       "[\n",
       "\t2\n",
       "\t2\n",
       "\t1\n",
       "\t1\n",
       "\t2\n",
       "\t1\n",
       "\t1\n",
       "\t2\n",
       "\t2\n",
       "\t1\n",
       "]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.str_leven(df[\"s1\"], df[\"s2\"]).head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2aab4a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}