diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/AdvancedDbType.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/AdvancedDbType.kt new file mode 100644 index 0000000000..31c824b725 --- /dev/null +++ b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/AdvancedDbType.kt @@ -0,0 +1,76 @@ +package org.jetbrains.kotlinx.dataframe.io.db + +import org.jetbrains.kotlinx.dataframe.DataColumn +import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema +import java.sql.ResultSet +import kotlin.reflect.KType + +/** + * Alternative version of [DbType] that allows to customize type mapping + * by initializing a [JdbcToDataFrameConverter] instance for each JDBC type. + * + * This can be helpful for JDBC databases that support structured data, like [DuckDb] + * or that need to a lot of type mapping. + */ +public abstract class AdvancedDbType(dbTypeInJdbcUrl: String) : DbType(dbTypeInJdbcUrl) { + + protected abstract fun generateConverter(tableColumnMetadata: TableColumnMetadata): AnyJdbcToDataFrameConverter + + private val converterCache = mutableMapOf() + + protected fun getConverter(tableColumnMetadata: TableColumnMetadata): AnyJdbcToDataFrameConverter = + converterCache.getOrPut(tableColumnMetadata) { + generateConverter(tableColumnMetadata) + } + + final override fun getExpectedJdbcType(tableColumnMetadata: TableColumnMetadata): KType = + getConverter(tableColumnMetadata).expectedJdbcType + + final override fun getPreprocessedValueType( + tableColumnMetadata: TableColumnMetadata, + expectedJdbcType: KType, + ): KType = getConverter(tableColumnMetadata).preprocessedValueType + + final override fun getTargetColumnSchema( + tableColumnMetadata: TableColumnMetadata, + expectedValueType: KType, + ): ColumnSchema = getConverter(tableColumnMetadata).targetSchema + + final override fun getValueFromResultSet( + rs: ResultSet, + columnIndex: Int, + tableColumnMetadata: TableColumnMetadata, + expectedJdbcType: KType, + ): J? = + getConverter(tableColumnMetadata).cast() + .getValueFromResultSetOrElse(rs, columnIndex) { + try { + rs.getObject(columnIndex + 1) + } catch (_: Throwable) { + // TODO? + rs.getString(columnIndex + 1) + } as J? + } + + final override fun preprocessValue( + value: J?, + tableColumnMetadata: TableColumnMetadata, + expectedJdbcType: KType, + expectedPreprocessedValueType: KType, + ): D? = getConverter(tableColumnMetadata).cast().preprocessOrCast(value) + + final override fun buildDataColumn( + name: String, + values: List, + tableColumnMetadata: TableColumnMetadata, + targetColumnSchema: ColumnSchema, + inferNullability: Boolean, + ): DataColumn = + getConverter(tableColumnMetadata).cast() + .buildDataColumnOrNull(name, values, inferNullability) + ?: values.toDataColumn( + name = name, + targetColumnSchema = targetColumnSchema, + inferNullability = inferNullability, + ) +} diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt index e4fb482d8f..3da18fe5dd 100644 --- a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt +++ b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt @@ -1,7 +1,15 @@ package org.jetbrains.kotlinx.dataframe.io.db +import kotlinx.datetime.LocalDateTime +import kotlinx.datetime.toKotlinLocalDateTime +import org.jetbrains.kotlinx.dataframe.AnyFrame +import org.jetbrains.kotlinx.dataframe.AnyRow import org.jetbrains.kotlinx.dataframe.DataColumn import org.jetbrains.kotlinx.dataframe.api.Infer +import org.jetbrains.kotlinx.dataframe.api.asDataColumn +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.schema +import org.jetbrains.kotlinx.dataframe.api.toDataFrame import org.jetbrains.kotlinx.dataframe.io.DbConnectionConfig import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema @@ -21,17 +29,21 @@ import java.sql.SQLXML import java.sql.Time import java.sql.Timestamp import java.sql.Types -import java.time.LocalDateTime import java.time.OffsetDateTime import java.time.OffsetTime -import java.util.Date import java.util.UUID +import kotlin.collections.toTypedArray import kotlin.reflect.KClass import kotlin.reflect.KType -import kotlin.reflect.full.createType -import kotlin.reflect.full.isSupertypeOf import kotlin.reflect.full.safeCast -import kotlin.reflect.full.starProjectedType +import kotlin.reflect.full.withNullability +import kotlin.reflect.typeOf +import kotlin.time.Instant +import kotlin.time.toKotlinInstant +import kotlin.uuid.Uuid +import kotlin.uuid.toKotlinUuid +import java.time.LocalDateTime as JavaLocalDateTime +import java.util.Date as JavaDate /** * The `DbType` class represents a database type used for reading dataframe from the database. @@ -39,6 +51,7 @@ import kotlin.reflect.full.starProjectedType * @property [dbTypeInJdbcUrl] The name of the database as specified in the JDBC URL. */ public abstract class DbType(public val dbTypeInJdbcUrl: String) { + /** * Represents the JDBC driver class name for a given database type. * @@ -82,10 +95,210 @@ public abstract class DbType(public val dbTypeInJdbcUrl: String) { */ public open val defaultQueryTimeout: Int? = null // null = no timeout + /** Default mapping of [Java SQL Types][Types] to [KType]. */ + protected val defaultJdbcTypeToKTypeMapping: Map = mapOf( + Types.BIT to typeOf(), + Types.TINYINT to typeOf(), + Types.SMALLINT to typeOf(), + Types.INTEGER to typeOf(), + Types.BIGINT to typeOf(), + Types.FLOAT to typeOf(), + Types.REAL to typeOf(), + Types.DOUBLE to typeOf(), + Types.NUMERIC to typeOf(), + Types.DECIMAL to typeOf(), + Types.CHAR to typeOf(), + Types.VARCHAR to typeOf(), + Types.LONGVARCHAR to typeOf(), + Types.DATE to typeOf(), + Types.TIME to typeOf