Pass Databricks's Spark session to a user defined module

Data Engineering

Join discussions on data engineering best practices, architectures, and optimization strategies within the Databricks Community. Exchange insights and solutions with fellow data engineers.

Hello everyone,

I have a .py file (not a notebook) where I have the following class with the following constructor:

class DataQualityChecker:
    def __init__(self, spark_session: SparkSession, df: DataFrame, quality_config_filepath: str) -> None:
        self.__spark = spark_session
        self.__df = df
        self.__quality_config = self.__import_quality_config(quality_config_filepath)
        self.__error_constraints = None

I'm trying to use this class in a notebook as follows:

from data_quality import DataQualityChecker

quality_checker = DataQualityChecker(spark, df, quality_config)

However, I am receiving the following error in the code snippet above:

TypeError: __init__() takes 3 positional arguments but 4 were given

Does anyone have any idea what it could be? I suspect the problem is trying to pass the spark variable as a parameter to the class constructor, but I don't know how to solve it.