{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39e1cac5",
   "metadata": {},
   "source": [
    "# How to read a Parquet file with Python\n",
    "\n",
    "Load Parquet into pandas with Rugo, read only the columns and rows you need, look inside a file before reading it, and query Parquet with SQL using Opteryx.\n",
    "\n",
    "This notebook accompanies [How to read a Parquet file with Python](https://docs.opteryx.app/docs/guides/read-a-parquet-file-with-python) on docs.opteryx.app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d8d573",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q rugo opteryx-core pandas pyarrow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49f9e218",
   "metadata": {},
   "source": [
    "## Get the data\n",
    "\n",
    "`space_missions.parquet`: 4,630 orbital launch attempts since Sputnik."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c987028",
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "\n",
    "urllib.request.urlretrieve(\"https://raw.githubusercontent.com/mabel-dev/opteryx-core/main/testdata/missions/space_missions.parquet\", \"space_missions.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "071c8567",
   "metadata": {},
   "source": [
    "## Load a Parquet file into pandas\n",
    "\n",
    "`read_parquet` yields one morsel per row group. Convert each to Arrow, concatenate, and convert once to pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82b9f268",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow\n",
    "from rugo import parquet\n",
    "\n",
    "with parquet.read_parquet(\"space_missions.parquet\") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df.dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a46d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "513aa657",
   "metadata": {},
   "source": [
    "## Read only the columns you need\n",
    "\n",
    "A column you don't ask for is never read from disk or decompressed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea6f390",
   "metadata": {},
   "outputs": [],
   "source": [
    "with parquet.read_parquet(\"space_missions.parquet\", columns=[\"Company\", \"Mission\", \"Price\"]) as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdb62bbc",
   "metadata": {},
   "source": [
    "## Filter while reading\n",
    "\n",
    "`predicates` is a list of `(column, operator, value)` tuples, combined with AND."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb27a221",
   "metadata": {},
   "outputs": [],
   "source": [
    "with parquet.read_parquet(\n",
    "    \"space_missions.parquet\",\n",
    "    columns=[\"Company\", \"Mission\", \"Price\"],\n",
    "    predicates=[(\"Company\", \"=\", \"SpaceX\"), (\"Price\", \"<\", 60)],\n",
    ") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daa8c9d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# `in`, `!=`, and a predicate column that isn't in `columns`\n",
    "with parquet.read_parquet(\n",
    "    \"space_missions.parquet\",\n",
    "    columns=[\"Company\", \"Mission\"],\n",
    "    predicates=[(\"Company\", \"in\", [\"SpaceX\", \"Rocket Lab\"]), (\"Mission_Status\", \"!=\", \"Success\")],\n",
    ") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c74eb45",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "with parquet.read_parquet(\n",
    "    \"space_missions.parquet\",\n",
    "    columns=[\"Mission\"],\n",
    "    predicates=[(\"Lauched_at\", \">=\", datetime.datetime(2020, 1, 1))],\n",
    ") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28b50abc",
   "metadata": {},
   "source": [
    "## Look inside a file first\n",
    "\n",
    "The `rugo` command reads the footer — schema, row counts, statistics — without reading the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1d7d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rugo info space_missions.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e993f74",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rugo schema space_missions.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343502c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rugo inspect space_missions.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdc5f88b",
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = parquet.read_metadata(\"space_missions.parquet\")\n",
    "print(meta.num_rows, \"rows in\", meta.num_row_groups, \"row group(s)\")\n",
    "for column in meta.schema_columns:\n",
    "    print(column.name, column.logical_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2ce03ea",
   "metadata": {},
   "source": [
    "## Files larger than memory\n",
    "\n",
    "Work through a file one morsel (one row group) at a time, keeping only what you need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4212ad3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "totals = {}\n",
    "with parquet.read_parquet(\"space_missions.parquet\", columns=[\"Company\", \"Price\"]) as reader:\n",
    "    for morsel in reader:\n",
    "        chunk = morsel.to_arrow().to_pandas()\n",
    "        for company, price in chunk.groupby(\"Company\")[\"Price\"].sum().items():\n",
    "            totals[company] = totals.get(company, 0) + price\n",
    "\n",
    "sorted(totals.items(), key=lambda kv: -kv[1])[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d813a16",
   "metadata": {},
   "source": [
    "## Query with SQL: Opteryx\n",
    "\n",
    "Aggregate inside the engine and only bring the answer into pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37f0ffe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opteryx\n",
    "\n",
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT Company,\n",
    "           COUNT(*)             AS launches,\n",
    "           ROUND(AVG(Price), 1) AS avg_price_musd\n",
    "      FROM READ_PARQUET('space_missions.parquet')\n",
    "     WHERE Mission_Status = 'Success'\n",
    "       AND Lauched_at >= CAST('2000-01-01' AS TIMESTAMP)\n",
    "     GROUP BY Company\n",
    "     ORDER BY launches DESC\n",
    "     LIMIT 5\n",
    "\"\"\")\n",
    "pyarrow.concat_tables(m.to_arrow() for m in morsels).to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3c34b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# EXPLAIN shows the filter pushed into the read\n",
    "morsels = opteryx.session().execute_to_morsels(\n",
    "    \"EXPLAIN SELECT Mission FROM READ_PARQUET('space_missions.parquet') WHERE Company = 'SpaceX'\"\n",
    ")\n",
    "pyarrow.concat_tables(m.to_arrow() for m in morsels).to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41be2c26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bind values as parameters rather than formatting them into the SQL\n",
    "morsels = opteryx.session().execute_to_morsels(\n",
    "    \"SELECT Mission, Price FROM READ_PARQUET('space_missions.parquet') WHERE Company = :company AND Price < :max_price\",\n",
    "    params={\"company\": \"SpaceX\", \"max_price\": 60},\n",
    ")\n",
    "pyarrow.concat_tables(m.to_arrow() for m in morsels).to_pandas().head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}