minor cleanup

Author: fqaiser94

sql/core/src/test/scala/org/apache/spark/sql/UpdateFieldsPerformanceSuite.scala

@@ -23,22 +23,22 @@ import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 
 class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
 
-  private def colName(d: Int, colNum: Int): String = s"nested${d}Col$colNum"
+  private def nestedColName(d: Int, colNum: Int): String = s"nested${d}Col$colNum"
 
   private def nestedStructType(
     depths: Seq[Int], colNums: Seq[Int], nullable: Boolean): StructType = {
     if (depths.length == 1) {
       StructType(colNums.map { colNum =>
-        StructField(colName(depths.head, colNum), IntegerType, nullable = false)
+        StructField(nestedColName(depths.head, colNum), IntegerType, nullable = false)
       })
     } else {
       val depth = depths.head
       val fields = colNums.foldLeft(Seq.empty[StructField]) {
         case (structFields, colNum) if colNum == 0 =>
-          val nested = nestedStructType(depths.tail, colNums, nullable = nullable)
-          structFields :+ StructField(colName(depth, colNum), nested, nullable = nullable)
+          val nested = nestedStructType(depths.tail, colNums, nullable)
+          structFields :+ StructField(nestedColName(depth, colNum), nested, nullable)
         case (structFields, colNum) =>
-          structFields :+ StructField(colName(depth, colNum), IntegerType, nullable = false)
+          structFields :+ StructField(nestedColName(depth, colNum), IntegerType, nullable = false)
       }
       StructType(fields)
     }
@@ -60,10 +60,10 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
    * Utility function for generating a DataFrame with nested columns.
    *
    * @param depth: The depth to which to create nested columns.
-   * @param numColsAtEachDepth: The number of columns to create at each depth. The columns names
-   *                          are in the format of nested${depth}Col${index}. The value of each
-   *                          column will be its index at that depth, or if the index of the column
-   *                          is 0, then the value could also be a struct.
+   * @param numColsAtEachDepth: The number of columns to create at each depth. The value of each
+   *                          column will be the same as its index (IntegerType) at that depth
+   *                          unless the index = 0, in which case it may be a StructType which
+   *                          represents the next depth.
    * @param nullable: This value is used to set the nullability of StructType columns.
    */
   private def nestedDf(
@@ -78,7 +78,7 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
 
     spark.createDataFrame(
       sparkContext.parallelize(Row(nestedColumn) :: Nil),
-      StructType(Seq(StructField(colName(0, 0), nestedColumnDataType, nullable = nullable))))
+      StructType(Seq(StructField(nestedColName(0, 0), nestedColumnDataType, nullable))))
   }
 
   test("nestedDf should generate nested DataFrames") {
@@ -139,21 +139,21 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
 
     // drop columns at the current depth
     val dropped = if (colNumsToDrop.nonEmpty) {
-      column.dropFields(colNumsToDrop.map(num => colName(depth, num)): _*)
+      column.dropFields(colNumsToDrop.map(num => nestedColName(depth, num)): _*)
     } else column
 
     // add columns at the current depth
     val added = colNumsToAdd.foldLeft(dropped) {
-      (res, num) => res.withField(colName(depth, num), lit(num))
+      (res, num) => res.withField(nestedColName(depth, num), lit(num))
     }
 
     if (depths.length == 1) {
       added
     } else {
       // add/drop columns at the next depth
-      val nestedColumn = col((0 to depth).map(d => s"`${colName(d, 0)}`").mkString("."))
+      val nestedColumn = col((0 to depth).map(d => s"`${nestedColName(d, 0)}`").mkString("."))
       added.withField(
-        colName(depth, 0),
+        nestedColName(depth, 0),
         addDropNestedColumns(nestedColumn, depths.tail, colNumsToAdd, colNumsToDrop))
     }
   }
@@ -165,16 +165,16 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
       val maxDepth = 20
 
       // dataframe with nested*Col0 to nested*Col4 at each of 20 depths
-      val inputDf = nestedDf(maxDepth, 5, nullable = nullable)
+      val inputDf = nestedDf(maxDepth, 5, nullable)
 
       // add nested*Col5 through nested*Col9 at each depth
       val resultDf = inputDf.select(addDropNestedColumns(
-        column = col(colName(0, 0)),
+        column = col(nestedColName(0, 0)),
         depths = 1 to maxDepth,
         colNumsToAdd = 5 to 9).as("nested0Col0"))
 
       // dataframe with nested*Col0 to nested*Col9 at each of 20 depths
-      val expectedDf = nestedDf(maxDepth, 10, nullable = nullable)
+      val expectedDf = nestedDf(maxDepth, 10, nullable)
       checkAnswer(resultDf, expectedDf.collect(), expectedDf.schema)
     }
 
@@ -183,16 +183,16 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
       val maxDepth = 20
 
       // dataframe with nested*Col0 to nested*Col9 at each of 20 depths
-      val inputDf = nestedDf(maxDepth, 10, nullable = nullable)
+      val inputDf = nestedDf(maxDepth, 10, nullable)
 
       // drop nested*Col5 to nested*Col9 at each of 20 depths
       val resultDf = inputDf.select(addDropNestedColumns(
-        column = col(colName(0, 0)),
+        column = col(nestedColName(0, 0)),
         depths = 1 to maxDepth,
         colNumsToDrop = 5 to 9).as("nested0Col0"))
 
       // dataframe with nested*Col0 to nested*Col4 at each of 20 depths
-      val expectedDf = nestedDf(maxDepth, 5, nullable = nullable)
+      val expectedDf = nestedDf(maxDepth, 5, nullable)
       checkAnswer(resultDf, expectedDf.collect(), expectedDf.schema)
     }
 
@@ -201,12 +201,12 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
       val maxDepth = 20
 
       // dataframe with nested*Col0 to nested*Col9 at each of 20 depths
-      val inputDf = nestedDf(maxDepth, 10, nullable = nullable)
+      val inputDf = nestedDf(maxDepth, 10, nullable)
 
       // add nested*Col10 through nested*Col14 at each depth
       // drop nested*Col5 through nested*Col9 at each depth
       val resultDf = inputDf.select(addDropNestedColumns(
-        column = col(colName(0, 0)),
+        column = col(nestedColName(0, 0)),
         depths = 1 to maxDepth,
         colNumsToAdd = 10 to 14,
         colNumsToDrop = 5 to 9).as("nested0Col0"))
@@ -217,11 +217,11 @@ class UpdateFieldsPerformanceSuite extends QueryTest with SharedSparkSession {
         val depths = 1 to maxDepth
         val numCols = (0 to 4) ++ (10 to 14)
         val nestedColumn = nestedRow(depths, numCols)
-        val nestedColumnDataType = nestedStructType(depths, numCols, nullable = nullable)
+        val nestedColumnDataType = nestedStructType(depths, numCols, nullable)
 
         spark.createDataFrame(
           sparkContext.parallelize(Row(nestedColumn) :: Nil),
-          StructType(Seq(StructField(colName(0, 0), nestedColumnDataType, nullable = nullable))))
+          StructType(Seq(StructField(nestedColName(0, 0), nestedColumnDataType, nullable))))
       }
       checkAnswer(resultDf, expectedDf.collect(), expectedDf.schema)
     }