0_notebook_return_in_parts.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7a957ef7-df70-40ae-ad35-46a7aa249c4b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "dbutils.widgets.text(\"sql_query\", \"\")\n",
    "dbutils.widgets.text(\"database\", \"\")\n",
    "sql_query = dbutils.widgets.get(\"sql_query\")\n",
    "database = dbutils.widgets.get(\"database\")\n",
    "\n",
    "if not database:\n",
    "  database = \"ifrs17_cur\"\n",
    "\n",
    "spark.sql(f\"USE {database}\")\n",
    "\n",
    "if not sql_query:\n",
    "  sql_query = \"SELECT * FROM data_ifrs17 LIMIT 10\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "32ef2628-62a0-4030-8f2e-a282a866d40d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.sql(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e87f7932-3ea9-4204-9fe2-52a6bcb30f18",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import hashlib\n",
    "import time\n",
    "\n",
    "# Generate a unique hash for the operation\n",
    "timestamp = str(int(time.time() * 10000))\n",
    "short_hash = hashlib.md5(timestamp.encode()).hexdigest()[:6]\n",
    "\n",
    "# Base folder for chunk files\n",
    "base_folder = f\"dbfs:/tmp/api_responses/result_{short_hash}.csv\"\n",
    "summary_folder = f\"dbfs:/tmp/api_responses/result_{short_hash}_summary.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f7201fc3-74ad-4595-b977-e3b111c37773",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Count total records\n",
    "record_count = df.count()\n",
    "\n",
    "# Calculate the number of partitions based on a fixed chunk size\n",
    "chunk_size = 1000\n",
    "num_partitions = max(1, -(-record_count // chunk_size))  # Ceiling division\n",
    "\n",
    "# Split the DataFrame into chunks of 1000 records\n",
    "df_repartitioned = df.repartition(num_partitions)  # Adjust the number of partitions based on data size\n",
    "\n",
    "# Write the chunk files to the base folder\n",
    "df_repartitioned.write.mode(\"overwrite\").option(\"header\", \"true\").csv(base_folder)\n",
    "\n",
    "# Collect the paths of chunk files\n",
    "chunk_files = [\n",
    "    file.path\n",
    "    for file in dbutils.fs.ls(base_folder)\n",
    "    if file.path.endswith(\".csv\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the summary CSV with file paths only\n",
    "summary_data = [(chunk_file,) for chunk_file in chunk_files]\n",
    "summary_df = spark.createDataFrame(summary_data, [\"csv_file_path\"])\n",
    "\n",
    "# Write the summary file to the summary folder\n",
    "summary_df.coalesce(1).write.mode(\"overwrite\").option(\"header\", \"true\").csv(summary_folder)\n",
    "\n",
    "# Collect the actual part file path of the summary\n",
    "summary_files = [\n",
    "    file.path\n",
    "    for file in dbutils.fs.ls(summary_folder)\n",
    "    if file.path.endswith(\".csv\")\n",
    "]\n",
    "\n",
    "# Ensure the notebook output points to the full path of the summary file part\n",
    "if len(summary_files) == 1:\n",
    "    dbutils.notebook.exit(summary_files[0])\n",
    "else:\n",
    "    raise ValueError(\"Error generating summary file.\")"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": null,
   "language": "python",
   "notebookMetadata": {
    "mostRecentlyExecutedCommandWithImplicitDF": {
     "commandId": 2697648055580681,
     "dataframes": [
      "_sqldf"
     ]
    },
    "pythonIndentUnit": 2
   },
   "notebookName": "return_query_results_api",
   "widgets": {
    "database": {
     "currentValue": "",
     "nuid": "b4d7ae36-22bf-4570-93e6-cd2d12dadc96",
     "typedWidgetInfo": {
      "autoCreated": false,
      "defaultValue": "",
      "label": null,
      "name": "database",
      "options": {
       "validationRegex": null,
       "widgetDisplayType": "Text"
      },
      "parameterDataType": "String"
     },
     "widgetInfo": {
      "defaultValue": "",
      "label": null,
      "name": "database",
      "options": {
       "autoCreated": null,
       "validationRegex": null,
       "widgetType": "text"
      },
      "widgetType": "text"
     }
    },
    "sql_query": {
     "currentValue": "",
     "nuid": "894807d3-e76e-4c16-b287-1c05d6d7e4f8",
     "typedWidgetInfo": {
      "autoCreated": false,
      "defaultValue": "",
      "label": null,
      "name": "sql_query",
      "options": {
       "validationRegex": null,
       "widgetDisplayType": "Text"
      },
      "parameterDataType": "String"
     },
     "widgetInfo": {
      "defaultValue": "",
      "label": null,
      "name": "sql_query",
      "options": {
       "autoCreated": null,
       "validationRegex": null,
       "widgetType": "text"
      },
      "widgetType": "text"
     }
    }
   }
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}