diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt index 7d83957bb7..dae7e29162 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt @@ -2,14 +2,22 @@ package org.jetbrains.kotlinx.dataframe.api import org.jetbrains.kotlinx.dataframe.ColumnsSelector import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.columns.ColumnReference import org.jetbrains.kotlinx.dataframe.impl.api.flattenImpl +import org.jetbrains.kotlinx.dataframe.impl.columns.toColumns +import kotlin.reflect.KProperty // region DataFrame public fun DataFrame.flatten(): DataFrame = flatten { all() } -public fun DataFrame.flatten( - columns: ColumnsSelector -): DataFrame = flattenImpl(columns) +public fun DataFrame.flatten(columns: ColumnsSelector): DataFrame = flattenImpl(columns) + +public fun DataFrame.flatten(vararg columns: String): DataFrame = flattenImpl { columns.toColumns() } + +public fun DataFrame.flatten(vararg columns: KProperty): DataFrame = flattenImpl { columns.toColumns() } + +public fun DataFrame.flatten(vararg columns: ColumnReference): DataFrame = + flattenImpl { columns.toColumns() } // endregion diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt index 1865a86d03..7a376daa07 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/flatten.kt @@ -1,6 +1,8 @@ package org.jetbrains.kotlinx.dataframe.api import io.kotest.matchers.shouldBe +import org.jetbrains.kotlinx.dataframe.DataRow +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema import org.junit.Test class FlattenTests { @@ -13,6 +15,41 @@ class FlattenTests { grouped.add("a") { 0 }.flatten().columnNames() shouldBe listOf("a1", "b", "c", "a") } + @DataSchema + interface TestRow { + val a: String + val b: String + val c: String + } + + @DataSchema + interface Grouped { + val d: DataRow + } + + @Test + fun `flatten access APIs`() { + val df = dataFrameOf("a", "b", "c")(1, 2, 3) + val grouped = df.group("a", "b").into("d") + + // String API + grouped.flatten("d") shouldBe df + val castedGroupedDF = grouped.cast() + + // KProperties API + castedGroupedDF.flatten(Grouped::d) shouldBe df + + // Extension properties API + castedGroupedDF.flatten { d } shouldBe df + + // Column accessors API + val d by columnGroup() + val a by d.column() + val b by d.column() + val c by d.column() + grouped.flatten(d) shouldBe df + } + @Test fun `flatten nested`() { val df = dataFrameOf("a", "b", "c", "d")(1, 2, 3, 4) 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 index d855c3ec63..7c32b9d3e0 100644 --- 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 @@ -1052,7 +1052,7 @@ class Modify : TestBase() { } @Test - fun flatten() { + fun flatten_properties() { // SampleStart // name.firstName -> firstName // name.lastName -> lastName @@ -1060,6 +1060,36 @@ class Modify : TestBase() { // SampleEnd } + @Test + fun flatten_strings() { + // SampleStart + // name.firstName -> firstName + // name.lastName -> lastName + df.flatten("name") + // SampleEnd + } + + @Test + fun flatten_accessors() { + // SampleStart + val name by columnGroup() + val firstName by name.column() + val lastName by name.column() + // name.firstName -> firstName + // name.lastName -> lastName + df.flatten(name) + // SampleEnd + } + + @Test + fun flatten_KProperties() { + // SampleStart + // name.firstName -> firstName + // name.lastName -> lastName + df.flatten(df::name) + // SampleEnd + } + @Test fun flattenAll() { // SampleStart diff --git a/docs/StardustDocs/topics/flatten.md b/docs/StardustDocs/topics/flatten.md index 929c819756..29ed6dc9f6 100644 --- a/docs/StardustDocs/topics/flatten.md +++ b/docs/StardustDocs/topics/flatten.md @@ -11,6 +11,8 @@ flatten [ { columns } ] Columns after flattening will keep their original names. Potential column name clashes are resolved by adding minimal possible name prefix from ancestor columns. + + ```kotlin // name.firstName -> firstName @@ -18,6 +20,27 @@ Columns after flattening will keep their original names. Potential column name c df.flatten { name } ``` + + + +```kotlin +val name by columnGroup() +val firstName by name.column() +val lastName by name.column() + +// name.firstName -> firstName +// name.lastName -> lastName +df.flatten(name) +``` + + + + +```kotlin +df.flatten("name") +``` + + To remove all column groupings in [`DataFrame`](DataFrame.md), invoke `flatten` without parameters: