Rename colsRange into select in ColumnsSelectionDsl. Fix bug with nested select in replace{}.with {}. Remove host and parent from ColumnWithPath. Add parentName into ColumnWithPath and ColumnPath`.

Author: nikitinas

src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataColumn.kt

@@ -15,7 +15,6 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.jetbrains.kotlinx.dataframe.columns.ValueColumn
-import org.jetbrains.kotlinx.dataframe.columns.size
 import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnGroupImpl
 import org.jetbrains.kotlinx.dataframe.impl.columns.FrameColumnImpl
 import org.jetbrains.kotlinx.dataframe.impl.columns.ValueColumnImpl
@@ -115,7 +114,7 @@ public interface DataColumn<out T> : BaseColumn<T> {
 
     override fun rename(newName: String): DataColumn<T>
 
-    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<T>? = this.addPath(context.df)
+    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<T>? = this.addPath()
 
     override operator fun getValue(thisRef: Any?, property: KProperty<*>): DataColumn<T> = super.getValue(thisRef, property) as DataColumn<T>
 

src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl.kt

@@ -10,17 +10,18 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
+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.FrameColumn
 import org.jetbrains.kotlinx.dataframe.columns.SingleColumn
-import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
 import org.jetbrains.kotlinx.dataframe.columns.renamedReference
-import org.jetbrains.kotlinx.dataframe.impl.DataFrameReceiver
 import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnsList
 import org.jetbrains.kotlinx.dataframe.impl.columns.DistinctColumnSet
+import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.allColumnsExcept
+import org.jetbrains.kotlinx.dataframe.impl.columns.changePath
 import org.jetbrains.kotlinx.dataframe.impl.columns.createColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.columns.getAt
 import org.jetbrains.kotlinx.dataframe.impl.columns.getChildrenAt
@@ -30,6 +31,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.top
 import org.jetbrains.kotlinx.dataframe.impl.columns.transform
 import org.jetbrains.kotlinx.dataframe.impl.columns.transformSingle
 import org.jetbrains.kotlinx.dataframe.impl.columns.tree.dfs
+import org.jetbrains.kotlinx.dataframe.impl.getColumnsWithPaths
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
 import kotlin.reflect.typeOf
@@ -63,13 +65,23 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
 
     public fun ColumnsContainer<*>.group(name: String): ColumnGroupReference = name.toColumnOf()
 
-    public operator fun String.rangeTo(endInclusive: String): ColumnSet<*> = colsRange(this, endInclusive)
-
-    public operator fun Column.rangeTo(endInclusive: Column): ColumnSet<*> {
-        val parent = path().parent()
-        require(parent != null)
-        require(endInclusive.path().parent() == parent) { "Start and end columns have different parent column paths" }
-        return parent.colsRange(name, endInclusive.name)
+    public operator fun String.rangeTo(endInclusive: String): ColumnSet<*> = toColumnAccessor().rangeTo(endInclusive.toColumnAccessor())
+
+    public operator fun Column.rangeTo(endInclusive: Column): ColumnSet<*> = object : ColumnSet<Any?> {
+        override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
+            val startPath = this@rangeTo.resolveSingle(context)!!.path
+            val endPath = endInclusive.resolveSingle(context)!!.path
+            val parentPath = startPath.parent()!!
+            require(parentPath == endPath.parent()) { "Start and end columns have different parent column paths" }
+            val parentCol = context.df.getColumnGroup(parentPath)
+            val startIndex = parentCol.getColumnIndex(startPath.name)
+            val endIndex = parentCol.getColumnIndex(endPath.name)
+            return (startIndex..endIndex).map {
+                parentCol.getColumn(it).let {
+                    it.addPath(parentPath + it.name)
+                }
+            }
+        }
     }
 
     public fun none(): ColumnSet<*> = ColumnsList<Any?>(emptyList())
@@ -94,24 +106,16 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
     public fun ColumnSet<*>.cols(range: IntRange): ColumnSet<Any?> =
         transform { it.flatMap { it.children().subList(range.start, range.endInclusive + 1) } }
 
-    // region colsRange
+    // region select
 
-    public fun <C, R> ColumnGroup<C>.colsRange(selector: ColumnsSelector<C, R>): ColumnSet<R> {
-        val receiver = object : DataFrameReceiver<C>(this, UnresolvedColumnsPolicy.Fail), ColumnsSelectionDsl<C> { }
-        return selector(receiver, receiver)
-    }
+    public fun <C> ColumnSet<DataRow<C>>.select(vararg columns: String): ColumnSet<*> = select { columns.toColumns() }
+
+    public fun <C, R> ColumnSet<DataRow<C>>.select(vararg columns: KProperty<R>): ColumnSet<R> = select { columns.toColumns() }
 
-    public fun ColumnSet<*>.colsRange(start: String, endInclusive: String): ColumnSet<Any?> {
-        val set = this
-        return set.transform {
-            it.flatMap {
-                val children = it.children()
-                val startIndex = children.indexOfFirst { it.name == start }
-                require(startIndex >= 0) { "Column `$start` not found" }
-                val endIndex = children.indexOfLast { it.name == endInclusive }
-                require(endIndex >= 0) { "Column `$endInclusive` not found" }
-                require(endIndex >= startIndex) { "End column `$endInclusive` is before start column `$start`" }
-                children.subList(startIndex, endIndex + 1)
+    public fun <C, R> ColumnSet<DataRow<C>>.select(selector: ColumnsSelector<C, R>): ColumnSet<R> = createColumnSet {
+        this@select.resolve(it).flatMap { group ->
+            group.asColumnGroup().getColumnsWithPaths(selector).map {
+                it.changePath(group.path + it.path)
             }
         }
     }

src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnPath.kt

@@ -39,6 +39,10 @@ public data class ColumnPath(val path: List<String>) : List<String> by path, Col
 
     override fun name(): String = path.last()
 
+    val columnName: String get() = name()
+
+    val parentName: String? get() = if (path.size > 1) path[path.size - 2] else null
+
     override fun rename(newName: String): ColumnPath = ColumnPath(path.dropLast(1) + newName)
 
     override fun getValue(row: AnyRow): Any? = row.owner[this][row.index()]

src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnReference.kt

@@ -34,7 +34,7 @@ public interface ColumnReference<out C> : SingleColumn<C> {
     public fun getValueOrNull(row: AnyRow): C? = resolveFor(row.df())?.get(row.index())
 
     override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>? {
-        return context.df.getColumn<C>(path(), context.unresolvedColumnsPolicy)?.addPath(path(), context.df)
+        return context.df.getColumn<C>(path(), context.unresolvedColumnsPolicy)?.addPath(path())
     }
 }
 

src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnWithPath.kt

@@ -1,6 +1,5 @@
 package org.jetbrains.kotlinx.dataframe.columns
 
-import org.jetbrains.kotlinx.dataframe.ColumnsContainer
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
@@ -10,15 +9,15 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.depth
 
 public interface ColumnWithPath<out T> : DataColumn<T> {
 
-    public val host: ColumnsContainer<*>
     public val data: DataColumn<T>
     public val path: ColumnPath
-    public val parent: ColumnWithPath<*>?
+    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(), host)
-    public fun getChild(name: String): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(name)?.addParentPath(path, host)
-    public fun getChild(index: Int): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(index)?.addParentPath(path, host)
-    public fun children(): List<ColumnWithPath<Any?>> = if (isColumnGroup()) data.asColumnGroup().columns().map { it.addParentPath(path, host) } else emptyList()
+    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 children(): List<ColumnWithPath<Any?>> = if (isColumnGroup()) data.asColumnGroup().columns().map { it.addParentPath(path) } else emptyList()
 
     override fun path(): ColumnPath = path
 

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/DataFrameReceiver.kt

@@ -17,15 +17,24 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.SingleColumn
 import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
 import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnGroupWithParent
+import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnGroupWithPathImpl
 import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.missing.MissingColumnGroup
 import org.jetbrains.kotlinx.dataframe.impl.columns.missing.MissingDataColumn
 import org.jetbrains.kotlinx.dataframe.nrow
 
+private fun <T> DataFrame<T>.unbox(): DataFrame<T> = when (this) {
+    is ColumnGroupWithParent -> source.unbox()
+    is ColumnGroupWithPathImpl -> column.unbox()
+    else -> this
+}
+
+internal abstract class DataFrameReceiverBase<T>(protected val df: DataFrame<T>) : DataFrameImpl<T>(df.columns(), df.nrow)
+
 internal open class DataFrameReceiver<T>(
-    val source: DataFrame<T>,
+    source: DataFrame<T>,
     private val unresolvedColumnsPolicy: UnresolvedColumnsPolicy
-) : DataFrameImpl<T>(source.columns(), source.nrow), SingleColumn<DataRow<T>> {
+) : DataFrameReceiverBase<T>(source.unbox()), SingleColumn<DataRow<T>> {
 
     private fun <R> DataColumn<R>?.check(path: ColumnPath): DataColumn<R>? =
         when (this) {
@@ -38,8 +47,8 @@ internal open class DataFrameReceiver<T>(
             else -> this
         }
 
-    override fun getColumnOrNull(name: String) = source.getColumnOrNull(name).check(pathOf(name))
-    override fun getColumnOrNull(index: Int) = source.getColumnOrNull(index).check(pathOf(""))
+    override fun getColumnOrNull(name: String) = df.getColumnOrNull(name).check(pathOf(name))
+    override fun getColumnOrNull(index: Int) = df.getColumnOrNull(index).check(pathOf(""))
 
     override fun <R> getColumnOrNull(column: ColumnReference<R>): DataColumn<R>? {
         val context = ColumnResolutionContext(this, unresolvedColumnsPolicy)
@@ -49,11 +58,11 @@ internal open class DataFrameReceiver<T>(
     override fun getColumnOrNull(path: ColumnPath) = super.getColumnOrNull(path).check(path)
     override fun <R> getColumnOrNull(column: ColumnSelector<T, R>) = getColumnsImpl(unresolvedColumnsPolicy, column).singleOrNull()
 
-    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<DataRow<T>>? = DataColumn.createColumnGroup("", source).addPath(emptyPath(), source)
+    override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<DataRow<T>>? = DataColumn.createColumnGroup("", df).addPath(emptyPath())
 
-    override fun columns() = source.columns().map { if (it.isColumnGroup()) ColumnGroupWithParent(null, it.asColumnGroup()) else it }
+    override fun columns() = df.columns().map { if (it.isColumnGroup()) ColumnGroupWithParent(null, it.asColumnGroup()) else it }
 
-    override fun columnNames() = source.columnNames()
+    override fun columnNames() = df.columnNames()
 
-    override fun columnTypes() = source.columnTypes()
+    override fun columnTypes() = df.columnTypes()
 }

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt

@@ -36,7 +36,7 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
     val parentPath = targetPath.dropLast(1)
     val toInsert = removeResult.removedColumns.map {
         val path = parentPath + it.name
-        ColumnToInsert(path, it.toColumnWithPath<C>(df).data, refNode)
+        ColumnToInsert(path, it.toColumnWithPath<C>().data, refNode)
     }
     return removeResult.df.insertImpl(toInsert)
 }
@@ -48,7 +48,7 @@ internal fun <T, C> MoveClause<T, C>.moveImpl(
     val receiver = object : DataFrameReceiver<T>(df, UnresolvedColumnsPolicy.Fail), ColumnsSelectionDsl<T> {}
     val removeResult = df.removeImpl(columns = columns)
     val columnsToInsert = removeResult.removedColumns.map {
-        val col = it.toColumnWithPath<C>(df)
+        val col = it.toColumnWithPath<C>()
         var path = newPathExpression(receiver, col).path()
         if (under) path += col.name()
         ColumnToInsert(path, col.data, it)

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/sort.kt

@@ -93,7 +93,7 @@ internal fun <C> ColumnWithPath<C>.addFlag(flag: SortFlag): ColumnWithPath<C> {
             }
         }
         else -> throw IllegalArgumentException("Can not apply sort flag to column kind ${col.kind}")
-    }.addPath(path, host)
+    }.addPath(path)
 }
 
 internal class ColumnsWithSortFlag<C>(val column: ColumnSet<C>, val flag: SortFlag) : ColumnSet<C> {

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/split.kt

@@ -29,7 +29,7 @@ internal fun <T, C, R> splitImpl(
 
     val toInsert = removeResult.removedColumns.flatMap { node ->
 
-        val column = node.toColumnWithPath<C>(clause.df)
+        val column = node.toColumnWithPath<C>()
         val columnCollectors = mutableListOf<ColumnDataCollector>()
         for (row in 0 until nrow) {
             val value = clause.transform(clause.df[row], column.data[row])

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/ColumnAccessorImpl.kt

@@ -29,7 +29,7 @@ internal class ColumnAccessorImpl<T>(val path: ColumnPath) : ColumnAccessor<T> {
                 df = col.asColumnGroup()
             }
         }
-        return col?.cast<T>()?.addPath(path, context.df)
+        return col?.cast<T>()?.addPath(path)
     }
 
     override fun rename(newName: String) = ColumnAccessorImpl<T>(path.dropLast(1) + newName)