<commit_before>package glint.models.client.buffered

import akka.util.Timeout
import breeze.linalg.Vector
import glint.models.client.BigMatrix

import scala.concurrent.{ExecutionContext, Future}
import scala.reflect.ClassTag

/**
  * A [[glint.models.client.BigMatrix BigMatrix]] that buffers numerous pushes before sending them in one go
  *
  * {{{
  *   matrix = client.matrix[Double](10000, 100)
  *   bufferedMatrix = new BufferedBigMatrix[Double](matrix, 1000)
  *   bufferedMatrix.bufferedPush(2, 6, 10.0)
  *   if (bufferedBigMatrix.isFull) {
  *     bufferedBigMatrix.flush()
  *   }
  * }}}
  *
  * @param underlying The underlying big matrix implementation
  * @param bufferSize The buffer size
  * @tparam V The type of values to store
  */
class BufferedBigMatrix[@specialized V: ClassTag](underlying: BigMatrix[V], bufferSize: Int) extends BigMatrix[V] {

  private val bufferRows = new Array[Long](bufferSize)
  private val bufferCols = new Array[Int](bufferSize)
  private val bufferValues = new Array[V](bufferSize)
  private var bufferIndex: Int = 0

  /**
    * Pulls a set of rows using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing the vectors representing the rows
    */
  override def pull(rows: Array[Long])(implicit timeout: Timeout, ec: ExecutionContext): Future[Array[Vector[V]]] = {
    underlying.pull(rows)
  }

  /**
    * Destroys the underlying big matrix and its resources on the parameter server
    *
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future whether the matrix was successfully destroyed
    */
  override def destroy()(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = underlying.destroy()

  /**
    * Pushes a set of values using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param cols The indices of the columns
    * @param values The values to update
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing either the success or failure of the operation
    */
  override def push(rows: Array[Long],
                    cols: Array[Int],
                    values: Array[V])(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = {
    underlying.push(rows, cols, values)
  }

  /**
    * Pushes a value into the buffer as long as there is space. Make sure to manually
    * [[glint.models.client.buffered.BufferedBigMatrix.flush flush]] the buffer when it is full, use
    * [[glint.models.client.buffered.BufferedBigMatrix.isFull isFull]] to check if the buffer is full. Attempting to
    * push to the buffer when it is full and not flushed will cause a java.lang.IndexOutOfBoundsException.
    *
    * @param row The row
    * @param col The column
    * @param value The value
    */
  def bufferedPush(row: Long, col: Int, value: V): Unit = {
    bufferRows(bufferIndex) = row
    bufferCols(bufferIndex) = col
    bufferValues(bufferIndex) = value
    bufferIndex += 1
  }

  /**
    * Flushes the buffer to the parameter server.
    *
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing the success or failure of the operation
    */
  def flush()(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = {
    if (bufferIndex == 0) {
      Future {
        true
      }
    } else {
      var index: Int = 0
      val pushRows = new Array[Long](bufferIndex)
      val pushCols = new Array[Int](bufferIndex)
      val pushValues = new Array[V](bufferIndex)
      while (index < bufferIndex) {
        pushRows(index) = bufferRows(index)
        pushCols(index) = bufferCols(index)
        pushValues(index) = bufferValues(index)
        index += 1
      }
      bufferIndex = 0
      underlying.push(pushRows, pushCols, pushValues)
    }
  }

  /**
    * @return True if the buffer is full, false otherwise
    */
  def isFull: Boolean = size == bufferSize

  /**
    * @return The size of the current buffer (in number of elements)
    */
  def size: Int = bufferIndex

  /**
    * Pulls a set of elements using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param cols The corresponding indices of the columns<commit_msg>Add functionality to automatically flush buffer on push<commit_after>package glint.models.client.buffered

import akka.util.Timeout
import breeze.linalg.Vector
import glint.models.client.BigMatrix

import scala.concurrent.{ExecutionContext, Future}
import scala.reflect.ClassTag

/**
  * A [[glint.models.client.BigMatrix BigMatrix]] that buffers numerous pushes before sending them in one go
  *
  * {{{
  *   matrix = client.matrix[Double](10000, 100)
  *   bufferedMatrix = new BufferedBigMatrix[Double](matrix, 1000)
  *   bufferedMatrix.bufferedPush(2, 6, 10.0)
  *   if (bufferedBigMatrix.isFull) {
  *     bufferedBigMatrix.flush()
  *   }
  * }}}
  *
  * @param underlying The underlying big matrix implementation
  * @param bufferSize The buffer size
  * @tparam V The type of values to store
  */
class BufferedBigMatrix[@specialized V: ClassTag](underlying: BigMatrix[V], bufferSize: Int) extends BigMatrix[V] {

  private val bufferRows = new Array[Long](bufferSize)
  private val bufferCols = new Array[Int](bufferSize)
  private val bufferValues = new Array[V](bufferSize)
  private var bufferIndex: Int = 0

  /**
    * Pulls a set of rows using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing the vectors representing the rows
    */
  override def pull(rows: Array[Long])(implicit timeout: Timeout, ec: ExecutionContext): Future[Array[Vector[V]]] = {
    underlying.pull(rows)
  }

  /**
    * Destroys the underlying big matrix and its resources on the parameter server
    *
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future whether the matrix was successfully destroyed
    */
  override def destroy()(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = underlying.destroy()

  /**
    * Pushes a set of values using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param cols The indices of the columns
    * @param values The values to update
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing either the success or failure of the operation
    */
  override def push(rows: Array[Long],
                    cols: Array[Int],
                    values: Array[V])(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = {
    underlying.push(rows, cols, values)
  }

  /**
    * Pushes a value into the buffer as long as there is space. If the buffer is full this method will automatically
    * flush the buffer to the parameter server. It will either return an option that is either set to a future when the
    * buffer was flushed or set to None when the buffer is not yet full.
    *
    * @param row The row
    * @param col The column
    * @param value The value
    * @return An option set to Future[Boolean] if a flush occurred or None if no flush occurred
    */
  def bufferedPush(row: Long,
                   col: Int,
                   value: V)(implicit timeout: Timeout, ec: ExecutionContext): Option[Future[Boolean]] = {
    bufferRows(bufferIndex) = row
    bufferCols(bufferIndex) = col
    bufferValues(bufferIndex) = value
    bufferIndex += 1
    flushIfFull()
  }

  /**
    * Flushes the buffer to the parameter server.
    *
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing the success or failure of the operation
    */
  def flush()(implicit timeout: Timeout, ec: ExecutionContext): Future[Boolean] = {
    if (bufferIndex == 0) {
      Future {
        true
      }
    } else {
      var index: Int = 0
      val pushRows = new Array[Long](bufferIndex)
      val pushCols = new Array[Int](bufferIndex)
      val pushValues = new Array[V](bufferIndex)
      while (index < bufferIndex) {
        pushRows(index) = bufferRows(index)
        pushCols(index) = bufferCols(index)
        pushValues(index) = bufferValues(index)
        index += 1
      }
      bufferIndex = 0
      underlying.push(pushRows, pushCols, pushValues)
    }
  }

  /**
    * Flushes the buffer to the parameter server if it is full.
    *
    * @param timeout The timeout for this request
    * @param ec The implicit execution context in which to execute the request
    * @return A future containing the success or failure of the operation
    */
  def flushIfFull()(implicit timeout: Timeout, ec: ExecutionContext): Option[Future[Boolean]] = {
    if (isFull) {
      Some(flush())
    } else {
      None
    }
  }

  /**
    * @return True if the buffer is full, false otherwise
    */
  def isFull: Boolean = size == bufferSize

  /**
    * @return The size of the current buffer (in number of elements)
    */
  def size: Int = bufferIndex

  /**
    * Pulls a set of elements using the underlying BigMatrix implementation
    *
    * @param rows The indices of the rows
    * @param cols The corresponding indices of the columns
{'fraction_non_alphanumeric': 0.08430818878580072, 'fraction_numerical': 0.00463896732553449, 'mean_word_length': 3.0844316309719932, 'pattern_counts': {'":': 0, '<': 5, '<?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': '4154382', 'n_tokens_mistral': 2770, 'n_tokens_neox': 2689, 'n_words': 1054}