Read parquet, simple transformations, write parquet using spark sql

ReadWriteParquet.py

from pyspark import SparkContext
from pyspark.sql import SparkSession
from pyspark.sql.functions import lit
from random import randint, seed

input_file = '/tmp/input.snappy.parquet'
output_file = '/tmp/output.snappy.parquet'

spark = SparkSession(SparkContext('local', 'make-recs-application'))


class ReadRewriteParquet:

    def __init__(self):
        self.id_ctr = 0

    def ids(self):
        df = spark.read.parquet(input_file)
        cnt = df.count()
        print(cnt)
        return df. \
            selectExpr("id as patientId"). \
            withColumn("amount", lit(40)). \
            withColumn("id", lit(randint(0, 1112)))
        # sort("]", ascending=True). \

    def make_record(self, patient_id):
        x = self.id_ctr
        record = [x, patient_id, 40.0]
        self.id_ctr = self.id_ctr + 1
        print(record)


if __name__ == '__main__':
    seed()
    f = ReadRewriteParquet()
    id_df = f.ids()
    id_df.write.parquet(output_file, compression='snappy')