Add convert asColumn operation as compiler plugin friendly variant of replace with

Author: koperagen

core/api/core.api

@@ -1434,10 +1434,13 @@ public final class org/jetbrains/kotlinx/dataframe/api/ConstructorsKt {
 public final class org/jetbrains/kotlinx/dataframe/api/Convert {
 	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)V
 	public final fun cast ()Lorg/jetbrains/kotlinx/dataframe/api/Convert;
+	public final fun getColumns ()Lkotlin/jvm/functions/Function2;
+	public final fun getDf ()Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public fun toString ()Ljava/lang/String;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/ConvertKt {
+	public static final fun asColumn (Lorg/jetbrains/kotlinx/dataframe/api/Convert;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun asFrame (Lorg/jetbrains/kotlinx/dataframe/api/Convert;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun convert (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/api/Convert;
 	public static final fun convert (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/api/Convert;

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/convert.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlinx.dataframe.annotations.Converter
 import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
 import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.annotations.Refine
+import org.jetbrains.kotlinx.dataframe.columns.BaseColumn
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
@@ -97,7 +98,10 @@ public inline fun <T, C, reified R> Convert<T, C?>.notNull(
     }
 
 @HasSchema(schemaArg = 0)
-public class Convert<T, out C>(internal val df: DataFrame<T>, internal val columns: ColumnsSelector<T, C>) {
+public class Convert<T, out C>(
+    @PublishedApi internal val df: DataFrame<T>,
+    @PublishedApi internal val columns: ColumnsSelector<T, C>,
+) {
     public fun <R> cast(): Convert<T, R> = Convert(df, columns as ColumnsSelector<T, R>)
 
     @Refine
@@ -129,6 +133,21 @@ public fun <T, C, R> Convert<T, DataRow<C>>.asFrame(
     body: ColumnsContainer<T>.(ColumnGroup<C>) -> DataFrame<R>,
 ): DataFrame<T> = to { body(this, it.asColumnGroup()).asColumnGroup(it.name()) }
 
+/**
+ * Compiler plugin-friendly variant of [ReplaceClause.with]
+ * [ReplaceClause.with] allows to change both column types and names.
+ * Tracking of column name changes in arbitrary lambda expression is unreliable and generally impossible
+ * to do statically.
+ * This function ensures that all column names remain as is and only their type changes to [R]
+ * Example:
+ * `df.convert { colsOf<String>() }.asColumn { it.asList().parallelStream().map { heavyIO(it) }.toList().toColumn() }`
+ */
+@Refine
+@Interpretable("ConvertAsColumn")
+public inline fun <T, C, R> Convert<T, C>.asColumn(
+    crossinline columnConverter: DataFrame<T>.(DataColumn<C>) -> BaseColumn<R>,
+): DataFrame<T> = df.replace(columns).with { columnConverter(df, it).rename(it.name()) }
+
 @Refine
 @Interpretable("PerRowCol")
 public inline fun <T, C, reified R> Convert<T, C>.perRowCol(

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt

@@ -53,6 +53,10 @@ public fun <T, C> ReplaceClause<T, C>.with(newColumns: List<AnyCol>): DataFrame<
 }
 
 // TODO: Issue #418: breaks if running on ColumnGroup and its child
+
+/**
+ * For an alternative supported in the compiler plugin use [Convert.asColumn]
+ */
 public fun <T, C> ReplaceClause<T, C>.with(transform: ColumnsContainer<T>.(DataColumn<C>) -> AnyBaseCol): DataFrame<T> {
     val removeResult = df.removeImpl(columns = columns)
     val toInsert = removeResult.removedColumns.map {

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.jetbrains.kotlinx.dataframe.api.ParserOptions
 import org.jetbrains.kotlinx.dataframe.api.add
 import org.jetbrains.kotlinx.dataframe.api.after
+import org.jetbrains.kotlinx.dataframe.api.asColumn
 import org.jetbrains.kotlinx.dataframe.api.asFrame
 import org.jetbrains.kotlinx.dataframe.api.asGroupBy
 import org.jetbrains.kotlinx.dataframe.api.at
@@ -87,6 +88,7 @@ import org.jetbrains.kotlinx.dataframe.api.sortWith
 import org.jetbrains.kotlinx.dataframe.api.split
 import org.jetbrains.kotlinx.dataframe.api.sum
 import org.jetbrains.kotlinx.dataframe.api.to
+import org.jetbrains.kotlinx.dataframe.api.toColumn
 import org.jetbrains.kotlinx.dataframe.api.toFloat
 import org.jetbrains.kotlinx.dataframe.api.toStart
 import org.jetbrains.kotlinx.dataframe.api.toMap
@@ -114,6 +116,7 @@ import org.junit.Test
 import java.net.URL
 import java.time.format.DateTimeFormatter
 import java.util.*
+import java.util.stream.Collectors
 
 @Suppress("ktlint:standard:chain-method-continuation", "ktlint:standard:argument-list-wrapping")
 class Modify : TestBase() {
@@ -235,6 +238,16 @@ class Modify : TestBase() {
         // SampleEnd
     }
 
+    @Test
+    @TransformDataFrameExpressions
+    fun convertAsColumn() {
+        // SampleStart
+        df.convert { name }.asColumn { col ->
+            col.toList().parallelStream().map { it.toString() }.collect(Collectors.toList()).toColumn()
+        }
+        // SampleEnd
+    }
+
     @Test
     @TransformDataFrameExpressions
     fun parseAll() {

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlinx.dataframe.api.addId
 import org.jetbrains.kotlinx.dataframe.api.after
 import org.jetbrains.kotlinx.dataframe.api.aggregate
 import org.jetbrains.kotlinx.dataframe.api.append
+import org.jetbrains.kotlinx.dataframe.api.asColumn
 import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.asDataFrame
 import org.jetbrains.kotlinx.dataframe.api.asFrame
@@ -82,7 +83,9 @@ import org.jetbrains.kotlinx.dataframe.api.single
 import org.jetbrains.kotlinx.dataframe.api.sortBy
 import org.jetbrains.kotlinx.dataframe.api.split
 import org.jetbrains.kotlinx.dataframe.api.sumOf
+import org.jetbrains.kotlinx.dataframe.api.toColumn
 import org.jetbrains.kotlinx.dataframe.api.toColumnAccessor
+import org.jetbrains.kotlinx.dataframe.api.toStr
 import org.jetbrains.kotlinx.dataframe.api.toTop
 import org.jetbrains.kotlinx.dataframe.api.under
 import org.jetbrains.kotlinx.dataframe.api.ungroup
@@ -100,7 +103,9 @@ import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.jetbrains.kotlinx.dataframe.columns.depth
 import org.jetbrains.kotlinx.dataframe.hasNulls
 import org.junit.Test
+import java.util.stream.Collectors
 import kotlin.reflect.typeOf
+import kotlin.streams.toList
 
 class DataFrameTreeTests : BaseTest() {
 
@@ -485,6 +490,15 @@ class DataFrameTreeTests : BaseTest() {
                 .sorted()
     }
 
+    @Test
+    fun `convert column expression ignoring name changes`() {
+        val res = df.convert { colsOf<Double?>() }.asColumn {
+            it.toList().parallelStream().map { it.toString() }.collect(Collectors.toList()).toColumn("123")
+        }
+
+        res shouldBe df.convert { colsOf<Double?>() }.toStr()
+    }
+
     @Test
     fun extensionPropertiesTest() {
         val code = CodeGenerator.create()