You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by mateiz <gi...@git.apache.org> on 2014/07/20 09:24:28 UTC
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
GitHub user mateiz opened a pull request:
https://github.com/apache/spark/pull/1499
(WIP) SPARK-2045 Sort-based shuffle
This adds a new ShuffleManager based on sorting, as described in https://issues.apache.org/jira/browse/SPARK-2045. The bulk of the code is in an ExternalSorter class that is similar to ExternalAppendOnlyMap, but sorts key-value pairs by partition ID and can be used to create a single sorted file with a map task's output. (Longer-term I think this can take on the remaining functionality in ExternalAppendOnlyMap and replace it so we don't have code duplication.)
The main TODOs still left are enabling ExternalSorter to merge *across* spilled files, though this may be of limited utility (if the data is highly combinable this will happen early on), and adding more tests (e.g. a version of our shuffle suite that runs on this).
Despite this though, this seems to work pretty well (running successfully in cases where the hash shuffle would OOM, such as 1000 reduce tasks on executors with only 1G memory), and it seems to be comparable in speed or faster than hash-based shuffle (it will create much fewer files for the OS to keep track of). So I'm posting it to get some early feedback.
After these TODOs are done, I'd also like to enable ExternalSorter to sort data within each partition by a key as well, which will allow us to use it to implement external spilling in reduce tasks in `sortByKey`.
You can merge this pull request into a Git repository by running:
$ git pull https://github.com/mateiz/spark sort-based-shuffle
Alternatively you can review and apply these changes as the patch at:
https://github.com/apache/spark/pull/1499.patch
To close this pull request, make a commit to your master/trunk branch
with (at least) the following in the commit message:
This closes #1499
----
commit cf014cb3ef394f8b34c717a2ab8c92f95edf6269
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-17T02:13:54Z
Scaffolding for sort-based shuffle
commit c1279c4b43121ad47d4e0c0cd46ba978fc29c7de
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-18T02:07:41Z
Some more partial work towards sort-based shuffle
commit 784e85a23c8bdea821dbbd2aa43f7e202ac39fe3
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-18T07:26:49Z
More partial work towards sort-based shuffle
commit 727685bd23fa83df18f471ece5ded657f8656f81
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-19T04:18:17Z
More work
commit abf052dbaafaea604f9b807e8ec5c38bd8e354fa
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-19T08:35:54Z
Add more error handling and tests for error cases
commit 066553bc99a389556eb072161fe5acc43dd73262
Author: Matei Zaharia <ma...@databricks.com>
Date: 2014-07-20T06:40:40Z
Add spill metrics to map tasks
----
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50669058
Jenkins, test this please
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49956204
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17078/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560994
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
--- End diff --
You could move this within `flush`
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49715947
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds no public classes<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16948/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49555343
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16878/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15494109
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
+
+ /** Shut down this ShuffleManager. */
+ override def stop(): Unit = {}
+
+ /** Get the location of a block in a map output file. Uses the index file we create for it. */
+ def getBlockLocation(blockId: ShuffleBlockId, diskManager: DiskBlockManager): FileSegment = {
+ // The block is actually going to be a range of a single map output file for this map, so
+ // figure out the ID of the consolidated file, then the offset within that from our index
+ val realId = ShuffleBlockId(blockId.shuffleId, blockId.mapId, 0)
+ val indexFile = diskManager.getFile(realId.name + ".index")
+ val in = new DataInputStream(new FileInputStream(indexFile))
--- End diff --
This index should be pretty small, so I suppose we could keep it in memory instead of writing it to a file. On the other hand, it might be likely to remain in the filesystem cache.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154344
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
--- End diff --
Do you consider this algorithm preferable to one where you immediately try to lookup a matching key in the buffer and combine in place? Such an algorithm would use nearly constant space and still maintain nearly linear time if collisions were rare (which they are in the hash case, though perhaps not with general partial orderings).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15544196
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
--- End diff --
comment on what this does
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49925284
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17055/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15494991
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
+
+ /** Shut down this ShuffleManager. */
+ override def stop(): Unit = {}
+
+ /** Get the location of a block in a map output file. Uses the index file we create for it. */
+ def getBlockLocation(blockId: ShuffleBlockId, diskManager: DiskBlockManager): FileSegment = {
+ // The block is actually going to be a range of a single map output file for this map, so
+ // figure out the ID of the consolidated file, then the offset within that from our index
+ val realId = ShuffleBlockId(blockId.shuffleId, blockId.mapId, 0)
+ val indexFile = diskManager.getFile(realId.name + ".index")
+ val in = new DataInputStream(new FileInputStream(indexFile))
--- End diff --
The problem is cleaning it up over time, it gets complicated (though we still have to clean up the files in ShuffleBlockManager, which is kind of messy).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50115862
@mridulm that should've been fixed recently in https://github.com/mateiz/spark/commit/9c299579f13f004f5fd1f4dd0b98b7d76cac2a55, which got rid of custom return types in ShuffledRDD for this reason (see the message there). Maybe you need to do a clean rebuild. In an earlier commit I also added a version of ShuffleSuite that runs all those tests with sort-based shuffle.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50669721
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17472/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50679871
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>class ShuffledRDD[K, V, C](<br>case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br>case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17471/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154006
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
--- End diff --
nit: For the purposes of slightly more defensive programming, consider putting this inside a `Utils.tryLog { }`. We probably don't want to leak files.
By the way, is this logic correct for an exception? It looks like we'll just keep merrily chugging along -- should we propagate the exception upwards?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154239
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
--- End diff --
Perhaps note that the ordering is assumed total -- it's probably obvious, but since we support partial ordering in this class, may not hurt to mention it
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49564210
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16882/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154080
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
--- End diff --
nit: it may be marginally clearer to put this inside the if statement where we call flush(), since it's really just a property tracked during the consumption of the iterator (though the name suggests something more general).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15603895
--- Diff: core/src/main/scala/org/apache/spark/storage/ShuffleBlockManager.scala ---
@@ -91,6 +97,20 @@ class ShuffleBlockManager(blockManager: BlockManager) extends Logging {
private val metadataCleaner =
new MetadataCleaner(MetadataCleanerType.SHUFFLE_BLOCK_MANAGER, this.cleanup, conf)
+ /**
+ * Register a completed map without getting a ShuffleWriterGroup. Used by sort-based shuffle
+ * because it just writes a single file by itself.
+ */
+ def addCompletedMap(shuffleId: Int, mapId: Int, numBuckets: Int): Unit = {
+ shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numBuckets))
--- End diff --
Yeah IMO we should open a JIRA to fix this later. It's only called once per map task but it is kind of wasteful.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50700807
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17526/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15259190
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
BTW, ExternalAppendOnlyMap is also broken in this respect : I added the comment above to explain why I observed the EOF exception.
This has been fixed in our internal branch, and so not observed without this patch :-)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15556144
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
--- End diff --
would be great to comment on why we are doing this batching (to avoid hash map growing too large right?)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15563349
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
--- End diff --
As you said in the comments, this is very similar to your code in `mergeSort`. It'll probably be cleaner to abstract this and add a `shoulCombine: Boolean` field in `mergeSort` for the extra logic here.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15324369
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala ---
@@ -120,8 +124,10 @@ class ExternalAppendOnlyMap[K, V, C](
val update: (Boolean, C) => C = (hadVal, oldVal) => {
if (hadVal) mergeValue(oldVal, value) else createCombiner(value)
}
- if (numPairsInMemory > trackMemoryThreshold && currentMap.atGrowThreshold) {
- val mapSize = currentMap.estimateSize()
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ currentMap.estimateSize() >= myMemoryThreshold)
--- End diff --
Did you mean the following:
```
if (elementsRead > trackMemoryThreshold &&
(elementsRead % 32 == 0 || currentMap.estimateSize() > myMemoryThreshold)
```
Otherwise say if the 33rd element was huge and we already went past our current memory threshold, we will have to wait for another 31 elements before we find a new threshold.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15508024
--- Diff: core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala ---
@@ -43,10 +44,10 @@ import org.apache.spark.{Logging, RangePartitioner}
*/
class OrderedRDDFunctions[K : Ordering : ClassTag,
V: ClassTag,
- P <: Product2[K, V] : ClassTag](
+ P <: Product2[K, V] : ClassTag] @DeveloperApi() (
--- End diff --
Yup, see the comments on https://github.com/mateiz/spark/commit/bf71388220fbe312afe7ef6b962afe6f96545dab.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15152216
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
--- End diff --
It appears that we'll either use the map or the buffer throughout the whole sorting process (rather than switch from one to another). In this case, would it make sense to put a `val usingMap/shouldCombine = aggregator.isDefined` out here, and then avoid passing it around elsewhere?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15208983
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
+
+ /** Shut down this ShuffleManager. */
+ override def stop(): Unit = {}
+
+ /** Get the location of a block in a map output file. Uses the index file we create for it. */
+ def getBlockLocation(blockId: ShuffleBlockId, diskManager: DiskBlockManager): FileSegment = {
+ // The block is actually going to be a range of a single map output file for this map, so
+ // so figure out the ID of the consolidated file, then the offset within that from our index
--- End diff --
extra "so"
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49961533
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds no public classes<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17078/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49924823
@colorant @mridulm for the serializing part, one issue is that as we merge streams we already need a heap of O(num streams) objects, so if you're worried about the objects being large, deferring the read won't help tremendously. I'm hoping to fix that by adding a max merge factor and supporting hierarchical merges later, but I think we can do it after this first PR. (If possible I would like to get it in 1.1 too, since it's not that much code. But I don't want to have a giant PR waiting that blocks other things.)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15544259
--- Diff: core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala ---
@@ -54,12 +55,16 @@ private[spark] class DiskBlockManager(shuffleManager: ShuffleBlockManager, rootD
addShutdownHook()
/**
- * Returns the physical file segment in which the given BlockId is located.
- * If the BlockId has been mapped to a specific FileSegment, that will be returned.
- * Otherwise, we assume the Block is mapped to a whole file identified by the BlockId directly.
+ * Returns the physical file segment in which the given BlockId is located. If the BlockId has
+ * been mapped to a specific FileSegment by the shuffle layer, that will be returned.
+ * Otherwise, we assume the Block is mapped to the whole file identified by the BlockId.
*/
def getBlockLocation(blockId: BlockId): FileSegment = {
- if (blockId.isShuffle && shuffleManager.consolidateShuffleFiles) {
+ val env = SparkEnv.get
+ if (blockId.isShuffle && env != null && env.shuffleManager.isInstanceOf[SortShuffleManager]) {
--- End diff --
add comment inline to explain env is null in tests
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562065
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
--- End diff --
I don't think this is true. We also use this if `ordering` is defined and `aggregator` is not.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15193928
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
--- End diff --
They should all have it
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562204
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in. These will be the indices of
+ // the next element we will read. We'll also store the last partition read so that
+ // readNextPartition() can figure out what partition that was from.
+ var partitionId = 0
+ var indexInPartition = 0L
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+ var lastPartitionId = 0
+
+ skipToNextPartition()
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ if (batchStreamsRead < spill.serializerBatchSizes.length) {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ } else {
+ // No more batches left; give an empty stream
+ bufferedStream
+ }
+ }
+
+ /**
+ * Update partitionId if we have reached the end of our current partition, possibly skipping
+ * empty partitions on the way.
+ */
+ private def skipToNextPartition() {
+ while (partitionId < numPartitions &&
+ indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0L
+ }
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ lastPartitionId = partitionId
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ skipToNextPartition()
+ // If we've finished reading the last partition, remember that we're done
+ if (partitionId == numPartitions) {
+ finished = true
+ deserStream.close()
+ }
+ (k, c)
+ }
+
+ var nextPartitionToRead = 0
+
+ def readNextPartition(): Iterator[Product2[K, C]] = new Iterator[Product2[K, C]] {
+ val myPartition = nextPartitionToRead
+ nextPartitionToRead += 1
+
+ override def hasNext: Boolean = {
+ if (nextItem == null) {
+ nextItem = readNextItem()
+ if (nextItem == null) {
+ return false
+ }
+ }
+ assert(lastPartitionId >= myPartition)
+ // Check that we're still in the right partition; note that readNextItem will have returned
+ // null at EOF above so we would've returned false there
+ lastPartitionId == myPartition
+ }
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val item = nextItem
+ nextItem = null
+ item
+ }
+ }
+ }
+
+ /**
+ * Return an iterator over all the data written to this object, grouped by partition and
+ * aggregated by the requested aggregator. For each partition we then have an iterator over its
+ * contents, and these are expected to be accessed in order (you can't "skip ahead" to one
+ * partition without reading the previous one). Guaranteed to return a key-value pair for each
+ * partition, in order of partition ID.
+ *
+ * For now, we just merge all the spilled files in once pass, but this can be modified to
+ * support hierarchical merging.
+ */
+ def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val usingMap = aggregator.isDefined
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ if (spills.isEmpty) {
+ // Special case: if we have only in-memory data, we don't need to merge streams, and perhaps
+ // we don't even need to sort by anything other than partition ID
+ if (!ordering.isDefined) {
+ // The user isn't requested sorted keys, so only sort by partition ID, not key
+ val partitionComparator = new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ groupByPartition(collection.destructiveSortedIterator(partitionComparator))
+ } else {
+ // We do need to sort by both partition ID and key
+ groupByPartition(collection.destructiveSortedIterator(partitionKeyComparator))
+ }
+ } else {
+ // General case: merge spilled and in-memory data
+ merge(spills, collection.destructiveSortedIterator(partitionKeyComparator))
+ }
+ }
+
+ /**
+ * Return an iterator over all the data written to this object, aggregated by our aggregator.
+ */
+ def iterator: Iterator[Product2[K, C]] = partitionedIterator.flatMap(pair => pair._2)
+
+ def stop(): Unit = {
+ spills.foreach(s => s.file.delete())
+ spills.clear()
+ }
+
+ def memoryBytesSpilled: Long = _memoryBytesSpilled
+
+ def diskBytesSpilled: Long = _diskBytesSpilled
+
+ /**
+ * Given a stream of ((partition, key), combiner) pairs *assumed to be sorted by partition ID*,
+ * group together the pairs for each for each partition into a sub-iterator.
--- End diff --
for each for each
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15194652
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
--- End diff --
That size is in bytes, not in objects. We could try to catch EOF but I've actually removed that, it seems cleaner with a count this way.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49938524
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds no public classes<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17055/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15620594
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,662 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ // Size of object batches when reading/writing from serializers.
+ //
+ // Objects are written in batches, with each batch using its own serialization stream. This
+ // cuts down on the size of reference-tracking maps constructed when deserializing a stream.
+ //
+ // NOTE: Setting this too low can cause excessive copying when serializing, since some serializers
+ // grow internal data structures by growing + copying every time the number of objects doubles.
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ private var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ private var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow aggregation or sorting.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator and no ordering are given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk if needed.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ // TODO: factor this out of both here and ExternalAppendOnlyMap
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle) [SPARK-2711]
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ val threadId = Thread.currentThread().getId
+ logInfo("Thread %d spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(threadId, memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ // Use the reverse of comparator.compare because PriorityQueue dequeues the max
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*) // Will contain only the iterators with hasNext = true
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering, so the objects with the same key are sequential.
+ new Iterator[Product2[K, C]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val elem = sorted.next()
+ val k = elem._1
+ var c = elem._2
+ while (sorted.hasNext && sorted.head._1 == k) {
+ c = mergeCombiners(c, sorted.head._2)
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in. These will be the indices of
+ // the next element we will read. We'll also store the last partition read so that
+ // readNextPartition() can figure out what partition that was from.
+ var partitionId = 0
+ var indexInPartition = 0L
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+ var lastPartitionId = 0
+
+ skipToNextPartition()
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ if (batchStreamsRead < spill.serializerBatchSizes.length) {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ } else {
+ // No more batches left; give an empty stream
+ bufferedStream
+ }
+ }
+
+ /**
+ * Update partitionId if we have reached the end of our current partition, possibly skipping
+ * empty partitions on the way.
+ */
+ private def skipToNextPartition() {
+ while (partitionId < numPartitions &&
+ indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0L
--- End diff --
Why don't we need to set `indexInBatch = 0` here as well? Is it because batches straddle partition boundaries (i.e. one batch can have elements from multiple partitions)?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15621714
--- Diff: core/src/test/scala/org/apache/spark/util/collection/ExternalSorterSuite.scala ---
@@ -0,0 +1,566 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.scalatest.FunSuite
+
+import org.apache.spark._
+import org.apache.spark.SparkContext._
+
+class ExternalSorterSuite extends FunSuite with LocalSparkContext {
+ test("empty data stream") {
+ val conf = new SparkConf(false)
+ conf.set("spark.shuffle.memoryFraction", "0.001")
+ conf.set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.SortShuffleManager")
+ sc = new SparkContext("local", "test", conf)
+
+ val agg = new Aggregator[Int, Int, Int](i => i, (i, j) => i + j, (i, j) => i + j)
+ val ord = implicitly[Ordering[Int]]
+
+ // Both aggregator and ordering
+ val sorter = new ExternalSorter[Int, Int, Int](
+ Some(agg), Some(new HashPartitioner(3)), Some(ord), None)
+ assert(sorter.iterator.toSeq === Seq())
+ sorter.stop()
+
+ // Only aggregator
+ val sorter2 = new ExternalSorter[Int, Int, Int](
+ Some(agg), Some(new HashPartitioner(3)), None, None)
+ assert(sorter2.iterator.toSeq === Seq())
+ sorter2.stop()
+
+ // Only ordering
+ val sorter3 = new ExternalSorter[Int, Int, Int](
+ None, Some(new HashPartitioner(3)), Some(ord), None)
+ assert(sorter3.iterator.toSeq === Seq())
+ sorter3.stop()
+
+ // Neither aggregator nor ordering
+ val sorter4 = new ExternalSorter[Int, Int, Int](
+ None, Some(new HashPartitioner(3)), None, None)
+ assert(sorter4.iterator.toSeq === Seq())
+ sorter4.stop()
+ }
+
+ test("few elements per partition") {
+ val conf = new SparkConf(false)
+ conf.set("spark.shuffle.memoryFraction", "0.001")
+ conf.set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.SortShuffleManager")
+ sc = new SparkContext("local", "test", conf)
+
+ val agg = new Aggregator[Int, Int, Int](i => i, (i, j) => i + j, (i, j) => i + j)
+ val ord = implicitly[Ordering[Int]]
+ val elements = Set((1, 1), (2, 2), (5, 5))
+ val expected = Set(
+ (0, Set()), (1, Set((1, 1))), (2, Set((2, 2))), (3, Set()), (4, Set()),
+ (5, Set((5, 5))), (6, Set()))
+
+ // Both aggregator and ordering
+ val sorter = new ExternalSorter[Int, Int, Int](
+ Some(agg), Some(new HashPartitioner(7)), Some(ord), None)
+ sorter.write(elements.iterator)
+ assert(sorter.partitionedIterator.map(p => (p._1, p._2.toSet)).toSet === expected)
+ sorter.stop()
+
+ // Only aggregator
+ val sorter2 = new ExternalSorter[Int, Int, Int](
+ Some(agg), Some(new HashPartitioner(7)), None, None)
+ sorter2.write(elements.iterator)
+ assert(sorter2.partitionedIterator.map(p => (p._1, p._2.toSet)).toSet === expected)
+ sorter2.stop()
+
+ // Only ordering
+ val sorter3 = new ExternalSorter[Int, Int, Int](
+ None, Some(new HashPartitioner(7)), Some(ord), None)
+ sorter3.write(elements.iterator)
+ assert(sorter3.partitionedIterator.map(p => (p._1, p._2.toSet)).toSet === expected)
+ sorter3.stop()
+
+ // Neither aggregator nor ordering
+ val sorter4 = new ExternalSorter[Int, Int, Int](
+ None, Some(new HashPartitioner(7)), None, None)
+ sorter4.write(elements.iterator)
+ assert(sorter4.partitionedIterator.map(p => (p._1, p._2.toSet)).toSet === expected)
+ sorter4.stop()
+ }
+
+ test("empty partitions with spilling") {
+ val conf = new SparkConf(false)
+ conf.set("spark.shuffle.memoryFraction", "0.001")
+ conf.set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.SortShuffleManager")
+ sc = new SparkContext("local", "test", conf)
+
+ val agg = new Aggregator[Int, Int, Int](i => i, (i, j) => i + j, (i, j) => i + j)
+ val ord = implicitly[Ordering[Int]]
+ val elements = Iterator((1, 1), (5, 5)) ++ (0 until 100000).iterator.map(x => (2, 2))
+
+ val sorter = new ExternalSorter[Int, Int, Int](
+ None, Some(new HashPartitioner(7)), None, None)
+ sorter.write(elements)
+ assert(sc.env.blockManager.diskBlockManager.getAllFiles().length > 0) // Make sure it spilled
+ val iter = sorter.partitionedIterator.map(p => (p._1, p._2.toList))
+ assert(iter.next() === (0, Nil))
+ assert(iter.next() === (1, List((1, 1))))
+ assert(iter.next() === (2, (0 until 100000).map(x => (2, 2)).toList))
+ assert(iter.next() === (3, Nil))
+ assert(iter.next() === (4, Nil))
+ assert(iter.next() === (5, List((5, 5))))
+ assert(iter.next() === (6, Nil))
+ sorter.stop()
+ }
+
+ test("spilling in local cluster") {
+ val conf = new SparkConf(true) // Load defaults, otherwise SPARK_HOME is not found
+ conf.set("spark.shuffle.memoryFraction", "0.001")
+ conf.set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.SortShuffleManager")
+ sc = new SparkContext("local-cluster[1,1,512]", "test", conf)
+
+ // reduceByKey - should spill ~8 times
+ val rddA = sc.parallelize(0 until 100000).map(i => (i/2, i))
+ val resultA = rddA.reduceByKey(math.max).collect()
+ assert(resultA.length == 50000)
+ resultA.foreach { case(k, v) =>
+ if (v != k * 2 + 1) {
+ fail(s"Value for ${k} was wrong: expected ${k * 2 + 1}, got ${v}")
+ }
+ }
+
+ // groupByKey - should spill ~17 times
+ val rddB = sc.parallelize(0 until 100000).map(i => (i/4, i))
+ val resultB = rddB.groupByKey().collect()
+ assert(resultB.length == 25000)
+ resultB.foreach { case(i, seq) =>
+ val expected = Set(i * 4, i * 4 + 1, i * 4 + 2, i * 4 + 3)
+ if (seq.toSet != expected) {
+ fail(s"Value for ${i} was wrong: expected ${expected}, got ${seq.toSet}")
+ }
+ }
+
+ // cogroup - should spill ~7 times
+ val rddC1 = sc.parallelize(0 until 10000).map(i => (i, i))
+ val rddC2 = sc.parallelize(0 until 10000).map(i => (i%1000, i))
+ val resultC = rddC1.cogroup(rddC2).collect()
+ assert(resultC.length == 10000)
+ resultC.foreach { case(i, (seq1, seq2)) =>
+ i match {
+ case 0 =>
+ assert(seq1.toSet == Set[Int](0))
+ assert(seq2.toSet == Set[Int](0, 1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000))
+ case 1 =>
+ assert(seq1.toSet == Set[Int](1))
+ assert(seq2.toSet == Set[Int](1, 1001, 2001, 3001, 4001, 5001, 6001, 7001, 8001, 9001))
+ case 5000 =>
+ assert(seq1.toSet == Set[Int](5000))
+ assert(seq2.toSet == Set[Int]())
+ case 9999 =>
+ assert(seq1.toSet == Set[Int](9999))
+ assert(seq2.toSet == Set[Int]())
+ case _ =>
+ }
+ }
+
+ // larger cogroup - should spill ~7 times
--- End diff --
strange that larger cogroup spills the same number of times than the previous cogroup
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49541073
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16864/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50522967
@mateiz please refer to changes here : https://github.com/apache/spark/pull/1609/files#diff-10
They should be relevant to this PR too
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15192894
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
--- End diff --
Good catch, fixed this and added a test
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15195282
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
--- End diff --
Yeah, this is confusing. We actually start indexInPartition at -1, so these point to the last read element. But the check for termination in hasNext was wrong (we just figure out we're at EOF due to the return null below), so that made it seem like this would be the next ID.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15494358
--- Diff: core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala ---
@@ -54,12 +55,16 @@ private[spark] class DiskBlockManager(shuffleManager: ShuffleBlockManager, rootD
addShutdownHook()
/**
- * Returns the physical file segment in which the given BlockId is located.
- * If the BlockId has been mapped to a specific FileSegment, that will be returned.
- * Otherwise, we assume the Block is mapped to a whole file identified by the BlockId directly.
+ * Returns the physical file segment in which the given BlockId is located. If the BlockId has
+ * been mapped to a specific FileSegment by the shuffle layer, that will be returned.
+ * Otherwise, we assume the Block is mapped to the whole file identified by the BlockId.
*/
def getBlockLocation(blockId: BlockId): FileSegment = {
--- End diff --
In this block of code, there's two variables named `shuffleManager`: the DickBlockManager's ShuffleBlockManager, and SparkEnv's ShuffleManager. I think it would help readability if we renamed DiskBlockManager's `shuffleManager` to `shuffleBlockManager`.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15148086
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // For now, just compare them by partition; later we can compare by key as well
+ private val comparator = new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: ArrayBuffer[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ // TODO: this could become a sparser data structure
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(comparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ */
+ def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[(K, C)] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): (K, C) = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ (p, readers.iterator.flatMap(_.readNextPartition()) ++ inMemIterator)
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
--- End diff --
i have no proof, but i think marking these as private[this] might help with perf. It is good to lock the visibility down anyway.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15147904
--- Diff: core/src/main/scala/org/apache/spark/util/collection/SizeTrackingBuffer.scala ---
@@ -0,0 +1,150 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.util.Arrays
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.util.SizeEstimator
+
+/**
+ * Append-only buffer that keeps track of its estimated size in bytes.
+ * We sample with a slow exponential back-off using the SizeEstimator to amortize the time,
+ * as each call to SizeEstimator can take a sizable amount of time (order of a few milliseconds).
+ *
+ * The tracking code is copied from SizeTrackingAppendOnlyMap -- we'll factor that out soon.
+ */
+private[spark] class SizeTrackingBuffer[T <: AnyRef](initialCapacity: Int = 64)
+ extends SizeTrackingCollection[T]
--- End diff --
This could be replaced with the buffer in #1165 except we also need destructiveSortedIterator. And to sort stuff in-place, that in turn requires T to be a subclass of AnyRef (otherwise we'd have to pass the Array[T] to a different Arrays.sort() method based on its type).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15556374
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
+ if (dep.mapSideCombine) {
+ if (!dep.aggregator.isDefined) {
+ throw new IllegalStateException("Aggregator is empty for map-side combine")
+ }
+ sorter = new ExternalSorter[K, V, C](
+ dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ } else {
+ // In this case we pass neither an aggregator nor an ordering to the sorter, because we
+ // don't care whether the keys get sorted in each partition; that will be done on the
+ // reduce side if the operation being run is sortByKey.
+ sorter = new ExternalSorter[K, V, V](
+ None, Some(dep.partitioner), None, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ }
+ }
+
+ // Create a single shuffle file with reduce ID 0 that we'll write all results to. We'll later
+ // serve different ranges of this file using an index file that we create at the end.
+ val blockId = ShuffleBlockId(dep.shuffleId, mapId, 0)
+ outputFile = blockManager.diskBlockManager.getFile(blockId)
+
+ // Track location of each range in the output file
+ val offsets = new Array[Long](numPartitions + 1)
+ val lengths = new Array[Long](numPartitions)
+
+ // Statistics
+ var totalBytes = 0L
+ var totalTime = 0L
+
+ for ((id, elements) <- partitions) {
+ if (elements.hasNext) {
+ val writer = blockManager.getDiskWriter(blockId, outputFile, ser, fileBufferSize)
+ for (elem <- elements) {
+ writer.write(elem)
+ }
+ writer.commit()
+ writer.close()
+ val segment = writer.fileSegment()
+ offsets(id + 1) = segment.offset + segment.length
+ lengths(id) = segment.length
+ totalTime += writer.timeWriting()
+ totalBytes += segment.length
+ } else {
+ // Don't create a new writer to avoid writing any headers and things like that
+ offsets(id + 1) = offsets(id)
+ }
+ }
+
+ val shuffleMetrics = new ShuffleWriteMetrics
+ shuffleMetrics.shuffleBytesWritten = totalBytes
+ shuffleMetrics.shuffleWriteTime = totalTime
+ context.taskMetrics.shuffleWriteMetrics = Some(shuffleMetrics)
+ context.taskMetrics.memoryBytesSpilled = sorter.memoryBytesSpilled
+ context.taskMetrics.diskBytesSpilled = sorter.diskBytesSpilled
--- End diff --
What if we spilled earlier in this task (e.g. while reducing results from the previous stage), won't this overwrite the bytes spilled there?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15148079
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // For now, just compare them by partition; later we can compare by key as well
+ private val comparator = new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: ArrayBuffer[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ // TODO: this could become a sparser data structure
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(comparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ */
+ def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
--- End diff --
I'd add @return to the javadoc to explain what the int represents.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154485
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
--- End diff --
maybe append ", possibly skipping zero-length partitions" or something
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by markhamstra <gi...@git.apache.org>.
Github user markhamstra commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50160338
After installing `hub` you can also do a bunch of new stuff on the command line, including `hub checkout https://github.com/apache/spark/pull/1499`
https://hub.github.com/
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49811607
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds no public classes<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16990/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15157645
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
--- End diff --
True, that would be better here; I'll try it
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50301634
@mateiz , yep, the map tasks did spill and it seems contribute most to the increased process time. though in my case only about 400K data been spilled to disk per task. But it lead to 15 seconds more time for process. Say the task time increased from 15 seconds to 30 seconds
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154163
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
--- End diff --
Is this TODO still true?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50296194
I observe that a groupby case is much slower ( several times? ) with this manager than the hash based one. not sure is there anything need to tune, but the data set I run is pretty small so should not have much thing to tune.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15277549
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading, possibly skipping zero-length
+ // partitions until we get to the next non-empty one or to EOF.
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
--- End diff --
Or maybe, to improve performance, we should do as mridulm suggested, set indexInPartition to 0, then in additition, do partition skip before actually read k,c object. thus, the caller will get correct partitionID for hasnext judgements for the last item. And then, this if judgement could be moved in to the while loop, and set the finish flag, and do close stream in the begining of next readNextItem iteration? so that this if judgement won't need to be done upon each item but each partition?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50116492
ah, thanks ! rerunning with 9c29957.
cant pull the pr - and manual merge is painful, hence delays in testing :-)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49568383
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16886/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15274240
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
--- End diff --
This causes a bug in the last partition offset by 1 when there are lot of empty partitions - particularly when first partition is empty.
We need to essentially initialize it to 0; and then apply same logic as what we have within readNextItem()
Something like :
```
private var indexInPartition = 0L // Just to make sure we start at index 0
skipToNextPartition()
private def skipToNextPartition() {
while (indexInPartition == spill.elementsPerPartition(partitionId)) {
partitionId += 1
indexInPartition = 0
}
}
```
; and in readNextItem :
```
indexInPartition += 1
skipToNextPartition()
```
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154289
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
--- End diff --
Is there a particular reason you chose to implement hasNext checks in certain iterators, but not others?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49558215
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16878/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15555693
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
--- End diff --
this should probably be info instead of warning level
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15277089
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading, possibly skipping zero-length
+ // partitions until we get to the next non-empty one or to EOF.
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
--- End diff --
maybe here should be indexInPartition >= spill.elementsPerPartition(partitionId) - 1, in case the last two partition is all 0 items, thus indexInPartition will be 0, and spill.elementsPerPartition(partitionId) - 1 will be -1
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49539437
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16864/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15620640
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,662 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ // Size of object batches when reading/writing from serializers.
+ //
+ // Objects are written in batches, with each batch using its own serialization stream. This
+ // cuts down on the size of reference-tracking maps constructed when deserializing a stream.
+ //
+ // NOTE: Setting this too low can cause excessive copying when serializing, since some serializers
+ // grow internal data structures by growing + copying every time the number of objects doubles.
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ private var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ private var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow aggregation or sorting.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator and no ordering are given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk if needed.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ // TODO: factor this out of both here and ExternalAppendOnlyMap
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle) [SPARK-2711]
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ val threadId = Thread.currentThread().getId
+ logInfo("Thread %d spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(threadId, memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ // Use the reverse of comparator.compare because PriorityQueue dequeues the max
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*) // Will contain only the iterators with hasNext = true
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering, so the objects with the same key are sequential.
+ new Iterator[Product2[K, C]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val elem = sorted.next()
+ val k = elem._1
+ var c = elem._2
+ while (sorted.hasNext && sorted.head._1 == k) {
+ c = mergeCombiners(c, sorted.head._2)
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in. These will be the indices of
+ // the next element we will read. We'll also store the last partition read so that
+ // readNextPartition() can figure out what partition that was from.
+ var partitionId = 0
+ var indexInPartition = 0L
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+ var lastPartitionId = 0
+
+ skipToNextPartition()
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ if (batchStreamsRead < spill.serializerBatchSizes.length) {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ } else {
+ // No more batches left; give an empty stream
+ bufferedStream
+ }
+ }
+
+ /**
+ * Update partitionId if we have reached the end of our current partition, possibly skipping
+ * empty partitions on the way.
+ */
+ private def skipToNextPartition() {
+ while (partitionId < numPartitions &&
+ indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0L
--- End diff --
I see, the answer is yes because when we spill we iterate through ((partition ID, k), v), not just (k, v) per partition.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49788955
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16983/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15544043
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
+ if (dep.mapSideCombine) {
+ if (!dep.aggregator.isDefined) {
+ throw new IllegalStateException("Aggregator is empty for map-side combine")
+ }
+ sorter = new ExternalSorter[K, V, C](
+ dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ } else {
+ // In this case we pass neither an aggregator nor an ordering to the sorter, because we
+ // don't care whether the keys get sorted in each partition; that will be done on the
+ // reduce side if the operation being run is sortByKey.
+ sorter = new ExternalSorter[K, V, V](
+ None, Some(dep.partitioner), None, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ }
+ }
+
+ // Create a single shuffle file with reduce ID 0 that we'll write all results to. We'll later
+ // serve different ranges of this file using an index file that we create at the end.
+ val blockId = ShuffleBlockId(dep.shuffleId, mapId, 0)
+ outputFile = blockManager.diskBlockManager.getFile(blockId)
+
+ // Track location of each range in the output file
+ val offsets = new Array[Long](numPartitions + 1)
+ val lengths = new Array[Long](numPartitions)
+
+ // Statistics
+ var totalBytes = 0L
+ var totalTime = 0L
+
+ for ((id, elements) <- partitions) {
+ if (elements.hasNext) {
+ val writer = blockManager.getDiskWriter(blockId, outputFile, ser, fileBufferSize)
+ for (elem <- elements) {
+ writer.write(elem)
+ }
+ writer.commit()
+ writer.close()
+ val segment = writer.fileSegment()
+ offsets(id + 1) = segment.offset + segment.length
+ lengths(id) = segment.length
+ totalTime += writer.timeWriting()
+ totalBytes += segment.length
+ } else {
+ // Don't create a new writer to avoid writing any headers and things like that
--- End diff --
I'd update the comment to also say explicitly this partition doesn't have any data
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50297730
Are map tasks spilling by any chance? There is one issue in this right now, which is that if your map task spills to disk, you need to spill multiple times with the sort-based shuffle, whereas with the old code you would directly open one file for each reduce task and write data to it. I think we can add the same behavior here if the # of reduce tasks is small enough (the problem before was that you'd open 10,000 files for 10,000 reduce tasks and allocate lots of space for buffers, not to mention do lots of random I/O). I would do that in a separate pull request though. For other things, like reduceByKey, this should have similar performance to the hash-based shuffle.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50115453
BTW, this is one of 5 failures from core.
I hope there are no merge issues though,
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560798
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
--- End diff --
It might be good to declare this on instantiation rather than in multiple places
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by lianhuiwang <gi...@git.apache.org>.
Github user lianhuiwang commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15352745
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
+ if (dep.mapSideCombine) {
+ if (!dep.aggregator.isDefined) {
+ throw new IllegalStateException("Aggregator is empty for map-side combine")
+ }
+ sorter = new ExternalSorter[K, V, C](
+ dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ } else {
+ // In this case we pass neither an aggregator nor an ordering to the sorter, because we
+ // don't care whether the keys get sorted in each partition; that will be done on the
+ // reduce side if the operation being run is sortByKey.
+ sorter = new ExternalSorter[K, V, V](
+ None, Some(dep.partitioner), None, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ }
+ }
+
+ // Create a single shuffle file with reduce ID 0 that we'll write all results to. We'll later
+ // serve different ranges of this file using an index file that we create at the end.
+ val blockId = ShuffleBlockId(dep.shuffleId, mapId, 0)
+ outputFile = blockManager.diskBlockManager.getFile(blockId)
+
+ // Track location of each range in the output file
+ val offsets = new Array[Long](numPartitions + 1)
+ val lengths = new Array[Long](numPartitions)
+
+ // Statistics
+ var totalBytes = 0L
+ var totalTime = 0L
+
+ for ((id, elements) <- partitions) {
+ if (elements.hasNext) {
+ val writer = blockManager.getDiskWriter(blockId, outputFile, ser, fileBufferSize)
+ for (elem <- elements) {
+ writer.write(elem)
+ }
+ writer.commit()
+ writer.close()
+ val segment = writer.fileSegment()
+ offsets(id + 1) = segment.offset + segment.length
+ lengths(id) = segment.length
+ totalTime += writer.timeWriting()
+ totalBytes += segment.length
+ } else {
+ // Don't create a new writer to avoid writing any headers and things like that
+ offsets(id + 1) = offsets(id)
+ }
+ }
+
+ val shuffleMetrics = new ShuffleWriteMetrics
+ shuffleMetrics.shuffleBytesWritten = totalBytes
+ shuffleMetrics.shuffleWriteTime = totalTime
+ context.taskMetrics.shuffleWriteMetrics = Some(shuffleMetrics)
+ context.taskMetrics.memoryBytesSpilled = sorter.memoryBytesSpilled
+ context.taskMetrics.diskBytesSpilled = sorter.diskBytesSpilled
+
+ // Write an index file with the offsets of each block, plus a final offset at the end for the
+ // end of the output file. This will be used by SortShuffleManager.getBlockLocation to figure
+ // out where each block begins and ends.
+
+ val diskBlockManager = blockManager.diskBlockManager
+ val indexFile = diskBlockManager.getFile(blockId.name + ".index")
+ val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexFile)))
+ try {
+ var i = 0
+ while (i < numPartitions + 1) {
+ out.writeLong(offsets(i))
+ i += 1
+ }
+ } finally {
+ out.close()
+ }
+
+ mapStatus = new MapStatus(blockManager.blockManagerId,
+ lengths.map(MapOutputTracker.compressSize))
--- End diff --
maybe can create new function instead of MapOutputTracker.compressSize? so here it do not refer to MapOutputTracker.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49707121
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16948/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15209013
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.util.collection.ExternalSorter
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.executor.ShuffleWriteMetrics
--- End diff --
reorder import list?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15333148
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
--- End diff --
Maybe I miss some code? But from what I understanding, when the cleaning is turned on, it seems to me, if by timestamp approaching, blockManager won't remove shuffle data,ShuffleBlockManager will do the work by itself. And if by auto clean approaching when doCleanShuffle is called, it go through the current ShuffleBlockManager interface. In neither case, these sortShuffleWritter generated file will be cleaned? Or do you mean, as long as there are no metadata in memory, shuffle file on disk is ok to not be removed until application exit?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15497588
--- Diff: core/src/test/scala/org/apache/spark/util/collection/ExternalSorterSuite.scala ---
@@ -0,0 +1,566 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.scalatest.FunSuite
+
+import org.apache.spark._
+import org.apache.spark.SparkContext._
+
+class ExternalSorterSuite extends FunSuite with LocalSparkContext {
--- End diff --
There's a number of tests cases in this file that only differ in a few parameters (e.g. the partial aggregation tests). These might be a candidate for ScalaTests's [table-driven property checks](http://www.scalatest.org/user_guide/table_driven_property_checks). Unfortunately, ScalaTest's `forAll(table)` short-circuits and exits on the first error rather than reporting all failing rows (there's an open issue to fix this: https://github.com/scalatest/scalatest/issues/178). In the meantime, there's still opportunities to factor apart the assertions and parameters.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by asfgit <gi...@git.apache.org>.
Github user asfgit closed the pull request at:
https://github.com/apache/spark/pull/1499
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49802023
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds no public classes<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16983/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15383498
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
+ if (dep.mapSideCombine) {
+ if (!dep.aggregator.isDefined) {
+ throw new IllegalStateException("Aggregator is empty for map-side combine")
+ }
+ sorter = new ExternalSorter[K, V, C](
+ dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ } else {
+ // In this case we pass neither an aggregator nor an ordering to the sorter, because we
+ // don't care whether the keys get sorted in each partition; that will be done on the
+ // reduce side if the operation being run is sortByKey.
+ sorter = new ExternalSorter[K, V, V](
+ None, Some(dep.partitioner), None, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ }
+ }
+
+ // Create a single shuffle file with reduce ID 0 that we'll write all results to. We'll later
+ // serve different ranges of this file using an index file that we create at the end.
+ val blockId = ShuffleBlockId(dep.shuffleId, mapId, 0)
+ outputFile = blockManager.diskBlockManager.getFile(blockId)
+
+ // Track location of each range in the output file
+ val offsets = new Array[Long](numPartitions + 1)
+ val lengths = new Array[Long](numPartitions)
+
+ // Statistics
+ var totalBytes = 0L
+ var totalTime = 0L
+
+ for ((id, elements) <- partitions) {
+ if (elements.hasNext) {
+ val writer = blockManager.getDiskWriter(blockId, outputFile, ser, fileBufferSize)
+ for (elem <- elements) {
+ writer.write(elem)
+ }
+ writer.commit()
+ writer.close()
+ val segment = writer.fileSegment()
+ offsets(id + 1) = segment.offset + segment.length
+ lengths(id) = segment.length
+ totalTime += writer.timeWriting()
+ totalBytes += segment.length
+ } else {
+ // Don't create a new writer to avoid writing any headers and things like that
+ offsets(id + 1) = offsets(id)
+ }
+ }
+
+ val shuffleMetrics = new ShuffleWriteMetrics
+ shuffleMetrics.shuffleBytesWritten = totalBytes
+ shuffleMetrics.shuffleWriteTime = totalTime
+ context.taskMetrics.shuffleWriteMetrics = Some(shuffleMetrics)
+ context.taskMetrics.memoryBytesSpilled = sorter.memoryBytesSpilled
+ context.taskMetrics.diskBytesSpilled = sorter.diskBytesSpilled
+
+ // Write an index file with the offsets of each block, plus a final offset at the end for the
+ // end of the output file. This will be used by SortShuffleManager.getBlockLocation to figure
+ // out where each block begins and ends.
+
+ val diskBlockManager = blockManager.diskBlockManager
+ val indexFile = diskBlockManager.getFile(blockId.name + ".index")
+ val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexFile)))
+ try {
+ var i = 0
+ while (i < numPartitions + 1) {
+ out.writeLong(offsets(i))
+ i += 1
+ }
+ } finally {
+ out.close()
+ }
+
+ mapStatus = new MapStatus(blockManager.blockManagerId,
+ lengths.map(MapOutputTracker.compressSize))
--- End diff --
This can be refactored later, there's a task to move MapOutputTracker around
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15555382
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
--- End diff --
you could do `orNull`
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50701304
I took another pass over the patch and the changes look ready to me. I also tested this locally and verified that the shuffle files were actually cleaned up. There is still a lot of very similar code in the ExternalSorter and the EAOM, though this is not of 1.1 priority to resolve. Instead, I have filed a [JIRA](https://issues.apache.org/jira/browse/SPARK-2761) for it.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15288486
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
--- End diff --
BTW, some obvious index offset issues to be guarded against in my proposal above :)
Additionally, can we defer deserializing the value until it is actually needed ?
This will bring down the actual number of deserialized objects to what is actually required (as opposed to required + num_inputs).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15382131
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
--- End diff --
Ah, I see, these files are never registered with the BlockManager so they won't be cleaned by it. I'll modify ShuffleBlockManager to clean them for now.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15331341
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
--- End diff --
shuffle output file in sortShuffleWritter do not get cleaned. We might need to add map to save registered shuffle handle, Then try to remove the data file in unregisterShuffle method. Though at present, in HashShuffleManager, this is also not implemented. But HashShuffleManager depends on shuffleBlockManager and the file will be cleaned there.
I have a PR to generalize shuffleBlockManager and hide it behind shuffleMananger. as in #1241 , and upon blockMananger do remove shuffle, will call into this unregisterShuffle method. Will rebase upon this PR been merged.
Should we fix this issue in my PR, or add the store/clean shuffleHandle logic here in this PR firstly?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49949660
No, we only read one value at a time. If you're doing something like groupByKey, the combine function may create an ArrayBuffer or stuff like that, but that's only for the key actively being read. As soon as one stream moves on to the next key, we read the first key-value pair in it but stop reading others.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50694691
Jenkins, retest this please.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49804353
We saw a bunch of EOF Exceptions from SpillReader.
java.io.EOFException
at java.io.ObjectInputStream$BlockDataInputStream.peekByte(ObjectInputStream.java:2577)
at java.io.ObjectInputStream.readObject0(ObjectInputStream.java:1315)
at java.io.ObjectInputStream.readObject(ObjectInputStream.java:369)
at org.apache.spark.serializer.JavaDeserializationStream.readObject(JavaSerializer.scala:79)
at org.apache.spark.util.collection.ExternalSorter$SpillReader.org$apache$spark$util$collection$ExternalSorter$SpillReader$$readNextItem(ExternalSorter.scala:526)
at org.apache.spark.util.collection.ExternalSorter$SpillReader$$anon$6.hasNext(ExternalSorter.scala:560)
at scala.collection.TraversableOnce$FlattenOps$$anon$1.hasNext(TraversableOnce.scala:432)
at scala.collection.Iterator$class.foreach(Iterator.scala:727)
at scala.collection.AbstractIterator.foreach(Iterator.scala:1157)
at org.apache.spark.shuffle.sort.SortShuffleWriter$$anonfun$write$2.apply(SortShuffleWriter.scala:93)
at org.apache.spark.shuffle.sort.SortShuffleWriter$$anonfun$write$2.apply(SortShuffleWriter.scala:89)
(Exact line might not match master, but the issue might be conveyed I hope)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154387
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
super nit: can just do the `+=` afterwards to avoid the `- 1`. I think this also makes logically an eensy bit more sense.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154433
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
+ // This is the last element, remember that we're done
+ finished = true
+ deserStream.close()
+ }
+ (k, c)
+ } catch {
+ case e: EOFException =>
+ finished = true
+ deserStream.close()
+ null
+ }
+ }
+
+ var nextPartitionToRead = 0
+
+ def readNextPartition(): Iterator[Product2[K, C]] = new Iterator[Product2[K, C]] {
+ val myPartition = nextPartitionToRead
+ nextPartitionToRead += 1
+
+ override def hasNext: Boolean = {
+ if (nextItem == null) {
+ nextItem = readNextItem()
+ if (nextItem == null) {
+ return false
+ }
+ }
+ // Check that we're still in the right partition; will be numPartitions at EOF
--- End diff --
grammar nit; I got confused by the comment and thought this check was only for EOF. Maybe just reorder second clause like "at EOF, partitionId will be numPartitions" or something.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49800606
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16990/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15558266
--- Diff: core/src/main/scala/org/apache/spark/storage/BlockId.scala ---
@@ -60,6 +60,12 @@ case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int)
}
@DeveloperApi
+case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int)
+ extends BlockId {
--- End diff --
nit: could you bump this one line
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560508
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
--- End diff --
This could just be private, unless you had a special reason for this
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560817
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
--- End diff --
Add SPARK-2711 to comment?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49565532
I've now updated this to support partial aggregation across spilled files and even if we don't have an Ordering, using hash code comparison similar to ExternalAppendOnlyMap. It also now fully implements the behavior in the docs, namely sorting the data if you pass an Ordering, etc.
It looks like Aaron found a problem with the size-tracking code -- will try to fix that in ExternalAppendOnlyMap as well. Once this class is in though it could replace ExternalAppendOnlyMap in most use cases, though its one downside is that it creates another object for each key of the in-memory collection (since we have `((Int, K), C)` pairs to allow sorting by partition). Might still be worth it long-term though. On the flip side, I think the hash based merging code here is more efficient, avoiding a bunch of `ArrayBuffer.remove` calls.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15388684
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
--- End diff --
Alright, fixed this in the latest commit.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15310488
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
--- End diff --
Good catch, thanks. I've updated it for this and added a test that would fail the previous one.
Not sure what you mean by deserializing the value here -- it seems like it will save at most one element.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15501936
--- Diff: core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala ---
@@ -43,10 +44,10 @@ import org.apache.spark.{Logging, RangePartitioner}
*/
class OrderedRDDFunctions[K : Ordering : ClassTag,
V: ClassTag,
- P <: Product2[K, V] : ClassTag](
+ P <: Product2[K, V] : ClassTag] @DeveloperApi() (
--- End diff --
Did you mean to add this annotation here?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15495220
--- Diff: core/src/main/scala/org/apache/spark/storage/ShuffleBlockManager.scala ---
@@ -91,6 +97,20 @@ class ShuffleBlockManager(blockManager: BlockManager) extends Logging {
private val metadataCleaner =
new MetadataCleaner(MetadataCleanerType.SHUFFLE_BLOCK_MANAGER, this.cleanup, conf)
+ /**
+ * Register a completed map without getting a ShuffleWriterGroup. Used by sort-based shuffle
+ * because it just writes a single file by itself.
+ */
+ def addCompletedMap(shuffleId: Int, mapId: Int, numBuckets: Int): Unit = {
+ shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numBuckets))
--- End diff --
In both TimeStampedHashMap and scala.collection.concurrent.Map, `putIfAbsent`'s value is call-by-value, not call-by-name:
```
def putIfAbsent(key: A, value: B): Option[B]
```
As a result, this code will end up calling `new ShuffleState(numBuckets)` on every completed map task. ShuffleState looks moderately expensive to construct since it contains a bunch of queues, etc.
It would be more efficient to extend TimeStampedHashMap with a `getOrElseUpdate` method whose `default` argument is call-by-name so that we don't end up constructing and discarding so many ShuffleStates.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49706936
@aarondav pushed an update to the grow code as well, which will now use estimateSize. I implemented the same thing in EAOM.
@colorant fixed those, thanks.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15604053
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
--- End diff --
I'd rather do this by refactoring the memory tracking in a later PR: https://issues.apache.org/jira/browse/SPARK-2711
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49973755
Added one more commit that fixes the type of ShuffledRDD, because in this new shuffle it's not possible to return a custom Product2 the way it's written now, and in the old one it wasn't possible anyway with aggregation. See https://github.com/mateiz/spark/commit/9c299579f13f004f5fd1f4dd0b98b7d76cac2a55 for details. @rxin you should look at this one.
With this commit there's now a version of ShuffleSuite that passes using SortBasedShuffle.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49565536
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16886/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15196170
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
--- End diff --
I'll add some tests with very few elements per partition
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15152201
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
--- End diff --
it appears that you could just do `Serializer.getSerializer(serializer)`, thanks to code you added last month :)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49949874
Let me know if I've missed something. That combining part happens in mergeWithAggregation. Actually for many operations map-side combine is disabled so we don't even do it (and ExternalSorter is not used on the reduce side right now).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49561754
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16882/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15152282
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // For now, just compare them by partition; later we can compare by key as well
+ private val comparator = new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: ArrayBuffer[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
--- End diff --
Doesn't this "collection.atGrowThreshold" thing assume that we're not just combining into a few large keys? collection.estimateSize() is very fast, and is intended to be possible to use it in a per-element manner.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15153152
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // For now, just compare them by partition; later we can compare by key as well
+ private val comparator = new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: ArrayBuffer[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
--- End diff --
I just copied this from ExternalAppendOnlyMap, but it's true that it's probably wrong. I guess the reason it was there is that there's a big non-linear jump in size when the collection doubles. However, with the current code, we may not notice if the values are really large. It's probably better to use estimateSize and leave enough buffer space that doubling just the array won't be a problem.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154182
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
--- End diff --
Did you prefer this style or the one where you indent the colon 4 spaces in?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50701589
Ok I merged this in master.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by colorant <gi...@git.apache.org>.
Github user colorant commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15276601
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
--- End diff --
Seems to me if modified so, the last item in each partition will be considered as belong to next partition? And for the last partition, the offset is already guarded by -1 in the termination judgement?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15310561
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,649 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
--- End diff --
(Well, one element per stream we're merging. I guess you're worried about overly large items? I was hoping to limit the merge factor instead and do hierarchical merges, since we also store these items in a heap and such later on, so there are still O(num mergers) of them.)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15555675
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
--- End diff --
inline param to explain what usingMap does (just copy from the next function)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15557939
--- Diff: core/src/main/scala/org/apache/spark/storage/ShuffleBlockManager.scala ---
@@ -91,6 +97,20 @@ class ShuffleBlockManager(blockManager: BlockManager) extends Logging {
private val metadataCleaner =
new MetadataCleaner(MetadataCleanerType.SHUFFLE_BLOCK_MANAGER, this.cleanup, conf)
+ /**
+ * Register a completed map without getting a ShuffleWriterGroup. Used by sort-based shuffle
+ * because it just writes a single file by itself.
+ */
+ def addCompletedMap(shuffleId: Int, mapId: Int, numBuckets: Int): Unit = {
+ shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numBuckets))
--- End diff --
Looks like we currently already do that in `forMapTask`. It hasn't bitten us yet, but it seems like a good idea to change it here and down there.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50292919
I've now rebased this on top of the SizeTracker class in #1165 -- should be ready to go in.
There is one issue left with both the ExternalSorter and ExternalAppendOnlyMap, which is that they assume they are the only spilling collection in the thread. Tasks that make two maps, for example, will overcommit memory. I'd like to fix that too but I've opened a separate JIRA for it (https://issues.apache.org/jira/browse/SPARK-2711).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154525
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
+ // This is the last element, remember that we're done
+ finished = true
+ deserStream.close()
+ }
+ (k, c)
+ } catch {
+ case e: EOFException =>
--- End diff --
In what situation would we expect to get an EOF rather than closing cleanly?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50303404
Jenkins, test this please
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50159041
Not sure if it helps, but you can get PRs via something like `git fetch
apache refs/pull/1499/head`. Easier workflow, perhaps.
On Jul 25, 2014 12:25 AM, "Mridul Muralidharan" <no...@github.com>
wrote:
> ah, thanks ! rerunning with 9c29957
> <https://github.com/apache/spark/commit/9c29957>.
> cant pull the pr - and manual merge is painful, hence delays in testing :-)
>
> —
> Reply to this email directly or view it on GitHub
> <https://github.com/apache/spark/pull/1499#issuecomment-50116492>.
>
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154203
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
--- End diff --
Does this still work if some of the iterators are empty? In particular because the comparator calls ".head"
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49833579
I had pulled about 20 mins after I mailed you ...
I have elaborated on why this occurs inline in the code - we can ignore it for now though, since it happens even in 'regular' case : we had fixed it so long ago, I had forgotten about it and assumed it was some other issue !
We can punt on resolving this for now, and address it when the consolidated shuffle pr comes out.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49949511
@mateiz The total memory overhead actually goes much higher than num_streams right ?
It should be order of num_streams + num_values for this key.
For fairly large values, the latter might fit into memory, but the former might not (particularly as number of mappers increases).
Or did I get this wrong from the PR ?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50704737
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>class ShuffledRDD[K, V, C](<br>case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br>case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17526/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15332387
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
--- End diff --
Good question, but files on disk actually get cleaned by BlockManager directly if cleaning is turned on, and sort-based shuffle is set up to only use files. ShuffleBlockManager had extra state in memory in the form of metadata, that's why it needs its own cleaner. But I don't think there's an issue here.
It will be important to do something for unregisterShuffle, but right now nothing calls that. That's part of the API because I wanted to move the MapOutputTracker behind the ShuffleManager as well. But any patches to move more stuff behind it are welcome, and I agree the ShuffleBlockManager should be specific to hash-based shuffle.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15147889
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
--- End diff --
Note that the Ordering is not currently used, but will allow us to sort keys within each partition so that we can use this in sortByKey
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49669437
Actually I missed one comment, the one on the condition for spilling; I'll have to fix that both in ExternalSorter and EAOM
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15153139
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
--- End diff --
I left it this way because we might want to disable partial aggregation partway through if we see it's not aggregating too much.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15563638
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
--- End diff --
Ah, right
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154509
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
--- End diff --
I am probably misreading this, but won't this increment partitionId right before hasNext has consumed the last element of a partition (thereby making the last element look like it's in the subsequent partition)?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49564761
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16885/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562761
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
--- End diff --
Maybe add a comment that communicates our assumption that all iterators on the heap `hasNext`
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562901
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
--- End diff --
Here we're assuming that the `ordering` provided to ExternalSorter has total ordering. Is this necessarily true?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560379
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
--- End diff --
This is also used in EAOM. Maybe we should abstract this to `Utils.getShuffleMemory`
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50700437
test this please!
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154478
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
--- End diff --
We have the actual batch size in spill.serializerBatchSizes, which we used to limit the stream anyway -- should we use that instead (in case we ever introduce variable batch sizes besides the last one)?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560695
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
--- End diff --
Maybe add a TODO here that says this is very similar to the EAOM code
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15192915
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
--- End diff --
Yeah I like to put the bracket that way, otherwise it sometimes gets lost
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50210371
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17202/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49669025
I've now updated it with I believe all of @aarondav's fixes and an optimization he suggested for in-memory-only data (don't bother sorting by key or merge-sorting).
I also tried running this on all our test suites by making sort-based shuffle the default, and it almost works, except there's a test with custom Product2 types in ShuffledRDD that breaks. However upon further investigation I'm not sure that feature was fully supported even in the current version (our current spilling code would also not return the right type), so we might remove that.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50668930
Thanks everyone, I think I addressed all the comments. Anything else before we merge this? I'd like to merge it fairly soon because there are a few other issues I'd like fixed for 1.1 that will depend on this code: https://issues.apache.org/jira/browse/SPARK-2711, https://issues.apache.org/jira/browse/SPARK-931
@mridulm I haven't updated the spilling code to match your fixes in https://github.com/apache/spark/pull/1609 yet but I will do that once that PR is merged (in case we refactor the fixes in any way)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15494423
--- Diff: core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala ---
@@ -54,12 +55,16 @@ private[spark] class DiskBlockManager(shuffleManager: ShuffleBlockManager, rootD
addShutdownHook()
/**
- * Returns the physical file segment in which the given BlockId is located.
- * If the BlockId has been mapped to a specific FileSegment, that will be returned.
- * Otherwise, we assume the Block is mapped to a whole file identified by the BlockId directly.
+ * Returns the physical file segment in which the given BlockId is located. If the BlockId has
+ * been mapped to a specific FileSegment by the shuffle layer, that will be returned.
+ * Otherwise, we assume the Block is mapped to the whole file identified by the BlockId.
*/
def getBlockLocation(blockId: BlockId): FileSegment = {
- if (blockId.isShuffle && shuffleManager.consolidateShuffleFiles) {
+ val env = SparkEnv.get
+ if (blockId.isShuffle && env != null && env.shuffleManager.isInstanceOf[SortShuffleManager]) {
--- End diff --
Why could `env` be null here? If we never expect it to be null, we should probably raise an error if it is.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15495007
--- Diff: core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala ---
@@ -54,12 +55,16 @@ private[spark] class DiskBlockManager(shuffleManager: ShuffleBlockManager, rootD
addShutdownHook()
/**
- * Returns the physical file segment in which the given BlockId is located.
- * If the BlockId has been mapped to a specific FileSegment, that will be returned.
- * Otherwise, we assume the Block is mapped to a whole file identified by the BlockId directly.
+ * Returns the physical file segment in which the given BlockId is located. If the BlockId has
+ * been mapped to a specific FileSegment by the shuffle layer, that will be returned.
+ * Otherwise, we assume the Block is mapped to the whole file identified by the BlockId.
*/
def getBlockLocation(blockId: BlockId): FileSegment = {
- if (blockId.isShuffle && shuffleManager.consolidateShuffleFiles) {
+ val env = SparkEnv.get
+ if (blockId.isShuffle && env != null && env.shuffleManager.isInstanceOf[SortShuffleManager]) {
--- End diff --
This is actually the case in tests
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49567523
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16885/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15555233
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer.getOrElse(null))
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
--- End diff --
can we mark this as private?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560873
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
--- End diff --
also the wording here is slightly different from the equivalent in EAOM. Not sure if this is intentional
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50303714
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17274/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50306524
QA results for PR 1499:<br>- This patch PASSES unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>class ShuffledRDD[K, V, C](<br>case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int)<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17274/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50115353
Running tests with
export SPARK_JAVA_OPTS="-Dspark.shuffle.manager=org.apache.spark.shuffle.sort.SortShuffleManager" causes :
'''
- sorting using mutable pairs *** FAILED ***
org.apache.spark.SparkException: Job aborted due to stage failure: Task 3.0:1 failed 4 times, most recent failure: Exception failure in TID 14 on host localhost: java.lang.ArrayStoreException: scala.Tuple2
scala.runtime.ScalaRunTime$.array_update(ScalaRunTime.scala:88)
scala.Array$.slowcopy(Array.scala:81)
scala.Array$.copy(Array.scala:107)
scala.collection.mutable.ResizableArray$class.copyToArray(ResizableArray.scala:77)
scala.collection.mutable.ArrayBuffer.copyToArray(ArrayBuffer.scala:47)
scala.collection.TraversableOnce$class.copyToArray(TraversableOnce.scala:241)
scala.collection.AbstractTraversable.copyToArray(Traversable.scala:105)
scala.collection.TraversableOnce$class.toArray(TraversableOnce.scala:249)
scala.collection.AbstractTraversable.toArray(Traversable.scala:105)
scala.collection.TraversableOnce$class.toArray(TraversableOnce.scala:252)
org.apache.spark.InterruptibleIterator.toArray(InterruptibleIterator.scala:28)
org.apache.spark.rdd.OrderedRDDFunctions$$anonfun$sortByKey$1.apply(OrderedRDDFunctions.scala:62)
org.apache.spark.rdd.OrderedRDDFunctions$$anonfun$sortByKey$1.apply(OrderedRDDFunctions.scala:61)
org.apache.spark.rdd.RDD$$anonfun$12.apply(RDD.scala:581)
org.apache.spark.rdd.RDD$$anonfun$12.apply(RDD.scala:581)
org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:35)
org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:262)
org.apache.spark.rdd.RDD.iterator(RDD.scala:229)
org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:112)
org.apache.spark.scheduler.Task.run(Task.scala:51)
org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:199)
java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1145)
java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:615)
java.lang.Thread.run(Thread.java:745)
Driver stacktrace:
at org.apache.spark.scheduler.DAGScheduler.org$apache$spark$scheduler$DAGScheduler$$failJobAndIndependentStages(DAGScheduler.scala:1055)
at org.apache.spark.scheduler.DAGScheduler$$anonfun$abortStage$1.apply(DAGScheduler.scala:1039)
at org.apache.spark.scheduler.DAGScheduler$$anonfun$abortStage$1.apply(DAGScheduler.scala:1037)
at scala.collection.mutable.ResizableArray$class.foreach(ResizableArray.scala:59)
at scala.collection.mutable.ArrayBuffer.foreach(ArrayBuffer.scala:47)
at org.apache.spark.scheduler.DAGScheduler.abortStage(DAGScheduler.scala:1037)
at org.apache.spark.scheduler.DAGScheduler$$anonfun$handleTaskSetFailed$1.apply(DAGScheduler.scala:641)
at org.apache.spark.scheduler.DAGScheduler$$anonfun$handleTaskSetFailed$1.apply(DAGScheduler.scala:641)
at scala.Option.foreach(Option.scala:236)
at org.apache.spark.scheduler.DAGScheduler.handleTaskSetFailed(DAGScheduler.scala:641)
...
'''
The actual classes are :
'''
classes = <[class scala.Tuple2] ;
[class scala.Tuple2,
class scala.Tuple2,
class scala.Tuple2]>
'''
obtained using :
'''
def collect(): Array[T] = {
println("classes = " + sc.runJob(this, (iter: Iterator[T]) =>
iter.map(v => if (v == null) "null" else v.getClass).mkString("[", ",\n\t\t", "]")).
mkString("<", " ; \n\t", ">"))
System.out.flush()
val results = sc.runJob(this, (iter: Iterator[T]) => iter.toArray)
''' in RDD.collect
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154451
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
+ // This is the last element, remember that we're done
+ finished = true
+ deserStream.close()
+ }
+ (k, c)
+ } catch {
+ case e: EOFException =>
+ finished = true
+ deserStream.close()
+ null
+ }
+ }
+
+ var nextPartitionToRead = 0
+
+ def readNextPartition(): Iterator[Product2[K, C]] = new Iterator[Product2[K, C]] {
+ val myPartition = nextPartitionToRead
+ nextPartitionToRead += 1
+
+ override def hasNext: Boolean = {
+ if (nextItem == null) {
+ nextItem = readNextItem()
+ if (nextItem == null) {
+ return false
+ }
+ }
+ // Check that we're still in the right partition; will be numPartitions at EOF
+ partitionId == myPartition
--- End diff --
I know it's part of the preconditions of the iterator to read only one partition at a time, but we could reinforce this with an
```scala
assert(partitionId >= myPartition)
```
before this line. This would check for misuse (i.e., calling readNextPartition() before the prior one was exhausted).
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50213171
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>class ShuffledRDD[K, V, C](<br>case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int)<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17202/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15605438
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
--- End diff --
Alright, just changed this to use mergeSort directly
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50678947
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>class ShuffledRDD[K, V, C](<br>case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br>case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17472/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-50669138
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/17471/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49826834
@mridulm which version of the code was that with? Right now line 526 of ExternalSorter is not calling readObject, so it's hard to debug. There might've been some fixes to this before too.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15603509
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
+ if (dep.mapSideCombine) {
+ if (!dep.aggregator.isDefined) {
+ throw new IllegalStateException("Aggregator is empty for map-side combine")
+ }
+ sorter = new ExternalSorter[K, V, C](
+ dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ } else {
+ // In this case we pass neither an aggregator nor an ordering to the sorter, because we
+ // don't care whether the keys get sorted in each partition; that will be done on the
+ // reduce side if the operation being run is sortByKey.
+ sorter = new ExternalSorter[K, V, V](
+ None, Some(dep.partitioner), None, dep.serializer)
+ sorter.write(records)
+ sorter.partitionedIterator
+ }
+ }
+
+ // Create a single shuffle file with reduce ID 0 that we'll write all results to. We'll later
+ // serve different ranges of this file using an index file that we create at the end.
+ val blockId = ShuffleBlockId(dep.shuffleId, mapId, 0)
+ outputFile = blockManager.diskBlockManager.getFile(blockId)
+
+ // Track location of each range in the output file
+ val offsets = new Array[Long](numPartitions + 1)
+ val lengths = new Array[Long](numPartitions)
+
+ // Statistics
+ var totalBytes = 0L
+ var totalTime = 0L
+
+ for ((id, elements) <- partitions) {
+ if (elements.hasNext) {
+ val writer = blockManager.getDiskWriter(blockId, outputFile, ser, fileBufferSize)
+ for (elem <- elements) {
+ writer.write(elem)
+ }
+ writer.commit()
+ writer.close()
+ val segment = writer.fileSegment()
+ offsets(id + 1) = segment.offset + segment.length
+ lengths(id) = segment.length
+ totalTime += writer.timeWriting()
+ totalBytes += segment.length
+ } else {
+ // Don't create a new writer to avoid writing any headers and things like that
+ offsets(id + 1) = offsets(id)
+ }
+ }
+
+ val shuffleMetrics = new ShuffleWriteMetrics
+ shuffleMetrics.shuffleBytesWritten = totalBytes
+ shuffleMetrics.shuffleWriteTime = totalTime
+ context.taskMetrics.shuffleWriteMetrics = Some(shuffleMetrics)
+ context.taskMetrics.memoryBytesSpilled = sorter.memoryBytesSpilled
+ context.taskMetrics.diskBytesSpilled = sorter.diskBytesSpilled
--- End diff --
Good point; I'll fix this in other places in the codebase too
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49539656
QA tests have started for PR 1499. This patch merges cleanly. <br>View progress: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16865/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mridulm <gi...@git.apache.org>.
Github user mridulm commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15259118
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
This is actually broken - the code above defends against reading too much data from the stream, but not against insufficient data being read. That is, it assumes that when nextBatchStream() is called, all data has been read from previous stream (and so buffered stream is at begining of next stream).
Specifically, it assumes that write/flush/close will not write data more than what is required to deserialize the key/value.
For example, I think an easy way to test is to use java serialization and set object reset to very low value (a reset results in TC_RESET byte getting written; which is consumed as part of next object's deser). This results in EOF getting thrown while deserializing the next block
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562939
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
--- End diff --
This is noted in the class comment.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154547
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ }
+
+ /**
+ * Return the next (K, C) pair from the deserialization stream and update partitionId,
+ * indexInPartition, indexInBatch and such to match its location.
+ *
+ * If the current batch is drained, construct a stream for the next batch and read from it.
+ * If no more pairs are left, return null.
+ */
+ private def readNextItem(): (K, C) = {
+ try {
+ if (finished) {
+ return null
+ }
+ val k = deserStream.readObject().asInstanceOf[K]
+ val c = deserStream.readObject().asInstanceOf[C]
+ // Start reading the next batch if we're done with this one
+ indexInBatch += 1
+ if (indexInBatch == serializerBatchSize) {
+ batchStream = nextBatchStream()
+ compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ deserStream = serInstance.deserializeStream(compressedStream)
+ indexInBatch = 0
+ }
+ // Update the partition location of the element we're reading
+ indexInPartition += 1
+ while (indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0
+ }
+ if (partitionId == numPartitions - 1 &&
+ indexInPartition == spill.elementsPerPartition(partitionId) - 1) {
+ // This is the last element, remember that we're done
+ finished = true
+ deserStream.close()
+ }
+ (k, c)
+ } catch {
+ case e: EOFException =>
+ finished = true
+ deserStream.close()
+ null
+ }
+ }
+
+ var nextPartitionToRead = 0
+
+ def readNextPartition(): Iterator[Product2[K, C]] = new Iterator[Product2[K, C]] {
+ val myPartition = nextPartitionToRead
+ nextPartitionToRead += 1
+
+ override def hasNext: Boolean = {
+ if (nextItem == null) {
+ nextItem = readNextItem()
+ if (nextItem == null) {
+ return false
+ }
+ }
+ // Check that we're still in the right partition; will be numPartitions at EOF
+ partitionId == myPartition
+ }
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val item = nextItem
+ nextItem = null
+ item
+ }
+ }
+ }
+
+ /**
+ * Return an iterator over all the data written to this object, grouped by partition. For each
--- End diff --
Should we mention the guarantee that all values for a given key are guaranteed to have already been combined (if aggregation is enabled)?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15557253
--- Diff: core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala ---
@@ -54,12 +55,16 @@ private[spark] class DiskBlockManager(shuffleManager: ShuffleBlockManager, rootD
addShutdownHook()
/**
- * Returns the physical file segment in which the given BlockId is located.
- * If the BlockId has been mapped to a specific FileSegment, that will be returned.
- * Otherwise, we assume the Block is mapped to a whole file identified by the BlockId directly.
+ * Returns the physical file segment in which the given BlockId is located. If the BlockId has
+ * been mapped to a specific FileSegment by the shuffle layer, that will be returned.
+ * Otherwise, we assume the Block is mapped to the whole file identified by the BlockId.
*/
def getBlockLocation(blockId: BlockId): FileSegment = {
- if (blockId.isShuffle && shuffleManager.consolidateShuffleFiles) {
+ val env = SparkEnv.get
+ if (blockId.isShuffle && env != null && env.shuffleManager.isInstanceOf[SortShuffleManager]) {
+ val sortShuffleManager = SparkEnv.get.shuffleManager.asInstanceOf[SortShuffleManager]
--- End diff --
you already declared `env` up there, so might as well use it
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154394
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
Is it also possible that spill.serializerBatchSizes is empty?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15560851
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
--- End diff --
We recently added thread ID to the log messages in EAOM in #1517. We should do the same here
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15555237
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
--- End diff --
and this?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by rxin <gi...@git.apache.org>.
Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15543949
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala ---
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{BufferedOutputStream, File, FileOutputStream, DataOutputStream}
+
+import org.apache.spark.{MapOutputTracker, SparkEnv, Logging, TaskContext}
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.shuffle.{ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.collection.ExternalSorter
+
+private[spark] class SortShuffleWriter[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Int,
+ context: TaskContext)
+ extends ShuffleWriter[K, V] with Logging {
+
+ private val dep = handle.dependency
+ private val numPartitions = dep.partitioner.numPartitions
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val ser = Serializer.getSerializer(dep.serializer.getOrElse(null))
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ private var sorter: ExternalSorter[K, V, _] = null
+ private var outputFile: File = null
+
+ private var stopping = false
+ private var mapStatus: MapStatus = null
+
+ /** Write a bunch of records to this task's output */
+ override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ val partitions: Iterator[(Int, Iterator[Product2[K, _]])] = {
--- End diff --
can u add inline comment explaining what the tuple represents? (i'm assuming int represents the partition id?)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15152231
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,390 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
--- End diff --
Let's make sure we have at least a comment saying this in the code, if the PR is to be merged while that's true.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15271480
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
Ah, got it. Sorry, I missed your comment in the code earlier.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15620674
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,662 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+
+ // Size of object batches when reading/writing from serializers.
+ //
+ // Objects are written in batches, with each batch using its own serialization stream. This
+ // cuts down on the size of reference-tracking maps constructed when deserializing a stream.
+ //
+ // NOTE: Setting this too low can cause excessive copying when serializing, since some serializers
+ // grow internal data structures by growing + copying every time the number of objects doubles.
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ private var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ private var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow aggregation or sorting.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator and no ordering are given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk if needed.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ // TODO: factor this out of both here and ExternalAppendOnlyMap
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle) [SPARK-2711]
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ val threadId = Thread.currentThread().getId
+ logInfo("Thread %d spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(threadId, memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ // Use the reverse of comparator.compare because PriorityQueue dequeues the max
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*) // Will contain only the iterators with hasNext = true
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to read all keys considered equal by the ordering at once and compare them.
+ new Iterator[Iterator[Product2[K, C]]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ // Buffers reused across elements to decrease memory allocation
+ val keys = new ArrayBuffer[K]
+ val combiners = new ArrayBuffer[C]
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ keys.clear()
+ combiners.clear()
+ val firstPair = sorted.next()
+ keys += firstPair._1
+ combiners += firstPair._2
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val pair = sorted.next()
+ var i = 0
+ var foundKey = false
+ while (i < keys.size && !foundKey) {
+ if (keys(i) == pair._1) {
+ combiners(i) = mergeCombiners(combiners(i), pair._2)
+ foundKey = true
+ }
+ i += 1
+ }
+ if (!foundKey) {
+ keys += pair._1
+ combiners += pair._2
+ }
+ }
+
+ // Note that we return an iterator of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ keys.iterator.zip(combiners.iterator)
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering, so the objects with the same key are sequential.
+ new Iterator[Product2[K, C]] {
+ val sorted = mergeSort(iterators, comparator).buffered
+
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val elem = sorted.next()
+ val k = elem._1
+ var c = elem._2
+ while (sorted.hasNext && sorted.head._1 == k) {
+ c = mergeCombiners(c, sorted.head._2)
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in. These will be the indices of
+ // the next element we will read. We'll also store the last partition read so that
+ // readNextPartition() can figure out what partition that was from.
+ var partitionId = 0
+ var indexInPartition = 0L
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+ var lastPartitionId = 0
+
+ skipToNextPartition()
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ if (batchStreamsRead < spill.serializerBatchSizes.length) {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
+ } else {
+ // No more batches left; give an empty stream
+ bufferedStream
+ }
+ }
+
+ /**
+ * Update partitionId if we have reached the end of our current partition, possibly skipping
+ * empty partitions on the way.
+ */
+ private def skipToNextPartition() {
+ while (partitionId < numPartitions &&
+ indexInPartition == spill.elementsPerPartition(partitionId)) {
+ partitionId += 1
+ indexInPartition = 0L
--- End diff --
Yeah exactly
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by JoshRosen <gi...@git.apache.org>.
Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15493355
--- Diff: core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala ---
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.shuffle.sort
+
+import java.io.{DataInputStream, FileInputStream}
+
+import org.apache.spark.shuffle._
+import org.apache.spark.{TaskContext, ShuffleDependency}
+import org.apache.spark.shuffle.hash.HashShuffleReader
+import org.apache.spark.storage.{DiskBlockManager, FileSegment, ShuffleBlockId}
+
+private[spark] class SortShuffleManager extends ShuffleManager {
+ /**
+ * Register a shuffle with the manager and obtain a handle for it to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ new BaseShuffleHandle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext): ShuffleReader[K, C] = {
+ // We currently use the same block store shuffle fetcher as the hash-based shuffle.
+ new HashShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
+ : ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Unit = {}
+
+ /** Shut down this ShuffleManager. */
+ override def stop(): Unit = {}
+
+ /** Get the location of a block in a map output file. Uses the index file we create for it. */
+ def getBlockLocation(blockId: ShuffleBlockId, diskManager: DiskBlockManager): FileSegment = {
+ // The block is actually going to be a range of a single map output file for this map, so
+ // figure out the ID of the consolidated file, then the offset within that from our index
+ val realId = ShuffleBlockId(blockId.shuffleId, blockId.mapId, 0)
--- End diff --
This could also be written as
```scala
val realId = blockId.copy(reduceId = 0)
```
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: SPARK-2045 Sort-based shuffle
Posted by andrewor14 <gi...@git.apache.org>.
Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15562851
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,667 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional Partitioner; if given, sort by partition ID and then key
+ * @param ordering optional Ordering to sort keys within each partition; should be a total ordering
+ * @param serializer serializer to use when spilling to disk
+ *
+ * Note that if an Ordering is given, we'll always sort using it, so only provide it if you really
+ * want the output keys to be sorted. In a map task without map-side combine for example, you
+ * probably want to pass None as the ordering to avoid extra sorting. On the other hand, if you do
+ * want to do combining, having an Ordering is more efficient than not having it.
+ *
+ * At a high level, this class works as follows:
+ *
+ * - We repeatedly fill up buffers of in-memory data, using either a SizeTrackingAppendOnlyMap if
+ * we want to combine by key, or an simple SizeTrackingBuffer if we don't. Inside these buffers,
+ * we sort elements of type ((Int, K), C) where the Int is the partition ID. This is done to
+ * avoid calling the partitioner multiple times on the same key (e.g. for RangePartitioner).
+ *
+ * - When each buffer reaches our memory limit, we spill it to a file. This file is sorted first
+ * by partition ID and possibly second by key or by hash code of the key, if we want to do
+ * aggregation. For each file, we track how many objects were in each partition in memory, so we
+ * don't have to write out the partition ID for every element.
+ *
+ * - When the user requests an iterator, the spilled files are merged, along with any remaining
+ * in-memory data, using the same sort order defined above (unless both sorting and aggregation
+ * are disabled). If we need to aggregate by key, we either use a total ordering from the
+ * ordering parameter, or read the keys with the same hash code and compare them with each other
+ * for equality to merge values.
+ *
+ * - Users are expected to call stop() at the end to delete all the intermediate files.
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val spillingEnabled = conf.getBoolean("spark.shuffle.spill", true)
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingPairBuffer[(Int, K), C]
+
+ // Number of pairs read from input since last spill; note that we count them even if a value is
+ // merged with a previous key in case we're doing something like groupBy where the result grows
+ private var elementsRead = 0L
+
+ // What threshold of elementsRead we start estimating map size at.
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.2)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // How much of the shared memory pool this collection has claimed
+ private var myMemoryThreshold = 0L
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ // A comparator for (Int, K) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[(Int, K)] = {
+ if (ordering.isDefined || aggregator.isDefined) {
+ // Sort by partition ID then key comparator
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ val partitionDiff = a._1 - b._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._2, b._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[(Int, K)] {
+ override def compare(a: (Int, K), b: (Int, K)): Int = {
+ a._1 - b._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer.insert((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C])
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ if (!spillingEnabled) {
+ return
+ }
+
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && elementsRead % 32 == 0 &&
+ collection.estimateSize() >= myMemoryThreshold)
+ {
+ // TODO: This logic doesn't work if there are two external collections being used in the same
+ // task (e.g. to read shuffle output and write it out into another shuffle).
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // us to double our threshold
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyClaimedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyClaimedMemory.getOrElse(0L))
+
+ // Try to allocate at least 2x more memory, otherwise spill
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ myMemoryThreshold = currentSize * 2
+ }
+ }
+ // Do not hold lock during spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingPairCollection[(Int, K), C] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0 // Objects written since the last flush
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ objectsWritten = 0
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ throw e
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingPairBuffer[(Int, K), C]
+ }
+
+ // Reset the amount of shuffle memory used by this map in the global pool
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+ shuffleMemoryMap.synchronized {
+ shuffleMemoryMap(Thread.currentThread().getId) = 0
+ }
+ myMemoryThreshold = 0
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new IteratorForPartition(p, inMemBuffered)
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
--- End diff --
Maybe add a comment here that looks like: "descending order because mutable.PriorityQueue dequeues the max, not the min" (taken from EAOM). It's not super clear why we need a minus sign here.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by SparkQA <gi...@git.apache.org>.
Github user SparkQA commented on the pull request:
https://github.com/apache/spark/pull/1499#issuecomment-49541476
QA results for PR 1499:<br>- This patch FAILED unit tests.<br>- This patch merges cleanly<br>- This patch adds the following public classes (experimental):<br>trait SizeTrackingCollection[T] extends Iterable[T] {<br><br>For more information see test ouptut:<br>https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/16865/consoleFull
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by mateiz <gi...@git.apache.org>.
Github user mateiz commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15157617
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
+ {
+ if (!totalOrder) {
+ // We only have a partial ordering, e.g. comparing the keys by hash code, which means that
+ // multiple distinct keys might be treated as equal by the ordering. To deal with this, we
+ // need to buffer every set of keys considered equal by the comparator in memory, then do
+ // another pass through them to find the truly equal ones.
+ val sorted = mergeSort(iterators, comparator).buffered
+ // Buffers reused across keys to decrease memory allocation
+ val buf = new ArrayBuffer[(K, C)]
+ val toReturn = new ArrayBuffer[(K, C)]
+ new Iterator[Iterator[Product2[K, C]]] {
+ override def hasNext: Boolean = sorted.hasNext
+
+ override def next(): Iterator[Product2[K, C]] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstPair = sorted.next()
+ buf += ((firstPair._1, firstPair._2)) // Copy it in case the Product2 object is reused
+ val key = firstPair._1
+ while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
+ val n = sorted.next()
+ buf += ((n._1, n._2))
+ }
+ // buf now contains all the elements with keys equal to our first one according to the
+ // partial ordering. Now we need to find which keys were "really" equal, which we do
+ // through linear scans through the buffer.
+ toReturn.clear()
+ while (!buf.isEmpty) {
+ val last = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ val k = last._1
+ var c = last._2
+ var i = 0
+ while (i < buf.size) {
+ while (i < buf.size && buf(i)._1 == k) {
+ c = mergeCombiners(c, buf(i)._2)
+ // Replace this element with the last one in the buffer
+ buf(i) = buf(buf.size - 1)
+ buf.reduceToSize(buf.size - 1)
+ }
+ i += 1
+ }
+ toReturn += ((k, c))
+ }
+ // Note that we return a *sequence* of elements since we could've had many keys marked
+ // equal by the partial order; we flatten this below to get a flat iterator of (K, C).
+ toReturn.iterator
+ }
+ }.flatMap(i => i)
+ } else {
+ // We have a total ordering. This means we can merge objects one by one as we read them
+ // from the iterators, without buffering all the ones that are "equal" to a given key.
+ // We do so with code similar to mergeSort, except our Iterator.next combines together all
+ // the elements with the given key.
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ val k = firstPair._1
+ var c = firstPair._2
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ var shouldStop = false
+ while (!heap.isEmpty && !shouldStop) {
+ shouldStop = true // Stop unless we find another element with the same key
+ val newBuf = heap.dequeue()
+ while (newBuf.hasNext && newBuf.head._1 == k) {
+ val elem = newBuf.next()
+ c = mergeCombiners(c, elem._2)
+ shouldStop = false
+ }
+ if (newBuf.hasNext) {
+ heap.enqueue(newBuf)
+ }
+ }
+ (k, c)
+ }
+ }
+ }
+ }
+
+ /**
+ * An internal class for reading a spilled file partition by partition. Expects all the
+ * partitions to be requested in order.
+ */
+ private[this] class SpillReader(spill: SpilledFile) {
+ val fileStream = new FileInputStream(spill.file)
+ val bufferedStream = new BufferedInputStream(fileStream, fileBufferSize)
+
+ // Track which partition and which batch stream we're in
+ var partitionId = 0
+ var indexInPartition = -1L // Just to make sure we start at index 0
+ var batchStreamsRead = 0
+ var indexInBatch = 0
+
+ // An intermediate stream that reads from exactly one batch
+ // This guards against pre-fetching and other arbitrary behavior of higher level streams
+ var batchStream = nextBatchStream()
+ var compressedStream = blockManager.wrapForCompression(spill.blockId, batchStream)
+ var deserStream = serInstance.deserializeStream(compressedStream)
+ var nextItem: (K, C) = null
+ var finished = false
+
+ /** Construct a stream that only reads from the next batch */
+ def nextBatchStream(): InputStream = {
+ batchStreamsRead += 1
+ ByteStreams.limit(bufferedStream, spill.serializerBatchSizes(batchStreamsRead - 1))
--- End diff --
Nope, there's always one at the end.
In the part above, we can't do the += after because we want to return a stream from the function. We'd have to assign it to a variable.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---
[GitHub] spark pull request: (WIP) SPARK-2045 Sort-based shuffle
Posted by aarondav <gi...@git.apache.org>.
Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/1499#discussion_r15154224
--- Diff: core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala ---
@@ -0,0 +1,573 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.util.collection
+
+import java.io._
+import java.util.Comparator
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.{Aggregator, SparkEnv, Logging, Partitioner}
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.storage.BlockId
+
+/**
+ * Sorts and potentially merges a number of key-value pairs of type (K, V) to produce key-combiner
+ * pairs of type (K, C). Uses a Partitioner to first group the keys into partitions, and then
+ * optionally sorts keys within each partition using a custom Comparator. Can output a single
+ * partitioned file with a different byte range for each partition, suitable for shuffle fetches.
+ *
+ * If combining is disabled, the type C must equal V -- we'll cast the objects at the end.
+ *
+ * @param aggregator optional Aggregator with combine functions to use for merging data
+ * @param partitioner optional partitioner; if given, sort by partition ID and then key
+ * @param ordering optional ordering to sort keys within each partition
+ * @param serializer serializer to use when spilling to disk
+ */
+private[spark] class ExternalSorter[K, V, C](
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ partitioner: Option[Partitioner] = None,
+ ordering: Option[Ordering[K]] = None,
+ serializer: Option[Serializer] = None) extends Logging {
+
+ private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
+ private val shouldPartition = numPartitions > 1
+
+ private val blockManager = SparkEnv.get.blockManager
+ private val diskBlockManager = blockManager.diskBlockManager
+ private val ser = Serializer.getSerializer(serializer)
+ private val serInstance = ser.newInstance()
+
+ private val conf = SparkEnv.get.conf
+ private val fileBufferSize = conf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
+ private val serializerBatchSize = conf.getLong("spark.shuffle.spill.batchSize", 10000)
+
+ private def getPartition(key: K): Int = {
+ if (shouldPartition) partitioner.get.getPartition(key) else 0
+ }
+
+ // Data structures to store in-memory objects before we spill. Depending on whether we have an
+ // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
+ // store them in an array buffer.
+ var map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ var buffer = new SizeTrackingBuffer[((Int, K), C)]
+
+ // Track how many elements we've read before we try to estimate memory. Ideally we'd use
+ // map.size or buffer.size for this, but because users' Aggregators can potentially increase
+ // the size of a merged element when we add values with the same key, it's safer to track
+ // elements read from the input iterator.
+ private var elementsRead = 0L
+ private val trackMemoryThreshold = 1000
+
+ // Spilling statistics
+ private var spillCount = 0
+ private var _memoryBytesSpilled = 0L
+ private var _diskBytesSpilled = 0L
+
+ // Collective memory threshold shared across all running tasks
+ private val maxMemoryThreshold = {
+ val memoryFraction = conf.getDouble("spark.shuffle.memoryFraction", 0.3)
+ val safetyFraction = conf.getDouble("spark.shuffle.safetyFraction", 0.8)
+ (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
+ }
+
+ // A comparator for keys K that orders them within a partition to allow partial aggregation.
+ // Can be a partial ordering by hash code if a total ordering is not provided through by the
+ // user. (A partial ordering means that equal keys have comparator.compare(k, k) = 0, but some
+ // non-equal keys also have this, so we need to do a later pass to find truly equal keys).
+ // Note that we ignore this if no aggregator is given.
+ private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
+ override def compare(a: K, b: K): Int = {
+ val h1 = if (a == null) 0 else a.hashCode()
+ val h2 = if (b == null) 0 else b.hashCode()
+ h1 - h2
+ }
+ })
+
+ private val sortWithinPartitions = ordering.isDefined || aggregator.isDefined
+
+ // A comparator for ((Int, K), C) elements that orders them by partition and then possibly by key
+ private val partitionKeyComparator: Comparator[((Int, K), C)] = {
+ if (sortWithinPartitions) {
+ // Sort by partition ID then key comparator
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ val partitionDiff = a._1._1 - b._1._1
+ if (partitionDiff != 0) {
+ partitionDiff
+ } else {
+ keyComparator.compare(a._1._2, b._1._2)
+ }
+ }
+ }
+ } else {
+ // Just sort it by partition ID
+ new Comparator[((Int, K), C)] {
+ override def compare(a: ((Int, K), C), b: ((Int, K), C)): Int = {
+ a._1._1 - b._1._1
+ }
+ }
+ }
+ }
+
+ // Information about a spilled file. Includes sizes in bytes of "batches" written by the
+ // serializer as we periodically reset its stream, as well as number of elements in each
+ // partition, used to efficiently keep track of partitions when merging.
+ private[this] case class SpilledFile(
+ file: File,
+ blockId: BlockId,
+ serializerBatchSizes: Array[Long],
+ elementsPerPartition: Array[Long])
+ private val spills = new ArrayBuffer[SpilledFile]
+
+ def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
+ // TODO: stop combining if we find that the reduction factor isn't high
+ val shouldCombine = aggregator.isDefined
+
+ if (shouldCombine) {
+ // Combine values in-memory first using our AppendOnlyMap
+ val mergeValue = aggregator.get.mergeValue
+ val createCombiner = aggregator.get.createCombiner
+ var kv: Product2[K, V] = null
+ val update = (hadValue: Boolean, oldValue: C) => {
+ if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
+ }
+ while (records.hasNext) {
+ elementsRead += 1
+ kv = records.next()
+ map.changeValue((getPartition(kv._1), kv._1), update)
+ maybeSpill(usingMap = true)
+ }
+ } else {
+ // Stick values into our buffer
+ while (records.hasNext) {
+ elementsRead += 1
+ val kv = records.next()
+ buffer += (((getPartition(kv._1), kv._1), kv._2.asInstanceOf[C]))
+ maybeSpill(usingMap = false)
+ }
+ }
+ }
+
+ private def maybeSpill(usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+
+ if (elementsRead > trackMemoryThreshold && collection.atGrowThreshold) {
+ // TODO: This is code from ExternalAppendOnlyMap that doesn't work if there are two external
+ // collections being used in the same task. However we'll just copy it for now.
+
+ val currentSize = collection.estimateSize()
+ var shouldSpill = false
+ val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
+
+ // Atomically check whether there is sufficient memory in the global pool for
+ // this map to grow and, if possible, allocate the required amount
+ shuffleMemoryMap.synchronized {
+ val threadId = Thread.currentThread().getId
+ val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
+ val availableMemory = maxMemoryThreshold -
+ (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))
+
+ // Assume map growth factor is 2x
+ shouldSpill = availableMemory < currentSize * 2
+ if (!shouldSpill) {
+ shuffleMemoryMap(threadId) = currentSize * 2
+ }
+ }
+ // Do not synchronize spills
+ if (shouldSpill) {
+ spill(currentSize, usingMap)
+ }
+ }
+ }
+
+ /**
+ * Spill the current in-memory collection to disk, adding a new file to spills, and clear it.
+ *
+ * @param usingMap whether we're using a map or buffer as our current in-memory collection
+ */
+ private def spill(memorySize: Long, usingMap: Boolean): Unit = {
+ val collection: SizeTrackingCollection[((Int, K), C)] = if (usingMap) map else buffer
+ val memorySize = collection.estimateSize()
+
+ spillCount += 1
+ logWarning("Spilling in-memory batch of %d MB to disk (%d spill%s so far)"
+ .format(memorySize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
+ val (blockId, file) = diskBlockManager.createTempBlock()
+ var writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ var objectsWritten = 0
+
+ // List of batch sizes (bytes) in the order they are written to disk
+ val batchSizes = new ArrayBuffer[Long]
+
+ // How many elements we have in each partition
+ val elementsPerPartition = new Array[Long](numPartitions)
+
+ // Flush the disk writer's contents to disk, and update relevant variables
+ def flush() = {
+ writer.commit()
+ val bytesWritten = writer.bytesWritten
+ batchSizes.append(bytesWritten)
+ _diskBytesSpilled += bytesWritten
+ objectsWritten = 0
+ }
+
+ try {
+ val it = collection.destructiveSortedIterator(partitionKeyComparator)
+ while (it.hasNext) {
+ val elem = it.next()
+ val partitionId = elem._1._1
+ val key = elem._1._2
+ val value = elem._2
+ writer.write(key)
+ writer.write(value)
+ elementsPerPartition(partitionId) += 1
+ objectsWritten += 1
+
+ if (objectsWritten == serializerBatchSize) {
+ flush()
+ writer.close()
+ writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize)
+ }
+ }
+ if (objectsWritten > 0) {
+ flush()
+ }
+ writer.close()
+ } catch {
+ case e: Exception =>
+ writer.close()
+ file.delete()
+ }
+
+ if (usingMap) {
+ map = new SizeTrackingAppendOnlyMap[(Int, K), C]
+ } else {
+ buffer = new SizeTrackingBuffer[((Int, K), C)]
+ }
+
+ spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
+ _memoryBytesSpilled += memorySize
+ }
+
+ /**
+ * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
+ * inside each partition. This can be used to either write out a new file or return data to
+ * the user.
+ *
+ * Returns an iterator over all the data written to this object, grouped by partition. For each
+ * partition we then have an iterator over its contents, and these are expected to be accessed
+ * in order (you can't "skip ahead" to one partition without reading the previous one).
+ * Guaranteed to return a key-value pair for each partition, in order of partition ID.
+ */
+ private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
+ : Iterator[(Int, Iterator[Product2[K, C]])] = {
+ // TODO: merge intermediate results if they are sorted by the comparator
+ val readers = spills.map(new SpillReader(_))
+ val inMemBuffered = inMemory.buffered
+ (0 until numPartitions).iterator.map { p =>
+ val inMemIterator = new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = {
+ inMemBuffered.hasNext && inMemBuffered.head._1._1 == p
+ }
+ override def next(): Product2[K, C] = {
+ val elem = inMemBuffered.next()
+ (elem._1._2, elem._2)
+ }
+ }
+ val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
+ if (aggregator.isDefined) {
+ // Perform partial aggregation across partitions
+ (p, mergeWithAggregation(
+ iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
+ } else if (ordering.isDefined) {
+ // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
+ // sort the elements without trying to merge them
+ (p, mergeSort(iterators, ordering.get))
+ } else {
+ (p, iterators.iterator.flatten)
+ }
+ }
+ }
+
+ /**
+ * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
+ */
+ private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K])
+ : Iterator[Product2[K, C]] =
+ {
+ val bufferedIters = iterators.map(_.buffered)
+ type Iter = BufferedIterator[Product2[K, C]]
+ val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
+ override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
+ })
+ heap.enqueue(bufferedIters: _*)
+ new Iterator[Product2[K, C]] {
+ override def hasNext: Boolean = !heap.isEmpty
+
+ override def next(): Product2[K, C] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val firstBuf = heap.dequeue()
+ val firstPair = firstBuf.next()
+ if (firstBuf.hasNext) {
+ heap.enqueue(firstBuf)
+ }
+ firstPair
+ }
+ }
+ }
+
+ /**
+ * Merge a sequence of (K, C) iterators by aggregating values for each key, assuming that each
+ * iterator is sorted by key with a given comparator. If the comparator is not a total ordering
+ * (e.g. when we sort objects by hash code and different keys may compare as equal although
+ * they're not), we still merge them by doing equality tests for all keys that compare as equal.
+ */
+ private def mergeWithAggregation(
+ iterators: Seq[Iterator[Product2[K, C]]],
+ mergeCombiners: (C, C) => C,
+ comparator: Comparator[K],
+ totalOrder: Boolean)
+ : Iterator[Product2[K, C]] =
--- End diff --
(oh man, 4 spaces AND bracket on the next line)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---