{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required Python packages\n",
    "!pip install orjson tqdm matplotlib pyarrow polars\n",
    "N = 10_000_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "\n",
    "if not os.path.exists(\"/mnt/ramdisk/identical_test.jsonl\"):\n",
    "    with open(\"/mnt/ramdisk/identical_test.jsonl\", \"w\") as f:\n",
    "        for i in tqdm(range(N)):\n",
    "            line = {\n",
    "                \"id\": i,\n",
    "                \"text\": f\"This is a test record {i}\",\n",
    "                \"tags\": [\"jazz\", \"funk\", f\"genre_{i}\"],\n",
    "                \"audio\": {\n",
    "                    \"transcript\": f\"This is a test transcript {i} \" * 50,\n",
    "                    \"audio_path\": f\"path/to/audio_{i}.wav\",\n",
    "                    \"audio_url\": f\"https://example.com/audio/{i}\",\n",
    "                    \"audio_duration\": 1.0 + (i % 100) / 10,\n",
    "                    \"audio_sample_rate\": 16000,\n",
    "                    \"energy\": [0.2 + (i % 10) / 100] * 60,\n",
    "                },\n",
    "            }\n",
    "            line_s = json.dumps(line)\n",
    "            f.write(line_s + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import orjson\n",
    "\n",
    "# this takes 2:07 minutes on the cluster\n",
    "\n",
    "python_metadata = []\n",
    "with open(\"/mnt/ramdisk/identical_test.jsonl\", \"r\") as f:\n",
    "    for line in tqdm(f, total=N):\n",
    "        python_data = orjson.loads(line)\n",
    "        python_metadata.append(python_data)\n",
    "print(python_metadata[-1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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, \"--input\", file_path, \"--in-memory\", \"--output-dir\", \"/mnt/ramdisk/identical_test_output\"],\n",
    "        text=True,\n",
    "        capture_output=True\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\", \"/mnt/ramdisk/identical_test.jsonl\")\n",
    "\n",
    "\n",
    "# Store for later use\n",
    "%store cpp_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "\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/ramdisk/identical_test_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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import List, Dict, Any, Tuple\n",
    "\n",
    "def convert_numpy_to_python(obj: Any) -> Any:\n",
    "    \"\"\"\n",
    "    Recursively converts NumPy types to Python native types.\n",
    "    \n",
    "    Args:\n",
    "        obj: Any object that might contain NumPy types\n",
    "        \n",
    "    Returns:\n",
    "        The same structure with NumPy types converted to Python types\n",
    "    \"\"\"\n",
    "    # Handle numpy arrays\n",
    "    if isinstance(obj, np.ndarray):\n",
    "        return convert_numpy_to_python(obj.tolist())\n",
    "    \n",
    "    # Handle numpy scalars\n",
    "    elif np.isscalar(obj) and isinstance(obj, np.generic):\n",
    "        return obj.item()\n",
    "    \n",
    "    # Handle dictionaries\n",
    "    elif isinstance(obj, dict):\n",
    "        return {k: convert_numpy_to_python(v) for k, v in obj.items()}\n",
    "    \n",
    "    # Handle lists or tuples\n",
    "    elif isinstance(obj, (list, tuple)):\n",
    "        return [convert_numpy_to_python(item) for item in obj]\n",
    "    \n",
    "    # Return other types unchanged\n",
    "    else:\n",
    "        return obj\n",
    "\n",
    "def compare_dicts(dict1: Dict, dict2: Dict) -> Tuple[bool, Dict]:\n",
    "    \"\"\"\n",
    "    Compares two dictionaries and returns a diff dictionary showing the differences.\n",
    "    \n",
    "    Args:\n",
    "        dict1: First dictionary\n",
    "        dict2: Second dictionary\n",
    "        \n",
    "    Returns:\n",
    "        Tuple of (is_equal, diff_dict)\n",
    "        - is_equal: Boolean indicating if dictionaries are equal\n",
    "        - diff_dict: Dictionary showing differences (only keys that differ)\n",
    "    \"\"\"\n",
    "    all_keys = set(dict1.keys()) | set(dict2.keys())\n",
    "    diff = {}\n",
    "    is_equal = True\n",
    "    \n",
    "    for key in all_keys:\n",
    "        # Check if key exists in both\n",
    "        if key not in dict1:\n",
    "            diff[key] = {\"in_dict1\": None, \"in_dict2\": dict2[key]}\n",
    "            is_equal = False\n",
    "        elif key not in dict2:\n",
    "            diff[key] = {\"in_dict1\": dict1[key], \"in_dict2\": None}\n",
    "            is_equal = False\n",
    "        # Check if values are equal\n",
    "        elif dict1[key] != dict2[key]:\n",
    "            diff[key] = {\"in_dict1\": dict1[key], \"in_dict2\": dict2[key]}\n",
    "            is_equal = False\n",
    "            \n",
    "    return is_equal, diff\n",
    "\n",
    "def compare_record_arrays(python_metadata: List[Dict], records: List[Dict]) -> List[Dict]:\n",
    "    \"\"\"\n",
    "    Compares two arrays of dictionaries element-wise after converting NumPy types.\n",
    "    \n",
    "    Args:\n",
    "        python_metadata: First array of dictionaries with Python native types\n",
    "        records: Second array of dictionaries with NumPy types\n",
    "        \n",
    "    Returns:\n",
    "        List of comparison results, each containing:\n",
    "        - index: The index in the arrays\n",
    "        - is_equal: Boolean indicating if dictionaries are equal\n",
    "        - differences: Dictionary showing differences (only for non-equal entries)\n",
    "    \"\"\"\n",
    "    # First, convert all NumPy types in records to Python native types\n",
    "    python_records = [convert_numpy_to_python(record) for record in records]\n",
    "    \n",
    "    print(\"finished converting numpy to python for records\")\n",
    "    # Check if arrays have the same length\n",
    "    if len(python_metadata) != len(python_records):\n",
    "        return False\n",
    "    \n",
    "    # Compare dictionaries element-wise\n",
    "    results = []\n",
    "    for i, (metadata, record) in enumerate(zip(python_metadata, python_records)):\n",
    "        is_equal, diff = compare_dicts(metadata, record)\n",
    "        # result = {\n",
    "        #     \"index\": i,\n",
    "        #     \"is_equal\": is_equal\n",
    "        # }\n",
    "        if not is_equal:\n",
    "            # result[\"differences\"] = diff\n",
    "            return False\n",
    "        else:\n",
    "            print(f\"Record {i} matches\")\n",
    "        # results.append(result)\n",
    "    \n",
    "    return True\n",
    "\n",
    "# Compare the arrays\n",
    "results = compare_record_arrays(python_metadata, records)\n",
    "print(results)\n",
    "\n",
    "# # Print results\n",
    "# for result in results:\n",
    "#     if result[\"is_equal\"]:\n",
    "#         print(f\"Record {result['index']} matches\")\n",
    "#     else:\n",
    "#         print(f\"Record {result['index']} has differences:\")\n",
    "#         for key, diff in result[\"differences\"].items():\n",
    "#             print(f\"  Key '{key}': {diff['in_dict1']} vs {diff['in_dict2']}\")\n",
    "\n",
    "# Example usage:\n",
    "# python_metadata = [{\"id\": 1, \"name\": \"Alice\"}, {\"id\": 2, \"name\": \"Bob\"}]\n",
    "# records = [{\"id\": np.int64(1), \"name\": \"Alice\"}, {\"id\": np.int64(2), \"name\": np.array(\"Bob\")}]\n",
    "# results = compare_record_arrays(python_metadata, records)\n",
    "# print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results)"
   ]
  }
 ],
 "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
}