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

The newest version!

@file:Suppress("INVISIBLE_REFERENCE")

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.impl.columns.ColumnGroupImpl
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation

fun PluginDataFrameSchema.asDataFrame(): DataFrame {
    val df = columns().map()
    return df
}

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

interface ConeTypesAdapter

private fun List.map(): DataFrame {
    val columns = map {
        it.asDataColumn()
    }
    return dataFrameOf(columns).cast()
}

@Suppress("INVISIBLE_REFERENCE")
fun SimpleCol.asDataColumn(): DataColumn<*> {
    val column = when (this) {
        is SimpleDataColumn -> DataColumn.createByType(this.name, listOf(this.type))
        is SimpleColumnGroup -> DataColumn.createColumnGroup(this.name, this.columns().map()) as ColumnGroupImpl<*>
        is SimpleFrameColumn -> DataColumn.createFrameColumn(this.name, listOf(this.columns().map()))
    }
    return column
}

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)
            }
        }
    }