<commit_before>package spark

import java.io.FileOutputStream
import java.io.ObjectOutputStream
import scala.collection.mutable.HashMap


class ShuffleMapTask(stageId: Int, rdd: RDD[_], dep: ShuffleDependency[_,_,_], val partition: Int, locs: Seq[String])
extends DAGTask[String](stageId) with Logging {
  val split = rdd.splits(partition)

  override def run (attemptId: Int): String = {
    val numOutputSplits = dep.partitioner.numPartitions
    val aggregator = dep.aggregator.asInstanceOf[Aggregator[Any, Any, Any]]
    val partitioner = dep.partitioner.asInstanceOf[Partitioner]
    val buckets = Array.tabulate(numOutputSplits)(_ => new HashMap[Any, Any])
    for (elem <- rdd.iterator(split)) {
      val (k, v) = elem.asInstanceOf[(Any, Any)]
      var bucketId = partitioner.getPartition(k)
      val bucket = buckets(bucketId)
      bucket(k) = bucket.get(k) match {
        case Some(c) => aggregator.mergeValue(c, v)
        case None => aggregator.createCombiner(v)
      }
    }
    val ser = SparkEnv.get.serializer.newInstance()
    for (i <- 0 until numOutputSplits) {
      val file = LocalFileShuffle.getOutputFile(dep.shuffleId, partition, i)
      val out = ser.outputStream(new FileOutputStream(file))
      buckets(i).foreach(pair => out.writeObject(pair))
      // TODO: have some kind of EOF marker
      out.close()
    }
    return LocalFileShuffle.getServerUri<commit_msg>Use BufferedOutputStream in ShuffleMapTask<commit_after>package spark

import java.io.BufferedOutputStream
import java.io.FileOutputStream
import java.io.ObjectOutputStream
import scala.collection.mutable.HashMap


class ShuffleMapTask(stageId: Int, rdd: RDD[_], dep: ShuffleDependency[_,_,_], val partition: Int, locs: Seq[String])
extends DAGTask[String](stageId) with Logging {
  val split = rdd.splits(partition)

  override def run (attemptId: Int): String = {
    val numOutputSplits = dep.partitioner.numPartitions
    val aggregator = dep.aggregator.asInstanceOf[Aggregator[Any, Any, Any]]
    val partitioner = dep.partitioner.asInstanceOf[Partitioner]
    val buckets = Array.tabulate(numOutputSplits)(_ => new HashMap[Any, Any])
    for (elem <- rdd.iterator(split)) {
      val (k, v) = elem.asInstanceOf[(Any, Any)]
      var bucketId = partitioner.getPartition(k)
      val bucket = buckets(bucketId)
      bucket(k) = bucket.get(k) match {
        case Some(c) => aggregator.mergeValue(c, v)
        case None => aggregator.createCombiner(v)
      }
    }
    val ser = SparkEnv.get.serializer.newInstance()
    for (i <- 0 until numOutputSplits) {
      val file = LocalFileShuffle.getOutputFile(dep.shuffleId, partition, i)
      val out = ser.outputStream(new BufferedOutputStream(new FileOutputStream(file)))
      buckets(i).foreach(pair => out.writeObject(pair))
      // TODO: have some kind of EOF marker
      out.close()
    }
    return LocalFileShuffle.getServerUri
{'fraction_non_alphanumeric': 0.11268579329415832, 'fraction_numerical': 0.000691323885240235, 'mean_word_length': 3.93015332197615, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, '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': '3622825', 'n_tokens_mistral': 915, 'n_tokens_neox': 900, 'n_words': 211}