Scala Extensions and Partial Functions

import org.apache.flink.streaming.api.scala.extensions._

class OnDataStream[T] {
  def mapWith[R: TypeInformation](fun: PartialFunction[T, R]): DataStream[R]
  def flatMapWith[R: TypeInformation](fun: PartialFunction[T, TraversableOnce[R]]): DataStream[R]
  def filterWith(fun: PartialFunction[T, Boolean]): DataStream[T]
  def keyingBy[K: TypeInformation](fun: PartialFunction[T, K]): KeyedStream[T, K]
}

class OnKeyedStream[T, K] {
  def mapWith[R: TypeInformation](fun: PartialFunction[T, R]): DataStream[R]
  def flatMapWith[R: TypeInformation](fun: PartialFunction[T, TraversableOnce[R]]): DataStream[R]
  def filterWith(fun: PartialFunction[T, Boolean]): DataStream[T]
  def reduceWith(fun: PartialFunction[(T, T), T]): DataStream[T]
}

import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.extensions._

val env = StreamExecutionEnvironment.getExecutionEnvironment
val events = env.fromElements("start", "data", "end", "error")

// Partial function mapping
val mapped = events.mapWith {
  case "start" => "BEGIN"
  case "end" => "FINISH"
  case data => s"PROCESS: $data"
}

// Partial function filtering
val filtered = events.filterWith {
  case "error" => false
  case _ => true
}