jt64gzCc

1# Introduction to DataFrames - Scala
2
3This gist demonstrates a number of common Spark DataFrame functions using Scala.
4
5## Create DataFrames
6
7```scala
8// Create the case classes for our domain
9case class Department(id: String, name: String)
10case class Employee(firstName: String, lastName: String, email: String, salary: Int)
11case class DepartmentWithEmployees(department: Department, employees: Seq[Employee])
12
13// Create the Departments
14val department1 = new Department("123456", "Computer Science")
15val department2 = new Department("789012", "Mechanical Engineering")
16val department3 = new Department("345678", "Theater and Drama")
17val department4 = new Department("901234", "Indoor Recreation")
18
19// Create the Employees
20val employee1 = new Employee("michael", "armbrust", "no-reply@berkeley.edu", 100000)
21val employee2 = new Employee("xiangrui", "meng", "no-reply@stanford.edu", 120000)
22val employee3 = new Employee("matei", null, "no-reply@waterloo.edu", 140000)
23val employee4 = new Employee(null, "wendell", "no-reply@princeton.edu", 160000)
24
25// Create the DepartmentWithEmployees instances from Departments and Employees
26val departmentWithEmployees1 = new DepartmentWithEmployees(department1, Seq(employee1, employee2))
27val departmentWithEmployees2 = new DepartmentWithEmployees(department2, Seq(employee3, employee4))
28val departmentWithEmployees3 = new DepartmentWithEmployees(department3, Seq(employee1, employee4))
29val departmentWithEmployees4 = new DepartmentWithEmployees(department4, Seq(employee2, employee3))
30```
31
32# Create DataFrames from a List of the Case Classe
33
34# Create the first DataFrame
35```scala
36val departmentsWithEmployeesSeq1 = Seq(departmentWithEmployees1, departmentWithEmployees2)
37val df1 = departmentsWithEmployeesSeq1.toDF()
38display(df1)
39```
40# Create the second DataFrame
41```scala
42val departmentsWithEmployeesSeq2 = Seq(departmentWithEmployees3, departmentWithEmployees4)
43val df2 = departmentsWithEmployeesSeq2.toDF()
44display(df2)
45```
46
47# Working with DataFrames
48
49## Union 2 DataFrames
50```scala
51val unionDF = df1.unionAll(df2)
52display(unionDF)
53```
54## Write the Unioned DataFrame to a Parquet file
55```scala
56// Remove the file if it exists
57dbutils.fs.rm("/tmp/example.parquet", true)
58unionDF.write.parquet("/tmp/example.parquet")
59```
60## Read a DataFrame from the Parquet file
61```scala
62val parquetDF = sqlContext.read.parquet("/tmp/example.parquet")
63```
64```scala
65val explodeDF = parquetDF.explode($"employees") {
66  case Row(employee: Seq[Row]) => employee.map{ employee =>
67    val firstName = employee(0).asInstanceOf[String]
68    val lastName = employee(1).asInstanceOf[String]
69    val email = employee(2).asInstanceOf[String]
70    val salary = employee(3).asInstanceOf[Int]
71    Employee(firstName, lastName, email, salary)
72  }
73}.cache()
74display(explodeDF)
75```
76```scala
77explodeDF
78```
79
80## Use `filter()` to return only the rows that match the given predicate
81```scala
82val filterDF = explodeDF
83  .filter($"firstName" === "xiangrui" || $"firstName" === "michael")
84  .sort($"lastName".asc)
85display(filterDF)
86```
87
88## The `where()` clause is equivalent to `filter()`
89```scala
90val whereDF = explodeDF.where(($"firstName" === "xiangrui") || ($"firstName" === "michael")).sort($"lastName".asc)
91display(whereDF)
92```
93## Replace `null` values with `–` using DataFrame Na functions
94
95```scala
96val naFunctions = explodeDF.na
97val nonNullDF = naFunctions.fill("--")
98display(nonNullDF)
99```
100
101## Retrieve only rows with missing firstName or lastName
102
103```scala
104val filterNonNullDF = nonNullDF.filter($"firstName" === "" || $"lastName" === "").sort($"email".asc)
105display(filterNonNullDF)
106```
107
108## Example aggregations using `agg()` and `countDistinct()`
109```scala
110import org.apache.spark.sql.functions._
111```
112```scala
113// Find the distinct (firstName, lastName) combinations
114val countDistinctDF = nonNullDF.select($"firstName", $"lastName")
115  .groupBy($"firstName", $"lastName")
116  .agg(countDistinct($"firstName") as "distinct_first_names")
117display(countDistinctDF)
118```
119
120## Compare the DataFrame and SQL Query Physical Plans
121
122```scala
123countDistinctDF.explain()
124```
125```scala
126// register the DataFrame as a temp table so that we can query it using SQL
127nonNullDF.registerTempTable("databricks_df_example")
128
129// Perform the same query as the DataFrame above and return ``explain``
130sqlContext.sql("""
131SELECT firstName, lastName, count(distinct firstName) as distinct_first_names
132FROM databricks_df_example
133GROUP BY firstName, lastName
134""").explain
135```
136```scala
137// Sum up all the salaries
138val salarySumDF = nonNullDF.agg("salary" -> "sum")
139display(salarySumDF)
140```
141## Print the summary statistics for the salaries
142```scala
143nonNullDF.describe("salary").show()
144```
145
146# Flattening
147
148If your data has several levels of nesting, here is a helper function to flatten your DataFrame to make it easier to work with
149
150```scala
151val veryNestedDF = Seq(("1", (2, (3, 4)))).toDF()
152```
153```scala
154import org.apache.spark.sql._
155import org.apache.spark.sql.functions._
156import org.apache.spark.sql.types._
157
158implicit class DataFrameFlattener(df: DataFrame) {
159  def flattenSchema: DataFrame = {
160    df.select(flatten(Nil, df.schema): _*)
161  }
162
163  protected def flatten(path: Seq[String], schema: DataType): Seq[Column] = schema match {
164    case s: StructType => s.fields.flatMap(f => flatten(path :+ f.name, f.dataType))
165    case other => col(path.map(n => s"`$n`").mkString(".")).as(path.mkString(".")) :: Nil
166  }
167}
168```
169```scala
170display(veryNestedDF)
171```
172```scala
173display(veryNestedDF.flattenSchema)
174```
175## Cleanup: Remove the parquet file
176```scala
177dbutils.fs.rm("/tmp/example.parquet", true)
178```
179
180# DataFrame FAQs
181
182This FAQ contains common use cases and example usage using the available APIs
183
184## Q: How can I get better performance with DataFrame UDFs? 
185## A: If the functionality exists in the available built-in functions, using these will perform better. Example usage below.
186
187We use the built-in functions and the withColumn() API to add new columns. We could have also used withColumnRenamed() to replace an existing column after the transformation. Note: Import the libraries in the first cell
188
189```scala
190import org.apache.spark.sql.functions._
191import org.apache.spark.sql.types._
192import org.apache.spark.sql._
193import org.apache.hadoop.io.LongWritable
194import org.apache.hadoop.io.Text
195import org.apache.hadoop.conf.Configuration
196import org.apache.hadoop.mapreduce.lib.input.TextInputFormat
197
198// Build an example DataFrame dataset to work with.
199dbutils.fs.rm("/tmp/dataframe_sample.csv", true)
200dbutils.fs.put("/tmp/dataframe_sample.csv", """
201id|end_date|start_date|location
2021|2015-10-14 00:00:00|2015-09-14 00:00:00|CA-SF
2032|2015-10-15 01:00:20|2015-08-14 00:00:00|CA-SD
2043|2015-10-16 02:30:00|2015-01-14 00:00:00|NY-NY
2054|2015-10-17 03:00:20|2015-02-14 00:00:00|NY-NY
2065|2015-10-18 04:30:00|2014-04-14 00:00:00|CA-LA
207""", true)
208
209val conf = new Configuration
210conf.set("textinputformat.record.delimiter", "\n")
211val rdd = sc.newAPIHadoopFile("/tmp/dataframe_sample.csv", classOf[TextInputFormat], classOf[LongWritable], classOf[Text], conf).map(_._2.toString).filter(_.nonEmpty)
212
213val header = rdd.first()
214// Parse the header line
215val rdd_noheader = rdd.filter(x => !x.contains("id"))
216// Convert the RDD[String] to an RDD[Rows]. Create an array using the delimiter and use Row.fromSeq()
217val row_rdd = rdd_noheader.map(x => x.split('|')).map(x => Row.fromSeq(x))
218
219val df_schema =
220  StructType(
221    header.split('|').map(fieldName => StructField(fieldName, StringType, true)))
222
223var df = sqlContext.createDataFrame(row_rdd, df_schema)
224df.printSchema
225```
226
227```scala
228// Instead of registering a UDF, call the builtin functions to perform operations on the columns.
229// This will provide a performance improvement as the builtins compile and run in the platform's JVM.
230
231// Convert to a Date type
232val timestamp2datetype: (Column) => Column = (x) => { to_date(x) }
233df = df.withColumn("date", timestamp2datetype(col("end_date")))
234
235// Parse out the date only
236val timestamp2date: (Column) => Column = (x) => { regexp_replace(x," (\\d+)[:](\\d+)[:](\\d+).*$", "") }
237df = df.withColumn("date_only", timestamp2date(col("end_date")))
238
239// Split a string and index a field
240val parse_city: (Column) => Column = (x) => { split(x, "-")(1) }
241df = df.withColumn("city", parse_city(col("location")))
242
243// Perform a date diff function
244val dateDiff: (Column, Column) => Column = (x, y) => { datediff(to_date(y), to_date(x)) }
245df = df.withColumn("date_diff", dateDiff(col("start_date"), col("end_date")))
246```
247
248```scala
249df.registerTempTable("sample_df")
250display(sql("select * from sample_df"))
251```
252
253## Q: I want to convert the DataFrame back to json strings to send back to Kafka
254## A: There is an underlying toJSON() function that returns an RDD of json strings using the column names and schema to produce the json records
255
256```scala
257val rdd_json = df.toJSON
258rdd_json.take(2).foreach(println)
259```
260
261## Q: My UDF takes a parameter including the column to operate on. How do I pass this parameter?
262## A: There is a function available called lit() that creates a static column
263
264```scala
265val add_n = udf((x: Integer, y: Integer) => x + y)
266
267// We register a UDF that adds a column to the DataFrame, and we cast the id column to an Integer type.
268df = df.withColumn("id_offset", add_n(lit(1000), col("id").cast("int")))
269display(df)
270```
271
272```scala
273val last_n_days = udf((x: Integer, y: Integer) => {
274  if (x < y) true else false
275})
276
277//last_n_days = udf(lambda x, y: True if x < y else False, BooleanType())
278
279val df_filtered = df.filter(last_n_days(col("date_diff"), lit(90)))
280display(df_filtered)
281```
282
283## Q: I have a table in the hive metastore and I’d like to access to table as a DataFrame. What’s the best way to define this? ## A: There’s multiple ways to define a DataFrame from a registered table. Syntax show below. Call table(tableName) or select and filter specific columns using an SQL query
284
285```scala
286// Both return DataFrame types
287val df_1 = table("sample_df")
288val df_2 = sqlContext.sql("select * from sample_df")
289```
290## Q: I’d like to clear all the cached tables on the current cluster. A: There’s an API available to do this at the global or per table level
291
292```scala
293sqlContext.clearCache()
294sqlContext.cacheTable("sample_df")
295sqlContext.uncacheTable("sample_df")
296```
297
298## Q: I’d like to compute aggregates on columns. What’s the best way to do this?
299## A: There’s a new API available named agg(*exprs) that takes a
300    list of column names and expressions for the type of aggregation you’d like to compute. You can leverage the built-in functions mentioned above as part of the expressions for each column
301    
302```scala
303// Provide the min, count, and avg and groupBy the location column. Diplay the results
304var agg_df = df.groupBy("location").agg(min("id"), count("id"), avg("date_diff"))
305display(agg_df)
306```
307
308## Q: I’d like to write out the DataFrames to Parquet, but would like to partition on a particular column. 
309## A: You can use the following APIs to accomplish this. Ensure the code does not create a large number of partitioned columns with the datasets otherwise the overhead of the metadata can cause significant slow downs. If there is a SQL table back by this directory, users will need to call refresh table _tableName_ to update the metadata prior to the query
310
311```scala
312df = df.withColumn("end_month", month(col("end_date")))
313df = df.withColumn("end_year", year(col("end_date")))
314dbutils.fs.rm("/tmp/sample_table", true)
315df.write.partitionBy("end_year", "end_month").parquet("/tmp/sample_table")
316display(dbutils.fs.ls("/tmp/sample_table"))
317```
318
319## Q: How do I properly handle cases where I want to filter out NULL data?
320## A: You can use filter() and provide similar syntax as you would with a SQL query
321
322```scala
323val null_item_schema = StructType(Array(StructField("col1", StringType, true),
324                               StructField("col2", IntegerType, true)))
325
326val null_dataset = sc.parallelize(Array(("test", 1 ), (null, 2))).map(x => Row.fromTuple(x))
327val null_df = sqlContext.createDataFrame(null_dataset, null_item_schema)
328display(null_df.filter("col1 IS NOT NULL"))
329```
330
331## Q: How do I infer the schema using the spark-csv or spark-avro libraries? 
332## A: Documented on the GitHub projects spark-csv, there is an inferSchema option flag. Providing a header would allow you to name the columns appropriately
333
334```scala
335val adult_df = sqlContext.read.
336    format("com.databricks.spark.csv").
337    option("header", "false").
338    option("inferSchema", "true").load("dbfs:/databricks-datasets/adult/adult.data")
339adult_df.printSchema()
340```