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

import java.io.{IOException, ByteArrayOutputStream, OutputStream, InputStream}

import com.google.cloud.dataflow.sdk.coders.{CoderException, AtomicCoder}
import com.google.cloud.dataflow.sdk.coders.Coder.Context
import com.google.cloud.dataflow.sdk.util.VarInt
import com.google.common.io.ByteStreams
import com.twitter.chill._
import org.apache.avro.specific.SpecificRecord

import scala.collection.JavaConversions.JIterableWrapper

private[dataflow] class KryoAtomicCoder extends AtomicCoder[Any] {

  @transient
  private lazy val kryo: Kryo = {
    val k = KryoSerializer.registered.newKryo()

    // java.lang.Iteable.asScala returns JIterableWrapper which causes problem.
    // Treat it as standard Iteable instead.
    k.register(classOf[JIterableWrapper[_]], new JIterableWrapperSerializer)

    k.forSubclass[SpecificRecord](new AvroSerializer)

    k.forClass(new KVSerializer)
    // TODO:
    // InstantCoder
    // TimestampedValueCoder

    k
  }

  override def encode(value: Any, outStream: OutputStream, context: Context): Unit = {
    if (value == null) {
      throw new CoderException("cannot encode a null value")
    }
    if (context.isWholeStream) {
      val output = new Output(outStream)
      kryo.writeClassAndObject(output, value)
      output.flush()
    } else {
      val s = new ByteArrayOutputStream()
      val output = new Output(s)
      kryo.writeClassAndObject(output, value)
      output.flush()<commit_msg>add TODO for scala 2.11<commit_after>package com.spotify.cloud.dataflow.coders

import java.io.{IOException, ByteArrayOutputStream, OutputStream, InputStream}

import com.google.cloud.dataflow.sdk.coders.{CoderException, AtomicCoder}
import com.google.cloud.dataflow.sdk.coders.Coder.Context
import com.google.cloud.dataflow.sdk.util.VarInt
import com.google.common.io.ByteStreams
import com.twitter.chill._
import org.apache.avro.specific.SpecificRecord

// TODO: need to fix for scala 2.11
import scala.collection.JavaConversions.JIterableWrapper

private[dataflow] class KryoAtomicCoder extends AtomicCoder[Any] {

  @transient
  private lazy val kryo: Kryo = {
    val k = KryoSerializer.registered.newKryo()

    // java.lang.Iteable.asScala returns JIterableWrapper which causes problem.
    // Treat it as standard Iteable instead.
    // TODO: need to fix for scala 2.11
    k.register(classOf[JIterableWrapper[_]], new JIterableWrapperSerializer)

    k.forSubclass[SpecificRecord](new AvroSerializer)

    k.forClass(new KVSerializer)
    // TODO:
    // InstantCoder
    // TimestampedValueCoder

    k
  }

  override def encode(value: Any, outStream: OutputStream, context: Context): Unit = {
    if (value == null) {
      throw new CoderException("cannot encode a null value")
    }
    if (context.isWholeStream) {
      val output = new Output(outStream)
      kryo.writeClassAndObject(output, value)
      output.flush()
    } else {
      val s = new ByteArrayOutputStream()
      val output = new Output(s)
      kryo.writeClassAndObject(output, value)
      output.flush()
{'fraction_non_alphanumeric': 0.09709677419354838, 'fraction_numerical': 0.002903225806451613, 'mean_word_length': 3.8302180685358254, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '6891589', 'n_tokens_mistral': 995, 'n_tokens_neox': 960, 'n_words': 215}