<commit_before>/*
 * Copyright 2016 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.scio

import java.nio.ByteBuffer

import com.spotify.scio.metrics._
import com.spotify.scio.util.ScioUtil
import com.twitter.algebird.Semigroup
import org.apache.beam.sdk.Pipeline.PipelineExecutionException
import org.apache.beam.sdk.PipelineResult.State
import org.apache.beam.sdk.io.FileSystems
import org.apache.beam.sdk.util.MimeTypes
import org.apache.beam.sdk.{PipelineResult, metrics => bm}

import scala.collection.JavaConverters._
import scala.concurrent.duration.Duration
import scala.concurrent.{Await, Future}
import scala.reflect.ClassTag
import scala.util.Try

/** Represent a Beam runner specific result. */
trait RunnerResult {
  /** Get a generic [[ScioResult]]. */
  def asScioResult: ScioResult
}

/** Represent a Scio pipeline result. */
abstract class ScioResult private[scio] (val internal: PipelineResult) {

  /** Get a Beam runner specific result. */
  def as[T <: RunnerResult : ClassTag]: T = {
    val cls = ScioUtil.classOf[T]
    try {
      cls.getConstructor(classOf[PipelineResult]).newInstance(internal)
    } catch {
      case e: Throwable =>
        throw new RuntimeException(s"Failed to construct runner specific result $cls", e)
    }
  }

  /**
   * `Future` for pipeline's final state. The `Future` will be completed once the pipeline
   * completes successfully.
   */
  val finalState: Future[State]

  /** Get metrics of the finished pipeline. */
  def getMetrics: Metrics

  /** Wait until the pipeline finishes. */
  def waitUntilFinish(duration: Duration = Duration.Inf): ScioResult = {
    Await.ready(finalState, duration)
    this
  }

  /**
   * Wait until the pipeline finishes with the State `DONE` (as opposed to `CANCELLED` or
   * `FAILED`). Throw exception otherwise.
   */
  def waitUntilDone(duration: Duration = Duration.Inf): ScioResult = {
    waitUntilFinish(duration)
    if (!this.state.equals(State.DONE)) {
      throw new PipelineExecutionException(new Exception(s"Job finished with state ${this.state}"))
    }
    this
  }

  /** Whether the pipeline is completed. */
  def isCompleted: Boolean = internal.getState.isTerminal

  /** Pipeline's current state. */
  def state: State = Try(internal.getState).getOrElse(State.UNKNOWN)

  /** Save metrics of the finished pipeline to a file. */
  def saveMetrics(filename: String): Unit = {
    require(isCompleted, "Pipeline has to be finished to save metrics.")
    val mapper = ScioUtil.getScalaJsonMapper
    val resourceId = FileSystems.matchNewResource(filename, false)
    val out = FileSystems.create(resourceId, MimeTypes.TEXT)
    try {
      out.write(ByteBuffer.wrap(mapper.writeValueAsBytes(getMetrics)))
    } finally {
      if (out != null) {
        out.close()
      }
    }
  }

  protected def getBeamMetrics: BeamMetrics = {
    require(isCompleted, "Pipeline has to be finished to get metrics.")

    def mkDist(d: bm.DistributionResult): BeamDistribution =
      BeamDistribution(d.sum(), d.count(), d.min(), d.max(), d.mean())
    def mkGauge(g: bm.GaugeResult): BeamGauge = BeamGauge(g.value(), g.timestamp())

    val beamCounters = allCounters.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(), v)
    }
    val beamDistributions = allDistributions.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(), MetricValue(mkDist(v.attempted), v.committed.map(mkDist)))
    }
    val beamGauges = allGauges.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(),
        MetricValue(mkGauge(v.attempted), v.committed.map(mkGauge)))
    }
    BeamMetrics(beamCounters, beamDistributions, beamGauges)
  }
<commit_msg>Fix Distribution counter publishing for batch jobs (issue 1029) (#1030)<commit_after>/*
 * Copyright 2016 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.scio

import java.nio.ByteBuffer

import com.spotify.scio.metrics._
import com.spotify.scio.util.ScioUtil
import com.twitter.algebird.Semigroup
import org.apache.beam.sdk.Pipeline.PipelineExecutionException
import org.apache.beam.sdk.PipelineResult.State
import org.apache.beam.sdk.io.FileSystems
import org.apache.beam.sdk.metrics.{DistributionResult, GaugeResult}
import org.apache.beam.sdk.util.MimeTypes
import org.apache.beam.sdk.{PipelineResult, metrics => bm}

import scala.collection.JavaConverters._
import scala.concurrent.duration.Duration
import scala.concurrent.{Await, Future}
import scala.reflect.ClassTag
import scala.util.Try

/** Represent a Beam runner specific result. */
trait RunnerResult {
  /** Get a generic [[ScioResult]]. */
  def asScioResult: ScioResult
}

/** Represent a Scio pipeline result. */
abstract class ScioResult private[scio] (val internal: PipelineResult) {

  /** Get a Beam runner specific result. */
  def as[T <: RunnerResult : ClassTag]: T = {
    val cls = ScioUtil.classOf[T]
    try {
      cls.getConstructor(classOf[PipelineResult]).newInstance(internal)
    } catch {
      case e: Throwable =>
        throw new RuntimeException(s"Failed to construct runner specific result $cls", e)
    }
  }

  /**
   * `Future` for pipeline's final state. The `Future` will be completed once the pipeline
   * completes successfully.
   */
  val finalState: Future[State]

  /** Get metrics of the finished pipeline. */
  def getMetrics: Metrics

  /** Wait until the pipeline finishes. */
  def waitUntilFinish(duration: Duration = Duration.Inf): ScioResult = {
    Await.ready(finalState, duration)
    this
  }

  /**
   * Wait until the pipeline finishes with the State `DONE` (as opposed to `CANCELLED` or
   * `FAILED`). Throw exception otherwise.
   */
  def waitUntilDone(duration: Duration = Duration.Inf): ScioResult = {
    waitUntilFinish(duration)
    if (!this.state.equals(State.DONE)) {
      throw new PipelineExecutionException(new Exception(s"Job finished with state ${this.state}"))
    }
    this
  }

  /** Whether the pipeline is completed. */
  def isCompleted: Boolean = internal.getState.isTerminal

  /** Pipeline's current state. */
  def state: State = Try(internal.getState).getOrElse(State.UNKNOWN)

  /** Save metrics of the finished pipeline to a file. */
  def saveMetrics(filename: String): Unit = {
    require(isCompleted, "Pipeline has to be finished to save metrics.")
    val mapper = ScioUtil.getScalaJsonMapper
    val resourceId = FileSystems.matchNewResource(filename, false)
    val out = FileSystems.create(resourceId, MimeTypes.TEXT)
    try {
      out.write(ByteBuffer.wrap(mapper.writeValueAsBytes(getMetrics)))
    } finally {
      if (out != null) {
        out.close()
      }
    }
  }

  protected def getBeamMetrics: BeamMetrics = {
    require(isCompleted, "Pipeline has to be finished to get metrics.")

    def mkDist(d: bm.DistributionResult): BeamDistribution = {
      val dist = Option(d).getOrElse(DistributionResult.ZERO)
      BeamDistribution(dist.sum(), dist.count(), dist.min(), dist.max(), dist.mean())
    }
    def mkGauge(g: bm.GaugeResult): BeamGauge = {
      val gauge = Option(g).getOrElse(GaugeResult.empty())
      BeamGauge(gauge.value(), gauge.timestamp())
    }
    val beamCounters = allCounters.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(), v)
    }
    val beamDistributions = allDistributions.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(), MetricValue(mkDist(v.attempted), v.committed.map(mkDist)))
    }
    val beamGauges = allGauges.map { case (k, v) =>
      BeamMetric(k.namespace(), k.name(),
        MetricValue(mkGauge(v.attempted), v.committed.map(mkGauge)))
    }
    BeamMetrics(beamCounters, beamDistributions, beamGauges)
  }

{'fraction_non_alphanumeric': 0.1139022711631108, 'fraction_numerical': 0.0027529249827942187, 'mean_word_length': 3.6550987720234915, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '7389247', 'n_tokens_mistral': 2692, 'n_tokens_neox': 2606, 'n_words': 746}