Aggregation.scala

import Aggregator.AggregationSyntax
import io.getquill.context.ZioJdbc.{QConnection, QDataSource}
import io.getquill.{PostgresZioJdbcContext, Query, SnakeCase}
import zio._
import zio.blocking.Blocking
import zio.console.putStrLn

object QuillContext extends PostgresZioJdbcContext(SnakeCase)
import QuillContext._

case class Person(id: Int, name: String, age: Int, houseId: Int)

case class Grade(id: Int, score: Int, personId: Int)

case class House(id: Int, address: String)

object Main extends App {
  val zioConn: ZLayer[Blocking, Throwable, QConnection] =
    QDataSource.fromPrefix("postgresDB") >>>
      QDataSource.toConnection

  override def run(args: List[String]): URIO[zio.ZEnv, ExitCode] = {
    val people: QuillContext.Quoted[Query[((House, Person), Grade)]] = quote {
      query[House]
        .join(
          query[Person]
        )
        .on(_.id == _.houseId)
        .join(query[Grade])
        .on { case ((_, p), g) => p.id == g.personId }
    }

    val program = for {
      ppl: Seq[((House, Person), Grade)] <- QuillContext
        .run(people)
        .tap(result => putStrLn(result.mkString("\n")))

      _ <- putStrLn("")

      aggTwo: Map[House, Seq[Grade]] =
        ppl.agg[House, Grade]

      _ <- putStrLn(aggTwo.mkString("\n") + "\n")

      aggThree: Map[House, Map[Person, Seq[Grade]]] =
        ppl.agg[House, Person, Grade]

      _ <- putStrLn(aggThree.mkString("\n"))
    } yield ()

    program
      .provideCustomLayer(zioConn)
      .exitCode
  }

}

trait TupleExtractor[A, B] {
  def apply(a: A): B
}

object TupleExtractor extends TupleExtractorLowPriorityImplicits {

  implicit final class TupleExtractorOps[A](val self: A) extends AnyVal {
    def extract[B](implicit tupleExtractor: TupleExtractor[A, B]): B =
      tupleExtractor.apply(self)
  }

  implicit def abA[A, B]: TupleExtractor[(A, B), A] = (a: (A, B)) => a._1

  implicit def abB[A, B]: TupleExtractor[(A, B), B] = (a: (A, B)) => a._2

  implicit def abcA[A, B, C]: TupleExtractor[(A, B, C), A] = (a: (A, B, C)) => a._1

  implicit def abcB[A, B, C]: TupleExtractor[(A, B, C), B] = (a: (A, B, C)) => a._2

  implicit def abcC[A, B, C]: TupleExtractor[(A, B, C), C] = (a: (A, B, C)) => a._3

  def main(args: Array[String]): Unit = {
    val value = (1, "hello").extract[String]
    println(value)
  }
}

trait TupleExtractorLowPriorityImplicits {
  implicit def nestedLeft[T, A, B](implicit tupleExtractor: TupleExtractor[T, A]): TupleExtractor[(T, B), A] =
    (a: (T, B)) => tupleExtractor.apply(a._1)

  implicit def nestedRight[T, A, B](implicit tupleExtractor: TupleExtractor[T, B]): TupleExtractor[(A, T), B] =
    (a: (A, T)) => tupleExtractor.apply(a._2)
}

trait Aggregator[S, A, B] {
  def getA(s: S): A

  def getB(s: S): B

  def apply(list: Seq[S]): Map[A, Seq[B]] = {
    list.groupMap(getA)(getB)
  }
}

object Aggregator {
  implicit final class AggregationSyntax[S](val self: Seq[S]) extends AnyVal {
    def agg[A, B](implicit
        exA: TupleExtractor[S, A],
        exB: TupleExtractor[S, B]
    ): Map[A, Seq[B]] =
      self.groupMap(exA.apply)(exB.apply)

    def agg[A, B, C](implicit
        exA: TupleExtractor[S, A],
        exB: TupleExtractor[S, B],
        exC: TupleExtractor[S, C]
    ): Map[A, Map[B, Seq[C]]] =
      self
        .groupBy(exA.apply)
        .map { case (a, s) =>
          a -> s.groupMap(exB.apply)(exC.apply)
        }
  }
}

When working with SQL in Scala, particularly with Slick, you often find yourself wanting to aggregate tabular data. Imagine the following hierarchy:

House 1
  Person 1
    Grade 1
    Grade 2
  Person 2
    Grade 3
House 2
  Person 3
    Grade 4
    Grade 5    

If you join through from house down to grade, this is going to get returned to you rather awkwardly.

House1 Person1 Grade1
House1 Person1 Grade2
House1 Person2 Grade3
etc...

I just hacked together a simple little type-level method for doing the aggregation automatically 😄

Given a val results: Seq[((House, Person), Grade)]

((House(1,521 House Street),Person(1,Bobby,13,1)),Grade(1,8,1))
((House(1,521 House Street),Person(1,Bobby,13,1)),Grade(2,4,1))
((House(1,521 House Street),Person(2,Kit,30,1)),Grade(3,5,2))
((House(2,327 Street Avenue),Person(3,Jill,13,2)),Grade(4,10,3))
((House(2,327 Street Avenue),Person(3,Jill,13,2)),Grade(5,6,3))
((House(1,521 House Street),Person(1,Bobby,13,1)),Grade(6,5,1))

I can call:

val map: Map[House, List[Grade]] = results.agg[House, Grade]

House(2,327 Street Avenue) -> List(Grade(4,10,3), Grade(5,6,3))
House(1,521 House Street) -> List(Grade(1,8,1), Grade(2,4,1), Grade(3,5,2), Grade(6,5,1))

// or

val map: Map[House, Map[Person, List[Grade]]] = results.agg[House, Person, Grade]

House(1,521 House Street) -> Map(Person(1,Bobby,13,1) -> List(Grade(1,8,1), Grade(2,4,1), Grade(6,5,1)), Person(2,Kit,30,1) -> List(Grade(3,5,2)))
House(2,327 Street Avenue) -> Map(Person(3,Jill,13,2) -> List(Grade(4,10,3), Grade(5,6,3)))