python/notebooks/druid_sketch_ingestion_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#\n",
    "# Copyright 2021 Rovio Entertainment Corporation\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# PRE-REQUISITES\n",
    "\n",
    "        AWS_PROFILE=smoke\n",
    "        JAR_BUCKET=<REPLACE THIS>\n",
    "\n",
    "## IF TESTING PYTHON CHANGES MANUALLY\n",
    "\n",
    "3. Build a zip of the python wrapper:\n",
    "\n",
    "        cd python \\\n",
    "          && zip --exclude='*.pyc' --exclude='*__pycache__*' --exclude='*~' --exclude='.pytest_cache' \\\n",
    "            -FSr ../target/rovio_ingest.zip rovio_ingest ; cd ..\n",
    "\n",
    "4. Copy the zip to s3:\n",
    "\n",
    "        aws s3 --profile $AWS_PROFILE cp \\\n",
    "          target/rovio_ingest.zip \\\n",
    "          s3://$JAR_BUCKET/tmp/juho/druid/python/rovio_ingest.zip\n",
    "\n",
    "Then invert the boolean in the cell below to use it in spark_conf.\n",
    "And skip the cell that would call install_pypi_package.\n",
    "\n",
    "## IF TESTING JAR CHANGES MANUALLY:\n",
    "\n",
    "1. Build the package (shaded jar) on command line:\n",
    "\n",
    "        mvn package -DskipTests\n",
    "\n",
    "2. A) Copy the shaded jar to s3:\n",
    "\n",
    "        aws s3 --profile $AWS_PROFILE cp \\\n",
    "          target/rovio-ingest-1.0.7_spark_3.0.1-SNAPSHOT.jar \\\n",
    "          s3://$JAR_BUCKET/tmp/vivek/druid/jars/rovio-ingest-1.0.7_spark_3.0.1-SNAPSHOT.jar\n",
    "\n",
    "2. B) Copy the plain jar to s3: \n",
    "\n",
    "        aws s3 --profile $AWS_PROFILE cp \\\n",
    "          target/original-rovio-ingest-1.0.7_spark_3.0.1-SNAPSHOT.jar \\\n",
    "          s3://$JAR_BUCKET/tmp/vivek/druid/jars/original-rovio-ingest-1.0.7_spark_3.0.1-SNAPSHOT.jar\n",
    "\n",
    "Then invert the boolean in the cell below to use it in spark_conf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext sparkmagic.magics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ENV = \"cloud\"\n",
    "PREFIX = \"tmp/vivek/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "ssm_client = boto3.session.Session(profile_name=ENV).client(service_name=\"ssm\")\n",
    "\n",
    "# secrets can be added at\n",
    "# https://console.aws.amazon.com/systems-manager/parameters/?region=us-east-1\n",
    "def get_param(secret_name: str) -> str:\n",
    "    return ssm_client.get_parameter(Name=\"/dataengineering/\" + secret_name)[\"Parameter\"][\"Value\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from IPython import get_ipython\n",
    "\n",
    "def set_spark_config(conf_dict):\n",
    "    get_ipython().run_cell_magic('spark', 'config', json.dumps(conf_dict))\n",
    "\n",
    "def create_spark_session_with_host(host):\n",
    "    get_ipython().run_line_magic('spark', 'add -l python -u http://{}:8998'.format(host))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "packages_bucket = get_param(\"rovio-ingest/packages_bucket\")\n",
    "\n",
    "spark_conf = {\n",
    "  \"conf\": {\n",
    "    \"spark.hadoop.mapreduce.input.fileinputformat.input.dir.recursive\": \"true\",\n",
    "    \"spark.sql.hive.caseSensitiveInferenceMode\": \"NEVER_INFER\",\n",
    "    \"spark.pyspark.python\": \"python3\",\n",
    "    \"spark.sql.session.timeZone\": \"UTC\",\n",
    "  }\n",
    "}\n",
    "\n",
    "# Assuming AWS EMR\n",
    "if True:\n",
    "    jars_base_path = \"s3://{packages_bucket}/{PREFIX}druid/jars\"\n",
    "    jars = (\n",
    "        f\"{jars_base_path}/rovio-ingest-1.0.7_spark_3.0.1-SNAPSHOT.jar,\"\n",
    "        f\"{jars_base_path}/datasketches-hive-1.2.0.jar,\"\n",
    "        f\"{jars_base_path}/datasketches-java-4.1.0.jar,\"\n",
    "        f\"{jars_base_path}/datasketches-memory-2.0.0.jar\"\n",
    "    )\n",
    "    \n",
    "    spark_conf[\"conf\"][\"spark.pyspark.python\"] = \"python3\"\n",
    "    spark_conf[\"conf\"][\"spark.pyspark.virtualenv.enabled\"] = \"true\"\n",
    "    spark_conf[\"conf\"][\"spark.pyspark.virtualenv.type\"] = \"native\"\n",
    "    spark_conf[\"conf\"][\"spark.pyspark.virtualenv.bin.path\"] = \"/usr/bin/virtualenv\"\n",
    "    spark_conf[\"conf\"][\"spark.jars\"] = jars\n",
    "         \n",
    "\n",
    "set_spark_config(spark_conf)\n",
    "create_spark_session_with_host(get_param(\"spark3/shared/host\"))\n",
    "\n",
    "# to debug problems in session creation, see livy session logs at http://{host}:8998/ui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%spark\n",
    "\n",
    "# This extension is provided by AWS EMR.\n",
    "# If not on EMR:\n",
    "#    A) install the module with pip on the cluster before creating the spark session\n",
    "#    B) build a zip & use with spark.submit.pyFiles\n",
    "# Use the latest stable release from PyPI.\n",
    "spark.sparkContext.install_pypi_package(\"rovio-ingest\")\n",
    "# Use a specific version to install a pre-release from PyPI.\n",
    "#spark.sparkContext.install_pypi_package(\"rovio-ingest==0.0.1.dev14\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%spark\n",
    "\n",
    "import boto3\n",
    "\n",
    "ssm_client = boto3.session.Session(region_name=\"us-east-1\").client(service_name=\"ssm\")\n",
    "\n",
    "def get_param(secret_name: str) -> str:\n",
    "    return ssm_client.get_parameter(Name=\"/dataengineering/\" + secret_name)[\"Parameter\"][\"Value\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%spark\n",
    "from datetime import datetime\n",
    "from pyspark.sql import functions as f, types as t, SparkSession\n",
    "\n",
    "from py4j.java_gateway import java_import\n",
    "from rovio_ingest import DRUID_SOURCE\n",
    "from rovio_ingest.extensions.dataframe_extension import ConfKeys, add_dataframe_druid_extension\n",
    "\n",
    "# fix df.explain on EMR 6\n",
    "java_import(spark._sc._jvm, \"org.apache.spark.sql.api.python.*\")\n",
    "\n",
    "add_dataframe_druid_extension()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sketches build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------+----------+---------+-------------------+---+\n",
      "|revenue|purchase_count|    app_id|player_id|         event_date|dau|\n",
      "+-------+--------------+----------+---------+-------------------+---+\n",
      "|   30.0|            10|testclient|       p1|2023-01-01 00:00:00|  1|\n",
      "|   15.0|             1|testclient|       p2|2023-01-01 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p3|2023-01-01 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p4|2023-01-01 00:00:00|  1|\n",
      "|    5.0|             2|testclient|       p1|2023-01-02 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p2|2023-01-02 00:00:00|  1|\n",
      "+-------+--------------+----------+---------+-------------------+---+"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "\n",
    "schema = 'revenue:DOUBLE, purchase_count:bigint, app_id:STRING, player_id STRING, event_date:TIMESTAMP'\n",
    "df = spark.createDataFrame([[30.0, 10, 'testclient', \"p1\", datetime(2023, 1, 1)],\n",
    "                            [15.0, 1, 'testclient', \"p2\", datetime(2023, 1, 1)],\n",
    "                            [0.0, 0, 'testclient', \"p3\", datetime(2023, 1, 1)],\n",
    "                            [0.0, 0, 'testclient', \"p4\", datetime(2023, 1, 1)],\n",
    "                            [5.0, 2, 'testclient', \"p1\", datetime(2023, 1, 2)],\n",
    "                            [0.0, 0, 'testclient', \"p2\", datetime(2023, 1, 2)]],\n",
    "                            schema)\n",
    "df = df.withColumn('dau', f.lit(1))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __PARTITION_NUM__#71]\n",
      "+- RepartitionByExpression [__PARTITION_TIME__#50, __PARTITION_NUM__#71]\n",
      "   +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62, cast((cast((__num_rows__#62 - 1) as double) / cast(5000000 as double)) as int) AS __PARTITION_NUM__#71]\n",
      "      +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62]\n",
      "         +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62, __num_rows__#62]\n",
      "            +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#50, __PARTITION_TIME__#50 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#62], [__PARTITION_TIME__#50], [__PARTITION_TIME__#50 ASC NULLS FIRST]\n",
      "               +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50]\n",
      "                  +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, cast((cast(if (isnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) null else UDF(knownnotnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) as double) / cast(1000 as double)) as timestamp) AS __PARTITION_TIME__#50]\n",
      "                     +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10]\n",
      "                        +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "revenue: double, purchase_count: bigint, app_id: string, player_id: string, event_date: timestamp, dau: int, __PARTITION_TIME__: timestamp, __PARTITION_NUM__: int\n",
      "Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __PARTITION_NUM__#71]\n",
      "+- RepartitionByExpression [__PARTITION_TIME__#50, __PARTITION_NUM__#71]\n",
      "   +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62, cast((cast((__num_rows__#62 - 1) as double) / cast(5000000 as double)) as int) AS __PARTITION_NUM__#71]\n",
      "      +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62]\n",
      "         +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50, __num_rows__#62, __num_rows__#62]\n",
      "            +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#50, __PARTITION_TIME__#50 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#62], [__PARTITION_TIME__#50], [__PARTITION_TIME__#50 ASC NULLS FIRST]\n",
      "               +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, __PARTITION_TIME__#50]\n",
      "                  +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, dau#10, cast((cast(if (isnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) null else UDF(knownnotnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) as double) / cast(1000 as double)) as timestamp) AS __PARTITION_TIME__#50]\n",
      "                     +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10]\n",
      "                        +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "RepartitionByExpression [__PARTITION_TIME__#50, __PARTITION_NUM__#71]\n",
      "+- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10, __PARTITION_TIME__#50, cast((cast((__num_rows__#62 - 1) as double) / 5000000.0) as int) AS __PARTITION_NUM__#71]\n",
      "   +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#50, __PARTITION_TIME__#50 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#62], [__PARTITION_TIME__#50], [__PARTITION_TIME__#50 ASC NULLS FIRST]\n",
      "      +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, if (isnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) null else cast((cast(UDF(knownnotnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) as double) / 1000.0) as timestamp) AS __PARTITION_TIME__#50]\n",
      "         +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- Exchange hashpartitioning(__PARTITION_TIME__#50, __PARTITION_NUM__#71, 1000), REPARTITION_BY_COL, [id=#35]\n",
      "   +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10, __PARTITION_TIME__#50, cast((cast((__num_rows__#62 - 1) as double) / 5000000.0) as int) AS __PARTITION_NUM__#71]\n",
      "      +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#50, __PARTITION_TIME__#50 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#62], [__PARTITION_TIME__#50], [__PARTITION_TIME__#50 ASC NULLS FIRST]\n",
      "         +- Sort [__PARTITION_TIME__#50 ASC NULLS FIRST, __PARTITION_TIME__#50 ASC NULLS FIRST], false, 0\n",
      "            +- Exchange hashpartitioning(__PARTITION_TIME__#50, 1000), ENSURE_REQUIREMENTS, [id=#31]\n",
      "               +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, if (isnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) null else cast((cast(UDF(knownnotnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) as double) / 1000.0) as timestamp) AS __PARTITION_TIME__#50]\n",
      "                  +- Scan ExistingRDD[revenue#0,purchase_count#1L,app_id#2,player_id#3,event_date#4]"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "\n",
    "from py4j.java_gateway import java_import\n",
    "from rovio_ingest import DRUID_SOURCE\n",
    "from rovio_ingest.extensions.dataframe_extension import ConfKeys, add_dataframe_druid_extension\n",
    "\n",
    "# fix df.explain on EMR 6\n",
    "java_import(spark._sc._jvm, \"org.apache.spark.sql.api.python.*\")\n",
    "\n",
    "add_dataframe_druid_extension()\n",
    "\n",
    "df_prepared = df.repartition_by_druid_segment_size('event_date', segment_granularity='DAY')\n",
    "df_prepared.explain(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- revenue: double (nullable = true)\n",
      " |-- purchase_count: long (nullable = true)\n",
      " |-- app_id: string (nullable = true)\n",
      " |-- player_id: string (nullable = true)\n",
      " |-- event_date: timestamp (nullable = true)\n",
      " |-- dau: integer (nullable = false)\n",
      " |-- __PARTITION_TIME__: timestamp (nullable = true)\n",
      " |-- __PARTITION_NUM__: integer (nullable = true)"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df_prepared.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------+----------+---------+-------------------+---+-------------------+-----------------+\n",
      "|revenue|purchase_count|    app_id|player_id|         event_date|dau| __PARTITION_TIME__|__PARTITION_NUM__|\n",
      "+-------+--------------+----------+---------+-------------------+---+-------------------+-----------------+\n",
      "|    0.0|             0|testclient|       p3|2023-01-01 00:00:00|  1|2023-01-01 00:00:00|                0|\n",
      "|   30.0|            10|testclient|       p1|2023-01-01 00:00:00|  1|2023-01-01 00:00:00|                0|\n",
      "|   15.0|             1|testclient|       p2|2023-01-01 00:00:00|  1|2023-01-01 00:00:00|                0|\n",
      "|    0.0|             0|testclient|       p4|2023-01-01 00:00:00|  1|2023-01-01 00:00:00|                0|\n",
      "|    0.0|             0|testclient|       p2|2023-01-02 00:00:00|  1|2023-01-02 00:00:00|                0|\n",
      "|    5.0|             2|testclient|       p1|2023-01-02 00:00:00|  1|2023-01-02 00:00:00|                0|\n",
      "+-------+--------------+----------+---------+-------------------+---+-------------------+-----------------+"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df_prepared.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{\"type\": \"longSum\", \"name\": \"dau\", \"fieldName\": \"dau\"}, {\"type\": \"longSum\", \"name\": \"purchase_count\", \"fieldName\": \"purchase_count\"}, {\"type\": \"doubleSum\", \"name\": \"revenue\", \"fieldName\": \"revenue\"}, {\"type\": \"thetaSketch\", \"name\": \"player_id_sketch\", \"fieldName\": \"player_id\", \"isInputThetaSketch\": false, \"size\": 4096}, {\"type\": \"HLLSketchBuild\", \"name\": \"player_id_hll\", \"fieldName\": \"player_id\"}, {\"type\": \"quantilesDoublesSketch\", \"name\": \"revenue_quantile\", \"fieldName\": \"revenue\"}, {\"type\": \"quantilesDoublesSketch\", \"name\": \"purchase_count_quantile\", \"fieldName\": \"purchase_count\"}]"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "import json\n",
    "\n",
    "metrics_spec = [{\"type\": \"longSum\", \"name\": \"dau\", \"fieldName\": \"dau\"},\n",
    "                {\"type\": \"longSum\", \"name\": \"purchase_count\", \"fieldName\": \"purchase_count\"},\n",
    "                {\"type\": \"doubleSum\", \"name\": \"revenue\", \"fieldName\": \"revenue\"},\n",
    "                {\n",
    "                    \"type\": \"thetaSketch\",\n",
    "                    \"name\": \"player_id_sketch\",\n",
    "                    \"fieldName\": \"player_id\",\n",
    "                    \"isInputThetaSketch\": False,\n",
    "                    \"size\": 4096\n",
    "                },\n",
    "                {\n",
    "                    \"type\": \"HLLSketchBuild\",\n",
    "                    \"name\": \"player_id_hll\",\n",
    "                    \"fieldName\": \"player_id\"\n",
    "                },\n",
    "                {\n",
    "                  \"type\" : \"quantilesDoublesSketch\",\n",
    "                  \"name\" : \"revenue_quantile\",\n",
    "                  \"fieldName\" : \"revenue\",\n",
    "                },\n",
    "                {\n",
    "                  \"type\" : \"quantilesDoublesSketch\",\n",
    "                  \"name\" : \"purchase_count_quantile\",\n",
    "                  \"fieldName\" : \"purchase_count\",\n",
    "                },\n",
    "              ]\n",
    "print(json.dumps(metrics_spec))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%spark\n",
    "import json\n",
    "\n",
    "DATA_SOURCE_NAME = \"tmp_vivek_sketch_build_test\"\n",
    "\n",
    "df_prepared \\\n",
    "    .write \\\n",
    "    .mode(\"overwrite\") \\\n",
    "    .format(DRUID_SOURCE) \\\n",
    "    .option(ConfKeys.DATA_SOURCE, DATA_SOURCE_NAME) \\\n",
    "    .option(ConfKeys.TIME_COLUMN, \"event_date\") \\\n",
    "    .option(ConfKeys.METADATA_DB_URI, get_param(\"druid/metadata_db/uri\")) \\\n",
    "    .option(ConfKeys.METADATA_DB_USERNAME, get_param(\"druid/metadata_db/username\")) \\\n",
    "    .option(ConfKeys.METADATA_DB_PASSWORD, get_param(\"druid/metadata_db/password\")) \\\n",
    "    .option(ConfKeys.DEEP_STORAGE_S3_BUCKET, get_param(\"druid/deep_storage/bucket\")) \\\n",
    "    .option(ConfKeys.DEEP_STORAGE_S3_BASE_KEY, \"druid/segments\") \\\n",
    "    .option(ConfKeys.METRICS_SPEC, json.dumps(metrics_spec)) \\\n",
    "    .option(\"druid.segment_storage.s3.disableacl\", \"true\") \\\n",
    "    .save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To list the written data you can run:\n",
    "\n",
    "    aws s3 --profile $AWS_PROFILE ls --recursive s3://{druid-deep-storage-bucket}/druid/segments/tmp_vivek_sketch_build_test/\n",
    "\n",
    "To see something like:\n",
    "\n",
    "\n",
    "    2023-07-18 20:41:04       1514 druid/segments/tmp_vivek_sketch_build_test/2023-01-01T00:00:00.000Z_2023-01-02T00:00:00.000Z/2023-07-18T17:40:43.280Z/0/index.zip\n",
    "    2023-07-18 20:35:42       1385 druid/segments/tmp_vivek_sketch_build_test/2023-01-02T00:00:00.000Z_2023-01-\n",
    "    2023-07-18 20:41:04       1466 druid/segments/tmp_vivek_sketch_build_test/2023-01-02T00:00:00.000Z_2023-01-03T00:00:00.000Z/2023-07-18T17:40:43.280Z/0/index.zip\n",
    "\n",
    "\n",
    "And run the following in druid-sql (JDBC)\n",
    "\n",
    "\n",
    "1.  raw data\n",
    "\n",
    "\n",
    "    SELECT * FROM tmp_vivek_sketch_build_test\n",
    "\n",
    "```\n",
    "    __time app_id dau player_id_hll player_id_sketch purchase_count purchase_count_quantile revenue revenue_quantile\n",
    "    2023-01-01T00:00:00.000Z testclient 4 \"AgEHDAMIBAAOmR4KEjEsCG3Y2QaJOGoG\" \"AgMDAAAazJMEAAAAAACAPxM98wEU890HhRF7xjua2g9FFIa9XWbPMMoiRAt1/IA3\" 11 \"AgMIGoAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAJEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA8D8AAAAAAAAkQA==\" 45.0 \"AgMIGoAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAPkAAAAAAAAAAAAAAAAAAAAAAAAAAAAAALkAAAAAAAAA+QA==\"\n",
    "    2023-01-02T00:00:00.000Z testclient 2 \"AgEHDAMIAgCJOGoGbdjZBg==\" \"AgMDAAAazJMCAAAAAACAPxM98wEU890HyiJEC3X8gDc=\" 2 \"AgMIGoAAAAACAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABA\" 5.0 \"AgMIGoAAAAACAAAAAAAAAAAAAAAAAAAAAAAAAAAAFEAAAAAAAAAAAAAAAAAAABRA\"\n",
    "```\n",
    "\n",
    "\n",
    "2.  aggregation with sketches\n",
    "  -- https://druid.apache.org/docs/latest/querying/sql-aggregations.html#sketch-functions\n",
    "\n",
    "\n",
    "    SELECT __time,\n",
    "      app_id,\n",
    "      sum(dau) dau,\n",
    "      APPROX_COUNT_DISTINCT_DS_HLL(player_id_hll) approx_hll_ds_players,\n",
    "      APPROX_COUNT_DISTINCT_DS_THETA(player_id_sketch) approx_theta_ds_players,\n",
    "      sum(revenue) revenue,\n",
    "      APPROX_QUANTILE_DS(revenue_quantile, 0.25) revenue_quantile_25,\n",
    "      APPROX_QUANTILE_DS(revenue_quantile, 0.50) revenue_quantile_50,\n",
    "      APPROX_QUANTILE_DS(revenue_quantile, 0.75) revenue_quantile_75,\n",
    "      APPROX_QUANTILE_DS(revenue_quantile, 0.99) revenue_quantile_99,\n",
    "      sum(purchase_count) purchase_count,\n",
    "      APPROX_QUANTILE_DS(purchase_count_quantile, 0.25) purchase_count_quantile_25,\n",
    "      APPROX_QUANTILE_DS(purchase_count_quantile, 0.50) purchase_count_quantile_50,\n",
    "      APPROX_QUANTILE_DS(purchase_count_quantile, 0.75) purchase_count_quantile_75,\n",
    "      APPROX_QUANTILE_DS(purchase_count_quantile, 0.99) purchase_count_quantile_99\n",
    "    FROM tmp_vivek_sketch_build_test\n",
    "    group by 1,2\n",
    "    order by 1,2\n",
    "\n",
    "```\n",
    "    __time app_id dau approx_hll_ds_players approx_theta_ds_players revenue revenue_quantile_25 revenue_quantile_50 revenue_quantile_75 revenue_quantile_99 purchase_count purchase_count_quantile_25 purchase_count_quantile_50 purchase_count_quantile_75 purchase_count_quantile_99\n",
    "    2023-01-01T00:00:00.000Z testclient 4 4 4 45.0 0.0 15.0 30.0 30.0 11 0.0 1.0 10.0 10.0\n",
    "    2023-01-02T00:00:00.000Z testclient 2 2 2 5.0 0.0 5.0 5.0 5.0 2 0.0 2.0 2.0 2.0\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sketches merge (when source df contains sketches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DataFrame[]"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "# based on https://datasketches.apache.org/docs/Theta/ThetaHiveUDFs.html\n",
    "# https://spark.apache.org/docs/latest/sql-ref-functions-udf-hive.html\n",
    "spark.sql(\"create temporary function data2sketch as 'org.apache.datasketches.hive.theta.DataToSketchUDAF'\")\n",
    "spark.sql(\"create temporary function unionSketches as 'org.apache.datasketches.hive.theta.UnionSketchUDAF'\")\n",
    "spark.sql(\"create temporary function estimate as 'org.apache.datasketches.hive.theta.EstimateSketchUDF'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------+----------+---------+-------------------+---+\n",
      "|revenue|purchase_count|    app_id|player_id|         event_date|dau|\n",
      "+-------+--------------+----------+---------+-------------------+---+\n",
      "|   30.0|            10|testclient|       p1|2023-01-01 00:00:00|  1|\n",
      "|   15.0|             1|testclient|       p2|2023-01-01 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p3|2023-01-01 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p4|2023-01-01 00:00:00|  1|\n",
      "|    5.0|             2|testclient|       p1|2023-01-02 00:00:00|  1|\n",
      "|    0.0|             0|testclient|       p2|2023-01-02 00:00:00|  1|\n",
      "+-------+--------------+----------+---------+-------------------+---+"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df.createOrReplaceTempView(\"tmp_df\")\n",
    "spark.sql(\"select * from tmp_df\").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------------+----------+-------+---+--------------+--------------------+\n",
      "|         event_date|    app_id|revenue|dau|purchase_count|     player_id_theta|\n",
      "+-------------------+----------+-------+---+--------------+--------------------+\n",
      "|2023-01-02 00:00:00|testclient|    5.0|  2|             2|AgMDAAAazJMCAAAAA...|\n",
      "|2023-01-01 00:00:00|testclient|   45.0|  4|            11|AgMDAAAazJMEAAAAA...|\n",
      "+-------------------+----------+-------+---+--------------+--------------------+"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df2 = spark.sql(\n",
    " \"\"\"\n",
    " select\n",
    "     event_date,\n",
    "     app_id,\n",
    "     sum(revenue) as revenue,\n",
    "     sum(dau) as dau,\n",
    "     sum(purchase_count) as purchase_count,\n",
    "     data2sketch(player_id) as player_id_theta\n",
    "  from tmp_df\n",
    "  group by 1,2\n",
    " \"\"\"\n",
    ")\n",
    "\n",
    "# Convert BinaryType column to Base64 encoded string.\n",
    "df2 = df2.withColumn(\"player_id_theta\", f.base64(f.col(\"player_id_theta\")))\n",
    "df2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- event_date: timestamp (nullable = true)\n",
      " |-- app_id: string (nullable = true)\n",
      " |-- revenue: double (nullable = true)\n",
      " |-- dau: long (nullable = true)\n",
      " |-- purchase_count: long (nullable = true)\n",
      " |-- player_id_theta: string (nullable = true)"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df2.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'[{\"type\": \"longSum\", \"name\": \"dau\", \"fieldName\": \"dau\"}, {\"type\": \"longSum\", \"name\": \"purchase_count\", \"fieldName\": \"purchase_count\"}, {\"type\": \"doubleSum\", \"name\": \"revenue\", \"fieldName\": \"revenue\"}, {\"type\": \"thetaSketch\", \"name\": \"player_id_theta\", \"fieldName\": \"player_id_theta\", \"isInputThetaSketch\": true, \"size\": 4096}]'"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "metrics_spec = [{\"type\": \"longSum\", \"name\": \"dau\", \"fieldName\": \"dau\"},\n",
    "                {\"type\": \"longSum\", \"name\": \"purchase_count\", \"fieldName\": \"purchase_count\"},\n",
    "                {\"type\": \"doubleSum\", \"name\": \"revenue\", \"fieldName\": \"revenue\"},\n",
    "                {\n",
    "                    \"type\": \"thetaSketch\",\n",
    "                    \"name\": \"player_id_theta\",\n",
    "                    \"fieldName\": \"player_id_theta\",\n",
    "                    \"isInputThetaSketch\": True,\n",
    "                    \"size\": 4096\n",
    "                }\n",
    "              ]\n",
    "json.dumps(metrics_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __PARTITION_NUM__#252]\n",
      "+- RepartitionByExpression [__PARTITION_TIME__#231, __PARTITION_NUM__#252]\n",
      "   +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243, cast((cast((__num_rows__#243 - 1) as double) / cast(5000000 as double)) as int) AS __PARTITION_NUM__#252]\n",
      "      +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243]\n",
      "         +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243, __num_rows__#243]\n",
      "            +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#231, __PARTITION_TIME__#231 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#243], [__PARTITION_TIME__#231], [__PARTITION_TIME__#231 ASC NULLS FIRST]\n",
      "               +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231]\n",
      "                  +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, cast((cast(if (isnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) null else UDF(knownnotnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) as double) / cast(1000 as double)) as timestamp) AS __PARTITION_TIME__#231]\n",
      "                     +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, base64(player_id_theta#164) AS player_id_theta#175]\n",
      "                        +- Aggregate [event_date#4, app_id#2], [event_date#4, app_id#2, sum(revenue#0) AS revenue#161, sum(dau#10) AS dau#162L, sum(purchase_count#1L) AS purchase_count#163L, data2sketch(data2sketch, HiveFunctionWrapper(org.apache.datasketches.hive.theta.DataToSketchUDAF,org.apache.datasketches.hive.theta.DataToSketchUDAF@6b90ac51,class org.apache.datasketches.hive.theta.DataToSketchUDAF), player_id#3, false, 0, 0) AS player_id_theta#164]\n",
      "                           +- SubqueryAlias tmp_df\n",
      "                              +- View (`tmp_df`, [revenue#0,purchase_count#1L,app_id#2,player_id#3,event_date#4,dau#10])\n",
      "                                 +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10]\n",
      "                                    +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "event_date: timestamp, app_id: string, revenue: double, dau: bigint, purchase_count: bigint, player_id_theta: string, __PARTITION_TIME__: timestamp, __PARTITION_NUM__: int\n",
      "Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __PARTITION_NUM__#252]\n",
      "+- RepartitionByExpression [__PARTITION_TIME__#231, __PARTITION_NUM__#252]\n",
      "   +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243, cast((cast((__num_rows__#243 - 1) as double) / cast(5000000 as double)) as int) AS __PARTITION_NUM__#252]\n",
      "      +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243]\n",
      "         +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, __num_rows__#243, __num_rows__#243]\n",
      "            +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#231, __PARTITION_TIME__#231 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#243], [__PARTITION_TIME__#231], [__PARTITION_TIME__#231 ASC NULLS FIRST]\n",
      "               +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231]\n",
      "                  +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, cast((cast(if (isnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) null else UDF(knownnotnull(cast((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * cast(1000 as bigint)) as bigint))) as double) / cast(1000 as double)) as timestamp) AS __PARTITION_TIME__#231]\n",
      "                     +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, base64(player_id_theta#164) AS player_id_theta#175]\n",
      "                        +- Aggregate [event_date#4, app_id#2], [event_date#4, app_id#2, sum(revenue#0) AS revenue#161, sum(dau#10) AS dau#162L, sum(purchase_count#1L) AS purchase_count#163L, data2sketch(data2sketch, HiveFunctionWrapper(org.apache.datasketches.hive.theta.DataToSketchUDAF,org.apache.datasketches.hive.theta.DataToSketchUDAF@6b90ac51,class org.apache.datasketches.hive.theta.DataToSketchUDAF), player_id#3, false, 0, 0) AS player_id_theta#164]\n",
      "                           +- SubqueryAlias tmp_df\n",
      "                              +- View (`tmp_df`, [revenue#0,purchase_count#1L,app_id#2,player_id#3,event_date#4,dau#10])\n",
      "                                 +- Project [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4, 1 AS dau#10]\n",
      "                                    +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "RepartitionByExpression [__PARTITION_TIME__#231, __PARTITION_NUM__#252]\n",
      "+- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, cast((cast((__num_rows__#243 - 1) as double) / 5000000.0) as int) AS __PARTITION_NUM__#252]\n",
      "   +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#231, __PARTITION_TIME__#231 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#243], [__PARTITION_TIME__#231], [__PARTITION_TIME__#231 ASC NULLS FIRST]\n",
      "      +- Aggregate [event_date#4, app_id#2], [event_date#4, app_id#2, sum(revenue#0) AS revenue#161, sum(1) AS dau#162L, sum(purchase_count#1L) AS purchase_count#163L, base64(data2sketch(data2sketch, HiveFunctionWrapper(org.apache.datasketches.hive.theta.DataToSketchUDAF,org.apache.datasketches.hive.theta.DataToSketchUDAF@6b90ac51,class org.apache.datasketches.hive.theta.DataToSketchUDAF), player_id#3, false, 0, 0)) AS player_id_theta#175, if (isnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) null else cast((cast(UDF(knownnotnull((unix_timestamp(event_date#4, yyyy-MM-dd HH:mm:ss, Some(UTC), false) * 1000))) as double) / 1000.0) as timestamp) AS __PARTITION_TIME__#231]\n",
      "         +- LogicalRDD [revenue#0, purchase_count#1L, app_id#2, player_id#3, event_date#4], false\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- Exchange hashpartitioning(__PARTITION_TIME__#231, __PARTITION_NUM__#252, 1000), REPARTITION_BY_COL, [id=#397]\n",
      "   +- Project [event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231, cast((cast((__num_rows__#243 - 1) as double) / 5000000.0) as int) AS __PARTITION_NUM__#252]\n",
      "      +- Window [row_number() windowspecdefinition(__PARTITION_TIME__#231, __PARTITION_TIME__#231 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __num_rows__#243], [__PARTITION_TIME__#231], [__PARTITION_TIME__#231 ASC NULLS FIRST]\n",
      "         +- Sort [__PARTITION_TIME__#231 ASC NULLS FIRST, __PARTITION_TIME__#231 ASC NULLS FIRST], false, 0\n",
      "            +- Exchange hashpartitioning(__PARTITION_TIME__#231, 1000), ENSURE_REQUIREMENTS, [id=#393]\n",
      "               +- ObjectHashAggregate(keys=[event_date#4, app_id#2], functions=[sum(revenue#0), sum(1), sum(purchase_count#1L), data2sketch(data2sketch, HiveFunctionWrapper(org.apache.datasketches.hive.theta.DataToSketchUDAF,org.apache.datasketches.hive.theta.DataToSketchUDAF@6b90ac51,class org.apache.datasketches.hive.theta.DataToSketchUDAF), player_id#3, false, 0, 0)], output=[event_date#4, app_id#2, revenue#161, dau#162L, purchase_count#163L, player_id_theta#175, __PARTITION_TIME__#231])\n",
      "                  +- Exchange hashpartitioning(event_date#4, app_id#2, 1000), ENSURE_REQUIREMENTS, [id=#390]\n",
      "                     +- ObjectHashAggregate(keys=[event_date#4, app_id#2], functions=[partial_sum(revenue#0), partial_sum(1), partial_sum(purchase_count#1L), partial_data2sketch(data2sketch, HiveFunctionWrapper(org.apache.datasketches.hive.theta.DataToSketchUDAF,org.apache.datasketches.hive.theta.DataToSketchUDAF@6b90ac51,class org.apache.datasketches.hive.theta.DataToSketchUDAF), player_id#3, false, 0, 0)], output=[event_date#4, app_id#2, sum#204, sum#271L, sum#206L, buf#207])\n",
      "                        +- Scan ExistingRDD[revenue#0,purchase_count#1L,app_id#2,player_id#3,event_date#4]"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df_prepared2 = df2.repartition_by_druid_segment_size('event_date', segment_granularity='DAY')\n",
    "df_prepared2.explain(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------------+----------+-------+---+--------------+----------------------------------------------------------------+-------------------+-----------------+\n",
      "|event_date         |app_id    |revenue|dau|purchase_count|player_id_theta                                                 |__PARTITION_TIME__ |__PARTITION_NUM__|\n",
      "+-------------------+----------+-------+---+--------------+----------------------------------------------------------------+-------------------+-----------------+\n",
      "|2023-01-01 00:00:00|testclient|45.0   |4  |11            |AgMDAAAazJMEAAAAAACAPxM98wEU890HhRF7xjua2g9FFIa9XWbPMMoiRAt1/IA3|2023-01-01 00:00:00|0                |\n",
      "|2023-01-02 00:00:00|testclient|5.0    |2  |2             |AgMDAAAazJMCAAAAAACAPxM98wEU890HyiJEC3X8gDc=                    |2023-01-02 00:00:00|0                |\n",
      "+-------------------+----------+-------+---+--------------+----------------------------------------------------------------+-------------------+-----------------+"
     ]
    }
   ],
   "source": [
    "%%spark\n",
    "df_prepared2.show(truncate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%spark\n",
    "\n",
    "DATA_SOURCE_NAME = \"tmp_vivek_sketch_build_test2\"\n",
    "\n",
    "df_prepared2 \\\n",
    "    .write \\\n",
    "    .mode(\"overwrite\") \\\n",
    "    .format(DRUID_SOURCE) \\\n",
    "    .option(ConfKeys.DATA_SOURCE, DATA_SOURCE_NAME) \\\n",
    "    .option(ConfKeys.TIME_COLUMN, \"event_date\") \\\n",
    "    .option(ConfKeys.METADATA_DB_URI, get_param(\"druid/metadata_db/uri\")) \\\n",
    "    .option(ConfKeys.METADATA_DB_USERNAME, get_param(\"druid/metadata_db/username\")) \\\n",
    "    .option(ConfKeys.METADATA_DB_PASSWORD, get_param(\"druid/metadata_db/password\")) \\\n",
    "    .option(ConfKeys.DEEP_STORAGE_S3_BUCKET, get_param(\"druid/deep_storage/bucket\")) \\\n",
    "    .option(ConfKeys.DEEP_STORAGE_S3_BASE_KEY, \"druid/segments\") \\\n",
    "    .option(ConfKeys.METRICS_SPEC, json.dumps(metrics_spec)) \\\n",
    "    .option(\"druid.segment_storage.s3.disableacl\", \"true\") \\\n",
    "    .save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To list the written data you can run:\n",
    "\n",
    "    aws s3 --profile AWS_PROFILE ls --recursive s3://{druid-deep-storage-bucket}/druid/segments/tmp_vivek_sketch_build_test2/\n",
    "\n",
    "To see something like:\n",
    "\n",
    "\n",
    "    2023-07-19 10:53:45       1274 druid/segments/tmp_vivek_sketch_build_test2/2023-01-01T00:00:00.000Z_2023-01-02T00:00:00.000Z/2023-07-19T07:53:41.372Z/0/index.zip\n",
    "    2023-07-19 10:53:45       1247 druid/segments/tmp_vivek_sketch_build_test2/2023-01-02T00:00:00.000Z_2023-01-03T00:00:00.000Z/2023-07-19T07:53:41.372Z/0/index.zip\n",
    "\n",
    "And run this in druid-sql (JDBC)\n",
    "\n",
    "1. raw data\n",
    "\n",
    "    SELECT * FROM tmp_vivek_sketch_build_test2\n",
    "\n",
    "```\n",
    "    time app_id dau player_id_theta purchase_count revenue\n",
    "    2023-01-01 00:00 testclient 4 \"AgMDAAAazJMEAAAAAACAPxM98wEU890HhRF7xjua2g9FFIa9XWbPMMoiRAt1/IA3\" 11 45\t\n",
    "    2023-01-02 00:00 testclient 2 \"AgMDAAAazJMCAAAAAACAPxM98wEU890HyiJEC3X8gDc=\" 2 5\t\n",
    "```\n",
    "\n",
    "2.  aggregation with sketches\n",
    "\n",
    "    Based on https://druid.apache.org/docs/latest/querying/sql-aggregations.html#sketch-functions\n",
    "\n",
    "    SELECT __time,\n",
    "      app_id,\n",
    "      sum(dau) dau,\n",
    "      APPROX_COUNT_DISTINCT_DS_THETA(player_id_theta) approx_theta_ds_players,\n",
    "      sum(revenue) revenue\n",
    "    FROM tmp_vivek_sketch_build_test2\n",
    "    group by 1,2\n",
    "    order by 1,2\n",
    "  \n",
    "```\n",
    "    time             app_id     dau approx_theta_ds_players revenue\n",
    "    2023-01-01 00:00 testclient 4 4 45\t\n",
    "    2023-01-02 00:00 testclient 2 2 5\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%spark cleanup"
   ]
  }
 ],
 "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}