read-trim-big-csv-to-dataframe-then-save-to-parquet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract $31^3$ samples from the HR4 catalog at $z=0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate edges \n",
    "import sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.spatial import cKDTree\n",
    "import gc\n",
    "\n",
    "# plot settings\n",
    "plt.rc('font', family='serif') \n",
    "plt.rc('font', serif='Times New Roman') \n",
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['mathtext.fontset'] = 'stix'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark import SparkContext   \n",
    "from pyspark.sql import SQLContext\n",
    "\n",
    "#sc = SparkContext(master='local[3]', appName='calgraph')\n",
    "sqlsc = SQLContext(sc)\n",
    "#sc.setCheckpointDir(\"./checkpoints\")\n",
    "#sc.setCheckpointDir(\"hdfs://localhost:8020/myhdfs/spark/checkpoints\")\n",
    "sc.setCheckpointDir(\"hdfs://master:54310/tmp/spark/checkpoints\")\n",
    "\n",
    "import pyspark.sql.functions as F\n",
    "import pyspark.sql.types as T\n",
    "from pyspark import Row\n",
    "from pyspark.sql.window import Window as W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.execution.arrow.enabled\", \"true\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* ###  Read the halo csv file to the dataframe `halodf`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "halo_schema = T.StructType([\\\n",
    "                            T.StructField('haloid', T.IntegerType(), False),\\\n",
    "                            T.StructField('px', T.FloatType(), False),\\\n",
    "                            T.StructField('py', T.FloatType(), False),\\\n",
    "                            T.StructField('pz', T.FloatType(), False),\\\n",
    "                            T.StructField('halomass', T.FloatType(), False),\\\n",
    "                          ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "halodf = sqlsc.read.csv(\"hdfs://master:54310/data/spark/hr4/hr4z0.csv\",\\\n",
    "                        header=True, schema = halo_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys.getsizeof(halodf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+---------+---------+---------+-------------+\n",
      "|haloid|       px|       py|       pz|     halomass|\n",
      "+------+---------+---------+---------+-------------+\n",
      "|     0|106.23875|2820.2603|310.53067|3.29161999E14|\n",
      "|     1|1015.0091| 3070.103|2687.5447|   5.79631E14|\n",
      "|     2|1150.7571| 656.3275|195.96417| 7.4869997E14|\n",
      "+------+---------+---------+---------+-------------+\n",
      "only showing top 3 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "halodf.show(3,truncate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------------+------------------+------------------+------------------+\n",
      "|summary|              haloid|                px|                py|                pz|\n",
      "+-------+--------------------+------------------+------------------+------------------+\n",
      "|  count|           362123180|         362123180|         362123180|         362123180|\n",
      "|   mean|1.8106158950000003E8|1574.9487699428887|1574.5065974890745|1575.7579631810697|\n",
      "| stddev|1.0453595787073924E8| 909.2939271893584|  909.506639813124| 909.4171871354729|\n",
      "|    min|                   0|         -3.628511|         -3.603208|          0.163826|\n",
      "|    max|           362123179|         3149.9463|         3149.9321|         3154.1765|\n",
      "+-------+--------------------+------------------+------------------+------------------+\n",
      "\n",
      "CPU times: user 8.9 ms, sys: 4.51 ms, total: 13.4 ms\n",
      "Wall time: 38.9 s\n"
     ]
    }
   ],
   "source": [
    "%time halodf.select('haloid','px','py','pz').describe().show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* ###  From the $3150^3$ box, we extract $31^3$ sub-boxes with the size, $L = 100 h^{-1}$Mpc\n",
    "\n",
    "> For each subboxes, we will take top 2000 most massive halos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "lsys = 3150.0\n",
    "mlength = 100.0\n",
    "mdiv = 31"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "subdf = halodf\\\n",
    ".withColumn('ixdiv',F.floor(halodf.px/mlength))\\\n",
    ".withColumn('iydiv',F.floor(halodf.py/mlength))\\\n",
    ".withColumn('izdiv',F.floor(halodf.pz/mlength))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "subdf = subdf\\\n",
    ".withColumn('ixdiv',F.when(subdf.ixdiv == -1, 0).otherwise(subdf.ixdiv))\\\n",
    ".withColumn('iydiv',F.when(subdf.iydiv == -1, 0).otherwise(subdf.iydiv))\\\n",
    ".withColumn('izdiv',F.when(subdf.izdiv == -1, 0).otherwise(subdf.izdiv))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataFrame[haloid: int, px: float, py: float, pz: float, halomass: float, ixdiv: bigint, iydiv: bigint, izdiv: bigint]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subdf.cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+---------+---------+---------+-------------+-----+-----+-----+\n",
      "|haloid|       px|       py|       pz|     halomass|ixdiv|iydiv|izdiv|\n",
      "+------+---------+---------+---------+-------------+-----+-----+-----+\n",
      "|     0|106.23875|2820.2603|310.53067|3.29161999E14|    1|   28|    3|\n",
      "|     1|1015.0091| 3070.103|2687.5447|   5.79631E14|   10|   30|   26|\n",
      "|     2|1150.7571| 656.3275|195.96417| 7.4869997E14|   11|    6|    1|\n",
      "|     3|2745.5166|2020.5542|839.78326|1.55997994E14|   27|   20|    8|\n",
      "|     4|1129.6155|1974.3558|1253.1699|3.26915999E14|   11|   19|   12|\n",
      "|     5|589.67566|2289.6702|2542.1628|2.77195998E14|    5|   22|   25|\n",
      "|     6|1486.9604| 2534.887| 2683.816|1.08108999E14|   14|   25|   26|\n",
      "|     7|3036.3958|13.892702|2686.4326|1.14171002E14|   30|    0|   26|\n",
      "|     8|1077.6719|230.39247|2675.5627| 7.3316702E14|   10|    2|   26|\n",
      "|     9|1666.9918|658.03186| 472.5904|5.46761014E14|   16|    6|    4|\n",
      "+------+---------+---------+---------+-------------+-----+-----+-----+\n",
      "only showing top 10 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "subdf.show(10,truncate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "idiv code min =  0   0   0\n",
      "CPU times: user 17.3 ms, sys: 8.74 ms, total: 26 ms\n",
      "Wall time: 59.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ixmin = subdf.agg({\"ixdiv\":\"min\"}).collect()[0][0]\n",
    "iymin = subdf.agg({\"iydiv\":\"min\"}).collect()[0][0]\n",
    "izmin = subdf.agg({\"izdiv\":\"min\"}).collect()[0][0]\n",
    "print \"idiv code min = \", ixmin,\" \",iymin,\" \",izmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "idiv code min =  31   31   31\n",
      "CPU times: user 6.49 ms, sys: 2.9 ms, total: 9.39 ms\n",
      "Wall time: 1.15 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ixmax = subdf.agg({\"ixdiv\":\"max\"}).collect()[0][0]\n",
    "iymax = subdf.agg({\"iydiv\":\"max\"}).collect()[0][0]\n",
    "izmax = subdf.agg({\"izdiv\":\"max\"}).collect()[0][0]\n",
    "print \"idiv code min = \", ixmax,\" \",iymax,\" \",izmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "29791\n"
     ]
    }
   ],
   "source": [
    "print 31**3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30]\n"
     ]
    }
   ],
   "source": [
    "print range(31)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the cell id_s, ix, iy, iz, from 0 to 30\n",
    "idmax = 31 # idmax=31 == [0, ... ,30]\n",
    "redf = pd.DataFrame()\n",
    "icounter = 0\n",
    "numhalo = 2000 # how many halos for the output\n",
    "for ix in range(idmax):\n",
    "    for iy in range(idmax):\n",
    "        for iz in range(idmax):\n",
    "            gc.collect()\n",
    "            tmppd = subdf.filter((F.col('ixdiv') ==ix) & (F.col('iydiv') ==iy) & (F.col('izdiv') ==iz))\\\n",
    "            .select('haloid','px','py','pz','halomass').toPandas()\n",
    "            tmppd = tmppd.sort_values(by='halomass',ascending=False)[0:numhalo] #sort the DF based on halomass\n",
    "            tmppd['isect'] = np.full(numhalo,icounter) #id for each section \n",
    "            #print \"ix, iy, iz, numdata : \",ix,\" \",iy,\" \", iz,\" \",len(tmppd.index)\n",
    "            #print tmppd.head(3)\n",
    "            redf = pd.concat([redf,tmppd])\n",
    "            icounter = icounter + 1\n",
    "            \n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#tmppd.sort_values(by='halomass',ascending=False)[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "59582000\n"
     ]
    }
   ],
   "source": [
    "redf = redf.reset_index(drop=True)\n",
    "print len(redf.index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* ###  write the result pandas-dataframe `redf` to a parquet file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "pq.write_table(pa.Table.from_pandas(redf), 'hrdata.parquet.snappy', compression='snappy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}