Add infer argument to map, convert and column API.

Author: nikitinas

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

@@ -36,8 +36,17 @@ public fun <T> ColumnGroupReference.column(): ColumnDelegate<T> = ColumnDelegate
 public fun <T> ColumnGroupReference.column(name: String): ColumnAccessor<T> = ColumnAccessorImpl(path() + name)
 public fun <T> ColumnGroupReference.column(path: ColumnPath): ColumnAccessor<T> = ColumnAccessorImpl(this.path() + path)
 
-public inline fun <reified T> column(name: String = "", noinline expression: RowExpression<Any?, T>): ColumnReference<T> = createComputedColumnReference(name, typeOf<T>(), expression)
-public inline fun <T, reified C> column(df: DataFrame<T>, name: String = "", noinline expression: RowExpression<T, C>): ColumnReference<C> = createComputedColumnReference(name, typeOf<C>(), expression as RowExpression<Any?, C>)
+public inline fun <reified T> column(
+    name: String = "",
+    infer: Infer = Infer.Nulls,
+    noinline expression: RowExpression<Any?, T>
+): ColumnReference<T> = createComputedColumnReference(name, typeOf<T>(), infer, expression)
+public inline fun <T, reified C> column(
+    df: DataFrame<T>,
+    name: String = "",
+    infer: Infer = Infer.Nulls,
+    noinline expression: RowExpression<T, C>
+): ColumnReference<C> = createComputedColumnReference(name, typeOf<C>(), infer, expression as RowExpression<Any?, C>)
 
 public fun columnGroup(): ColumnDelegate<AnyRow> = column()
 public fun columnGroup(name: String): ColumnAccessor<AnyRow> = column(name)

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

@@ -47,23 +47,26 @@ public fun <T, C> DataFrame<T>.convert(vararg columns: ColumnReference<C>): Conv
 public inline fun <T, C, reified R> DataFrame<T>.convert(
     firstCol: ColumnReference<C>,
     vararg cols: ColumnReference<C>,
+    infer: Infer = Infer.Nulls,
     noinline expression: RowValueExpression<T, C, R>
 ): DataFrame<T> =
-    convert(*headPlusArray(firstCol, cols)).with(inferType = false, expression)
+    convert(*headPlusArray(firstCol, cols)).with(infer, expression)
 
 public inline fun <T, C, reified R> DataFrame<T>.convert(
     firstCol: KProperty<C>,
     vararg cols: KProperty<C>,
+    infer: Infer = Infer.Nulls,
     noinline expression: RowValueExpression<T, C, R>
 ): DataFrame<T> =
-    convert(*headPlusArray(firstCol, cols)).with(inferType = false, expression)
+    convert(*headPlusArray(firstCol, cols)).with(infer, expression)
 
 public inline fun <T, reified R> DataFrame<T>.convert(
     firstCol: String,
     vararg cols: String,
+    infer: Infer = Infer.Nulls,
     noinline expression: RowValueExpression<T, Any?, R>
 ): DataFrame<T> =
-    convert(*headPlusArray(firstCol, cols)).with(inferType = false, expression)
+    convert(*headPlusArray(firstCol, cols)).with(infer, expression)
 
 public inline fun <T, C, reified R> Convert<T, C?>.notNull(crossinline expression: RowValueExpression<T, C, R>): DataFrame<T> =
     with {
@@ -80,17 +83,16 @@ public data class Convert<T, C>(val df: DataFrame<T>, val columns: ColumnsSelect
 public fun <T> Convert<T, *>.to(type: KType): DataFrame<T> = to { it.convertTo(type) }
 
 public inline fun <T, C, reified R> Convert<T, C>.with(
-    // TODO: replace with `Infer`
-    inferType: Boolean = false,
+    infer: Infer = Infer.Nulls,
     noinline rowConverter: RowValueExpression<T, C, R>
 ): DataFrame<T> =
-    withRowCellImpl(if (inferType) null else typeOf<R>(), rowConverter)
+    withRowCellImpl(typeOf<R>(), infer, rowConverter)
 
 public inline fun <T, C, reified R> Convert<T, C>.perRowCol(
-    inferType: Boolean = false,
+    infer: Infer = Infer.Nulls,
     noinline expression: RowColumnExpression<T, C, R>
 ): DataFrame<T> =
-    convertRowColumnImpl(if (inferType) null else typeOf<R>(), expression)
+    convertRowColumnImpl(typeOf<R>(), infer, expression)
 
 public fun <T, C> Convert<T, C>.to(columnConverter: DataFrame<T>.(DataColumn<C>) -> AnyCol): DataFrame<T> =
     df.replace(columns).with { columnConverter(df, it) }

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

@@ -9,30 +9,24 @@ import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
-import org.jetbrains.kotlinx.dataframe.columns.size
-import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.columns.createComputedColumnReference
 import org.jetbrains.kotlinx.dataframe.impl.columns.newColumn
-import org.jetbrains.kotlinx.dataframe.impl.columns.newColumnWithActualType
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
 import kotlin.reflect.typeOf
 
 // region ColumnReference
 
-public inline fun <C, reified R> ColumnReference<C>.map(noinline transform: (C) -> R): ColumnReference<R> =
-    map(typeOf<R>(), transform)
-
-public fun <C, R> ColumnReference<C>.map(tartypeOf: KType?, transform: (C) -> R): ColumnReference<R> =
-    createComputedColumnReference(this.name, tartypeOf) { transform(this@map()) }
+public inline fun <C, reified R> ColumnReference<C>.map(infer: Infer = Infer.Nulls, noinline transform: (C) -> R): ColumnReference<R> =
+    createComputedColumnReference(name(), typeOf<R>(), infer) { transform(this@map()) }
 
 // endregion
 
 // region DataColumn
 
 public inline fun <T, reified R> DataColumn<T>.map(
-    infer: Infer = if (typeOf<R>().isMarkedNullable) Infer.Nulls else Infer.None,
+    infer: Infer = Infer.Nulls,
     crossinline transform: (T) -> R
 ): DataColumn<R> {
     val newValues = Array(size()) { transform(get(it)) }.asList()
@@ -41,7 +35,7 @@ public inline fun <T, reified R> DataColumn<T>.map(
 
 public fun <T, R> DataColumn<T>.mapTo(
     type: KType,
-    infer: Infer = if (type.isMarkedNullable) Infer.Nulls else Infer.None,
+    infer: Infer = Infer.Nulls,
     transform: (T) -> R
 ): DataColumn<R> {
     val values = Array<Any?>(size()) { transform(get(it)) }.asList()
@@ -62,11 +56,7 @@ public inline fun <T, reified R> ColumnsContainer<T>.map(
     name: String,
     infer: Infer = Infer.Nulls,
     noinline body: AddExpression<T, R>
-): DataColumn<R> = when (infer) {
-    Infer.Type -> newColumnWithActualType(name, body)
-    Infer.Nulls -> newColumn(typeOf<R>(), name, true, body)
-    Infer.None -> newColumn(typeOf<R>(), name, false, body)
-}
+): DataColumn<R> = newColumn(typeOf<R>(), name, infer, body)
 
 public inline fun <T, reified R> ColumnsContainer<T>.map(
     column: ColumnAccessor<R>,

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

@@ -10,19 +10,22 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.impl.api.removeImpl
 import org.jetbrains.kotlinx.dataframe.impl.api.withRowCellImpl
 import org.jetbrains.kotlinx.dataframe.impl.columns.toColumns
+import org.jetbrains.kotlinx.dataframe.impl.getListType
 import org.jetbrains.kotlinx.dataframe.impl.nameGenerator
 import kotlin.reflect.KProperty
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
 
 public fun <T, C> DataFrame<T>.merge(selector: ColumnsSelector<T, C>): Merge<T, C, List<C>> =
-    Merge(this, selector, false, { it })
+    Merge(this, selector, false, { it }, typeOf<Any?>(), Infer.Type)
 
 public fun <T> DataFrame<T>.merge(vararg columns: String): Merge<T, Any?, List<Any?>> =
     merge { columns.toColumns() }
 
-public fun <T, C> DataFrame<T>.merge(vararg columns: ColumnReference<C>): Merge<T, C, List<C>> =
+public inline fun <T, reified C> DataFrame<T>.merge(vararg columns: ColumnReference<C>): Merge<T, C, List<C>> =
     merge { columns.toColumns() }
 
-public fun <T, C> DataFrame<T>.merge(vararg columns: KProperty<C>): Merge<T, C, List<C>> =
+public inline fun <T, reified C> DataFrame<T>.merge(vararg columns: KProperty<C>): Merge<T, C, List<C>> =
     merge { columns.toColumns() }
 
 public data class Merge<T, C, R>(
@@ -34,12 +37,16 @@ public data class Merge<T, C, R>(
     internal val notNull: Boolean,
     @PublishedApi
     internal val transform: DataRow<T>.(List<C>) -> R,
+    @PublishedApi
+    internal val resultType: KType,
+    @PublishedApi
+    internal val infer: Infer,
 )
 
 public fun <T, C, R> Merge<T, C, R>.notNull(): Merge<T, C, R> = copy(notNull = true)
 
 public fun <T, C, R> Merge<T, C, R>.into(columnName: String): DataFrame<T> = into(pathOf(columnName))
-public fun <T, C, R> Merge<T, C, R>.into(column: ColumnAccessor<R>): DataFrame<T> = into(column.path())
+public fun <T, C, R> Merge<T, C, R>.into(column: ColumnAccessor<*>): DataFrame<T> = into(column.path())
 
 public fun <T, C, R> Merge<T, C, R>.intoList(): List<R> =
     df.select(selector).rows().map { transform(it, it.values() as List<C>) }
@@ -51,7 +58,7 @@ public fun <T, C, R> Merge<T, C, R>.into(path: ColumnPath): DataFrame<T> {
     // move columns into group
     val grouped = df.move(selector).under { mergePath }
 
-    var res = grouped.convert { getColumnGroup(mergePath) }.withRowCellImpl(null) {
+    var res = grouped.convert { getColumnGroup(mergePath) }.withRowCellImpl(resultType, infer) {
         val srcRow = df[index()]
         var values = it.values() as List<C>
         if (notNull) {
@@ -77,15 +84,22 @@ public fun <T, C, R> Merge<T, C, R>.by(
     limit: Int = -1,
     truncated: CharSequence = "..."
 ): Merge<T, C, String> =
-    Merge(df, selector, notNull) {
-        it.joinToString(
-            separator = separator,
-            prefix = prefix,
-            postfix = postfix,
-            limit = limit,
-            truncated = truncated
-        )
-    }
+    Merge(
+        df, selector, notNull,
+        transform = {
+            it.joinToString(
+                separator = separator,
+                prefix = prefix,
+                postfix = postfix,
+                limit = limit,
+                truncated = truncated
+            )
+        },
+        typeOf<String>(), Infer.Nulls
+    )
 
-public inline fun <T, C, R, reified V> Merge<T, C, R>.by(crossinline transform: DataRow<T>.(R) -> V): Merge<T, C, V> =
-    Merge(df, selector, notNull) { transform(this@by.transform(this, it)) }
+public inline fun <T, C, R, reified V> Merge<T, C, R>.by(
+    infer: Infer = Infer.Nulls,
+    crossinline transform: DataRow<T>.(R) -> V
+): Merge<T, C, V> =
+    Merge(df, selector, notNull, { transform(this@by.transform(this, it)) }, typeOf<V>(), infer)

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

@@ -16,7 +16,7 @@ import org.jetbrains.kotlinx.dataframe.impl.asList
 import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.columns.toColumns
-import org.jetbrains.kotlinx.dataframe.impl.createTypeWithArgument
+import org.jetbrains.kotlinx.dataframe.impl.getListType
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
 import kotlin.reflect.typeOf
@@ -274,8 +274,8 @@ public inline fun <T, C : Iterable<R>, reified R> Split<T, C>.intoRows(dropEmpty
 public fun <T, C : AnyFrame> Split<T, C>.intoRows(dropEmpty: Boolean = true): DataFrame<T> =
     by { it.rows() }.intoRows(dropEmpty)
 
-internal fun <T, C, R> Convert<T, C?>.splitInplace(tartypeOf: KType, transform: DataRow<T>.(C) -> Iterable<R>) =
-    withRowCellImpl(List::class.createTypeWithArgument(tartypeOf)) { if (it == null) emptyList() else transform(it).asList() }
+internal fun <T, C, R> Convert<T, C?>.splitInplace(type: KType, transform: DataRow<T>.(C) -> Iterable<R>) =
+    withRowCellImpl(getListType(type), Infer.None) { if (it == null) emptyList() else transform(it).asList() }
 
 public fun <T, C, R> SplitWithTransform<T, C, R>.intoRows(dropEmpty: Boolean = true): DataFrame<T> {
     val paths = df.getColumnPaths(columns).toColumnSet()

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

@@ -4,4 +4,5 @@ import kotlin.reflect.KType
 
 private val listTypes = mutableMapOf<KType, KType>()
 
+@PublishedApi
 internal fun getListType(valueType: KType): KType = listTypes.getOrPut(valueType) { List::class.createTypeWithArgument(valueType) }

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.RowColumnExpression
 import org.jetbrains.kotlinx.dataframe.RowValueExpression
 import org.jetbrains.kotlinx.dataframe.api.Convert
+import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.api.ParserOptions
 import org.jetbrains.kotlinx.dataframe.api.name
 import org.jetbrains.kotlinx.dataframe.api.to
@@ -39,15 +40,20 @@ import kotlin.reflect.full.withNullability
 import kotlin.reflect.jvm.jvmErasure
 
 @PublishedApi
-internal fun <T, C, R> Convert<T, C>.withRowCellImpl(type: KType?, rowConverter: RowValueExpression<T, C, R>): DataFrame<T> =
-    to { col -> df.newColumn(type, col.name) { rowConverter(it, it[col]) } }
+internal fun <T, C, R> Convert<T, C>.withRowCellImpl(
+    type: KType,
+    infer: Infer,
+    rowConverter: RowValueExpression<T, C, R>
+): DataFrame<T> =
+    to { col -> df.newColumn(type, col.name, infer) { rowConverter(it, it[col]) } }
 
 @PublishedApi
 internal fun <T, C, R> Convert<T, C>.convertRowColumnImpl(
-    type: KType?,
+    type: KType,
+    infer: Infer,
     rowConverter: RowColumnExpression<T, C, R>
 ): DataFrame<T> =
-    to { col -> df.newColumn(type, col.name) { rowConverter(it, col) } }
+    to { col -> df.newColumn(type, col.name, infer) { rowConverter(it, col) } }
 
 internal fun AnyCol.convertToTypeImpl(newType: KType): AnyCol {
     val from = type

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

@@ -2,28 +2,30 @@ package org.jetbrains.kotlinx.dataframe.impl.columns
 
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
+import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import kotlin.reflect.KType
 
 internal class ComputedColumnReference<R>(
     val name: String,
-    val type: KType?,
+    val type: KType,
+    val infer: Infer,
     val compute: RowExpression<Any?, R>
 ) :
     ColumnReference<R> {
 
     override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<R> {
-        return context.df.newColumn(type, name) { compute(it, it) }.addPath(context.df)
+        return context.df.newColumn(type, name, infer) { compute(it, it) }.addPath(context.df)
     }
 
     override fun name() = name
 
-    override fun rename(newName: String) = ComputedColumnReference(newName, type, compute)
+    override fun rename(newName: String) = ComputedColumnReference(newName, type, infer, compute)
 
     override fun getValue(row: AnyRow) = compute(row, row)
 }
 
 @PublishedApi
-internal fun <R> createComputedColumnReference(name: String, type: KType?, compute: RowExpression<Any?, R>): ColumnReference<R> = ComputedColumnReference(name, type, compute)
+internal fun <R> createComputedColumnReference(name: String, type: KType, infer: Infer, compute: RowExpression<Any?, R>): ColumnReference<R> = ComputedColumnReference(name, type, infer, compute)

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

@@ -42,17 +42,16 @@ import kotlin.reflect.full.withNullability
 
 @PublishedApi
 internal fun <T, R> ColumnsContainer<T>.newColumn(
-    type: KType?,
+    type: KType,
     name: String = "",
-    inferNulls: Boolean = true,
+    infer: Infer = Infer.Nulls,
     expression: AddExpression<T, R>
 ): DataColumn<R> {
     val (nullable, values) = computeValues(this as DataFrame<T>, expression)
-    if (type == null) return guessColumnType(name, values)
-    return when (type.classifier) {
-        DataFrame::class -> DataColumn.createFrameColumn(name, values as List<AnyFrame>) as DataColumn<R>
-        DataRow::class -> DataColumn.createColumnGroup(name, (values as List<AnyRow>).concat()) as DataColumn<R>
-        else -> DataColumn.createValueColumn(name, values, if (inferNulls) type.withNullability(nullable) else type)
+    return when (infer) {
+        Infer.Nulls -> DataColumn.create(name, values, type.withNullability(nullable), Infer.None)
+        Infer.Type -> DataColumn.createWithTypeInference(name, values, nullable)
+        Infer.None -> DataColumn.create(name, values, type, Infer.None)
     }
 }
 

src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/inferType.kt

@@ -22,7 +22,7 @@ class InferTypeTests {
     fun `infer type with argument`() {
         val col by columnOf(1)
         val df = dataFrameOf(col)
-        val converted = df.convert(col).with(inferType = false) {
+        val converted = df.convert(col).with(Infer.None) {
             B(it) as A<Int>
         }
         converted[col].type() shouldBe typeOf<A<Int>>()