Scala: How to write your own collect and map functions

By Alvin Alexander. Last updated: February 28, 2024

As a brief note today, I just wrote a Scala collect function for my free Scala video training courses, and wrote it as shown in this code:

/**
 * `collect` is like `map`, and you can think
 * of `collect` being written like this.
 */
def collect[A, B](xs: Seq[A], pf: PartialFunction[A, B]): Seq[B] =
    for
        x <- xs
        if pf.isDefinedAt(x)
    yield
        pf(x)

val divide = new PartialFunction[Int, Int] {
    def apply(x: Int) =
        42 / x
    def isDefinedAt(x: Int) =
        x != 0
}

Using that 'collect' function

You can now use that collect function like this:

val xs = List(0, 1, 2, 3)
val rez = collect(xs, divide)   // rez: List(42, 21, 14)

Note that the Scala collections types have a built-in collect method; I just wrote this to show how you can write your own. Here’s the built-in collect method on a Scala List:

val xs = List(0, 1, 2, 3)
val rez = xs.collect(divide)   // rez: List(42, 21, 14)

If you try to use 'map'

Also note that if you try to use map like that, you’ll get this error:

scala> xs.map(divide)
java.lang.ArithmeticException: / by zero
  at rs$line$1$$anon$1.apply$mcII$sp(rs$line$1:2)
  at rs$line$1$$anon$1.apply(rs$line$1:2)
  at rs$line$1$$anon$1.apply(rs$line$1:2)
  at scala.collection.immutable.List.map(List.scala:246)
  ... 32 elided

Write your own 'map' function

As a final comparison, this shows how you can write your own map and collect functions:

def collect[A, B](xs: Seq[A], pf: PartialFunction[A, B]): Seq[B] =
    for
        x <- xs
        if pf.isDefinedAt(x)
    yield
        pf(x)

def map[A, B](xs: Seq[A], f: A => B): Seq[B] =
    for
        x <- xs
    yield
        f(x)

As shown, the big differences are that (a) the second argument to collect is a partial function (PartialFunction), and (b) I call isDefinedAt in the for expression.

If you ever wanted to see how to write your own collect and map functions in Scala, I hope these examples are helpful.

scala

for loop