Repository: OpenLineage/OpenLineage

Issue #181: [INTEGRATION] [Spark] Making the spark integration extensible
Created: 2021-04-06T20:03:51Z
Context

The SparkListener approach was already implemented in the Spark integration prior to the code being moved from Marquez to OpenLineage. Much of the context here is obsolete. However, the extensibility issue still needs to be addressed.

The Spark integration relies on instrumentation achieved by adding a jar to the javaagent argument on the command line when starting a Spark job. When the JVM starts up, the application ClassLoader loads both the instrumentation jar as well as the Spark, Hadoop, and other dependency jars that are on the classpath. When the SparkContext is loaded, Spark parses its own configuration and finds additional jars that have been specified, either on the command line or in a configuration map that can be populated at runtime. The jars are (sometimes) downloaded and added to a MutableURLClassLoader that is owned and managed by the SparkContext.
Since the Marquez integration is loaded by the top level application ClassLoader, and Java’s ClassLoaders are ignorant of their child ClassLoaders, the integration is only able to find classes that are loaded by that top level ClassLoader. BigQuery and other dependencies may be specified in the spark.jars configuration, so that they are loaded by that child MutableURLClassLoader, hence invisible to the Marquez agent classes. While it is possible from the integration to detect that BigQuery classes are available, attempting to instantiate the classes that actually deal with those classes cause NoClassDefFoundErrors, since the ClassLoader can’t find the classes loaded by the child loader.
The current workaround relies on instrumentation to invoke methods from the parent ClassLoader without needing to directly access the BigQueryRelation class or related classes.
I think the instrumentation approach has the following drawbacks

We lose type safety and the ability to easily unit test the instrumentation code
Failure to instrument the code at runtime leaves users blind to the missing dataset nodes
Isn't extensible to support other integrations

To be fair, it has the following advantages

Only a single jar is needed at runtime
It's possible for a single jar to support multiple, statically non-compatible versions of the BigQuery integration using conditional checks (e.g., if ( class.hasMethod("X")) generateXSpecificCode())

Reflection has been suggested as an alternative strategy, but I think the pros/cons are largely the same as the instrumentation approach- it uses indirect method invocation to circumvent static type checking.
Proposal
I propose a more extensible method of loading and discovering node visitors. The basic strategy is

Create a common interface for visiting logical plan nodes- something that extends the PartitalFunction<LogicalPlan, List<Dataset>> interface we already use
Take advantage of the JVM's ServiceLoader capabilities to enable implementations on the classpath to be found at runtime
Package node visitors in different subprojects that generate their own jars
Entrust Spark users to declare dependencies on the node visitors they need using the spark.jars.packages or spark.jars configuration properties

The main advantages I see here are

Static typing, easy to read code, and easy unit testing
Extensible framework to support new integrations
This point, I feel, is extremely important. While the single jar deployment is convenient, it creates an unnecessary barrier for users who want to take part in the OpenLineage ecosystem, but can't/won't contribute directly to Marquez or OpenLineage. This may be because


their company has strict open source contribution policies
they need to work with some proprietary, custom datasource. E.g., a robotics company might have a datasource specific to ROS topics that are particular to that company
they want to deploy their integration without going through the community review/release process
Forcing developers to fork the repo in order to add custom integrations is a deterrent to building an ecosystem.


Adding integrations is less complex, as developers don't need to worry about, e.g., conflicting versions of Guava required by the BigQueryRelation and the HBaseRelation.

The tradeoff here is that users need to declare the MarquezOpenLineage integration dependencies and deploy them separately from the main jar. Statically typed dataset integrations should be loaded by the Spark ClassLoader (e.g., declared in the spark.jars.packages or spark.jars configuration) in order to link to the Datasource classes they depend on. In terms of developer experience, I don't think this pattern is foreign to Spark users today.
Notably, nothing prohibits users from bundling their own superjars that contain all of their dependencies and deploying them to their Spark clusters. Similarly, nothing prevents integration authors from using reflection in their code. Users are given the freedom to extend and package their integrations in a way that suits them without enforcing an opaque, black-box model.
As a separate issue, we can consider changing the MarquezOpenLineage agent to be loaded by the spark.extraListeners configuration. This would allow users to reference all of the MarquezOpenLineage integration dependencies by specifying the listener and the node visitors in spark.jars.packages so that all of the classes are loaded by the Spark ClassLoader. That eliminates any confusion about which components need to be loaded by which ClassLoader. It also simplifies testing and could enable integration in serverless environments, such as AWS Glue, where users don't have access to modify the JVM arguments.

This has already been implemented

Packaging Spark Integrations
Spark integration packages need to create a file in META-INF/services/marquez.spark.agent.lifecycle.plan.DatasetSource which declares a list of classes that implement the openlineage.spark.agent.lifecycle.plan.DatasetSource interface. During startup, the main MarquezOpenLineage agent uses the JVM ServiceLoader to find all service implementations on the classpath (we need to specifically check the Spark URLClassloader to find implementations loaded by Spark at runtime). Combining jars requires concatenating the contents of the META-INF/services/marquez.spark.agent.lifecycle.plan.DatasetSource files so that all implementations are listed.
Deploying Dependencies
Currently, users deploy the Marquez integration by adding the javaagent:<jar>=http://marquez_url/path/to/job argument to their spark-submit args. That doesn't change with this proposal.
What does change is that users would be required to specify the integrations they require in their spark.jars or spark.jars.packages configuration. As an example, the airflow example in the README would look like this:
from uuid import uuid4
import os
...
job_name = 'job_name'


properties = {
  'spark.extraListeners': 'openlineage.spark.agent.OpenLineageSparkListener',
  'spark.jar.packages': 'io.openlineage:openlineage-spark:0.0.1',
  'spark.openlineage.host': os.environ.get('MARQUEZ_URL'),
  'spark.openlineage.namespace': os.getenv('MARQUEZ_NAMESPACE', 'default'),
  'spark.openlineage.apiKey': os.environ.get('MARQUEZ_API_KEY'),
  'spark.openlineage.parentRunId': "{{task_run_id(run_id, task)}}",
  'spark.openlineage.parentJobName': '{{ task.task_id }}'
}

t1 = DataProcPySparkOperator(
    task_id=job_name,
    gcp_conn_id='google_cloud_default',
    project_id='project_id',
    cluster_name='cluster-name',
    region='us-west1',
    main='gs://bucket/your-prog.py',
    job_name=job_name,
    dataproc_pyspark_properties=properties,

    #Both bigquery relation and bigquery openlineage integration declared here
    dataproc_pyspark_jars='gs://spark-lib/bigquery/spark-bigquery-latest_2.12.jar,gs://my_bucket/path/to/bigquery-integration.jar', 
    dag=dag)
Alternatively, a cluster can be created by loading the correct jars into the Spark jars directory, as described in the BigQuery connector docs. In the same way, the MarquezOpenLineage integrations would need to be installed on the driver node. In either case, the person responsible for loading the bigquery jars on the classpath is also responsible for loading the Marquez integration. Dependencies can be bundled so that the Marquez integration and the BigQueryRelation class are all loaded in the same jar.

Duplicate issue #410: [INTEGRATION][spark] QueryPlanVisitor should support optional dependencies
Created: 2021-11-16T17:32:13Z
BigQueryNodeVisitor has method hasBigQueryClasses() that checks for existence of bigquery classes at runtime. If bigquery dependency isn't there, the visitor is skipped.
This pattern happens or should happen in many cases where logical plan node isn't bundled with core spark sql. We should add generalized method to QueryPlanVisitor interface and review visitors whether they need to have to implement this method.
Discussed here: #387

Closing comment (2021-11-18T01:43:21Z): Isn't this similar to #181 as well?