How to match an RDD[ParentClass] with RDD[Subclass] in apache spark?

问题

I have to match an rdd with its types.

trait Fruit

case class Apple(price:Int) extends Fruit
case class Mango(price:Int) extends  Fruit

Now a dstream of type DStream[Fruit] is coming. It is either Apple or Mango.

How to perform operation based on the subclass? Something like the below (which doesn't work):

dStream.foreachRDD{rdd:RDD[Fruit] =>
     rdd match {
       case rdd: RDD[Apple] =>
         //do something

       case rdd: RDD[Mango] =>
         //do something

       case _ =>
         println(rdd.count() + "<<<< not matched anything")
     }
    }

回答1:

Since we have an RDD[Fruit], any row can be either Apple or Mango. When using foreachRDD, each RDD will contain a mix of these (and possible other) types.

To differentiate between the different types, we can use collect[U](f: PartialFunction[T, U]): RDD[U] (this is not to be confused with collect(): Array[T] that returns a list with the elements from the RDD). This function will return an RDD that contains all matching values by applying a function f (in this case, we can use a pattern match here).

Below follows a small illustrative example (adding Orange to the fruits as well).

Setup:

val ssc = new StreamingContext(spark.sparkContext, Seconds(1))
val inputData: Queue[RDD[Fruit]] = Queue()
val dStream: InputDStream[Fruit] = ssc.queueStream(inputData)

inputData += spark.sparkContext.parallelize(Seq(Apple(5), Apple(5), Mango(11)))
inputData += spark.sparkContext.parallelize(Seq(Mango(10), Orange(1), Orange(3)))

This creates a stream of RDD[Fruit] with two separate RDDs.

dStream.foreachRDD{rdd: RDD[Fruit] =>
  val mix = rdd.collect{
    case row: Apple => ("APPLE", row.price) // do any computation on apple rows
    case row: Mango => ("MANGO", row.price) // do any computation on mango rows
    //case _@row => do something with other rows (will be removed by default).
  }
  mix foreach println
}

In the above collect, we change each row slightly (removing the class) and then prints the resulting RDD. Result:

// First RDD
(MANGO,11)
(APPLE,5)
(APPLE,5)

// Second RDD
(MANGO,10)

As can be seen, the pattern match have kept and changed the rows containing Apple and Mango while removing all Orange classes.

---

Separate RDDs

If wanted, it is also possible to separate the two subclasses into their own RDDs as follows. Any computations can then be performed on these two RDDs.

val apple = rdd.collect{case row: Apple => row}
val mango = rdd.collect{case row: Mango => row}

---

Complete example code

trait Fruit
case class Apple(price:Int) extends Fruit
case class Mango(price:Int) extends  Fruit
case class Orange(price:Int) extends  Fruit

object Test {
  def main(args: Array[String]) {
    val spark = SparkSession.builder.master("local[*]").getOrCreate()

    val ssc = new StreamingContext(spark.sparkContext, Seconds(1))
    val inputData: Queue[RDD[Fruit]] = Queue()
    val inputStream: InputDStream[Fruit] = ssc.queueStream(inputData)

    inputData += spark.sparkContext.parallelize(Seq(Apple(5), Apple(5), Mango(11)))
    inputData += spark.sparkContext.parallelize(Seq(Mango(10), Orange(1), Orange(3)))

    inputStream.foreachRDD{rdd:RDD[Fruit] =>
      val mix = rdd.collect{
        case row: Apple => ("APPLE", row.price) // do any computation on apple rows
        case row: Mango => ("MANGO", row.price) // do any computation on mango rows
        //case _@row => do something with other rows (will be removed by default).
      }
      mix foreach println
    }

    ssc.start()
    ssc.awaitTermination()
  }
}

来源：https://stackoverflow.com/questions/61444973/how-to-match-an-rddparentclass-with-rddsubclass-in-apache-spark