added test and small fix to verify that primary constructor arguments…

… keep their order as dataframe columns, even when they have @ColumnName annotations.
Kotlin · Dec 16, 2024 · 5cd3c45 · 5cd3c45
1 parent 59ebff8
commit 5cd3c45
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 9 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/Utils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/Utils.kt
@@ -381,19 +381,37 @@ internal fun KCallable<*>.isGetterLike(): Boolean =
         else -> false
     }
 
+/** @include [KCallable.getterName] */
+internal val KFunction<*>.getterName: String
+    get() = name
+        .removePrefix("get")
+        .removePrefix("is")
+        .replaceFirstChar { it.lowercase() }
+
+/** @include [KCallable.getterName] */
+internal val KProperty<*>.getterName: String
+    get() = name
+
+/**
+ * Returns the getter name for this callable.
+ * The name of the callable is returned with proper getter-trimming if it's a [KFunction].
+ */
+internal val KCallable<*>.getterName: String
+    get() = when (this) {
+        is KFunction<*> -> getterName
+        is KProperty<*> -> getterName
+        else -> name
+    }
+
 /** @include [KCallable.columnName] */
 @PublishedApi
 internal val KFunction<*>.columnName: String
-    get() = findAnnotation<ColumnName>()?.name
-        ?: name
-            .removePrefix("get")
-            .removePrefix("is")
-            .replaceFirstChar { it.lowercase() }
+    get() = findAnnotation<ColumnName>()?.name ?: getterName
 
 /** @include [KCallable.columnName] */
 @PublishedApi
 internal val KProperty<*>.columnName: String
-    get() = findAnnotation<ColumnName>()?.name ?: name
+    get() = findAnnotation<ColumnName>()?.name ?: getterName
 
 /**
  * Returns the column name for this callable.
@@ -405,5 +423,5 @@ internal val KCallable<*>.columnName: String
     get() = when (this) {
         is KFunction<*> -> columnName
         is KProperty<*> -> columnName
-        else -> findAnnotation<ColumnName>()?.name ?: name
+        else -> findAnnotation<ColumnName>()?.name ?: getterName
     }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/schema/Utils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/schema/Utils.kt
@@ -15,6 +15,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ValueColumn
 import org.jetbrains.kotlinx.dataframe.hasNulls
 import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import org.jetbrains.kotlinx.dataframe.impl.getterName
 import org.jetbrains.kotlinx.dataframe.impl.isGetterLike
 import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
 import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
@@ -192,8 +193,8 @@ internal fun <T> Iterable<KCallable<T>>.sortWithConstructor(klass: KClass<*>): L
     // else sort the ones in the primary constructor first according to the order in there
     // leave the rest at the end in lexicographical order
     val (propsInConstructor, propsNotInConstructor) =
-        lexicographicalColumns.partition { it.columnName in primaryConstructorOrder.keys }
+        lexicographicalColumns.partition { it.getterName in primaryConstructorOrder.keys }
 
     return propsInConstructor
-        .sortedBy { primaryConstructorOrder[it.columnName] } + propsNotInConstructor
+        .sortedBy { primaryConstructorOrder[it.getterName] } + propsNotInConstructor
 }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
@@ -5,6 +5,7 @@ import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
 import org.jetbrains.kotlinx.dataframe.kind
@@ -79,14 +80,35 @@ class CreateDataFrameTests {
 
     @Test
     fun `preserve fields order`() {
+        // constructor properties will keep order, so x, c
         class B(val x: Int, val c: String, d: Double) {
+            // then child properties will be sorted lexicographically by column name, so a, b
             val b: Int = x
             val a: Double = d
         }
 
         listOf(B(1, "a", 2.0)).toDataFrame().columnNames() shouldBe listOf("x", "c", "a", "b")
     }
 
+    @Test
+    fun `preserve fields order with @ColumnName`() {
+        // constructor properties will keep order, so z, y
+        class B(
+            @ColumnName("z") val x: Int,
+            @ColumnName("y") val c: String,
+            d: Double,
+        ) {
+            // then child properties will be sorted lexicographically by column name, so w, x
+            @ColumnName("x")
+            val a: Double = d
+
+            @ColumnName("w")
+            val b: Int = x
+        }
+
+        listOf(B(1, "a", 2.0)).toDataFrame().columnNames() shouldBe listOf("z", "y", "w", "x")
+    }
+
     @DataSchema
     data class A(val v: Int)