[SPARK-18891][SQL] Support for Scala Map collection types #16986
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,6 +22,7 @@ import java.lang.reflect.Modifier | |
| import scala.collection.mutable.Builder | ||
| import scala.language.existentials | ||
| import scala.reflect.ClassTag | ||
| import scala.util.Try | ||
|
|
||
| import org.apache.spark.{SparkConf, SparkEnv} | ||
| import org.apache.spark.serializer._ | ||
|
|
@@ -652,6 +653,218 @@ case class MapObjects private( | |
| } | ||
| } | ||
|
|
||
| object CollectObjectsToMap { | ||
| private val curId = new java.util.concurrent.atomic.AtomicInteger() | ||
|
|
||
| /** | ||
| * Construct an instance of CollectObjects case class. | ||
| * | ||
| * @param keyFunction The function applied on the key collection elements. | ||
| * @param valueFunction The function applied on the value collection elements. | ||
| * @param inputData An expression that when evaluated returns a map object. | ||
| * @param collClass The type of the resulting collection. | ||
| */ | ||
| def apply( | ||
| keyFunction: Expression => Expression, | ||
| valueFunction: Expression => Expression, | ||
| inputData: Expression, | ||
| collClass: Class[_]): CollectObjectsToMap = { | ||
| val id = curId.getAndIncrement() | ||
| val keyLoopValue = s"CollectObjectsToMap_keyLoopValue$id" | ||
| val mapType = inputData.dataType.asInstanceOf[MapType] | ||
| val keyLoopVar = LambdaVariable(keyLoopValue, "", mapType.keyType, nullable = false) | ||
| val valueLoopValue = s"CollectObjectsToMap_valueLoopValue$id" | ||
| val valueLoopIsNull = s"CollectObjectsToMap_valueLoopIsNull$id" | ||
| val valueLoopVar = LambdaVariable(valueLoopValue, valueLoopIsNull, mapType.valueType) | ||
| val tupleLoopVar = s"CollectObjectsToMap_tupleLoopValue$id" | ||
| val builderValue = s"CollectObjectsToMap_builderValue$id" | ||
|
||
| CollectObjectsToMap( | ||
| keyLoopValue, keyFunction(keyLoopVar), | ||
| valueLoopValue, valueLoopIsNull, valueFunction(valueLoopVar), | ||
| inputData, | ||
| tupleLoopVar, collClass, builderValue) | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * An equivalent to the [[MapObjects]] case class but returning an ObjectType containing | ||
|
||
| * a Scala collection constructed using the associated builder, obtained by calling `newBuilder` | ||
| * on the collection's companion object. | ||
| * | ||
| * @param keyLoopValue the name of the loop variable that is used when iterating over the key | ||
| * collection, and which is used as input for the `keyLambdaFunction` | ||
| * @param keyLambdaFunction A function that takes the `keyLoopVar` as input, and is used as | ||
| * a lambda function to handle collection elements. | ||
| * @param valueLoopValue the name of the loop variable that is used when iterating over the value | ||
| * collection, and which is used as input for the `valueLambdaFunction` | ||
| * @param valueLoopIsNull the nullability of the loop variable that is used when iterating over | ||
| * the value collection, and which is used as input for the | ||
| * `valueLambdaFunction` | ||
| * @param valueLambdaFunction A function that takes the `valueLoopVar` as input, and is used as | ||
| * a lambda function to handle collection elements. | ||
| * @param inputData An expression that when evaluated returns a map object. | ||
| * @param tupleLoopValue the name of the loop variable that holds the tuple to be added to the | ||
| * resulting map (used only for Scala Map) | ||
| * @param collClass The type of the resulting collection. | ||
| * @param builderValue The name of the builder variable used to construct the resulting collection. | ||
| */ | ||
| case class CollectObjectsToMap private( | ||
| keyLoopValue: String, | ||
| keyLambdaFunction: Expression, | ||
| valueLoopValue: String, | ||
| valueLoopIsNull: String, | ||
| valueLambdaFunction: Expression, | ||
| inputData: Expression, | ||
| tupleLoopValue: String, | ||
| collClass: Class[_], | ||
| builderValue: String) extends Expression with NonSQLExpression { | ||
|
|
||
| override def nullable: Boolean = inputData.nullable | ||
|
|
||
| override def children: Seq[Expression] = | ||
| keyLambdaFunction :: valueLambdaFunction :: inputData :: Nil | ||
|
|
||
| override def eval(input: InternalRow): Any = | ||
| throw new UnsupportedOperationException("Only code-generated evaluation is supported") | ||
|
|
||
| override def dataType: DataType = ObjectType(collClass) | ||
|
|
||
| override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { | ||
| val mapType = inputData.dataType.asInstanceOf[MapType] | ||
| val keyElementJavaType = ctx.javaType(mapType.keyType) | ||
| ctx.addMutableState(keyElementJavaType, keyLoopValue, "") | ||
| val genKeyFunction = keyLambdaFunction.genCode(ctx) | ||
| val valueElementJavaType = ctx.javaType(mapType.valueType) | ||
| ctx.addMutableState("boolean", valueLoopIsNull, "") | ||
| ctx.addMutableState(valueElementJavaType, valueLoopValue, "") | ||
| val genValueFunction = valueLambdaFunction.genCode(ctx) | ||
| val genInputData = inputData.genCode(ctx) | ||
| val dataLength = ctx.freshName("dataLength") | ||
| val loopIndex = ctx.freshName("loopIndex") | ||
|
|
||
| val keyArray = ctx.freshName("keyArray") | ||
| val valueArray = ctx.freshName("valueArray") | ||
|
|
||
| // The data with PythonUserDefinedType are actually stored with the data type of its sqlType. | ||
| // When we want to apply MapObjects on it, we have to use it. | ||
| def inputDataType(dataType: DataType) = dataType match { | ||
|
||
| case p: PythonUserDefinedType => p.sqlType | ||
| case _ => dataType | ||
| } | ||
|
|
||
| def lengthAndLoopVar(elementType: DataType, genInputData: ExprCode, method: String, | ||
|
||
| array: String) = | ||
| s"${genInputData.value}.$method().numElements()" -> | ||
| ctx.getValue(s"${genInputData.value}.$method()", elementType, loopIndex) | ||
|
|
||
| val ((getKeyLength, getKeyLoopVar), (getValueLength, getValueLoopVar)) = ( | ||
| lengthAndLoopVar(inputDataType(mapType.keyType), genInputData, "keyArray", keyArray), | ||
| lengthAndLoopVar(inputDataType(mapType.valueType), genInputData, "valueArray", valueArray) | ||
| ) | ||
|
|
||
| // Make a copy of the data if it's unsafe-backed | ||
| def makeCopyIfInstanceOf(clazz: Class[_ <: Any], value: String) = | ||
| s"$value instanceof ${clazz.getSimpleName}? $value.copy() : $value" | ||
| def genFunctionValue(lambdaFunction: Expression, genFunction: ExprCode) = | ||
| lambdaFunction.dataType match { | ||
| case StructType(_) => makeCopyIfInstanceOf(classOf[UnsafeRow], genFunction.value) | ||
| case ArrayType(_, _) => makeCopyIfInstanceOf(classOf[UnsafeArrayData], genFunction.value) | ||
| case MapType(_, _, _) => makeCopyIfInstanceOf(classOf[UnsafeMapData], genFunction.value) | ||
| case _ => genFunction.value | ||
| } | ||
| val genKeyFunctionValue = genFunctionValue(keyLambdaFunction, genKeyFunction) | ||
| val genValueFunctionValue = genFunctionValue(valueLambdaFunction, genValueFunction) | ||
|
|
||
| val valueLoopNullCheck = | ||
| s"$valueLoopIsNull = ${genInputData.value}.valueArray().isNullAt($loopIndex);" | ||
|
||
|
|
||
| val constructBuilder = collClass match { | ||
| // Scala Map | ||
| case cls if classOf[scala.collection.Map[_, _]].isAssignableFrom(cls) => | ||
| val builderClass = classOf[Builder[_, _]].getName | ||
| s""" | ||
| $builderClass $builderValue = ${collClass.getName}$$.MODULE$$.newBuilder(); | ||
| $builderValue.sizeHint($dataLength); | ||
| """ | ||
| // Java Map, AbstractMap => HashMap | ||
|
||
| case cls if classOf[java.util.Map[_, _]] == cls || | ||
| classOf[java.util.AbstractMap[_, _]] == cls => | ||
| val builderClass = classOf[java.util.HashMap[_, _]].getName | ||
| s"${collClass.getName} $builderValue = new $builderClass($dataLength);" | ||
| // Java SortedMap, NavigableMap => TreeMap | ||
| case cls if classOf[java.util.SortedMap[_, _]] == cls || | ||
| classOf[java.util.NavigableMap[_, _]] == cls => | ||
| val builderClass = classOf[java.util.TreeMap[_, _]].getName | ||
| s"${collClass.getName} $builderValue = new $builderClass();" | ||
| // Java ConcurrentMap => ConcurrentHashMap | ||
| case cls if classOf[java.util.concurrent.ConcurrentMap[_, _]] == cls => | ||
| val builderClass = classOf[java.util.concurrent.ConcurrentHashMap[_, _]].getName | ||
| s"${collClass.getName} $builderValue = new $builderClass();" | ||
| // Java ConcurrentNavigableMap => ConcurrentSkipListMap | ||
| case cls if classOf[java.util.concurrent.ConcurrentNavigableMap[_, _]] == cls => | ||
| val builderClass = classOf[java.util.concurrent.ConcurrentSkipListMap[_, _]].getName | ||
| s"${collClass.getName} $builderValue = new $builderClass();" | ||
| // Java concrete Map implementation | ||
| case cls => | ||
| // Check for constructor with capacity specification | ||
| if (Try(cls.getConstructor(Integer.TYPE)).isSuccess) { | ||
| s"${collClass.getName} $builderValue = new ${cls.getName}($dataLength);" | ||
| } else { | ||
| s"${collClass.getName} $builderValue = new ${cls.getName}();" | ||
| } | ||
| } | ||
|
|
||
| val (appendToBuilder, getBuilderResult) = | ||
| if (classOf[scala.collection.Map[_, _]].isAssignableFrom(collClass)) { | ||
|
||
| val tupleClass = classOf[(_, _)].getName | ||
| s""" | ||
| $tupleClass $tupleLoopValue; | ||
|
|
||
| if (${genValueFunction.isNull}) { | ||
| $tupleLoopValue = new $tupleClass($genKeyFunctionValue, null); | ||
| } else { | ||
| $tupleLoopValue = new $tupleClass($genKeyFunctionValue, $genValueFunctionValue); | ||
| } | ||
|
|
||
| $builderValue.$$plus$$eq($tupleLoopValue); | ||
|
||
| """ -> s"${ev.value} = (${collClass.getName}) $builderValue.result();" | ||
| } else { | ||
| s"$builderValue.put($genKeyFunctionValue, $genValueFunctionValue);" -> | ||
| s"${ev.value} = $builderValue;" | ||
| } | ||
|
|
||
| val code = s""" | ||
| ${genInputData.code} | ||
| ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)}; | ||
|
|
||
| if (!${genInputData.isNull}) { | ||
| if ($getKeyLength != $getValueLength) { | ||
|
||
| throw new RuntimeException("Invalid state: Inconsistent lengths of key-value arrays"); | ||
| } | ||
| int $dataLength = $getKeyLength; | ||
| $constructBuilder | ||
|
|
||
| int $loopIndex = 0; | ||
| while ($loopIndex < $dataLength) { | ||
| $keyLoopValue = ($keyElementJavaType) ($getKeyLoopVar); | ||
| $valueLoopValue = ($valueElementJavaType) ($getValueLoopVar); | ||
| $valueLoopNullCheck | ||
|
Contributor
There was a problem hiding this comment. we can also inline this. The principle is, we should inline these simple codes as many as possible, then when you look at this code block, it's more clear what's going on. |
||
|
|
||
| ${genKeyFunction.code} | ||
| ${genValueFunction.code} | ||
|
|
||
| $appendToBuilder | ||
|
|
||
| $loopIndex += 1; | ||
| } | ||
|
|
||
| $getBuilderResult | ||
| } | ||
| """ | ||
| ev.copy(code = code, isNull = genInputData.isNull) | ||
| } | ||
| } | ||
|
|
||
| object ExternalMapToCatalyst { | ||
| private val curId = new java.util.concurrent.atomic.AtomicInteger() | ||
|
|
||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
CollectObjects->CollectObjectsToMap