TransformableColumnSet and TransformableSingleColumn are now generalized versions on which transformers like "recursively" can be called.

Author: Jolanrensen

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

@@ -6,10 +6,14 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.impl.columns.*
 
 /**
+ * ## ColumnSet
  * Entity that can be resolved into a list of [columns][DataColumn].
  *
  * Used as a return type of [ColumnsSelector].
  * @param C common type of resolved columns
+ * @see [SingleColumn]
+ * @see [TransformableColumnSet]
+ * @see [TransformableSingleColumn]
  */
 public interface ColumnSet<out C> {
 
@@ -21,36 +25,6 @@ public interface ColumnSet<out C> {
     public fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>>
 }
 
-public interface ColumnSetWithRecursively<out C> : ColumnSet<C> {
-
-    public fun resolveRecursively(
-        context: ColumnResolutionContext,
-        includeGroups: Boolean = true,
-        includeTopLevel: Boolean = true,
-    ): List<ColumnWithPath<C>>
-}
-
-internal fun <C> ColumnSetWithRecursively<C>.recursivelyImpl(
-    includeGroups: Boolean = true,
-    includeTopLevel: Boolean = true,
-): ColumnSet<C> = object : ColumnSet<C> {
-
-    override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
-        this@recursivelyImpl.resolveRecursively(
-            context = context,
-            includeGroups = includeGroups,
-            includeTopLevel = includeTopLevel,
-        )
-}
-
-
-public fun interface ColumnSetTransformer {
-
-    public fun transform(columnSet: ColumnSet<*>): ColumnSet<*>
-}
-
-public operator fun ColumnSetTransformer.invoke(columnSet: ColumnSet<*>): ColumnSet<*> = transform(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

@@ -2,13 +2,17 @@ package org.jetbrains.kotlinx.dataframe.columns
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
+import org.jetbrains.kotlinx.dataframe.impl.columns.*
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
 /**
  * Entity that can be [resolved][resolveSingle] into [DataColumn].
  *
  * @param C Column [type][BaseColumn.type] of resolved column.
+ * @see [ColumnSet]
+ * @see [TransformableColumnSet]
+ * @see [TransformableSingleColumn]
  */
 public interface SingleColumn<out C> : ColumnSet<C> {
 
@@ -19,29 +23,6 @@ public interface SingleColumn<out C> : ColumnSet<C> {
     public fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>?
 }
 
-internal fun <C> SingleColumnWithRecursively<C>.recursivelyImpl(
-    includeGroups: Boolean = true,
-    includeTopLevel: Boolean = true,
-): SingleColumn<C> = object : SingleColumn<C> {
-
-    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>? =
-        this@recursivelyImpl.resolveSingleRecursively(
-            context = context,
-            includeGroups = includeGroups,
-            includeTopLevel = includeTopLevel,
-        )
-}
-
-public interface SingleColumnWithRecursively<out C> : SingleColumn<C> {
-
-    public fun resolveSingleRecursively(
-        context: ColumnResolutionContext,
-        includeGroups: Boolean = true,
-        includeTopLevel: Boolean = true,
-    ): ColumnWithPath<C>?
-}
-
-
 @OptIn(ExperimentalContracts::class)
 public fun ColumnSet<*>.isSingleColumn(): Boolean {
     contract {

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

@@ -9,7 +9,6 @@ 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.*
-import org.jetbrains.kotlinx.dataframe.columns.ColumnSetTransformer
 import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.GroupByReceiverImpl
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.receivers.AggregateInternalDsl

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlinx.dataframe.api.getFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.update
 import org.jetbrains.kotlinx.dataframe.api.with
 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.impl.columns.addPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.assertIsComparable

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

@@ -2,8 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.columns
 
 import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
 import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
-import org.jetbrains.kotlinx.dataframe.columns.ColumnSetTransformer
-import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 
 internal class ColumnsList<C>(val columns: List<ColumnSet<C>>) : ColumnSet<C> {
     constructor(vararg columns: ColumnSet<C>) : this(columns.toList())

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

@@ -2,7 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.columns
 
 import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
 import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
-import org.jetbrains.kotlinx.dataframe.columns.ColumnSetTransformer
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 
 internal class DistinctColumnSet<T>(val src: ColumnSet<T>) : ColumnSet<T> {

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

@@ -0,0 +1,93 @@
+package org.jetbrains.kotlinx.dataframe.impl.columns
+
+import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
+import org.jetbrains.kotlinx.dataframe.columns.*
+import org.jetbrains.kotlinx.dataframe.impl.columns.tree.flattenRecursively
+
+internal fun <C> TransformableColumnSet<C>.recursivelyImpl(
+    includeGroups: Boolean = true,
+    includeTopLevel: Boolean = true,
+): ColumnSet<C> = object : ColumnSet<C> {
+
+    override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
+        this@recursivelyImpl.transformResolve(
+            context = context,
+            transformer = RecursivelyTransformer(
+                includeGroups = includeGroups,
+                includeTopLevel = includeTopLevel,
+            ),
+        )
+}
+
+internal fun <C> TransformableSingleColumn<C>.recursivelyImpl(
+    includeGroups: Boolean = true,
+    includeTopLevel: Boolean = true,
+): SingleColumn<C> = object : SingleColumn<C> {
+
+    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>? =
+        this@recursivelyImpl.transformResolveSingle(
+            context = context,
+            transformer = RecursivelyTransformer(
+                includeGroups = includeGroups,
+                includeTopLevel = includeTopLevel,
+            ),
+        )
+}
+
+private class RecursivelyTransformer(
+    val includeGroups: Boolean = true,
+    val includeTopLevel: Boolean = true,
+) : ColumnSetTransformer {
+
+    override fun transform(columnSet: ColumnSet<*>): ColumnSet<*> =
+        columnSet.flattenRecursively(
+            includeGroups = includeGroups,
+            includeTopLevel = includeTopLevel,
+        )
+
+    override fun transformSingle(singleColumn: SingleColumn<*>): ColumnSet<*> =
+        singleColumn.flattenRecursivelySingle(
+            includeGroups = includeGroups,
+            includeTopLevel = includeTopLevel,
+        )
+}
+
+internal fun ColumnSet<*>.flattenRecursively(
+    includeGroups: Boolean = true,
+    includeTopLevel: Boolean = true,
+): ColumnSet<*> = transform { list ->
+    val cols =
+        if (isSingleColumnGroup(list)) {
+            list.single().children()
+        } else {
+            list
+        }
+
+    if (includeTopLevel) {
+        cols.flattenRecursively()
+    } else {
+        cols
+            .filter { it.isColumnGroup() }
+            .flatMap { it.children().flattenRecursively() }
+    }.filter { includeGroups || !it.isColumnGroup() }
+}
+
+internal fun SingleColumn<*>.flattenRecursivelySingle(
+    includeGroups: Boolean = true,
+    includeTopLevel: Boolean = true,
+): ColumnSet<*> = transformSingle {
+    val cols =
+        if (isSingleColumnGroup(listOf(it))) {
+            it.children()
+        } else {
+            listOf(it)
+        }
+
+    if (includeTopLevel) {
+        cols.flattenRecursively()
+    } else {
+        cols
+            .filter { it.isColumnGroup() }
+            .flatMap { it.children().flattenRecursively() }
+    }.filter { includeGroups || !it.isColumnGroup() }
+}

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

@@ -0,0 +1,53 @@
+package org.jetbrains.kotlinx.dataframe.impl.columns
+
+import org.jetbrains.kotlinx.dataframe.columns.*
+import org.jetbrains.kotlinx.dataframe.api.*
+
+/**
+ * ## Transformable ColumnSet
+ * This type of [ColumnSet] can be [transformed][transformResolve] before being resolved.
+ *
+ * This is especially useful for calls like
+ * [cols { }][ColumnsSelectionDsl.cols].[recursively()][ColumnsSelectionDsl.recursively],
+ * where [recursively][ColumnsSelectionDsl.recursively] modifies the [ColumnSet][ColumnSet]
+ * that [cols { }][ColumnsSelectionDsl.cols] operates on before it's evaluated.
+ *
+ * @see [ColumnSet]
+ * @see [TransformableSingleColumn]
+ * @see [SingleColumn]
+ */
+public interface TransformableColumnSet<out C> : ColumnSet<C> {
+
+    public fun transformResolve(
+        context: ColumnResolutionContext,
+        transformer: ColumnSetTransformer,
+    ): List<ColumnWithPath<C>>
+}
+
+/**
+ * ## Transformable SingleColumn
+ * This type of [SingleColumn] can be [transformed][transformResolveSingle] before being resolved.
+ *
+ * This is especially useful for calls like
+ * [first { }][ColumnsSelectionDsl.first].[recursively()][ColumnsSelectionDsl.recursively],
+ * where [recursively][ColumnsSelectionDsl.recursively] modifies the [ColumnSet][ColumnSet]
+ * that [first { }][ColumnsSelectionDsl.first] operates on before it's evaluated.
+ *
+ * @see [SingleColumn]
+ * @see [TransformableColumnSet]
+ * @see [ColumnSet]
+ */
+public interface TransformableSingleColumn<out C> : SingleColumn<C> {
+
+    public fun transformResolveSingle(
+        context: ColumnResolutionContext,
+        transformer: ColumnSetTransformer,
+    ): ColumnWithPath<C>?
+}
+
+public interface ColumnSetTransformer {
+
+    public fun transform(columnSet: ColumnSet<*>): ColumnSet<*>
+
+    public fun transformSingle(singleColumn: SingleColumn<*>): ColumnSet<*>
+}