update OneDAL.scala

Signed-off-by: minmingzhu <minming.zhu@intel.com>
oap-project · Jun 29, 2023 · ce644d4 · ce644d4
1 parent a862ef2
commit ce644d4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala
@@ -501,7 +501,7 @@ object OneDAL {
     }
 
     // Get dimensions for each partition
-    val partitionDims = Utils.getPartitionDims(labeledPoints.select(featuresCol).rdd.map { row =>
+    val partitionDims = Utils.getPartitionDims(dataForConversion.select(featuresCol).rdd.map{ row =>
       val vector = row.getAs[Vector](0)
       vector
     })
@@ -510,7 +510,7 @@ object OneDAL {
     // "No partitions or no locations for partitions found".
     // TODO: ML-312: Improve ExecutorInProcessCoalescePartitioner
     val nonEmptyPartitions = dataForConversion.select(labelCol, featuresCol).toDF().rdd.mapPartitionsWithIndex {
-      (index: Int, it: Iterator[Row]) => Iterator(Tuple3( partitionDims(index)._1,index, it))
+      (index: Int, it: Iterator[Row]) => Iterator(Tuple3(partitionDims(index)._1, index, it))
     }.filter {
       _._1 > 0
     }.flatMap {