{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tmlt.analytics.query_builder import QueryBuilder\n",
    "from tmlt.analytics.privacy_budget import PureDPBudget\n",
    "from tmlt.analytics.session import Session\n",
    "from tmlt.analytics.keyset import KeySet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "import pandas as pd\n",
    "spark = SparkSession.builder.getOrCreate()\n",
    "private_data = spark.createDataFrame(pd.DataFrame([[\"0\", 1, 0], [\"1\", 0, 1]], columns=[\"A\", \"B\", \"X\"]))\n",
    "public_data = spark.createDataFrame(pd.DataFrame([[\"0\", 1], [\"1\", 0]], columns=[\"A\", \"Y\"]))\n",
    "total_budget = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = Session.from_dataframe(\n",
    "    privacy_budget=PureDPBudget(total_budget),\n",
    "    source_id=\"my_private_data\",\n",
    "    dataframe=private_data,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.private_sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.add_public_dataframe(\"my_public_data\", public_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.public_sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create View"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "join_query = QueryBuilder(\"my_private_data\").join_public(\"my_public_data\")\n",
    "sess.create_view(join_query, \"joined_data\", cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.private_sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.get_schema(\"joined_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_query = QueryBuilder(\"joined_data\").groupby(KeySet.from_dict({\"Y\": [0, 1]})).count()\n",
    "sess.evaluate(count_query, privacy_budget=PureDPBudget(total_budget/5)).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.remaining_privacy_budget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partition the Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions = sess.partition_and_create(\n",
    "    \"joined_data\", \n",
    "    privacy_budget=PureDPBudget(3),\n",
    "    column=\"Y\", \n",
    "    splits={\"partition_0\": 0, \"partition_1\": 1}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.remaining_privacy_budget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions[\"partition_0\"].remaining_privacy_budget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions[\"partition_0\"].private_sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions[\"partition_0\"].get_schema(\"partition_0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions[\"partition_1\"].remaining_privacy_budget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sessions[\"partition_1\"].private_sources"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
