{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Generate parquet output with PySpark\n\nParquet is a columnar-format with desirable properties for larger datasets. We may for example only be interested in a subset of columns. Parquet allows for reading individual columns without having to parse every line.\n\nI'll use PySpark to do this, for my personal reference on installing Spark and adjusting memory usage on a Kaggle instance."},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":false,"_kg_hide-output":true},"cell_type":"code","source":"! pip install pyspark","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pathlib import Path\nfrom pyspark.sql import SparkSession\n\n\nspark = (\n    SparkSession.builder\n    .config(\"spark.driver.memory\", \"12g\")\n    .getOrCreate()\n)\nspark.conf.set(\"spark.sql.shuffle.partitions\", spark.sparkContext.defaultParallelism*2)\nspark.conf.get(\"spark.driver.memory\")\n\n\nprefix = \"../input/riiid-test-answer-prediction\"\nfor path in Path(prefix).glob(\"*.csv\"):\n    name = path.name.split(\".\")[0]\n    print(f\"writing {path} to {name}\")\n    %time df = spark.read.csv(path.as_posix(), header=True, inferSchema=True)\n    df.printSchema()\n    \n    if \"train\" in path.name:\n        # dataset is 5.7GB, so each partition should be ~1GB with 4 parts\n        %time df.repartitionByRange(4, \"user_id\", \"timestamp\").write.parquet(name, mode=\"overwrite\")\n    else:\n        %time df.repartition(1).write.parquet(name, mode=\"overwrite\")","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}