poc implemented, needs more tests

Author: Jolanrensen

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataColumn.kt

@@ -107,11 +107,11 @@ public interface DataColumn<out T> : BaseColumn<T> {
 
     override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<T>? = this.addPath()
 
-    override fun resolveSingleAfter(
+    override fun resolveSingleAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): ColumnWithPath<T>? =
-        transform(this).cast<T>()
+        transformer.transformRemainingSingle(this).cast<T>()
             .let { it as DataColumn<T> }
             .addPath()
 

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl.kt

@@ -111,7 +111,7 @@ public class ColumnMatch<C>(public val left: ColumnReference<C>, public val righ
 
     override fun resolveAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): List<ColumnWithPath<C>> {
         throw UnsupportedOperationException()
     }

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/join.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.columns.RenamedColumnReference
 import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.getColumn
-import org.jetbrains.kotlinx.dataframe.impl.columns.transform
 import kotlin.reflect.KProperty
 
 /**
@@ -38,11 +37,11 @@ public interface ColumnReference<out C> : SingleColumn<C> {
             .getColumn<C>(path(), context.unresolvedColumnsPolicy)
             ?.addPath(path())
 
-    override fun resolveSingleAfter(
+    override fun resolveSingleAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): ColumnWithPath<C>? =
-        transform(context.df.asColumnGroup()).cast<C>()
+        transformer.transformRemainingSingle(context.df.asColumnGroup()).cast<C>()
             .resolve(context)
             .toDataFrame()
             .getColumn<C>(path(), context.unresolvedColumnsPolicy)

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnReference.kt

@@ -21,15 +21,22 @@ public interface ColumnSet<out C> {
     public fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>>
 
     /**
-     * Resolves this [ColumnSet] as a [List]<[ColumnWithPath]<[C]>> after applying [transform] to the parent
+     * Resolves this [ColumnSet] as a [List]<[ColumnWithPath]<[C]>> after applying [transformer] to the parent
      * [ColumnSet]. This essentially injects a call right before the current in the [ColumnSet.resolve] chain.
      */
     public fun resolveAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): List<ColumnWithPath<C>>
 }
 
+public interface ColumnSetTransformer {
+
+    public fun transformRemainingSingle(singleColumn: SingleColumn<*>): SingleColumn<*>
+
+    public fun transform(columnSet: ColumnSet<*>): ColumnSet<*>
+}
+
 public class ColumnResolutionContext internal constructor(
     internal val df: DataFrame<*>,
     internal val unresolvedColumnsPolicy: UnresolvedColumnsPolicy,

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnSet.kt

@@ -15,12 +15,22 @@ public interface ColumnWithPath<out T> : DataColumn<T> {
     public val name: String get() = name()
     public val parentName: String? get() = path.parentName
     public fun depth(): Int = path.depth()
-    public fun <C> getChild(accessor: ColumnReference<C>): ColumnWithPath<C>? = asColumnGroup().getColumnOrNull(accessor)?.addPath(path + accessor.path())
-    public fun getChild(name: String): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(name)?.addParentPath(path)
-    public fun getChild(index: Int): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(index)?.addParentPath(path)
-    public fun <C> getChild(accessor: KProperty<C>): ColumnWithPath<C>? = asColumnGroup().getColumnOrNull(accessor)?.addParentPath(path)
+    public fun <C> getChild(accessor: ColumnReference<C>): ColumnWithPath<C>? =
+        asColumnGroup().getColumnOrNull(accessor)?.addPath(path + accessor.path())
 
-    public fun children(): List<ColumnWithPath<Any?>> = if (isColumnGroup()) data.asColumnGroup().columns().map { it.addParentPath(path) } else emptyList()
+    public fun getChild(name: String): ColumnWithPath<Any?>? =
+        asColumnGroup().getColumnOrNull(name)?.addParentPath(path)
+
+    public fun getChild(index: Int): ColumnWithPath<Any?>? =
+        asColumnGroup().getColumnOrNull(index)?.addParentPath(path)
+    public fun <C> getChild(accessor: KProperty<C>): ColumnWithPath<C>? =
+        asColumnGroup().getColumnOrNull(accessor)?.addParentPath(path)
+
+    public fun children(): List<ColumnWithPath<Any?>> =
+        if (isColumnGroup())
+            data.asColumnGroup().columns().map { it.addParentPath(path) }
+        else
+            emptyList()
 
     override fun path(): ColumnPath = path
 

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnWithPath.kt

@@ -15,13 +15,13 @@ public interface SingleColumn<out C> : ColumnSet<C> {
 
     override fun resolveAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
-    ): List<ColumnWithPath<C>> = resolveSingleAfter(context, transform)?.let { listOf(it) } ?: emptyList()
+        transformer: ColumnSetTransformer,
+    ): List<ColumnWithPath<C>> = resolveSingleAfterTransform(context, transformer)?.let { listOf(it) } ?: emptyList()
 
     public fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>?
 
-    public fun resolveSingleAfter(
+    public fun resolveSingleAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): ColumnWithPath<C>?
 }

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/SingleColumn.kt

@@ -55,14 +55,14 @@ internal open class DataFrameReceiver<T>(
     override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<DataRow<T>>? =
         DataColumn.createColumnGroup("", df).addPath(emptyPath())
 
-    override fun resolveSingleAfter(
+    override fun resolveSingleAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): ColumnWithPath<DataRow<T>>? =
         DataColumn.createColumnGroup(
             name = "",
             df = createColumnSet { df.columns().map { it.addPath() } }
-                .let(transform).cast<T>()
+                .let(transformer::transform).cast<T>()
                 .resolve(context)
                 .toDataFrame() as DataFrame<T>,
         ).addPath(emptyPath())

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/DataFrameReceiver.kt

@@ -2,7 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation
 
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.columns.*
-import org.jetbrains.kotlinx.dataframe.impl.columns.transform
 
 internal class ConfiguredAggregateColumn<C> private constructor(
     val columns: ColumnSet<C>,
@@ -32,9 +31,9 @@ internal class ConfiguredAggregateColumn<C> private constructor(
 
     override fun resolveAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): List<ColumnWithPath<C>> =
-        resolve(context, transform(columns).cast<C>())
+        resolve(context, transformer.transform(columns).cast<C>())
 
     companion object {
 

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/ConfiguredAggregateColumn.kt

@@ -8,12 +8,9 @@ import org.jetbrains.kotlinx.dataframe.api.PivotColumnsSelector
 import org.jetbrains.kotlinx.dataframe.api.forEach
 import org.jetbrains.kotlinx.dataframe.api.groupBy
 import org.jetbrains.kotlinx.dataframe.api.toPath
-import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
-import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
-import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
-import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.columns.*
+import org.jetbrains.kotlinx.dataframe.columns.ColumnSetTransformer
 import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
-import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.GroupByReceiverImpl
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.receivers.AggregateInternalDsl
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.receivers.AggregatePivotDslImpl
@@ -53,7 +50,7 @@ internal class PivotChainColumnSet<C>(val first: ColumnSet<C>, val second: Colum
 
     override fun resolveAfterTransform(
         context: ColumnResolutionContext,
-        transform: (ColumnSet<*>) -> ColumnSet<*>,
+        transformer: ColumnSetTransformer,
     ): List<ColumnWithPath<C>> {
         throw UnsupportedOperationException("Not implemented")
     }