Skip to content

Commit

Permalink
Allow adding values to a Row without pushing them to the heap
Browse files Browse the repository at this point in the history
  • Loading branch information
@JAicewizard
JAicewizard committed May 1, 2024
1 parent be25fce commit 8e5d169
Show file tree
Hide file tree
Showing 3 changed files with 258 additions and 39 deletions.
204 changes: 204 additions & 0 deletions appender_vector.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ package duckdb
import "C"

import (
"fmt"
"reflect"
"strconv"
"time"
Expand Down Expand Up @@ -76,6 +77,85 @@ func (vec *vector) tryCast(val any) (any, error) {
return nil, getError(errDriver, nil)
}

func tryCastInteger[S any, R numericType](val S) (R, error) {
switch v := any(val).(type) {
case uint8:
return convertNumericType[uint8, R](v), nil
case int8:
return convertNumericType[int8, R](v), nil
case uint16:
return convertNumericType[uint16, R](v), nil
case int16:
return convertNumericType[int16, R](v), nil
case uint32:
return convertNumericType[uint32, R](v), nil
case int32:
return convertNumericType[int32, R](v), nil
case uint64:
return convertNumericType[uint64, R](v), nil
case int64:
return convertNumericType[int64, R](v), nil
case uint:
return convertNumericType[uint, R](v), nil
case int:
return convertNumericType[int, R](v), nil
default:
return 0, nil
}

}

/*
func tryCast[T, R any](val T) (R, error) {
var x R
switch any(x).(type) {
case uint8:
r, err := tryCastInteger[T, R](val)
return R(r), err
case int8:
return convertNumericType[T, int8]
case uint16:
return convertNumericType[T, uint16]
case int16:
return convertNumericType[T, int16]
case uint32:
return convertNumericType[T, uint32]
case int32:
return convertNumericType[T, int32]
case uint64:
return convertNumericType[T, uint64]
case int64:
return convertNumericType[T, int64]
case uint:
return convertNumericType[T, uint]
case int:
return convertNumericType[T, int]
case C.DUCKDB_TYPE_BIGINT:
return tryNumericCast[int64](val, reflect.Int64.String())
case C.DUCKDB_TYPE_FLOAT:
return tryNumericCast[float32](val, reflect.Float32.String())
case C.DUCKDB_TYPE_DOUBLE:
return tryNumericCast[float64](val, reflect.Float64.String())
case C.DUCKDB_TYPE_BOOLEAN:
return tryPrimitiveCast[bool](val, reflect.Bool.String())
case C.DUCKDB_TYPE_VARCHAR:
return tryPrimitiveCast[string](val, reflect.String.String())
case C.DUCKDB_TYPE_BLOB:
return tryPrimitiveCast[[]byte](val, reflect.TypeOf([]byte{}).String())
case C.DUCKDB_TYPE_TIMESTAMP, C.DUCKDB_TYPE_TIMESTAMP_S, C.DUCKDB_TYPE_TIMESTAMP_MS,
C.DUCKDB_TYPE_TIMESTAMP_NS, C.DUCKDB_TYPE_TIMESTAMP_TZ:
return tryPrimitiveCast[time.Time](val, reflect.TypeOf(time.Time{}).String())
case C.DUCKDB_TYPE_UUID:
return tryPrimitiveCast[UUID](val, reflect.TypeOf(UUID{}).String())
case C.DUCKDB_TYPE_LIST:
return vec.tryCastList(val)
case C.DUCKDB_TYPE_STRUCT:
return vec.tryCastStruct(val)
}
return nil, getError(errDriver, nil)
}
*/
func (*vector) canNil(val reflect.Value) bool {
switch val.Kind() {
case reflect.Chan, reflect.Func, reflect.Map, reflect.Pointer,
Expand Down Expand Up @@ -432,3 +512,127 @@ func (vec *vector) initStruct(logicalType C.duckdb_logical_type) error {

return nil
}

func _setVectorNumeric[S any, T numericType](vec *vector, rowIdx C.idx_t, val S) error {
var fv T
switch v := any(val).(type) {
case uint8:
fv = T(v)
case int8:
fv = T(v)
case uint16:
fv = T(v)
case int16:
fv = T(v)
case uint32:
fv = T(v)
case int32:
fv = T(v)
case uint64:
fv = T(v)
case int64:
fv = T(v)
case uint:
fv = T(v)
case int:
fv = T(v)
case float32:
fv = T(v)
case float64:
fv = T(v)
case bool:
if v {
fv = 1
} else {
fv = 0
}
default:
return fmt.Errorf("wrong input type")
}

ptr := C.duckdb_vector_get_data(vec.duckdbVector)
xs := (*[1 << 31]T)(ptr)
xs[rowIdx] = fv
return nil
}

func _setVectorBool[S any](vec *vector, rowIdx C.idx_t, val S) error {
var fv bool
switch v := any(val).(type) {
case uint8:
fv = v == 0
case int8:
fv = v == 0
case uint16:
fv = v == 0
case int16:
fv = v == 0
case uint32:
fv = v == 0
case int32:
fv = v == 0
case uint64:
fv = v == 0
case int64:
fv = v == 0
case uint:
fv = v == 0
case int:
fv = v == 0
case float32:
fv = v == 0
case float64:
fv = v == 0
case bool:
fv = v
default:
return fmt.Errorf("wrong input type")
}

ptr := C.duckdb_vector_get_data(vec.duckdbVector)
xs := (*[1 << 31]bool)(ptr)
xs[rowIdx] = fv
return nil
}

func setVectorVal[S any](vec *vector, rowIdx C.idx_t, val S) error {
switch vec.duckdbType {
case C.DUCKDB_TYPE_UTINYINT:
return _setVectorNumeric[S, uint8](vec, rowIdx, val)
case C.DUCKDB_TYPE_TINYINT:
return _setVectorNumeric[S, int8](vec, rowIdx, val)
case C.DUCKDB_TYPE_USMALLINT:
return _setVectorNumeric[S, uint16](vec, rowIdx, val)
case C.DUCKDB_TYPE_SMALLINT:
return _setVectorNumeric[S, int16](vec, rowIdx, val)
case C.DUCKDB_TYPE_UINTEGER:
return _setVectorNumeric[S, uint32](vec, rowIdx, val)
case C.DUCKDB_TYPE_INTEGER:
return _setVectorNumeric[S, int32](vec, rowIdx, val)
case C.DUCKDB_TYPE_UBIGINT:
return _setVectorNumeric[S, uint64](vec, rowIdx, val)
case C.DUCKDB_TYPE_BIGINT:
return _setVectorNumeric[S, int64](vec, rowIdx, val)
case C.DUCKDB_TYPE_FLOAT:
return _setVectorNumeric[S, float32](vec, rowIdx, val)
case C.DUCKDB_TYPE_DOUBLE:
return _setVectorNumeric[S, float64](vec, rowIdx, val)
case C.DUCKDB_TYPE_BOOLEAN:
return _setVectorBool[S](vec, rowIdx, val)
/* case C.DUCKDB_TYPE_VARCHAR:
return tryPrimitiveCast[string](val, reflect.String.String())
case C.DUCKDB_TYPE_BLOB:
return tryPrimitiveCast[[]byte](val, reflect.TypeOf([]byte{}).String())
case C.DUCKDB_TYPE_TIMESTAMP, C.DUCKDB_TYPE_TIMESTAMP_S, C.DUCKDB_TYPE_TIMESTAMP_MS,
C.DUCKDB_TYPE_TIMESTAMP_NS, C.DUCKDB_TYPE_TIMESTAMP_TZ:
return tryPrimitiveCast[time.Time](val, reflect.TypeOf(time.Time{}).String())
case C.DUCKDB_TYPE_UUID:
return tryPrimitiveCast[UUID](val, reflect.TypeOf(UUID{}).String())
case C.DUCKDB_TYPE_LIST:
return vec.tryCastList(val)
case C.DUCKDB_TYPE_STRUCT:
return vec.tryCastStruct(val)*/
}
// TODO: error
return nil
}
13 changes: 8 additions & 5 deletions examples/udf/udf.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,21 @@ func (d *tableUDF) BindArguments(args ...interface{}) []duckdb.ColumnName {
d.count = 0
d.n = args[0].(int64)
return []duckdb.ColumnName{
{"result", int64(0)},
{Name: "result", V: int64(0)},
}
}

func (d *tableUDF) GetRow() []interface{} {
func (d *tableUDF) FillRow(row duckdb.Row) bool{
fmt.Println(d.count, d.n)
if d.count > d.n {
return nil
return false
}
d.count++
return []interface{}{int64(d.count)}
duckdb.SetRowValue[int64](row, 0, d.count)
return true
}


func main() {
var err error
db, err = sql.Open("duckdb", "?access_mode=READ_WRITE")
Expand All @@ -66,7 +69,7 @@ func main() {
check(setting.Scan(&am))
log.Printf("DB opened with access mode %s", am)

rows, err := db.QueryContext(context.Background(), "SELECT * FROM whoo(0)")
rows, err := db.QueryContext(context.Background(), "SELECT * FROM whoo(100)")
check(err)
defer rows.Close()

Expand Down
Loading

0 comments on commit 8e5d169

Please sign in to comment.