diff --git a/src/main/scala/com.ariskk.flink4s/WindowedStream.scala b/src/main/scala/com.ariskk.flink4s/WindowedStream.scala index 386fb87..664eb58 100644 --- a/src/main/scala/com.ariskk.flink4s/WindowedStream.scala +++ b/src/main/scala/com.ariskk.flink4s/WindowedStream.scala @@ -1,8 +1,10 @@ package com.ariskk.flink4s +import cats.Semigroup +import cats.kernel.Monoid import org.apache.flink.streaming.api.datastream.{WindowedStream => JavaWStream} import org.apache.flink.streaming.api.windowing.windows.Window -import org.apache.flink.api.common.functions.ReduceFunction +import org.apache.flink.api.common.functions.{AggregateFunction, ReduceFunction} import org.apache.flink.api.common.typeinfo.TypeInformation final case class WindowedStream[T, K, W <: Window](stream: JavaWStream[T, K, W])(implicit @@ -16,4 +18,22 @@ final case class WindowedStream[T, K, W <: Window](stream: JavaWStream[T, K, W]) DataStream(stream.reduce(reducer)) } + def aggregate[A, O](agg: A, mergeF: (A, A) => A)(aggregateF: (A, T) => A)(outputF: A => O)( + implicit + aggTypeInformation: TypeInformation[A], + typeInformation: TypeInformation[O] + ): DataStream[O] = { + val reducer = new AggregateFunction[T, A, O] { + + override def createAccumulator(): A = agg + + override def add(value: T, accumulator: A): A = aggregateF(accumulator, value) + + override def getResult(accumulator: A): O = outputF(accumulator) + + override def merge(a: A, b: A): A = mergeF(a, b) + } + DataStream(stream.aggregate[A, O](reducer, aggTypeInformation, typeInformation)) + } + } diff --git a/src/test/scala/com.ariskk.flink4s/WindowedStreamSpec.scala b/src/test/scala/com.ariskk.flink4s/WindowedStreamSpec.scala index 5c06477..97db5c1 100644 --- a/src/test/scala/com.ariskk.flink4s/WindowedStreamSpec.scala +++ b/src/test/scala/com.ariskk.flink4s/WindowedStreamSpec.scala @@ -1,11 +1,12 @@ package com.ariskk.flink4s -import scala.collection.mutable.{Buffer => MutableBuffer} +import cats.Monoid +import cats.kernel.Semigroup +import scala.collection.mutable.{Buffer => MutableBuffer} import org.scalatest.funspec.AnyFunSpec import org.scalatest.matchers.should.Matchers import org.apache.flink.streaming.api.functions.sink.SinkFunction - import com.ariskk.flink4s.TypeInfo.{intTypeInfo, stringTypeInfo} final class WindowedStreamSpec extends AnyFunSpec with Matchers { @@ -41,6 +42,53 @@ final class WindowedStreamSpec extends AnyFunSpec with Matchers { results.size should equal(4) results shouldBe List(50, 100, 100, 100) } + + it("should apply aggregation to count window with slide") { + val env = FlinkExecutor.newEnv(parallelism = 1) + val stream = env.fromCollection((1 to 200).toList.map(_ => 1)) + val results = stream + .keyBy(identity) + .countWindow(100, 50) + .aggregate[Int, Int]((agg, i) => agg + i)(identity(_)) + .runAndCollect + + results.size should equal(4) + results shouldBe List(50, 100, 100, 100) + } + + + it("should apply aggregation if aggregator and out are different types") { + val env = FlinkExecutor.newEnv(parallelism = 1) + val stream = env.fromCollection((1 to 200).toList.map(_ => 1)) + val results = stream + .keyBy(identity) + .countWindow(100, 50) + .aggregate[Int, String]((agg, i) => agg + i)(_.toString) + .runAndCollect + + results.size should equal(4) + results shouldBe List("50", "100", "100", "100") + } + + it("should apply aggregation based on Monoid") { + val env = FlinkExecutor.newEnv(parallelism = 1) + val stream = env.fromCollection((1 to 200).toList.map(_ => 1)) + implicit val semigroup = new Monoid[Int] { + + override def empty: Int = 5 + + override def combine(x: Int, y: Int): Int = x + y + } + + val results = stream + .keyBy(identity) + .countWindow(100, 50) + .aggregate[Int, Int]((agg, i) => agg + i)(identity(_)) + .runAndCollect + + results.size should equal(4) + results shouldBe List(55, 105, 105, 105) + } } }