added sorting functions and tests

Kotlin · Mar 9, 2022 · 8dc0b0f · 8dc0b0f
1 parent 5f7536c
commit 8dc0b0f
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/kotlin-spark-api/3.2/src/main/kotlin/org/jetbrains/kotlinx/spark/api/Dataset.kt b/kotlin-spark-api/3.2/src/main/kotlin/org/jetbrains/kotlinx/spark/api/Dataset.kt
@@ -309,6 +309,30 @@ inline fun <reified L : Any?, reified R : Any?> Dataset<L>.fullJoin(
  */
 inline fun <reified T> Dataset<T>.sort(columns: (Dataset<T>) -> Array<Column>): Dataset<T> = sort(*columns(this))
 
+/** Returns a dataset sorted by the first (`_1`) value of each [Tuple2] inside. */
+@JvmName("sortByTuple2Key")
+fun <T1, T2> Dataset<Tuple2<T1, T2>>.sortByKey(): Dataset<Tuple2<T1, T2>> = sort("_1")
+
+/** Returns a dataset sorted by the second (`_2`) value of each [Tuple2] inside. */
+@JvmName("sortByTuple2Value")
+fun <T1, T2> Dataset<Tuple2<T1, T2>>.sortByValue(): Dataset<Tuple2<T1, T2>> = sort("_2")
+
+/** Returns a dataset sorted by the first (`_1`) value of each [Arity2] inside. */
+@JvmName("sortByArity2Key")
+fun <T1, T2> Dataset<Arity2<T1, T2>>.sortByKey(): Dataset<Arity2<T1, T2>> = sort("_1")
+
+/** Returns a dataset sorted by the second (`_2`) value of each [Arity2] inside. */
+@JvmName("sortByArity2Value")
+fun <T1, T2> Dataset<Arity2<T1, T2>>.sortByValue(): Dataset<Arity2<T1, T2>> = sort("_2")
+
+/** Returns a dataset sorted by the first (`first`) value of each [Pair] inside. */
+@JvmName("sortByPairKey")
+fun <T1, T2> Dataset<Pair<T1, T2>>.sortByKey(): Dataset<Pair<T1, T2>> = sort("first")
+
+/** Returns a dataset sorted by the second (`second`) value of each [Pair] inside. */
+@JvmName("sortByPairValue")
+fun <T1, T2> Dataset<Pair<T1, T2>>.sortByValue(): Dataset<Pair<T1, T2>> = sort("second")
+
 /**
  * This function creates block, where one can call any further computations on already cached dataset
  * Data will be unpersisted automatically at the end of computation

diff --git a/kotlin-spark-api/3.2/src/test/kotlin/org/jetbrains/kotlinx/spark/api/DatasetFunctionTest.kt b/kotlin-spark-api/3.2/src/test/kotlin/org/jetbrains/kotlinx/spark/api/DatasetFunctionTest.kt
@@ -282,6 +282,42 @@ class DatasetFunctionTest : ShouldSpec({
 
                 dataset6.toList<List<Int>>() shouldBe listOf(listOf(1, 2, 3), listOf(4, 5, 6))
             }
+
+            should("Sort Arity2 Dataset") {
+                val list = listOf(
+                    c(1, 6),
+                    c(2, 5),
+                    c(3, 4),
+                )
+                val dataset = list.toDS()
+
+                dataset.sortByKey().collectAsList() shouldBe list.sortedBy { it._1 }
+                dataset.sortByValue().collectAsList() shouldBe list.sortedBy { it._2 }
+            }
+
+            should("Sort Tuple2 Dataset") {
+                val list = listOf(
+                    Tuple2(1, 6),
+                    Tuple2(2, 5),
+                    Tuple2(3, 4),
+                )
+                val dataset = list.toDS()
+
+                dataset.sortByKey().collectAsList() shouldBe list.sortedBy { it._1 }
+                dataset.sortByValue().collectAsList() shouldBe list.sortedBy { it._2 }
+            }
+
+            should("Sort Pair Dataset") {
+                val list = listOf(
+                    Pair(1, 6),
+                    Pair(2, 5),
+                    Pair(3, 4),
+                )
+                val dataset = list.toDS()
+
+                dataset.sortByKey().collectAsList() shouldBe list.sortedBy { it.first }
+                dataset.sortByValue().collectAsList() shouldBe list.sortedBy { it.second }
+            }
         }
     }
 
@@ -401,6 +437,7 @@ class DatasetFunctionTest : ShouldSpec({
                 b.count() shouldBe 1
             }
 
+
         }
     }
 })