PySpark for Data Science Workflows | by Ben Weber | Towards Data Science

Demonstrated experience in PySpark is one of the most desirable competencies that employers are looking for when building data science teams, because it enables these teams to own live data products. While I’ve previously blogged about PySpark, Parallelization, and UDFs, I wanted to provide a proper overview of this topic as a book chapter. I’m sharing this complete chapter, because I want to encourage the adoption of PySpark as a tool for data scientists. All code examples from this post are available here, and all prerequisites are covered in the sample chapters here. You might want to grab some snacks before diving in!
Spark is a general-purpose computing framework that can scale to massive data volumes. It builds upon prior big data tools such as Hadoop and MapReduce, while providing significant improvements in the expressivity of the languages it supports. One of the core components of Spark is resilient distributed datasets (RDD), which enable clusters of machines to perform workloads in a coordinated, and fault-tolerant process. In more recent versions of Spark, the Dataframe API provides an abstraction on top of RDDs that resembles the same data structure in R and Pandas. PySpark is the Python interface to Spark, and it provides an API for working with large-scale datasets in a distributed computing environment.
PySpark is an extremely valuable tool for data scientists, because it can streamline the process for translating prototype models into production-grade model workflows. At Zynga, our data science team owns a number of production-grade systems that provide useful signals to our game and marketing teams. By using PySpark, we’ve been able to reduce the amount of support we need from engineering teams to scale up models from concept to production.
Up until now in this book, all of the models we’ve built and deployed have been targeted at single machines. While we are able to scale up model serving to multiple machines using Lambda, ECS, and GKS, these containers worked in isolation and there was no coordination among nodes in these environments. With PySpark, we can build model workflows that are designed to operate in cluster environments for both model training and model serving. The result is that data scientists can now tackle much larger-scale problems than previously possible using prior Python tools. PySpark provides a nice tradeoff between an expressive programming language and APIs to Spark versus more legacy options such as MapReduce. A general trend is that the use of Hadoop is dropping as more data science and engineering teams are switching to Spark ecosystems. In Chapter 7 we’ll explore another distributed computing ecosystem for data science called Cloud Dataflow, but for now Spark is the open-source leader in this space. PySpark was one of the main motivations for me to switch from R to Python for data science workflows.
The goal of this chapter is to provide an introduction to PySpark for Python programmers that shows how to build large-scale model pipelines for batch scoring applications, where you may have billions of records and millions of users that need to be scored. While production-grade systems will typically push results to application databases, in this chapter we’ll focus on batch processes that pull in data from a lake and push results back to the data lake for other systems to use. We’ll explore pipelines that perform model applications for both AWS and GCP. While the data sets used in this Chapter rely on AWS and GCP for storage, the Spark environment does not have to run on either of these platforms and instead can run on Azure, other clouds, or on-pem Spark clusters.
We’ll cover a variety of different topics in this chapter to show different use cases of PySpark for scalable model pipelines. After showing how to make data available to Spark on S3, we’ll cover some of the basics of PySpark focusing on Dataframe operations. Next, we’ll build out a predictive model pipeline that reads in data from S3, performs batch model predictions, and then writes the results to S3. We’ll follow this by showing off how a newer feature called Pandas UDFs can be used with PySpark to perform distributed deep learning and feature engineering. To conclude, we’ll build another batch model pipeline now using GCP and then discuss how to productize workflows in a Spark ecosystem.
There’s a variety of ways to both configure Spark clusters and submit commands to a cluster for execution. When getting started with PySpark as a data scientist, my recommendation is to use a freely-available notebook environment for getting up and running with Spark as quick as possible. While PySpark may not perform quite as well as Java or Scala for large-scale workflows, the ease of development in an interactive programming environment is worth the tradeoff.
Based on your organization, you may be starting from scratch for Spark or using an existing solution. Here are the types of Spark deployments I’ve seen in practice:
Self Hosted: An engineering team manages a set of clusters and provides console and notebook access.
Cloud Solutions: AWS provides a managed Spark option called EMR and GCP has Cloud DataProc.
Vendor Solutions: Databricks, Cloudera, and other vendors provide fully-managed Spark environments.
There’s a number of different factors to consider when choosing a Spark ecosystem, including cost, scalability, and feature sets. As you scale the size of the team using Spark, additional considerations are whether an ecosystem supports multi-tenancy, where multiple jobs can run concurrently on the same cluster, and isolation where one job failing should not kill other jobs. Self-hosted solutions require significant engineering work to support these additional considerations, so many organizations use cloud or vendor solutions for Spark. In this book, we’ll use the Databricks Community Edition, which provides all of the baseline features needed for learning Spark in a collaborative notebook environment.
Spark is a rapidly evolving ecosystem, and it’s difficult to author books about this subject that do not quickly become out of date as the platform evolves. Another issue is that many books target Scala rather than Python for the majority of coding examples. My advice for readers that want to dig deeper into the Spark ecosystem is to explore books based on the broader Spark ecosystem, such as (Karau et al. 2015). You’ll likely need to read through Scala or Java code examples, but the majority of content covered will be relevant to PySpark.
A Spark environment is a cluster of machines with a single driver node and one or more worker nodes. The driver machine is the master node in the cluster and is responsible for coordinating the workloads to perform. In general, workloads will be distributed across the worker nodes when performing operations on Spark dataframe. However, when working with native Python objects, such as lists or dictionaries, objects will be instantiated on the driver node.
Ideally, you want all of your workloads to be operating on worker nodes, so that the execution of the steps to perform is distributed across the cluster, and not bottlenecked by the driver node. However, there are some types of operations in PySpark where the driver has to perform all of the work. The most common situation where this happens is when using Pandas dataframes in your workloads. When you use toPandas or other commands to convert a data set to a Pandas object, all of the data is loaded into memory on the driver node, which can crash the driver node when working with large data sets.
In PySpark, the majority of commands are lazily executed, meaning that an operation is not performed until an output is explicitly needed. For example, a join operation between two Spark dataframes will not immediately cause the join operation to be performed, which is how Pandas works. Instead, the join is performed once an output is added to the chain of operations to perform, such as displaying a sample of the resulting dataframe. One of the key differences between Pandas operations, where operations are eagerly performed and pulled into memory, is that PySpark operations are lazily performed and not pulled into memory until needed. One of the benefits of this approach is that the graph of operations to perform can be optimized before being sent to the cluster to execute.
In general, nodes in a Spark cluster should be considered ephemeral, because a cluster can be resized during execution. Additionally, some vendors may spin up a new cluster when scheduling a job to run. This means that common operations in Python, such as saving files to disk, do not map directly to PySpark. Instead, using a distributed computing environment means that you need to use a persistent file store such as S3 when saving data. This is important for logging, because a worker node may crash and it may not be possible to ssh into the node for debugging. Most Spark deployments have a logging system set up to help with this issue, but it’s good practice to log workflow status to persistent storage.
One of the quickest ways to get up and running with PySpark is to use a hosted notebook environment. Databricks is the largest Spark vendor and provides a free version for getting started called Community Edition [13]. We’ll use this environment to get started with Spark and build AWS and GCP model pipelines.
The first step is to create a login on the Databricks website for the community edition. Next, perform the following steps to spin up a test cluster after logging in:
Click on “Clusters” on the left navigation barClick “Create Cluster”Assign a name, “DSP”Select the most recent runtime (non-beta)Click “Create Cluster”
Click on “Clusters” on the left navigation bar
Click “Create Cluster”
Assign a name, “DSP”
Select the most recent runtime (non-beta)
Click “Create Cluster”
After a few minutes we’ll have a cluster set up that we can use for submitting Spark commands. Before attaching a notebook to the cluster, we’ll first set up the libraries that we’ll use throughout this chapter. Instead of using pip to install libraries, we’ll use the Databricks UI, which makes sure that every node in the cluster has the same set of libraries installed. We’ll use both Maven and PyPI to install libraries on the clusters. To install the BigQuery connector, perform the following steps:
Click on “Clusters” on the left navigation barSelect the “DSP” clusterClick on the “Libraries” tabSelect “Install New”Click on the “Maven” tab.Set coordinates to com.spotify:spark-bigquery_2.11:0.2.2Click install
Click on “Clusters” on the left navigation bar
Select the “DSP” cluster
Click on the “Libraries” tab
Select “Install New”
Click on the “Maven” tab.
Set coordinates to com.spotify:spark-bigquery_2.11:0.2.2
Click install
The UI will then show the status as resolving, and then installing, and then installed. We also need to attach a few Python libraries that are not pre-installed on a new Databricks cluster. Standard libraries such as Pandas are installed, but you might need to upgrade to a more recent version since the libraries pre-installed by Databricks can lag significantly.
To install a Python library on Databricks, perform the same steps as before up to step 5. Next, instead of selecting “Maven” choose “PyPI”. Under Package, specify the package you want to install and then click “Install”. To follow along with all of the sections in this chapter, you’ll need to install the following Python packages:
koalas — for Dataframe conversion
featuretools — for feature generation
tensorflow — for a deep learning backend
keras — for a deep learning model
You’ll now have a cluster set up capable of performing distributed feature engineering and deep learning. We’ll start with basic Spark commands, show off newer functionality such as the Koalas library, and then dig into these more advanced topics. After set up, your cluster library setup should look like Figure 6.1. To ensure that everything is set up successfully, restart the cluster and check the status of the installed libraries.
Now that we have provisioned a cluster and set up the required libraries, we can create a notebook to start submitting commands to the cluster. To create a new notebook, perform the following steps:
Click on “Databricks” on the left navigation barUnder “Common Tasks”, select “New Notebook”Assign a name “CH6”Select “Python” as the languageselect “DSP” as the clusterClick “Create”
Click on “Databricks” on the left navigation bar
Under “Common Tasks”, select “New Notebook”
Assign a name “CH6”
Select “Python” as the language
select “DSP” as the cluster
Click “Create”
The result will be a notebook environment where you can start running Python and PySpark commands, such as print("Hello World!"). An example notebook running this command is shown in Figure 6.2. We now have a PySpark environment up and running that we can use to build distributed model pipelines.
Data is essential for PySpark workflows. Spark supports a variety of methods for reading in data sets, including connecting to data lakes and data warehouses, as well as loading sample data sets from libraries, such as the Boston housing data set. Since the theme of this book is building scalable pipelines, we’ll focus on using data layers that work with distributed workflows. To get started with PySpark, we’ll stage input data for a modeling pipeline on S3, and then read in the data set as a Spark dataframe.
This section will show how to stage data to S3, set up credentials for accessing the data from Spark, and fetching the data from S3 into a Spark dataframe. The first step is to set up a bucket on S3 for storing the data set we want to load. To perform this step, run the following operations on the command line.
aws s3api create-bucket --bucket dsp-ch6 --region us-east-1aws s3 ls
After running the command to create a new bucket, we use the ls command to verify that the bucket has been successfully created. Next, we’ll download the games data set to the EC2 instance and then move the file to S3 using the cp command, as shown in the snippet below.
wget https://github.com/bgweber/Twitch/raw/master/                            Recommendations/games-expand.csvaws s3 cp games-expand.csv s3://dsp-ch6/csv/games-expand.csv
In addition to staging the games data set to S3, we’ll also copy a subset of the CSV files from the Kaggle NHL data set, which we set up in Section 1.5.2. Run the following commands to stage the plays and stats CSV files from the NHL data set to S3.
aws s3 cp game_plays.csv s3://dsp-ch6/csv/game_plays.csvaws s3 cp game_skater_stats.csv                   s3://dsp-ch6/csv/game_skater_stats.csvaws s3 ls  s3://dsp-ch6/csv/
We now have all of the data sets needed for the code examples in this chapter. In order to read in these data sets from Spark, we’ll need to set up S3 credentials for interacting with S3 from the Spark cluster.
For production environments, it is better to use IAM roles to manage access instead of using access keys. However, the community edition of Databricks constrains how much configuration is allowed, so we’ll use access keys to get up and running with the examples in this chapter. We already set up a user for accessing S3 from an EC2 instance. To create a set of credentials for accessing S3 programmatically, perform the following steps from the AWS console:
Search for and select “IAM”Click on “Users”Select the user created in Section 3.3.2, “S3_Lambda”Click “Security Credentials”Click “Create Access Key”
Search for and select “IAM”
Click on “Users”
Select the user created in Section 3.3.2, “S3_Lambda”
Click “Security Credentials”
Click “Create Access Key”
The result will be an access key and a secret key enabling access to S3. Save these values in a secure location, as we’ll use them in the notebook to connect to the data sets on S3. Once you are done with this chapter, it is recommended to revoke these credentials.
Now that we have credentials set up for access, we can return to the Databricks notebook to read in the data set. To enable access to S3, we need to set the access key and secret key in the Hadoop configuration of the cluster. To set these keys, run the PySpark commands shown in the snippet below. You’ll need to replace the access and secret keys with the credentials we just create for the S3_Lambda role.
AWS_ACCESS_KEY = "AK..."AWS_SECRET_KEY = "dC..."sc._jsc.hadoopConfiguration().set(                      "fs.s3n.awsAccessKeyId", AWS_ACCESS_KEY)sc._jsc.hadoopConfiguration().set(                      "fs.s3n.awsSecretAccessKey", AWS_SECRET_KEY)
We can now read the data set into a Spark dataframe using the read command, as shown below. This command uses the spark context to issue a read command and reads in the data set using the CSV input reader. We also specify that the CSV file includes a header row and that we want Spark to infer the data types for the columns. When reading in CSV files, Spark eagerly fetches the data set into memory, which can cause issues for larger data sets. When working with large CSV files, it’s a best practice to split up large data sets into multiple files and then read in the files using a wildcard in the input path. When using other file formats, such as Parquet or AVRO, Spark lazily fetches the data sets.
games_df = spark.read.csv("s3://dsp-ch6/csv/games-expand.csv",                              header=True, inferSchema = True)display(games_df)
The display command in the snippet above is a utility function provided by Databricks that samples the input dataframe and shows a table representation of the frame, as shown in Figure 6.3. It is similar to the head function in Pandas, but provides additional functionality such as transforming the sampled dataframe into a plot. We’ll explore the plotting functionality in Section 6.3.3.
Now that we have data loaded into a Spark dataframe, we can begin exploring the PySpark language, which enables data scientists to build production-grade model pipelines.
PySpark is a powerful language for both exploratory analysis and building machine learning pipelines. The core data type in PySpark is the Spark dataframe, which is similar to Pandas dataframes, but is designed to execute in a distributed environment. While the Spark Dataframe API does provide a familiar interface for Python programmers, there are significant differences in the way that commands issued to these objects are executed. A key difference is that Spark commands are lazily executed, which means that commands such as iloc are not available on these objects. While working with Spark dataframes can seem constraining, the benefit is that PySpark can scale to much larger data sets than Pandas.
This section will walk through common operations for Spark dataframes, including persisting data, converting between different dataframe types, transforming dataframes, and using user-defined functions. We’ll use the NHL stats data set, which provides user-level summaries of player performance for each game. To load this data set as a Spark dataframe, run the commands in the snippet below.
stats_df = spark.read.csv("s3://dsp-ch6/csv/game_skater_stats.csv",                                    header=True, inferSchema = True)display(stats_df)
A common operation in PySpark is saving a dataframe to persistent storage, or reading in a data set from a storage layer. While PySpark can work with databases such as Redshift, it performs much better when using distributed files stores such as S3 or GCS. In this chapter we’ll use these types of storage layers as the outputs of model pipelines, but it’s also useful to stage data to S3 as intermediate steps within a workflow. For example, in the AutoModel [14] system at Zynga, we stage the output of the feature generation step to S3 before using MLlib to train and apply model predictions.
The data storage layer to use will depend on your cloud platform. For AWS, S3 works well with Spark for distributed data reads and writes. When using S3 or other data lakes, Spark supports a variety of different file formats for persisting data. Parquet is typically the industry standard when working with Spark, but we’ll also explore Avro and ORC in addition to CSV. Avro is a better format for streaming data pipelines, and ORC is useful when working with legacy data pipelines.
To show the range of data formats supported by Spark, we’ll take the stats data set and write it to AVRO, then Parquet, then ORC, and finally CSV. After performing this round trip of data IO, we’ll end up with our initial Spark dataframe. To start, we’ll save the stats dataframe in Avro format, using the code snippet shown below. This code writes the dataframe to S3 in Avro format using the Databricks Avro writer, and then reads in the results using the same library. The result of performing these steps is that we now have a Spark dataframe pointing to the Avro files on S3. Since PySpark lazily evaluates operations, the Avro files are not pulled to the Spark cluster until an output needs to be created from this data set.
# AVRO writeavro_path = "s3://dsp-ch6/avro/game_skater_stats/"stats_df.write.mode('overwrite').format(                       "com.databricks.spark.avro").save(avro_path)# AVRO read avro_df = sqlContext.read.format(                       "com.databricks.spark.avro").load(avro_path)
Avro is a distributed file format that is record based, while the Parquet and OR formats are column based. It is useful for the streaming workflows that we’ll explore in Chapter 9, because it compresses records for distributed data processing. The output of saving the stats dataframe in Avro format is shown in the snippet below, which shows a subset of the status files and data files generated when persisting a dataframe to S3 as Avro. Like most scalable data formats, Avro will write records to several files, based on partitions if specified, in order to enable efficient read and write operations.
aws s3 ls s3://dsp-ch6/avro/game_skater_stats/2019-11-27 23:02:43       1455 _committed_15886175782508531572019-11-27 22:36:31       1455 _committed_16007797309378807952019-11-27 23:02:40          0 _started_15886175782508531572019-11-27 23:31:42          0 _started_69420741361908385862019-11-27 23:31:47    1486327 part-00000-tid-6942074136190838586-               c6806d0e-9e3d-40fc-b212-61c3d45c1bc3-15-1-c000.avro2019-11-27 23:31:43      44514 part-00007-tid-6942074136190838586-               c6806d0e-9e3d-40fc-b212-61c3d45c1bc3-22-1-c000.avro
Parquet on S3 is currently the standard approach for building data lakes on AWS, and tools such as Delta Lake are leveraging this format to provide highly-scalable data platforms. Parquet is a columnar-oriented file format that is designed for efficient reads when only a subset of columns are being accessed for an operation, such as when using Spark SQL. Parquet is a native format for Spark, which means that PySpark has built-in functions for both reading and writing files in this format.
An example of writing the stats dataframe as Parquet files and reading in the result as a new dataframe is shown in the snippet below. In this example, we haven’t set a partition key, but as with Avro, the dataframe will be split up into multiple files in order to support highly-performant read and write operations. When working with large-scale data sets, it’s useful to set partition keys for the file export using the repartition function. After this section, we’ll use Parquet as the primary file format when working with Spark.
# parquet outparquet_path = "s3a://dsp-ch6/games-parquet/"avro_df.write.mode('overwrite').parquet(parquet_path)# parquet inparquet_df = sqlContext.read.parquet(parquet_path)
ORC is a another columnar format that works well with Spark. The main benefit over Parquet is that it can support improved compression, at the cost of additional compute cost. I’m including it in this chapter, because some legacy systems still use this format. An example of writing the stats dataframe to ORC and reading the results back into a Spark dataframe is shown in the snippet below. Like the Avro format, the ORC write command will distribute the dataframe to multiple files based on the size.
# orc outorc_path = "s3a://dsp-ch6/games-orc/"parquet_df.write.mode('overwrite').orc(orc_path)# orc inorc_df = sqlContext.read.orc(orc_path)
To complete our round trip of file formats, we’ll write the results back to S3 in the CSV format. To make sure that we write a single file rather than a batch of files, we’ll use the coalesce command to collect the data to a single node before exporting it.This is a command that will fail with large data sets, and in general it’s best to avoid using the CSV format when using Spark. However, CSV files are still a common format for sharing data, so it’s useful to understand how to export to this format.
# CSV outcsv_path = "s3a://dsp-ch6/games-csv-out/"orc_df.coalesce(1).write.mode('overwrite').format( "com.databricks.spark.csv").option("header","true").save(csv_path)  # and CSV to finish the round trip csv_df = spark.read.csv(csv_path, header=True, inferSchema = True)
The resulting dataframe is the same as the dataframe that we first read in from S3, but if the data types are not trivial to infer, then the CSV format can cause problems. When persisting data with PySpark, it’s best to use file formats that describe the schema of the data being persisted.
While it’s best to work with Spark dataframes when authoring PySpark workloads, it’s often necessary to translate between different formats based on your use case. For example, you might need to perform a Pandas operation, such as selecting a specific element from a dataframe. When this is required, you can use the toPandas function to pull a Spark dataframe into the memory of the driver node. The code snippet below shows how to perform this task, display the results, and then convert the Pandas dataframe back to a Spark dataframe. In general, it’s best to avoid Pandas when authoring PySpark workflows, because it prevents distribution and scale, but it’s often the best way of expressing a command to execute.
stats_pd = stats_df.toPandas()stats_df = sqlContext.createDataFrame(stats_pd)
To bridge the gap between Pandas and Spark dataframes, Databricks introduced a new library called Koalas that resembles the Pandas API for Spark-backed dataframes. The result is that you can author Python code that works with Pandas commands that can scale to Spark-scale data sets. An example of converting a Spark dataframe to Koalas and back to Spark is shown in the following snippet. After converting the stats dataframe to Koalas, the snippet shows how to calculate the average time on ice as well as index into the Koalas frame. The intent of Koalas is to provide a Pandas interface to Spark dataframes, and as the Koalas library matures more Python modules may be able to take advantage of Spark. The output from the snippet shows that the average time on ice was 993 seconds per game.
import databricks.koalas as ksstats_ks = stats_df.to_koalas()stats_df = stats_ks.to_spark()print(stats_ks['timeOnIce'].mean())print(stats_ks.iloc[:1, 1:2])
During the development of this book, Koalas is still preliminary and only partially implemented, but it looks to provide a familiar interface for Python coders. Both Pandas and Spark dataframes can work with Koalas, and the snippet below shows how to go from Spark to Koalas to Pandas to Spark, and Spark to Pandas to Koalas to Spark.
# spark -> koalas -> pandas -> sparkdf = sqlContext.createDataFrame(stats_df.to_koalas().toPandas())# spark -> pandas -> koalas -> sparkdf = ks.from_pandas(stats_df.toPandas()).to_spark()
In general, you’ll be working with Spark dataframes when authoring code in a PySpark environment. However, it’s useful to be able to work with different object types as necessary to build model workflows. Koalas and Pandas UDFs provide powerful tools for porting workloads to large-scale data ecosystems.
The PySpark Dataframe API provides a variety of useful functions for aggregating, filtering, pivoting, and summarizing data. While some of this functionality maps well to Pandas operations, my recommendation for quickly getting up and running with munging data in PySpark is to use the SQL interface to dataframes in Spark called Spark SQL. If you’re already using the pandasql or framequery libraries, then Spark SQL should provide a familiar interface. If you’re new to these libraries, then the SQL interface still provides an approachable way of working with the Spark ecosystem. We’ll cover the Dataframe API later in this section, but first start with the SQL interface to get up and running.
Exploratory data analysis (EDA) is one of the key steps in a data science workflow for understanding the shape of a data set. To work through this process in PySpark, we’ll load the stats data set into a dataframe, expose it as a view, and then calculate summary statistics. The snippet below shows how to load the NHL stats data set, expose it as a view to Spark, and then run a query against the dataframe. The aggregated dataframe is then visualized using the display command in Databricks.
stats_df = spark.read.csv("s3://dsp-ch6/csv/game_skater_stats.csv",                                   header=True, inferSchema = True)stats_df.createOrReplaceTempView("stats")new_df = spark.sql("""  select player_id, sum(1) as games, sum(goals) as goals  from stats  group by 1  order by 3 desc  limit 5""")display(new_df)
An output of this code block is shown in Figure 6.4. It shows the highest scoring players in the NHL dataset by ranking the results based on the total number of goals. One of the powerful features of Spark is that the SQL query will not operate against the dataframe until a result set is needed. This means that commands in a notebook can set up multiple data transformation steps for Spark dataframes, which are not performed until a later step needs to execute the graph of operations defined in a code block.
Spark SQL is expressive, fast, and my go-to method for working with big data sets in a Spark environment. While prior Spark versions performed better with the Dataframe API versus Spark SQL, the difference in performance is now trivial and you should use the transformation tools that provide the best iteration speed for working with large data sets. With Spark SQL, you can join dataframes, run nested queries, set up temp tables, and mix expressive Spark operations with SQL operations. For example, if you want to look at the distribution of goals versus shots in the NHL stats data, you can run the following command on the dataframe.
display(spark.sql("""  select cast(goals/shots * 50 as int)/50.0 as Goals_per_shot      ,sum(1) as Players   from (    select player_id, sum(shots) as shots, sum(goals) as goals    from stats    group by 1    having goals >= 5  )    group by 1  order by 1"""))
This query restricts the ratio of goals to shots to players with more than 5 goals, to prevent outliers such as goalies scoring during power plays. We’ll use the display command to output the result set as a table and then use Databricks to display the output as a graph. Many Spark ecosystems have ways of visualizing results, and the Databricks environment provides this capability through the display command, which works well with both tabular and pivot table data. After running the above command, you can click on the chart icon and choose dimensions and measures which show the distribution of goals versus shots, as visualized in Figure 6.5.
While I’m an advocate of using SQL to transform data, since it scales to different programming environments, it’s useful to get familiar with some of the basic dataframe operations in PySpark. The code snippet below shows how to perform common operations including dropping columns, selecting a subset of columns, and adding new columns to a Spark dataframe. Like prior commands, all of these operations are lazily performed. There are some syntax differences from Pandas, but the general commands used for transforming data sets should be familiar.
from pyspark.sql.functions import lit# dropping columnscopy_df = stats_df.drop('game_id', 'player_id')# selection columns copy_df = copy_df.select('assists', 'goals', 'shots')# adding columnscopy_df = copy_df.withColumn("league", lit('NHL'))display(copy_df)
One of the common operations to perform between dataframes is a join. This is easy to express in a Spark SQL query, but sometimes it is preferable to do this programmatically with the Dataframe API. The code snippet below shows how to join two dataframes together, when joining on the game_id and player_id fields. The league column which is a literal will be joined with the rest of the stats dataframe. This is a trivial example where we are adding a new column onto a small dataframe, but the join operation from the Dataframe API can scale to massive data sets.
copy_df = stats_df.select('game_id', 'player_id').                          withColumn("league", lit('NHL'))df = copy_df.join(stats_df, ['game_id', 'player_id'])display(df)
The result set from the join operation above is shown in Figure 6.6. Spark supports a variety of different join types, and in this example we used an inner join to append league the league column to the stats dataframe.
It’s also possible to perform aggregation operations on a dataframe, such as calculating sums and averages of columns. An example of computing the average time on ice for players in the stats data set, and total number of goals scored is shown in the snippet below. The groupBy command uses the player_id as the column for collapsing the data set, and the agg command specifies the aggregations to perform.
summary_df = stats_df.groupBy("player_id").agg(                             {'timeOnIce':'avg', 'goals':'sum'})display(summary_df)
The snippet creates a dataframe with player_id, timeOnIce, and goals columns. We’ll again use the plotting functionality in Databricks to visualize the results, but this time select the scatter plot option. The resulting plot of goals versus time on ice is shown in Figure 6.7.
We’ve worked through introductory examples to get up and running with dataframes in PySpark, focusing on operations that are useful for munging data prior to training machine learning models. These types of operations, in combination with reading and writing dataframes provides a useful set of skills for performing exploratory analysis on massive data sets.
While PySpark provides a great deal of functionality for working with dataframes, it often lacks core functionality provided in Python libraries, such as curve-fitting functions in SciPy. While it’s possible to use the toPandas function to convert dataframes into the Pandas format for Python libraries, this approach breaks when using large data sets. Pandas UDFs are a newer feature in PySpark that help data scientists work around this problem, by distributing the Pandas conversion across the worker nodes in a Spark cluster. With a Pandas UDF, you define a group by operation to partition the data set into dataframes that are small enough to fit into the memory of worker nodes, and then author a function that takes in a Pandas dataframe as the input parameter and returns a transformed Pandas dataframe as the result. Behind the scenes, PySpark uses the PyArrow library to efficiently translate dataframes from Spark to Pandas and back from Pandas to Spark. This approach enables Python libraries, such as Keras, to be scaled up to a large cluster of machines.
This section will walk through an example problem where we need to use an existing Python library and show how to translate the workflow into a scalable solution using Pandas UDFs. The question we are looking to answer is understanding if there is a positive or negative relationship between the shots and hits attributes in the stats data set. To calculate this relationship, we can use the leastsq function in SciPy, as shown in the snippet below. This example creates a Pandas dataframe for a single player_id, and then fits a simple linear regression between these attributes. The output is the coefficients used to fit the least-squares operation, and in this case the number of shots was not strongly correlated with the number of hits.
sample_pd = spark.sql("""  select * from stats  where player_id = 8471214""").toPandas()# Import python libraries from scipy.optimize import leastsqimport numpy as np# Define a function to fitdef fit(params, x, y):    return (y - (params[0] + x * params[1] ))  # Fit the curve and show the results result = leastsq(fit, [1,0], args=(sample_pd.shots,sample_pd.hits))print(result)
Now we want to perform this operation for every player in the stats data set. To scale to this volume, we’ll first partition by player_id, as shown by the groupBy operation in the code snippet below. Next, We’ll run the analyze_player function for each of these partitioned data sets using the apply command. While the stats_df dataframe used as input to this operation and the players_df dataframe returned are Spark dataframes, the sampled_pd dataframe and the dataframe returned by the analyze player function are Pandas. The Pandas UDF annotation provides a hint to PySpark for how to distribute this workload so that it can scale the operation across the cluster of worker nodes rather than eagerly pulling all of the data to the driver node. Like most Spark operations, Pandas UDFs are lazily evaluated and will not be executed until and output value is needed.
Our initial example now translated to use Pandas UDFs is shown below. After defining additional modules to include, we specify the schema of the dataframe that will be returned from the operation. The schema object defines the structure of the Spark dataframe that will return from applying the analyze player function. The next step in the code block lists an annotation that defines this function as a grouped map operation, which means that it works on dataframes rather than scalar values. As before, we’ll use the leastsq function to fit the shots and hits attributes. After calculating the coefficients for this curve fitting, we create a new Pandas dataframe with the player id, and regression coefficients. The display command at the end of this code block will force the Pandas UDF to execute, which will create a partition for each of the players in the data set, apply the least squares operation, and merge the results back together into a large Spark dataframe.
# Load necessary librariesfrom pyspark.sql.functions import pandas_udf, PandasUDFTypefrom pyspark.sql.types import *import pandas as pd# Create the schema for the resulting data frameschema = StructType([StructField('ID', LongType(), True),                     StructField('p0', DoubleType(), True),                     StructField('p1', DoubleType(), True)])# Define the UDF, input and outputs are Pandas DFs@pandas_udf(schema, PandasUDFType.GROUPED_MAP)def analize_player(sample_pd):        # return empty params in not enough data    if (len(sample_pd.shots) <= 1):        return pd.DataFrame({'ID': [sample_pd.player_id[0]],                                    'p0': [ 0 ], 'p1': [ 0 ]})         # Perform curve fitting         result = leastsq(fit, [1, 0], args=(sample_pd.shots,                                   sample_pd.hits))        # Return the parameters as a Pandas DF     return pd.DataFrame({'ID': [sample_pd.player_id[0]],                        'p0': [result[0][0]], 'p1': [result[0][1]]})    # perform the UDF and show the results player_df = stats_df.groupby('player_id').apply(analyze_player)display(player_df)
The key capability that Pandas UDFs provide is that they enable Python libraries to be used in a distributed environment, as long as you have a good way of partitioning your data. This means that libraries such as Featuretools, which were not initially designed to work in a distributed environment, can be scaled up to a large cluster. The result of applying the Pandas UDF from above on the stats data set is shown in Figure 6.8. This feature enables a mostly seamless translation between different dataframe formats.
To further demonstrate the value of Pandas UDFs, we’ll apply them to distributing a feature generation pipeline and a deep learning pipeline. However, there are some issues when using Pandas UDFs in workflows, because they can make debugging more of a challenge and sometimes fail due to data type mismatches between Spark and Pandas.
While PySpark provides a familiar environment for Python programmers, it’s good to follow a few best practices to make sure you are using Spark efficiently. Here are a set of recommendations I’ve compiled based on my experience porting a few projections from Python to PySpark:
Avoid dictionaries: Using Python data types such as dictionaries means that the code might not be executable in a distributed mode. Instead of using keys to index values in a dictionary, consider adding another column to a dataframe that can be used as a filter. This recommendation applies to other Python types including lists that are not distributable in PySpark.
Limit Pandas usage: Calling toPandas will cause all data to be loaded into memory on the driver node, and prevents operations from being performed in a distributed mode. It’s fine to use this function when data has already been aggregated and you want to make use of familiar Python plotting tools, but it should not be used for large dataframes.
Avoid loops: Instead of using for loops, it’s often possible to use functional approaches such as group by and apply to achieve the same result. Using this pattern means that code can be parallelized by supported execution environments. I’ve noticed that focusing on using this pattern in Python has also resulted in cleaning code that is easier to translate to PySpark.
Minimize eager operations: In order for your pipeline to be as scalable as possible, it’s good to avoid eager operations that pull full dataframes into memory. For example, reading in CSVs is an eager operation, and my work around is to stage the dataframe to S3 as Parquet before using it in later pipeline steps.
Use SQL: There are libraries that provide SQL operations against dataframes in both Python and PySpark. If you’re working with someone else’s Python code, it can be tricky to decipher what some of the Pandas operations are achieving. If you plan on porting your code from Python to PySpark, then using a SQL library for Pandas can make this translation easier.
By following these best practices when writing PySpark code, I’ve been able to improve both my Python and PySpark data science workflows.
Now that we’ve covered loading and transforming data with PySpark, we can now use the machine learning libraries in PySpark to build a predictive model. The core library for building predictive models in PySpark is called MLlib. This library provides a suite of supervised and unsupervised algorithms. While this library does not have complete coverage of all of the algorithms in sklearn, it provides functionality for the majority of the types of operations needed for data science workflows. In this chapter, we’ll show how to apply MLlib to a classification problem and save the outputs from the model application to a data lake.
games_df = spark.read.csv("s3://dsp-ch6/csv/games-expand.csv",                                 header=True, inferSchema = True)games_df.createOrReplaceTempView("games_df")games_df = spark.sql("""  select *, row_number() over (order by rand()) as user_id    ,case when rand() > 0.7 then 1 else 0 end as test  from games_df""")
The first step in the pipeline is loading the data set that we want to use for model training. The snippet above shows how to load the games data set, and append two additional attributes to the loaded dataframe using Spark SQL. The result of running this query is that about 30% of users will be assigned a test label which we’ll use for model application, and each record is assigned a unique user ID which we’ll use when saving the model predictions.
The next step is splitting up the data set into train and test dataframes. For this pipeline, we’ll use the test dataframe as the data set for model application, where we predict user behavior. An example of splitting up the dataframes using the test column is shown in the snippet below. This should result in roughly 16.1k training users and 6.8k test users.
trainDF = games_df.filter("test == 0")testDF = games_df.filter("test == 1")print("Train " + str(trainDF.count()))print("Test " + str(testDF.count()))
MLlib requires that the input data is formatted using vector data types in Spark. To transform our dataframe into this format, we can use the VectorAssembler class to combine a range of columns into a single vector column. The code snippet below shows how to use this class to merge the first 10 columns in the dataframe into a new vector column called features. After applying this command to the training dataframe using the transform function, we use the select function to only retrieve the values we need from the dataframe for model training and application. For the training dataframe, we only need the label and features, and with the test dataframe we also select the user ID.
from pyspark.ml.feature import VectorAssembler# create a vector representationassembler = VectorAssembler(        inputCols= trainDF.schema.names[0:10],        outputCol="features" )trainVec = assembler.transform(trainDF).select('label', 'features')testVec = assembler.transform(testDF).select(                                    'label', 'features', 'user_id')display(testVec)
The display command shows the result of transforming our test dataset into vector types usable by MLlib. The output dataframe is visualized in Figure 6.9.
Now that we have prepared our training and test data sets, we can use the logistic regression algorithm provided by MLlib to fit the training dataframe. We first create a logistic regression object and define the columns to use as labels and features. Next, we use the fit function to train the model on the training data set. In the last step in the snippet below, we use the transform function to apply the model to our test data set.
from pyspark.ml.classification import LogisticRegression# specify the columns for the modellr = LogisticRegression(featuresCol='features', labelCol='label')# fit on training datamodel = lr.fit(trainVec)# predict on test data predDF = model.transform(testVec)
The resulting dataframe now has a probability column, as shown in Figure 6.10. This column is a 2-element array with the probabilities for class 0 and 1. To test the accuracy of the logistic regression model on the test data set, we can use the binary classification evaluator in Mlib to calculate the ROC metric, as shown in the snippet below. For my run of the model, the ROC metric has a value of 0.761.
from pyspark.ml.evaluation import BinaryClassificationEvaluatorroc = BinaryClassificationEvaluator().evaluate(predDF)print(roc)
In a production pipeline, there will not be labels for the users that need predictions, meaning that you’ll need to perform cross validation to select the best model for making predictions. An example of this approach is covered in Section 6.7. In this case, we are using a single data set to keep code examples short, but a similar pipeline can be used in a production workflows.
Now that we have the model predictions for our test users, we need to retrieve the predicted label in order to create a dataframe to persist to a data lake. Since the probability column created by MLlib is an array, we’ll need to define a UDF that retrieves the second element as our propensity column, as shown in the snippet below.
from pyspark.sql.functions import udffrom pyspark.sql.types import FloatType# split out the array into a column secondElement = udf(lambda v:float(v[1]),FloatType())predDF = predDF.select("*",                 secondElement("probability").alias("propensity"))display(predDF) 
After running this code block, the dataframe will have an additional column called propensity as shown in Figure 6.10. The final step in this batch prediction pipeline is to save the results to S3. We’ll use the select function to retrieve the relevant columns from the predictions dataframe, and then use the write function on the dataframe to persist the results as Parquet on S3.
# save results to S3results_df = predDF.select("user_id", "propensity")results_path = "s3a://dsp-ch6/game-predictions/"results_df.write.mode('overwrite').parquet(results_path)
We now have all of the building blocks needed to create a PySpark pipeline that can fetch data from a storage layer, train a predictive model, and write the results to persistent storage. We’ll cover how to schedule this type of pipeline in Section 6.8.
When developing models, it’s useful to inspect the output to see if the distribution of model predictions matches expectations. We can use Spark SQL to perform an aggregation on the model outputs and then use the display command to perform this process directly in Databricks, as shown in the snippet below. The result of performing these steps on the model predictions is shown in Figure 6.11.
# plot the predictions predDF .createOrReplaceTempView("predDF ")plotDF = spark.sql("""  select cast(propensity*100 as int)/100 as propensity,          label, sum(1) as users  from predDF   group by 1, 2  order by 1, 2  """)# table outputdisplay(plotDF) 
MLlib can be applied to a wide variety of problems using a large suite of algorithms. While we explored logistic regression in this section, the libraries provides a number of different classification approaches, and there are other types of operations supported including regression and clustering.
While MLlib provides scalable implementations for classic machine learning algorithms, it does not natively support deep learning libraries such as Tensorflow and PyTorch. There are libraries that parallelize the training of deep learning models on Spark, but the data set needs to be able to fit in memory on each worker node, and these approaches are best used for distributed hyperparameter tuning on medium-sized data sets.
For the model application stage, where we already have a deep learning model trained but need to apply the resulting model to a large user base, we can use Pandas UDFs. With Pandas UDFs, we can partition and distribute our data set, run the resulting dataframes against a Keras model, and then compile the results back into a single large Spark dataframe. This section will show how we can take the Keras model that we built in Section 1.6.3, and scale it to larger data sets using PySpark and Pandas UDFs. However, we still have the requirement that the data used for training the model can fit into memory on the driver node.
We’ll use the same data sets from the prior section, where we split the games data set into training and test sets of users. This is a relatively small data set, so we can use the toPandas operation to load the dataframe onto the driver node, as shown in the snippet below. The result is a dataframe and list that we can provide as input to train a Keras deep learning model.
# build model on the driver node train_pd = trainDF.toPandas()x_train = train_pd.iloc[:,0:10]y_train = train_pd['label']
When using PyPI to install TensorFlow on the Spark cluster, the installed version of the library should be 2.0 or greater. This differs from Version 1 of TensorFlow, which we used in prior chapters. The main impact in terms of the code snippets is that Tensorflow 2 now has a built-in AUC function that no longer requires the workflow we previously applied.
We’ll use the same approach as before to train a Keras model. The code snippet below shows how to set up a network with an input layer, dropout later, single hidden layer, and an output layer, optimized with rmsprop and a cross entropy loss. In the model application phase, we’ll reuse the model object in a Pandas UDFs to distribute the workload.
import tensorflow as tfimport kerasfrom keras import models, layersmodel = models.Sequential()model.add(layers.Dense(64, activation='relu', input_shape=(10,)))model.add(layers.Dropout(0.1))model.add(layers.Dense(64, activation='relu'))model.add(layers.Dense(1, activation='sigmoid'))model.compile(optimizer='rmsprop', loss='binary_crossentropy')history = model.fit(x_train, y_train, epochs=100, batch_size=100,                  validation_split = .2, verbose=0)
To test for overfitting, we can plot the results of the training and validation data sets, as shown in Figure 6.12. The snippet below shows how to use matplotlib to display the losses over time for these data sets. While the training loss continued to decrease over additional epochs, the validation loss stopped improving after 20 epochs, but did not noticeably increase over time.
import matplotlib.pyplot as pltloss = history.history['loss']val_loss = history.history['val_loss']epochs = range(1, len(loss) + 1)fig = plt.figure(figsize=(10,6) )plt.plot(epochs, loss, 'bo', label='Training Loss')plt.plot(epochs, val_loss, 'b', label='Validation Loss')plt.legend()plt.show()display(fig)
Now that we have a trained deep learning model, we can use PySpark to apply it in a scalable pipeline. The first step is determining how to partition the set of users that need to be scored. For this data set, we can assign the user base into 100 different buckets, as shown in the snippet below. This randomly assigns each user into 1 of 100 buckets, which means that after applying the group by step, each dataframe that gets translated to Pandas will be roughly 1% the size of the original dataframe. If you have a large data set, you may need to use thousands of buckets to distribute the data set, and maybe more.
# set up partitioning for the train data frametestDF.createOrReplaceTempView("testDF ")partitionedDF = spark.sql("""  select *, cast(rand()*100 as int) as partition_id  from testDF """)
The next step is to define the Pandas UDF that will apply the Keras model. We’ll define an output schema of a user ID and propensity score, as shown below. The UDF uses the predict function on the model object we previously trained to create a prediction column on the passed in dataframe. The return command selects the two relevant columns that we defined for the schema object. The group by command partitions the data set using our bucketing approach, and the apply command performs the Keras model application across the cluster of worker nodes. The result is a Spark dataframe visualized with the display command, as shown in Figure 6.13.
from pyspark.sql.functions import pandas_udf, PandasUDFTypefrom pyspark.sql.types import *schema = StructType([StructField('user_id', LongType(), True),                     StructField('propensity', DoubleType(),True)])@pandas_udf(schema, PandasUDFType.GROUPED_MAP)def apply_keras(pd):    pd['propensity'] = model.predict(pd.iloc[:,0:10])    return pd[['user_id', 'propensity']]results_df=partitionedDF.groupby('partition_id').apply(apply_keras)display(results_df)
One thing to note is that there are limitations on the types of objects that you can reference in a Pandas UDFs. In this example, we referenced the model object, which was created on the driver node when training the model. When variables in PySpark are transferred from the driver node to workers nodes for distributed operations, a copy of the variable is made, because synchronizing variables across a cluster would be inefficient. This means that any changes made to a variable within a Pandas UDF will not apply to the original object. It’s also why data types such as Python lists and dictionaries should be avoided when using UDFs. Functions work in a similar way, in Section 6.3.4 we used the fit function in a Pandas UDF where the function was initially defined on the driver node. Spark also provides broadcast variables for sharing variables in a cluster, but ideally distributed code segments should avoid sharing state through variables if possible.
Feature engineering is a key step in a data science workflow, and sometimes it is necessary to use Python libraries to implement this functionality. For example, the AutoModel system at Zynga uses the Featuretools library to generate hundreds of features from raw tracking events, which are then used as input to classification models. To scale up the automated feature engineering approach that we first explored in Section 1.7, we can use Pandas UDFs to distribute the feature application process. Like the prior section, we need to sample data when determining the transformation to perform, but when applying the transformation we can scale to massive data sets.
For this section, we’ll use the game plays data set from the NHL Kaggle example, which includes detailed play-by-play descriptions of the events that occurred during each match. Our goal is to transform the deep and narrow dataframe into a shallow and wide dataframe that summarizes each game as a single record with hundreds of columns. An example of loading this data in PySpark and selecting the relevant columns is shown in the snippet below. Before calling toPandas, we use the filter function to sample 0.3% of the records, and then cast the result to a Pandas frame, which has a shape of 10,717 rows and 16 columns.
plays_df = spark.read.csv("s3://dsp-ch6/csv/game_plays.csv",             header=True, inferSchema = True).drop(            'secondaryType', 'periodType', 'dateTime', 'rink_side')plays_pd = plays_df.filter("rand() < 0.003").toPandas()plays_pd.shape
We’ll use the same two-step process covered in Section 1.7 where we first one-hot encode the categorical features in the dataframe, and then apply deep feature synthesis to the data set. The code snippet below shows how to perform the encoding process using the Featuretools library. The output is a transformation of the initial dataframe that now has 20 dummy variables instead of the event and description variables.
import featuretools as ftfrom featuretools import Featurees = ft.EntitySet(id="plays")es = es.entity_from_dataframe(entity_id="plays",dataframe=plays_pd,                    index="play_id", variable_types = {                    "event": ft.variable_types.Categorical,                    "description": ft.variable_types.Categorical })f1 = Feature(es["plays"]["event"])f2 = Feature(es["plays"]["description"])encoded, defs = ft.encode_features(plays_pd, [f1, f2], top_n=10)encoded.reset_index(inplace=True)
The next step is using the dfs function to perform deep feature synthesis on our encoded dataframe. The input dataframe will have a record per play, while the output dataframe will have a single record per game after collapsing the detailed events into a wide column representation using a variety of different aggregations.
es = ft.EntitySet(id="plays")es = es.entity_from_dataframe(entity_id="plays",                                 dataframe=encoded, index="play_id")es = es.normalize_entity(base_entity_id="plays",                             new_entity_id="games", index="game_id")features, transform=ft.dfs(entityset=es,                                  target_entity="games",max_depth=2)features.reset_index(inplace=True)
One of the new steps that we need to perform versus the prior approach, is that we need to determine what the schema will be for the generated features, since this is needed as an input to the Pandas UDF annotation. To figure out what the generated schema is for the generated dataframe, we can create a Spark dataframe and then retrieve the schema from the Spark dataframe. Before converting the Pandas dataframe, we need to modify the column names in the generated dataframe to remove special characters, as shown in the snippet below. The resulting Spark schema for the feature application step is displayed in Figure 6.14.
features.columns = features.columns.str.replace("[(). =]", "")schema = sqlContext.createDataFrame(features).schemafeatures.columns
We now have the required schema for defining a Pandas UDF. Unlike the past UDFs we defined, the schema may change between different runs based on the feature transformation aggregations selected by Featuretools. In these steps, we also created a defs object that defines the feature transformations to use for encoding and a transform object that defines the transformations to perform deep feature synthesis. Like the model object in the past section, copies of these objects will be passed to the Pandas UDF executing on worker nodes.
To enable our approach to scale across a cluster of worker nodes, we need to define a column to use for partitioning. Like the prior section, we can bucket events into different sets of data to ensure that the UDF process can scale. One difference from before is that we need all of the plays from a specific game to be grouped into the same partition. To achieve this result, we can partition by the game_id rather than the player_id. An example of this approach in shown in the code snippet below. Additionally, we can use the hash function on the game ID to randomize the value, resulting in more balanced buckets.
# bucket IDs plays_df.createOrReplaceTempView("plays_df")plays_df = spark.sql("""  select *, abs(hash(game_id))%1000 as partition_id   from plays_df """)
We can now apply feature transformation to the full data set, using the Pandas UDF defined below. The plays dataframe is partitioned by the bucket before being passed to the generate features function. This function uses the previously generated feature transformations to ensure that the same transformation is applied across all of the worker nodes. The input Pandas dataframe is a narrow and deep representation of play data, while the returned dataframe is a shallow and wide representation of game summaries.
from pyspark.sql.functions import pandas_udf, PandasUDFType@pandas_udf(schema, PandasUDFType.GROUPED_MAP)def gen_features(plays_pd):    es = ft.EntitySet(id="plays")    es = es.entity_from_dataframe(entity_id="plays",            dataframe=plays_pd, index="play_id", variable_types = {                    "event": ft.variable_types.Categorical,                    "description": ft.variable_types.Categorical })    encoded_features = ft.calculate_feature_matrix(defs, es)        encoded_features.reset_index(inplace=True)      es = ft.EntitySet(id="plays")    es = es.entity_from_dataframe(entity_id="plays",                                 dataframe=encoded, index="play_id")    es = es.normalize_entity(base_entity_id="plays",                            new_entity_id="games", index="game_id")    generated = ft.calculate_feature_matrix(transform,es).fillna(0)        generated.reset_index(inplace=True)    generated.columns = generated.columns.str.replace("[(). =]","")    return generated   features_df = plays_df.groupby('partition_id').apply(gen_features)display(features_df)
The output of the display command is shown in Figure 6.15. We’ve now worked through feature generation and deep learning in scalable model pipelines. Now that we have a transformed data set, we can join the result with additional features, such as the label that we are looking to predict, and develop a complete model pipeline.
A common workflow for batch model pipelines is reading input data from a lake, applying a machine learning model, and then writing the results to an application database. In GCP, BigQuery serves as the data lake and Cloud Bigtable can serve as an application database. We’ll build and end-to-end pipeline with these components in the next chapter, but for now we’ll get hands on with a subset of the GCP components directly in Spark.
While there is a Spark connector for BigQuery [15], enabling large-scale PySpark pipelines to be built using BigQuery directly, there are some issues with this library that make it quite complicated to set up for our Databricks environment. For example, we would need to rebuild some of the JAR files and shade the dependencies. One alternative is to use the Python BigQuey connector that we explored in Section 5.1, but this approach is not distributed and will eagerly pull the query results to the driver node as a Pandas dataframe. For this chapter, we’ll explore a workflow where we unload query results to Cloud Storage, and then read in the data set from GCS as the initial step in the pipeline. Similarly, for model output we’ll save the results to GCS, where the output is available for pushing to Bigtable. To productize this type of workflow, Airflow could be used to chain these different actions together.
The first step we’ll perform is exporting the results of a BigQuery query to GCS, which can be performed manually using the BigQuery UI. This is possible to perform directly in Spark, but as I mentioned the setup is quite involved to configure with the current version of the connector library. We’ll use the natality data set for this pipeline, which lists attributes about a child delivery, such as birth weight.
create table dsp_demo.natality as (  select *  from `bigquery-public-data.samples.natality`  order by rand()  limit 10000 )
To create a data set, we’ll sample 10k records from the natality public data set in BigQuery. To export this result set to GCS, we need to create a table on BigQuery with the data that we want to export. The SQL for creating this data sample is shown in the snippet above. To export this data to GCS, perform the following steps:
Browse to the GCP ConsoleSearch for “BigQuery”Paste the Query from the snippet above into the query editorClick RunIn the left navigation pane, select the new table, “dsp_demo.natality”Click “Export”, and then “Export to GCS”Set the location, “/dsp_model_store/natality/avro”Use “Avro” as export formatClick “Export”
Browse to the GCP Console
Search for “BigQuery”
Paste the Query from the snippet above into the query editor
Click Run
In the left navigation pane, select the new table, “dsp_demo.natality”
Click “Export”, and then “Export to GCS”
Set the location, “/dsp_model_store/natality/avro”
Use “Avro” as export format
Click “Export”
After performing these steps, the sampled natality data will be saved to GCS in Avro format. The confirmation dialog from exporting the data set is shown in Figure 6.16. We now have the data saved to GCS in a format that works well with Spark.
We now have a data set that we can use as input to a PySpark pipeline, but we don’t yet have access to the bucket on GCS from our Spark environment. With AWS, we were able to set up programmatic access to S3 using an access and secret key. With GCP, the process is a bit more complicated because we need to move the json credentials file to the driver node of the cluster in order to read and write files on GCS. One of the challenges with using Spark is that you may not have SSH access to the driver node, which means that we’ll need to use persistent storage to move the file to the driver machine. This isn’t recommended for production environments, but instead is being shown as a proof of concept. The best practice for managing credentials in a production environment is to use IAM roles.
aws s3 cp dsdemo.json s3://dsp-ch6/secrets/dsdemo.jsonaws s3 ls  s3://dsp-ch6/secrets/
To move the json file to the driver node, we can first copy the credentials file to S3, as shown in the snippet above. Now we can switch back to Databricks and author the model pipeline. To copy the file to the driver node, we can read in the file using the sc Spark context to read the file line by line. This is different from all of our prior operations where we have read in data sets as dataframes. After reading the file, we then create a file on the driver node using the Python open and write functions. Again, this is an unusual action to perform in Spark, because you typically want to write to persistent storage rather than local storage. The result of performing these steps is that the credentials file will now be available locally on the driver node in the cluster.
creds_file = '/databricks/creds.json'creds = sc.textFile('s3://dsp-ch6/secrets/dsdemo.json')with open(creds_file, 'w') as file:    for line in creds.take(100):        file.write(line + "\n")
Now that we have the json credentials file moved to the driver local storage, we can set up the Hadoop configuration needed to access data on GCS. The code snippet below shows how to configure the project ID, file system implementation, and credentials file location. After running these commands, we now have access to read and write files on GCS.
sc._jsc.hadoopConfiguration().set("fs.gs.impl",            "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem")sc._jsc.hadoopConfiguration().set("fs.gs.project.id",                                                  "your_project_id")sc._jsc.hadoopConfiguration().set(         "mapred.bq.auth.service.account.json.keyfile", creds_file)sc._jsc.hadoopConfiguration().set(             "fs.gs.auth.service.account.json.keyfile", creds_file)
To read in the natality data set, we can use the read function with the Avro setting to fetch the data set. Since we are using the Avro format, the dataframe will be lazily loaded and the data is not retrieved until the display command is used to sample the data set, as shown in the snippet below.
natality_path = "gs://dsp_model_store/natality/avro"natality_df = spark.read.format("avro").load(natality_path)display(natality_df)
Before we can use MLlib to build a regression model, we need to perform a few transformations on the data set to select a subset of the features, cast data types, and split records into training and test groups. We’ll also use the fillna function as shown below in order to replace any null values in the dataframe with zeros. For this modeling exercise, we’ll build a regression model that predicts the birth weight of a baby using a few different features including the marriage status of the mother and parent ages. The prepared dataframe is shown in Figure 6.17.
natality_df.createOrReplaceTempView("natality_df")natality_df = spark.sql("""SELECT year, plurality, apgar_5min,        mother_age, father_age,           gestation_weeks, ever_born       ,case when mother_married = true              then 1 else 0 end as mother_married       ,weight_pounds as weight       ,case when rand() < 0.5 then 1 else 0 end as testfrom natality_df       """).fillna(0)trainDF = natality_df.filter("test == 0")testDF = natality_df.filter("test == 1")display(natality_df)
Next, we’ll translate our dataframe into the vector data types that MLlib requires as input. The process for transforming the natality data set is shown in the snippet below. After executing the transform function, we now have training and test data sets we can use as input to a regression model. The label we are building a model to predict is the weight column.
from pyspark.ml.feature import VectorAssembler# create a vector representationassembler = VectorAssembler(inputCols= trainDF.schema.names[0:8],                            outputCol="features" )trainVec = assembler.transform(trainDF).select('weight','features')testVec = assembler.transform(testDF).select('weight', 'features')
MLlib provides a set of utilities for performing cross validation and hyperparameter tuning in a model workflow. The code snippet below shows how to perform this process for a random forest regression model. Instead of calling fit directly on the model object, we wrap the model object with a cross validator object that explores different parameter settings, such as tree depth and number of trees. This workflow is similar to the grid search functions in sklearn. After searching through the parameter space, and using cross validation based on the number of folds, the random forest model is retrained on the complete training data set before being applied to make predictions on the test data set. The result is a dataframe with the actual weight and predicted birth weight.
from pyspark.ml.tuning import ParamGridBuilder from pyspark.ml.regression import RandomForestRegressorfrom pyspark.ml.tuning import CrossValidatorfrom pyspark.ml.evaluation import RegressionEvaluatorfolds = 3rf_trees = [ 50, 100  ]rf_depth = [ 4, 5 ] rf= RandomForestRegressor(featuresCol='features',labelCol='weight')paramGrid = ParamGridBuilder().addGrid(rf.numTrees, rf_trees).                              ddGrid(rf.maxDepth, rf_depth).build()crossval = CrossValidator(estimator=rf, estimatorParamMaps =                         paramGrid,  evaluator=RegressionEvaluator(                              labelCol='weight'), numFolds = folds) rfModel = crossval.fit(trainVec)   predsDF = rfModel.transform(testVec).select("weight", "prediction") 
In the final step of our GCP model pipeline, we’ll save the results to GCS, so that other applications or processes in a workflow can make use of the predictions. The code snippet below shows how to write the dataframe to GCS in Avro format. To ensure that different runs of the pipeline do not overwrite past predictions, we append a timestamp to the export path.
import timeout_path = "gs://dsp_model_store/natality/preds-{time}/".                             format(time = int(time.time()*1000))predsDF.write.mode('overwrite').format("avro").save(out_path)print(out_path)
Using GCP components with PySpark took a bit of effort to configure, but in this case we are running Spark in a different cloud provider than where we are reading and writing data. In a production environment, you’ll most likely be running Spark in the same cloud as where you are working with data sets, which means that you can leverage IAM roles for properly managing access to different services.
Once you’ve tested a batch model pipeline in a notebook environment, there are a few different ways of scheduling the pipeline to run on a regular schedule. For example, you may want a churn prediction model for a mobile game to run every morning and publish the scores to an application database. Similar to the workflow tools we covered in Chapter 5, a PySpark pipeline should have monitoring in place for any failures that may occur. There’s a few different approaches for scheduling PySpark jobs to run:
Workflow Tools: Airflow, Azkaban, and Luigi all support running spark jobs as part of a workflow.
Cloud Tools: EMR on AWS and Dataproc on GCP support scheduled Spark jobs.
Vendor Tools: Databricks supports setting up job schedules with monitoring through the web UI.
Spark Submit: If you already have a cluster provisioned, use can issue spark-submit commands using a tool such as crontab.
Vendor and cloud tools are typically easier to get up and running, because they provide options for provisioning clusters as part of the workflow. For example, with Databricks you can define the type of cluster to spin up for running a notebook on a schedule. When using a workflow tool, such as Airflow, you’ll need to add additional steps to your workflow in order to spin up and terminate clusters. Most workflow tools provide connectors to EMR for managing clusters as part of a workflow. The Spark submit option is useful when first getting started with scheduling Spark jobs, but it’s doesn’t support managing clusters as part of a workflow.
Spark jobs can run on ephemeral or persistent clusters. An ephemeral cluster is a Spark cluster that is provisioned to perform a set of tasks and then terminated, such as running a churn model pipeline. A persistent cluster is a long-running cluster than may support interactive notebooks, such as the Databricks cluster we set up at the start of this chapter. Persistent clusters are useful for development, but can be expensive if the hardware spun up for the cluster is under utilized. Some vendors support auto scaling of clusters to reduce the costs of long-running persistent clusters. Ephemeral clusters are useful, because spinning up a new cluster to perform a task enables isolation of failure across tasks, and it means that different model pipelines can use different library versions.
In addition to setting up tools for scheduling jobs and alerting on job failures, it’s useful to set up additional data and model quality checks for Spark model pipelines. For example, I’ve set up Spark jobs that perform audit tasks, such as making sure that an application database has predictions for the current day, and trigger alerts if prediction data is stale. It’s also a good practice to log metrics, such as the ROC of a cross-validated model, as part of a Spark pipeline. We’ll cover this in more detail in Chapter 11.
PySpark is a powerful tool for data scientists to build scalable analyses and model pipelines. It a highly desirable skill set for companies, because it enables data science teams to own more of the process of building and owning data products. There’s a variety of ways to set up an environment for PySpark, and in this chapter we explored a free notebook environment from one of the popular Spark vendors.
This chapter focused on batch model pipelines, where the goal is to create a set of predictions for a large number of users on a regular schedule. We explored pipelines for both AWS and GCP deployments, where the data sources and data outputs are data lakes. One of the issues with these types of pipelines is that predictions may be quite stale by the time that a prediction is used. In Chapter 9, we’ll explore streaming pipelines for PySpark, where the latency of model predictions is minimized.
PySpark is a highly expressive language for authoring model pipelines, because it supports all Python functionality, but does require some workarounds to get code to execute across a cluster of workers nodes. In the next chapter, we’ll explore Dataflow, a runtime for the Apache Beam library, which also enables large-scale distributed Python pipelines, but is more constrained in the types of operations that you can perform.
Karau, Holden, Andy Konwinski, Patrick Wendell, and Matei Zaharia. 2015. Learning Spark: Lightning-Fast Big Data Analysis. 1st ed. O’Reilly Media.
13. https://community.cloud.databricks.com/↩︎14. https://www.gamasutra.com/blogs/BenWeber/20190426/340293/↩︎15. https://github.com/spotify/spark-bigquery/↩︎
Ben Weber is a distinguished data scientist at Zynga. We are hiring!