Support writing ColumnGroup into Arrow as StructVector

Author: koperagen

dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowWriterImpl.kt

@@ -26,6 +26,7 @@ import org.apache.arrow.vector.TinyIntVector
 import org.apache.arrow.vector.VarCharVector
 import org.apache.arrow.vector.VariableWidthVector
 import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.types.DateUnit
 import org.apache.arrow.vector.types.FloatingPointPrecision
 import org.apache.arrow.vector.types.pojo.ArrowType
@@ -49,8 +50,10 @@ import org.jetbrains.kotlinx.dataframe.api.convertToShort
 import org.jetbrains.kotlinx.dataframe.api.convertToString
 import org.jetbrains.kotlinx.dataframe.api.forEachIndexed
 import org.jetbrains.kotlinx.dataframe.api.map
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.exceptions.CellConversionException
 import org.jetbrains.kotlinx.dataframe.exceptions.TypeConverterNotFoundException
+import org.jetbrains.kotlinx.dataframe.indices
 import org.jetbrains.kotlinx.dataframe.name
 import org.jetbrains.kotlinx.dataframe.values
 import kotlin.reflect.full.isSubtypeOf
@@ -72,7 +75,15 @@ internal class ArrowWriterImpl(
     private fun allocateVector(vector: FieldVector, size: Int, totalBytes: Long? = null) {
         when (vector) {
             is FixedWidthVector -> vector.allocateNew(size)
+
             is VariableWidthVector -> totalBytes?.let { vector.allocateNew(it, size) } ?: vector.allocateNew(size)
+
+            is StructVector -> {
+                vector.childrenFromFields.forEach { child ->
+                    allocateVector(child, size)
+                }
+            }
+
             else -> throw IllegalArgumentException("Can not allocate ${vector.javaClass.canonicalName}")
         }
     }
@@ -138,6 +149,8 @@ internal class ArrowWriterImpl(
 
             is ArrowType.Time -> column.convertToLocalTime()
 
+            is ArrowType.Struct -> column
+
             else ->
                 throw NotImplementedError(
                     "Saving ${targetFieldType.javaClass.canonicalName} is currently not implemented",
@@ -277,6 +290,18 @@ internal class ArrowWriterImpl(
                         } ?: vector.setNull(i)
                     }
 
+            is StructVector -> {
+                require(column is ColumnGroup<*>) {
+                    "StructVector expects ColumnGroup, but got ${column::class.simpleName}"
+                }
+
+                column.columns().forEach { childColumn ->
+                    infillVector(vector.getChild(childColumn.name()), childColumn)
+                }
+
+                column.indices.forEach { i -> vector.setIndexDefined(i) }
+            }
+
             else -> {
                 // TODO implement other vector types from [readField] (VarBinaryVector, UIntVector, DurationVector, StructVector) and may be others (ListVector, FixedSizeListVector etc)
                 throw NotImplementedError("Saving to ${vector.javaClass.canonicalName} is currently not implemented")

dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowTypesMatching.kt

@@ -11,6 +11,7 @@ import org.apache.arrow.vector.types.pojo.Field
 import org.apache.arrow.vector.types.pojo.FieldType
 import org.apache.arrow.vector.types.pojo.Schema
 import org.jetbrains.kotlinx.dataframe.AnyCol
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.typeClass
 import kotlin.reflect.full.isSubtypeOf
 import kotlin.reflect.typeOf
@@ -27,6 +28,15 @@ public fun AnyCol.toArrowField(mismatchSubscriber: (ConvertingMismatch) -> Unit
     val columnType = column.type()
     val nullable = columnType.isMarkedNullable
     return when {
+        column is ColumnGroup<*> -> {
+            val childFields = column.columns().map { it.toArrowField(mismatchSubscriber) }
+            Field(
+                column.name(),
+                FieldType(nullable, ArrowType.Struct(), null),
+                childFields,
+            )
+        }
+
         columnType.isSubtypeOf(typeOf<String?>()) ->
             Field(
                 column.name(),

dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowKtTest.kt

@@ -38,6 +38,7 @@ import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.map
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.remove
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.exceptions.TypeConverterNotFoundException
 import org.junit.Assert
 import org.junit.Test
@@ -724,4 +725,89 @@ internal class ArrowKtTest {
 
         dataFrame.rowsCount() shouldBe 900
     }
+
+    @Test
+    fun testColumnGroupRoundtrip() {
+        val original = dataFrameOf(
+            "outer" to columnOf("x", "y", "z"),
+            "inner" to columnOf(
+                "nested1" to columnOf("a", "b", "c"),
+                "nested2" to columnOf(1, 2, 3),
+            ),
+        )
+
+        val featherBytes = original.saveArrowFeatherToByteArray()
+        val fromFeather = DataFrame.readArrowFeather(featherBytes)
+        fromFeather shouldBe original
+
+        val ipcBytes = original.saveArrowIPCToByteArray()
+        val fromIpc = DataFrame.readArrowIPC(ipcBytes)
+        fromIpc shouldBe original
+    }
+
+    @Test
+    fun testNestedColumnGroupRoundtrip() {
+        val deeplyNested by columnOf(
+            "level2" to columnOf(
+                "level3" to columnOf(1, 2, 3),
+            ),
+        )
+        val original = dataFrameOf(deeplyNested)
+
+        val bytes = original.saveArrowFeatherToByteArray()
+        val restored = DataFrame.readArrowFeather(bytes)
+
+        restored shouldBe original
+    }
+
+    @Test
+    fun testColumnGroupWithNulls() {
+        val group by columnOf(
+            "a" to columnOf("x", null, "z"),
+            "b" to columnOf(1, 2, null),
+        )
+        val original = dataFrameOf(group)
+
+        val bytes = original.saveArrowFeatherToByteArray()
+        val restored = DataFrame.readArrowFeather(bytes)
+
+        restored shouldBe original
+    }
+
+    @Test
+    fun testReadParquetWithNestedStruct() {
+        val resourceUrl = testResource("books.parquet")
+        val resourcePath = resourceUrl.toURI().toPath()
+
+        val df = DataFrame.readParquet(resourcePath)
+
+        df.columnNames() shouldBe listOf("id", "title", "author", "genre", "publisher")
+
+        val authorGroup = df["author"] as ColumnGroup<*>
+        authorGroup.columnNames() shouldBe listOf("id", "firstName", "lastName")
+
+        df["id"].type() shouldBe typeOf<Int>()
+        df["title"].type() shouldBe typeOf<String>()
+        df["genre"].type() shouldBe typeOf<String>()
+        df["publisher"].type() shouldBe typeOf<String>()
+        authorGroup["id"].type() shouldBe typeOf<Int>()
+        authorGroup["firstName"].type() shouldBe typeOf<String>()
+        authorGroup["lastName"].type() shouldBe typeOf<String>()
+    }
+
+    @Test
+    fun testParquetNestedStructRoundtrip() {
+        val resourceUrl = testResource("books.parquet")
+        val resourcePath = resourceUrl.toURI().toPath()
+
+        val original = DataFrame.readParquet(resourcePath)
+
+        val featherBytes = original.saveArrowFeatherToByteArray()
+        val fromFeather = DataFrame.readArrowFeather(featherBytes)
+        fromFeather shouldBe original
+
+        val ipcBytes = original.saveArrowIPCToByteArray()
+        val fromIpc = DataFrame.readArrowIPC(ipcBytes)
+        fromIpc shouldBe original
+    }
 }