{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: orjson in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (3.10.14)\n",
      "Requirement already satisfied: tqdm in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (4.67.1)\n",
      "Requirement already satisfied: matplotlib in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (3.10.0)\n",
      "Requirement already satisfied: pyarrow in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (19.0.0)\n",
      "Requirement already satisfied: polars in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (1.24.0)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.3.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (4.55.3)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.4.8)\n",
      "Requirement already satisfied: numpy>=1.23 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/rider/.local/lib/python3.10/site-packages (from matplotlib) (24.2)\n",
      "Requirement already satisfied: pillow>=8 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (10.4.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (3.2.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (2.9.0.post0)\n",
      "Requirement already satisfied: six>=1.5 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "# Install required Python packages\n",
    "!pip install orjson tqdm matplotlib pyarrow polars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating sample JSONL file with 10000000 records...\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import orjson\n",
    "import random\n",
    "import os\n",
    "import time\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "# Create a sample JSONL file if it doesn't exist\n",
    "file_path = \"/tmp/test_data.jsonl\"\n",
    "N = 10000000  # Number of records\n",
    "print(f\"Creating sample JSONL file with {N} records...\")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# line = {\n",
    "#     \"id\": 0,\n",
    "#     \"text\": \"This is a test record 0\",\n",
    "#     \"tags\": [\"jazz\", \"funk\", \"soul\"],\n",
    "#     \"audio\": {\n",
    "#         \"transcript\": \"This is a test transcript 0\" * 50,\n",
    "#         \"audio_path\": \"This is a test audio path 0\",\n",
    "#         \"audio_url\": \"This is a test audio url 0\",\n",
    "#         \"audio_duration\": 1.0,\n",
    "#         \"audio_sample_rate\": 16000,\n",
    "#         \"energy\": [0.2] * 60,\n",
    "#     },\n",
    "# }\n",
    "# line_s = json.dumps(line)\n",
    "# with open(\"/tmp/test_data.jsonl\", \"w\") as f:\n",
    "#     for i in tqdm(range(N)):\n",
    "#         f.write(line_s + \"\\n\")\n",
    "\n",
    "# assert len(metadata) == N\n",
    "# assert metadata[-1][\"audio\"][\"transcript\"] == \"This is a test transcript 0\" * 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building C++ parser...\n",
      "CMAKE output:\n",
      "-- Found libgomp: /usr/lib/gcc/x86_64-linux-gnu/11/libgomp.so\n",
      "-- Arrow version: 19.0.0\n",
      "-- Found the Arrow shared library: /home/rider/anaconda3/envs/suno_env/lib/libarrow.so.1900.0.0\n",
      "-- Found the Arrow import library: ARROW_IMPORT_LIB-NOTFOUND\n",
      "-- Found the Arrow static library: \n",
      "-- ArrowAcero version: 19.0.0\n",
      "-- Found the ArrowAcero shared library: /home/rider/anaconda3/envs/suno_env/lib/libarrow_acero.so.1900.0.0\n",
      "-- Found the ArrowAcero import library: ARROW_ACERO_IMPORT_LIB-NOTFOUND\n",
      "-- Found the ArrowAcero static library: \n",
      "-- Parquet version: 19.0.0\n",
      "-- Found the Parquet shared library: /home/rider/anaconda3/envs/suno_env/lib/libparquet.so.1900.0.0\n",
      "-- Found the Parquet import library: PARQUET_IMPORT_LIB-NOTFOUND\n",
      "-- Found the Parquet static library: \n",
      "-- ArrowDataset version: 19.0.0\n",
      "-- Found the ArrowDataset shared library: /home/rider/anaconda3/envs/suno_env/lib/libarrow_dataset.so.1900.0.0\n",
      "-- Found the ArrowDataset import library: ARROW_DATASET_IMPORT_LIB-NOTFOUND\n",
      "-- Found the ArrowDataset static library: \n",
      "-- Building with OpenMP support\n",
      "-- Explicitly linking with system libgomp\n",
      "-- Configuring done\n",
      "-- Generating done\n",
      "-- Build files have been written to: /home/rider/suno-data-processing/jsonl-parser/build\n",
      "\n",
      "MAKE output:\n",
      "\u001b[35m\u001b[1mConsolidate compiler generated dependencies of target parser\u001b[0m\n",
      "[100%] Built target parser\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# build and execute the cpp parser\n",
    "\n",
    "import subprocess\n",
    "import os\n",
    "\n",
    "## 4. C++ Implementation\n",
    "\n",
    "# Step 1: Build the C++ parser\n",
    "print(\"Building C++ parser...\")\n",
    "\n",
    "# Navigate to the build directory and run cmake\n",
    "build_dir = \"jsonl-parser/build\"\n",
    "os.makedirs(build_dir, exist_ok=True)\n",
    "\n",
    "# Run cmake\n",
    "cmake_result = subprocess.run(\n",
    "    [\"cmake\", \"..\"], \n",
    "    cwd=build_dir,\n",
    "    capture_output=True,\n",
    "    text=True\n",
    ")\n",
    "print(f\"CMAKE output:\\n{cmake_result.stdout}\")\n",
    "if cmake_result.returncode != 0:\n",
    "    print(f\"CMAKE error:\\n{cmake_result.stderr}\")\n",
    "\n",
    "# Run make\n",
    "make_result = subprocess.run(\n",
    "    [\"make\", \"-j\"], \n",
    "    cwd=build_dir,\n",
    "    capture_output=True,\n",
    "    text=True\n",
    ")\n",
    "print(f\"MAKE output:\\n{make_result.stdout}\")\n",
    "if make_result.returncode != 0:\n",
    "    print(f\"MAKE error:\\n{make_result.stderr}\")\n",
    "# copy the parser executable to the current directory\n",
    "os.system(\"cp jsonl-parser/build/parser .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File size: 18205 MB\n",
      "Using disk-based processing\n",
      "Using 32 threads\n",
      "Finding line indices using parallel scan...\n",
      "Found 10000000 lines in 0.676 seconds\n",
      "Inferring Arrow schema from sample data...\n",
      "Schema inferred successfully, fields: 4\n",
      "Processing in Arrow format and writing to disk\n",
      "Progress: 100.0% - Processed 10000000 of 10000000 records (2187068 records/sec)                      \n",
      "\n",
      "Processed 10000000 records in 5.7 seconds\n",
      "Processing performance: 1761493.7 records/second\n",
      "Merging Arrow files from disk...\n",
      "Merging Arrow files: 100.0% - 288/288 batches | 4.0 batches/sec | ETA: calculating...        \n",
      "Merged 32 Arrow files with 288 batches in 18.3 seconds (15.77 batches/s)\n",
      "Merged Arrow files in 18.26 seconds\n",
      "Final output saved to /mnt/localdisk/parser-output/parsed_results_2025-03-14_12-32.arrow.arrow\n",
      "Total execution time: 25.18 seconds\n",
      "\n",
      "Last record: {\"id\": 0, \"text\": \"This is a test record 0\", \"tags\": [\"jazz\", \"funk\", \"soul\"], \"audio\": {\"transcript...\n",
      "Stored 'cpp_results' (dict)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages/IPython/extensions/storemagic.py:229: UserWarning: This is now an optional IPython functionality, setting autorestore/cpp_results requires you to install the `pickleshare` library.\n",
      "  db[ 'autorestore/' + arg ] = obj\n"
     ]
    }
   ],
   "source": [
    "import subprocess\n",
    "import json\n",
    "\n",
    "def run_cpp_and_capture_jsonl(executable_path, file_path):\n",
    "    # Run the C++ program and capture its stderr output\n",
    "    result = subprocess.run(\n",
    "        [executable_path, file_path, \"--arrow\", \"--merge\"],\n",
    "        text=False,\n",
    "        capture_output=False\n",
    "    )\n",
    "    \n",
    "    # Create a results object with both the parsed data and metadata\n",
    "    cpp_results = {\n",
    "        'return_code': result.returncode,\n",
    "        'stdout': result.stdout,\n",
    "        'stderr_raw': result.stderr\n",
    "    }\n",
    "    \n",
    "    return cpp_results\n",
    "\n",
    "# Usage:\n",
    "cpp_results = run_cpp_and_capture_jsonl(\"./parser\", \"/tmp/test_data.jsonl\")\n",
    "\n",
    "\n",
    "# Store for later use\n",
    "%store cpp_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Arrow table load: 0.02 seconds\n",
      "Pandas conversion: 17.03 seconds\n",
      "Dict conversion: 22.22 seconds\n",
      "Converted 10000000 records in 39.27 seconds (254656 records/sec)\n",
      "First record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n",
      "100th record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n",
      "1000th record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n",
      "10000th record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n"
     ]
    }
   ],
   "source": [
    "import pyarrow as pa\n",
    "import pyarrow.ipc as ipc\n",
    "import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "import polars as pl\n",
    "import pyarrow.ipc as ipc\n",
    "\n",
    "def fast_arrow_to_dict(file_path):\n",
    "    \"\"\"Ultra-fast conversion using column-wise operations\"\"\"\n",
    "    start_time = time.time()\n",
    "    \n",
    "    t1 = time.time()\n",
    "    with pa.memory_map(file_path, 'r') as source:\n",
    "        reader = ipc.open_file(source)\n",
    "        table = reader.read_all()\n",
    "    t2 = time.time()\n",
    "    print(f\"Arrow table load: {t2-t1:.2f} seconds\")\n",
    "    \n",
    "    df = table.to_pandas()\n",
    "    t3 = time.time()\n",
    "    print(f\"Pandas conversion: {t3-t2:.2f} seconds\")\n",
    "    \n",
    "    records = df.to_dict('records')\n",
    "    t4 = time.time()\n",
    "    print(f\"Dict conversion: {t4-t3:.2f} seconds\")\n",
    "    \n",
    "    duration = time.time() - start_time\n",
    "    print(f\"Converted {len(records)} records in {duration:.2f} seconds ({len(records)/duration:.0f} records/sec)\")\n",
    "    return records\n",
    "\n",
    "# Usage\n",
    "records = fast_arrow_to_dict(\"/mnt/localdisk/parser-output/parsed_results_2025-03-14_03-15.arrow\")\n",
    "print(f\"First record: {records[0]}\")\n",
    "print(f\"100th record: {records[100]}\")\n",
    "print(f\"1000th record: {records[1000]}\")\n",
    "print(f\"10000th record: {records[10000]}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100000th record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n",
      "1000000th record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n",
      "last record: {'id': 0, 'text': 'This is a test record 0', 'tags': array(['jazz', 'funk', 'soul'], dtype=object), 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,\n",
      "       0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2])}}\n"
     ]
    }
   ],
   "source": [
    "print(f\"100000th record: {records[100000]}\")\n",
    "print(f\"1000000th record: {records[1000000]}\")\n",
    "print(f\"last record: {records[-1]}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: orjson in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (3.10.14)\n",
      "Requirement already satisfied: tqdm in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (4.67.1)\n",
      "Requirement already satisfied: matplotlib in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (3.10.0)\n",
      "Requirement already satisfied: pyarrow in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (19.0.0)\n",
      "Requirement already satisfied: polars in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (1.24.0)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.3.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (4.55.3)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.4.8)\n",
      "Requirement already satisfied: numpy>=1.23 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/rider/.local/lib/python3.10/site-packages (from matplotlib) (24.2)\n",
      "Requirement already satisfied: pillow>=8 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (10.4.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (3.2.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from matplotlib) (2.9.0.post0)\n",
      "Requirement already satisfied: six>=1.5 in /home/rider/anaconda3/envs/suno_env/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] Failed to open local file '/mnt/localdisk/parser-output/parsed_results_2025-03-14_12-32.arrow'. Detail: [errno 2] No such file or directory",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 7\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpolars\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpl\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpyarrow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mipc\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mipc\u001b[39;00m\n\u001b[0;32m----> 7\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmemory_map\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/mnt/localdisk/parser-output/parsed_results_2025-03-14_12-32.arrow\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mr\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m source:\n\u001b[1;32m      8\u001b[0m     reader \u001b[38;5;241m=\u001b[39m ipc\u001b[38;5;241m.\u001b[39mopen_file(source)\n\u001b[1;32m      9\u001b[0m     table \u001b[38;5;241m=\u001b[39m reader\u001b[38;5;241m.\u001b[39mread_all()\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/pyarrow/io.pxi:1147\u001b[0m, in \u001b[0;36mpyarrow.lib.memory_map\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/pyarrow/io.pxi:1094\u001b[0m, in \u001b[0;36mpyarrow.lib.MemoryMappedFile._open\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/pyarrow/error.pxi:155\u001b[0m, in \u001b[0;36mpyarrow.lib.pyarrow_internal_check_status\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/pyarrow/error.pxi:92\u001b[0m, in \u001b[0;36mpyarrow.lib.check_status\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] Failed to open local file '/mnt/localdisk/parser-output/parsed_results_2025-03-14_12-32.arrow'. Detail: [errno 2] No such file or directory"
     ]
    }
   ],
   "source": [
    "# With polars (even faster)\n",
    "# Install required Python packages\n",
    "!pip install orjson tqdm matplotlib pyarrow polars\n",
    "import pyarrow as pa\n",
    "import polars as pl\n",
    "import pyarrow.ipc as ipc\n",
    "with pa.memory_map(\"/mnt/localdisk/parser-output/parsed_results_2025-03-14_12-32.arrow.arrow\", 'r') as source:\n",
    "    reader = ipc.open_file(source)\n",
    "    table = reader.read_all()\n",
    "df = pl.from_arrow(table)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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: (3, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>name</th><th>value</th><th>tags</th><th>metadata</th></tr><tr><td>i64</td><td>str</td><td>f64</td><td>list[str]</td><td>struct[3]</td></tr></thead><tbody><tr><td>9999</td><td>&quot;Item 9999&quot;</td><td>875.012753</td><td>[&quot;tag4&quot;, &quot;tag1&quot;, … &quot;tag2&quot;]</td><td>{1.7418e9,&quot;D&quot;,true}</td></tr><tr><td>10000</td><td>&quot;Item 10000&quot;</td><td>600.986184</td><td>[&quot;tag2&quot;, &quot;tag1&quot;, &quot;tag2&quot;]</td><td>{1.7418e9,&quot;D&quot;,true}</td></tr><tr><td>10001</td><td>&quot;Item 10001&quot;</td><td>501.927344</td><td>[&quot;tag4&quot;, &quot;tag2&quot;, … &quot;tag3&quot;]</td><td>{1.7418e9,&quot;B&quot;,true}</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (3, 5)\n",
       "┌───────┬────────────┬────────────┬────────────────────────────┬─────────────────────┐\n",
       "│ id    ┆ name       ┆ value      ┆ tags                       ┆ metadata            │\n",
       "│ ---   ┆ ---        ┆ ---        ┆ ---                        ┆ ---                 │\n",
       "│ i64   ┆ str        ┆ f64        ┆ list[str]                  ┆ struct[3]           │\n",
       "╞═══════╪════════════╪════════════╪════════════════════════════╪═════════════════════╡\n",
       "│ 9999  ┆ Item 9999  ┆ 875.012753 ┆ [\"tag4\", \"tag1\", … \"tag2\"] ┆ {1.7418e9,\"D\",true} │\n",
       "│ 10000 ┆ Item 10000 ┆ 600.986184 ┆ [\"tag2\", \"tag1\", \"tag2\"]   ┆ {1.7418e9,\"D\",true} │\n",
       "│ 10001 ┆ Item 10001 ┆ 501.927344 ┆ [\"tag4\", \"tag2\", … \"tag3\"] ┆ {1.7418e9,\"B\",true} │\n",
       "└───────┴────────────┴────────────┴────────────────────────────┴─────────────────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.slice(99999, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# This takes a minute\n",
    "records = df.to_dicts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25319c6bc4114069bfd562e0b90da0f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10000000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': 0, 'text': 'This is a test record 0', 'tags': ['jazz', 'funk', 'soul'], 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]}}\n"
     ]
    }
   ],
   "source": [
    "import orjson\n",
    "\n",
    "# this takes 2:07 minutes on the cluster\n",
    "\n",
    "metadata = []\n",
    "with open(\"/mnt/localdisk/test_data.jsonl\", \"r\") as f:\n",
    "    for line in tqdm(f, total=N):\n",
    "        data = orjson.loads(line)\n",
    "        metadata.append(data)\n",
    "print(metadata[-1])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
