Skip to content

Commit efcef69

Browse files
committed
Add Convert.asFrame function
1 parent f5a06eb commit efcef69

File tree

5 files changed

+532
-1
lines changed

5 files changed

+532
-1
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,15 @@ import kotlinx.datetime.toLocalDateTime
88
import org.jetbrains.kotlinx.dataframe.AnyBaseCol
99
import org.jetbrains.kotlinx.dataframe.AnyCol
1010
import org.jetbrains.kotlinx.dataframe.AnyFrame
11+
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
1112
import org.jetbrains.kotlinx.dataframe.ColumnsSelector
1213
import org.jetbrains.kotlinx.dataframe.DataColumn
1314
import org.jetbrains.kotlinx.dataframe.DataFrame
15+
import org.jetbrains.kotlinx.dataframe.DataRow
1416
import org.jetbrains.kotlinx.dataframe.RowColumnExpression
1517
import org.jetbrains.kotlinx.dataframe.RowValueExpression
1618
import org.jetbrains.kotlinx.dataframe.annotations.*
19+
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
1720
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
1821
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
1922
import org.jetbrains.kotlinx.dataframe.dataTypes.IFRAME
@@ -111,6 +114,9 @@ public inline fun <T, C, reified R> Convert<T, C>.with(
111114
noinline rowConverter: RowValueExpression<T, C, R>
112115
): DataFrame<T> = with(Infer.Nulls, rowConverter)
113116

117+
public fun <T, C, R> Convert<T, DataRow<C>>.asFrame(body: ColumnsContainer<T>.(ColumnGroup<C>) -> DataFrame<R>): DataFrame<T> =
118+
to { body(this, it.asColumnGroup()).asColumnGroup(it.name()) }
119+
114120
public inline fun <T, C, reified R> Convert<T, C>.perRowCol(
115121
infer: Infer = Infer.Nulls,
116122
noinline expression: RowColumnExpression<T, C, R>,

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,14 @@ class Modify : TestBase() {
210210
// SampleEnd
211211
}
212212

213+
@Test
214+
@TransformDataFrameExpressions
215+
fun convertAsFrame() {
216+
// SampleStart
217+
df.convert { name }.asFrame { it.add("fullName") { "$firstName $lastName" } }
218+
// SampleEnd
219+
}
220+
213221
@Test
214222
@TransformDataFrameExpressions
215223
fun parseAll() {

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -661,4 +661,10 @@ class DataFrameTreeTests : BaseTest() {
661661
val a: DataFrame<NameAndCity> = typed2["nameAndCity"].cast<DataRow<NameAndCity>>().asDataFrame()
662662
val b: DataFrame<NameAndCity> = typed2[nameAndCity].cast<NameAndCity>().asDataFrame()
663663
}
664+
665+
@Test
666+
fun `convert ColumnGroup as Frame`() {
667+
val df = typed2.convert { nameAndCity }.asFrame { it.remove { city } }
668+
df.nameAndCity.columns() shouldBe typed2.nameAndCity.remove { city }.columns()
669+
}
664670
}

0 commit comments

Comments
 (0)