{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19fb32da",
   "metadata": {},
   "source": [
    "# How to read an ORC file with Python\n",
    "\n",
    "Read an ORC file's footer, load it into pandas with Rugo, select columns, see how nested types come across, and convert to Parquet to query with Opteryx.\n",
    "\n",
    "This notebook accompanies [How to read an ORC file with Python](https://docs.opteryx.app/docs/guides/read-an-orc-file-with-python) on docs.opteryx.app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b8de3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q rugo opteryx-core pandas pyarrow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3baea5f",
   "metadata": {},
   "source": [
    "## Get the data\n",
    "\n",
    "`space_missions.orc`: 4,630 orbital launch attempts, with each rocket as a nested struct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c853ea12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "\n",
    "urllib.request.urlretrieve(\"https://docs.opteryx.app/data/space_missions.orc\", \"space_missions.orc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2fa7f3e",
   "metadata": {},
   "source": [
    "## Read the footer first\n",
    "\n",
    "The footer is at the end of the file, so nothing is decoded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77582b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rugo.orc import read_metadata, read_orc\n",
    "\n",
    "meta = read_metadata(\"space_missions.orc\")\n",
    "print(meta.compression, meta.num_rows, meta.stripes)\n",
    "for column in meta.columns:\n",
    "    print(column[\"name\"], column[\"kind\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b5aa615",
   "metadata": {},
   "source": [
    "## Load an ORC file into pandas\n",
    "\n",
    "`read_orc` yields one morsel per stripe. Types come from the footer; the nested `Rocket` struct arrives as JSON text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "773304b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow\n",
    "\n",
    "df = pyarrow.concat_tables(morsel.to_arrow() for morsel in read_orc(\"space_missions.orc\")).to_pandas()\n",
    "\n",
    "df.dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e797046",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c3808e",
   "metadata": {},
   "source": [
    "## Read only the columns you need\n",
    "\n",
    "Streams for the columns you leave out are skipped rather than decoded. `columns` takes top-level names only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16a10649",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pyarrow.concat_tables(\n",
    "    morsel.to_arrow()\n",
    "    for morsel in read_orc(\"space_missions.orc\", columns=[\"Company\", \"Mission\", \"Rocket\"])\n",
    ").to_pandas()\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47dbf106",
   "metadata": {},
   "source": [
    "## Take a field out of a nested struct\n",
    "\n",
    "A dotted name such as `Rocket.Name` is refused, so read `Rocket` and parse the JSON text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3e25a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "df[\"Rocket_Name\"] = df[\"Rocket\"].map(lambda r: json.loads(r)[\"Name\"] if r else None)\n",
    "df[[\"Mission\", \"Rocket_Name\"]].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb9b57e",
   "metadata": {},
   "source": [
    "## Filtering rows\n",
    "\n",
    "There is no `predicates` argument, so filter after reading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61458a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "spacex = df[df[\"Company\"] == \"SpaceX\"]\n",
    "len(spacex)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3266a992",
   "metadata": {},
   "source": [
    "## Query with SQL\n",
    "\n",
    "Opteryx reads ORC directly with `READ_ORC`, with `->>` reaching into the nested struct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2cce97",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opteryx\n",
    "\n",
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT Company, COUNT(*) AS launches\n",
    "      FROM READ_ORC('space_missions.orc')\n",
    "     WHERE Rocket->>'Status' = 'Active'\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": "markdown",
   "id": "8748d156",
   "metadata": {},
   "source": [
    "## Convert to Parquet for repeated queries\n",
    "\n",
    "If you will query the data more than once, convert it once with Rugo and query the Parquet copy. Rugo writes one row group per stripe, with the types carried across."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7168fefe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rugo import parquet\n",
    "\n",
    "with open(\"space_missions.parquet\", \"wb\") as f:\n",
    "    with parquet.open_parquet_writer(f.write) as writer:\n",
    "        for morsel in read_orc(\"space_missions.orc\"):\n",
    "            writer.write_row_group(morsel)\n",
    "\n",
    "!rugo schema space_missions.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1fffb0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT Company, COUNT(*) AS launches\n",
    "      FROM READ_PARQUET('space_missions.parquet')\n",
    "     WHERE Rocket->>'Status' = 'Active'\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()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}