twitter · ianoc · Jan 7, 2017 · Jan 7, 2017 · Jan 7, 2017 · Jan 7, 2017
diff --git a/algebird-core/src/main/scala/com/twitter/algebird/Aggregator.scala b/algebird-core/src/main/scala/com/twitter/algebird/Aggregator.scala
@@ -512,6 +512,16 @@ trait MonoidAggregator[-A, B, +C] extends Aggregator[A, B, C] { self =>
       def present(bs: (B, B2)) = (self.present(bs._1), that.present(bs._2))
     }
 
+  /**
+   * Only transform values where the function is defined, else discard
+   */
+  def collectBefore[A2](fn: PartialFunction[A2, A]): MonoidAggregator[A2, B, C] =
+    new MonoidAggregator[A2, B, C] {
+      def prepare(a: A2) = if (fn.isDefinedAt(a)) self.prepare(fn(a)) else self.monoid.zero
+      def monoid = self.monoid
+      def present(b: B) = self.present(b)
+    }
+
   /**
    * Only aggregate items that match a predicate
    */

diff --git a/algebird-test/src/test/scala/com/twitter/algebird/AggregatorLaws.scala b/algebird-test/src/test/scala/com/twitter/algebird/AggregatorLaws.scala
@@ -97,6 +97,95 @@ class AggregatorLaws extends CheckProperties {
     }
   }
 
+  property("Aggregator.count is like List.count") {
+    forAll { (in: List[Int], fn: Int => Boolean) =>
+      in.count(fn) == (Aggregator.count(fn)(in))
+    }
+  }
+  property("Aggregator.exists is like List.exists") {
+    forAll { (in: List[Int], fn: Int => Boolean) =>
+      in.exists(fn) == (Aggregator.exists(fn)(in))
+    }
+  }
+  property("Aggregator.forall is like List.forall") {
+    forAll { (in: List[Int], fn: Int => Boolean) =>
+      in.forall(fn) == (Aggregator.forall(fn)(in))
+    }
+  }
+  property("Aggregator.head is like List.head") {
+    forAll { (in: List[Int]) =>
+      in.headOption == (Aggregator.head.applyOption(in))
+    }
+  }
+  property("Aggregator.last is like List.last") {
+    forAll { (in: List[Int]) =>
+      in.lastOption == (Aggregator.last.applyOption(in))
+    }
+  }
+  property("Aggregator.maxBy is like List.maxBy") {
+    forAll { (head: Int, in: List[Int], fn: Int => Int) =>
+      val nonempty = head :: in
+      nonempty.maxBy(fn) == (Aggregator.maxBy(fn).apply(nonempty))
+    }
+  }
+  property("Aggregator.minBy is like List.minBy") {
+    forAll { (head: Int, in: List[Int], fn: Int => Int) =>
+      val nonempty = head :: in
+      nonempty.minBy(fn) == (Aggregator.minBy(fn).apply(nonempty))
+    }
+  }
+  property("Aggregator.sortedTake same as List.sorted.take") {
+    forAll { (in: List[Int], t0: Int) =>
+      val t = math.max(t0, 1)
+      val l = in.sorted.take(t)
+      val a = (Aggregator.sortedTake[Int](t).apply(in))
+      l == a
+    }
+  }
+  property("Aggregator.sortByTake same as List.sortBy(fn).take") {
+    forAll { (in: List[Int], t0: Int, fn: Int => Int) =>
+      val t = math.max(t0, 1)
+      val l = in.sortBy(fn).take(t)
+      val a = (Aggregator.sortByTake(t)(fn).apply(in))
+      // since we considered two things equivalent under fn,
+      // we have to use that here:
+      val ord = Ordering.Iterable(Ordering.by(fn))
+      ord.equiv(l, a)
+    }
+  }
+  property("Aggregator.sortByReverseTake same as List.sortBy(fn).reverse.take") {
+    forAll { (in: List[Int], t0: Int, fn: Int => Int) =>
+      val t = math.max(t0, 1)
+      val l = in.sortBy(fn).reverse.take(t)
+      val a = (Aggregator.sortByReverseTake(t)(fn).apply(in))
+      // since we considered two things equivalent under fn,
+      // we have to use that here:
+      val ord = Ordering.Iterable(Ordering.by(fn))
+      ord.equiv(l, a)
+    }
+  }
+  property("Aggregator.immutableSortedTake same as List.sorted.take") {
+    forAll { (in: List[Int], t0: Int) =>
+      val t = math.max(t0, 1)
+      val l = in.sorted.take(t)
+      val a = (Aggregator.immutableSortedTake[Int](t).apply(in))
+      l == a
+    }
+  }
+  property("Aggregator.immutableSortedReverseTake same as List.sorted.reverse.take") {
+    forAll { (in: List[Int], t0: Int) =>
+      val t = math.max(t0, 1)
+      val l = in.sorted.reverse.take(t)
+      val a = (Aggregator.immutableSortedReverseTake[Int](t).apply(in))
+      l == a
+    }
+  }
+  property("Aggregator.toList is identity on lists") {
+    forAll { (in: List[Int]) =>
+      in == Aggregator.toList(in)
+    }
+  }
+
   property("MonoidAggregator.sumBefore is correct") {
     forAll{ (in: List[List[Int]], ag: MonoidAggregator[Int, Int, Int]) =>
       val liftedAg = ag.sumBefore
@@ -126,4 +215,12 @@ class AggregatorLaws extends CheckProperties {
       ag.filterBefore(fn).apply(in) == ag.apply(in.filter(fn))
     }
   }
+
+  property("MonoidAggregator.collectBefore is like filter + compose") {
+    forAll { (in: List[Int], ag: MonoidAggregator[Int, Int, Int], fn: Int => Option[Int]) =>
+      val cp = ag.collectBefore[Int] { case x if fn(x).isDefined => fn(x).get }
+      val fp = ag.composePrepare[Int](fn(_).get).filterBefore[Int](fn(_).isDefined)
+      cp(in) == fp(in)
+    }
+  }
 }