org.jetbrains.kotlinx.dataframe.plugin.impl.DataFrameAdapter.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of compiler-plugin-all Show documentation

Data processing in Kotlin

There is a newer version: 0.16.0-dev-5328

package org.jetbrains.kotlinx.dataframe.plugin.impl

import org.jetbrains.kotlinx.dataframe.AnyCol
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation

fun PluginDataFrameSchema.asDataFrame(): DataFrame = columns().map()

fun DataFrame.toPluginDataFrameSchema() = PluginDataFrameSchema(columns().mapBack())

interface ConeTypesAdapter

private fun List.map(): DataFrame = map {
    when (it) {
        is SimpleDataColumn -> DataColumn.create(it.name, listOf(it.type))
        is SimpleColumnGroup -> DataColumn.createColumnGroup(it.name, it.columns().map())
        is SimpleFrameColumn -> DataColumn.createFrameColumn(it.name, listOf(it.columns().map()))
    }
}.let { dataFrameOf(it).cast() }

private fun List.mapBack(): List = map {
    when (it) {
        is ColumnGroup<*> -> {
            SimpleColumnGroup(it.name(), it.columns().mapBack())
        }
        is FrameColumn<*> -> {
            SimpleFrameColumn(it.name(), it[0].columns().mapBack())
        }
        else -> {
            SimpleDataColumn(it.name(), it[0] as TypeApproximation)
        }
    }
}