Skip to content

Commit

Permalink
support missing timestamp types
Browse files Browse the repository at this point in the history
  • Loading branch information
@taniabogatsch
taniabogatsch committed Apr 2, 2024
1 parent 143da09 commit f1c994b
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 65 deletions.
91 changes: 63 additions & 28 deletions appender_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,51 +158,73 @@ func TestAppenderPrimitive(t *testing.T) {
uint64 UBIGINT,
int64 BIGINT,
timestamp TIMESTAMP,
timestampS TIMESTAMP_S,
timestampMS TIMESTAMP_MS,
timestampNS TIMESTAMP_NS,
timestampTZ TIMESTAMPTZ,
float REAL,
double DOUBLE,
string VARCHAR,
bool BOOLEAN
)`)

type row struct {
ID int64
UInt8 uint8
Int8 int8
UInt16 uint16
Int16 int16
UInt32 uint32
Int32 int32
UInt64 uint64
Int64 int64
Timestamp time.Time
Float float32
Double float64
String string
Bool bool
ID int64
UInt8 uint8
Int8 int8
UInt16 uint16
Int16 int16
UInt32 uint32
Int32 int32
UInt64 uint64
Int64 int64
Timestamp time.Time
TimestampS time.Time
TimestampMS time.Time
TimestampNS time.Time
TimestampTZ time.Time
Float float32
Double float64
String string
Bool bool
}

// Get the timestamp for all TS columns.
IST, err := time.LoadLocation("Asia/Kolkata")
require.NoError(t, err)

const longForm = "2006-01-02 15:04:05 MST"
ts, err := time.ParseInLocation(longForm, "2016-01-17 20:04:05 IST", IST)
require.NoError(t, err)

rowsToAppend := make([]row, numAppenderTestRows)
for i := 0; i < numAppenderTestRows; i++ {

u64 := rand.Uint64()
// Go SQL does not support uint64 values with their high bit set (see for example https://github.com/lib/pq/issues/72).
if u64 > 9223372036854775807 {
u64 = 9223372036854775807
}

rowsToAppend[i] = row{
ID: int64(i),
UInt8: uint8(randInt(0, 255)),
Int8: int8(randInt(-128, 127)),
UInt16: uint16(randInt(0, 65535)),
Int16: int16(randInt(-32768, 32767)),
UInt32: uint32(randInt(0, 4294967295)),
Int32: int32(randInt(-2147483648, 2147483647)),
UInt64: u64,
Int64: rand.Int63(),
Timestamp: time.UnixMilli(randInt(0, time.Now().UnixMilli())).UTC(),
Float: rand.Float32(),
Double: rand.Float64(),
String: randString(int(randInt(0, 128))),
Bool: rand.Int()%2 == 0,
ID: int64(i),
UInt8: uint8(randInt(0, 255)),
Int8: int8(randInt(-128, 127)),
UInt16: uint16(randInt(0, 65535)),
Int16: int16(randInt(-32768, 32767)),
UInt32: uint32(randInt(0, 4294967295)),
Int32: int32(randInt(-2147483648, 2147483647)),
UInt64: u64,
Int64: rand.Int63(),
Timestamp: ts,
TimestampS: ts,
TimestampMS: ts,
TimestampNS: ts,
TimestampTZ: ts,
Float: rand.Float32(),
Double: rand.Float64(),
String: randString(int(randInt(0, 128))),
Bool: rand.Int()%2 == 0,
}

require.NoError(t, a.AppendRow(
Expand All @@ -216,6 +238,10 @@ func TestAppenderPrimitive(t *testing.T) {
rowsToAppend[i].UInt64,
rowsToAppend[i].Int64,
rowsToAppend[i].Timestamp,
rowsToAppend[i].TimestampS,
rowsToAppend[i].TimestampMS,
rowsToAppend[i].TimestampNS,
rowsToAppend[i].TimestampTZ,
rowsToAppend[i].Float,
rowsToAppend[i].Double,
rowsToAppend[i].String,
Expand Down Expand Up @@ -246,11 +272,20 @@ func TestAppenderPrimitive(t *testing.T) {
&r.UInt64,
&r.Int64,
&r.Timestamp,
&r.TimestampS,
&r.TimestampMS,
&r.TimestampNS,
&r.TimestampTZ,
&r.Float,
&r.Double,
&r.String,
&r.Bool,
))
rowsToAppend[i].Timestamp = rowsToAppend[i].Timestamp.UTC()
rowsToAppend[i].TimestampS = rowsToAppend[i].TimestampS.UTC()
rowsToAppend[i].TimestampMS = rowsToAppend[i].TimestampMS.UTC()
rowsToAppend[i].TimestampNS = rowsToAppend[i].TimestampNS.UTC()
rowsToAppend[i].TimestampTZ = rowsToAppend[i].TimestampTZ.UTC()
require.Equal(t, rowsToAppend[i], r)
i++
}
Expand Down
66 changes: 33 additions & 33 deletions types.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,42 +14,42 @@ import (
)

var unsupportedAppenderTypeMap = map[C.duckdb_type]string{
C.DUCKDB_TYPE_INVALID: "INVALID",
C.DUCKDB_TYPE_DATE: "DATE",
C.DUCKDB_TYPE_TIME: "TIME",
C.DUCKDB_TYPE_INTERVAL: "INTERVAL",
C.DUCKDB_TYPE_HUGEINT: "HUGEINT",
C.DUCKDB_TYPE_UHUGEINT: "UHUGEINT",
C.DUCKDB_TYPE_DECIMAL: "DECIMAL",
C.DUCKDB_TYPE_TIMESTAMP_S: "TIMESTAMP_S",
C.DUCKDB_TYPE_TIMESTAMP_MS: "TIMESTAMP_MS",
C.DUCKDB_TYPE_TIMESTAMP_NS: "TIMESTAMP_NS",
C.DUCKDB_TYPE_ENUM: "ENUM",
C.DUCKDB_TYPE_MAP: "MAP",
C.DUCKDB_TYPE_UNION: "UNION",
C.DUCKDB_TYPE_BIT: "BIT",
C.DUCKDB_TYPE_TIME_TZ: "TIME_TZ",
C.DUCKDB_TYPE_TIMESTAMP_TZ: "TIMESTAMP_TZ",
C.DUCKDB_TYPE_INVALID: "INVALID",
C.DUCKDB_TYPE_DATE: "DATE",
C.DUCKDB_TYPE_TIME: "TIME",
C.DUCKDB_TYPE_INTERVAL: "INTERVAL",
C.DUCKDB_TYPE_HUGEINT: "HUGEINT",
C.DUCKDB_TYPE_UHUGEINT: "UHUGEINT",
C.DUCKDB_TYPE_DECIMAL: "DECIMAL",
C.DUCKDB_TYPE_ENUM: "ENUM",
C.DUCKDB_TYPE_MAP: "MAP",
C.DUCKDB_TYPE_UNION: "UNION",
C.DUCKDB_TYPE_BIT: "BIT",
C.DUCKDB_TYPE_TIME_TZ: "TIME_TZ",
}

var appenderTypeIdMap = map[C.duckdb_type]string{
C.DUCKDB_TYPE_BOOLEAN: "bool",
C.DUCKDB_TYPE_TINYINT: "int8",
C.DUCKDB_TYPE_SMALLINT: "int16",
C.DUCKDB_TYPE_INTEGER: "int32",
C.DUCKDB_TYPE_BIGINT: "int64",
C.DUCKDB_TYPE_UTINYINT: "uint8",
C.DUCKDB_TYPE_USMALLINT: "uint16",
C.DUCKDB_TYPE_UINTEGER: "uint32",
C.DUCKDB_TYPE_UBIGINT: "uint64",
C.DUCKDB_TYPE_FLOAT: "float32",
C.DUCKDB_TYPE_DOUBLE: "float64",
C.DUCKDB_TYPE_VARCHAR: "string",
C.DUCKDB_TYPE_BLOB: "[]uint8",
C.DUCKDB_TYPE_TIMESTAMP: "time.Time",
C.DUCKDB_TYPE_UUID: "duckdb.UUID",
C.DUCKDB_TYPE_LIST: "slice",
C.DUCKDB_TYPE_STRUCT: "struct",
C.DUCKDB_TYPE_BOOLEAN: "bool",
C.DUCKDB_TYPE_TINYINT: "int8",
C.DUCKDB_TYPE_SMALLINT: "int16",
C.DUCKDB_TYPE_INTEGER: "int32",
C.DUCKDB_TYPE_BIGINT: "int64",
C.DUCKDB_TYPE_UTINYINT: "uint8",
C.DUCKDB_TYPE_USMALLINT: "uint16",
C.DUCKDB_TYPE_UINTEGER: "uint32",
C.DUCKDB_TYPE_UBIGINT: "uint64",
C.DUCKDB_TYPE_FLOAT: "float32",
C.DUCKDB_TYPE_DOUBLE: "float64",
C.DUCKDB_TYPE_VARCHAR: "string",
C.DUCKDB_TYPE_BLOB: "[]uint8",
C.DUCKDB_TYPE_TIMESTAMP: "time.Time",
C.DUCKDB_TYPE_TIMESTAMP_S: "time.Time",
C.DUCKDB_TYPE_TIMESTAMP_MS: "time.Time",
C.DUCKDB_TYPE_TIMESTAMP_NS: "time.Time",
C.DUCKDB_TYPE_UUID: "duckdb.UUID",
C.DUCKDB_TYPE_LIST: "slice",
C.DUCKDB_TYPE_STRUCT: "struct",
C.DUCKDB_TYPE_TIMESTAMP_TZ: "time.Time",
}

type UUID [16]byte
Expand Down
30 changes: 26 additions & 4 deletions vector.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,13 +62,21 @@ func (vec *vector) init(logicalType C.duckdb_logical_type, colIdx int) error {
case C.DUCKDB_TYPE_BLOB:
vec.initBlob()
case C.DUCKDB_TYPE_TIMESTAMP:
vec.initTS(C.DUCKDB_TYPE_TIMESTAMP)
vec.initTS(duckdbType)
case C.DUCKDB_TYPE_TIMESTAMP_S:
vec.initTS(duckdbType)
case C.DUCKDB_TYPE_TIMESTAMP_MS:
vec.initTS(duckdbType)
case C.DUCKDB_TYPE_TIMESTAMP_NS:
vec.initTS(duckdbType)
case C.DUCKDB_TYPE_UUID:
vec.initUUID()
case C.DUCKDB_TYPE_LIST:
err = vec.initList(logicalType, colIdx)
case C.DUCKDB_TYPE_STRUCT:
err = vec.initStruct(logicalType)
case C.DUCKDB_TYPE_TIMESTAMP_TZ:
vec.initTS(duckdbType)
default:
name, found := unsupportedAppenderTypeMap[duckdbType]
if !found {
Expand Down Expand Up @@ -122,9 +130,9 @@ func (vec *vector) setCString(rowIdx C.idx_t, value string, len int) {
C.free(unsafe.Pointer(str))
}

func (vec *vector) setTime(rowIdx C.idx_t, value time.Time) {
func (vec *vector) setTime(rowIdx C.idx_t, ticks int64) {
var ts C.duckdb_timestamp
ts.micros = C.int64_t(value.UTC().UnixMicro())
ts.micros = C.int64_t(ticks)
setPrimitive[C.duckdb_timestamp](vec, rowIdx, ts)
}

Expand Down Expand Up @@ -202,7 +210,21 @@ func (vec *vector) initBlob() {

func (vec *vector) initTS(duckdbType C.duckdb_type) {
vec.fn = func(vec *vector, rowIdx C.idx_t, val any) {
vec.setTime(rowIdx, val.(time.Time))
v := val.(time.Time)
var ticks int64
switch duckdbType {
case C.DUCKDB_TYPE_TIMESTAMP:
ticks = v.UTC().UnixMicro()
case C.DUCKDB_TYPE_TIMESTAMP_S:
ticks = v.UTC().Unix()
case C.DUCKDB_TYPE_TIMESTAMP_MS:
ticks = v.UTC().UnixMilli()
case C.DUCKDB_TYPE_TIMESTAMP_NS:
ticks = v.UTC().UnixNano()
case C.DUCKDB_TYPE_TIMESTAMP_TZ:
ticks = v.UTC().UnixMicro()
}
vec.setTime(rowIdx, ticks)
}
vec.duckdbType = duckdbType
}
Expand Down

0 comments on commit f1c994b

Please sign in to comment.