[CT-431] Support for Pyspark driver

[cccs-jc]

Describe the feature

Add a forth connection option to the dbt-adapter. This forth connection would create a pyspark context and utilize the spark.sql(<sql>) function to execute sql statements.

Describe alternatives you've considered

The alternative is to run your own thrift server which is difficult to setup.

Who will this benefit?

A pyspark context gives you control over the spark application configuration and amount of resources to use in the spark cluster. It also enables the user to register custom pyspark UDF functions and to create custom views based on pyspark Dataframes.

1. control spark job creation
2. register custom pyspark UDF
3. register custom views based on pypsark Dataframes

Are you interested in contributing this feature?

Yes for sure. I have an implementation of this feature which I would like to contribute to this project.

Here's how it works from the user's point of view.

Specify a new connection method pyspark

profiles.yml

project1:
  target: dev
  outputs:
    dev:
      type: spark
      method: pyspark
      python_module: spark.spark_context

The python_module points to a python file found in the PYTHONPATH. In this case spark/spark_context.py. The pyspark adapter will call the create_spark_context() function found in this file. The user can thus create their own spark context (either a local instance or one that leverages a cluster). This hook also lets you create custom pyspark UDF registrations.

spark/spark_context.py

def create_spark_context():
    return SparkSession.builder.getOrCreate()

A second hook is available in the the sources. This hook lets the user register custom pyspark Dataframe as an sql view df.createOrReplaceTempView(<view name>). The user is free to create any Dataframe logic they need.

sources.yml

sources:
  - name: custom
    tables:
    - name: raw_users
      meta:
        python_module: models.staging.raw_users

The python_module is loaded the same way except that the hook function has a different signature.

def create_dataframe(spark, start_time, end_time)

Using pyspark you can work around the limitations of SparkSQL when reading datafiles. For example using pyspark you can load the schema of the json you are going to read. Thus avoiding sparks schema discovery process and making sure that the data is read in the schema you want.

models/staging/raw_users.py

def create_dataframe(spark, start_time, end_time):
    data = < load schema from a file>
    schema = StructType.fromJson(json.loads(data))

    df = (spark
	.read
	.schema(schema)
	.format("json")
	.load(f"{WAREHOUSE}/{YEAR}/{MONTH}/{DAY}/users/*.logs")
	)
    return df

The implementation of this feature consist of a new PysparkConnectionWrapper which executes sql statements via the spark context

def execute(self, sql, bindings=None):
    self.result = self.spark.sql(sql)

And a new method on the SparkRelation which register pyspark dataframes as sql views.

def load_python_module(self, start_time, end_time):
    path = self.meta.get('python_module')
    module = importlib.import_module(path)
    create_dataframe = getattr(module, "create_dataframe")
    df = create_dataframe(spark, start_time, end_time)
    df.createOrReplaceTempView(self.identifier)

Registration of pyspark views is initiated by a modified source macro

{% macro source(source_name, identifier, start_dt = None, end_dt = None) %}
    {%- set relation = builtins.source(source_name, identifier) -%}
    {%- if execute and (relation.source_meta.python_module or relation.meta.python_module) -%}

        {%- do relation.load_python_module(start_dt, end_dt) -%}

        {%- do return(relation.identifier) -%}
    {% else -%}
        {%- do return(relation) -%}
    {% endif -%}
{% endmacro %}

[cccs-jc]

created a pull request #308

[jtcohen6]

@cccs-jc Really interesting! Thanks for opening the issue, and the accompanying PR.

We just merged a new connection method (session) that can leverage a PySpark session (#279). I'd be very curious to get your eyes on that implementation (and @JCZuurmond's eyes on yours, in #308). What are the crucial differences between the two approaches? For my part, I'm not convinced that dbt's role should be to manage the Spark session's instantiation/configuration, versus simply being able to connect to one that already exists.

As always, the questions are:

• Which generally-applicable guardrails do we want to keep in place, (separation of concerns between dbt + database "infrastructure")
• Where does Spark blur those boundaries, in a way that gets us more than what's available elsewhere (all-on-one-machine "unit" testing of macro/materialization/etc functionality)

models/staging/raw_users.py

Neat :) We're definitely thinking about what it might look like to implement "dataframe" support in dbt + Spark/Databricks. I don't think a dbt-managed pyspark is session is the way we'd want to go, though.

[cccs-jc]

@JCZuurmond's implementation is very close to mine. The main differences are that mine has these additional features

A hook to create custom spark session

• In this hook you can create a spark context which specifies RAM/CPU to use from a cluster, configure iceberg catalogs etc
• In this hook you can also register custom UDF which can be use in SQL transformations

The only way I know how to register UDF via SQL requires the python package to be installed in a pre-existing server which is very inconvenient.

A hook to register a custom view

• This hook lets you register a pyspark dataframe as a dbt source. This lets you work around the limitations of SparkSQL. For example using pyspark you can load the schema of the json you are going to read. Thus avoiding spark's schema discovery process and making sure that the data is read in the schema you want.

Creating views in SQL (dbt-externals) is possible but it has limitations which the programmatic approach does not.

Being able to create custom python UDF and views is important to us. Curious to know what alternatives you are thinking about?

[JCZuurmond]

@cccs-jc : Great minds think alike! 💯

A hook to create custom spark session.

You solve this by first creating the Spark session with the config and configuration that you prefer. Then, SparkSession.builder.getOrCreate() will get this context for you. There is no hook needed.

A hook to register a custom view

Why not do this outside of dbt or right before you call dbt programmatically with the Spark session approach? The example you give about the json schema is solved by using dbt external tables.

If you have improvements to the current implementation, feel free to add changes to the session module!

[cccs-jc]

Our initial proof of concept used the approach you mention. That is a modified dbt launch script which creates a spark context and registers views. The in the dbt adapter retrieve the session and execute sql.

The issues with this approach are:

• You need a modified dbt launch script so if you want to use the vscode-power-dbt extension you need a modified dbt-core package.
• Since dbt-adapter retrieves the spark session and is already aware of spark sessions why not have it use a hook to create one.
• Registering views can take a bit of time. When you have a few dozen views and you launch dbt for a specific model why register these views. Perhaps the model you run does not even use any of the views. Seems to make more sense to create the views as they are needed. Thus the hook into the source processing.

I'm aware of dbt-external tables however it's not as flexible as pyspark code. We built pyspark libraries to create dataframe out of our raw data (hundreds of tables) and would like to leverage them within dbt.

Also I'm not sure dbt-external could handle use cases like reading from Azure Kusto. https://docs.microsoft.com/en-us/azure/data-explorer/spark-connector

The improvement I would like to contribute is hooks into spark session creation and view creation.