dfs -> allChildren

topDfs -> topmostChildren
top -> roots

Author: Jolanrensen

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

@@ -4531,36 +4531,75 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
     public fun <C> ColumnSet<C>.take(n: Int): ColumnSet<C> = transform { it.take(n) }
     public fun <C> ColumnSet<C>.dropLast(n: Int = 1): ColumnSet<C> = transform { it.dropLast(n) }
     public fun <C> ColumnSet<C>.takeLast(n: Int): ColumnSet<C> = transform { it.takeLast(n) }
-    public fun <C> ColumnSet<C>.top(): ColumnSet<C> = transform { it.top() }
+
+    @Deprecated("Use roots() instead", ReplaceWith("roots()"))
+    public fun <C> ColumnSet<C>.top(): ColumnSet<C> = roots()
+
+    /**
+     * ## Roots
+     *
+     * Returns a sub-set of columns that are roots of the trees of columns.
+     *
+     * In practice, this means that if a column in [this] is a child of another column in [this],
+     * it will not be included in the result.
+     *
+     * TODO example and SingleColumn case
+     */
+    public fun <C> ColumnSet<C>.roots(): ColumnSet<C> = transform { it.roots() }
+
     public fun <C> ColumnSet<C>.takeWhile(predicate: ColumnFilter<C>): ColumnSet<C> =
         transform { it.takeWhile(predicate) }
 
     public fun <C> ColumnSet<C>.takeLastWhile(predicate: ColumnFilter<C>): ColumnSet<C> =
         transform { it.takeLastWhile(predicate) }
 
-    public fun <C> ColumnSet<C>.filter(predicate: ColumnFilter<C>): ColumnSet<C> =
-        colsInternal(predicate as ColumnFilter<*>) as ColumnSet<C>
+    public fun <C> ColumnSet<C>.filter(predicate: ColumnFilter<C>): TransformableColumnSet<C> =
+        colsInternal(predicate as ColumnFilter<*>) as TransformableColumnSet<C>
 
-    public fun SingleColumn<*>.nameContains(text: CharSequence): ColumnSet<*> = cols { it.name.contains(text) }
-    public fun <C> ColumnSet<C>.nameContains(text: CharSequence): ColumnSet<C> = cols { it.name.contains(text) }
-    public fun SingleColumn<*>.nameContains(regex: Regex): ColumnSet<*> = cols { it.name.contains(regex) }
-    public fun <C> ColumnSet<C>.nameContains(regex: Regex): ColumnSet<C> = cols { it.name.contains(regex) }
-    public fun SingleColumn<*>.startsWith(prefix: CharSequence): ColumnSet<*> = cols { it.name.startsWith(prefix) }
-    public fun <C> ColumnSet<C>.startsWith(prefix: CharSequence): ColumnSet<C> = cols { it.name.startsWith(prefix) }
-    public fun SingleColumn<*>.endsWith(suffix: CharSequence): ColumnSet<*> = cols { it.name.endsWith(suffix) }
-    public fun <C> ColumnSet<C>.endsWith(suffix: CharSequence): ColumnSet<C> = cols { it.name.endsWith(suffix) }
+    public fun SingleColumn<*>.nameContains(text: CharSequence): TransformableColumnSet<*> =
+        cols { it.name.contains(text) }
 
-    public fun <C> ColumnSet<C>.except(vararg other: ColumnSet<*>): ColumnSet<*> = except(other.toColumnSet())
-    public fun <C> ColumnSet<C>.except(vararg other: String): ColumnSet<*> = except(other.toColumnSet())
+    public fun <C> ColumnSet<C>.nameContains(text: CharSequence): TransformableColumnSet<C> =
+        cols { it.name.contains(text) }
 
-    public fun <C> ColumnSet<C?>.withoutNulls(): ColumnSet<C> =
-        transform { it.filter { !it.hasNulls() } } as ColumnSet<C>
+    public fun SingleColumn<*>.nameContains(regex: Regex): TransformableColumnSet<*> = cols { it.name.contains(regex) }
+    public fun <C> ColumnSet<C>.nameContains(regex: Regex): TransformableColumnSet<C> = cols { it.name.contains(regex) }
+    public fun SingleColumn<*>.startsWith(prefix: CharSequence): TransformableColumnSet<*> =
+        cols { it.name.startsWith(prefix) }
+
+    public fun <C> ColumnSet<C>.startsWith(prefix: CharSequence): TransformableColumnSet<C> =
+        cols { it.name.startsWith(prefix) }
+
+    public fun SingleColumn<*>.endsWith(suffix: CharSequence): TransformableColumnSet<*> =
+        cols { it.name.endsWith(suffix) }
+
+    public fun <C> ColumnSet<C>.endsWith(suffix: CharSequence): TransformableColumnSet<C> =
+        cols { it.name.endsWith(suffix) }
+
+    public fun <C> ColumnSet<C>.except(vararg other: ColumnSet<*>): TransformableColumnSet<*> = except(other.toColumnSet())
+    public fun <C> ColumnSet<C>.except(vararg other: String): TransformableColumnSet<*> = except(other.toColumnSet())
+
+    public fun <C> ColumnSet<C?>.withoutNulls(): TransformableColumnSet<C> =
+        transform { it.filter { !it.hasNulls() } } as TransformableColumnSet<C>
+
+    public infix fun <C> ColumnSet<C>.except(other: ColumnSet<*>): TransformableColumnSet<*> =
+//        createColumnSet {
+//            [email protected](it).allColumnsExcept(other.resolve(it))
+//        }
+        createTransformableColumnSet(
+            resolver = { context ->
+                this@except.resolve(context).allColumnsExcept(other.resolve(context))
+            },
+            transformResolve = { context, transformer ->
+                transformer.transform(this@except)
+                    .resolve(context)
+                    .allColumnsExcept(other.resolve(context))
+            },
+        )
 
-    public infix fun <C> ColumnSet<C>.except(other: ColumnSet<*>): ColumnSet<*> =
-        createColumnSet { resolve(it).allColumnsExcept(other.resolve(it)) }
 
-    public infix fun <C> ColumnSet<C>.except(selector: ColumnsSelector<T, *>): ColumnSet<C> =
-        except(selector.toColumns()) as ColumnSet<C>
+    public infix fun <C> ColumnSet<C>.except(selector: ColumnsSelector<T, *>): TransformableColumnSet<C> =
+        except(selector.toColumns()) as TransformableColumnSet<C>
 
     public operator fun <C> ColumnsSelector<T, C>.invoke(): ColumnSet<C> =
         this(this@ColumnsSelectionDsl, this@ColumnsSelectionDsl)
@@ -4971,7 +5010,10 @@ private interface CommonColsOfDocs {
  * @param [filter] an optional filter function that takes a column of type [C] and returns `true` if the column should be included.
  * @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns of given type that were included by [filter].
  */
-public fun <C> ColumnSet<*>.colsOf(type: KType, filter: (DataColumn<C>) -> Boolean = { true }): TransformableColumnSet<C> =
+public fun <C> ColumnSet<*>.colsOf(
+    type: KType,
+    filter: (DataColumn<C>) -> Boolean = { true },
+): TransformableColumnSet<C> =
     colsInternal { it.isSubtypeOf(type) && filter(it.cast()) } as TransformableColumnSet<C>
 
 /**
@@ -5053,7 +5095,10 @@ public inline fun <reified C> ColumnSet<*>.colsOf(noinline filter: (DataColumn<C
  * @param [filter] an optional filter function that takes a column of type [C] and returns `true` if the column should be included.
  * @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns of given type that were included by [filter].
  */
-public fun <C> SingleColumn<*>.colsOf(type: KType, filter: (DataColumn<C>) -> Boolean = { true }): TransformableColumnSet<C> =
+public fun <C> SingleColumn<*>.colsOf(
+    type: KType,
+    filter: (DataColumn<C>) -> Boolean = { true },
+): TransformableColumnSet<C> =
     colsInternal { it.isSubtypeOf(type) && filter(it.cast()) } as TransformableColumnSet<C>
 
 /**

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

@@ -15,7 +15,9 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnSet
 internal fun <T, C> DataFrame<T>.flattenImpl(
     columns: ColumnsSelector<T, C>,
 ): DataFrame<T> {
-    val rootColumns = getColumnsWithPaths { columns.toColumnSet().filter { it.isColumnGroup() }.top() }
+    val rootColumns = getColumnsWithPaths {
+        columns.toColumnSet().filter { it.isColumnGroup() }.roots()
+    }
     val rootPrefixes = rootColumns.map { it.path }.toSet()
     val nameGenerators = rootPrefixes.map { it.dropLast() }.distinct().associateWith { path ->
         val usedNames = get(path).asColumnGroup().columns().filter { path + it.name() !in rootPrefixes }.map { it.name() }

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

@@ -167,25 +167,31 @@ internal fun <C> ColumnsContainer<*>.getColumn(path: ColumnPath, policy: Unresol
             UnresolvedColumnsPolicy.Create -> DataColumn.empty().cast<C>()
         }
 
-internal fun <T> List<ColumnWithPath<T>>.top(): List<ColumnWithPath<T>> {
-    val root = TreeNode.createRoot<ColumnWithPath<T>?>(null)
-    forEach { root.put(it.path, it) }
-    return root.topDfs { it.data != null }.map { it.data!! }
+/**
+ * Returns a sub-list of columns that are roots of the trees of columns.
+ *
+ * In practice, this means that if a column in [this] is a child of another column in [this],
+ * it will not be included in the result.
+ */
+internal fun <T> List<ColumnWithPath<T>>.roots(): List<ColumnWithPath<T>> {
+    val emptyRoot = TreeNode.createRoot<ColumnWithPath<T>?>(data = null)
+    this.forEach { emptyRoot.put(it.path, it) }
+    return emptyRoot.topmostChildren { it.data != null }.map { it.data!! }
 }
 
 internal fun List<ColumnWithPath<*>>.allColumnsExcept(columns: Iterable<ColumnWithPath<*>>): List<ColumnWithPath<*>> {
     if (isEmpty()) return emptyList()
     val fullTree = collectTree()
     columns.forEach {
         var node = fullTree.getOrPut(it.path).asNullable()
-        node?.dfs()?.forEach { it.data = null }
+        node?.allChildren()?.forEach { it.data = null }
         while (node != null) {
             node.data = null
             node = node.parent
         }
     }
-    val dfs = fullTree.topDfs { it.data != null }
-    return dfs.map { it.data!!.addPath(it.pathFromRoot()) }
+    val subtrees = fullTree.topmostChildren { it.data != null }
+    return subtrees.map { it.data!!.addPath(it.pathFromRoot()) }
 }
 
 internal fun KType.toColumnKind(): ColumnKind = jvmErasure.let {

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

@@ -144,10 +144,24 @@ internal fun <T> createColumn(values: Iterable<T>, suggestedType: KType, guessTy
 
 // region create Columns
 
-internal fun <C> createColumnSet(resolver: (ColumnResolutionContext) -> List<ColumnWithPath<C>>): ColumnSet<C> =
-    object : ColumnSet<C> {
-        override fun resolve(context: ColumnResolutionContext) = resolver(context)
-    }
+internal fun <C> createColumnSet(
+    resolver: (context: ColumnResolutionContext) -> List<ColumnWithPath<C>>,
+): ColumnSet<C> = object : ColumnSet<C> {
+    override fun resolve(context: ColumnResolutionContext) = resolver(context)
+}
+
+internal fun <C> createTransformableColumnSet(
+    resolver: (context: ColumnResolutionContext) -> List<ColumnWithPath<C>>,
+    transformResolve: (context: ColumnResolutionContext, transformer: ColumnSetTransformer) -> List<ColumnWithPath<C>>,
+): TransformableColumnSet<C> = object : TransformableColumnSet<C> {
+    override fun resolve(context: ColumnResolutionContext) = resolver(context)
+
+    override fun transformResolve(
+        context: ColumnResolutionContext,
+        transformer: ColumnSetTransformer,
+    ): List<ColumnWithPath<C>> = transformResolve(context, transformer)
+}
+
 
 // region toColumnSet
 

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

@@ -6,7 +6,7 @@ internal class TreeNode<T>(
     override val name: String,
     override val depth: Int,
     override var data: T,
-    override val parent: TreeNode<T>? = null
+    override val parent: TreeNode<T>? = null,
 ) : ReadonlyTreeNode<T> {
 
     companion object {
@@ -28,7 +28,7 @@ internal class TreeNode<T>(
     fun pathFromRoot(): ColumnPath {
         val path = mutableListOf<String>()
         var node: TreeNode<T>? = this
-        while (node != null && node.parent != null) {
+        while (node?.parent != null) {
             path.add(node.name)
             node = node.parent
         }
@@ -48,19 +48,34 @@ internal class TreeNode<T>(
         return addChild(childName, createData())
     }
 
-    fun dfs(enterCondition: (TreeNode<T>) -> Boolean = { true }, yieldCondition: (TreeNode<T>) -> Boolean = { true }): List<TreeNode<T>> {
+    @Deprecated("Use allChildren instead", ReplaceWith("allChildren(enterCondition, yieldCondition)"))
+    fun dfs(
+        enterCondition: (TreeNode<T>) -> Boolean = { true },
+        yieldCondition: (TreeNode<T>) -> Boolean = { true },
+    ): List<TreeNode<T>> = allChildren(enterCondition, yieldCondition)
+
+    /**
+     * Traverses the tree in depth-first order and returns all nodes that satisfy [yieldCondition].
+     * If [enterCondition] returns false for a node, its children are not traversed.
+     * By default, all nodes are traversed and all nodes are returned.
+     */
+    fun allChildren(
+        enterCondition: (TreeNode<T>) -> Boolean = { true },
+        yieldCondition: (TreeNode<T>) -> Boolean = { true },
+    ): List<TreeNode<T>> {
         val result = mutableListOf<TreeNode<T>>()
-        fun doDfs(node: TreeNode<T>) {
+
+        fun traverse(node: TreeNode<T>) {
             if (yieldCondition(node)) {
                 result.add(node)
             }
             if (enterCondition(node)) {
                 node.children.forEach {
-                    doDfs(it)
+                    traverse(it)
                 }
             }
         }
-        doDfs(this)
+        traverse(this)
         return result
     }
 }

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

@@ -29,9 +29,22 @@ internal fun <T> TreeNode<T>.getOrPut(path: ColumnPath, createData: (ColumnPath)
     return node
 }
 
-internal fun <T> TreeNode<T>.topDfs(yieldCondition: (TreeNode<T>) -> Boolean): List<TreeNode<T>> = dfs(enterCondition = { !yieldCondition(it) }, yieldCondition = yieldCondition)
+/**
+ * Traverses all children in the tree in depth-first order and returns the top-most nodes that satisfy
+ * [yieldCondition]. This means that if a node satisfies [yieldCondition], its children are not traversed, regardless of
+ * whether they satisfy [yieldCondition] or not.
+ */
+internal fun <T> TreeNode<T>.topmostChildren(yieldCondition: (TreeNode<T>) -> Boolean): List<TreeNode<T>> =
+    allChildren(
+        enterCondition = { !yieldCondition(it) },
+        yieldCondition = yieldCondition,
+    )
 
-internal fun <T> TreeNode<T>.topDfsExcluding(excludeRoot: TreeNode<*>): List<TreeNode<T>> {
+@Deprecated("Use topmostChildren instead", ReplaceWith("topmostChildren(yieldCondition)"))
+internal fun <T> TreeNode<T>.topDfs(yieldCondition: (TreeNode<T>) -> Boolean): List<TreeNode<T>> =
+    topmostChildren(yieldCondition)
+
+internal fun <T> TreeNode<T>.topmostChildrenExcluding(excludeRoot: TreeNode<*>): List<TreeNode<T>> {
     val result = mutableListOf<TreeNode<T>>()
     fun doDfs(node: TreeNode<T>, exclude: TreeNode<*>) {
         if (exclude.children.isNotEmpty()) {
@@ -48,19 +61,31 @@ internal fun <T> TreeNode<T>.topDfsExcluding(excludeRoot: TreeNode<*>): List<Tre
     return result
 }
 
-internal fun <T> TreeNode<T?>.dfsNotNull() = dfs { it.data != null }.map { it as TreeNode<T> }
-internal fun <T> TreeNode<T?>.dfsTopNotNull() = dfs(enterCondition = { it.data == null }, yieldCondition = { it.data != null }).map { it as TreeNode<T> }
+internal fun <T> TreeNode<T?>.allChildrenNotNull(): List<TreeNode<T>> =
+    allChildren { it.data != null } as List<TreeNode<T>>
+
+internal fun <T> TreeNode<T?>.topmostChildrenNotNull() =
+    topmostChildren { it.data != null } as List<TreeNode<T>>
+
+internal fun TreeNode<ColumnPosition>.allRemovedColumns() =
+    allChildren { it.data.wasRemoved && it.data.column != null }
+
+internal fun TreeNode<ColumnPosition>.allWithColumns() =
+    allChildren { it.data.column != null }
 
-internal fun TreeNode<ColumnPosition>.allRemovedColumns() = dfs { it.data.wasRemoved && it.data.column != null }
-internal fun TreeNode<ColumnPosition>.allWithColumns() = dfs { it.data.column != null }
 internal fun Iterable<ColumnWithPath<*>>.flattenRecursively(): List<ColumnWithPath<*>> {
     val result = mutableListOf<ColumnWithPath<*>>()
+
     fun flattenRecursively(cols: Iterable<ColumnWithPath<*>>) {
         cols.forEach {
             result.add(it)
             val path = it.path
             if (it.data.isColumnGroup()) {
-                flattenRecursively(it.data.asColumnGroup().columns().map { it.addPath(path + it.name()) })
+                flattenRecursively(
+                    it.data.asColumnGroup()
+                        .columns()
+                        .map { it.addPath(path + it.name()) }
+                )
             }
         }
     }