<commit_before>package com.spotify.cloud.dataflow.util

import java.lang.{Iterable => JIterable}
import java.util.{List => JList}

import com.google.cloud.dataflow.sdk.coders.{Coder, CoderRegistry}
import com.google.cloud.dataflow.sdk.transforms.Combine.CombineFn
import com.google.cloud.dataflow.sdk.transforms.{SerializableFunction, DoFn}
import com.google.cloud.dataflow.sdk.transforms.Partition.PartitionFn
import com.google.common.collect.Lists
import com.spotify.cloud.dataflow.coders.KryoAtomicCoder
import com.twitter.algebird.{Monoid, Semigroup}

import scala.collection.JavaConverters._
import scala.reflect.ClassTag

private[dataflow] object Functions {

  private val BUFFER_SIZE = 20

  private abstract class KryoCombineFn[VI, VA, VO] extends CombineFn[VI, VA, VO] {

    override def getAccumulatorCoder(registry: CoderRegistry, inputCoder: Coder[VI]): Coder[VA] =
      new KryoAtomicCoder().asInstanceOf[Coder[VA]]

    override def getDefaultOutputCoder(registry: CoderRegistry, inputCoder: Coder[VI]): Coder[VO] =
      new KryoAtomicCoder().asInstanceOf[Coder[VO]]

  }

  def aggregateFn[T, U: ClassTag](zeroValue: U)(seqOp: (U, T) => U, combOp: (U, U) => U)
      : CombineFn[T, (U, JList[T]), U] = new KryoCombineFn[T, (U, JList[T]), U] {

    // defeat closure
    val s = ClosureCleaner(seqOp)
    val c = ClosureCleaner(combOp)

    private def fold(accumulator: (U, JList[T])): U = {
      val (a, l) = accumulator
      l.asScala.foldLeft(a)(s)
    }

    override def createAccumulator(): (U, JList[T]) = (zeroValue, Lists.newArrayList())

    override def addInput(accumulator: (U, JList[T]), input: T): (U, JList[T]) = {
      val (a, l) = accumulator
      l.add(input)
      if (l.size() >= BUFFER_SIZE) {
        (fold(accumulator), Lists.newArrayList())
      } else {
        accumulator
      }
    }

    override def extractOutput(accumulator: (U, JList[T])): U = accumulator._2.asScala.foldLeft(accumulator._1)(s)

    override def mergeAccumulators(accumulators: JIterable[(U, JList[T])]): (U, JList[T]) = {
      val combined = accumulators.asScala.map(fold).reduce(c)
      (combined, Lists.newArrayList())
    }

  }

  def combineFn[T, C: ClassTag](createCombiner: T => C, mergeValue: (C, T) => C, mergeCombiners: (C, C) => C)
      : CombineFn[T, (Option[C], JList[T]), C] = new KryoCombineFn[T, (Option[C], JList[T]), C] {
<commit_msg>style fixes<commit_after>package com.spotify.cloud.dataflow.util

import java.lang.{Iterable => JIterable}
import java.util.{List => JList}

import com.google.cloud.dataflow.sdk.coders.{Coder, CoderRegistry}
import com.google.cloud.dataflow.sdk.transforms.Combine.CombineFn
import com.google.cloud.dataflow.sdk.transforms.{SerializableFunction, DoFn}
import com.google.cloud.dataflow.sdk.transforms.Partition.PartitionFn
import com.google.common.collect.Lists
import com.spotify.cloud.dataflow.coders.KryoAtomicCoder
import com.twitter.algebird.{Monoid, Semigroup}

import scala.collection.JavaConverters._
import scala.reflect.ClassTag

private[dataflow] object Functions {

  private val BUFFER_SIZE = 20

  private abstract class KryoCombineFn[VI, VA, VO] extends CombineFn[VI, VA, VO] {

    override def getAccumulatorCoder(registry: CoderRegistry, inputCoder: Coder[VI]): Coder[VA] =
      new KryoAtomicCoder().asInstanceOf[Coder[VA]]

    override def getDefaultOutputCoder(registry: CoderRegistry, inputCoder: Coder[VI]): Coder[VO] =
      new KryoAtomicCoder().asInstanceOf[Coder[VO]]

  }

  def aggregateFn[T, U: ClassTag](zeroValue: U)(seqOp: (U, T) => U, combOp: (U, U) => U)
  : CombineFn[T, (U, JList[T]), U] = new KryoCombineFn[T, (U, JList[T]), U] {

    // defeat closure
    val s = ClosureCleaner(seqOp)
    val c = ClosureCleaner(combOp)

    private def fold(accumulator: (U, JList[T])): U = {
      val (a, l) = accumulator
      l.asScala.foldLeft(a)(s)
    }

    override def createAccumulator(): (U, JList[T]) = (zeroValue, Lists.newArrayList())

    override def addInput(accumulator: (U, JList[T]), input: T): (U, JList[T]) = {
      val (a, l) = accumulator
      l.add(input)
      if (l.size() >= BUFFER_SIZE) {
        (fold(accumulator), Lists.newArrayList())
      } else {
        accumulator
      }
    }

    override def extractOutput(accumulator: (U, JList[T])): U = accumulator._2.asScala.foldLeft(accumulator._1)(s)

    override def mergeAccumulators(accumulators: JIterable[(U, JList[T])]): (U, JList[T]) = {
      val combined = accumulators.asScala.map(fold).reduce(c)
      (combined, Lists.newArrayList())
    }

  }

  def combineFn[T, C: ClassTag](createCombiner: T => C, mergeValue: (C, T) => C, mergeCombiners: (C, C) => C)
  : CombineFn[T, (Option[C], JList[T]), C] = new KryoCombineFn[T, (Option[C], JList[T]), C] {

{'fraction_non_alphanumeric': 0.1489272191838452, 'fraction_numerical': 0.0016827934371055953, 'mean_word_length': 4.191048034934497, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6615030', 'n_tokens_mistral': 1781, 'n_tokens_neox': 1705, 'n_words': 349}