Skip to content

Commit

Permalink
Implement TimeStampMicroTZVector for parquet isAdjustedToUTC timestam…
Browse files Browse the repository at this point in the history
…p columns Kotlin#926
  • Loading branch information
Laurent Perez committed Oct 17, 2024
1 parent 54d00e6 commit b5c614c
Showing 1 changed file with 16 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
package org.jetbrains.kotlinx.dataframe.io

import kotlinx.datetime.Instant
import kotlinx.datetime.LocalDate
import kotlinx.datetime.LocalDateTime
import kotlinx.datetime.LocalTime
import kotlinx.datetime.TimeZone
import kotlinx.datetime.toKotlinLocalDate
import kotlinx.datetime.toKotlinLocalDateTime
import kotlinx.datetime.toKotlinLocalTime
import kotlinx.datetime.toLocalDateTime
import org.apache.arrow.memory.RootAllocator
import org.apache.arrow.vector.BigIntVector
import org.apache.arrow.vector.BitVector
Expand All @@ -21,6 +24,7 @@ import org.apache.arrow.vector.LargeVarBinaryVector
import org.apache.arrow.vector.LargeVarCharVector
import org.apache.arrow.vector.NullVector
import org.apache.arrow.vector.SmallIntVector
import org.apache.arrow.vector.TimeStampMicroTZVector
import org.apache.arrow.vector.TimeMicroVector
import org.apache.arrow.vector.TimeMilliVector
import org.apache.arrow.vector.TimeNanoVector
Expand Down Expand Up @@ -179,6 +183,16 @@ private fun TimeStampMicroVector.values(range: IntRange): List<LocalDateTime?> =
}
}

private fun TimeStampMicroTZVector.values(range: IntRange): List<LocalDateTime?> =
range.mapIndexed { i, it ->
if (isNull(i)) {
null
} else {
Instant.fromEpochMilliseconds(getObject(it) / 1000)
.toLocalDateTime(TimeZone.of(this.timeZone))
}
}

private fun TimeStampMilliVector.values(range: IntRange): List<LocalDateTime?> =
range.mapIndexed { i, it ->
if (isNull(i)) {
Expand Down Expand Up @@ -345,6 +359,8 @@ private fun readField(root: VectorSchemaRoot, field: Field, nullability: Nullabi

is TimeStampMicroVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)

is TimeStampMicroTZVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)

is TimeStampMilliVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)

is TimeStampSecVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
Expand Down

0 comments on commit b5c614c

Please sign in to comment.