PySpark schema save/load example

pyspark_schema_util.py

import json
from pyspark.sql.types import *

def save_schema_as_json(df, schema_file):
  """
    Saves dataframe schema as json
  """
  schema = df.schema.json()
  schema = json.dumps(json.loads(schema), indent=4)
  with open(schema_file, "w") as f:
    f.write(schema)

def load_schema_as_json(schema_file, path):
  """
    Loads json data using schema file
  """
  with open(schema_file) as f:
    schema_json = f.read()
    
  schema = StructType.fromJson(json.loads(schema_json))
  return spark.read.schema(schema).json(path)