Oi João! Tudo bem?
A mensagem de erro "JAVA_HOME is not set" indica que a variável de ambiente JAVA_HOME não está configurada corretamente. Essa variável é necessária porque o Spark depende do Java para funcionar.
Quando você instala o pyspark no Colab, ele geralmente consegue localizar o Java padrão do sistema sem que você precise configurar manualmente a variável JAVA_HOME.
Então, antes de tentar a solução abaixo eu recomendo que você reinicie o seu notebook. Às vezes, o Colab pode apresentar instabilidades temporárias na leitura. Abaixo deixo os passos:
No menu de ferramentas do Google Colaboratory, clicar em "Ambientes de execução";
Selecionar a opção "Reiniciar sessão e executar tudo".
Para realizar essas etapas, você pode acompanhar a imagem abaixo:
Verifique se há algum trecho de código anterior ausente, você pode comparar o seu código com o notebook do instrutor
Caso seguindo o passo acima e ainda assim o erro persistir, você pode instalar o Java executando o seguinte comando em uma célula de código:
!apt-get install openjdk-8-jdk-headless -qq > /dev/null
Após instalar o Java, você precisa definir a variável de ambiente JAVA_HOME. Você pode fazer isso com o seguinte comando:
import os
os.environ["JAVA_HOME"] = "/usr/lib/jvm/java-8-openjdk-amd64"
Depois de configurar a variável JAVA_HOME, tente reiniciar a sessão novamente.
Espero ter ajudado.
Qualquer dúvida que surgir, compartilhe no fórum. Abraços e bons estudos!
Caso este post tenha lhe ajudado, por favor, marcar como solucionado ✓. Bons Estudos!