attempt to fix #1221 with simple cast, unsufficient for renaming, #663 handles it better

Author: Jolanrensen

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

@@ -99,3 +99,5 @@ public fun <C> TransformableColumnSet<*>.cast(): TransformableColumnSet<C> = thi
 public fun <C> TransformableSingleColumn<*>.cast(): TransformableSingleColumn<C> = this as TransformableSingleColumn<C>
 
 public fun <C> ColumnReference<*>.cast(): ColumnReference<C> = this as ColumnReference<C>
+
+public fun <T> GroupBy<*, *>.cast(): GroupBy<T, T> = this as GroupBy<T, T>

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

@@ -115,4 +115,5 @@ public class ReducedGroupBy<T, G>(
 }
 
 @PublishedApi
-internal fun <T, G> GroupBy<T, G>.reduce(reducer: Selector<DataFrame<G>, DataRow<G>?>) = ReducedGroupBy(this, reducer)
+internal fun <T, G> GroupBy<T, G>.reduce(reducer: Selector<DataFrame<G>, DataRow<G>?>): ReducedGroupBy<T, G> =
+    ReducedGroupBy(this, reducer)

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/ReplCodeGeneratorImpl.kt

@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.api.schema
 import org.jetbrains.kotlinx.dataframe.codeGen.Code
 import org.jetbrains.kotlinx.dataframe.codeGen.CodeGenerator
@@ -40,6 +41,7 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
         when (type.classifier) {
             DataFrame::class -> type.arguments[0].type?.jvmErasure
             DataRow::class -> type.arguments[0].type?.jvmErasure
+            GroupBy::class -> type.arguments[0].type?.jvmErasure
             else -> null
         }
 

dataframe-jupyter/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/Integration.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlinx.dataframe.api.SplitWithTransform
 import org.jetbrains.kotlinx.dataframe.api.Update
 import org.jetbrains.kotlinx.dataframe.api.asDataFrame
 import org.jetbrains.kotlinx.dataframe.api.columnsCount
+import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
 import org.jetbrains.kotlinx.dataframe.codeGen.CodeGenerator
 import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
@@ -153,6 +154,17 @@ internal class Integration(private val notebook: Notebook, private val options:
             null
         }
 
+    private fun KotlinKernelHost.updateGroupByVariable(
+        groupBy: GroupBy<*, *>,
+        property: KProperty<*>,
+        codeGen: ReplCodeGenerator,
+    ): VariableName? =
+        execute(
+            codeWithConverter = codeGen.process(groupBy.concat(), property),
+            property = property,
+            type = GroupBy::class.createStarProjectedType(false),
+        )
+
     override fun Builder.onLoaded() {
         if (version != null) {
             if (enableExperimentalCsv?.toBoolean() == true) {
@@ -283,12 +295,14 @@ internal class Integration(private val notebook: Notebook, private val options:
 
         addTypeConverter(object : FieldHandler {
             override val execution = FieldHandlerFactory.createUpdateExecution<Any> { instance, property ->
+                // TODO check #1245 here
                 when (instance) {
                     is AnyCol -> updateAnyColVariable(instance, property, codeGen)
                     is ColumnGroup<*> -> updateColumnGroupVariable(instance, property, codeGen)
                     is AnyRow -> updateAnyRowVariable(instance, property, codeGen)
                     is AnyFrame -> updateAnyFrameVariable(instance, property, codeGen)
                     is ImportDataSchema -> updateImportDataSchemaVariable(instance, property)
+                    is GroupBy<*, *> -> updateGroupByVariable(instance, property, codeGen)
                     else -> error("${instance::class} should not be handled by Dataframe field handler")
                 }
             }
@@ -298,7 +312,8 @@ internal class Integration(private val notebook: Notebook, private val options:
                     value is ColumnGroup<*> ||
                     value is AnyRow ||
                     value is AnyFrame ||
-                    value is ImportDataSchema
+                    value is ImportDataSchema ||
+                    value is GroupBy<*, *>
         })
 
         fun KotlinKernelHost.addDataSchemas(classes: List<KClass<*>>) {

dataframe-jupyter/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/CodeGenerationTests.kt

@@ -92,4 +92,39 @@ class CodeGenerationTests : DataFrameJupyterTest() {
             df2.group.a
             """.checkCompilation()
     }
+
+    // Issue #1221
+    @Test
+    fun `GroupBy code generation`() {
+        @Language("kts")
+        val a = """
+            val ab = dataFrameOf("a", "b")(1, 2)
+            ab.groupBy { a }.aggregate { sum { b } into "bSum" }
+        """.checkCompilation()
+
+        @Language("kts")
+        val b = """
+            val ab = dataFrameOf("a", "b")(1, 2)
+            val grouped = ab.groupBy { a }
+            grouped.aggregate { sum { b } into "bSum" }
+        """.checkCompilation()
+
+        @Language("kts")
+        val c = """
+            val grouped = dataFrameOf("a", "b")(1, 2).groupBy("a")
+            grouped.aggregate { sum { b } into "bSum" }
+        """.checkCompilation()
+
+        @Language("kts")
+        val d = """
+            val grouped = dataFrameOf("a", "b")(1, 2).groupBy("a")
+            grouped.keys.a
+        """.checkCompilation()
+
+        @Language("kts")
+        val e = """
+            val grouped = dataFrameOf("a", "b")(1, 2).groupBy { "a"<Int>() named "k" }
+            grouped.keys.k
+        """.checkCompilation()
+    }
 }