providing unfold {}.by {} and replace {}.byUnfolding {} sample API

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 <reified T> DataColumn<T>.unfold(vararg props: KProperty<*>, maxDepth: Int = 0): AnyCol =
-    unfoldImpl(skipPrimitive = true) { properties(roots = props, maxDepth = maxDepth) }
+// region DataColumn
 
-public inline fun <reified T> DataColumn<T>.unfold(noinline body: CreateDataFrameDsl<T>.() -> Unit): AnyCol =
+public fun <T> DataColumn<T>.unfold(
+    vararg props: KProperty<*>,
+    maxDepth: Int = 0,
+): AnyCol = unfoldImpl(skipPrimitive = true) { properties(roots = props, maxDepth = maxDepth) }
+
+public fun <T> DataColumn<T>.unfold(body: CreateDataFrameDsl<T>.() -> Unit): AnyCol =
     unfoldImpl(skipPrimitive = false, body)
 
-public inline fun <T, reified C> ReplaceClause<T, C>.unfold(vararg props: KProperty<*>, maxDepth: Int = 0): DataFrame<T> =
-    with { it.unfold(props = props, maxDepth) }
+// endregion
+// region DataFrame
 
-@Refine
-@Interpretable("ReplaceUnfold1")
-public inline fun <T, reified C> ReplaceClause<T, C>.unfold(noinline body: CreateDataFrameDsl<C>.() -> Unit): DataFrame<T> =
-    with { it.unfoldImpl(skipPrimitive = false, body) }
+public fun <T, C> DataFrame<T>.unfold(columns: ColumnsSelector<T, C>): UnfoldingDataFrame<T, C> =
+    UnfoldingDataFrame(
+        originalDf = this,
+        unfoldedDf = replace(columns).with { it.unfold() },
+        columns = columns,
+    )
+
+public fun <T> DataFrame<T>.unfold(vararg columns: String): UnfoldingDataFrame<T, *> = unfold { columns.toColumnSet() }
+
+public fun <T, C> DataFrame<T>.unfold(vararg columns: ColumnReference<C>): UnfoldingDataFrame<T, C> =
+    unfold { columns.toColumnSet() }
+
+public fun <T, C> DataFrame<T>.unfold(vararg columns: KProperty<C>): UnfoldingDataFrame<T, C> =
+    unfold { columns.toColumnSet() }
 
-public fun <T> DataFrame<T>.unfold(columns: ColumnsSelector<T, *>): DataFrame<T> = replace(columns).with { it.unfold() }
+public class UnfoldingDataFrame<T, C>(
+    internal val originalDf: DataFrame<T>,
+    private val unfoldedDf: DataFrame<T>,
+    internal val columns: ColumnsSelector<T, C>,
+) : DataFrame<T> by unfoldedDf
 
-public fun <T> DataFrame<T>.unfold(vararg columns: String): DataFrame<T> = unfold { columns.toColumnSet() }
+public fun <T, C> UnfoldingDataFrame<T, C>.by(body: CreateDataFrameDsl<C>.() -> Unit): DataFrame<T> =
+    originalDf.replace(columns).with { it.unfoldImpl(skipPrimitive = false, body) }
 
-public fun <T> DataFrame<T>.unfold(vararg columns: AnyColumnReference): DataFrame<T> = unfold { columns.toColumnSet() }
+public fun <T, C> UnfoldingDataFrame<T, C>.by(
+    vararg props: KProperty<*>,
+    maxDepth: Int = 0,
+): DataFrame<T> = originalDf.replace(*props).with { it.unfold(props = props, maxDepth) }
+
+// endregion
+// region replace
+
+public fun <T, C> ReplaceClause<T, C>.byUnfolding(
+    vararg props: KProperty<*>,
+    maxDepth: Int = 0,
+): DataFrame<T> = with { it.unfold(props = props, maxDepth) }
+
+@Refine
+@Interpretable("ReplaceUnfold1")
+public fun <T, C> ReplaceClause<T, C>.byUnfolding(body: CreateDataFrameDsl<C>.() -> Unit): DataFrame<T> =
+    with { it.unfoldImpl(skipPrimitive = false, body) }
 
-public fun <T> DataFrame<T>.unfold(vararg columns: KProperty<*>): DataFrame<T> = unfold { columns.toColumnSet() }
+// endregion

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<String>()
@@ -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<String>()

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<String>()
+        b.values() shouldBe listOf("123")
+
+        val c = conv["a"]["c"]["prop"]
+        c.type() shouldBe typeOf<Int>()
+        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<String>()
+        a.values() shouldBe listOf("1")
+
+        val b = conv["col"]["b"]
+        b.type() shouldBe typeOf<Int>()
+        b.values() shouldBe listOf(123)
+
+        val d = conv["col"]["bb"]["d"]
+        d.type() shouldBe typeOf<Double>()
+        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<String>()
+        a.values() shouldBe listOf("1", "2")
+
+        val b = conv["col2"]["d"]
+        b.type() shouldBe typeOf<Double>()
+        b.values() shouldBe listOf(1.0, 2.0)
+    }
+}