fix: validate GroupBy nested sort paths

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt

@@ -177,8 +177,7 @@ public fun <T, G> GroupBy<T, G>.sortByDesc(vararg cols: KProperty<Comparable<*>?
     sortByDesc { cols.toColumnSet() }
 
 public fun <T, G, C> GroupBy<T, G>.sortByDesc(selector: SortColumnsSelector<G, C>): GroupBy<T, G> {
-    val set = selector.toColumnSet()
-    return sortByImpl { set.desc() }
+    return sortByImpl(selector, SortFlag.Reversed)
 }
 
 private fun <T, G, C> GroupBy<T, G>.createColumnFromGroupExpression(

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/sort.kt

@@ -5,12 +5,14 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.api.SortColumnsSelector
 import org.jetbrains.kotlinx.dataframe.api.asGroupBy
+import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.castFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.getFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.api.update
 import org.jetbrains.kotlinx.dataframe.api.with
 import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
+import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnsResolver
@@ -28,21 +30,52 @@ import org.jetbrains.kotlinx.dataframe.kind
 import org.jetbrains.kotlinx.dataframe.nrow
 
 @Suppress("UNCHECKED_CAST", "RemoveExplicitTypeArguments")
-internal fun <T, G> GroupBy<T, G>.sortByImpl(columns: SortColumnsSelector<G, *>): GroupBy<T, G> =
-    toDataFrame()
+internal fun <T, G> GroupBy<T, G>.sortByImpl(
+    columns: SortColumnsSelector<G, *>,
+    sortFlag: SortFlag? = null,
+): GroupBy<T, G> {
+    val groupedDf = toDataFrame()
+    val groupsToValidate = groups.values().toList().ifEmpty {
+        listOf(DataFrame.empty(groups.schema.value).cast())
+    }
+    validateSortColumns(groupsToValidate, groupedDf, columns)
+    return groupedDf
         // sort the individual groups by the columns specified
         .update { groups }
-        .with { it.sortByImpl(UnresolvedColumnsPolicy.Skip, columns) }
+        .with { it.sortByImpl(UnresolvedColumnsPolicy.Skip, columns, sortFlag) }
         // sort the groups by the columns specified (must be either be the keys column or "groups")
         // will do nothing if the columns specified are not the keys column or "groups"
-        .sortByImpl(UnresolvedColumnsPolicy.Skip, columns as SortColumnsSelector<T, *>)
+        .sortByImpl(UnresolvedColumnsPolicy.Skip, columns as SortColumnsSelector<T, *>, sortFlag)
         .asGroupBy { it.getFrameColumn(groups.name()).castFrameColumn<G>() }
+}
+
+@Suppress("UNCHECKED_CAST")
+private fun <T, G> validateSortColumns(
+    groups: List<DataFrame<G>>,
+    groupedDf: DataFrame<T>,
+    columns: SortColumnsSelector<G, *>,
+) {
+    val missingInGroups = groups
+        .map { columns.missingPaths(it) }
+        .reduce { missingInAllGroups, missingInGroup -> missingInAllGroups.intersect(missingInGroup) }
+    val missingInGroupedDf = (columns as SortColumnsSelector<T, *>).missingPaths(groupedDf)
+    val missingInBoth = missingInGroups.intersect(missingInGroupedDf)
+    missingInBoth.firstOrNull()?.let { missingPath ->
+        groups.first().getSortColumns({ missingPath }, UnresolvedColumnsPolicy.Fail)
+    }
+}
+
+private fun <T> SortColumnsSelector<T, *>.missingPaths(df: DataFrame<T>): Set<ColumnPath> =
+    toColumnSet().resolve(df, UnresolvedColumnsPolicy.Create)
+        .mapNotNull { (it.data as? MissingColumnGroup<*>)?.path }
+        .toSet()
 
 internal fun <T, C> DataFrame<T>.sortByImpl(
     unresolvedColumnsPolicy: UnresolvedColumnsPolicy = UnresolvedColumnsPolicy.Fail,
     columns: SortColumnsSelector<T, C>,
+    sortFlag: SortFlag? = null,
 ): DataFrame<T> {
-    val sortColumns = getSortColumns(columns, unresolvedColumnsPolicy)
+    val sortColumns = getSortColumns(columns, unresolvedColumnsPolicy, sortFlag)
     if (sortColumns.isEmpty()) return this
 
     val compChain = sortColumns.map {
@@ -71,6 +104,7 @@ internal fun AnyCol.createComparator(nullsLast: Boolean): java.util.Comparator<I
 internal fun <T, C> DataFrame<T>.getSortColumns(
     columns: SortColumnsSelector<T, C>,
     unresolvedColumnsPolicy: UnresolvedColumnsPolicy,
+    sortFlag: SortFlag? = null,
 ): List<SortColumnDescriptor<*>> =
     columns.toColumnSet().resolve(this, unresolvedColumnsPolicy)
         // can appear using [DataColumn<R>?.check] with UnresolvedColumnsPolicy.Skip
@@ -82,6 +116,7 @@ internal fun <T, C> DataFrame<T>.getSortColumns(
                 else -> throw IllegalStateException("Can not use ${col.kind} as sort column")
             }
         }
+        .map { if (sortFlag == null) it else it.addFlag(sortFlag) }
 
 internal enum class SortFlag { Reversed, NullsLast }
 
@@ -110,7 +145,10 @@ internal fun <C> ColumnWithPath<C>.addFlag(flag: SortFlag): ColumnWithPath<C> {
 }
 
 internal class ColumnSetWithSortFlag<C>(val column: ColumnsResolver<C>, val flag: SortFlag) : ColumnSet<C> {
-    override fun resolve(context: ColumnResolutionContext) = column.resolve(context).map { it.addFlag(flag) }
+    override fun resolve(context: ColumnResolutionContext) =
+        column.resolve(context).map {
+            if (context.allowMissingColumns && it.data is MissingColumnGroup<*>) it else it.addFlag(flag)
+        }
 }
 
 internal class SortColumnDescriptor<C>(
@@ -119,6 +157,12 @@ internal class SortColumnDescriptor<C>(
     val nullsLast: Boolean = false,
 ) : ValueColumnInternal<C> by column.internalValueColumn()
 
+private fun SortColumnDescriptor<*>.addFlag(flag: SortFlag): SortColumnDescriptor<*> =
+    when (flag) {
+        SortFlag.Reversed -> SortColumnDescriptor(column, direction.reversed(), nullsLast)
+        SortFlag.NullsLast -> SortColumnDescriptor(column, direction, true)
+    }
+
 internal enum class SortDirection { Asc, Desc }
 
 internal fun SortDirection.reversed(): SortDirection =

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/FrameColumnImpl.kt

@@ -53,11 +53,14 @@ internal open class FrameColumnImpl<T> constructor(
 
     override fun forceResolve() = ResolvingFrameColumn(this)
 
-    override fun get(indices: Iterable<Int>): FrameColumn<T> =
-        DataColumn.createFrameColumn(
+    override fun get(indices: Iterable<Int>): FrameColumn<T> {
+        val indicesList = indices.toList()
+        return DataColumn.createFrameColumn(
             name = name,
-            groups = indices.map { values[it] },
+            groups = indicesList.map { values[it] },
+            schema = schema.takeIf { indicesList.isEmpty() },
         )
+    }
 
     override fun get(columnName: String) =
         throw UnsupportedOperationException("Can not get nested column '$columnName' from FrameColumn '$name'")

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/sort.kt

@@ -1,7 +1,10 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.assertions.throwables.shouldThrowMessage
 import io.kotest.matchers.shouldBe
+import io.kotest.matchers.string.shouldContain
+import io.kotest.matchers.string.shouldNotContain
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.io.readCsv
@@ -99,4 +102,83 @@ class SortDataColumn {
             aggregate.sortBy { pathOf("L", "extra") }
         }
     }
+
+    @Test
+    fun `group by sort validates nested paths`() {
+        val df = DataFrame.readCsv(testResource("ds_salaries.csv")).cast<DsSalaries>()
+        val grouped = df.group { salaryInUsd }.into("group").groupBy { companyLocation }
+        val salaryPath = pathOf("group", "group", "salary_in_usd")
+
+        grouped.sortBy { salaryPath }.groups.values().forEach {
+            val salaries = it[pathOf("group", "salary_in_usd")].values().map { value -> value as Int }
+            salaries shouldBe salaries.sorted()
+        }
+        grouped.sortByDesc { salaryPath }.groups.values().forEach {
+            val salaries = it[pathOf("group", "salary_in_usd")].values().map { value -> value as Int }
+            salaries shouldBe salaries.sortedDescending()
+        }
+
+        val invalidPath = pathOf("group", "salaryInUsd")
+        val ascendingError = shouldThrow<IllegalStateException> {
+            grouped.sortBy { invalidPath }
+        }.message
+        val descendingError = shouldThrow<IllegalStateException> {
+            grouped.sortByDesc { invalidPath }
+        }.message
+
+        ascendingError shouldBe descendingError
+        ascendingError.orEmpty() shouldContain "group/salaryInUsd"
+        ascendingError.orEmpty() shouldNotContain "Can not apply sort flag to column kind"
+
+        shouldThrowMessage(ascendingError.orEmpty()) {
+            grouped.sortBy { salaryPath and invalidPath }
+        }
+    }
+
+    @Test
+    fun `group by sort preserves group and key sorting`() {
+        val df = DataFrame.readCsv(testResource("ds_salaries.csv")).cast<DsSalaries>()
+
+        df.groupBy { companyLocation }.sortBy { salaryInUsd }.groups.values().forEach {
+            val salaries = it[salaryInUsd].values()
+            salaries shouldBe salaries.sorted()
+        }
+
+        val grouped = df.group { salaryInUsd }.into("group").groupBy { companyLocation }
+        val sortedKeys = grouped.sortBy { companyLocation }.keys[companyLocation].values()
+        sortedKeys shouldBe sortedKeys.sorted()
+    }
+
+    @Test
+    fun `group by sort accepts columns present in later heterogeneous groups`() {
+        val grouped = dataFrameOf("key", "groups")(
+            "first",
+            dataFrameOf("other")(2),
+            "second",
+            dataFrameOf("value")(2, 1),
+        ).asGroupBy("groups")
+
+        grouped.sortBy("value").groups.values().toList()[1]["value"].values() shouldBe listOf(1, 2)
+    }
+
+    @Test
+    fun `group by sort validates paths when no groups are present`() {
+        val emptySource = dataFrameOf("value") { emptyList<Int>() }.groupBy("value")
+        val filteredGroups = dataFrameOf("value")(1).groupBy("value").filter { false }
+        val invalidPath = pathOf("missing", "nested")
+
+        shouldThrow<IllegalStateException> {
+            emptySource.sortBy { invalidPath }
+        }.message.orEmpty() shouldContain "missing/nested"
+        shouldThrow<IllegalStateException> {
+            filteredGroups.sortBy { invalidPath }
+        }.message.orEmpty() shouldContain "missing/nested"
+    }
+
+    @Test
+    fun `group by sort preserves schema after filtering all groups`() {
+        val filteredGroups = dataFrameOf("key", "value")(1, 2).groupBy("key").filter { false }
+
+        filteredGroups.sortBy("value").toDataFrame() shouldBe filteredGroups.toDataFrame()
+    }
 }