- CodegenSupport(接口)
相邻Operator通过Produce-Consume模式生成代码。
Produce生成整体处理的框架代码，例如aggregation生成的代码框架如下:
if (!initialized) {
# create a hash map, then build the aggregation hash map
# call child.produce()
initialized = true;
}
while (hashmap.hasNext()) {
row = hashmap.next();
# build the aggregation results
# create variables for results
# call consume(), which will call parent.doConsume()
if (shouldStop()) return;
}
Consume生成当前节点处理上游输入的Row的逻辑。如Filter生成代码如下:
# code to evaluate the predicate expression, result is isNull1 and value2
if (!isNull1 && value2) {
# call consume(), which will call parent.doConsume()
}
- WholeStageCodegenExec(类)
CodegenSupport的实现类之一，Stage内部所有相邻的实现CodegenSupport接口的Operator的融合，产出的代码把所有被融合的Operator的执行逻辑封装到一个Wrapper类中，该Wrapper类作为Janino即时compile的入参。 - InputAdapter(类)
CodegenSupport的实现类之一，胶水类，用来连接WholeStageCodegenExec节点和未实现CodegenSupport的上游节点。 - BufferedRowIterator(接口)
WholeStageCodegenExec生成的java代码的父类，重要方法:
public InternalRow next() // 返回下一条Row public void append(InternalRow row) // append一条RowSimple call graph
Produce-Consume Pattern
doProduce() doConsume() 会被子类覆写
produce() consume() 均为 trait CodegenSupport extends SparkPlan 的 final 方法
CodegenSupport的实现类之一，胶水类，用来连接WholeStageCodegenExec节点和未实现CodegenSupport的上游节点。
private def insertInputAdapter(plan: SparkPlan): SparkPlan = {
plan match {
case p if !supportCodegen(p) =>
// collapse them recursively
InputAdapter(insertWholeStageCodegen(p))
case j: SortMergeJoinExec =>
// The children of SortMergeJoin should do codegen separately.
j.withNewChildren(j.children.map(
child => InputAdapter(insertWholeStageCodegen(child))))
case j: ShuffledHashJoinExec =>
// The children of ShuffledHashJoin should do codegen separately.
j.withNewChildren(j.children.map(
child => InputAdapter(insertWholeStageCodegen(child))))
case p => p.withNewChildren(p.children.map(insertInputAdapter))
}
}
UT
test("range/filter should be combined") {
val df = spark.range(10).filter("id = 1").selectExpr("id + 1")
val plan = df.queryExecution.executedPlan
assert(plan.find(_.isInstanceOf[WholeStageCodegenExec]).isDefined)
assert(df.collect() === Array(Row(2)))
df.explain(false)
df.queryExecution.debug.codegen
}
11:32:34.837 WARN org.apache.hadoop.util.NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
== Physical Plan ==
*(1) Project [(id#0L + 1) AS (id + 1)#4L]
+- *(1) Filter (id#0L = 1)
+- *(1) Range (0, 10, step=1, splits=2)
Found 1 WholeStageCodegen subtrees.
== Subtree 1 / 1 (maxMethodCodeSize:301; maxConstantPoolSize:177(0.27% used); numInnerClasses:0) ==
*(1) Project [(id#0L + 1) AS (id + 1)#4L]
+- *(1) Filter (id#0L = 1)
+- *(1) Range (0, 10, step=1, splits=2)
Generated code:
public Object generate(Object[] references) {
return new GeneratedIteratorForCodegenStage1(references);
}
// codegenStageId=1
final class GeneratedIteratorForCodegenStage1 extends org.apache.spark.sql.execution.BufferedRowIterator {
private Object[] references;
private scala.collection.Iterator[] inputs;
private boolean range_initRange_0;
private long range_nextIndex_0;
private TaskContext range_taskContext_0;
private InputMetrics range_inputMetrics_0;
private long range_batchEnd_0;
private long range_numElementsTodo_0;
private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] range_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[3];
public GeneratedIteratorForCodegenStage1(Object[] references) {
this.references = references;
}
public void init(int index, scala.collection.Iterator[] inputs) {
partitionIndex = index;
this.inputs = inputs;
range_taskContext_0 = TaskContext.get();
range_inputMetrics_0 = range_taskContext_0.taskMetrics().inputMetrics();
range_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
range_mutableStateArray_0[1] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
range_mutableStateArray_0[2] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
}
private void initRange(int idx) {
java.math.BigInteger index = java.math.BigInteger.valueOf(idx);
java.math.BigInteger numSlice = java.math.BigInteger.valueOf(2L);
java.math.BigInteger numElement = java.math.BigInteger.valueOf(10L);
java.math.BigInteger step = java.math.BigInteger.valueOf(1L);
java.math.BigInteger start = java.math.BigInteger.valueOf(0L);
long partitionEnd;
java.math.BigInteger st = index.multiply(numElement).divide(numSlice).multiply(step).add(start);
if (st.compareTo(java.math.BigInteger.valueOf(Long.MAX_VALUE)) > 0) {
range_nextIndex_0 = Long.MAX_VALUE;
} else if (st.compareTo(java.math.BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
range_nextIndex_0 = Long.MIN_VALUE;
} else {
range_nextIndex_0 = st.longValue();
}
range_batchEnd_0 = range_nextIndex_0;
java.math.BigInteger end = index.add(java.math.BigInteger.ONE).multiply(numElement).divide(numSlice)
.multiply(step).add(start);
if (end.compareTo(java.math.BigInteger.valueOf(Long.MAX_VALUE)) > 0) {
partitionEnd = Long.MAX_VALUE;
} else if (end.compareTo(java.math.BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
partitionEnd = Long.MIN_VALUE;
} else {
partitionEnd = end.longValue();
}
java.math.BigInteger startToEnd = java.math.BigInteger.valueOf(partitionEnd).subtract(
java.math.BigInteger.valueOf(range_nextIndex_0));
range_numElementsTodo_0 = startToEnd.divide(step).longValue();
if (range_numElementsTodo_0 < 0) {
range_numElementsTodo_0 = 0;
} else if (startToEnd.remainder(step).compareTo(java.math.BigInteger.valueOf(0L)) != 0) {
range_numElementsTodo_0++;
}
}
protected void processNext() throws java.io.IOException {
// initialize Range
if (!range_initRange_0) {
range_initRange_0 = true;
initRange(partitionIndex);
}
while (true) {
if (range_nextIndex_0 == range_batchEnd_0) {
long range_nextBatchTodo_0;
if (range_numElementsTodo_0 > 1000L) {
range_nextBatchTodo_0 = 1000L;
range_numElementsTodo_0 -= 1000L;
} else {
range_nextBatchTodo_0 = range_numElementsTodo_0;
range_numElementsTodo_0 = 0;
if (range_nextBatchTodo_0 == 0) break;
}
range_batchEnd_0 += range_nextBatchTodo_0 * 1L;
}
int range_localEnd_0 = (int)((range_batchEnd_0 - range_nextIndex_0) / 1L);
for (int range_localIdx_0 = 0; range_localIdx_0 < range_localEnd_0; range_localIdx_0++) {
long range_value_0 = ((long)range_localIdx_0 * 1L) + range_nextIndex_0;
do {
boolean filter_value_0 = false;
filter_value_0 = range_value_0 == 1L;
if (!filter_value_0) continue;
((org.apache.spark.sql.execution.metric.SQLMetric) references[1] ).add(1);
// common sub-expressions
long project_value_0 = -1L;
project_value_0 = range_value_0 + 1L;
range_mutableStateArray_0[2].reset();
range_mutableStateArray_0[2].write(0, project_value_0);
append((range_mutableStateArray_0[2].getRow()));
} while(false);
if (shouldStop()) {
range_nextIndex_0 = range_value_0 + 1L;
((org.apache.spark.sql.execution.metric.SQLMetric) references[0] ).add(range_localIdx_0 + 1);
range_inputMetrics_0.incRecordsRead(range_localIdx_0 + 1);
return;
}
}
range_nextIndex_0 = range_batchEnd_0;
((org.apache.spark.sql.execution.metric.SQLMetric) references[0] ).add(range_localEnd_0);
range_inputMetrics_0.incRecordsRead(range_localEnd_0);
range_taskContext_0.killTaskIfInterrupted();
}
}
}
11:32:40.126 WARN org.apache.spark.sql.execution.WholeStageCodegenSuite:
===== POSSIBLE THREAD LEAK IN SUITE o.a.s.sql.execution.WholeStageCodegenSuite, thread names: rpc-boss-3-1, shuffle-boss-6-1 =====
Process finished with exit code 0
一个简单的实验
case class FilterExec(condition: Expression, child: SparkPlan)
在FilterExec 增加这一行
override def supportCodegen: Boolean = false
会拆成两个wholestagecodegen
16:27:42.332 WARN org.apache.hadoop.util.NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
== Physical Plan ==
*(2) Project [(id#0L + 1) AS (id + 1)#4L]
+- Filter (id#0L = 1)
+- *(1) Range (0, 10, step=1, splits=2)
Found 2 WholeStageCodegen subtrees.
== Subtree 1 / 2 (maxMethodCodeSize:282; maxConstantPoolSize:175(0.27% used); numInnerClasses:0) ==
*(1) Range (0, 10, step=1, splits=2)
Generated code:
public Object generate(Object[] references) {
return new GeneratedIteratorForCodegenStage1(references);
}
// codegenStageId=1
final class GeneratedIteratorForCodegenStage1 extends org.apache.spark.sql.execution.BufferedRowIterator {
private Object[] references;
private scala.collection.Iterator[] inputs;
private boolean range_initRange_0;
private long range_nextIndex_0;
private TaskContext range_taskContext_0;
private InputMetrics range_inputMetrics_0;
private long range_batchEnd_0;
private long range_numElementsTodo_0;
private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] range_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[1];
public GeneratedIteratorForCodegenStage1(Object[] references) {
this.references = references;
}
public void init(int index, scala.collection.Iterator[] inputs) {
partitionIndex = index;
this.inputs = inputs;
range_taskContext_0 = TaskContext.get();
range_inputMetrics_0 = range_taskContext_0.taskMetrics().inputMetrics();
range_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
}
private void initRange(int idx) {
java.math.BigInteger index = java.math.BigInteger.valueOf(idx);
java.math.BigInteger numSlice = java.math.BigInteger.valueOf(2L);
java.math.BigInteger numElement = java.math.BigInteger.valueOf(10L);
java.math.BigInteger step = java.math.BigInteger.valueOf(1L);
java.math.BigInteger start = java.math.BigInteger.valueOf(0L);
long partitionEnd;
java.math.BigInteger st = index.multiply(numElement).divide(numSlice).multiply(step).add(start);
if (st.compareTo(java.math.BigInteger.valueOf(Long.MAX_VALUE)) > 0) {
range_nextIndex_0 = Long.MAX_VALUE;
} else if (st.compareTo(java.math.BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
range_nextIndex_0 = Long.MIN_VALUE;
} else {
range_nextIndex_0 = st.longValue();
}
range_batchEnd_0 = range_nextIndex_0;
java.math.BigInteger end = index.add(java.math.BigInteger.ONE).multiply(numElement).divide(numSlice)
.multiply(step).add(start);
if (end.compareTo(java.math.BigInteger.valueOf(Long.MAX_VALUE)) > 0) {
partitionEnd = Long.MAX_VALUE;
} else if (end.compareTo(java.math.BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
partitionEnd = Long.MIN_VALUE;
} else {
partitionEnd = end.longValue();
}
java.math.BigInteger startToEnd = java.math.BigInteger.valueOf(partitionEnd).subtract(
java.math.BigInteger.valueOf(range_nextIndex_0));
range_numElementsTodo_0 = startToEnd.divide(step).longValue();
if (range_numElementsTodo_0 < 0) {
range_numElementsTodo_0 = 0;
} else if (startToEnd.remainder(step).compareTo(java.math.BigInteger.valueOf(0L)) != 0) {
range_numElementsTodo_0++;
}
}
protected void processNext() throws java.io.IOException {
// initialize Range
if (!range_initRange_0) {
range_initRange_0 = true;
initRange(partitionIndex);
}
while (true) {
if (range_nextIndex_0 == range_batchEnd_0) {
long range_nextBatchTodo_0;
if (range_numElementsTodo_0 > 1000L) {
range_nextBatchTodo_0 = 1000L;
range_numElementsTodo_0 -= 1000L;
} else {
range_nextBatchTodo_0 = range_numElementsTodo_0;
range_numElementsTodo_0 = 0;
if (range_nextBatchTodo_0 == 0) break;
}
range_batchEnd_0 += range_nextBatchTodo_0 * 1L;
}
int range_localEnd_0 = (int)((range_batchEnd_0 - range_nextIndex_0) / 1L);
for (int range_localIdx_0 = 0; range_localIdx_0 < range_localEnd_0; range_localIdx_0++) {
long range_value_0 = ((long)range_localIdx_0 * 1L) + range_nextIndex_0;
range_mutableStateArray_0[0].reset();
range_mutableStateArray_0[0].write(0, range_value_0);
append((range_mutableStateArray_0[0].getRow()));
if (shouldStop()) {
range_nextIndex_0 = range_value_0 + 1L;
((org.apache.spark.sql.execution.metric.SQLMetric) references[0] ).add(range_localIdx_0 + 1);
range_inputMetrics_0.incRecordsRead(range_localIdx_0 + 1);
return;
}
}
range_nextIndex_0 = range_batchEnd_0;
((org.apache.spark.sql.execution.metric.SQLMetric) references[0] ).add(range_localEnd_0);
range_inputMetrics_0.incRecordsRead(range_localEnd_0);
range_taskContext_0.killTaskIfInterrupted();
}
}
}
== Subtree 2 / 2 (maxMethodCodeSize:89; maxConstantPoolSize:91(0.14% used); numInnerClasses:0) ==
*(2) Project [(id#0L + 1) AS (id + 1)#4L]
+- Filter (id#0L = 1)
+- *(1) Range (0, 10, step=1, splits=2)
Generated code:
public Object generate(Object[] references) {
return new GeneratedIteratorForCodegenStage2(references);
}
// codegenStageId=2
final class GeneratedIteratorForCodegenStage2 extends org.apache.spark.sql.execution.BufferedRowIterator {
private Object[] references;
private scala.collection.Iterator[] inputs;
private scala.collection.Iterator inputadapter_input_0;
private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] project_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[1];
public GeneratedIteratorForCodegenStage2(Object[] references) {
this.references = references;
}
public void init(int index, scala.collection.Iterator[] inputs) {
partitionIndex = index;
this.inputs = inputs;
inputadapter_input_0 = inputs[0];
project_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
}
protected void processNext() throws java.io.IOException {
while ( inputadapter_input_0.hasNext()) {
InternalRow inputadapter_row_0 = (InternalRow) inputadapter_input_0.next();
// common sub-expressions
long inputadapter_value_0 = inputadapter_row_0.getLong(0);
long project_value_0 = -1L;
project_value_0 = inputadapter_value_0 + 1L;
project_mutableStateArray_0[0].reset();
project_mutableStateArray_0[0].write(0, project_value_0);
append((project_mutableStateArray_0[0].getRow()));
if (shouldStop()) return;
}
}
}
16:27:47.464 WARN org.apache.spark.sql.execution.WholeStageCodegenSuite:
===== POSSIBLE THREAD LEAK IN SUITE o.a.s.sql.execution.WholeStageCodegenSuite, thread names: rpc-boss-3-1, shuffle-boss-6-1 =====
Process finished with exit code 0
最里侧的operator
比如rangeExec， 肯定要实现 doproduce方法，但 consume不需要实现，直接调用父类的consume()
case class RangeExec(range: org.apache.spark.sql.catalyst.plans.logical.Range)
extends LeafExecNode with CodegenSupport {
val start: Long = range.start
val end: Long = range.end
val step: Long = range.step
val numSlices: Int = range.numSlices.getOrElse(sparkContext.defaultParallelism)
val numElements: BigInt = range.numElements
val isEmptyRange: Boolean = start == end || (start < end ^ 0 < step)
override val output: Seq[Attribute] = range.output
override def outputOrdering: Seq[SortOrder] = range.outputOrdering
override def outputPartitioning: Partitioning = {
if (numElements > 0) {
if (numSlices == 1) {
SinglePartition
} else {
RangePartitioning(outputOrdering, numSlices)
}
} else {
UnknownPartitioning(0)
}
}
override lazy val metrics = Map(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
override def doCanonicalize(): SparkPlan = {
RangeExec(range.canonicalized.asInstanceOf[org.apache.spark.sql.catalyst.plans.logical.Range])
}
override def inputRDDs(): Seq[RDD[InternalRow]] = {
val rdd = if (isEmptyRange) {
new EmptyRDD[InternalRow](sqlContext.sparkContext)
} else {
sqlContext.sparkContext.parallelize(0 until numSlices, numSlices).map(i => InternalRow(i))
}
rdd :: Nil
}
protected override def doProduce(ctx: CodegenContext): String = {
val numOutput = metricTerm(ctx, "numOutputRows")
val initTerm = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initRange")
val nextIndex = ctx.addMutableState(CodeGenerator.JAVA_LONG, "nextIndex")
val value = ctx.freshName("value")
val ev = ExprCode.forNonNullValue(JavaCode.variable(value, LongType))
val BigInt = classOf[java.math.BigInteger].getName
// Inline mutable state since not many Range operations in a task
val taskContext = ctx.addMutableState("TaskContext", "taskContext",
v => s"$v = TaskContext.get();", forceInline = true)
val inputMetrics = ctx.addMutableState("InputMetrics", "inputMetrics",
v => s"$v = $taskContext.taskMetrics().inputMetrics();", forceInline = true)
// In order to periodically update the metrics without inflicting performance penalty, this
// operator produces elements in batches. After a batch is complete, the metrics are updated
// and a new batch is started.
// In the implementation below, the code in the inner loop is producing all the values
// within a batch, while the code in the outer loop is setting batch parameters and updating
// the metrics.
// Once nextIndex == batchEnd, it's time to progress to the next batch.
val batchEnd = ctx.addMutableState(CodeGenerator.JAVA_LONG, "batchEnd")
// How many values should still be generated by this range operator.
val numElementsTodo = ctx.addMutableState(CodeGenerator.JAVA_LONG, "numElementsTodo")
// How many values should be generated in the next batch.
val nextBatchTodo = ctx.freshName("nextBatchTodo")
// The default size of a batch, which must be positive integer
val batchSize = 1000
val initRangeFuncName = ctx.addNewFunction("initRange",
s"""
| private void initRange(int idx) {
| $BigInt index = $BigInt.valueOf(idx);
| $BigInt numSlice = $BigInt.valueOf(${numSlices}L);
| $BigInt numElement = $BigInt.valueOf(${numElements.toLong}L);
| $BigInt step = $BigInt.valueOf(${step}L);
| $BigInt start = $BigInt.valueOf(${start}L);
| long partitionEnd;
|
| $BigInt st = index.multiply(numElement).divide(numSlice).multiply(step).add(start);
| if (st.compareTo($BigInt.valueOf(Long.MAX_VALUE)) > 0) {
| $nextIndex = Long.MAX_VALUE;
| } else if (st.compareTo($BigInt.valueOf(Long.MIN_VALUE)) < 0) {
| $nextIndex = Long.MIN_VALUE;
| } else {
| $nextIndex = st.longValue();
| }
| $batchEnd = $nextIndex;
|
| $BigInt end = index.add($BigInt.ONE).multiply(numElement).divide(numSlice)
| .multiply(step).add(start);
| if (end.compareTo($BigInt.valueOf(Long.MAX_VALUE)) > 0) {
| partitionEnd = Long.MAX_VALUE;
| } else if (end.compareTo($BigInt.valueOf(Long.MIN_VALUE)) < 0) {
| partitionEnd = Long.MIN_VALUE;
| } else {
| partitionEnd = end.longValue();
| }
|
| $BigInt startToEnd = $BigInt.valueOf(partitionEnd).subtract(
| $BigInt.valueOf($nextIndex));
| $numElementsTodo = startToEnd.divide(step).longValue();
| if ($numElementsTodo < 0) {
| $numElementsTodo = 0;
| } else if (startToEnd.remainder(step).compareTo($BigInt.valueOf(0L)) != 0) {
| $numElementsTodo++;
| }
| }
""".stripMargin)
val localIdx = ctx.freshName("localIdx")
val localEnd = ctx.freshName("localEnd")
val stopCheck = if (parent.needStopCheck) {
s"""
|if (shouldStop()) {
| $nextIndex = $value + ${step}L;
| $numOutput.add($localIdx + 1);
| $inputMetrics.incRecordsRead($localIdx + 1);
| return;
|}
""".stripMargin
} else {
"// shouldStop check is eliminated"
}
val loopCondition = if (limitNotReachedChecks.isEmpty) {
"true"
} else {
limitNotReachedChecks.mkString(" && ")
}
// An overview of the Range processing.
//
// For each partition, the Range task needs to produce records from partition start(inclusive)
// to end(exclusive). For better performance, we separate the partition range into batches, and
// use 2 loops to produce data. The outer while loop is used to iterate batches, and the inner
// for loop is used to iterate records inside a batch.
//
// `nextIndex` tracks the index of the next record that is going to be consumed, initialized
// with partition start. `batchEnd` tracks the end index of the current batch, initialized
// with `nextIndex`. In the outer loop, we first check if `nextIndex == batchEnd`. If it's true,
// it means the current batch is fully consumed, and we will update `batchEnd` to process the
// next batch. If `batchEnd` reaches partition end, exit the outer loop. Finally we enter the
// inner loop. Note that, when we enter inner loop, `nextIndex` must be different from
// `batchEnd`, otherwise we already exit the outer loop.
//
// The inner loop iterates from 0 to `localEnd`, which is calculated by
// `(batchEnd - nextIndex) / step`. Since `batchEnd` is increased by `nextBatchTodo * step` in
// the outer loop, and initialized with `nextIndex`, so `batchEnd - nextIndex` is always
// divisible by `step`. The `nextIndex` is increased by `step` during each iteration, and ends
// up being equal to `batchEnd` when the inner loop finishes.
//
// The inner loop can be interrupted, if the query has produced at least one result row, so that
// we don't buffer too many result rows and waste memory. It's ok to interrupt the inner loop,
// because `nextIndex` will be updated before interrupting.
s"""
| // initialize Range
| if (!$initTerm) {
| $initTerm = true;
| $initRangeFuncName(partitionIndex);
| }
|
| while ($loopCondition) {
| if ($nextIndex == $batchEnd) {
| long $nextBatchTodo;
| if ($numElementsTodo > ${batchSize}L) {
| $nextBatchTodo = ${batchSize}L;
| $numElementsTodo -= ${batchSize}L;
| } else {
| $nextBatchTodo = $numElementsTodo;
| $numElementsTodo = 0;
| if ($nextBatchTodo == 0) break;
| }
| $batchEnd += $nextBatchTodo * ${step}L;
| }
|
| int $localEnd = (int)(($batchEnd - $nextIndex) / ${step}L);
| for (int $localIdx = 0; $localIdx < $localEnd; $localIdx++) {
| long $value = ((long)$localIdx * ${step}L) + $nextIndex;
| ${consume(ctx, Seq(ev))}
| $stopCheck
| }
| $nextIndex = $batchEnd;
| $numOutput.add($localEnd);
| $inputMetrics.incRecordsRead($localEnd);
| $taskContext.killTaskIfInterrupted();
| }
""".stripMargin
}
protected override def doExecute(): RDD[InternalRow] = {
val numOutputRows = longMetric("numOutputRows")
if (isEmptyRange) {
new EmptyRDD[InternalRow](sqlContext.sparkContext)
} else {
sqlContext
.sparkContext
.parallelize(0 until numSlices, numSlices)
.mapPartitionsWithIndex { (i, _) =>
val partitionStart = (i * numElements) / numSlices * step + start
val partitionEnd = (((i + 1) * numElements) / numSlices) * step + start
def getSafeMargin(bi: BigInt): Long =
if (bi.isValidLong) {
bi.toLong
} else if (bi > 0) {
Long.MaxValue
} else {
Long.MinValue
}
val safePartitionStart = getSafeMargin(partitionStart)
val safePartitionEnd = getSafeMargin(partitionEnd)
val rowSize = UnsafeRow.calculateBitSetWidthInBytes(1) + LongType.defaultSize
val unsafeRow = UnsafeRow.createFromByteArray(rowSize, 1)
val taskContext = TaskContext.get()
val iter = new Iterator[InternalRow] {
private[this] var number: Long = safePartitionStart
private[this] var overflow: Boolean = false
private[this] val inputMetrics = taskContext.taskMetrics().inputMetrics
override def hasNext =
if (!overflow) {
if (step > 0) {
number < safePartitionEnd
} else {
number > safePartitionEnd
}
} else false
override def next() = {
val ret = number
number += step
if (number < ret ^ step < 0) {
// we have Long.MaxValue + Long.MaxValue < Long.MaxValue
// and Long.MinValue + Long.MinValue > Long.MinValue, so iff the step causes a step
// back, we are pretty sure that we have an overflow.
overflow = true
}
numOutputRows += 1
inputMetrics.incRecordsRead(1)
unsafeRow.setLong(0, ret)
unsafeRow
}
}
new InterruptibleIterator(taskContext, iter)
}
}
}
override def simpleString(maxFields: Int): String = {
s"Range ($start, $end, step=$step, splits=$numSlices)"
