Can you reference a spark instance within a custom python package you import into a databricks notebook?

Question

I want to create a custom connector class in a python package that I import into a databricks notebook. Something like:

class snowflake_read():
    
    def __init__(self, format, options):
        self.options = options
        self.format = format


    def sf_query(self, query):
        df = spark.read.format(self.format).options(self.options).options("query", query).load()
        return df

While I am able to import the package into the notebook and initialize the snowflake_read() class, I am not able to use the sf_query() function as spark is undefined.

import read_package

sf_read = snowflake_read(<format>, <options>)
sf_read.sf_query('<query>')

This code produces an undefined variable ('spark') error.

How do I define spark within the init function? Spark is automatically defined/initialized within the datbricks notebook but clearly isn't within the scope of the imported class.

Answer 1

spark is automatically defined only in the notebooks directly, but not in the packages. You either need to pass that instance explicitly as a function parameter, or inside the function add following before calling the spark.read (import could be put outside of the function):

from pyspark.sql import SparkSession
spark = SparkSession.getActiveSession()

See documentation for SparkSession.getActiveSession function.