custom project context class

Description

I would like to use centralized Spark configuration file under conf/base/spark.yml and would like to pass these to SparkSession.

Context

I followed the steps mentioned here in the documentation. https://kedro.readthedocs.io/en/stable/11_tools_integration/01_pyspark.html?highlight=SparkSession#initialise-a-sparksession-in-custom-project-context-class

First Issue: KedroContext Which is not imported in the given example.

Second Issue: If I import like this from kedro.framework.context import KedroContext I am getting the below error: TypeError: init() got an unexpected keyword argument ‘package_name’

File "/root/user_files/envs/xpvu4/lib/python3.7/site-packages/kedro/framework/session/session.py", line 226, in create
    session._setup_logging()
  File "/root/user_files/envs/xpvu4/lib/python3.7/site-packages/kedro/framework/session/session.py", line 244, in _setup_logging
    conf_logging = self._get_logging_config()
  File "/root/user_files/envs/xpvu4/lib/python3.7/site-packages/kedro/framework/session/session.py", line 230, in _get_logging_config
    context = self.load_context()
  File "/root/user_files/envs/xpvu4/lib/python3.7/site-packages/kedro/framework/session/session.py", line 306, in load_context
    extra_params=extra_params,
TypeError: __init__() got an unexpected keyword argument 'package_name'

-- Separate them if you have more than one.

Your Environment

Include as many relevant details about the environment in which you experienced the bug:

• Kedro version used (pip show kedro or kedro -V): kedro, version 0.17.0
• Python version used (python -V): Python 3.7.9
• Operating system and version:

[samhiscoxqb]

Hello! The example CustomContext needs a little tweak to include project_name as an __init__ argument:

from typing import Any, Dict, Union
from pathlib import Path

from pyspark import SparkConf
from pyspark.sql import SparkSession

from kedro.framework.context import KedroContext


class CustomContext(KedroContext):

    def __init__(
        self,
        package_name: str,                                                 # <= Add package_name here
        project_path: Union[Path, str],
        env: str = None,
        extra_params: Dict[str, Any] = None,
    ):
        super().__init__(package_name, project_path, env, extra_params).   # <= and here!
        self.init_spark_session()

    def init_spark_session(self) -> None:
        """Initialises a SparkSession using the config defined in project's conf folder."""

        # Load the spark configuration in spark.yaml using the config loader
        parameters = self.config_loader.get("spark*", "spark*/**")
        spark_conf = SparkConf().setAll(parameters.items())

        # Initialise the spark session
        spark_session_conf = (
            SparkSession.builder
            .appName(self.package_name)
            .enableHiveSupport()
            .config(conf=spark_conf)
        )
        _spark_session = spark_session_conf.getOrCreate()
        _spark_session.sparkContext.setLogLevel("WARN")

I’ll open a PR to update this - thanks for raising both 👍

[remiromi]

@samhiscoxqb thanks a lot for your help, it worked! I had also to edit another line that you wrote correctly but in the doc is not (-> https://kedro.readthedocs.io/en/stable/11_tools_integration/01_pyspark.html ) In the doc the code is

# Initialise the spark session
        spark_session_conf = (
            SparkSession.builder
            .appName(self.project_name)      # Error Here
            .enableHiveSupport()
            .config(conf=spark_conf)
        )

While I had to edit the appName line as

# Initialise the spark session
        spark_session_conf = (
            SparkSession.builder
            .appName(self.package_name)        # Corrected line
            .enableHiveSupport()
            .config(conf=spark_conf)
        )

since it was giving me AttributeError on the self.project_name call.

Thanks again!