init research

2026-02-08 11:20:43 -10:00
commit bdf064f54d
3041 changed files with 1592200 additions and 0 deletions
@@ -0,0 +1,156 @@
+package org.jetbrains.kotlinx.dataframe.explainer
+
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.api.*
+
+@TransformDataFrameExpressions
+fun callChain(df: DataFrame<*>) {
+    val df1 = df
+        .filter { "age"<Int>() > 20 }
+        .groupBy("value")
+        .sum()
+}
+
+interface Person
+
+class Wrapper {
+    val df = dataFrameOf("name", "age", "city", "weight")(
+        "Alice", 15, "London", 54,
+        "Bob", 45, "Dubai", 87,
+        "Charlie", 20, "Moscow", null,
+        "Charlie", 40, "Milan", null,
+        "Bob", 30, "Tokyo", 68,
+        "Alice", 20, null, 55,
+        "Charlie", 30, "Moscow", 90
+    )
+
+    val typed: DataFrame<Person> = df.cast()
+
+    @TransformDataFrameExpressions
+    fun ff() {
+        val name by column<String>()
+        val aggregated = typed.groupBy { name() }.aggregate {
+            (if (name().first().startsWith("A")) first() else null) into "agg"
+        }["agg"]
+    }
+}
+
+@TransformDataFrameExpressions
+fun aggregateDf() {
+    val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
+        "Alice", "Cooper", 15, "London", 54, true,
+        "Bob", "Dylan", 45, "Dubai", 87, true,
+        "Charlie", "Daniels", 20, "Moscow", null, false,
+        "Charlie", "Chaplin", 40, "Milan", null, true,
+        "Bob", "Marley", 30, "Tokyo", 68, true,
+        "Alice", "Wolf", 20, null, 55, false,
+        "Charlie", "Byrd", 30, "Moscow", 90, true
+    ).group("firstName", "lastName").into("name")
+
+    df.groupBy("city").aggregate {
+        count() into "total"
+        count { "age"<Int>() > 18 } into "adults"
+        median("age") into "median age"
+        min("age") into "min age"
+        maxBy("age")["name"] into "oldest"
+    }
+}
+
+@TransformDataFrameExpressions
+fun move() {
+    val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
+        "Alice", "Cooper", 15, "London", 54, true,
+        "Bob", "Dylan", 45, "Dubai", 87, true,
+        "Charlie", "Daniels", 20, "Moscow", null, false,
+        "Charlie", "Chaplin", 40, "Milan", null, true,
+        "Bob", "Marley", 30, "Tokyo", 68, true,
+        "Alice", "Wolf", 20, null, 55, false,
+        "Charlie", "Byrd", 30, "Moscow", 90, true
+    ).group("firstName", "lastName").into("name")
+
+    df.move("age", "weight").into { pathOf("info", it.name()) }
+}
+
+fun interface PluginCallback {
+    fun doAction(
+        source: String,
+        name: String,
+        df: Any,
+        id: String,
+        receiverId: String?,
+        containingClassFqName: String?,
+        containingFunName: String?,
+        statementIndex: Int
+    )
+}
+
+object PluginCallbackProxy : PluginCallback {
+
+    var action: PluginCallback = PluginCallback { _, _, _, _, _, _, _, _  -> Unit }
+
+    override fun doAction(source: String, name: String, df: Any, id: String, receiverId: String?, containingClassFqName: String?, containingFunName: String?, statemenIndex: Int) {
+        action.doAction(source, name, df, id, receiverId, containingClassFqName, containingFunName, statemenIndex)
+    }
+}
+
+//fun callChainTransformed(df: DataFrame<*>) {
+//    val df1 = df
+//        .filter { "age"<Int>() > 20 }
+//        .also { PluginCallbackProxy.doAction(""".filter { "age"<Int>() > 20 }""", "filter", it) }
+//}
+
+//fun callChainTransformed(df: DataFrame<*>) {
+//    val df1 = df
+//        .filter { "age"<Int>() > 20 }
+//        .also { PluginCallbackProxy.action(""".filter { "age"<Int>() > 20 }""", it) }
+//        .groupBy("something")
+//        .sum()
+//        .also { PluginCallbackProxy.action(""".groupBy("something").sum()""", it) }
+//}
+
+annotation class TransformDataFrameExpressions
+
+fun box(): String {
+    val age by columnOf(10, 21, 30, 1)
+    val value by columnOf("a", "b", "c", "c")
+    val df = dataFrameOf(age, value)
+    val expressions = mutableListOf<String>()
+    PluginCallbackProxy.action = PluginCallback { source, _, df, id, receiverId, containingClassFqName, containingFunName, statementIndex ->
+        println("== Call ==")
+        expressions += source
+        if (df is AnyFrame) {
+            println(source)
+            df.print()
+            println(id)
+            println(receiverId)
+            println(containingClassFqName)
+            println(containingFunName)
+            println("statementIndex = ${statementIndex}")
+        } else {
+            println(df::class)
+        }
+        println("== End ==")
+    }
+
+    println("CallChain")
+    callChain(df)
+    println("expressions = ${expressions}")
+    expressions.clear()
+
+    println("ff")
+    Wrapper().ff()
+    println("expressions = ${expressions}")
+    expressions.clear()
+//    callChainTransformed(df)
+
+    println("aggregateDf")
+    aggregateDf()
+    println("expressions = ${expressions}")
+    expressions.clear()
+
+    println("move")
+    move()
+    println("expressions = ${expressions}")
+    expressions.clear()
+    return "OK"
+}