How to create a data frame in a for loop with the variable that is iterating in loop

user_11672756 0
Viewed 4106

So I have a huge data frame which is combination of individual tables, it has an identifier column at the end which specifies the table number as shown below

+----------------------------+
| col1 col2 .... table_num   |
+----------------------------+
| x     y            1       |
| a     b            1       |
| .     .            .       |
| .     .            .       |
| q     p            2       |
+----------------------------+

(original table)

I have to split this into multiple little dataframes based on table num. The number of tables combined to create this is pretty large so it's not feasible to individually create the disjoint subset dataframes, so I was thinking if I made a for loop iterating over min to max values of table_num I could achieve this task but I can't seem to do it, any help is appreciated.

This is what I came up with 

for (x < min(table_num) to max(table_num)) {

var df(x)= spark.sql("select * from df1 where state = x")
df(x).collect()

but I don't think the declaration is right.

so essentially what I need is df's that look like this 

+-----------------------------+
| col1  col2  ...   table_num |
+-----------------------------+
| x      y             1      |
| a      b             1      |
+-----------------------------+


+------------------------------+
| col1   col2  ...   table_num |
+------------------------------+
| xx      xy             2     |
| aa      bb             2     |
+------------------------------+

+-------------------------------+
| col1    col2  ...   table_num |
+-------------------------------+
| xxy      yyy             3    |
| aaa      bbb             3    |
+-------------------------------+

... and so on ...

(how I would like the Dataframes split)

apache-spark-sql apache-spark scala
3 Answers

In Spark Arrays can be almost in data type. When made as vars you can dynamically add and remove elements from them. Below I am going to isolate the table nums into their own array, this is so I can easily iterate through them. After isolated I go through a while loop to add each table as a unique element to the DF Holder Array. To query the elements of the array use DFHolderArray(n-1) where n is the position you want to query with 0 being the first element. 

//This will go and turn the distinct row nums in a queriable (this is 100% a word) array
val tableIDArray = inputDF.selectExpr("table_num").distinct.rdd.map(x=>x.mkString.toInt).collect

//Build the iterator
var iterator = 1  

//holders for DF and transformation step
var tempDF = spark.sql("select 'foo' as bar")
var interimDF = tempDF

//This will be an array for dataframes
var DFHolderArray : Array[org.apache.spark.sql.DataFrame] = Array(tempDF) 

//loop while the you have note reached end of array
while(iterator<=tableIDArray.length) {
  //Call the table that is stored in that location of the array
  tempDF = spark.sql("select * from df1 where state = '" + tableIDArray(iterator-1) + "'")
  //Fluff
  interimDF = tempDF.withColumn("User_Name", lit("Stack_Overflow"))
  //If logic to overwrite or append the DF
  DFHolderArray = if (iterator==1) {
    Array(interimDF)
  } else {
    DFHolderArray ++ Array(interimDF)
  }
  iterator = iterator + 1
}

//To query the data
DFHolderArray(0).show(10,false)
DFHolderArray(1).show(10,false)
DFHolderArray(2).show(10,false)
//....
user_98315220

Approach is to collect all unique keys and build respective data frames. I added some functional flavor to it.

Sample dataset:

  name,year,country,id
  Bayern Munich,2014,Germany,7747
  Bayern Munich,2014,Germany,7747
  Bayern Munich,2014,Germany,7746
  Borussia Dortmund,2014,Germany,7746
  Borussia Mönchengladbach,2014,Germany,7746
  Schalke 04,2014,Germany,7746
  Schalke 04,2014,Germany,7753
  Lazio,2014,Germany,7753

Code:

  val df = spark.read.format(source = "csv")
    .option("header", true)
    .option("delimiter", ",")
    .option("inferSchema", true)
    .load("groupby.dat")

  import spark.implicits._

  //collect data for each key into a data frame
  val uniqueIds = df.select("id").distinct().map(x => x.mkString.toInt).collect()
  // List buffer to hold separate data frames
  var dataframeList: ListBuffer[org.apache.spark.sql.DataFrame] = ListBuffer()
  println(uniqueIds.toList)

  // filter data
  uniqueIds.foreach(x => {
    val tempDF = df.filter(col("id") === x)
    dataframeList += tempDF
  })

  //show individual data frames
  for (tempDF1 <- dataframeList) {
    tempDF1.show()
  }
user_97555410

One approach would be to write the DataFrame as partitioned Parquet files and read them back into a Map, as shown below:

import org.apache.spark.sql.functions._
import spark.implicits._

val df = Seq(
  ("a", "b", 1), ("c", "d", 1), ("e", "f", 1), 
  ("g", "h", 2), ("i", "j", 2)
).toDF("c1", "c2", "table_num")

val filePath = "/path/to/parquet/files"

df.write.partitionBy("table_num").parquet(filePath)

val tableNumList = df.select("table_num").distinct.map(_.getAs[Int](0)).collect
// tableNumList: Array[Int] = Array(1, 2)

val dfMap = ( for { n <- tableNumList } yield
    (n, spark.read.parquet(s"$filePath/table_num=$n").withColumn("table_num", lit(n)))
  ).toMap

To access the individual DataFrames from the Map:

dfMap(1).show
// +---+---+---------+
// | c1| c2|table_num|
// +---+---+---------+
// |  a|  b|        1|
// |  c|  d|        1|
// |  e|  f|        1|
// +---+---+---------+

dfMap(2).show
// +---+---+---------+
// | c1| c2|table_num|
// +---+---+---------+
// |  g|  h|        2|
// |  i|  j|        2|
// +---+---+---------+
user_63165080
Related