fixed implementations where allExcept was used to select nested columns

Kotlin · Jan 28, 2025 · 21c6a37 · 21c6a37
1 parent b576fc4
commit 21c6a37
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 17 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/groupBy.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/groupBy.kt
@@ -61,7 +61,9 @@ public fun <T> Pivot<T>.groupBy(vararg columns: KProperty<*>): PivotGroupBy<T> =
 public fun <T> Pivot<T>.groupByOther(): PivotGroupBy<T> {
     val impl = this as PivotImpl<T>
     val pivotColumns = df.getPivotColumnPaths(columns).toColumnSet()
-    return impl.toGroupedPivot(moveToTop = false) { allExcept(pivotColumns) }
+    return impl.toGroupedPivot(moveToTop = false) {
+        (this as DataFrame<T>).remove { pivotColumns }.columns().toColumnSet()
+    }
 }
 
 // endregion

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/PivotGroupByImpl.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/PivotGroupByImpl.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation
 
+import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateBody
@@ -9,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
 import org.jetbrains.kotlinx.dataframe.api.aggregate
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.firstOrNull
+import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.GroupByImpl
 import org.jetbrains.kotlinx.dataframe.impl.api.aggregatePivot
@@ -33,10 +35,11 @@ internal data class PivotGroupByImpl<T>(
         df.groups.firstOrNull()
             ?.getPivotColumnPaths(columns).orEmpty()
             .let { pivotPaths ->
-                {
-                    all().except(
-                        pivotPaths.toColumnSet() and (df as GroupByImpl).keyColumnsInGroups.toColumnSet(),
-                    )
+                return@let {
+                    (this as AnyFrame)
+                        .remove { pivotPaths.toColumnSet() and (df as GroupByImpl).keyColumnsInGroups.toColumnSet() }
+                        .columns()
+                        .toColumnSet()
                 }
             }
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/implode.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/implode.kt
@@ -8,16 +8,20 @@ import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.dropNA
 import org.jetbrains.kotlinx.dataframe.api.groupBy
 import org.jetbrains.kotlinx.dataframe.api.map
+import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.api.replace
 import org.jetbrains.kotlinx.dataframe.api.with
 import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.columns.asAnyFrameColumn
 import org.jetbrains.kotlinx.dataframe.impl.columns.extractDataFrame
 import org.jetbrains.kotlinx.dataframe.impl.getListType
 import kotlin.reflect.typeOf
 
 internal fun <T, C> DataFrame<T>.implodeImpl(dropNA: Boolean = false, columns: ColumnsSelector<T, C>): DataFrame<T> =
-    groupBy { allExcept(columns) }.updateGroups {
+    groupBy {
+        (this as DataFrame<T>).remove(columns).columns().toColumnSet()
+    }.updateGroups {
         replace(columns).with { column ->
             val (value, type) = when (column.kind()) {
                 ColumnKind.Value -> (if (dropNA) column.dropNA() else column).toList() to getListType(column.type())

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
@@ -2482,17 +2482,6 @@ class DataFrameTests : BaseTest() {
     @Test
     fun `except in columns selector`() {
         typed.select { allExcept { age and weight } } shouldBe typed.select { name and city }
-
-        typed
-            .group { age and weight and city }.into("info")
-            .alsoDebug()
-            .select { allExcept { "info"["age"] } }
-            .alsoDebug()
-            .let {
-                it.name shouldBe typed.name
-                it["info"]["weight"] shouldBe typed.weight
-                it["info"]["city"] shouldBe typed.city
-            }
     }
 
     @Test