如何在火花数据框中明智地访问元素列

Question

I have a text file which contains the following data
3 5
10 20 30 40 50
0 0 0 2 5
5 10 10 10 10
Question:
first line of the file will give us number of rows and number of columns of data
print the sum of column if every element of the column is not a prime otherwise print zero
Output:
0
30
40
0
0
Explanation: (0 because column 10 0 5 has prime number 5)
(30 because column 20 0 10 has no prime number so print 20+0+10=30) like wise apply for all columns.
suggest us the method to access the dataframe in column wise manner

Answer 1

General idea: Just zip every value with an index, create a pairRDD then apply a reduceByKey (the key here is the index) verifying at each step the number is a prime number.

val rows = spark.sparkContext.parallelize(
    Seq(
      Array(10,20,30,40,50),
      Array(0,0,0,2,5),
      Array(5,10,10,10,10)
    )
  )

def isPrime(i: Int): Boolean = i>=2 && ! ((2 until i-1) exists (i % _ == 0))

val result = rows.flatMap{arr => arr.map(Option(_)).zipWithIndex.map(_.swap)}
  .reduceByKey{
    case (None, _) | (_, None) => None
    case (Some(a),Some(b)) if isPrime(a) | isPrime(b) => None
    case (Some(a),Some(b)) => Some(a+b)
}.map{case (k,v) => k -> v.getOrElse(0)}

result.foreach(println)

Output (you'll have to collect data in order to sort by column index):

(3,0)
(0,0)
(4,0)
(2,40)
(1,30)