|
| 1 | +package org.jetbrains.kotlinx.dataframe.examples.exposed |
| 2 | + |
| 3 | +import org.jetbrains.exposed.v1.core.BiCompositeColumn |
| 4 | +import org.jetbrains.exposed.v1.core.Column |
| 5 | +import org.jetbrains.exposed.v1.core.Expression |
| 6 | +import org.jetbrains.exposed.v1.core.ExpressionAlias |
| 7 | +import org.jetbrains.exposed.v1.core.ResultRow |
| 8 | +import org.jetbrains.exposed.v1.core.Table |
| 9 | +import org.jetbrains.exposed.v1.jdbc.Query |
| 10 | +import org.jetbrains.kotlinx.dataframe.AnyFrame |
| 11 | +import org.jetbrains.kotlinx.dataframe.DataFrame |
| 12 | +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema |
| 13 | +import org.jetbrains.kotlinx.dataframe.api.convertTo |
| 14 | +import org.jetbrains.kotlinx.dataframe.api.toDataFrame |
| 15 | +import org.jetbrains.kotlinx.dataframe.codeGen.NameNormalizer |
| 16 | +import org.jetbrains.kotlinx.dataframe.impl.schema.DataFrameSchemaImpl |
| 17 | +import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema |
| 18 | +import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema |
| 19 | +import kotlin.reflect.KProperty1 |
| 20 | +import kotlin.reflect.full.isSubtypeOf |
| 21 | +import kotlin.reflect.full.memberProperties |
| 22 | +import kotlin.reflect.typeOf |
| 23 | + |
| 24 | +/** |
| 25 | + * Retrieves all columns of any [Iterable][Iterable]`<`[ResultRow][ResultRow]`>`, like [Query][Query], |
| 26 | + * from Exposed row by row and converts the resulting [Map] into a [DataFrame], cast to type [T]. |
| 27 | + * |
| 28 | + * In notebooks, the untyped version works just as well due to runtime inference :) |
| 29 | + */ |
| 30 | +inline fun <reified T : Any> Iterable<ResultRow>.convertToDataFrame(): DataFrame<T> = |
| 31 | + convertToDataFrame().convertTo<T>() |
| 32 | + |
| 33 | +/** |
| 34 | + * Retrieves all columns of an [Iterable][Iterable]`<`[ResultRow][ResultRow]`>` from Exposed, like [Query][Query], |
| 35 | + * row by row and converts the resulting [Map] of lists into a [DataFrame] by calling |
| 36 | + * [Map.toDataFrame]. |
| 37 | + */ |
| 38 | +@JvmName("convertToAnyFrame") |
| 39 | +fun Iterable<ResultRow>.convertToDataFrame(): AnyFrame { |
| 40 | + val map = mutableMapOf<String, MutableList<Any?>>() |
| 41 | + for (row in this) { |
| 42 | + for (expression in row.fieldIndex.keys) { |
| 43 | + map.getOrPut(expression.readableName) { |
| 44 | + mutableListOf() |
| 45 | + } += row[expression] |
| 46 | + } |
| 47 | + } |
| 48 | + return map.toDataFrame() |
| 49 | +} |
| 50 | + |
| 51 | +/** |
| 52 | + * Retrieves a simple column name from [this] [Expression]. |
| 53 | + * |
| 54 | + * Might need to be expanded with multiple types of [Expression]. |
| 55 | + */ |
| 56 | +val Expression<*>.readableName: String |
| 57 | + get() = when (this) { |
| 58 | + is Column<*> -> name |
| 59 | + is ExpressionAlias<*> -> alias |
| 60 | + is BiCompositeColumn<*, *, *> -> getRealColumns().joinToString("_") { it.readableName } |
| 61 | + else -> toString() |
| 62 | + } |
| 63 | + |
| 64 | +/** |
| 65 | + * Creates a [DataFrameSchema] from the declared [Table] instance. |
| 66 | + * |
| 67 | + * This is not needed for conversion, but it can be useful to create a DataFrame [@DataSchema][DataSchema] instance. |
| 68 | + * |
| 69 | + * @param columnNameToAccessor Optional [MutableMap] which will be filled with entries mapping |
| 70 | + * the SQL column name to the accessor name from the [Table]. |
| 71 | + * This can be used to define a [NameNormalizer] later. |
| 72 | + * @see toDataFrameSchemaWithNameNormalizer |
| 73 | + */ |
| 74 | +@Suppress("UNCHECKED_CAST") |
| 75 | +fun Table.toDataFrameSchema(columnNameToAccessor: MutableMap<String, String> = mutableMapOf()): DataFrameSchema { |
| 76 | + // we use reflection to go over all `Column<*>` properties in the Table object |
| 77 | + val columns = this::class.memberProperties |
| 78 | + .filter { it.returnType.isSubtypeOf(typeOf<Column<*>>()) } |
| 79 | + .associate { prop -> |
| 80 | + prop as KProperty1<Table, Column<*>> |
| 81 | + |
| 82 | + // retrieve the SQL column name |
| 83 | + val columnName = prop.get(this).name |
| 84 | + // store the SQL column name together with the accessor name in the map |
| 85 | + columnNameToAccessor[columnName] = prop.name |
| 86 | + |
| 87 | + // get the column type from `val a: Column<Type>` |
| 88 | + val type = prop.returnType.arguments.first().type!! |
| 89 | + |
| 90 | + // and we add the name and column shema type to the `columns` map :) |
| 91 | + columnName to ColumnSchema.Value(type) |
| 92 | + } |
| 93 | + return DataFrameSchemaImpl(columns) |
| 94 | +} |
| 95 | + |
| 96 | +/** |
| 97 | + * Creates a [DataFrameSchema] from the declared [Table] instance with a [NameNormalizer] to |
| 98 | + * convert the SQL column names to the corresponding Kotlin property names. |
| 99 | + * |
| 100 | + * This is not needed for conversion, but it can be useful to create a DataFrame [@DataSchema][DataSchema] instance. |
| 101 | + * |
| 102 | + * @see toDataFrameSchema |
| 103 | + */ |
| 104 | +fun Table.toDataFrameSchemaWithNameNormalizer(): Pair<DataFrameSchema, NameNormalizer> { |
| 105 | + val columnNameToAccessor = mutableMapOf<String, String>() |
| 106 | + return Pair(toDataFrameSchema(), NameNormalizer { columnNameToAccessor[it] ?: it }) |
| 107 | +} |
0 commit comments