small fix for one failing test in #713, where GroupBy.aggregate { thi…

…s into "" } would break typing
Kotlin · Jul 31, 2024 · cb374a4 · cb374a4
1 parent 3037a97
commit cb374a4
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/GroupByReceiverImpl.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/GroupByReceiverImpl.kt
@@ -6,6 +6,8 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedDsl
 import org.jetbrains.kotlinx.dataframe.aggregation.NamedValue
 import org.jetbrains.kotlinx.dataframe.api.asDataFrame
+import org.jetbrains.kotlinx.dataframe.api.cast
+import org.jetbrains.kotlinx.dataframe.api.convertTo
 import org.jetbrains.kotlinx.dataframe.api.toDataFrameFromPairs
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
@@ -118,7 +120,18 @@ internal class GroupByReceiverImpl<T>(override val df: DataFrame<T>, override va
                 pivot.aggregator.values.clear()
             }
 
-            is AggregateInternalDsl<*> -> yield(value.copy(value = value.value.df))
+            is AggregateInternalDsl<*> -> {
+                // Attempt to create DataFrame<Type> from AggregateInternalDsl<Type>
+                val dfType = value.type?.arguments?.firstOrNull()?.type?.let {
+                    DataFrame::class.createTypeWithArgument(it)
+                }
+                yield(
+                    value.copy(
+                        value = value.value.df,
+                        type = dfType ?: value.type,
+                    ),
+                )
+            }
 
             else -> values.add(value)
         }