From d0ce92d1028c87f4ec293b34449c418b653f700a Mon Sep 17 00:00:00 2001
From: Mikita Sakalouski <38785549+mikita-sakalouski@users.noreply.github.com>
Date: Thu, 30 May 2024 00:08:49 +0200
Subject: [PATCH] chore: Update pyproject.toml dependencies for pyspark and
 spark-expectations

---
 pyproject.toml                                         | 10 ++++++++--
 .../integrations/spark/dq/spark_expectations.py        | 10 ++++++----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 5811a235..6abf886b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -278,10 +278,16 @@ version = ["pyspark35"]
 
 [tool.hatch.envs.hatch-test.overrides]
 matrix.version.extra-dependencies = [
-  { value = "pyspark>=3.3,<3.4; spark-expectations>=2.1.0", if = [
+  { value = "pyspark>=3.3,<3.4", if = [
     "pyspark33",
   ] },
-  { value = "pyspark>=3.4,<3.5; spark-expectations>=2.1.0", if = [
+  { value = "spark-expectations>=2.1.0", if = [
+    "pyspark33",
+  ] },
+  { value = "pyspark>=3.4,<3.5", if = [
+    "pyspark34",
+  ] },
+  { value = "spark-expectations>=2.1.0", if = [
     "pyspark34",
   ] },
   { value = "pyspark>=3.5,<3.6", if = [
diff --git a/src/koheesio/integrations/spark/dq/spark_expectations.py b/src/koheesio/integrations/spark/dq/spark_expectations.py
index 6a01f91f..325ccaf5 100644
--- a/src/koheesio/integrations/spark/dq/spark_expectations.py
+++ b/src/koheesio/integrations/spark/dq/spark_expectations.py
@@ -4,19 +4,21 @@
 
 from typing import Any, Dict, Optional, Union
 
+import pyspark
+from pydantic import Field
+from pyspark.sql import DataFrame
 from spark_expectations.config.user_config import Constants as user_config
 from spark_expectations.core.expectations import (
     SparkExpectations,
     WrappedDataFrameWriter,
 )
 
-from pydantic import Field
-
-from pyspark.sql import DataFrame
-
 from koheesio.spark.transformations import Transformation
 from koheesio.spark.writers import BatchOutputMode
 
+if pyspark.__version__.startswith("3.5"):
+    raise ImportError("Spark Expectations is not supported for Spark 3.5")
+
 
 class SparkExpectationsTransformation(Transformation):
     """