[原始碼解析]為什麼mapPartition比map更高效
0x00 摘要
自從函數語言程式設計和響應式程式設計逐漸進入到程式設計師的生活之後，map函式作為其中一個重要運算元也為大家所熟知，無論是前端web開發，手機開發還是後端伺服器開發，都很難逃過它的手心。而在大資料領域中又往往可以見到另外一個運算元mapPartition的身影。在效能調優中，經常會被建議儘量用 mappartition 操作去替代 map 操作。本文將從Flink原始碼和示例入手，為大家解析為什麼mapPartition比map更高效。
0x01 map vs mapPartition
1.1 map
Map的作用是將資料流上每個元素轉換為另外的元素，比如data.map { x => x.toInt }。它把陣列流中的每一個值，使用所提供的函式執行一遍，一一對應。得到與元素個數相同的陣列流。然後返回這個新資料流。
1.2 mapPartition
MapPartition的作用是單個函式呼叫並行分割槽，比如data.mapPartition { in => in map { (_, 1) } }。該函式將分割槽作為“迭代器”，可以產生任意數量的結果。每個分割槽中的元素數量取決於並行度和以前的operations。
1.3 異同
其實，兩者完成的業務操作是一樣的，本質上都是將資料流上每個元素轉換為另外的元素。
區別主要在兩點。
從邏輯實現來講，
- map邏輯實現簡單，就是在函式中簡單一一轉換，map函式的輸入和輸入都是單個元素。
- mapPartition相對複雜，函式的輸入有兩個，一般格式為
void mapPartition(Iterable<T> values, Collector<O> out)。其中values是需要對映轉換的所有記錄，out是用來傳送結果的collector。具體返回什麼，如何操作out來返回結果，則完全依賴於業務邏輯。
從呼叫次數來說，
- 資料有多少個元素，map就會被呼叫多少次。
- 資料有多少分割槽，mapPartition就會被呼叫多少次。
為什麼MapPartition有這麼高效呢，下面我們將具體論證。
0x02 程式碼
首先我們給出示例程式碼，從下文中我們可以看出，map就是簡單的轉換，而mapPartition則不但要做轉換，程式設計師還需要手動操作如何返回結果：
public class IteratePi {
public static void main(String[] args) throws Exception {
final ExecutionEnvironment env=ExecutionEnvironment.getExecutionEnvironment();
//迭代次數
int iterativeNum=10;
DataSet<Integer> wordList = env.fromElements(1, 2, 3);
IterativeDataSet<Integer> iterativeDataSet=wordList.iterate(iterativeNum);
DataSet<Integer> mapResult=iterativeDataSet
.map(new MapFunction<Integer, Integer>() {
@Override
public Integer map(Integer value) throws Exception {
value += 1;
return value;
}
});
//迭代結束的條件
DataSet<Integer> result=iterativeDataSet.closeWith(mapResult);
result.print();
MapPartitionOperator<Integer, Integer> mapPartitionResult = iterativeDataSet
.mapPartition(new MapPartitionFunction<Integer, Integer>() {
@Override
public void mapPartition(Iterable<Integer> values, Collector<Integer> out) {
for (Integer value : values) {
// 這裡需要程式設計師自行決定如何返回，即呼叫collect操作。
out.collect(value + 2);
}
} }
);
//迭代結束的條件
DataSet<Integer> partitionResult=iterativeDataSet.closeWith(mapPartitionResult);
partitionResult.print();
}
}
0x03 Flink的傳輸機制
世界上很少有沒有來由的愛，也少見免費的午餐。mapPartition之所以高效，其所依賴的基礎就是Flink的傳輸機制。所以我們下面就講解下為什麼。
大家都知道，Spark是用微批處理來模擬流處理，就是說，spark還是一批一批的傳輸和處理資料，所以我們就能理解mapPartition的機制就是基於這一批資料做統一處理。這樣確實可以高效。
但是Flink號稱是純流，即Flink是每來一個輸入record，就進行一次業務處理，然後返回給下游運算元。
有的兄弟就會產生疑問：每次都只是處理單個記錄，怎麼能夠讓mapPartition做到批次處理呢。其實這就是Flink的微妙之處：即Flink確實是每次都處理一個輸入record，但是在上下游傳輸時候，Flink還是把records累積起來做批量傳輸的。也可以這麼理解：從傳輸的角度講，Flink是微批處理的。
3.1 傳輸機制概述
Flink 的網路棧是組成 flink-runtime 模組的核心元件之一，也是 Flink 作業的核心部分。所有來自 TaskManager 的工作單元（子任務）都通過它來互相連線。流式傳輸資料流都要經過網路棧，所以它對 Flink 作業的效能表現（包括吞吐量和延遲指標）至關重要。與通過 Akka 使用 RPC 的 TaskManager 和 JobManager 之間的協調通道相比，TaskManager 之間的網路棧依賴的是更底層的，基於 Netty 的 API。
3.2 遠端通訊
一個執行的application的tasks在持續交換資料。TaskManager負責做資料傳輸。不同任務之間的每個（遠端）網路連線將在 Flink 的網路棧中獲得自己的 TCP 通道。但是如果同一任務的不同子任務被安排到了同一個 TaskManager，則它們與同一個 TaskManager 的網路連線將被多路複用，並共享一個 TCP 通道以減少資源佔用。
每個TaskManager有一組網路緩衝池（預設每個buffer是32KB），用於傳送與接受資料。如傳送端和接收端位於不同的TaskManager程式中，則它們需要通過作業系統的網路棧進行交流。流應用需要以管道的模式進行資料交換，也就是說，每對TaskManager會維持一個永久的TCP連線用於做資料交換。在shuffle連線模式下（多個sender與多個receiver），每個sender task需要向每個receiver task傳送資料，此時TaskManager需要為每個receiver task都分配一個緩衝區。
一個記錄被建立並傳遞之後（例如通過 Collector.collect()），它會被遞交到RecordWriter，其將來自 Java 物件的記錄序列化為一個位元組序列，後者最終成為網路快取。RecordWriter 首先使用SpanningRecordSerializer將記錄序列化為一個靈活的堆上位元組陣列。然後它嘗試將這些位元組寫入目標網路通道的關聯網路快取。
因為如果逐個傳送會降低每個記錄的開銷並帶來更高的吞吐量，所以為了取得高吞吐量，TaskManager的網路元件首先從緩衝buffer中收集records，然後再傳送。也就是說，records並不是一個接一個的傳送，而是先放入緩衝，然後再以batch的形式傳送。這個技術可以高效使用網路資源，並達到高吞吐。類似於網路或磁碟 I/O 協議中使用的緩衝技術。
接收方網路棧（netty）將接收到的快取寫入適當的輸入通道。最後（流式）任務的執行緒從這些佇列中讀取並嘗試在RecordReader的幫助下，通過Deserializer將積累的資料反序列化為 Java 物件。
3.3 TaskManager程式內傳輸
若sender與receiver任務都執行在同一個TaskManager程式，則sender任務會將傳送的條目做序列化，並存入一個位元組緩衝。然後將緩衝放入一個佇列，直到佇列被填滿。
Receiver任務從佇列中獲取緩衝，並反序列化輸入的條目。所以，在同一個TaskManager內，任務之間的資料傳輸並不經過網路互動。
即在同一個TaskManager程式內，也是批量傳輸。
3.4 原始碼分析
我們基於Flink優化的結果進行分析驗證，看看Flink是不是把記錄寫入到buffer中，這種情況下執行的是CountingCollector和ChainedMapDriver。
copyFromSerializerToTargetChannel:153, RecordWriter (org.apache.flink.runtime.io.network.api.writer)
emit:116, RecordWriter (org.apache.flink.runtime.io.network.api.writer)
emit:60, ChannelSelectorRecordWriter (org.apache.flink.runtime.io.network.api.writer)
collect:65, OutputCollector (org.apache.flink.runtime.operators.shipping)
collect:35, CountingCollector (org.apache.flink.runtime.operators.util.metrics)
collect:79, ChainedMapDriver (org.apache.flink.runtime.operators.chaining)
collect:35, CountingCollector (org.apache.flink.runtime.operators.util.metrics)
invoke:196, DataSourceTask (org.apache.flink.runtime.operators)
doRun:707, Task (org.apache.flink.runtime.taskmanager)
run:532, Task (org.apache.flink.runtime.taskmanager)
run:748, Thread (java.lang)
當執行完使用者定義的map函式之後，系統執行在 ChainedMapDriver.collect 函式。
public void collect(IT record) {
this.outputCollector.collect(this.mapper.map(record));// mapper就是使用者程式碼
}
然後呼叫到了CountingCollector.collect
public void collect(OUT record) {
this.collector.collect(record);// record就是使用者轉換後的記錄
}
OutputCollector.collect函式會把記錄傳送給所有的writers。
this.delegate.setInstance(record);// 先把record設定到SerializationDelegate中
for (RecordWriter<SerializationDelegate<T>> writer : writers) { // 所有的writer
writer.emit(this.delegate); // 傳送record
}
RecordWriter負責把資料序列化，然後寫入到快取中。它有兩個實現類：
BroadcastRecordWriter: 維護了多個下游channel，傳送資料到下游所有的channel中。ChannelSelectorRecordWriter: 通過channelSelector物件判斷資料需要發往下游的哪個channel。我們用的正是這個RecordWriter。
這裡我們分析下ChannelSelectorRecordWriter的emit方法：
public void emit(T record) throws IOException, InterruptedException {
emit(record, channelSelector.selectChannel(record));
}
這裡使用了channelSelector.selectChannel方法。該方法為record尋找到對應下游channel id。
public class OutputEmitter<T> implements ChannelSelector<SerializationDelegate<T>> {
public final int selectChannel(SerializationDelegate<T> record) {
switch (strategy) {
case FORWARD:
return forward(); // 我們程式碼用到了這種情況。這裡 return 0;
......
}
}
}
接下來我們又回到了父類RecordWriter.emit。
protected void emit(T record, int targetChannel) throws IOException, InterruptedException {
serializer.serializeRecord(record);
// Make sure we don't hold onto the large intermediate serialization buffer for too long
if (copyFromSerializerToTargetChannel(targetChannel)) {
serializer.prune();
}
}
關鍵的邏輯在於copyFromSerializerToTargetChannel。此方法從序列化器中複製資料到目標channel，我們可以看出來，每條記錄都是寫入到buffer中。
protected boolean copyFromSerializerToTargetChannel(int targetChannel) throws IOException, InterruptedException {
// We should reset the initial position of the intermediate serialization buffer before
// copying, so the serialization results can be copied to multiple target buffers.
// 此處Serializer為SpanningRecordSerializer
// reset方法將serializer內部的databuffer position重置為0
serializer.reset();
boolean pruneTriggered = false;
// 獲取目標channel的bufferBuilder
// bufferBuilder內維護了MemorySegment，即記憶體片段
// Flink的記憶體管理依賴MemorySegment，可實現堆內堆外記憶體的管理
// RecordWriter內有一個bufferBuilder陣列，長度和下游channel數目相同
// 該陣列以channel ID為下標，儲存和channel對應的bufferBuilder
// 如果對應channel的bufferBuilder尚未建立，呼叫requestNewBufferBuilder申請一個新的bufferBuilder
BufferBuilder bufferBuilder = getBufferBuilder(targetChannel);
// 複製serializer的資料到bufferBuilder中
SerializationResult result = serializer.copyToBufferBuilder(bufferBuilder);
// 迴圈直到result完全被寫入到buffer
// 一條資料可能會被寫入到多個快取中
// 如果快取不夠用，會申請新的快取
// 資料完全寫入完畢之時，當前正在操作的快取是沒有寫滿的
// 因此返回true，表明需要壓縮該buffer的空間
while (result.isFullBuffer()) {
finishBufferBuilder(bufferBuilder);
// If this was a full record, we are done. Not breaking out of the loop at this point
// will lead to another buffer request before breaking out (that would not be a
// problem per se, but it can lead to stalls in the pipeline).
if (result.isFullRecord()) {
pruneTriggered = true;
emptyCurrentBufferBuilder(targetChannel);
break;
}
bufferBuilder = requestNewBufferBuilder(targetChannel);
result = serializer.copyToBufferBuilder(bufferBuilder);
}
checkState(!serializer.hasSerializedData(), "All data should be written at once");
// 如果buffer超時時間為0，需要flush目標channel的資料
if (flushAlways) {
flushTargetPartition(targetChannel);
}
return pruneTriggered;
}
0x04 runtime
4.1 Driver
Driver是Flink runtime的一個重要概念，是在一個task中執行的使用者業務邏輯元件，具體實現了批量操作程式碼。其內部API包括初始化，清除，執行，取消等邏輯。
public interface Driver<S extends Function, OT> {
......
void setup(TaskContext<S, OT> context);
void run() throws Exception;
void cleanup() throws Exception;
void cancel() throws Exception;
}
具體在 org.apache.flink.runtime.operators 目錄下，我們能夠看到各種Driver的實現，基本的運算元都有自己的Driver。
......
CoGroupDriver.java
FlatMapDriver.java
FullOuterJoinDriver.java
GroupReduceCombineDriver.java
GroupReduceDriver.java
JoinDriver.java
LeftOuterJoinDriver.java
MapDriver.java
MapPartitionDriver.java
......
4.2 MapDriver
map運算元對應的就是MapDriver。
結合上節我們知道，上游資料是通過batch方式批量傳入的。所以，在run函式會遍歷輸入，每次取出一個record，然後呼叫使用者自定義函式function.map對這個record做map操作。
public class MapDriver<IT, OT> implements Driver<MapFunction<IT, OT>, OT> {
@Override
public void run() throws Exception {
final MutableObjectIterator<IT> input = this.taskContext.getInput(0);
.....
else {
IT record = null;
// runtime主動進行迴圈，這樣導致大量函式呼叫
while (this.running && ((record = input.next()) != null)) {
numRecordsIn.inc();
output.collect(function.map(record)); // function是使用者函式
}
}
}
}
4.3 MapPartitionDriver
MapPartitionDriver是mapPartition的具體元件。系統會把得到的批量資料inIter一次性的都傳給使用者自定義函式，由使用者程式碼來進行遍歷操作。
public class MapPartitionDriver<IT, OT> implements Driver<MapPartitionFunction<IT, OT>, OT> {
@Override
public void run() throws Exception {
final MutableObjectIterator<IT> input = new CountingMutableObjectIterator<>(this.taskContext.<IT>getInput(0), numRecordsIn);
......
} else {
final NonReusingMutableToRegularIteratorWrapper<IT> inIter = new NonReusingMutableToRegularIteratorWrapper<IT>(input, this.taskContext.<IT>getInputSerializer(0).getSerializer());
// runtime不參與迴圈，這樣可以減少函式呼叫
function.mapPartition(inIter, output);
}
}
}
4.4 效率區別
我們能夠看到map和mapPartition的input都是MutableObjectIterator
- 假設一共有60個資料需要轉換，map會在runtime中呼叫使用者函式60次。
- runtime把資料分成6個partition操作，則mapPartition在runtime中會呼叫使用者函式6次，在每個使用者函式中分別迴圈10次。對於runtime來說，map操作會多出54次使用者函式呼叫。
- 如果使用者業務中需要頻繁建立額外的物件或者外部資源操作，mapPartition的優勢更可以體現。 例如將資料寫入Mysql, 那麼map需要為每個元素建立一個資料庫連線，而mapPartition為每個partition建立一個連結。
假設有上億個資料需要map，這資源佔用和執行速度效率差別會相當大。
0x05 優化和ChainedMapDriver
之前提到了優化，這裡我們再詳細深入下如何優化map運算元。
Flink有一個關鍵的優化技術稱為任務鏈，用於（在某些情況下）減少本地通訊的過載。為了滿足任務鏈的條件，至少兩個以上的operator必須配置為同一並行度，並且使用本地向前的（local forwad）方式連線。任務鏈可以被認為是一種管道。
當管道以任務鏈的方式執行時候，Operators的函式被融合成單個任務，並由一個單獨的執行緒執行。一個function產生的records，通過使用一個簡單的方法呼叫，被遞交給下一個function。所以這裡在方法之間的records傳遞中，基本沒有序列化以及通訊消耗。
針對優化後的Operator Chain，runtime對應的Driver則是ChainedMapDriver。這是通過 MAP(MapDriver.class, ChainedMapDriver.class, PIPELINED, 0), 對映得到的。
我們可以看到，因為是任務鏈，所以每個record是直接在管道中流淌 ，ChainedMapDriver連迴圈都省略了，直接map轉換後丟給下游去也。
public class ChainedMapDriver<IT, OT> extends ChainedDriver<IT, OT> {
private MapFunction<IT, OT> mapper; // 使用者函式
@Override
public void collect(IT record) {
try {
this.numRecordsIn.inc();
this.outputCollector.collect(this.mapper.map(record));
} catch (Exception ex) {
throw new ExceptionInChainedStubException(this.taskName, ex);
}
}
}
// 這時的呼叫棧如下
map:23, UserFunc$1 (com.alibaba.alink)
collect:79, ChainedMapDriver (org.apache.flink.runtime.operators.chaining)
collect:35, CountingCollector (org.apache.flink.runtime.operators.util.metrics)
invoke:196, DataSourceTask (org.apache.flink.runtime.operators)
doRun:707, Task (org.apache.flink.runtime.taskmanager)
run:532, Task (org.apache.flink.runtime.taskmanager)
run:748, Thread (java.lang)
0x06 總結
map和mapPartition實現的基礎是Flink的資料傳輸機制 ：Flink確實是每次都處理一個輸入record，但是在上下游之間傳輸時候，Flink還是把records累積起來做批量傳輸。即可以認為從資料傳輸模型角度講，Flink是微批次的。
對於資料流轉換，因為是批量傳輸，所以對於積累的records，map是在runtime Driver程式碼中進行迴圈，mapPartition在使用者程式碼中進行迴圈。
map的函式呼叫次數要遠高於mapPartition。如果在使用者函式中涉及到頻繁建立額外的物件或者外部資源操作，則mapPartition效能遠遠高出。
如果沒有connection之類的操作，則通常效能差別並不大，通常不會成為瓶頸，也沒有想象的那麼嚴重。