How Databricks Runtime 14.x destroyed 3d-party PySpark packages compatibility

We all know that Python is not the fastest programming language, but it works very well as a glue for compiled languages like the JVM family. And this is exactly how PySpark worked before Spark Connect: we have a JVM when all the Spark routines are running and we have Py4j acting as a bridge. If you take a look at the SparkSession object, you can see the following private attributes:

self._sc = sparkContext
self._jsc = self._sc._jsc
self._jvm = self._sc._jvm

So, we are having an access to the underlying JVM and underlying JavaSparkContext object. Another important thing is an __init__ method of pyspark.sql.DataFrame object:

def __init__(
    self,
    jdf: JavaObject,
    sql_ctx: Union["SQLContext", "SparkSession"],
):

It allows developers to take a py4j.java_gateway.JavaObject that represents an org.apache.spark.sql.DataFrame JVM object and create from it a PySpark DataFrame. Back converting is even easier: pyspark.sql.DataFrame has an attribute _jdf that is a reference to the underlying org.apache.spark.sql.DataFrame object:

self._jdf: JavaObject = jdf

One may already understand the obvious pattern of creating a PySpark package:

1. We are writing a core part of the package in Java/Scala with access to all Apache Spark APIs in fast compilable language
2. We are creating a python library that just provide bindings to underlying JAR-package via py4j.

For example, recently I was happy to create such a PySpark wrapper for a new huge Graph storage format library: my PR with PySpark bindings to GraphAr scala package.

The answer is quite obvious: because Spark Connect decouples the driver from the user's code, the user no longer has access to the underlying JVM. Of course, the topic of importance of py4j for 3d-party libs is very specific and many users of PySpark did not even notice the importance of this change. But for me it was the first question.

I will be honest if I say that I tried to find an alternative way for 3d-parties to interact with the underlying JVM with Spark Connect. But there is almost nothing about it in all the promotional material from Databricks and in the Apache Spark documentation. The only theoretically possible workaround I could find is this tiny Markdown page, hidden very deep in the source code of Apache Spark. Based on this, I can imagine that the Spark Connect protocol was designed to be extensible, but again: there is no documentation on how to do it!

Synapse ML is a well know (4.9k stars) spark extension, focused on applying ML/DL on Apache Spark clusters. One may know it as MMLSpark. There core part of the library is written in Scala, but APIs for R, Python, #NET and Java are provided. If one make a look how a Python API is organized under the hood they would see the described above "py4j-pattern":

class DiscreteHyperParam(object):
    """
    Specifies a discrete list of values.
    """

    def __init__(self, values, seed=0):
        ctx = SparkContext.getOrCreate()
        self.jvm = ctx.getOrCreate()._jvm
        self.hyperParam = self.jvm.com.microsoft.azure.synapse.ml.automl.HyperParamUtils.getDiscreteHyperParam(
            values,
            seed,
        )

link to the code above

Due to the popularity of that library they already faced issues from Databricks users: [BUG] Databricks 14.3 LTS usage of internal _jvm variable is no longer supported #2167. And I have zero ideas how they are going to fix it because to make it work with Spark Connect they need to rewrite all the logic in pure Python/PySpark.

PyDeequ is a popular (625 stars) Data Quality library that is native to Apache Spark because its core is written in Scala. Again, if one make a look on how is it implemented under the hood they will see "py4j-pattern" again:

class _AnalyzerObject:
    """
    Analyzer base object to pass and accumulate the analyzers of the run with respect to the JVM
    """

    def _set_jvm(self, jvm):
        self._jvm = jvm
        return self

    @property
    def _deequAnalyzers(self):
        if self._jvm:
            return self._jvm.com.amazon.deequ.analyzers
        raise AttributeError(
            "JVM not set, please run _set_jvm() method first."
        )  # TODO: Test that this exception gets raised

link to the code above

Spark-NLP is one of the most popular (3.6k stars) way to run LLMs on Apache Spark clusters. Let's again go the source code and see how it works. Oops, looks like we found using of _jdf / _jvm again:

class RecursiveEstimator(JavaEstimator, ABC):

    def _fit_java(self, dataset, pipeline=None):
        self._transfer_params_to_java()
        if pipeline:
            return self._java_obj.recursiveFit(dataset._jdf, pipeline._to_java())
        else:
            return self._java_obj.fit(dataset._jdf)

link to the code above

spark-extensions is relative popular (155 stars) and actively maintained library, that contains a lot of small helpers and extensions of Apache Spark/PySpark. Under the hood its PySpark part is based on the "py4j-pattern" (yes, again):

func = sc._jvm.uk.co.gresearch.spark.__getattr__("package$").__getattr__("MODULE$").dotNetTicksToTimestamp

link to the code above

Sparkling Water is an official way to run H2O models on Apache Spark cluster. Repository has 955 stars and is actively maintained. Under the hood one may again find "py4j-pattern" that is based on _jvm / _jdf:

class H2OTargetEncoderModel(H2OTargetEncoderMOJOParams, JavaModel, JavaMLWritable):

    def transform(self, dataset):
        callerFrame = inspect.stack()[1]
        inTrainingMode = (callerFrame[3] == '_fit') & callerFrame[1].endswith('pyspark/ml/pipeline.py')
        if inTrainingMode:
            return self.transformTrainingDataset(dataset)
        else:
            return super(H2OTargetEncoderModel, self).transform(dataset)

    def transformTrainingDataset(self, dataset):
        self._transfer_params_to_java()
        return DataFrame(self._java_obj.transformTrainingDataset(dataset._jdf), dataset.sql_ctx)

link to the code above

gimel is a quite popular (239 stars) framework that is built on top of Apache Spark. In the documentation they directly recommend to use "py4j-pattern":

# import DataFrame and SparkSession
from pyspark.sql import DataFrame, SparkSession, SQLContext

# fetch reference to the class in JVM
ScalaDataSet = sc._jvm.com.paypal.gimel.DataSet

# fetch reference to java SparkSession
jspark = spark._jsparkSession

# initiate dataset
dataset = ScalaDataSet.apply(jspark)

# Read Data | kafka semantics abstracted for user
df = dataset.read("kafka_dataset")

# Apply transformations (business logic | abstracted for Gimel)
transformed_df = df(...transformations...)

# Write Data | Elastic semantics abstracted for user
dataset.write("elastic_dataset",df)

link to the code above

HNSWlib is a quite popular (240 stars) and modern JVM library for an Approximate Nearest Neighbors Search. hnswlib-spark is an Apache Spark/PySpark wrapper on top of the main library. And under the hood PySpark part is partially based on a "py4j-pattern" by using SparkContext constructor:

def __init__(self):
    spark_conf = SparkConf()
    spark_conf.setAppName(spark_nlp_config.app_name)
    spark_conf.setMaster(spark_nlp_config.master)
    spark_conf.set("spark.driver.memory", memory)
    spark_conf.set("spark.serializer", spark_nlp_config.serializer)
    spark_conf.set("spark.kryo.registrator", spark_nlp_config.registrator)
    spark_conf.set("spark.jars.packages", spark_nlp_config.maven_spark)
    spark_conf.set("spark.hnswlib.settings.index.cache_folder", cache_folder)

    # Make the py4j JVM stdout and stderr available without buffering
    popen_kwargs = {
        'stdout': subprocess.PIPE,
        'stderr': subprocess.PIPE,
        'bufsize': 0
    }

    # Launch the gateway with our custom settings
    self.gateway = launch_gateway(conf=spark_conf, popen_kwargs=popen_kwargs)
    self.process = self.gateway.proc
    # Use the gateway we launched
    spark_context = SparkContext(gateway=self.gateway)
    self.spark_session = SparkSession(spark_context)

    self.out_thread = threading.Thread(target=self.output_reader)
    self.error_thread = threading.Thread(target=self.error_reader)
    self.std_background_listeners()

link to the code above

AUT is a tool and a library to analyze Web Archives on Apache Spark clusters. Its PySpark part uses the same "py4j-pattern":

class WebArchive:
    def __init__(self, sc, sqlContext, path):
        self.sc = sc
        self.sqlContext = sqlContext
        self.loader = sc._jvm.io.archivesunleashed.df.DataFrameLoader(sc._jsc.sc())
        self.path = path

link to the code above

Linkis is a top-level Apache project (3.2k stars). It's PySpark part is heavily based on the same "py4j-pattern":

jsc = intp.getJavaSparkContext()
jconf = intp.getSparkConf()
conf = SparkConf(_jvm = gateway.jvm, _jconf = jconf)
sc = SparkContext(jsc=jsc, gateway=gateway, conf=conf)
sqlc = HiveContext(sc, intp.sqlContext())
sqlContext = sqlc
spark = SparkSession(sc, intp.getSparkSession())

link to the code above

spark-dgraph-connector is an another project from G-Research. It's PySpark part uses the same "py4j-pattern":

class DgraphReader:
    def __init__(self, reader: DataFrameReader):
        super().__init__()
        self._jvm = reader._spark._jvm
        self._spark = reader._spark
        self._reader = self._jvm.uk.co.gresearch.spark.dgraph.connector.DgraphReader(reader._jreader)

link to the code above

And finally a project where I'm a contributor and maintainer of PySpark part: GraphAr. GraphAr is a novel way to store huge Graph data in DataLake or LakeHouse solutions. The whole PySpark part is based on "py4j-pattern". Mostly because Synapse ML and PyDeequ were main sources of inspiration for me when I worked on the implementation…

I'm more than sure that many companies using Databricks have their own in-house helpers, libraries, etc. And I'm more than sure that a lot of these in-house projects rely on the same "py4j-pattern".