{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf18d811",
   "metadata": {},
   "source": [
    "# How to read an Avro file with Python\n",
    "\n",
    "Read an Avro file's schema, load it into pandas with Rugo, select fields including nested ones, read through a reader schema, and convert to Parquet to query with Opteryx.\n",
    "\n",
    "This notebook accompanies [How to read an Avro file with Python](https://docs.opteryx.app/docs/guides/read-an-avro-file-with-python) on docs.opteryx.app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07903004",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q rugo opteryx-core pandas pyarrow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab37ff4",
   "metadata": {},
   "source": [
    "## Get the data\n",
    "\n",
    "`space_missions.avro`: 4,630 orbital launch attempts, with each rocket as a nested record."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4971fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "\n",
    "urllib.request.urlretrieve(\"https://raw.githubusercontent.com/mabel-dev/opteryx-core/main/testdata/avro/space_missions.avro\", \"space_missions.avro\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb80839",
   "metadata": {},
   "source": [
    "## Read the schema first\n",
    "\n",
    "The schema is in the header, so nothing is decoded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30268224",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rugo.avro import read_avro, read_metadata\n",
    "\n",
    "meta = read_metadata(\"space_missions.avro\")\n",
    "print(meta.codec)\n",
    "for field in meta.schema[\"fields\"]:\n",
    "    print(field[\"name\"], field[\"type\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08602fd8",
   "metadata": {},
   "source": [
    "## Load an Avro file into pandas\n",
    "\n",
    "Logical types are decoded; the nested `Rocket` record arrives as JSON text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27cb6b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow\n",
    "\n",
    "with read_avro(\"space_missions.avro\") 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": "245ce598",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36da48d4",
   "metadata": {},
   "source": [
    "## Read only the fields you need\n",
    "\n",
    "A dotted name reads one field out of a nested record."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea4c3da4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with read_avro(\n",
    "    \"space_missions.avro\",\n",
    "    columns=[\"Company\", \"Mission\", \"Rocket.Name\", \"Rocket.Status\"],\n",
    ") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf30163",
   "metadata": {},
   "source": [
    "## Filtering rows\n",
    "\n",
    "Avro has no statistics, so there's no `predicates` argument. Filter each morsel as it arrives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e15642a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "parts = []\n",
    "with read_avro(\"space_missions.avro\", columns=[\"Company\", \"Mission\", \"Price\"]) as reader:\n",
    "    for morsel in reader:\n",
    "        chunk = morsel.to_arrow().to_pandas()\n",
    "        parts.append(chunk[(chunk[\"Company\"] == \"SpaceX\") & (chunk[\"Price\"] < 60)])\n",
    "\n",
    "df = pd.concat(parts, ignore_index=True)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd98510",
   "metadata": {},
   "source": [
    "## Read through a reader schema\n",
    "\n",
    "Drop fields, add fields with defaults, and reshape nested records."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf39e7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader_schema = {\n",
    "    \"type\": \"record\",\n",
    "    \"name\": \"Launch\",\n",
    "    \"fields\": [\n",
    "        {\"name\": \"Mission\", \"type\": \"string\"},\n",
    "        {\"name\": \"Price\", \"type\": [\"null\", \"double\"], \"default\": None},\n",
    "        {\"name\": \"Source\", \"type\": \"string\", \"default\": \"space_missions\"},\n",
    "        {\"name\": \"Rocket\", \"type\": {\"type\": \"record\", \"name\": \"Rocket\", \"fields\": [\n",
    "            {\"name\": \"Name\", \"type\": \"string\"},\n",
    "            {\"name\": \"Stages\", \"type\": [\"null\", \"int\"], \"default\": None},\n",
    "        ]}},\n",
    "    ],\n",
    "}\n",
    "\n",
    "with read_avro(\"space_missions.avro\", reader_schema=reader_schema) as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea273e4",
   "metadata": {},
   "source": [
    "## Query with SQL\n",
    "\n",
    "Opteryx reads Avro directly with `READ_AVRO`, with `->>` reaching into the nested record."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b7f6b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opteryx\n",
    "\n",
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT Company, COUNT(*) AS launches\n",
    "      FROM READ_AVRO('space_missions.avro')\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": "31471ad9",
   "metadata": {},
   "source": [
    "## Convert to Parquet for repeated queries\n",
    "\n",
    "`READ_AVRO` reads every row, because Avro has no statistics to skip data with. If you will query the data more than once, convert it once with Rugo and query the Parquet copy. From the shell this is `rugo convert space_missions.avro space_missions.parquet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1be345c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rugo import parquet\n",
    "\n",
    "with read_avro(\"space_missions.avro\") as reader:\n",
    "    with open(\"space_missions.parquet\", \"wb\") as f:\n",
    "        with parquet.open_parquet_writer(f.write) as writer:\n",
    "            for morsel in reader:\n",
    "                writer.write_row_group(morsel)\n",
    "\n",
    "!rugo schema space_missions.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207b9cc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opteryx\n",
    "\n",
    "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
}