diff --git a/sql/core/src/main/scala/org/apache/spark/sql/api/r/SQLUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/api/r/SQLUtils.scala index b0120a8d0dc4ffacd92e0188937e2d051d8cb5df..b3f134614c6bb813bbeae23f14bddfc10b898bb0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/api/r/SQLUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/api/r/SQLUtils.scala @@ -130,16 +130,18 @@ private[r] object SQLUtils { } def dfToCols(df: DataFrame): Array[Array[Any]] = { - // localDF is Array[Row] - val localDF = df.collect() + val localDF: Array[Row] = df.collect() val numCols = df.columns.length + val numRows = localDF.length - // result is Array[Array[Any]] - (0 until numCols).map { colIdx => - localDF.map { row => - row(colIdx) + val colArray = new Array[Array[Any]](numCols) + for (colNo <- 0 until numCols) { + colArray(colNo) = new Array[Any](numRows) + for (rowNo <- 0 until numRows) { + colArray(colNo)(rowNo) = localDF(rowNo)(colNo) } - }.toArray + } + colArray } def saveMode(mode: String): SaveMode = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/api/r/SQLUtilsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/api/r/SQLUtilsSuite.scala new file mode 100644 index 0000000000000000000000000000000000000000..f54e23e3aa6cb67b51bf76f204862a53b9685199 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/api/r/SQLUtilsSuite.scala @@ -0,0 +1,38 @@ +/* +* Licensed to the Apache Software Foundation (ASF) under one or more +* contributor license agreements. See the NOTICE file distributed with +* this work for additional information regarding copyright ownership. +* The ASF licenses this file to You under the Apache License, Version 2.0 +* (the "License"); you may not use this file except in compliance with +* the License. You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +package org.apache.spark.sql.api.r + +import org.apache.spark.sql.test.SharedSQLContext + +class SQLUtilsSuite extends SharedSQLContext { + + import testImplicits._ + + test("dfToCols should collect and transpose a data frame") { + val df = Seq( + (1, 2, 3), + (4, 5, 6) + ).toDF + assert(SQLUtils.dfToCols(df) === Array( + Array(1, 4), + Array(2, 5), + Array(3, 6) + )) + } + +}