Stop using Pandas and start using Spark with Scala
By: Chloe Connor
Tags: Scala, Spark, Towards Data Science, Pandas, Data Engineering

Data transformations

Most (if not all) of the data transformations you can apply to Pandas DataFrames, are available in Spark. There are of course differences in syntax, and sometimes additional things to be aware of, some of which we’ll go through now.

In general, I’ve found Spark more consistent in notation compared with Pandas and because Scala is statically typed, you can often just do myDataset. and wait for your compiler to tell you what methods are available!

Let’s start with a simple transformation, where we just want to add a new column to our Dataset, and assign it constant value. In Pandas this looks like:

Pandas df_teams['sport'] = 'football'

There’s a small difference in Spark, besides syntax, and that’s that adding a constant value to this new field requires us to import a spark function called lit.

Spark import org.apache.spark.sql.functions.lit val newTeams = teams.withColumn("sport", lit("football"))

Note that we’ve created a new object as our original teams dataset is a val, which means it’s immutable. This is a good thing as we know that whenever we use our teams Dataset, we always get the same object.

Now let’s add a column based on a function. In Pandas this will look like:

Pandas def is_prem(league):

if league == 'Premier League':

return True

else:

return False

df_teams['premier_league'] = df_teams['league'].apply(lambda x:

is_prem(x))



To do the same in Spark, we need to serialise the function so that Spark can apply it. This is done using something called UserDefinedFunctions. We’ve also used a case match, as this is a nicer implementation in Scala than the if-else, but either will work.

We will also need to import another useful spark function, col, which is used to refer to a column.

Spark import org.apache.spark.sql.functions.col def isPrem(league: String): Boolean =

league match {

case "Premier League" => true

case _ => false

}

val isPremUDF: UserDefinedFunction =

udf[Boolean, String](isPrem)

val teamsWithLeague: DataFrame = teams.withColumn("premier_league",

isPremUDF(col("league")))

Now that we’ve added a new column that isn’t in our case class, this will convert it back to a DataFrame. So we either need to add another field to our original case class (and allow it to be nullable, using Options), or create a new case class.

An Option in Scala just means the field is nullable. If the value is null we use None , and if populated we use Some("value") . An example of an optional string:

val optionalString : Option[String] = Some("something")

To get the string from this we can call optionalString.get() , and this will just return "something" . Note that if we’re not sure whether it will be null or not, we can use optionalString.getOrElse("nothing") which will return the string "nothing" if null.

Filtering a Dataset is another common requirement, which is a good example of where Spark is more consistent than Pandas, as it follows the same pattern as other transformations, where we do dataset “dot” transformation (ie dataset.filter(...) ).

Pandas df_teams = df_teams[df_teams['goals_this_season'] > 50]

Spark val filteredTeams = teams.filter(col("goals_this_season") > 50)

We are likely to need to perform some aggregations on our dataset, which is very similar in Pandas and Spark.

Pandas df_teams.groupby(['league']).count()

Spark teams.groupBy("league").count()

For multiple aggregations, we can again do something similar to Pandas, with a map of field to aggregation. If we want to do our own aggregations we can use UserDefinedAggregations.

teams.agg(Map(

"matches_played" -> "avg",

"goals_this_season" -> "count"))

Often we also want to combine multiple Datasets, which may be with union:

Pandas pd.concat([teams, another_teams], ignore_index=True)

Spark teams.unionByName(anotherTeams)

… or with a join:

val players: Dataset[Player] = spark

.createDataset(Seq(neilMaupey, sergioAguero))

teams.join(players,

teams.col("top_goal_scorer") === players.col("player_name"),

"left"

).drop("player_name")

In this example we have also created a new Dataset, this time using a case class called Player. Note that this case class has a field injury, which can be null.

case class Player(player_name: String, goals_scored: Int,

injury: Option[String])

Notice that we’ve dropped the player_name column as this will be a duplicate of top_goal_scorer.

We may also want parts of our code to just use Scala native data structures such as Arrays, Lists, etc. To get one of our columns as an Array, we need to map to our value and call .collect() .

val teamNames: Array[String] = teams.map(team => team.name)

.collect()

Note that we’re able to use the case class’ inbuilt getters to return the name field, and this won’t compile if name is not a field in our class FootballTeam.

As an aside, we can add functions to our case classes too, and both values and functions will come up as options for autocompletion when using an IDE such as IntelliJ or vs code with Metals plugin.

To filter our Dataset based on whether it exists in this Array we need to treat it as a sequence of args, by calling _*.