Skip to content

Commit

Permalink
Specify an ordering for fields in SubmitDriverRequestMessage
Browse files Browse the repository at this point in the history
Previously APP_ARGs, SPARK_PROPERTYs and ENVIRONMENT_VARIABLEs
will appear in the JSON at random places. Now they are grouped
together at the end of the JSON blob.
  • Loading branch information
Andrew Or committed Jan 20, 2015
1 parent 6ff088d commit 484bd21
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,8 @@ private[spark] class SubmitDriverRequestMessage extends SubmitRestProtocolMessag
SubmitDriverRequestField.ACTION,
SubmitDriverRequestField.requiredFields) {

import SubmitDriverRequestField._

// Ensure continuous range of app arg indices starting from 0
override def validate(): this.type = {
import SubmitDriverRequestField._
Expand All @@ -93,6 +95,22 @@ private[spark] class SubmitDriverRequestMessage extends SubmitRestProtocolMessag
}
super.validate()
}

// List the fields in the following order:
// ACTION < SPARK_VERSION < * < APP_ARG < SPARK_PROPERTY < ENVIRONMENT_VARIABLE < MESSAGE
protected override def sortedFields: Seq[(SubmitRestProtocolField, String)] = {
fields.toSeq.sortBy { case (k, _) =>
k match {
case ACTION => 0
case SPARK_VERSION => 1
case APP_ARG(index) => 10 + index
case SPARK_PROPERTY(propKey) => 100
case ENVIRONMENT_VARIABLE(envKey) => 1000
case MESSAGE => Int.MaxValue
case _ => 2
}
}
}
}

private[spark] object SubmitDriverRequestMessage extends SubmitRestProtocolMessageCompanion {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,16 +24,17 @@ import org.json4s.jackson.JsonMethods._
import org.json4s.JsonAST._

import org.apache.spark.{Logging, SparkException}
import org.apache.spark.util.Utils
import org.apache.spark.util.{JsonProtocol, Utils}

/**
* A field used in a SubmitRestProtocolMessage.
* Three special fields ACTION, SPARK_VERSION, and MESSAGE are common across all messages.
*/
private[spark] abstract class SubmitRestProtocolField
private[spark] object SubmitRestProtocolField {
/** Return whether the provided field name refers to the ACTION field. */
def isActionField(field: String): Boolean = field == "ACTION"
def isSparkVersionField(field: String): Boolean = field == "SPARK_VERSION"
def isMessageField(field: String): Boolean = field == "MESSAGE"
}

/**
Expand Down Expand Up @@ -125,23 +126,26 @@ private[spark] abstract class SubmitRestProtocolMessage(

/** Return the JSON representation of this message. */
def toJson: String = {
val stringFields = fields
val jsonFields = sortedFields
.filter { case (_, v) => v != null }
.map { case (k, v) => (k.toString, v) }
val jsonFields = fieldsToJson(stringFields)
pretty(render(jsonFields))
.map { case (k, v) => JField(k.toString, JString(v)) }
.toList
pretty(render(JObject(jsonFields)))
}

/**
* Return the JSON representation of the message fields, putting ACTION first.
* This assumes that applying `org.apache.spark.util.JsonProtocol.mapFromJson`
* to the result yields the original input.
* Return a list of (field, value) pairs with the following ordering:
* ACTION < SPARK_VERSION < * < MESSAGE
*/
private def fieldsToJson(fields: Map[String, String]): JValue = {
val jsonFields = fields.toList
.sortBy { case (k, _) => if (isActionField(k)) 0 else 1 }
.map { case (k, v) => JField(k, JString(v)) }
JObject(jsonFields)
protected def sortedFields: Seq[(SubmitRestProtocolField, String)] = {
fields.toSeq.sortBy { case (k, _) =>
k.toString match {
case x if isActionField(x) => 0
case x if isSparkVersionField(x) => 1
case x if isMessageField(x) => Int.MaxValue
case _ => 2
}
}
}
}

Expand All @@ -155,7 +159,7 @@ private[spark] object SubmitRestProtocolMessage {
* If such a field does not exist in the JSON, throw an exception.
*/
def fromJson(json: String): SubmitRestProtocolMessage = {
val fields = org.apache.spark.util.JsonProtocol.mapFromJson(parse(json))
val fields = JsonProtocol.mapFromJson(parse(json))
val action = fields
.flatMap { case (k, v) => if (isActionField(k)) Some(v) else None }
.headOption
Expand Down

0 comments on commit 484bd21

Please sign in to comment.