From 09b7f50f70f8b0634a5a686af46217b13312dcb9 Mon Sep 17 00:00:00 2001
From: svittoz <137794505+svittoz@users.noreply.github.com>
Date: Tue, 28 May 2024 13:57:58 +0200
Subject: [PATCH] pyarrow fix (#65)

pyarrow path fix for executor
---
 changelog.md                         |  3 +++
 eds_scikit/__init__.py               | 14 +++----------
 eds_scikit/io/__init__.py            |  8 ++++++-
 eds_scikit/io/improve_performance.py | 31 ++++++++++++++++++++++++++++
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/changelog.md b/changelog.md
index e1534b5b..613ad5b3 100644
--- a/changelog.md
+++ b/changelog.md
@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+### Fixed
+- pyarrow fix did not work on spark executors
+
 ## v0.1.7 (2024-04-12)
 ### Changed
 - Support for pyarrow > 0.17.0
diff --git a/eds_scikit/__init__.py b/eds_scikit/__init__.py
index 6da07cb7..b7b168af 100644
--- a/eds_scikit/__init__.py
+++ b/eds_scikit/__init__.py
@@ -12,33 +12,25 @@
 import importlib
 import os
 import pathlib
-import sys
 import time
 from packaging import version
 from typing import List, Tuple
 from pathlib import Path
 
 import pandas as pd
-import pyarrow
-import pyarrow.ipc
 import pyspark
 from loguru import logger
 from pyspark import SparkContext
 from pyspark.sql import SparkSession
 
 import eds_scikit.biology  # noqa: F401 --> To register functions
-from eds_scikit.io import improve_performances
-
-pyarrow.open_stream = pyarrow.ipc.open_stream
-
-sys.path.insert(
-    0, (pathlib.Path(__file__).parent / "package-override").absolute().as_posix()
-)
-os.environ["PYTHONPATH"] = ":".join(sys.path)
+from eds_scikit.io import improve_performances, pyarrow_fix
 
 # Remove SettingWithCopyWarning
 pd.options.mode.chained_assignment = None
 
+pyarrow_fix()
+
 logger.warning(
     """To improve performances when using Spark and Koalas, please call `eds_scikit.improve_performances()`
 This function optimally configures Spark. Use it as:
diff --git a/eds_scikit/io/__init__.py b/eds_scikit/io/__init__.py
index bc01ea06..7fb4e35e 100644
--- a/eds_scikit/io/__init__.py
+++ b/eds_scikit/io/__init__.py
@@ -2,7 +2,13 @@
 from .files import PandasData
 from .hive import HiveData
 from .postgres import PostgresData
-from .improve_performance import improve_performances, koalas_options, load_koalas
+from .improve_performance import (
+    improve_performances,
+    koalas_options,
+    load_koalas,
+    pyarrow_fix,
+)
+
 
 __all__ = [
     "BaseData",
diff --git a/eds_scikit/io/improve_performance.py b/eds_scikit/io/improve_performance.py
index 42b6f19a..afed9f27 100644
--- a/eds_scikit/io/improve_performance.py
+++ b/eds_scikit/io/improve_performance.py
@@ -6,6 +6,7 @@
 from typing import List, Tuple
 
 import pyarrow
+import pyarrow.ipc
 import pyspark
 from packaging import version
 from pyspark import SparkContext
@@ -50,6 +51,36 @@ def set_env_variables() -> None:
         os.environ["PYARROW_IGNORE_TIMEZONE"] = "0"
 
 
+def pyarrow_fix():
+    """
+    Fixing error 'pyarrow has no attributes open_stream' due to PySpark 2 incompatibility with pyarrow > 0.17
+    """
+
+    # Setting path to our patched pyarrow module
+    pyarrow.open_stream = pyarrow.ipc.open_stream
+
+    sys.path.insert(
+        0, (Path(__file__).parent / "package-override").absolute().as_posix()
+    )
+    os.environ["PYTHONPATH"] = ":".join(sys.path)
+
+    # Setting this path for Pyspark executors
+    global spark, sc, sql
+
+    spark = SparkSession.builder.getOrCreate()
+
+    conf = spark.sparkContext.getConf()
+    conf.set(
+        "spark.executorEnv.PYTHONPATH",
+        f"{Path(__file__).parent.parent}/package-override:{conf.get('spark.executorEnv.PYTHONPATH')}",
+    )
+    spark = SparkSession.builder.enableHiveSupport().config(conf=conf).getOrCreate()
+
+    sc = spark.sparkContext
+
+    sql = spark.sql
+
+
 def improve_performances(
     to_add_conf: List[Tuple[str, str]] = [],
     quiet_spark: bool = True,