From ed73b4ac4acd00bfe1bb7ccee15cab73f4a2db19 Mon Sep 17 00:00:00 2001 From: Anshul Khandelwal Date: Tue, 7 Nov 2023 16:52:40 +0530 Subject: [PATCH 1/2] scan enum to string --- duckdb_test.go | 13 ++++++++++++- rows.go | 25 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/duckdb_test.go b/duckdb_test.go index 1c1470b4..5db2b399 100644 --- a/duckdb_test.go +++ b/duckdb_test.go @@ -498,9 +498,20 @@ func TestENUMs(t *testing.T) { var name string var env string - require.NoError(t, db.QueryRow("SELECT name, CAST(environment AS text) FROM vehicles WHERE environment = ?", "Air").Scan(&name, &env)) + require.NoError(t, db.QueryRow("SELECT name, environment FROM vehicles WHERE environment = ?", "Air").Scan(&name, &env)) require.Equal(t, "Aircraft", name) require.Equal(t, "Air", env) + + // enum list + _, err = db.Exec("CREATE TABLE all_enums (environments element[])") + require.NoError(t, err) + + _, err = db.Exec("INSERT INTO all_enums VALUES (['Air', 'Sea', 'Land'])") + require.NoError(t, err) + + var row Composite[[]string] + require.NoError(t, db.QueryRow("SELECT environments FROM all_enums").Scan(&row)) + require.Equal(t, 3, len(row.Get())) } func TestHugeInt(t *testing.T) { diff --git a/rows.go b/rows.go index fb025ea8..04841d65 100644 --- a/rows.go +++ b/rows.go @@ -147,6 +147,8 @@ func scan(vector C.duckdb_vector, rowIdx C.idx_t) (any, error) { return hugeIntToNative(hi), nil case C.DUCKDB_TYPE_VARCHAR: return scanString(vector, rowIdx), nil + case C.DUCKDB_TYPE_ENUM: + return scanENUM(ty, vector, rowIdx) case C.DUCKDB_TYPE_BLOB: return scanBlob(vector, rowIdx), nil case C.DUCKDB_TYPE_DECIMAL: @@ -211,6 +213,8 @@ func (r *rows) ColumnTypeScanType(index int) reflect.Type { return reflect.TypeOf(big.NewInt(0)) case C.DUCKDB_TYPE_VARCHAR: return reflect.TypeOf("") + case C.DUCKDB_TYPE_ENUM: + return reflect.TypeOf("") case C.DUCKDB_TYPE_BLOB: return reflect.TypeOf([]byte{}) case C.DUCKDB_TYPE_DECIMAL: @@ -401,6 +405,27 @@ func scanInterval(vector C.duckdb_vector, rowIdx C.idx_t) (Interval, error) { return data, nil } +func scanENUM(ty C.duckdb_logical_type, vector C.duckdb_vector, rowIdx C.idx_t) (string, error) { + var nativeValue uint64 + int_type := C.duckdb_enum_internal_type(ty) + switch int_type { + case C.DUCKDB_TYPE_UTINYINT: + nativeValue = uint64(get[uint8](vector, rowIdx)) + case C.DUCKDB_TYPE_USMALLINT: + nativeValue = uint64(get[uint16](vector, rowIdx)) + case C.DUCKDB_TYPE_UINTEGER: + nativeValue = uint64(get[uint32](vector, rowIdx)) + case C.DUCKDB_TYPE_UBIGINT: + nativeValue = get[uint64](vector, rowIdx) + default: + return "", errInvalidType + } + + val := C.duckdb_enum_dictionary_value(ty, (C.idx_t)(nativeValue)) + defer C.duckdb_free(unsafe.Pointer(val)) + return C.GoString(val), nil +} + var ( errInvalidType = errors.New("invalid data type") errMissingKeyOrValue = errors.New("missing key and/or value for map item") From 492d863c025c1751ad4787a76e68865f18bdc5cb Mon Sep 17 00:00:00 2001 From: Anshul Khandelwal Date: Tue, 7 Nov 2023 16:56:44 +0530 Subject: [PATCH 2/2] scan enum to string --- duckdb_test.go | 2 +- rows.go | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/duckdb_test.go b/duckdb_test.go index 5db2b399..72048339 100644 --- a/duckdb_test.go +++ b/duckdb_test.go @@ -511,7 +511,7 @@ func TestENUMs(t *testing.T) { var row Composite[[]string] require.NoError(t, db.QueryRow("SELECT environments FROM all_enums").Scan(&row)) - require.Equal(t, 3, len(row.Get())) + require.ElementsMatch(t, []string{"Air", "Land", "Sea"}, row.Get()) } func TestHugeInt(t *testing.T) { diff --git a/rows.go b/rows.go index 04841d65..55275a33 100644 --- a/rows.go +++ b/rows.go @@ -406,22 +406,22 @@ func scanInterval(vector C.duckdb_vector, rowIdx C.idx_t) (Interval, error) { } func scanENUM(ty C.duckdb_logical_type, vector C.duckdb_vector, rowIdx C.idx_t) (string, error) { - var nativeValue uint64 - int_type := C.duckdb_enum_internal_type(ty) - switch int_type { + var idx uint64 + internalType := C.duckdb_enum_internal_type(ty) + switch internalType { case C.DUCKDB_TYPE_UTINYINT: - nativeValue = uint64(get[uint8](vector, rowIdx)) + idx = uint64(get[uint8](vector, rowIdx)) case C.DUCKDB_TYPE_USMALLINT: - nativeValue = uint64(get[uint16](vector, rowIdx)) + idx = uint64(get[uint16](vector, rowIdx)) case C.DUCKDB_TYPE_UINTEGER: - nativeValue = uint64(get[uint32](vector, rowIdx)) + idx = uint64(get[uint32](vector, rowIdx)) case C.DUCKDB_TYPE_UBIGINT: - nativeValue = get[uint64](vector, rowIdx) + idx = get[uint64](vector, rowIdx) default: return "", errInvalidType } - val := C.duckdb_enum_dictionary_value(ty, (C.idx_t)(nativeValue)) + val := C.duckdb_enum_dictionary_value(ty, (C.idx_t)(idx)) defer C.duckdb_free(unsafe.Pointer(val)) return C.GoString(val), nil }