Merge pull request #781 from Kotlin/convert-as-frame

Add Convert.asFrame function
Kotlin · Jul 17, 2024 · 0e97e5b · 0e97e5b
2 parents b253e1c + efcef69
commit 0e97e5b
Show file tree

Hide file tree

Showing 5 changed files with 532 additions and 1 deletion.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/convert.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/convert.kt
@@ -8,12 +8,15 @@ import kotlinx.datetime.toLocalDateTime
 import org.jetbrains.kotlinx.dataframe.AnyBaseCol
 import org.jetbrains.kotlinx.dataframe.AnyCol
 import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.ColumnsContainer
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowColumnExpression
 import org.jetbrains.kotlinx.dataframe.RowValueExpression
 import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.dataTypes.IFRAME
@@ -111,6 +114,9 @@ public inline fun <T, C, reified R> Convert<T, C>.with(
     noinline rowConverter: RowValueExpression<T, C, R>
 ): DataFrame<T> = with(Infer.Nulls, rowConverter)
 
+public fun <T, C, R> Convert<T, DataRow<C>>.asFrame(body: ColumnsContainer<T>.(ColumnGroup<C>) -> DataFrame<R>): DataFrame<T> =
+    to { body(this, it.asColumnGroup()).asColumnGroup(it.name()) }
+
 public inline fun <T, C, reified R> Convert<T, C>.perRowCol(
     infer: Infer = Infer.Nulls,
     noinline expression: RowColumnExpression<T, C, R>,

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt
@@ -210,6 +210,14 @@ class Modify : TestBase() {
         // SampleEnd
     }
 
+    @Test
+    @TransformDataFrameExpressions
+    fun convertAsFrame() {
+        // SampleStart
+        df.convert { name }.asFrame { it.add("fullName") { "$firstName $lastName" } }
+        // SampleEnd
+    }
+
     @Test
     @TransformDataFrameExpressions
     fun parseAll() {

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt
@@ -661,4 +661,10 @@ class DataFrameTreeTests : BaseTest() {
         val a: DataFrame<NameAndCity> = typed2["nameAndCity"].cast<DataRow<NameAndCity>>().asDataFrame()
         val b: DataFrame<NameAndCity> = typed2[nameAndCity].cast<NameAndCity>().asDataFrame()
     }
+
+    @Test
+    fun `convert ColumnGroup as Frame`() {
+        val df = typed2.convert { nameAndCity }.asFrame { it.remove { city } }
+        df.nameAndCity.columns() shouldBe typed2.nameAndCity.remove { city }.columns()
+    }
 }