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, 12)</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></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></tr></thead><tbody><tr><td>0.0</td><td>0</td><td>&quot;a&quot;</td><td>0</td><td>0.718061</td><td>&quot;a&quot;</td><td>0.121634</td><td>0.763113</td><td>0.732562</td><td>0.377112</td><td>0.375955</td><td>-0.851582</td></tr><tr><td>0.841471</td><td>1</td><td>&quot;a&quot;</td><td>0</td><td>0.699118</td><td>&quot;a&quot;</td><td>0.962587</td><td>0.096575</td><td>0.673011</td><td>0.738888</td><td>0.100335</td><td>-0.83615</td></tr><tr><td>0.909297</td><td>2</td><td>&quot;a&quot;</td><td>1</td><td>0.805626</td><td>&quot;a&quot;</td><td>0.813999</td><td>0.460178</td><td>0.514103</td><td>0.51099</td><td>0.785711</td><td>-0.510919</td></tr><tr><td>0.14112</td><td>3</td><td>&quot;a&quot;</td><td>1</td><td>0.00557</td><td>&quot;a&quot;</td><td>0.977954</td><td>0.687947</td><td>0.828845</td><td>0.067208</td><td>0.464131</td><td>-0.890152</td></tr><tr><td>-0.756802</td><td>4</td><td>&quot;a&quot;</td><td>0</td><td>0.668774</td><td>&quot;a&quot;</td><td>0.743661</td><td>0.900544</td><td>0.401849</td><td>0.629593</td><td>0.745616</td><td>-0.221055</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 12)\n",
       "┌───────────┬──────────┬───────┬────────┬───┬──────────┬──────────┬──────────┬───────────┐\n",
       "│ f         ┆ time_idx ┆ dummy ┆ actual ┆ … ┆ x3       ┆ a        ┆ b        ┆ y         │\n",
       "│ ---       ┆ ---      ┆ ---   ┆ ---    ┆   ┆ ---      ┆ ---      ┆ ---      ┆ ---       │\n",
       "│ f64       ┆ i64      ┆ str   ┆ i32    ┆   ┆ f64      ┆ f64      ┆ f64      ┆ f64       │\n",
       "╞═══════════╪══════════╪═══════╪════════╪═══╪══════════╪══════════╪══════════╪═══════════╡\n",
       "│ 0.0       ┆ 0        ┆ a     ┆ 0      ┆ … ┆ 0.732562 ┆ 0.377112 ┆ 0.375955 ┆ -0.851582 │\n",
       "│ 0.841471  ┆ 1        ┆ a     ┆ 0      ┆ … ┆ 0.673011 ┆ 0.738888 ┆ 0.100335 ┆ -0.83615  │\n",
       "│ 0.909297  ┆ 2        ┆ a     ┆ 1      ┆ … ┆ 0.514103 ┆ 0.51099  ┆ 0.785711 ┆ -0.510919 │\n",
       "│ 0.14112   ┆ 3        ┆ a     ┆ 1      ┆ … ┆ 0.828845 ┆ 0.067208 ┆ 0.464131 ┆ -0.890152 │\n",
       "│ -0.756802 ┆ 4        ┆ a     ┆ 0      ┆ … ┆ 0.401849 ┆ 0.629593 ┆ 0.745616 ┆ -0.221055 │\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",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6f98453-34cd-4afc-b35d-db58fa60a69a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>x1</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────┐\n",
       "│ x1  │\n",
       "│ --- │\n",
       "│ f64 │\n",
       "╞═════╡\n",
       "│ 0.0 │\n",
       "└─────┘"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Column-wise Jaccard Similarity. Result should be 0 as they are distinct\n",
    "df.select(\n",
    "    pds.query_jaccard_col(\"x1\", pl.col(\"x2\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "416d5346-e75b-4769-a953-e898d6a4d84c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th></tr><tr><td>array[f64, 2]</td></tr></thead><tbody><tr><td>[1.939505, 0.0]</td></tr><tr><td>[1.939506, 0.000209]</td></tr><tr><td>[1.939508, 0.000418]</td></tr><tr><td>[1.939512, 0.000627]</td></tr><tr><td>[1.939518, 0.000835]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────┐\n",
       "│ f                    │\n",
       "│ ---                  │\n",
       "│ array[f64, 2]        │\n",
       "╞══════════════════════╡\n",
       "│ [1.939505, 0.0]      │\n",
       "│ [1.939506, 0.000209] │\n",
       "│ [1.939508, 0.000418] │\n",
       "│ [1.939512, 0.000627] │\n",
       "│ [1.939518, 0.000835] │\n",
       "└──────────────────────┘"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. First is real part, second is complex part\n",
    "# By default, this behaves the same as np's rfft, which returns a non-redundant \n",
    "# compact representation of fft output.\n",
    "df.select(\n",
    "    pds.rfft(\"f\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "71c76353",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. But return the full length\n",
    "df.select(\n",
    "    pds.rfft(\"f\", return_full=True)\n",
    ").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cd6662d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th><th>a</th><th>b</th></tr><tr><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>1.3944e-15</td><td>-0.377112</td><td>-0.375955</td></tr><tr><td>-0.841471</td><td>-0.738888</td><td>-0.100335</td></tr><tr><td>-0.909297</td><td>-0.51099</td><td>-0.785711</td></tr><tr><td>-0.14112</td><td>-0.067208</td><td>-0.464131</td></tr><tr><td>0.756802</td><td>-0.25248</td><td>-0.369661</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.377112 ┆ -0.375955 │\n",
       "│ -0.841471  ┆ -0.738888 ┆ -0.100335 │\n",
       "│ -0.909297  ┆ -0.51099  ┆ -0.785711 │\n",
       "│ -0.14112   ┆ -0.067208 ┆ -0.464131 │\n",
       "│ 0.756802   ┆ -0.25248  ┆ -0.369661 │\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>lstsq_coeffs</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.504704, -0.337103]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌────────────────────────┐\n",
       "│ lstsq_coeffs           │\n",
       "│ ---                    │\n",
       "│ list[f64]              │\n",
       "╞════════════════════════╡\n",
       "│ [-0.504704, -0.337103] │\n",
       "└────────────────────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Least Square (Linear Regression)\n",
    "df.select(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e6da23d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (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.219255</td><td>0.001702</td><td>128.82251</td><td>0.0</td><td>0.215919</td><td>0.222592</td></tr><tr><td>&quot;exp(x2)&quot;</td><td>0.175347</td><td>0.000682</td><td>257.08593</td><td>0.0</td><td>0.17401</td><td>0.176684</td></tr><tr><td>&quot;sin(x3)&quot;</td><td>-1.745287</td><td>0.001352</td><td>-1290.855454</td><td>0.0</td><td>-1.747937</td><td>-1.742636</td></tr><tr><td>&quot;__bias__&quot;</td><td>-0.108876</td><td>0.001515</td><td>-71.843245</td><td>0.0</td><td>-0.111847</td><td>-0.105905</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.219255  ┆ 0.001702 ┆ 128.82251    ┆ 0.0   ┆ 0.215919  ┆ 0.222592  │\n",
       "│ exp(x2)  ┆ 0.175347  ┆ 0.000682 ┆ 257.08593    ┆ 0.0   ┆ 0.17401   ┆ 0.176684  │\n",
       "│ sin(x3)  ┆ -1.745287 ┆ 0.001352 ┆ -1290.855454 ┆ 0.0   ┆ -1.747937 ┆ -1.742636 │\n",
       "│ __bias__ ┆ -0.108876 ┆ 0.001515 ┆ -71.843245   ┆ 0.0   ┆ -0.111847 ┆ -0.105905 │\n",
       "└──────────┴───────────┴──────────┴──────────────┴───────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.query_lstsq_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": 9,
   "id": "4c34e7b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>lstsq_coeffs</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.504704, -0.337103]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌────────────────────────┐\n",
       "│ lstsq_coeffs           │\n",
       "│ ---                    │\n",
       "│ list[f64]              │\n",
       "╞════════════════════════╡\n",
       "│ [-0.504704, -0.337103] │\n",
       "└────────────────────────┘"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.lazy().select(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = \"y\", # We can either put pl.col(\"y\") here or just the string \"y\"\n",
    "        add_bias=False\n",
    "    )\n",
    ").collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "16511624-fc7f-45fc-b28e-ad9a91c1bfe9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (10_000, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>lstsq_coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&hellip;</td><td>&hellip;</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10_000, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ lstsq_coeffs           │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ …     ┆ …                      │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    \"dummy\",\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    ).over(pl.col(\"dummy\"))\n",
    ")"
   ]
  },
  {
   "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.121634</td><td>0.763113</td><td>-0.851582</td><td>-0.318637</td><td>-0.532946</td></tr><tr><td>0.962587</td><td>0.096575</td><td>-0.83615</td><td>-0.518377</td><td>-0.317772</td></tr><tr><td>0.813999</td><td>0.460178</td><td>-0.510919</td><td>-0.565955</td><td>0.055036</td></tr><tr><td>0.977954</td><td>0.687947</td><td>-0.890152</td><td>-0.725486</td><td>-0.164666</td></tr><tr><td>0.743661</td><td>0.900544</td><td>-0.221055</td><td>-0.678904</td><td>0.457849</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.121634 ┆ 0.763113 ┆ -0.851582 ┆ -0.318637 ┆ -0.532946 │\n",
       "│ 0.962587 ┆ 0.096575 ┆ -0.83615  ┆ -0.518377 ┆ -0.317772 │\n",
       "│ 0.813999 ┆ 0.460178 ┆ -0.510919 ┆ -0.565955 ┆ 0.055036  │\n",
       "│ 0.977954 ┆ 0.687947 ┆ -0.890152 ┆ -0.725486 ┆ -0.164666 │\n",
       "│ 0.743661 ┆ 0.900544 ┆ -0.221055 ┆ -0.678904 ┆ 0.457849  │\n",
       "└──────────┴──────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# If you want prediction and residue instead of coefficients\n",
    "df.select(\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    \"y\",\n",
    "    pds.query_lstsq(\n",
    "        \"x1\", pl.col(\"x2\"),\n",
    "        target = \"y\",\n",
    "        add_bias=False, \n",
    "        return_pred=True\n",
    "    ).alias(\"prediction\")\n",
    ").unnest(\"prediction\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0e9fb061-340d-423d-9107-772387006ff2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>lstsq_coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>[-0.499809, -0.347422]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.509563, -0.32691]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ lstsq_coeffs           │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.499809, -0.347422] │\n",
       "│ b     ┆ [-0.509563, -0.32691]  │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy\").agg(\n",
    "    pds.query_lstsq(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "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>lstsq_coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;b&quot;</td><td>[-0.336282, -0.155478]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.336463, -0.167543]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ lstsq_coeffs           │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ b     ┆ [-0.336282, -0.155478] │\n",
       "│ a     ┆ [-0.336463, -0.167543] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lasso\n",
    "df.group_by(\"dummy\").agg(\n",
    "    pds.query_lstsq(\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;b&quot;</td><td>-0.543494</td></tr><tr><td>&quot;a&quot;</td><td>-0.534178</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬───────────┐\n",
       "│ dummy ┆ lasso_r2  │\n",
       "│ ---   ┆ ---       │\n",
       "│ str   ┆ f64       │\n",
       "╞═══════╪═══════════╡\n",
       "│ b     ┆ -0.543494 │\n",
       "│ a     ┆ -0.534178 │\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.query_lstsq(\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.851582</td><td>0.121634</td><td>0.763113</td><td>null</td><td>null</td></tr><tr><td>-0.83615</td><td>0.962587</td><td>0.096575</td><td>null</td><td>null</td></tr><tr><td>-0.510919</td><td>0.813999</td><td>0.460178</td><td>null</td><td>null</td></tr><tr><td>-0.890152</td><td>0.977954</td><td>0.687947</td><td>null</td><td>null</td></tr><tr><td>-0.221055</td><td>0.743661</td><td>0.900544</td><td>[-0.536204, -0.362532]</td><td>-0.72523</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>-0.019252</td><td>0.086137</td><td>0.78473</td><td>[-0.859638, -0.002479]</td><td>-0.075992</td></tr><tr><td>-0.564748</td><td>0.577775</td><td>0.371494</td><td>[-1.009086, 0.089103]</td><td>-0.549923</td></tr><tr><td>-0.661185</td><td>0.895924</td><td>0.83152</td><td>[-0.957789, 0.219304]</td><td>-0.67575</td></tr><tr><td>-0.724577</td><td>0.207569</td><td>0.373584</td><td>[-0.847493, 0.016208]</td><td>-0.169858</td></tr><tr><td>-0.223634</td><td>0.315226</td><td>0.192244</td><td>[-0.799903, -0.096319]</td><td>-0.270667</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.851582 ┆ 0.121634 ┆ 0.763113 ┆ null                   ┆ null      │\n",
       "│ -0.83615  ┆ 0.962587 ┆ 0.096575 ┆ null                   ┆ null      │\n",
       "│ -0.510919 ┆ 0.813999 ┆ 0.460178 ┆ null                   ┆ null      │\n",
       "│ -0.890152 ┆ 0.977954 ┆ 0.687947 ┆ null                   ┆ null      │\n",
       "│ -0.221055 ┆ 0.743661 ┆ 0.900544 ┆ [-0.536204, -0.362532] ┆ -0.72523  │\n",
       "│ …         ┆ …        ┆ …        ┆ …                      ┆ …         │\n",
       "│ -0.019252 ┆ 0.086137 ┆ 0.78473  ┆ [-0.859638, -0.002479] ┆ -0.075992 │\n",
       "│ -0.564748 ┆ 0.577775 ┆ 0.371494 ┆ [-1.009086, 0.089103]  ┆ -0.549923 │\n",
       "│ -0.661185 ┆ 0.895924 ┆ 0.83152  ┆ [-0.957789, 0.219304]  ┆ -0.67575  │\n",
       "│ -0.724577 ┆ 0.207569 ┆ 0.373584 ┆ [-0.847493, 0.016208]  ┆ -0.169858 │\n",
       "│ -0.223634 ┆ 0.315226 ┆ 0.192244 ┆ [-0.799903, -0.096319] ┆ -0.270667 │\n",
       "└───────────┴──────────┴──────────┴────────────────────────┴───────────┘"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Rolling regression\n",
    "df.select(\n",
    "    \"y\",\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    pds.query_rolling_lstsq(\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.003027, 28.800535, 28.766516]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────────────────────────┐\n",
       "│ a                               │\n",
       "│ ---                             │\n",
       "│ list[f64]                       │\n",
       "╞═════════════════════════════════╡\n",
       "│ [29.003027, 28.800535, 28.7665… │\n",
       "└─────────────────────────────────┘"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Only want singular values (principal values?)\n",
    "df.select(\n",
    "    pds.query_singular_values(\"a\", \"b\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 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>28.989289</td><td>[-0.344595, 0.938751]</td></tr><tr><td>28.798623</td><td>[0.938751, 0.344595]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌────────────────┬───────────────────────┐\n",
       "│ singular_value ┆ weight_vector         │\n",
       "│ ---            ┆ ---                   │\n",
       "│ f64            ┆ list[f64]             │\n",
       "╞════════════════╪═══════════════════════╡\n",
       "│ 28.989289      ┆ [-0.344595, 0.938751] │\n",
       "│ 28.798623      ┆ [0.938751, 0.344595]  │\n",
       "└────────────────┴───────────────────────┘"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Singular values + The principal components\n",
    "df.select(\n",
    "    pds.query_pca(\"a\", \"b\")\n",
    ").unnest(\"a\")"
   ]
  },
  {
   "cell_type": "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.081069</td></tr><tr><td>-0.464474</td></tr><tr><td>0.257457</td></tr><tr><td>0.108498</td></tr><tr><td>0.178948</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ pc1       │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ -0.081069 │\n",
       "│ -0.464474 │\n",
       "│ 0.257457  │\n",
       "│ 0.108498  │\n",
       "│ 0.178948  │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# PC1\n",
    "df.select(\n",
    "    pds.query_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.333429</td><td>1.002914</td><td>0.520892</td><td>0.496904</td><td>0.508616</td><td>0.524247</td><td>0.499725</td></tr><tr><td>&quot;a&quot;</td><td>0.332551</td><td>0.994944</td><td>0.48794</td><td>0.505324</td><td>0.49648</td><td>0.489441</td><td>0.502887</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.333429 ┆ 1.002914 ┆ 0.520892  ┆ 0.496904 ┆ 0.508616 ┆ 0.524247       ┆ 0.499725 │\n",
       "│ a            ┆ 0.332551 ┆ 0.994944 ┆ 0.48794   ┆ 0.505324 ┆ 0.49648  ┆ 0.489441       ┆ 0.502887 │\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;church&quot;</td></tr><tr><td>&quot;going&quot;</td></tr><tr><td>&quot;to&quot;</td></tr><tr><td>&quot;world&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",
       "│ church │\n",
       "│ going  │\n",
       "│ to     │\n",
       "│ world  │\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;world&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;&quot;</td></tr><tr><td>&quot;go&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ world  │\n",
       "│ hello  │\n",
       "│ church │\n",
       "│        │\n",
       "│ go     │\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": "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": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    pds.similar_to_vocab(\n",
    "        pl.col(\"word\"),\n",
    "        vocab = gibberish,\n",
    "        threshold = 0.5,\n",
    "        metric = \"lv\", # Levenshtein similarity. Other options: dleven, osa, jw\n",
    "        strategy = \"any\" # True if the word is similar to any word in vocab. Other options: \"all\", \"avg\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "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": 30,
     "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": 31,
   "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.012492</td></tr><tr><td>0.595561</td></tr><tr><td>0.185681</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ a         │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ null      │\n",
       "│ null      │\n",
       "│ -0.012492 │\n",
       "│ 0.595561  │\n",
       "│ 0.185681  │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 31,
     "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": 32,
   "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>2.978394</td><td>null</td></tr><tr><td>null</td><td>0.970464</td><td>null</td></tr><tr><td>-0.012492</td><td>-0.072767</td><td>0.292304</td></tr><tr><td>0.595561</td><td>0.286816</td><td>1.661764</td></tr><tr><td>0.185681</td><td>0.591648</td><td>1.802518</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      ┆ 2.978394      ┆ null                            │\n",
       "│ null      ┆ 0.970464      ┆ null                            │\n",
       "│ -0.012492 ┆ -0.072767     ┆ 0.292304                        │\n",
       "│ 0.595561  ┆ 0.286816      ┆ 1.661764                        │\n",
       "│ 0.185681  ┆ 0.591648      ┆ 1.802518                        │\n",
       "└───────────┴───────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 32,
     "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": 33,
   "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;8kwTY&quot;</td><td>null</td></tr><tr><td>null</td><td>&quot;kMp&quot;</td><td>null</td></tr><tr><td>-0.012492</td><td>&quot;Rc37s&quot;</td><td>&quot;W1&quot;</td></tr><tr><td>0.595561</td><td>&quot;5&quot;</td><td>&quot;Xj5GN&quot;</td></tr><tr><td>0.185681</td><td>&quot;v7fD&quot;</td><td>&quot;1buvo&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      ┆ 8kwTY      ┆ null                            │\n",
       "│ null      ┆ kMp        ┆ null                            │\n",
       "│ -0.012492 ┆ Rc37s      ┆ W1                              │\n",
       "│ 0.595561  ┆ 5          ┆ Xj5GN                           │\n",
       "│ 0.185681  ┆ v7fD       ┆ 1buvo                           │\n",
       "└───────────┴────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 33,
     "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": 34,
   "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.012492</td><td>&quot;kecis&quot;</td></tr><tr><td>0.595561</td><td>&quot;u9Vh3&quot;</td></tr><tr><td>0.185681</td><td>&quot;0adxI&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.012492 ┆ kecis      │\n",
       "│ 0.595561  ┆ u9Vh3      │\n",
       "│ 0.185681  ┆ 0adxI      │\n",
       "└───────────┴────────────┘"
      ]
     },
     "execution_count": 34,
     "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": 35,
   "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>-0.261042</td><td>null</td><td>-0.261344</td></tr><tr><td>null</td><td>-0.527867</td><td>null</td><td>-0.527969</td></tr><tr><td>-0.012492</td><td>-0.631845</td><td>2.562054</td><td>-0.631728</td></tr><tr><td>0.595561</td><td>0.673208</td><td>1.030199</td><td>0.672906</td></tr><tr><td>0.185681</td><td>0.572215</td><td>2.065594</td><td>0.572634</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      ┆ -0.261042 ┆ null     ┆ -0.261344       │\n",
       "│ null      ┆ -0.527867 ┆ null     ┆ -0.527969       │\n",
       "│ -0.012492 ┆ -0.631845 ┆ 2.562054 ┆ -0.631728       │\n",
       "│ 0.595561  ┆ 0.673208  ┆ 1.030199 ┆ 0.672906        │\n",
       "│ 0.185681  ┆ 0.572215  ┆ 2.065594 ┆ 0.572634        │\n",
       "└───────────┴───────────┴──────────┴─────────────────┘"
      ]
     },
     "execution_count": 35,
     "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": 36,
   "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.106948</td><td>0.427388</td><td>5</td></tr><tr><td>null</td><td>0.095822</td><td>-0.798821</td><td>9</td></tr><tr><td>-0.012492</td><td>0.388392</td><td>-0.185242</td><td>2</td></tr><tr><td>0.595561</td><td>0.042763</td><td>0.802807</td><td>4</td></tr><tr><td>0.185681</td><td>0.321888</td><td>0.277879</td><td>2</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.106948 ┆ 0.427388  ┆ 5                │\n",
       "│ null      ┆ 0.095822 ┆ -0.798821 ┆ 9                │\n",
       "│ -0.012492 ┆ 0.388392 ┆ -0.185242 ┆ 2                │\n",
       "│ 0.595561  ┆ 0.042763 ┆ 0.802807  ┆ 4                │\n",
       "│ 0.185681  ┆ 0.321888 ┆ 0.277879  ┆ 2                │\n",
       "└───────────┴──────────┴───────────┴──────────────────┘"
      ]
     },
     "execution_count": 36,
     "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": 37,
   "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.57873</td><td>0.562857</td><td>7.284288</td><td>0.026196</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.57873             ┆ 0.562857        ┆ 7.284288                   ┆ 0.026196               │\n",
       "└─────────────────────┴─────────────────┴────────────────────────────┴────────────────────────┘"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate 2 random sample, both normally distributed\n",
    "# Run Welch's t test on them, p value should be big since they have equal mean\n",
    "# Run a normality test. Again, p value should be big since they are normally distributed \n",
    "\n",
    "df.with_columns(\n",
    "    pds.random_normal(0.5, 1.0).alias(\"test1\"),\n",
    "    pds.random_normal(0.5, 2.0).alias(\"test2\"),\n",
    ").select(\n",
    "    pds.query_ttest_ind(\"test1\", \"test2\", equal_var=False).alias(\"t-test\"),\n",
    "    pds.normal_test(\"test1\").alias(\"normality_test\")\n",
    ").select(\n",
    "    pl.col(\"t-test\").struct.field(\"statistic\").alias(\"t-tests: statistics\")\n",
    "    , pl.col(\"t-test\").struct.field(\"pvalue\").alias(\"t-tests: pvalue\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"statistic\").alias(\"normality_test: statistics\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"pvalue\").alias(\"normality_test: pvalue\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "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.099163</td><td>0.187826</td><td>3</td><td>1</td></tr><tr><td>1</td><td>0.212648</td><td>0.759222</td><td>4</td><td>2</td></tr><tr><td>2</td><td>0.698327</td><td>0.597246</td><td>1</td><td>6</td></tr><tr><td>0</td><td>0.445486</td><td>0.3729</td><td>4</td><td>4</td></tr><tr><td>1</td><td>0.649741</td><td>0.974461</td><td>4</td><td>9</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.099163 ┆ 0.187826 ┆ 3          ┆ 1          │\n",
       "│ 1         ┆ 0.212648 ┆ 0.759222 ┆ 4          ┆ 2          │\n",
       "│ 2         ┆ 0.698327 ┆ 0.597246 ┆ 1          ┆ 6          │\n",
       "│ 0         ┆ 0.445486 ┆ 0.3729   ┆ 4          ┆ 4          │\n",
       "│ 1         ┆ 0.649741 ┆ 0.974461 ┆ 4          ┆ 9          │\n",
       "└───────────┴──────────┴──────────┴────────────┴────────────┘"
      ]
     },
     "execution_count": 38,
     "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": 39,
   "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.03261,0.973986}</td><td>{27.618037,0.840533}</td><td>{0.849663,0.493529}</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.03261,0.973986} ┆ {27.618037,0.840533} ┆ {0.849663,0.493529} │\n",
       "└─────────────────────┴──────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In dataframe statistical tests!\n",
    "df.select(\n",
    "    pds.query_ttest_ind(\"var1\", \"var2\", equal_var=True).alias(\"t-test\"),\n",
    "    pds.query_chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "    pds.query_f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "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         ┆ {1.320059,0.186906}  ┆ {36.690713,0.436664} ┆ {0.2116,0.932081}   │\n",
      "│ 1         ┆ {-0.123313,0.901867} ┆ {29.961059,0.750506} ┆ {1.363618,0.244265} │\n",
      "│ 2         ┆ {-1.237752,0.215895} ┆ {26.494933,0.876345} ┆ {0.36564,0.833211}  │\n",
      "└───────────┴──────────────────────┴──────────────────────┴─────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Can also be done in group by context\n",
    "print(\n",
    "    df.group_by(\"market_id\").agg(\n",
    "        pds.query_ttest_ind(\"var1\", \"var2\", equal_var=False).alias(\"t-test\"),\n",
    "        pds.query_chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "        pds.query_f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "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>567</td><td>0.1134</td></tr><tr><td>558</td><td>0.1116</td></tr><tr><td>567</td><td>0.1134</td></tr><tr><td>564</td><td>0.1128</td></tr><tr><td>544</td><td>0.1088</td></tr><tr><td>588</td><td>0.1176</td></tr><tr><td>529</td><td>0.1058</td></tr><tr><td>548</td><td>0.1096</td></tr><tr><td>535</td><td>0.107</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (9, 2)\n",
       "┌─────────────────┬──────────────────────────┐\n",
       "│ first_digit_cnt ┆ first_digit_distribution │\n",
       "│ ---             ┆ ---                      │\n",
       "│ u32             ┆ f64                      │\n",
       "╞═════════════════╪══════════════════════════╡\n",
       "│ 567             ┆ 0.1134                   │\n",
       "│ 558             ┆ 0.1116                   │\n",
       "│ 567             ┆ 0.1134                   │\n",
       "│ 564             ┆ 0.1128                   │\n",
       "│ 544             ┆ 0.1088                   │\n",
       "│ 588             ┆ 0.1176                   │\n",
       "│ 529             ┆ 0.1058                   │\n",
       "│ 548             ┆ 0.1096                   │\n",
       "│ 535             ┆ 0.107                    │\n",
       "└─────────────────┴──────────────────────────┘"
      ]
     },
     "execution_count": 41,
     "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": 42,
   "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": 43,
   "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.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td><td>14</td></tr><tr><td>1</td><td>0.176238</td><td>0.691061</td><td>0.38883</td><td>0.873326</td><td>9.814184</td><td>22</td></tr><tr><td>2</td><td>0.201282</td><td>0.876521</td><td>0.566095</td><td>0.046258</td><td>3.17553</td><td>20</td></tr><tr><td>3</td><td>0.859807</td><td>0.967636</td><td>0.785815</td><td>0.804368</td><td>9.882445</td><td>11</td></tr><tr><td>4</td><td>0.064274</td><td>0.618677</td><td>0.463781</td><td>0.232584</td><td>8.62535</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.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 ┆ 14           │\n",
       "│ 1   ┆ 0.176238 ┆ 0.691061 ┆ 0.38883  ┆ 0.873326 ┆ 9.814184 ┆ 22           │\n",
       "│ 2   ┆ 0.201282 ┆ 0.876521 ┆ 0.566095 ┆ 0.046258 ┆ 3.17553  ┆ 20           │\n",
       "│ 3   ┆ 0.859807 ┆ 0.967636 ┆ 0.785815 ┆ 0.804368 ┆ 9.882445 ┆ 11           │\n",
       "│ 4   ┆ 0.064274 ┆ 0.618677 ┆ 0.463781 ┆ 0.232584 ┆ 8.62535  ┆ 20           │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────┘"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get neighbor count. The point itself is always considered a neighbor to itself.\n",
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        0.1, # radius \n",
    "        pl.col(\"var1\"), \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        dist = \"inf\", # L Infinity distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l_inf_cnt\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "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.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td><td>233</td></tr><tr><td>1</td><td>0.176238</td><td>0.691061</td><td>0.38883</td><td>0.873326</td><td>9.814184</td><td>968</td></tr><tr><td>2</td><td>0.201282</td><td>0.876521</td><td>0.566095</td><td>0.046258</td><td>3.17553</td><td>1</td></tr><tr><td>3</td><td>0.859807</td><td>0.967636</td><td>0.785815</td><td>0.804368</td><td>9.882445</td><td>421</td></tr><tr><td>4</td><td>0.064274</td><td>0.618677</td><td>0.463781</td><td>0.232584</td><td>8.62535</td><td>32</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.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 ┆ 233         │\n",
       "│ 1   ┆ 0.176238 ┆ 0.691061 ┆ 0.38883  ┆ 0.873326 ┆ 9.814184 ┆ 968         │\n",
       "│ 2   ┆ 0.201282 ┆ 0.876521 ┆ 0.566095 ┆ 0.046258 ┆ 3.17553  ┆ 1           │\n",
       "│ 3   ┆ 0.859807 ┆ 0.967636 ┆ 0.785815 ┆ 0.804368 ┆ 9.882445 ┆ 421         │\n",
       "│ 4   ┆ 0.064274 ┆ 0.618677 ┆ 0.463781 ┆ 0.232584 ┆ 8.62535  ┆ 32          │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────────┘"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        pl.col(\"r\"), # radius be an expression too\n",
    "        \"var1\", \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        dist = \"l1\", # L 1 distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l1_r_cnt\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td><td>[0, 1220, … 61]</td></tr><tr><td>1</td><td>0.176238</td><td>0.691061</td><td>0.38883</td><td>0.873326</td><td>9.814184</td><td>[1, 1919, … 1438]</td></tr><tr><td>2</td><td>0.201282</td><td>0.876521</td><td>0.566095</td><td>0.046258</td><td>3.17553</td><td>[2, 1530, … 1511]</td></tr><tr><td>3</td><td>0.859807</td><td>0.967636</td><td>0.785815</td><td>0.804368</td><td>9.882445</td><td>[3, 1734, … 838]</td></tr><tr><td>4</td><td>0.064274</td><td>0.618677</td><td>0.463781</td><td>0.232584</td><td>8.62535</td><td>[4, 348, … 307]</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.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 ┆ [0, 1220, … 61]   │\n",
       "│ 1   ┆ 0.176238 ┆ 0.691061 ┆ 0.38883  ┆ 0.873326 ┆ 9.814184 ┆ [1, 1919, … 1438] │\n",
       "│ 2   ┆ 0.201282 ┆ 0.876521 ┆ 0.566095 ┆ 0.046258 ┆ 3.17553  ┆ [2, 1530, … 1511] │\n",
       "│ 3   ┆ 0.859807 ┆ 0.967636 ┆ 0.785815 ┆ 0.804368 ┆ 9.882445 ┆ [3, 1734, … 838]  │\n",
       "│ 4   ┆ 0.064274 ┆ 0.618677 ┆ 0.463781 ┆ 0.232584 ┆ 8.62535  ┆ [4, 348, … 307]   │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴───────────────────┘"
      ]
     },
     "execution_count": 45,
     "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": 46,
   "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, 1220, … 687]  ┆ 4                  │\n",
      "│ 1   ┆ [1, 1919, … 254]  ┆ 14                 │\n",
      "│ 2   ┆ [2, 1530, … 1457] ┆ 8                  │\n",
      "│ 3   ┆ [3, 1734, … 820]  ┆ 5                  │\n",
      "│ 4   ┆ [4, 348, … 1888]  ┆ 13                 │\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": 47,
   "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.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td><td>[0, 1220, … 61]</td><td>[0.0, 0.033199, … 0.060468]</td></tr><tr><td>1</td><td>0.176238</td><td>0.691061</td><td>0.38883</td><td>0.873326</td><td>9.814184</td><td>[1, 1919, … 1438]</td><td>[0.0, 0.022711, … 0.06655]</td></tr><tr><td>2</td><td>0.201282</td><td>0.876521</td><td>0.566095</td><td>0.046258</td><td>3.17553</td><td>[2, 1530, … 1511]</td><td>[0.0, 0.067087, … 0.076166]</td></tr><tr><td>3</td><td>0.859807</td><td>0.967636</td><td>0.785815</td><td>0.804368</td><td>9.882445</td><td>[3, 1734, … 838]</td><td>[0.0, 0.023593, … 0.07114]</td></tr><tr><td>4</td><td>0.064274</td><td>0.618677</td><td>0.463781</td><td>0.232584</td><td>8.62535</td><td>[4, 348, … 307]</td><td>[0.0, 0.062754, … 0.076836]</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.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 ┆ [0, 1220, … 61]  ┆ [0.0, 0.033199,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.060468]      │\n",
       "│ 1   ┆ 0.176238 ┆ 0.691061 ┆ 0.38883  ┆ 0.873326 ┆ 9.814184 ┆ [1, 1919, …      ┆ [0.0, 0.022711,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆ 1438]            ┆ … 0.06655]       │\n",
       "│ 2   ┆ 0.201282 ┆ 0.876521 ┆ 0.566095 ┆ 0.046258 ┆ 3.17553  ┆ [2, 1530, …      ┆ [0.0, 0.067087,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆ 1511]            ┆ … 0.076166]      │\n",
       "│ 3   ┆ 0.859807 ┆ 0.967636 ┆ 0.785815 ┆ 0.804368 ┆ 9.882445 ┆ [3, 1734, … 838] ┆ [0.0, 0.023593,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.07114]       │\n",
       "│ 4   ┆ 0.064274 ┆ 0.618677 ┆ 0.463781 ┆ 0.232584 ┆ 8.62535  ┆ [4, 348, … 307]  ┆ [0.0, 0.062754,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.076836]      │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┴──────────────────┘"
      ]
     },
     "execution_count": 47,
     "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": 48,
   "id": "b5c69ae4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0</td><td>0.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td></tr><tr><td>1</td><td>0.176238</td><td>0.691061</td><td>0.38883</td><td>0.873326</td><td>9.814184</td></tr><tr><td>5</td><td>0.712239</td><td>0.209799</td><td>0.535259</td><td>0.505202</td><td>7.753107</td></tr><tr><td>7</td><td>0.854959</td><td>0.645341</td><td>0.465969</td><td>0.924045</td><td>8.752052</td></tr><tr><td>8</td><td>0.797774</td><td>0.567559</td><td>0.411024</td><td>0.193817</td><td>9.037423</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 0   ┆ 0.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 │\n",
       "│ 1   ┆ 0.176238 ┆ 0.691061 ┆ 0.38883  ┆ 0.873326 ┆ 9.814184 │\n",
       "│ 5   ┆ 0.712239 ┆ 0.209799 ┆ 0.535259 ┆ 0.505202 ┆ 7.753107 │\n",
       "│ 7   ┆ 0.854959 ┆ 0.645341 ┆ 0.465969 ┆ 0.924045 ┆ 8.752052 │\n",
       "│ 8   ┆ 0.797774 ┆ 0.567559 ┆ 0.411024 ┆ 0.193817 ┆ 9.037423 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 48,
     "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": 49,
   "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>0</td><td>0.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td></tr><tr><td>40</td><td>0.41957</td><td>0.46582</td><td>0.289675</td><td>0.797387</td><td>1.012019</td></tr><tr><td>61</td><td>0.492211</td><td>0.532409</td><td>0.747621</td><td>0.607236</td><td>6.230372</td></tr><tr><td>66</td><td>0.445782</td><td>0.557399</td><td>0.566395</td><td>0.87715</td><td>5.277144</td></tr><tr><td>138</td><td>0.512836</td><td>0.538711</td><td>0.918922</td><td>0.563493</td><td>1.969114</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 0   ┆ 0.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 │\n",
       "│ 40  ┆ 0.41957  ┆ 0.46582  ┆ 0.289675 ┆ 0.797387 ┆ 1.012019 │\n",
       "│ 61  ┆ 0.492211 ┆ 0.532409 ┆ 0.747621 ┆ 0.607236 ┆ 6.230372 │\n",
       "│ 66  ┆ 0.445782 ┆ 0.557399 ┆ 0.566395 ┆ 0.87715  ┆ 5.277144 │\n",
       "│ 138 ┆ 0.512836 ┆ 0.538711 ┆ 0.918922 ┆ 0.563493 ┆ 1.969114 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 49,
     "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": 50,
   "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>0</td><td>0.481702</td><td>0.566173</td><td>0.69857</td><td>0.458351</td><td>7.900392</td></tr><tr><td>61</td><td>0.492211</td><td>0.532409</td><td>0.747621</td><td>0.607236</td><td>6.230372</td></tr><tr><td>180</td><td>0.522403</td><td>0.494358</td><td>0.795411</td><td>0.971752</td><td>7.740899</td></tr><tr><td>377</td><td>0.520169</td><td>0.449136</td><td>0.262714</td><td>0.088863</td><td>8.283747</td></tr><tr><td>480</td><td>0.472071</td><td>0.455911</td><td>0.734145</td><td>0.448297</td><td>6.728847</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 0   ┆ 0.481702 ┆ 0.566173 ┆ 0.69857  ┆ 0.458351 ┆ 7.900392 │\n",
       "│ 61  ┆ 0.492211 ┆ 0.532409 ┆ 0.747621 ┆ 0.607236 ┆ 6.230372 │\n",
       "│ 180 ┆ 0.522403 ┆ 0.494358 ┆ 0.795411 ┆ 0.971752 ┆ 7.740899 │\n",
       "│ 377 ┆ 0.520169 ┆ 0.449136 ┆ 0.262714 ┆ 0.088863 ┆ 8.283747 │\n",
       "│ 480 ┆ 0.472071 ┆ 0.455911 ┆ 0.734145 ┆ 0.448297 ┆ 6.728847 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 50,
     "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": 51,
   "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, 911, 1130]</td><td>3</td></tr><tr><td>1</td><td>[1, 1919]</td><td>2</td></tr><tr><td>2</td><td>[2]</td><td>1</td></tr><tr><td>3</td><td>[3, 548, 1734]</td><td>3</td></tr><tr><td>4</td><td>[4, 981, … 343]</td><td>5</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, 911, 1130]  ┆ 3     │\n",
       "│ 1   ┆ [1, 1919]       ┆ 2     │\n",
       "│ 2   ┆ [2]             ┆ 1     │\n",
       "│ 3   ┆ [3, 548, 1734]  ┆ 3     │\n",
       "│ 4   ┆ [4, 981, … 343] ┆ 5     │\n",
       "└─────┴─────────────────┴───────┘"
      ]
     },
     "execution_count": 51,
     "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": "9f8eabf1",
   "metadata": {},
   "source": [
    "# String Nearest Neighbors\n",
    "\n",
    "This might be very slow for very large vocab / column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f2c9ed80",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame({\n",
    "    \"a\":[\"AAAAA\", \"ABCABC\", \"AAAADDD\", \"ADSDSDS\", \"WORD\"],\n",
    "    \"b\":[\"AAAAT\", \"ABCACD\", \"ADSSD\", \"APPLES\", \"WORLD\"] \n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "062fd694",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;AAAAT&quot;</td></tr><tr><td>&quot;ABCACD&quot;</td></tr><tr><td>&quot;AAAAT&quot;</td></tr><tr><td>&quot;ADSSD&quot;</td></tr><tr><td>&quot;WORLD&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ str                      │\n",
       "╞══════════════════════════╡\n",
       "│ AAAAT                    │\n",
       "│ ABCACD                   │\n",
       "│ AAAAT                    │\n",
       "│ ADSSD                    │\n",
       "│ WORLD                    │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use Levenshtein to find the nearest neighbor in vocab to word in column a\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 1, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "29437b8b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>list[str]</td></tr></thead><tbody><tr><td>[&quot;AAAAT&quot;, &quot;ADSSD&quot;]</td></tr><tr><td>[&quot;ABCACD&quot;, &quot;AAAAT&quot;]</td></tr><tr><td>[&quot;AAAAT&quot;, &quot;ABCACD&quot;]</td></tr><tr><td>[&quot;ADSSD&quot;, &quot;APPLES&quot;]</td></tr><tr><td>[&quot;WORLD&quot;, &quot;ADSSD&quot;]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ list[str]                │\n",
       "╞══════════════════════════╡\n",
       "│ [\"AAAAT\", \"ADSSD\"]       │\n",
       "│ [\"ABCACD\", \"AAAAT\"]      │\n",
       "│ [\"AAAAT\", \"ABCACD\"]      │\n",
       "│ [\"ADSSD\", \"APPLES\"]      │\n",
       "│ [\"WORLD\", \"ADSSD\"]       │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use Levenshtein to find 2 nearest neighbors\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 2, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9150bece",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>similar_words_from_vocab</th></tr><tr><td>list[str]</td></tr></thead><tbody><tr><td>[&quot;AAAAT&quot;, &quot;ADSSD&quot;]</td></tr><tr><td>[&quot;ABCACD&quot;]</td></tr><tr><td>[]</td></tr><tr><td>[]</td></tr><tr><td>[]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────────┐\n",
       "│ similar_words_from_vocab │\n",
       "│ ---                      │\n",
       "│ list[str]                │\n",
       "╞══════════════════════════╡\n",
       "│ [\"AAAAT\", \"ADSSD\"]       │\n",
       "│ [\"ABCACD\"]               │\n",
       "│ []                       │\n",
       "│ []                       │\n",
       "│ []                       │\n",
       "└──────────────────────────┘"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Currently only Levenshtein and hamming are implemented for this\n",
    "# Empty means nothing in vocab can be compared in the hamming sense with the corresponding word in a\n",
    "df.select(\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = pl.col(\"b\"),\n",
    "        k = 2, \n",
    "        threshold = 4,\n",
    "        metric = \"hamming\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "b35776a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>similar_words_from_vocab</th></tr><tr><td>str</td><td>list[str]</td></tr></thead><tbody><tr><td>&quot;AAAAA&quot;</td><td>[&quot;AAAAA&quot;, &quot;ZIV&quot;, &quot;WORLD&quot;]</td></tr><tr><td>&quot;ABCABC&quot;</td><td>[&quot;ABCDEFG&quot;, &quot;AAAAA&quot;, &quot;ZIV&quot;]</td></tr><tr><td>&quot;AAAADDD&quot;</td><td>[&quot;AAAAA&quot;, &quot;WORLD&quot;, &quot;ABCDEFG&quot;]</td></tr><tr><td>&quot;ADSDSDS&quot;</td><td>[&quot;ABCDEFG&quot;, &quot;WORLD&quot;, &quot;AAAAA&quot;]</td></tr><tr><td>&quot;WORD&quot;</td><td>[&quot;WORLD&quot;, &quot;ZIV&quot;, &quot;TQQQ&quot;]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────┬───────────────────────────────┐\n",
       "│ a       ┆ similar_words_from_vocab      │\n",
       "│ ---     ┆ ---                           │\n",
       "│ str     ┆ list[str]                     │\n",
       "╞═════════╪═══════════════════════════════╡\n",
       "│ AAAAA   ┆ [\"AAAAA\", \"ZIV\", \"WORLD\"]     │\n",
       "│ ABCABC  ┆ [\"ABCDEFG\", \"AAAAA\", \"ZIV\"]   │\n",
       "│ AAAADDD ┆ [\"AAAAA\", \"WORLD\", \"ABCDEFG\"] │\n",
       "│ ADSDSDS ┆ [\"ABCDEFG\", \"WORLD\", \"AAAAA\"] │\n",
       "│ WORD    ┆ [\"WORLD\", \"ZIV\", \"TQQQ\"]      │\n",
       "└─────────┴───────────────────────────────┘"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# You may provide a vocab like this\n",
    "df.select(\n",
    "    pl.col(\"a\"),\n",
    "    pds.query_similar_words(\n",
    "        \"a\",\n",
    "        vocab = [\"WORLD\", \"AAAAA\", \"ABCDEFG\", \"ZIV\", \"TQQQ\"],\n",
    "        k = 3, \n",
    "        metric = \"lv\"\n",
    "    ).alias(\"similar_words_from_vocab\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98264a05",
   "metadata": {},
   "source": [
    "# Using PDS Expressions On Series / NumPy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "0a42a771",
   "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>actual</th><th>predicted</th><th>0-2</th><th>0-9</th></tr><tr><td>f64</td><td>f64</td><td>i32</td><td>i32</td></tr></thead><tbody><tr><td>1.0</td><td>0.333752</td><td>1</td><td>9</td></tr><tr><td>1.0</td><td>0.766684</td><td>0</td><td>7</td></tr><tr><td>0.0</td><td>0.977014</td><td>2</td><td>2</td></tr><tr><td>1.0</td><td>0.768644</td><td>2</td><td>2</td></tr><tr><td>0.0</td><td>0.173281</td><td>0</td><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌────────┬───────────┬─────┬─────┐\n",
       "│ actual ┆ predicted ┆ 0-2 ┆ 0-9 │\n",
       "│ ---    ┆ ---       ┆ --- ┆ --- │\n",
       "│ f64    ┆ f64       ┆ i32 ┆ i32 │\n",
       "╞════════╪═══════════╪═════╪═════╡\n",
       "│ 1.0    ┆ 0.333752  ┆ 1   ┆ 9   │\n",
       "│ 1.0    ┆ 0.766684  ┆ 0   ┆ 7   │\n",
       "│ 0.0    ┆ 0.977014  ┆ 2   ┆ 2   │\n",
       "│ 1.0    ┆ 0.768644  ┆ 2   ┆ 2   │\n",
       "│ 0.0    ┆ 0.173281  ┆ 0   ┆ 2   │\n",
       "└────────┴───────────┴─────┴─────┘"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "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, 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, 1)\n",
       "┌─────────────┐\n",
       "│ jaccard_col │\n",
       "│ ---         │\n",
       "│ f64         │\n",
       "╞═════════════╡\n",
       "│ 0.3         │\n",
       "└─────────────┘"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds.eval_series(\n",
    "    df[\"0-2\"], df[\"0-9\"], # use series as args\n",
    "    expr = \"query_jaccard_col\" # name of the pds expression\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "e8c6d51b",
   "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, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>precision</th><th>recall</th><th>f</th><th>average_precision</th><th>roc_auc</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.497343</td><td>0.497721</td><td>0.497532</td><td>0.49837</td><td>0.497449</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 5)\n",
       "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
       "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
       "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
       "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
       "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
       "│ 0.497343  ┆ 0.497721 ┆ 0.497532 ┆ 0.49837           ┆ 0.497449 │\n",
       "└───────────┴──────────┴──────────┴───────────────────┴──────────┘"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds.eval_series(\n",
    "    df[\"actual\"], df[\"predicted\"], # use series as args\n",
    "    expr = \"query_binary_metrics\" # name of the pds expression\n",
    ").unnest(\"binary_metrics\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "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>cnt&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.202925</td><td>0.2</td><td>0.2</td><td>0.0</td></tr><tr><td>0.401657</td><td>0.2</td><td>0.188</td><td>0.000743</td></tr><tr><td>0.605492</td><td>0.2</td><td>0.223</td><td>0.002504</td></tr><tr><td>0.812192</td><td>0.2</td><td>0.206</td><td>0.000177</td></tr><tr><td>inf</td><td>0.2</td><td>0.183</td><td>0.00151</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌──────────┬──────────────┬────────────┬──────────┐\n",
       "│ cnt<=    ┆ baseline_pct ┆ actual_pct ┆ psi_bin  │\n",
       "│ ---      ┆ ---          ┆ ---        ┆ ---      │\n",
       "│ f64      ┆ f64          ┆ f64        ┆ f64      │\n",
       "╞══════════╪══════════════╪════════════╪══════════╡\n",
       "│ 0.202925 ┆ 0.2          ┆ 0.2        ┆ 0.0      │\n",
       "│ 0.401657 ┆ 0.2          ┆ 0.188      ┆ 0.000743 │\n",
       "│ 0.605492 ┆ 0.2          ┆ 0.223      ┆ 0.002504 │\n",
       "│ 0.812192 ┆ 0.2          ┆ 0.206      ┆ 0.000177 │\n",
       "│ inf      ┆ 0.2          ┆ 0.183      ┆ 0.00151  │\n",
       "└──────────┴──────────────┴────────────┴──────────┘"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds.eval_series(\n",
    "    # can also use NumPy\n",
    "    np.random.random(size = 1000), \n",
    "    np.random.random(size = 1000), \n",
    "    expr = \"query_psi\", # name of the pds expression\n",
    "    n_bins = 5, \n",
    "    return_report = True\n",
    ").unnest(\"psi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>cid_ce</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>12.610733</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌───────────┐\n",
       "│ cid_ce    │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ 12.610733 │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds.eval_series(\n",
    "    np.random.random(size = 1000), \n",
    "    expr = \"query_cid_ce\", # name of the pds expression\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "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, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>c3_stats</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.128507</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌──────────┐\n",
       "│ c3_stats │\n",
       "│ ---      │\n",
       "│ f64      │\n",
       "╞══════════╡\n",
       "│ 0.128507 │\n",
       "└──────────┘"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds.eval_series(\n",
    "    pl.Series(values=np.random.random(size = 1000)), \n",
    "    expr = \"query_c3_stats\", # name of the pds expression\n",
    "    lag = 3\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc7ee650",
   "metadata": {},
   "source": [
    "# Examples using Ball Tree Features\n",
    "\n",
    "Not currently available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "cc89e553",
   "metadata": {},
   "outputs": [],
   "source": [
    "# size = 1_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",
    "# )\n",
    "# df = df.with_columns(pl.int_range(0, size).cast(dtype=pl.UInt32).alias(\"index\"))\n",
    "# df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "728a9bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Pointwise Nearest Neighbors\n",
    "# df.with_columns(\n",
    "#     pds.query_bt_knn_ptwise(\n",
    "#         pl.col(\"x2\"), \n",
    "#         pl.col(\"x3\"), \n",
    "#         index=pl.col(\"index\"),\n",
    "#         r=999.0,\n",
    "#         k =5,\n",
    "#         distance_metric=\"haversine\",\n",
    "#         parallel=True).alias(\"ball_tree_knn_ptwise\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "74f1e763",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Pointwise Nearest Neighbors with distances\n",
    "# # we get a struct with an ids column and a dist column\n",
    "\n",
    "# df.with_columns(\n",
    "#     pds.query_bt_knn_ptwise(\n",
    "#         pl.col(\"x2\"), \n",
    "#         pl.col(\"x3\"), \n",
    "#         index=pl.col(\"index\"),\n",
    "#         r=999.0,\n",
    "#         k =5,\n",
    "#         distance_metric=\"haversine\",\n",
    "#         return_dist=True,\n",
    "#         parallel=True).alias(\"ball_tree_knn_ptwise\")\n",
    "# ).unnest(\"ball_tree_knn_ptwise\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "d9ce8e89",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Frequency Count\n",
    "# df.with_columns(\n",
    "#     pds.query_bt_knn_radius_freq_cnt(\n",
    "#         pl.col(\"x2\"), \n",
    "#         pl.col(\"x3\"), \n",
    "#         index=pl.col(\"index\"), \n",
    "#         r=999.0, \n",
    "#         k =5, \n",
    "#         distance_metric=\"haversine\", \n",
    "#         parallel=True\n",
    "#         ).alias(\"ball_tree_knn_radius_freq_cnt\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "ddf08cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Average distance of k nearest neighbors\n",
    "# df.with_columns(\n",
    "#     pds.query_bt_knn_avg(\n",
    "#         pl.col(\"x2\"), \n",
    "#         pl.col(\"x3\"), \n",
    "#         pl.col(\"x1\"), \n",
    "#         index=pl.col(\"index\"), \n",
    "#         r=999.0, \n",
    "#         k =1,\n",
    "#         distance_metric=\"euclidean\", \n",
    "#         parallel=True).alias(\"btree_knn_avg\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "195210f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Count neighbors within radius\n",
    "# df.with_columns(\n",
    "#     pds.query_bt_nb_cnt(\n",
    "#         pl.col(\"x2\"),\n",
    "#         pl.col(\"x3\"),\n",
    "#         index=pl.col(\"index\"),\n",
    "#         r=18.0, \n",
    "#         distance_metric=\"haversine\",\n",
    "#         parallel=True).alias(\"btree_within_radius\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "026350c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # returns true if the row is within k nearest neighbors of the given point\n",
    "# df.with_columns(\n",
    "#     pds.bt_within_dist_from(\n",
    "#         pl.col(\"x2\"),\n",
    "#         pl.col(\"x3\"),\n",
    "#         pt=[123,78.99],\n",
    "#         r=999999, \n",
    "#         distance_metric=\"haversine\",\n",
    "#         parallel=True).alias(\"btree_within_dist_from\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "dbe0983f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # check if a point is included in the k nearest neighbors of the given point\n",
    "# # Note: This is an exact match so your point could be quite close but now show up\n",
    "# # We can pass in an EPSILON to account for this. THis is defaulted to the EPSILON of the data type in Rust\n",
    "\n",
    "# # pick a random point\n",
    "# point = df.select(pl.col(\"x2\"), pl.col(\"x3\"))\n",
    "# rr = list(point.row(2))\n",
    "# df.with_columns(\n",
    "#     pds.is_bt_knn_from(\n",
    "#         pl.col(\"x2\"),\n",
    "#         pl.col(\"x3\"),\n",
    "#         pt=rr, \n",
    "#         k=56, \n",
    "#         distance_metric=\"haversine\",\n",
    "#         parallel=True, epsilon=0.5).alias(\"btree_nn_within\")\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cc8f20",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}