/* * Modifications copyright (C) 2018 MemVerge Inc. * * Use TmpShuffleFile interface instead of raw file. * * 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 import org.apache.spark.{Aggregator, TaskContext} class SplashAggregator[K, V, C]( agg: Aggregator[K, V, C]) extends Aggregator[K, V, C]( agg.createCombiner, agg.mergeValue, agg.mergeCombiners) { override def combineValuesByKey( iter: Iterator[_ <: Product2[K, V]], context: TaskContext): Iterator[(K, C)] = { val combiners = new SplashAppendOnlyMap[K, V, C](createCombiner, mergeValue, mergeCombiners) combiners.insertAll(iter) updateMetrics(context, combiners) combiners.iterator } override def combineCombinersByKey( iter: Iterator[_ <: Product2[K, C]], context: TaskContext): Iterator[(K, C)] = { val combiners = new SplashAppendOnlyMap[K, C, C](identity, mergeCombiners, mergeCombiners) combiners.insertAll(iter) updateMetrics(context, combiners) combiners.iterator } /** Update task metrics after populating the external map. */ private def updateMetrics(context: TaskContext, map: SplashAppendOnlyMap[_, _, _]): Unit = { Option(context).foreach { c => c.taskMetrics().incMemoryBytesSpilled(map.memoryBytesSpilled) c.taskMetrics().incDiskBytesSpilled(map.bytesSpilled) c.taskMetrics().incPeakExecutionMemory(map.peakMemoryUsedBytes) } } }