<commit_before>package com.Alteryx.sparkGLM

import breeze.linalg._
import breeze.numerics._

import edu.berkeley.cs.amplab.mlmatrix.{NormalEquations, RowPartitionedMatrix}
import edu.berkeley.cs.amplab.mlmatrix.util.Utils
import edu.berkeley.cs.amplab.mlmatrix.NormalEquations._

import org.apache.spark.sql.{DataFrame, Column}
import org.apache.spark.sql.functions.{col, lit}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.types.DoubleType

object utils {

  //
  // DataFrame utility functions
  //

  def matchCols(estDF: DataFrame, scoreDF: DataFrame): DataFrame = {
    val missingCols: Array[Column] = estDF.columns.diff(scoreDF.columns).map { field =>
      lit(0).cast(DoubleType).as(field)
    }
    scoreDF.select(missingCols :+ col("*"):_*)
  }

  def matchCols(xnames: Array[String], scoreDF: DataFrame): DataFrame = {
    val missingCols: Array[Column] = xnames.diff(scoreDF.columns).map { field =>
      lit(0).cast(DoubleType).as(field)
    }
    scoreDF.select(missingCols :+ col("*"):_*)
  }

  // HOLDING ON TO THIS FOR NOW, IT DEPENDS IF THINGS MAKE IT ACROSS FROM mlmatrix
  def dataFrameToMatrix(df: DataFrame): RDD[DenseMatrix[Double]] = {
    val matrixRDD = df.map(x => x.toSeq.toArray).map(y => y.map(z => z.asInstanceOf[Double]))
    RowPartitionedMatrix.arrayToMatrix(matrixRDD)
  }

  // Convert a single partitioned DataFrame into a DenseMatrix
  def dfToDenseMatrix(df: DataFrame): DenseMatrix[Double] = {
    require(df.rdd.partitions.size == 1,
      "The DataFrame must be in a single partition")
    require(df.dtypes.forall(_._2 == "DoubleType"),
      "The provided DataFrame must contain all 'DoubleType' columns")
    val dataArray = dataFrameToMatrix(df).collect
    dataArray(0)
  }


  //
  // Modelling support methods
  //

  // From ml-matrix, the reducer function to sum together the partitons of X'X
  // and X'y
  def reduceNormal(
    a: (DenseMatrix[Double], DenseMatrix[Double]),
    b: (DenseMatrix[Double], DenseMatrix[Double])): (DenseMatrix[Double], DenseMatrix[Double]) = {
    a._1 :+= b._1
    a._2 :+= b._2
    a
  }

  // A case class for a minimal weighted least squares fit object
  case class WLSObj(coefs: DenseMatrix[Double], diagDesign: DenseVector[Double])

  // A method to do a weighted least squares fit to data in a single partition
  def wlsSingle (
      X: DenseMatrix[Double],
      y: DenseMatrix[Double],
      w: DenseMatrix[Double]): WLSObj = {
    val W = diag(w.toDenseVector)
    val XtWXi = inv(X.t * W * X)
    val XtWy = X.t * W * y
    val coefs = XtWXi * XtWy
    val diagDesign = sqrt(diag(XtWXi))
    new WLSObj(coefs = coefs, diagDesign = diagDesign)
  }

  // A method to get the coponents needed to do WLS on RowPartionedMatrix objects
  def wlsComponents (
      X: RowPartitionedMatrix,
      y: RowPartitionedMatrix,
      w: RowPartitionedMatrix): (DenseMatrix[Double], DenseMatrix[Double]) = {
    val Xyw = X.rdd.zip(y.rdd).zip(w.rdd).map {
      case((a, b), c) => (a.mat, b.mat, c.mat)
    }
    val XtWX_XtWy = Xyw.map { part =>
      (part._1.t * diag(part._3.toDenseVector) * part._1,
        part._1.t * diag(part._3.toDenseVector) * part._2)
    }

    val treeBranchingFactor = X.rdd.context.getConf.getInt("spark.mlmatrix.treeBranchingFactor", 2).toInt
    val depth = math.ceil(math.log(XtWX_XtWy.partitions.size)/math.log(treeBranchingFactor)).toInt
    val reduced = edu.berkeley.cs.amplab.mlmatrix.util.Utils.treeReduce(XtWX_XtWy, utils.reduceNormal, depth=depth)

    reduced
  }<commit_msg>Added WLS changes from @dputler<commit_after>package com.Alteryx.sparkGLM

import breeze.linalg._
import breeze.numerics._

import edu.berkeley.cs.amplab.mlmatrix.{NormalEquations, RowPartitionedMatrix}
import edu.berkeley.cs.amplab.mlmatrix.util.Utils
import edu.berkeley.cs.amplab.mlmatrix.NormalEquations._

import org.apache.spark.sql.{DataFrame, Column}
import org.apache.spark.sql.functions.{col, lit}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.types.DoubleType

object utils {

  //
  // DataFrame utility functions
  //

  def matchCols(estDF: DataFrame, scoreDF: DataFrame): DataFrame = {
    val missingCols: Array[Column] = estDF.columns.diff(scoreDF.columns).map { field =>
      lit(0).cast(DoubleType).as(field)
    }
    scoreDF.select(missingCols :+ col("*"):_*)
  }

  def matchCols(xnames: Array[String], scoreDF: DataFrame): DataFrame = {
    val missingCols: Array[Column] = xnames.diff(scoreDF.columns).map { field =>
      lit(0).cast(DoubleType).as(field)
    }
    scoreDF.select(missingCols :+ col("*"):_*)
  }

  // HOLDING ON TO THIS FOR NOW, IT DEPENDS IF THINGS MAKE IT ACROSS FROM mlmatrix
  def dataFrameToMatrix(df: DataFrame): RDD[DenseMatrix[Double]] = {
    val matrixRDD = df.map(x => x.toSeq.toArray).map(y => y.map(z => z.asInstanceOf[Double]))
    RowPartitionedMatrix.arrayToMatrix(matrixRDD)
  }

  // Convert a single partitioned DataFrame into a DenseMatrix
  def dfToDenseMatrix(df: DataFrame): DenseMatrix[Double] = {
    require(df.rdd.partitions.size == 1,
      "The DataFrame must be in a single partition")
    require(df.dtypes.forall(_._2 == "DoubleType"),
      "The provided DataFrame must contain all 'DoubleType' columns")
    val dataArray = dataFrameToMatrix(df).collect
    dataArray(0)
  }


  //
  // Modelling support methods
  //

  // From ml-matrix, the reducer function to sum together the partitons of X'X
  // and X'y
  def reduceNormal(
    a: (DenseMatrix[Double], DenseMatrix[Double]),
    b: (DenseMatrix[Double], DenseMatrix[Double])): (DenseMatrix[Double], DenseMatrix[Double]) = {
    a._1 :+= b._1
    a._2 :+= b._2
    a
  }

  // A method to multiple a left-hand side matrix by a diagonal matrix
  // without actually creating the diagnol matrix that may not fit into memory
  def leftMultDiag(
    left: DenseMatrix[Double],
    dcol: DenseMatrix[Double]): DenseMatrix[Double] = {
      require(dcol.cols == 1,
        "The argument dcol must contain only a single column")
      val nCols = left.cols - 1
      val dcolArray = dcol.toArray
      val newLeft = new DenseMatrix(left.rows, left.cols, left.toArray)
      for (i <- 0 to nCols) {
        newLeft(::, i) :*= dcolArray(i)
      }
      newLeft
  }
  // A case class for a minimal weighted least squares fit object
  case class WLSObj(coefs: DenseMatrix[Double], diagDesign: DenseVector[Double])

  // A method to do a weighted least squares fit to data in a single partition
  def wlsSingle (
      X: DenseMatrix[Double],
      y: DenseMatrix[Double],
      w: DenseMatrix[Double]): WLSObj = {
    val XtW = leftMultDiag(X.t, w)
    val XtWXi = inv(XtW * X)
    val XtWy = XtW * y
    val coefs = XtWXi * XtWy
    val diagDesign = sqrt(diag(XtWXi))
    new WLSObj(coefs = coefs, diagDesign = diagDesign)
  }

  // A method to get the coponents needed to do WLS on RowPartionedMatrix objects
  def wlsComponents (
      X: RowPartitionedMatrix,
      y: RowPartitionedMatrix,
      w: RowPartitionedMatrix): (DenseMatrix[Double], DenseMatrix[Double]) = {
    val Xyw = X.rdd.zip(y.rdd).zip(w.rdd).map {
      case((a, b), c) => (a.mat, b.mat, c.mat)
    }
    val XtWX_XtWy = Xyw.map { part =>
      (leftMultDiag(part._1.t, part._3) * part._1,
        leftMultDiag(part._1.t, part._3) * part._2)
    }

    val treeBranchingFactor = X.rdd.context.getConf.getInt("spark.mlmatrix.treeBranchingFactor", 2).toInt
    val depth = math.ceil(math.log(XtWX_XtWy.partitions.size)/math.log(treeBranchingFactor)).toInt
    val reduced = edu.berkeley.cs.amplab.mlmatrix.util.Utils.treeReduce(XtWX_XtWy, utils.reduceNormal, depth=depth)

    reduced
  }
{'fraction_non_alphanumeric': 0.12621486734961912, 'fraction_numerical': 0.004596795376937221, 'mean_word_length': 3.4558221181977764, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 17, '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': '4377299', 'n_tokens_mistral': 2761, 'n_tokens_neox': 2670, 'n_words': 609}