Adding NotebookVersion Parameter as specified in official AWS Docs https://docs.aws.amazon.com/athena/latest/APIReference/API_StartSession.html#athena-StartSession-request-NotebookVersion, this parameter is necessary to create session using Spark

Author: Davide Bisso

awswrangler/athena/_spark.py

@@ -94,6 +94,7 @@ def create_spark_session(
     default_executor_dpu_size: int = 1,
     additional_configs: dict[str, Any] | None = None,
     spark_properties: dict[str, Any] | None = None,
+    notebook_version: str = 'Athena notebook version 1',
     idle_timeout: int = 15,
     boto3_session: boto3.Session | None = None,
 ) -> str:
@@ -116,6 +117,9 @@ def create_spark_session(
     spark_properties: Dict[str, Any], optional
         Contains SparkProperties in the form of key-value pairs.Specifies custom jar files and Spark properties
         for use cases like cluster encryption, table formats, and general Spark tuning.
+    notebook_version: str
+        The notebook version. This value is supplied automatically for notebook sessions in the Athena console and is not required for programmatic session access. 
+        The only valid notebook version is Athena notebook version 1. If you specify a value for NotebookVersion, you must also specify a value for NotebookId
     idle_timeout : int, optional
          The idle timeout in minutes for the session. The default is 15.
     boto3_session : boto3.Session(), optional
@@ -146,6 +150,7 @@ def create_spark_session(
         WorkGroup=workgroup,
         EngineConfiguration=engine_configuration,
         SessionIdleTimeoutInMinutes=idle_timeout,
+        NotebookVersion=notebook_version,
     )
     _logger.info("Session info:\n%s", response)
     session_id: str = response["SessionId"]
@@ -166,6 +171,7 @@ def run_spark_calculation(
     default_executor_dpu_size: int = 1,
     additional_configs: dict[str, Any] | None = None,
     spark_properties: dict[str, Any] | None = None,
+    notebook_version: str = 'Athena notebook version 1',
     idle_timeout: int = 15,
     boto3_session: boto3.Session | None = None,
 ) -> dict[str, Any]:
@@ -192,6 +198,9 @@ def run_spark_calculation(
     spark_properties: Dict[str, Any], optional
         Contains SparkProperties in the form of key-value pairs.Specifies custom jar files and Spark properties
         for use cases like cluster encryption, table formats, and general Spark tuning.
+    notebook_version: str
+        The notebook version. This value is supplied automatically for notebook sessions in the Athena console and is not required for programmatic session access. 
+        The only valid notebook version is Athena notebook version 1. If you specify a value for NotebookVersion, you must also specify a value for NotebookId
     idle_timeout : int, optional
         The idle timeout in minutes for the session. The default is 15.
     boto3_session : boto3.Session(), optional
@@ -221,6 +230,7 @@ def run_spark_calculation(
             default_executor_dpu_size=default_executor_dpu_size,
             additional_configs=additional_configs,
             spark_properties=spark_properties,
+            notebook_version=notebook_version,
             idle_timeout=idle_timeout,
             boto3_session=boto3_session,
         )