Steps to setup colab notebook pyspark

setup_notebook.py

from subprocess import STDOUT, check_call, check_output

def install_java_scala():
          try:
          java_ver = check_output(['java', '-version'], stderr=STDOUT)
        except:
          java_ver = b''
        try:
          scala_ver = check_output(['scala', '-version'], stderr=STDOUT)
        except:
          scala_ver = b''
        if b'1.8.0_232' not in java_ver:
          java_8_install = ['apt-get', '--quiet', 'install',
                            '-y', 'openjdk-8-jdk-headless']
          java_set_alt = ['update-alternatives', '--set', 'java', 
                          '/usr/lib/jvm/java-8-openjdk-amd64/jre/bin/java' ] 
          check_call(java_8_install, stdout=open(os.devnull, 'wb'), 
                     stderr=STDOUT)
          os.environ["JAVA_HOME"] = "/usr/lib/jvm/java-8-openjdk-amd64" 
          check_call(java_set_alt)  
        if b'2.11.12' not in scala_ver:
          scala_install = ['apt-get', '--quiet', 'install', 'scala']
          check_call(scala_install)