Github user mgaido91 commented on a diff in the pull request:

    https://github.com/apache/spark/pull/21045#discussion_r188882871
  
    --- Diff: 
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
 ---
    @@ -90,6 +90,132 @@ case class MapKeys(child: Expression)
       override def prettyName: String = "map_keys"
     }
     
    +@ExpressionDescription(
    +  usage = """_FUNC_(a1, a2, ...) - Returns a merged array containing in 
the N-th position the
    +  N-th value of each array given.""",
    +  examples = """
    +    Examples:
    +      > SELECT _FUNC_(array(1, 2, 3), array(2, 3, 4));
    +        [[1, 2], [2, 3], [3, 4]]
    +      > SELECT _FUNC_(array(1, 2), array(2, 3), array(3, 4));
    +        [[1, 2, 3], [2, 3, 4]]
    +  """,
    +  since = "2.4.0")
    +case class Zip(children: Seq[Expression]) extends Expression with 
ExpectsInputTypes {
    +  override def inputTypes: Seq[AbstractDataType] = 
Seq.fill(children.length)(ArrayType)
    +
    +  def mountSchema(): StructType = {
    +    val arrayAT = children.map(_.dataType.asInstanceOf[ArrayType])
    +    val fields = arrayAT.zipWithIndex.foldRight(List[StructField]()) {
    +      (item, list) => {
    +        val (arr, idx) = item
    +        StructField(s"_$idx", arr.elementType, arr.containsNull) :: list
    +      }
    +    }
    +    StructType(fields)
    +  }
    +
    +  override def dataType: DataType = ArrayType(mountSchema())
    +
    +  override def prettyName: String = "zip"
    +
    +  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    +    val genericArrayData = classOf[GenericArrayData].getName
    +    val genericInternalRow = classOf[GenericInternalRow].getName
    +
    +    val evals = children.map(_.genCode(ctx))
    +    val numArrs = evals.length
    +
    +    val arrCardinality = ctx.freshName("arrCardinality")
    +    val arrVals = ctx.freshName("arrVals")
    +
    +    val arrayTypes = 
children.map(_.dataType.asInstanceOf[ArrayType].elementType)
    +
    +    val inputs = evals.zipWithIndex.map { case (eval, index) =>
    +      s"""
    +        |${eval.code}
    +        |if (!${eval.isNull}) {
    +        |  $arrVals[$index] = ${eval.value};
    +        |  $arrCardinality[$index] = ${eval.value}.numElements();
    +        |} else {
    +        |  $arrVals[$index] = null;
    +        |  $arrCardinality[$index] = 0;
    +        |}
    +      """.stripMargin
    +    }.mkString("\n")
    +
    +    val myobject = ctx.freshName("myobject")
    +    val biggestCardinality = ctx.freshName("biggestCardinality")
    +    val j = ctx.freshName("j")
    +    val i = ctx.freshName("i")
    +    val args = ctx.freshName("args")
    +
    +    val fillValue = evals.zipWithIndex.map { case (eval, index) =>
    +      s"""
    +      |if ($j == ${index}) {
    +      |  $myobject[$j] = ${CodeGenerator.getValue(s"$arrVals[$j]", 
arrayTypes(index), i)};
    +      |}
    +      """.stripMargin
    +    }.mkString("\n")
    +
    +    ev.copy(s"""
    +      |ArrayData[] $arrVals = new ArrayData[$numArrs];
    +      |int[] $arrCardinality = new int[$numArrs];
    +      |$inputs
    +      |int $biggestCardinality = 0;
    +      |for (int $i = 0; $i < $numArrs; $i ++) {
    +      |  $biggestCardinality = Math.max($biggestCardinality, 
$arrCardinality[$i]);
    +      |}
    +      |Object[] $args = new Object[$biggestCardinality];
    +      |for (int $i = 0; $i < $biggestCardinality; $i ++) {
    +      |  Object[] $myobject = new Object[$numArrs];
    +      |  for (int $j = 0; $j < $numArrs; $j ++) {
    +      |    if ($arrVals[$j] != null && $arrCardinality[$j] > $i) {
    +      |      $fillValue;
    +      |    } else {
    +      |      $myobject[$j] = null;
    +      |    }
    +      |  }
    +      |  $args[$i] = new $genericInternalRow($myobject);
    +      |}
    +      |boolean ${ev.isNull} = false;
    +      |$genericArrayData ${ev.value} = new $genericArrayData($args);
    +    """.stripMargin)
    +  }
    +
    +  override def nullable: Boolean = children.forall(_.nullable)
    +
    +  override def eval(input: InternalRow): Any = {
    +    val inputArrays = children.map(_.eval(input).asInstanceOf[ArrayData])
    +    val arrayTypes = 
children.map(_.dataType.asInstanceOf[ArrayType].elementType)
    +    val numberOfArrays = children.length
    --- End diff --
    
    this is fixed for each input. Can we compute it as a lazy val in in order 
to avoid recomputation?


---

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org

Reply via email to