From 5b5f7245cf402ae87c7d0a402eb40781f548528c Mon Sep 17 00:00:00 2001
From: svittoz <simon.vittoz-ext@aphp.fr>
Date: Fri, 26 Apr 2024 17:05:45 +0000
Subject: [PATCH] pyarrow path fix for executor

---
 eds_scikit/__init__.py               | 15 ++++--------
 eds_scikit/io/__init__.py            |  8 ++++++-
 eds_scikit/io/improve_performance.py | 34 +++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/eds_scikit/__init__.py b/eds_scikit/__init__.py
index 6da07cb7..832e9aba 100644
--- a/eds_scikit/__init__.py
+++ b/eds_scikit/__init__.py
@@ -12,36 +12,29 @@
 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:
 `spark, sc, sql = eds_scikit.improve_performances()`
 The functions respectively returns a SparkSession, a SparkContext and an sql method"""
 )
+
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..47186d04 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}/package-override"
+    )
+
+    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,
@@ -88,13 +119,14 @@ def improve_performances(
     tz = os.environ.get("TZ", "UTC")
     os.environ["TZ"] = tz
     time.tzset()
-
+    
     to_add_conf.extend(
         [
             ("spark.app.name", f"{os.environ.get('USER')}_{app_name}_scikit"),
             ("spark.sql.session.timeZone", tz),
             ("spark.sql.execution.arrow.enabled", "true"),
             ("spark.sql.execution.arrow.pyspark.enabled", "true"),
+            ("spark.executorEnv.PYTHONPATH", f"{BASE_DIR}/package-override")
         ]
     )