From 20cab7a12968babcb7f9d8ce8a24e3ca0917eda1 Mon Sep 17 00:00:00 2001 From: Jolan Rensen Date: Mon, 1 Jul 2024 11:49:19 +0200 Subject: [PATCH] providing unfold {}.by {} and replace {}.byUnfolding {} sample API --- .../jetbrains/kotlinx/dataframe/api/unfold.kt | 63 ++++++++++++---- .../kotlinx/dataframe/api/replace.kt | 6 +- .../jetbrains/kotlinx/dataframe/api/unfold.kt | 73 +++++++++++++++++++ 3 files changed, 125 insertions(+), 17 deletions(-) create mode 100644 core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt index 00831c7859..fb364efa02 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt @@ -1,34 +1,69 @@ package org.jetbrains.kotlinx.dataframe.api import org.jetbrains.kotlinx.dataframe.AnyCol -import org.jetbrains.kotlinx.dataframe.AnyColumnReference import org.jetbrains.kotlinx.dataframe.ColumnsSelector import org.jetbrains.kotlinx.dataframe.DataColumn import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.annotations.Interpretable import org.jetbrains.kotlinx.dataframe.annotations.Refine +import org.jetbrains.kotlinx.dataframe.columns.ColumnReference import org.jetbrains.kotlinx.dataframe.columns.toColumnSet import org.jetbrains.kotlinx.dataframe.impl.api.unfoldImpl import kotlin.reflect.KProperty -public inline fun DataColumn.unfold(vararg props: KProperty<*>, maxDepth: Int = 0): AnyCol = - unfoldImpl(skipPrimitive = true) { properties(roots = props, maxDepth = maxDepth) } +// region DataColumn -public inline fun DataColumn.unfold(noinline body: CreateDataFrameDsl.() -> Unit): AnyCol = +public fun DataColumn.unfold( + vararg props: KProperty<*>, + maxDepth: Int = 0, +): AnyCol = unfoldImpl(skipPrimitive = true) { properties(roots = props, maxDepth = maxDepth) } + +public fun DataColumn.unfold(body: CreateDataFrameDsl.() -> Unit): AnyCol = unfoldImpl(skipPrimitive = false, body) -public inline fun ReplaceClause.unfold(vararg props: KProperty<*>, maxDepth: Int = 0): DataFrame = - with { it.unfold(props = props, maxDepth) } +// endregion +// region DataFrame -@Refine -@Interpretable("ReplaceUnfold1") -public inline fun ReplaceClause.unfold(noinline body: CreateDataFrameDsl.() -> Unit): DataFrame = - with { it.unfoldImpl(skipPrimitive = false, body) } +public fun DataFrame.unfold(columns: ColumnsSelector): UnfoldingDataFrame = + UnfoldingDataFrame( + originalDf = this, + unfoldedDf = replace(columns).with { it.unfold() }, + columns = columns, + ) + +public fun DataFrame.unfold(vararg columns: String): UnfoldingDataFrame = unfold { columns.toColumnSet() } + +public fun DataFrame.unfold(vararg columns: ColumnReference): UnfoldingDataFrame = + unfold { columns.toColumnSet() } + +public fun DataFrame.unfold(vararg columns: KProperty): UnfoldingDataFrame = + unfold { columns.toColumnSet() } -public fun DataFrame.unfold(columns: ColumnsSelector): DataFrame = replace(columns).with { it.unfold() } +public class UnfoldingDataFrame( + internal val originalDf: DataFrame, + private val unfoldedDf: DataFrame, + internal val columns: ColumnsSelector, +) : DataFrame by unfoldedDf -public fun DataFrame.unfold(vararg columns: String): DataFrame = unfold { columns.toColumnSet() } +public fun UnfoldingDataFrame.by(body: CreateDataFrameDsl.() -> Unit): DataFrame = + originalDf.replace(columns).with { it.unfoldImpl(skipPrimitive = false, body) } -public fun DataFrame.unfold(vararg columns: AnyColumnReference): DataFrame = unfold { columns.toColumnSet() } +public fun UnfoldingDataFrame.by( + vararg props: KProperty<*>, + maxDepth: Int = 0, +): DataFrame = originalDf.replace(*props).with { it.unfold(props = props, maxDepth) } + +// endregion +// region replace + +public fun ReplaceClause.byUnfolding( + vararg props: KProperty<*>, + maxDepth: Int = 0, +): DataFrame = with { it.unfold(props = props, maxDepth) } + +@Refine +@Interpretable("ReplaceUnfold1") +public fun ReplaceClause.byUnfolding(body: CreateDataFrameDsl.() -> Unit): DataFrame = + with { it.unfoldImpl(skipPrimitive = false, body) } -public fun DataFrame.unfold(vararg columns: KProperty<*>): DataFrame = unfold { columns.toColumnSet() } +// endregion diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt index 09c038eabc..063ffca037 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt @@ -21,7 +21,7 @@ class ReplaceTests { val a by columnOf("123") val df = dataFrameOf(a) - val conv = df.replace { a }.unfold { + val conv = df.replace { a }.byUnfolding { "b" from { it } "c" from { DataRow.readJsonStr("""{"prop": 1}""") } } @@ -39,7 +39,7 @@ class ReplaceTests { fun `unfold properties`() { val col by columnOf(A("1", 123, B(3.0))) val df1 = dataFrameOf(col) - val conv = df1.replace { col }.unfold(maxDepth = 2) + val conv = df1.replace { col }.byUnfolding(maxDepth = 2) val a = conv["col"]["a"] a.type() shouldBe typeOf() @@ -62,7 +62,7 @@ class ReplaceTests { val col1 by columnOf("1", "2") val col2 by columnOf(B(1.0), B(2.0)) val df1 = dataFrameOf(col1, col2) - val conv = df1.replace { nameStartsWith("col") }.unfold() + val conv = df1.replace { nameStartsWith("col") }.byUnfolding() val a = conv["col1"] a.type() shouldBe typeOf() diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt new file mode 100644 index 0000000000..268b6aa369 --- /dev/null +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt @@ -0,0 +1,73 @@ +package org.jetbrains.kotlinx.dataframe.api + +import io.kotest.matchers.shouldBe +import org.jetbrains.kotlinx.dataframe.DataRow +import org.jetbrains.kotlinx.dataframe.io.readJsonStr +import org.junit.Test +import kotlin.reflect.typeOf + +class UnfoldTests { + @Test + fun `unfold primitive`() { + val a by columnOf("123") + val df = dataFrameOf(a) + + val conv = df.unfold { a }.by { + "b" from { it } + "c" from { DataRow.readJsonStr("""{"prop": 1}""") } + } + + val b = conv["a"]["b"] + b.type() shouldBe typeOf() + b.values() shouldBe listOf("123") + + val c = conv["a"]["c"]["prop"] + c.type() shouldBe typeOf() + c.values() shouldBe listOf(1) + } + + @Test + fun aaa() { + val col by columnOf(A("1", 123, B(3.0))) + col.unfold().print() + } + + @Test + fun `unfold properties`() { + val col by columnOf(A("1", 123, B(3.0))) + val df1 = dataFrameOf(col) + // TODO `df1.replace { col }.with { it.unfold() }` breaks now + val conv = df1.unfold { col }.by(maxDepth = 2) + + val a = conv["col"]["a"] + a.type() shouldBe typeOf() + a.values() shouldBe listOf("1") + + val b = conv["col"]["b"] + b.type() shouldBe typeOf() + b.values() shouldBe listOf(123) + + val d = conv["col"]["bb"]["d"] + d.type() shouldBe typeOf() + d.values() shouldBe listOf(3.0) + } + + class B(val d: Double) + class A(val a: String, val b: Int, val bb: B) + + @Test + fun `skip primitive`() { + val col1 by columnOf("1", "2") + val col2 by columnOf(B(1.0), B(2.0)) + val df1 = dataFrameOf(col1, col2) + val conv = df1.unfold { nameStartsWith("col") } + + val a = conv["col1"] + a.type() shouldBe typeOf() + a.values() shouldBe listOf("1", "2") + + val b = conv["col2"]["d"] + b.type() shouldBe typeOf() + b.values() shouldBe listOf(1.0, 2.0) + } +}