Don't import spark session implicits (#25)

cwienberg · Aug 26, 2022 · b3dd37d · b3dd37d
1 parent 2e530b4
commit b3dd37d
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/main/scala/net/gonzberg/spark/sorting/SecondarySortGroupByKeyDatasetFunctions.scala b/src/main/scala/net/gonzberg/spark/sorting/SecondarySortGroupByKeyDatasetFunctions.scala
@@ -110,22 +110,21 @@ final class SecondarySortGroupByKeyDatasetFunctions[K, V](
     keyOrdering: Ordering[K],
     kaEncoder: Encoder[(K, A)]
   ): Dataset[(K, A)] = {
-    import dataset.sparkSession.implicits._
     val repartitionedStartValues = SecondarySortGroupByKeyDatasetFunctions
       .repartitionAndSort(startValues, numPartitions)
       .rdd
     val repartitionedValues = SecondarySortGroupByKeyDatasetFunctions
       .repartitionAndSort(dataset, numPartitions, orderExprs)
       .rdd
-    repartitionedStartValues
+    val kaRDD = repartitionedStartValues
       .zipPartitions(
         repartitionedValues.mapPartitions(
           new GroupByKeyIterator(_),
           preservesPartitioning = true
         ),
         preservesPartitioning = true
       )(joinAndFold(op))
-      .toDS()
+    dataset.sparkSession.createDataset(kaRDD)
   }
 
   /** Groups by key and applies a binary operation using foldLeft