@@ -340,174 +340,6 @@ case class MapValues(child: Expression)
340340 override def prettyName : String = " map_values"
341341}
342342
343- /**
344- * Returns an unordered array of all entries in the given map.
345- */
346- @ ExpressionDescription (
347- usage = " _FUNC_(map) - Returns an unordered array of all entries in the given map." ,
348- examples = """
349- Examples:
350- > SELECT _FUNC_(map(1, 'a', 2, 'b'));
351- [{"key":1,"value":"a"},{"key":2,"value":"b"}]
352- """ ,
353- since = " 2.4.0" )
354- case class MapEntries (child : Expression ) extends UnaryExpression with ExpectsInputTypes {
355-
356- override def inputTypes : Seq [AbstractDataType ] = Seq (MapType )
357-
358- @ transient private lazy val childDataType : MapType = child.dataType.asInstanceOf [MapType ]
359-
360- override def dataType : DataType = {
361- ArrayType (
362- StructType (
363- StructField (" key" , childDataType.keyType, false ) ::
364- StructField (" value" , childDataType.valueType, childDataType.valueContainsNull) ::
365- Nil ),
366- false )
367- }
368-
369- override protected def nullSafeEval (input : Any ): Any = {
370- val childMap = input.asInstanceOf [MapData ]
371- val keys = childMap.keyArray()
372- val values = childMap.valueArray()
373- val length = childMap.numElements()
374- val resultData = new Array [AnyRef ](length)
375- var i = 0
376- while (i < length) {
377- val key = keys.get(i, childDataType.keyType)
378- val value = values.get(i, childDataType.valueType)
379- val row = new GenericInternalRow (Array [Any ](key, value))
380- resultData.update(i, row)
381- i += 1
382- }
383- new GenericArrayData (resultData)
384- }
385-
386- override protected def doGenCode (ctx : CodegenContext , ev : ExprCode ): ExprCode = {
387- nullSafeCodeGen(ctx, ev, c => {
388- val arrayData = ctx.freshName(" arrayData" )
389- val numElements = ctx.freshName(" numElements" )
390- val keys = ctx.freshName(" keys" )
391- val values = ctx.freshName(" values" )
392- val isKeyPrimitive = CodeGenerator .isPrimitiveType(childDataType.keyType)
393- val isValuePrimitive = CodeGenerator .isPrimitiveType(childDataType.valueType)
394-
395- val wordSize = UnsafeRow .WORD_SIZE
396- val structSize = UnsafeRow .calculateBitSetWidthInBytes(2 ) + wordSize * 2
397- val (isPrimitive, elementSize) = if (isKeyPrimitive && isValuePrimitive) {
398- (true , structSize + wordSize)
399- } else {
400- (false , - 1 )
401- }
402-
403- val allocation =
404- s """
405- |ArrayData $arrayData = ArrayData.allocateArrayData(
406- | $elementSize, $numElements, " $prettyName failed.");
407- """ .stripMargin
408-
409- val code = if (isPrimitive) {
410- val genCodeForPrimitive = genCodeForPrimitiveElements(
411- ctx, arrayData, keys, values, ev.value, numElements, structSize)
412- s """
413- |if ( $arrayData instanceof UnsafeArrayData) {
414- | $genCodeForPrimitive
415- |} else {
416- | ${genCodeForAnyElements(ctx, arrayData, keys, values, ev.value, numElements)}
417- |}
418- """ .stripMargin
419- } else {
420- s " ${genCodeForAnyElements(ctx, arrayData, keys, values, ev.value, numElements)}"
421- }
422-
423- s """
424- |final int $numElements = $c.numElements();
425- |final ArrayData $keys = $c.keyArray();
426- |final ArrayData $values = $c.valueArray();
427- | $allocation
428- | $code
429- """ .stripMargin
430- })
431- }
432-
433- private def getKey (varName : String , index : String ) =
434- CodeGenerator .getValue(varName, childDataType.keyType, index)
435-
436- private def getValue (varName : String , index : String ) =
437- CodeGenerator .getValue(varName, childDataType.valueType, index)
438-
439- private def genCodeForPrimitiveElements (
440- ctx : CodegenContext ,
441- arrayData : String ,
442- keys : String ,
443- values : String ,
444- resultArrayData : String ,
445- numElements : String ,
446- structSize : Int ): String = {
447- val unsafeArrayData = ctx.freshName(" unsafeArrayData" )
448- val baseObject = ctx.freshName(" baseObject" )
449- val unsafeRow = ctx.freshName(" unsafeRow" )
450- val structsOffset = ctx.freshName(" structsOffset" )
451- val offset = ctx.freshName(" offset" )
452- val z = ctx.freshName(" z" )
453- val calculateHeader = " UnsafeArrayData.calculateHeaderPortionInBytes"
454-
455- val baseOffset = Platform .BYTE_ARRAY_OFFSET
456- val wordSize = UnsafeRow .WORD_SIZE
457- val structSizeAsLong = s " ${structSize}L "
458-
459- val setKey = CodeGenerator .setColumn(unsafeRow, childDataType.keyType, 0 , getKey(keys, z))
460-
461- val valueAssignmentChecked = CodeGenerator .createArrayAssignment(
462- unsafeRow, childDataType.valueType, values, " 1" , z, childDataType.valueContainsNull)
463-
464- s """
465- |UnsafeArrayData $unsafeArrayData = (UnsafeArrayData) $arrayData;
466- |Object $baseObject = $unsafeArrayData.getBaseObject();
467- |final int $structsOffset = $calculateHeader( $numElements) + $numElements * $wordSize;
468- |UnsafeRow $unsafeRow = new UnsafeRow(2);
469- |for (int $z = 0; $z < $numElements; $z++) {
470- | long $offset = $structsOffset + $z * $structSizeAsLong;
471- | $unsafeArrayData.setLong( $z, ( $offset << 32) + $structSizeAsLong);
472- | $unsafeRow.pointTo( $baseObject, $baseOffset + $offset, $structSize);
473- | $setKey;
474- | $valueAssignmentChecked
475- |}
476- | $resultArrayData = $arrayData;
477- """ .stripMargin
478- }
479-
480- private def genCodeForAnyElements (
481- ctx : CodegenContext ,
482- arrayData : String ,
483- keys : String ,
484- values : String ,
485- resultArrayData : String ,
486- numElements : String ): String = {
487- val z = ctx.freshName(" z" )
488- val isValuePrimitive = CodeGenerator .isPrimitiveType(childDataType.valueType)
489- val getValueWithCheck = if (childDataType.valueContainsNull && isValuePrimitive) {
490- s " $values.isNullAt( $z) ? null : (Object) ${getValue(values, z)}"
491- } else {
492- getValue(values, z)
493- }
494-
495- val rowClass = classOf [GenericInternalRow ].getName
496- val genericArrayDataClass = classOf [GenericArrayData ].getName
497- val genericArrayData = ctx.freshName(" genericArrayData" )
498- val rowObject = s " new $rowClass(new Object[]{ ${getKey(keys, z)}, $getValueWithCheck}) "
499- s """
500- | $genericArrayDataClass $genericArrayData = ( $genericArrayDataClass) $arrayData;
501- |for (int $z = 0; $z < $numElements; $z++) {
502- | $genericArrayData.update( $z, $rowObject);
503- |}
504- | $resultArrayData = $arrayData;
505- """ .stripMargin
506- }
507-
508- override def prettyName : String = " map_entries"
509- }
510-
511343/**
512344 * Returns the union of all the given maps.
513345 */
0 commit comments