small additions and fixes

Author: Jolanrensen

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

@@ -60,7 +60,7 @@ public interface ColumnsContainer<out T> {
     public operator fun <R> get(column: KProperty<DataRow<R>>): ColumnGroup<R> = get(column.columnName).asColumnGroup().cast()
     public operator fun <R> get(column: KProperty<DataFrame<R>>): FrameColumn<R> = get(column.columnName).asAnyFrameColumn().castFrameColumn()
 
-    public operator fun <C> get(columns: ColumnsSelector<T, C>): List<DataColumn<C>>
+    public fun <C> get(columns: ColumnsSelector<T, C>): List<DataColumn<C>>
     public operator fun <C> get(column: ColumnSelector<T, C>): DataColumn<C> = get(column as ColumnsSelector<T, C>).single()
 
     // endregion

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

@@ -41,7 +41,9 @@ internal class ColumnGroupWithParent<T>(override val parent: ColumnGroupReferenc
     override fun getColumnOrNull(index: Int) = super.getColumnOrNull(index)?.addParent(this)
     override fun getColumnOrNull(name: String) = super.getColumnOrNull(name)?.addParent(this)
     override fun <R> getColumnOrNull(column: ColumnReference<R>) = super.getColumnOrNull(column)?.addParent(this)
-    override fun getColumnOrNull(path: ColumnPath) = path.fold(this as AnyCol?) { col, name -> col?.asColumnGroup()?.getColumnOrNull(name) }
+    override fun getColumnOrNull(path: ColumnPath) = path.fold(this as AnyCol?) { col, name ->
+        col?.asColumnGroup()?.getColumnOrNull(name)
+    }
     override fun <R> getColumnOrNull(column: ColumnSelector<T, R>) = getColumnOrNull(getColumnWithPath(column).path)?.cast<R>()
 
     override operator fun get(columnName: String): AnyCol = getColumn(columnName)

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

@@ -104,6 +104,11 @@ internal fun <A, B> ColumnSet<A>.transform(converter: (List<ColumnWithPath<A>>)
         override fun resolve(context: ColumnResolutionContext) = converter(this@transform.resolve(context))
     }
 
+internal fun <A, B> ColumnSet<A>.transformWithContext(converter: ColumnResolutionContext.(List<ColumnWithPath<A>>) -> List<ColumnWithPath<B>>): ColumnSet<B> =
+    object : ColumnSet<B> {
+        override fun resolve(context: ColumnResolutionContext) = converter(context, this@transformWithContext.resolve(context))
+    }
+
 internal fun <T> ColumnSet<T>.singleImpl() = object : SingleColumn<T> {
     override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<T>? {
         return this@singleImpl.resolve(context).singleOrNull()

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

@@ -43,6 +43,12 @@ class ColumnsSelectionDslTests : TestBase() {
                 it.any { it == "Alice" }
             }
         } shouldBe df.select { name.firstName }
+
+        df.select {
+            pathOf("name").first {
+                it.any { it == "Alice" }
+            }
+        } shouldBe df.select { name.firstName }
     }
 
     @Test
@@ -72,6 +78,12 @@ class ColumnsSelectionDslTests : TestBase() {
                 it.any { it == "Alice" }
             }
         } shouldBe df.select { name.firstName }
+
+        df.select {
+            pathOf("name").last {
+                it.any { it == "Alice" }
+            }
+        } shouldBe df.select { name.firstName }
     }
 
     @Test
@@ -103,6 +115,12 @@ class ColumnsSelectionDslTests : TestBase() {
                 it.any { it == "Alice" }
             }
         } shouldBe df.select { name.firstName }
+
+        df.select {
+            pathOf("name").single {
+                it.any { it == "Alice" }
+            }
+        } shouldBe df.select { name.firstName }
     }
 
     @Test
@@ -145,8 +163,6 @@ class ColumnsSelectionDslTests : TestBase() {
                 .asColumnGroup(firstNames)
         }
 
-        dfGroup.print(columnTypes = true, title = true)
-
         dfGroup.select { colGroup("name") } shouldBe dfGroup.select { name }
         dfGroup.select { colGroup<String>("name") } shouldBe dfGroup.select { name }
         dfGroup.select { colGroup(pathOf("name")) } shouldBe dfGroup.select { name }
@@ -158,6 +174,12 @@ class ColumnsSelectionDslTests : TestBase() {
         dfGroup.select { colGroup("name").colGroup(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
         dfGroup.select { colGroup("name").colGroup<String>(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
         dfGroup.select { colGroup("name").colGroup(MyName::firstNames) } shouldBe dfGroup.select { name[firstNames] }
+
+        dfGroup.select {
+            "name"["firstNames"]["firstName", "secondName"]
+        } shouldBe dfGroup.select {
+            name[firstNames]["firstName"] and name[firstNames]["secondName"]
+        }
     }
 
     @DataSchema

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

@@ -60,7 +60,7 @@ public interface ColumnsContainer<out T> {
     public operator fun <R> get(column: KProperty<DataRow<R>>): ColumnGroup<R> = get(column.columnName).asColumnGroup().cast()
     public operator fun <R> get(column: KProperty<DataFrame<R>>): FrameColumn<R> = get(column.columnName).asAnyFrameColumn().castFrameColumn()
 
-    public operator fun <C> get(columns: ColumnsSelector<T, C>): List<DataColumn<C>>
+    public fun <C> get(columns: ColumnsSelector<T, C>): List<DataColumn<C>>
     public operator fun <C> get(column: ColumnSelector<T, C>): DataColumn<C> = get(column as ColumnsSelector<T, C>).single()
 
     // endregion