Skip to content

Navigation Menu

Sign in
Sign up

[AURON #2307] Add @transient to large fields and use NativePartition to reduce serialization overhead #2515

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Deegue wants to merge 1 commit into apache:master
base: master
Choose a base branch
Loading
from Deegue:auron#2307_perf-add-transient-reduce-serialization
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ case class NativeShuffleExchangeExec(
"Input iterator must be empty (SPARK-44605: adapt to Spark 4+ ShuffleWriteProcessor API changes)")

val rdd = dep.asInstanceOf[AuronShuffleDependency[_, _, _]].inputRdd
val partition = rdd.partitions(mapIndex)
val partition = dep.asInstanceOf[AuronShuffleDependency[_, _, _]].rddPartitions(mapIndex)
internalWrite(rdd, dep, mapId, context, partition)
}

Expand Down
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import org.apache.auron.util.SparkVersionUtil
class NativeRDD(
@transient private val rddSparkContext: SparkContext,
val metrics: SparkMetricNode,
private val rddPartitions: Array[Partition],
@transient private val rddPartitions: Array[Partition],
private val rddPartitioner: Option[Partitioner],
private val rddDependencies: Seq[Dependency[_]],
private val rddShuffleReadFull: Boolean,
Expand All @@ -46,6 +46,8 @@ class NativeRDD(
with Logging
with Serializable {

private val numRddPartitions: Int = rddPartitions.length

// use serializable wrapper to avoid serializing nativePlan
val nativePlanWrapper = new NativePlanWrapper(nativePlan)

Expand All @@ -60,7 +62,10 @@ class NativeRDD(
def isShuffleReadFull: Boolean = Shims.get.getRDDShuffleReadFull(this)
Shims.get.setRDDShuffleReadFull(this, rddShuffleReadFull)

override protected def getPartitions: Array[Partition] = rddPartitions
// Spark 4+ needs the number of partitions
override protected def getPartitions: Array[Partition] =
if (rddPartitions != null) rddPartitions
else Array.tabulate(numRddPartitions)(i => new Partition { override def index: Int = i })
override protected def getDependencies: Seq[Dependency[_]] = rddDependencies
override val partitioner: Option[Partitioner] = rddPartitioner

Expand Down Expand Up @@ -101,6 +106,15 @@ class EmptyNativeRDD(@transient private val rddSparkContext: SparkContext)

}

case class NativePartition[P](override val index: Int, payload: P) extends Partition {}

object NativePartition {
def unwrap(partition: Partition): Partition = partition match {
case np: NativePartition[_] => np.payload.asInstanceOf[Partition]
case other => other
}
}

class NativePlanWrapper(var p: (Partition, TaskContext) => PhysicalPlanNode)
extends Serializable {
def plan(split: Partition, context: TaskContext): PhysicalPlanNode = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,11 @@ import scala.collection.immutable.SortedMap
import scala.jdk.CollectionConverters._

import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.internal.Logging
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.auron.Shims
Expand Down Expand Up @@ -174,18 +176,22 @@ abstract class NativeAggBase(
val nativeAggrModes = this.nativeAggrModes
val nativeAggrs = this.nativeAggrs
val nativeGroupingExprs = this.nativeGroupingExprs
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}

new NativeRDD(
sparkContext,
nativeMetrics,
rddPartitions = inputRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = inputRDD.partitioner,
rddDependencies = new OneToOneDependency(inputRDD) :: Nil,
inputRDD.isShuffleReadFull,
(partition, taskContext) => {
val inputPartition = NativePartition.unwrap(partition)

lazy val inputPlan =
inputRDD.nativePlan(inputRDD.partitions(partition.index), taskContext)
inputRDD.nativePlan(inputPartition, taskContext)
pb.PhysicalPlanNode
.newBuilder()
.setAgg(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,9 @@ import org.apache.auron.{protobuf => pb, sparkver}
import org.apache.auron.jni.JniBridge
import org.apache.auron.metric.SparkMetricNode

abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val child: SparkPlan)
abstract class NativeBroadcastExchangeBase(
@transient mode: BroadcastMode,
override val child: SparkPlan)
extends BroadcastExchangeLike
with NativeSupports {

Expand Down Expand Up @@ -245,7 +247,7 @@ abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val chi
metrics("dataSize") += byteArray.length
})

val input = inputRDD.nativePlan(inputRDD.partitions(split.index), context)
val input = inputRDD.nativePlan(split, context)
val nativeIpcWriterExec = pb.PhysicalPlanNode
.newBuilder()
.setIpcWriter(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.auron.Shims
Expand Down Expand Up @@ -153,25 +154,30 @@ abstract class NativeBroadcastJoinBase(
Seq(FullOuter, LeftOuter, LeftSemi, LeftAnti).contains(joinType)
})

val nativePartitions = probedRDD.partitions.map { p =>
NativePartition[Partition](p.index, p)
}

new NativeRDD(
sparkContext,
nativeMetrics,
probedRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = probedRDD.partitioner,
rddDependencies = new OneToOneDependency(probedRDD) :: Nil,
probedShuffleReadFull,
(partition, context) => {
val partition0 = new Partition() {
override def index: Int = 0
}
val probedPartition = NativePartition.unwrap(partition)
val (leftChild, rightChild) = broadcastSide match {
case JoinBuildLeft =>
(
leftRDD.nativePlan(partition0, context),
rightRDD.nativePlan(rightRDD.partitions(partition.index), context))
rightRDD.nativePlan(probedPartition, context))
case JoinBuildRight =>
(
leftRDD.nativePlan(leftRDD.partitions(partition.index), context),
leftRDD.nativePlan(probedPartition, context),
rightRDD.nativePlan(partition0, context))
}
val cachedBuildHashMapId = s"bhm_stage${context.stageId}_rdd${builtRDD.id}"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,12 @@ package org.apache.spark.sql.execution.auron.plan
import scala.collection.mutable

import org.apache.spark.OneToOneDependency
import org.apache.spark.sql.auron.{NativeHelper, NativeRDD, NativeSupports, Shims}
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.auron.Shims
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.physical.{SinglePartition, UnknownPartitioning}
Expand Down Expand Up @@ -56,16 +61,19 @@ abstract class NativeCollectLimitBase(limit: Int, offset: Int, override val chil
// merge all LocalLimit child partitions into a single partition
val shuffled = Shims.get.createNativeShuffleExchangeExec(SinglePartition, partial)
val singlePartitionRDD = NativeHelper.executeNative(shuffled)
val nativePartitions = singlePartitionRDD.partitions.map { p =>
NativePartition[Partition](p.index, p)
}

new NativeRDD(
sparkContext,
SparkMetricNode(metrics, singlePartitionRDD.metrics :: Nil),
singlePartitionRDD.partitions,
rddPartitions = nativePartitions.toArray,
singlePartitionRDD.partitioner,
new OneToOneDependency(singlePartitionRDD) :: Nil,
rddShuffleReadFull = false,
(partition, taskContext) => {
val inputPartition = singlePartitionRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeLimitExec = LimitExecNode
.newBuilder()
.setInput(singlePartitionRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,10 @@ import scala.collection.immutable.SortedMap
import scala.jdk.CollectionConverters._

import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, SortOrder}
Expand Down Expand Up @@ -75,16 +77,19 @@ abstract class NativeExpandBase(
val nativeMetrics = SparkMetricNode(metrics, inputRDD.metrics :: Nil)
val nativeSchema = this.nativeSchema
val nativeProjections = this.nativeProjections
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}

new NativeRDD(
sparkContext,
nativeMetrics,
rddPartitions = inputRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = inputRDD.partitioner,
rddDependencies = new OneToOneDependency(inputRDD) :: Nil,
inputRDD.isShuffleReadFull,
(partition, taskContext) => {
val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeExpandExec = ExpandExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ abstract class NativeFileSourceScanBase(basedFileScan: FileSourceScanExec)
override val output: Seq[Attribute] = basedFileScan.output
override val outputPartitioning: Partitioning = basedFileScan.outputPartitioning

@transient
protected val inputFileScanRDD: FileScanRDD = {
MethodUtils.invokeMethod(basedFileScan, true, "prepare")
MethodUtils.invokeMethod(basedFileScan, true, "waitForSubqueries")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,10 @@ import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._

import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.expressions.And
Expand Down Expand Up @@ -91,15 +93,18 @@ abstract class NativeFilterBase(condition: Expression, override val child: Spark
val inputRDD = NativeHelper.executeNative(child)
val nativeMetrics = SparkMetricNode(metrics, inputRDD.metrics :: Nil)
val nativeFilterExprs = this.nativeFilterExprs
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}
new NativeRDD(
sparkContext,
nativeMetrics,
rddPartitions = inputRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = inputRDD.partitioner,
rddDependencies = new OneToOneDependency(inputRDD) :: Nil,
inputRDD.isShuffleReadFull,
(partition, taskContext) => {
val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeFilterExec = FilterExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@ import scala.jdk.CollectionConverters._

import com.google.protobuf.ByteString
import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.expressions.Attribute
Expand Down Expand Up @@ -134,16 +136,19 @@ abstract class NativeGenerateBase(
val nativeGenerator = this.nativeGenerator
val nativeGeneratorOutput = this.nativeGeneratorOutput
val nativeRequiredChildOutput = this.nativeRequiredChildOutput
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}

new NativeRDD(
sparkContext,
nativeMetrics,
rddPartitions = inputRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = inputRDD.partitioner,
rddDependencies = new OneToOneDependency(inputRDD) :: Nil,
inputRDD.isShuffleReadFull,
(partition, taskContext) => {
val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeGenerateExec = pb.GenerateExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@ package org.apache.spark.sql.execution.auron.plan
import scala.collection.immutable.SortedMap

import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.expressions.Attribute
Expand Down Expand Up @@ -52,16 +54,19 @@ abstract class NativeGlobalLimitBase(limit: Int, offset: Int, override val child
override def doExecuteNative(): NativeRDD = {
val inputRDD = NativeHelper.executeNative(child)
val nativeMetrics = SparkMetricNode(metrics, inputRDD.metrics :: Nil)
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}

new NativeRDD(
sparkContext,
nativeMetrics,
inputRDD.partitions,
nativePartitions.toArray,
inputRDD.partitioner,
new OneToOneDependency(inputRDD) :: Nil,
rddShuffleReadFull = false,
(partition, taskContext) => {
val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeLimitExec = LimitExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@ package org.apache.spark.sql.execution.auron.plan
import scala.collection.immutable.SortedMap

import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.expressions.Attribute
Expand Down Expand Up @@ -49,16 +51,19 @@ abstract class NativeLocalLimitBase(limit: Int, override val child: SparkPlan)
override def doExecuteNative(): NativeRDD = {
val inputRDD = NativeHelper.executeNative(child)
val nativeMetrics = SparkMetricNode(metrics, inputRDD.metrics :: Nil)
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}

new NativeRDD(
sparkContext,
nativeMetrics,
rddPartitions = inputRDD.partitions,
rddPartitions = nativePartitions.toArray,
rddPartitioner = inputRDD.partitioner,
rddDependencies = new OneToOneDependency(inputRDD) :: Nil,
rddShuffleReadFull = false,
(partition, taskContext) => {
val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val nativeLimitExec = LimitExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, taskContext))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,11 @@ import org.apache.hadoop.hive.ql.plan.TableDesc
import org.apache.hadoop.mapred.JobConf
import org.apache.hadoop.mapreduce.Job
import org.apache.spark.OneToOneDependency
import org.apache.spark.Partition
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativePartition
import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.auron.NativeSupports
import org.apache.spark.sql.catalyst.catalog.CatalogTable
Expand Down Expand Up @@ -85,10 +87,13 @@ abstract class NativeOrcSinkBase(
val inputRDD = NativeHelper.executeNative(child)
val nativeMetrics = SparkMetricNode(metrics, inputRDD.metrics :: Nil)
val nativeDependencies = new OneToOneDependency(inputRDD) :: Nil
val nativePartitions = inputRDD.partitions.map { inputPartition =>
NativePartition[Partition](inputPartition.index, inputPartition)
}
new NativeRDD(
sparkSession.sparkContext,
nativeMetrics,
inputRDD.partitions,
nativePartitions.toArray,
inputRDD.partitioner,
nativeDependencies,
inputRDD.isShuffleReadFull,
Expand Down Expand Up @@ -116,7 +121,7 @@ abstract class NativeOrcSinkBase(
.setValue(entry.getValue)
.build())

val inputPartition = inputRDD.partitions(partition.index)
val inputPartition = NativePartition.unwrap(partition)
val orcSink = OrcSinkExecNode
.newBuilder()
.setInput(inputRDD.nativePlan(inputPartition, context))
Expand Down
Loading
Loading

AltStyle によって変換されたページ (->オリジナル) /