{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "569d101c",
   "metadata": {},
   "source": [
    "# How to read a JSON Lines file with Python\n",
    "\n",
    "Load JSONL into pandas with Rugo, get the column types right, filter while reading, handle malformed lines, and query JSONL with SQL using Opteryx.\n",
    "\n",
    "This notebook accompanies [How to read a JSON Lines file with Python](https://docs.opteryx.app/docs/guides/read-a-json-lines-file-with-python) on docs.opteryx.app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e30e8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q rugo opteryx-core pandas pyarrow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a24cd9c6",
   "metadata": {},
   "source": [
    "## Get the data\n",
    "\n",
    "Download `space_missions.parquet` (4,630 orbital launch attempts) and convert it to JSONL with the `rugo` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f919fd",
   "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\")\n",
    "\n",
    "!rugo convert space_missions.parquet space_missions.jsonl\n",
    "!head -n 2 space_missions.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7b83735",
   "metadata": {},
   "source": [
    "## Load a JSONL file into pandas\n",
    "\n",
    "Look at the types: `Price` should be a number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f285708",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow\n",
    "from rugo.jsonl import read_jsonl\n",
    "\n",
    "with read_jsonl(\"space_missions.jsonl\") as reader:\n",
    "    df = pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas()\n",
    "\n",
    "df.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdcbbc57",
   "metadata": {},
   "source": [
    "## Getting the types right\n",
    "\n",
    "Types are inferred from the first 5 records. The first launches have no price, so `Price` falls back to text. Look at more records, or declare the type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8634dd9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with read_jsonl(\"space_missions.jsonl\", infer_sample_size=500) 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": "22b0c655",
   "metadata": {},
   "outputs": [],
   "source": [
    "with read_jsonl(\"space_missions.jsonl\", explicit_schema={\"Price\": \"DOUBLE\"}) 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": "1cb571a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a declared column is parsed strictly\n",
    "with open(\"na.jsonl\", \"w\") as f:\n",
    "    f.write('{\"Price\": 1.5}\\n{\"Price\": \"N/A\"}\\n')\n",
    "\n",
    "try:\n",
    "    with read_jsonl(\"na.jsonl\", explicit_schema={\"Price\": \"DOUBLE\"}) as reader:\n",
    "        list(reader)\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c6e1a9",
   "metadata": {},
   "source": [
    "## Timestamps\n",
    "\n",
    "Timestamps arrive as text. These carry a `+00:00` offset, so convert them after reading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80345018",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df[\"Lauched_at\"] = pd.to_datetime(df[\"Lauched_at\"], utc=True)\n",
    "df[\"Lauched_at\"].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a187d7",
   "metadata": {},
   "source": [
    "## Read only the columns and rows you need\n",
    "\n",
    "Non-matching records never become columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49cff84b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with read_jsonl(\n",
    "    \"space_missions.jsonl\",\n",
    "    columns=[\"Company\", \"Mission\", \"Price\"],\n",
    "    predicates=[(\"Company\", \"==\", \"SpaceX\"), (\"Price\", \"<\", 60)],\n",
    "    infer_sample_size=500,\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": "a587230a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a predicate value must match the column's JSON type; it is never converted\n",
    "try:\n",
    "    with read_jsonl(\"space_missions.jsonl\", predicates=[(\"Price\", \"<\", \"60\")]) as reader:\n",
    "        list(reader)\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d214226f",
   "metadata": {},
   "source": [
    "## Malformed lines\n",
    "\n",
    "By default a bad line fails the read. `fail_on_error=False` skips it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "095721ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"events.jsonl\", \"w\") as f:\n",
    "    f.write('{\"a\": 1}\\n{\"a\": 2\\n{\"a\": 3}\\n')\n",
    "\n",
    "try:\n",
    "    with read_jsonl(\"events.jsonl\") as reader:\n",
    "        list(reader)\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3f4d17",
   "metadata": {},
   "outputs": [],
   "source": [
    "with read_jsonl(\"events.jsonl\", fail_on_error=False) as reader:\n",
    "    print(pyarrow.concat_tables(morsel.to_arrow() for morsel in reader).to_pandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "199dbac4",
   "metadata": {},
   "source": [
    "## Query with SQL: Opteryx\n",
    "\n",
    "`READ_JSONL` uses the same inference, with the same options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e30fa61d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opteryx\n",
    "\n",
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT Mission_Status, COUNT(*) AS launches\n",
    "      FROM READ_JSONL('space_missions.jsonl')\n",
    "     WHERE Company = 'SpaceX'\n",
    "     GROUP BY Mission_Status\n",
    "     ORDER BY launches DESC\n",
    "\"\"\")\n",
    "pyarrow.concat_tables(m.to_arrow() for m in morsels).to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9842f104",
   "metadata": {},
   "outputs": [],
   "source": [
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT COUNT(*) AS cheap_launches\n",
    "      FROM READ_JSONL('space_missions.jsonl', infer_sample_size => 500)\n",
    "     WHERE Price < 60\n",
    "\"\"\")\n",
    "pyarrow.concat_tables(m.to_arrow() for m in morsels).to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b73cb490",
   "metadata": {},
   "outputs": [],
   "source": [
    "morsels = opteryx.session().execute_to_morsels(\"\"\"\n",
    "    SELECT EXTRACT(year FROM CAST(Lauched_at AS TIMESTAMP)) AS year, COUNT(*) AS launches\n",
    "      FROM READ_JSONL('space_missions.jsonl')\n",
    "     GROUP BY year\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": "885bf886",
   "metadata": {},
   "source": [
    "## Convert to Parquet\n",
    "\n",
    "Fix the types once, with Rugo reading the JSONL and writing the Parquet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf1a917",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rugo import parquet\n",
    "\n",
    "with read_jsonl(\"space_missions.jsonl\", explicit_schema={\"Price\": \"DOUBLE\"}) as reader:\n",
    "    with open(\"space_missions_typed.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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6404dfb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rugo schema space_missions_typed.parquet"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}