Adding Relational Algebra style joins to Scala Traversables.

TraversableExtensions

import scala.collection.generic.CanBuildFrom
import scala.collection.TraversableLike
import scala.collection.SeqLike
import scala.collection.Seq
import Option.empty

object TraversableExtensions {

  implicit class TraversableWithJoin[A, Repr <: Traversable[A]](val xs: TraversableLike[A, Repr]) extends AnyVal {

		def ⋈[B, J](fA: A => J)(t: Traversable[B])(fB: B => J) {
			innerJoin(fA)(t)(fB)
		}

		def innerJoin[That, B, J](fA: A => J)(t: Traversable[B])(fB: B => J)(implicit bf: CanBuildFrom[Repr, (A, B), That]): That = {
			val builder = bf(xs.asInstanceOf[Repr])
			val left = xs.groupBy(fA);
			val right = t.groupBy(fB);
			val keys = left.keySet intersect right.keySet

			for (key <- keys) {
				left(key).foreach(l => right(key).foreach(r => builder += Tuple2(l, r)))
			}
			builder.result
		}

		def ⟕[B, J](fA: A => J)(t: Traversable[B])(fB: B => J) {
			leftJoin(fA)(t)(fB)
		}

		def leftJoin[That, B, J](fA: A => J)(t: Traversable[B])(fB: B => J)(implicit bf: CanBuildFrom[Repr, (A, Option[B]), That]): That = {
			val builder = bf(xs.asInstanceOf[Repr])
			val left = xs.groupBy(fA);
			val right = t.groupBy(fB);

			for (leftGroup <- left) {
				val key = leftGroup._1
				val leftElements = leftGroup._2

				if (right.contains(key)) {
					leftElements.foreach(l => right(key).foreach(r => builder += Tuple2(l, Option(r))))
				} else {
					leftElements.foreach(l => builder += Tuple2(l, empty))
				}
			}
			builder.result
		}

		def ⟖[B, J](fA: A => J)(t: Traversable[B])(fB: B => J) {
			rightJoin(fA)(t)(fB)
		}

		def rightJoin[That, B, J](fA: A => J)(t: Traversable[B])(fB: B => J)(implicit bf: CanBuildFrom[Repr, (Option[A], B), That]): That = {
			val builder = bf(xs.asInstanceOf[Repr])
			val left = xs.groupBy(fA);
			val right = t.groupBy(fB);

			for (rightGroup <- right) {
				val key = rightGroup._1
				val rightElements = rightGroup._2

				if (left.contains(key)) {
					rightElements.foreach(r => left(key).foreach(l => builder += Tuple2(Option(l), r)))
				} else {
					rightElements.foreach(r => builder += Tuple2(empty, r))
				}
			}
			builder.result
		}

		def ⟗[B, J](fA: A => J)(t: Traversable[B])(fB: B => J) {
			outerJoin(fA)(t)(fB)
		}

		def outerJoin[That, B, J](fA: A => J)(t: Traversable[B])(fB: B => J)(implicit bf: CanBuildFrom[Repr, (Option[A], Option[B]), That]): That = {
			val builder = bf(xs.asInstanceOf[Repr])
			val left = xs.groupBy(fA);
			val right = t.groupBy(fB);
			val keys = left.keySet ++ right.keySet

			for (key <- keys) {
				if (left.contains(key) && right.contains(key)) {
					left(key).foreach(l => right(key).foreach(r => builder += Tuple2(Option(l), Option(r))))
				} else if (left.contains(key)) {
					left(key).foreach(l => builder += Tuple2(Option(l), empty))
				} else {
					right(key).foreach(r => builder += Tuple2(empty, Option(r)))
				}
			}
			builder.result
		}
	}

}