<commit_before>package nest.sparkle.datastream

import scala.collection.mutable.ArrayBuffer
import scala.concurrent.{Future, Promise}
import scala.util.{Failure, Success}

import org.joda.time.{Interval => JodaInterval}
import rx.lang.scala.{Notification, Observable}
import spire.implicits._
import spire.math.Numeric

import nest.sparkle.measure.Span
import nest.sparkle.util.BooleanOption._
import nest.sparkle.util.{Log, PeriodWithZone}

/** Result of reducing a stream of data. Includes a stream of reduced results and
  * a state object that can be used to resume the reduction on a subsequent stream. */
case class PeriodsResult[K,V,F]
    ( reducedStream:DataStream[K,Option[V]],
      override val finishState: Future[Option[FrozenProgress[K,F]]] )
    extends ReductionResult[K,V,FrozenProgress[K,F]]

/** State to continue the reduction between stream */
case class FrozenProgress[K, F](periodsIterator: Iterator[JodaInterval],
                                periodsUsed: Int,
                                periodStart: K,
                                periodEnd: K,
                                reductionProgress: F)

/** functions for reducing a DataStream by time period */
trait DataStreamPeriodReduction[K,V] extends Log {
  self:DataStream[K,V] =>
  implicit def keyTypeTag_ = keyTypeTag
  implicit def valueTypeTag_ = valueTypeTag

  val defaultMaxPeriods = 1000

  /** Reduce an Observable of DataArrays into a smaller Observable by dividing
    * the pair data into partitions based on joda time period, and reducing
    * each partition's pair data with a supplied reduction function. The keys
    * for each partition are the start of the period.
    *
    * The values for each partition are returned optionally, None is returned
    * for partitions that contain no pair data.
    *
    * Note that the key data is interpreted as epoch milliseconds. LATER make this configurable.
    *
    * @param maxPeriods reduce into at most this many time periods
    */
  def reduceByPeriod[F] // format: OFF
      ( periodWithZone: PeriodWithZone,
        range: SoftInterval[K],
        reduction: Reduction[V, F],
        maxPeriods: Int = defaultMaxPeriods,
        optPrevious: Option[FrozenProgress[K, F]] = None)
      ( implicit numericKey: Numeric[K], parentSpan:Span )
      : PeriodsResult[K,V,F] = { // format: ON

    val state = new State(periodWithZone, reduction, range, maxPeriods, optPrevious)
    val finishState = Promise[Option[FrozenProgress[K,F]]]()
    val reduced = data.materialize.flatMap { notification =>
      notification match {
        case Notification.OnNext(dataArray) =>
          val produced = state.processArray(dataArray)
          Observable.from(produced)
        case Notification.OnCompleted =>
          finishState.complete(Success(Some(state.frozenState)))
          state.remaining match {
            case Some(remaining) =>
              Observable.from(Seq(remaining))
            case None            =>
              Observable.empty
          }
        case Notification.OnError(err) =>
          finishState.complete(Failure(err))
          Observable.error(err)
      }
    }

    val reducedStream = new DataStream(reduced)
    PeriodsResult(reducedStream, finishState.future)
  }

  /** tracks progress in iterating through time periods */
  private case class PeriodProgress[F]
      ( periodWithZone: PeriodWithZone,
        maxPeriods:Int,
        targetStart:K,
        until:Option[K] )
      ( implicit numericKey:Numeric[K] ) {
    var start: K = 0.asInstanceOf[K]
    var end: K = 0.asInstanceOf[K]
    var periodsUsed = 0 // total number of time periods visited (to compare against maxPeriods)
    private var periodsIterator: Iterator[JodaInterval] =      // iterates through time periods
      PeriodGroups.jodaIntervals(periodWithZone, targetStart)

    /** advance period iteration to the next time period */
    def toNextPeriod(): Boolean = {

      def clipToRequestedEnd(proposedEnd:K):K = {
        until match {
          case Some(untilEnd) if untilEnd < proposedEnd => untilEnd
          case _                                        => proposedEnd
        }
      }

      if (periodsIterator.hasNext && periodsUsed < maxPeriods) {
        val interval = periodsIterator.next()
        start = numericKey.fromLong(interval.getStartMillis)
        end = clipToRequestedEnd(numericKey.fromLong(interval.getEndMillis))
        periodsUsed += 1
        true
      } else {
        false
      }
    }

    /** move period iteration forward until it reaches a period containing the provided key */
    def advanceTo(key:K): Boolean = {
      var reachedMax = false
      while (start < key && !reachedMax) {
        reachedMax = toNextPeriod()
      }
      reachedMax
    }

    /** return a freeze-dried copy of the current period iteration and total accumulation,
      * (so that we can reconstruct the current state of progress on a later stream).
      */
    def frozen[F](reductionProgress:F):FrozenProgress[K,F] = {
      FrozenProgress(periodsIterator, periodsUsed, start, end, reductionProgress)
    }

    /** return a new PeriodProgress based on the current state combined with a previously frozen one */
    def unfreeze(frozenProgress: FrozenProgress[K,F]): PeriodProgress[F] = {
      val copied:PeriodProgress[F] = this.copy()
      copied.start = frozenProgress.periodStart
      copied.end = frozenProgress.periodEnd
      copied.periodsUsed = frozenProgress.periodsUsed
      copied.periodsIterator = frozenProgress.periodsIterator
      copied
    }
  }


  /** holds the accumulated reduction results */
  private class Results {
    private val keys = ArrayBuffer[K]()
    private val values = ArrayBuffer[Option[V]]()

    /** store a key,Option[value] pair in the results buffer */
    def add(key:K, value:Option[V]): Unit = {<commit_msg>fix bug in rounding the last item in the period sequence<commit_after>package nest.sparkle.datastream

import scala.collection.mutable.ArrayBuffer
import scala.concurrent.{Future, Promise}
import scala.util.{Failure, Success}

import org.joda.time.{Interval => JodaInterval}
import rx.lang.scala.{Notification, Observable}
import spire.implicits._
import spire.math.Numeric

import nest.sparkle.measure.Span
import nest.sparkle.util.BooleanOption._
import nest.sparkle.util.{Log, PeriodWithZone}

/** Result of reducing a stream of data. Includes a stream of reduced results and
  * a state object that can be used to resume the reduction on a subsequent stream. */
case class PeriodsResult[K,V,F]
    ( reducedStream:DataStream[K,Option[V]],
      override val finishState: Future[Option[FrozenProgress[K,F]]] )
    extends ReductionResult[K,V,FrozenProgress[K,F]]

/** State to continue the reduction between stream */
case class FrozenProgress[K, F](periodsIterator: Iterator[JodaInterval],
                                periodsUsed: Int,
                                periodStart: K,
                                periodEnd: K,
                                reductionProgress: F)

/** functions for reducing a DataStream by time period */
trait DataStreamPeriodReduction[K,V] extends Log {
  self:DataStream[K,V] =>
  implicit def keyTypeTag_ = keyTypeTag
  implicit def valueTypeTag_ = valueTypeTag

  val defaultMaxPeriods = 1000

  /** Reduce an Observable of DataArrays into a smaller Observable by dividing
    * the pair data into partitions based on joda time period, and reducing
    * each partition's pair data with a supplied reduction function. The keys
    * for each partition are the start of the period.
    *
    * The values for each partition are returned optionally, None is returned
    * for partitions that contain no pair data.
    *
    * Note that the key data is interpreted as epoch milliseconds. LATER make this configurable.
    *
    * @param maxPeriods reduce into at most this many time periods
    */
  def reduceByPeriod[F] // format: OFF
      ( periodWithZone: PeriodWithZone,
        range: SoftInterval[K],
        reduction: Reduction[V, F],
        maxPeriods: Int = defaultMaxPeriods,
        optPrevious: Option[FrozenProgress[K, F]] = None)
      ( implicit numericKey: Numeric[K], parentSpan:Span )
      : PeriodsResult[K,V,F] = { // format: ON

    val state = new State(periodWithZone, reduction, range, maxPeriods, optPrevious)
    val finishState = Promise[Option[FrozenProgress[K,F]]]()
    val reduced = data.materialize.flatMap { notification =>
      notification match {
        case Notification.OnNext(dataArray) =>
          val produced = state.processArray(dataArray)
          Observable.from(produced)
        case Notification.OnCompleted =>
          finishState.complete(Success(Some(state.frozenState)))
          state.remaining match {
            case Some(remaining) =>
              Observable.from(Seq(remaining))
            case None            =>
              Observable.empty
          }
        case Notification.OnError(err) =>
          finishState.complete(Failure(err))
          Observable.error(err)
      }
    }

    val reducedStream = new DataStream(reduced)
    PeriodsResult(reducedStream, finishState.future)
  }

  /** tracks progress in iterating through time periods */
  private case class PeriodProgress[F]
      ( periodWithZone: PeriodWithZone,
        maxPeriods:Int,
        targetStart:K,
        until:Option[K] )
      ( implicit numericKey:Numeric[K] ) {
    var start: K = 0.asInstanceOf[K]
    var end: K = 0.asInstanceOf[K]
    var periodsUsed = 0 // total number of time periods visited (to compare against maxPeriods)
    private var periodsIterator: Iterator[JodaInterval] =      // iterates through time periods
      PeriodGroups.jodaIntervals(periodWithZone, targetStart)

    /** advance period iteration to the next time period
      * return true if the iterator was advanced, false if iteration has reached its end */
    def toNextPeriod(): Boolean = {

      def clipToRequestedEnd(proposedEnd:K):K = {
        until match {
          case Some(untilEnd) if untilEnd < proposedEnd => untilEnd
          case _                                        => proposedEnd
        }
      }

      if (periodsIterator.hasNext && periodsUsed < maxPeriods) {
        val interval = periodsIterator.next()
        start = numericKey.fromLong(interval.getStartMillis)
        end = clipToRequestedEnd(numericKey.fromLong(interval.getEndMillis))
        periodsUsed += 1
        true
      } else {
        false
      }
    }

    /** move period iteration forward until it reaches a period containing the provided key
      * unless iteration is blocked by maxPeriods. returns true if the the advance is successful */
    def advanceTo(key:K): Boolean = {
      var unBlocked = true
      while (start < key && unBlocked) {
        unBlocked = toNextPeriod()
      }
      unBlocked
    }

    /** return a freeze-dried copy of the current period iteration and total accumulation,
      * (so that we can reconstruct the current state of progress on a later stream).
      */
    def frozen[F](reductionProgress:F):FrozenProgress[K,F] = {
      FrozenProgress(periodsIterator, periodsUsed, start, end, reductionProgress)
    }

    /** return a new PeriodProgress based on the current state combined with a previously frozen one */
    def unfreeze(frozenProgress: FrozenProgress[K,F]): PeriodProgress[F] = {
      val copied:PeriodProgress[F] = this.copy()
      copied.start = frozenProgress.periodStart
      copied.end = frozenProgress.periodEnd
      copied.periodsUsed = frozenProgress.periodsUsed
      copied.periodsIterator = frozenProgress.periodsIterator
      copied
    }
  }


  /** holds the accumulated reduction results */
  private class Results {
    private val keys = ArrayBuffer[K]()
    private val values = ArrayBuffer[Option[V]]()

    /** store a key,Option[value] pair in the results buffer */
    def add(key:K, value:Option[V]): Unit = {
{'fraction_non_alphanumeric': 0.08512597304762702, 'fraction_numerical': 0.0013392483468653218, 'mean_word_length': 2.81238034460753, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4673012', 'n_tokens_mistral': 3252, 'n_tokens_neox': 3178, 'n_words': 978}