pyarrow fix (#65)

pyarrow path fix for executor

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

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:

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",

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,