[SPARK-23628][SQL] Fix calculateParamLength

mgaido91 · mgaido91 · commit fbe9db6b9897 · 2018-03-09T10:34:57.000+01:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -1253,14 +1253,15 @@ class CodegenContext {
    */
   def calculateParamLength(params: Seq[Expression]): Int = {
     def paramLengthForExpr(input: Expression): Int = {
-      // For a nullable expression, we need to pass in an extra boolean parameter.
-      (if (input.nullable) 1 else 0) + javaType(input.dataType) match {
+      val javaParamLength = javaType(input.dataType) match {
         case JAVA_LONG | JAVA_DOUBLE => 2
         case _ => 1
       }
+      // For a nullable expression, we need to pass in an extra boolean parameter.
+      (if (input.nullable) 1 else 0) + javaParamLength
     }
     // Initial value is 1 for `this`.
-    1 + params.map(paramLengthForExpr(_)).sum
+    1 + params.map(paramLengthForExpr).sum
   }
 
   /**
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala
@@ -436,4 +436,10 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {
     ctx.addImmutableStateIfNotExists("String", mutableState2)
     assert(ctx.inlinedMutableStates.length == 2)
   }
+
+  test("SPARK-23628: calculateParamLength should compute properly the param length") {
+    val ctx = new CodegenContext
+    assert(ctx.calculateParamLength(Seq.range(0, 100).map(Literal(_))) == 101)
+    assert(ctx.calculateParamLength(Seq.range(0, 100).map(x => Literal(x.toLong))) == 201)
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
 import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec
 import org.apache.spark.sql.execution.joins.SortMergeJoinExec
 import org.apache.spark.sql.expressions.scalalang.typed
-import org.apache.spark.sql.functions.{avg, broadcast, col, max}
+import org.apache.spark.sql.functions.{avg, broadcast, col, lit, max}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
@@ -249,12 +249,12 @@ class WholeStageCodegenSuite extends QueryTest with SharedSQLContext {
   }
 
   test("Skip splitting consume function when parameter number exceeds JVM limit") {
-    import testImplicits._
-
-    Seq((255, false), (254, true)).foreach { case (columnNum, hasSplit) =>
+    // since every field is nullable we have 2 params for each input column (one for the value
+    // and one for the isNull variable)
+    Seq((128, false), (127, true)).foreach { case (columnNum, hasSplit) =>
       withTempPath { dir =>
         val path = dir.getCanonicalPath
-        spark.range(10).select(Seq.tabulate(columnNum) {i => ('id + i).as(s"c$i")} : _*)
+        spark.range(10).select(Seq.tabulate(columnNum) {i => lit(i).as(s"c$i")} : _*)
           .write.mode(SaveMode.Overwrite).parquet(path)
 
         withSQLConf(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "255",
@@ -263,10 +263,10 @@ class WholeStageCodegenSuite extends QueryTest with SharedSQLContext {
           val df = spark.read.parquet(path).selectExpr(projection: _*)
 
           val plan = df.queryExecution.executedPlan
-          val wholeStageCodeGenExec = plan.find(p => p match {
-            case wp: WholeStageCodegenExec => true
+          val wholeStageCodeGenExec = plan.find {
+            case _: WholeStageCodegenExec => true
             case _ => false
-          })
+          }
           assert(wholeStageCodeGenExec.isDefined)
           val code = wholeStageCodeGenExec.get.asInstanceOf[WholeStageCodegenExec].doCodeGen()._2
           assert(code.body.contains("project_doConsume") == hasSplit)

Original file line number	Diff line number	Diff line change
`@@ -1253,14 +1253,15 @@ class CodegenContext {`
`1253`	`1253`	`*/`
`1254`	`1254`	`def calculateParamLength(params: Seq[Expression]): Int = {`
`1255`	`1255`	`def paramLengthForExpr(input: Expression): Int = {`
`1256`		`- // For a nullable expression, we need to pass in an extra boolean parameter.`
`1257`		`- (if (input.nullable) 1 else 0) + javaType(input.dataType) match {`
	`1256`	`+ val javaParamLength = javaType(input.dataType) match {`
`1258`	`1257`	`case JAVA_LONG \| JAVA_DOUBLE => 2`
`1259`	`1258`	`case _ => 1`
`1260`	`1259`	`}`
	`1260`	`+ // For a nullable expression, we need to pass in an extra boolean parameter.`
	`1261`	`+ (if (input.nullable) 1 else 0) + javaParamLength`
`1261`	`1262`	`}`
`1262`	`1263`	// Initial value is 1 for `this`.
`1263`		`- 1 + params.map(paramLengthForExpr(_)).sum`
	`1264`	`+ 1 + params.map(paramLengthForExpr).sum`
`1264`	`1265`	`}`
`1265`	`1266`
`1266`	`1267`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -436,4 +436,10 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {`
`436`	`436`	`ctx.addImmutableStateIfNotExists("String", mutableState2)`
`437`	`437`	`assert(ctx.inlinedMutableStates.length == 2)`
`438`	`438`	`}`
	`439`	`+`
	`440`	`+ test("SPARK-23628: calculateParamLength should compute properly the param length") {`
	`441`	`+ val ctx = new CodegenContext`
	`442`	`+ assert(ctx.calculateParamLength(Seq.range(0, 100).map(Literal(_))) == 101)`
	`443`	`+ assert(ctx.calculateParamLength(Seq.range(0, 100).map(x => Literal(x.toLong))) == 201)`
	`444`	`+ }`
`439`	`445`	`}`