diff --git a/src/koheesio/spark/utils.py b/src/koheesio/spark/utils.py index 7d42b0c..b382c4b 100644 --- a/src/koheesio/spark/utils.py +++ b/src/koheesio/spark/utils.py @@ -188,12 +188,16 @@ def import_pandas_based_on_pyspark_version(): """ try: import pandas as pd + pyspark_version = get_spark_minor_version() pandas_version = pd.__version__ - if (pyspark_version < 3.4 and pandas_version >= '2') or (pyspark_version >= 3.4 and pandas_version < '2'): - raise ImportError(f"For PySpark {pyspark_version}, please install Pandas version {'< 2' if pyspark_version < 3.4 else '>= 2'}") + if (pyspark_version < 3.4 and pandas_version >= "2") or (pyspark_version >= 3.4 and pandas_version < "2"): + raise ImportError( + f"For PySpark {pyspark_version}, " + f"please install Pandas version {'< 2' if pyspark_version < 3.4 else '>= 2'}" + ) return pd - except ImportError: - raise ImportError("Pandas module is not installed.") + except ImportError as e: + raise ImportError("Pandas module is not installed.") from e