From 6831cc1a2aa8a2eb72e989d891044e35a22082d0 Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Wed, 25 Mar 2026 01:37:25 -0700
Subject: [PATCH] Add column selector overloads for GroupBy.sortByKey

Adds ColumnsSelector overloads for sortByKey and sortByKeyDesc
so users can specify which key columns to sort by instead of
sorting by all keys with the same order.

Fixes #1523
---
 .../org/jetbrains/kotlinx/dataframe/api/sort.kt       | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt
index bb7e05607c..fc1425be97 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataFrameExpression
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
@@ -220,4 +221,14 @@ public fun <T, G> GroupBy<T, G>.sortByKey(nullsLast: Boolean = false): GroupBy<T
         .sortBy { keys.columns().toColumnSet().nullsLast(nullsLast) }
         .asGroupBy(groups)
 
+public fun <T, G, C> GroupBy<T, G>.sortByKey(columns: ColumnsSelector<T, C>): GroupBy<T, G> =
+    toDataFrame()
+        .sortBy { columns.toColumnSet() }
+        .asGroupBy(groups)
+
+public fun <T, G, C> GroupBy<T, G>.sortByKeyDesc(columns: ColumnsSelector<T, C>): GroupBy<T, G> =
+    toDataFrame()
+        .sortBy { columns.toColumnSet().desc() }
+        .asGroupBy(groups)
+
 // endregion