Optimize type inference in joinWith

Author: koperagen

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

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.impl.api
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.api.JoinExpression
 import org.jetbrains.kotlinx.dataframe.api.JoinType
 import org.jetbrains.kotlinx.dataframe.api.JoinedDataRow
@@ -12,6 +13,7 @@ import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.count
 import org.jetbrains.kotlinx.dataframe.api.indices
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
 import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
 import org.jetbrains.kotlinx.dataframe.impl.DataRowImpl
 
@@ -91,8 +93,26 @@ internal fun <A, B> DataFrame<A>.joinWithImpl(
         }
     }
 
+    val leftColumns = columns()
+    val rightColumns = if (addNewColumns) right.columns() else emptyList()
     val df: DataFrame<*> = outputData.mapIndexed { index, values ->
-        DataColumn.createByInference(generator.names[index], values)
+        val srcColumn = if (index < leftColumns.size) {
+            leftColumns[index]
+        } else {
+            rightColumns[index - leftColumns.size]
+        }
+        // let's optimize an easy case.
+        // handling introduction of nulls into ColumnGroup and FrameColumn is not straightforward
+        when (srcColumn.kind()) {
+            ColumnKind.Value -> DataColumn.createByType(
+                name = generator.names[index],
+                values = values,
+                type = srcColumn.type(),
+                infer = Infer.Nulls,
+            )
+
+            else -> DataColumn.createByInference(generator.names[index], values)
+        }
     }.toDataFrame()
 
     return df.cast()

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.api.rightJoinWith
 import org.jetbrains.kotlinx.dataframe.api.select
 import org.junit.Test
+import kotlin.reflect.typeOf
 
 @Suppress("ktlint:standard:argument-list-wrapping")
 class JoinWithTests : BaseJoinTest() {
@@ -34,9 +35,12 @@ class JoinWithTests : BaseJoinTest() {
         res.columnsCount() shouldBe 8
         res.rowsCount() shouldBe 7
         res["age1"].hasNulls() shouldBe false
+        res["age1"].type() shouldBe typeOf<String>()
+        res["age1"].values().all { it != null } shouldBe true
         res.count { name == "Charlie" && city == "Moscow" } shouldBe 4
         res.select { city and name }.distinct().rowsCount() shouldBe 3
         res[Person2::grade].hasNulls() shouldBe false
+        res.age.type() shouldBe typeOf<Int>()
     }
 
     @Test
@@ -49,6 +53,8 @@ class JoinWithTests : BaseJoinTest() {
         res.select { city and name }.distinct().rowsCount() shouldBe 6
         res.count { it["grade"] == null } shouldBe 3
         res.age.hasNulls() shouldBe false
+        res.age.type() shouldBe typeOf<Int>()
+        res["age1"].type() shouldBe typeOf<String?>()
     }
 
     @Test
@@ -62,6 +68,7 @@ class JoinWithTests : BaseJoinTest() {
         res.select { city and name }.distinct().rowsCount() shouldBe 4
         res[Person2::grade].hasNulls() shouldBe false
         res.age.hasNulls() shouldBe true
+        res.age.type() shouldBe typeOf<Int?>()
         val newEntries = res.filter { it["age"] == null }
         newEntries.rowsCount() shouldBe 2
         newEntries.all { it["name1"] == "Bob" && it["origin"] == "Paris" && weight == null } shouldBe true
@@ -78,6 +85,8 @@ class JoinWithTests : BaseJoinTest() {
         val distinct = res.select { name and age and city and weight }.distinct()
         val expected = typed.append(null, null, null, null)
         distinct shouldBe expected
+        res.age.type() shouldBe typeOf<Int?>()
+        res["age1"].type() shouldBe typeOf<String?>()
     }
 
     @Test
@@ -103,7 +112,7 @@ class JoinWithTests : BaseJoinTest() {
     }
 
     @Test
-    fun rightJoin() {
+    fun `exclude join`() {
         val df = dataFrameOf("a", "b")(
             1, "a",
             2, "b",
@@ -116,7 +125,17 @@ class JoinWithTests : BaseJoinTest() {
             2, "II",
             3, "III",
         )
+
         df.append(4, "e").excludeJoin(df1).print()
+
+        val res = df.append(4, "e").excludeJoin(df1)
+
+        res.rowsCount() shouldBe 1
+        res["a"].values() shouldBe listOf(4)
+        res["b"].values() shouldBe listOf("e")
+        res.columnsCount() shouldBe 2
+        res["a"].type() shouldBe typeOf<Int>()
+        res["b"].type() shouldBe typeOf<String>()
     }
 
     @Test