Appender fixes (#147)

* appender fixes

* appender fixes

* some refactor

* Fix indentation

* review comments - use separate type for uuid

* review comments - move UUID to types.go

---------

Co-authored-by: Marc Boeker <[email protected]>

Author: k-anshul

appender.go

@@ -67,25 +67,23 @@ func (a *Appender) Error() error {
 
 // Flush the appender to the underlying table and clear the internal cache.
 func (a *Appender) Flush() error {
-	// set the size of the current chunk to the current row
-	C.duckdb_data_chunk_set_size(a.chunks[a.currentChunkIdx], C.uint64_t(a.currentRow))
+	if a.currentChunkIdx == 0 && a.currentRow == 0 {
+		return nil
+	}
 
-	// append all chunks to the appender and destroy them
-	var state C.duckdb_state
-	for i, chunk := range a.chunks {
-		state = C.duckdb_append_data_chunk(*a.appender, chunk)
-		if state == C.DuckDBError {
-			dbErr := C.GoString(C.duckdb_appender_error(*a.appender))
-			return fmt.Errorf("duckdb error appending chunk %d of %d: %s", i+1, a.currentChunkIdx+1, dbErr)
-		}
-		C.duckdb_destroy_data_chunk(&chunk)
+	err := a.appendChunks()
+	if err != nil {
+		return err
 	}
 
 	if state := C.duckdb_appender_flush(*a.appender); state == C.DuckDBError {
 		dbErr := C.GoString(C.duckdb_appender_error(*a.appender))
 		return errors.New(dbErr)
 	}
 
+	a.currentRow = 0
+	a.currentChunkIdx = 0
+	a.chunks = a.chunks[:0]
 	return nil
 }
 
@@ -96,6 +94,12 @@ func (a *Appender) Close() error {
 	}
 
 	a.closed = true
+	// append chunks if not already done via flush
+	if a.currentChunkIdx != 0 || a.currentRow != 0 {
+		if err := a.appendChunks(); err != nil {
+			return err
+		}
+	}
 
 	if state := C.duckdb_appender_destroy(a.appender); state == C.DuckDBError {
 		dbErr := C.GoString(C.duckdb_appender_error(*a.appender))
@@ -161,6 +165,8 @@ func (a *Appender) initializeChunkTypes(args []driver.Value) {
 			tmpChunkTypes[i] = C.duckdb_create_logical_type(C.DUCKDB_TYPE_BOOLEAN)
 		case []byte:
 			tmpChunkTypes[i] = C.duckdb_create_logical_type(C.DUCKDB_TYPE_BLOB)
+		case UUID:
+			tmpChunkTypes[i] = C.duckdb_create_logical_type(C.DUCKDB_TYPE_UUID)
 		case string:
 			tmpChunkTypes[i] = C.duckdb_create_logical_type(C.DUCKDB_TYPE_VARCHAR)
 		case time.Time:
@@ -236,6 +242,8 @@ func (a *Appender) appendRowArray(args []driver.Value) error {
 			set[bool](a.chunkVectors[i], a.currentRow, v)
 		case []byte:
 			set[[]byte](a.chunkVectors[i], a.currentRow, v)
+		case UUID:
+			set[C.duckdb_hugeint](a.chunkVectors[i], a.currentRow, uuidToHugeInt(v))
 		case string:
 			str := C.CString(v)
 			C.duckdb_vector_assign_string_element(a.chunkVectors[i], C.uint64_t(a.currentRow), str)
@@ -253,4 +261,21 @@ func (a *Appender) appendRowArray(args []driver.Value) error {
 	return nil
 }
 
+func (a *Appender) appendChunks() error {
+	// set the size of the current chunk to the current row
+	C.duckdb_data_chunk_set_size(a.chunks[a.currentChunkIdx], C.uint64_t(a.currentRow))
+
+	// append all chunks to the appender and destroy them
+	var state C.duckdb_state
+	for i, chunk := range a.chunks {
+		state = C.duckdb_append_data_chunk(*a.appender, chunk)
+		if state == C.DuckDBError {
+			dbErr := C.GoString(C.duckdb_appender_error(*a.appender))
+			return fmt.Errorf("duckdb error appending chunk %d of %d: %s", i+1, a.currentChunkIdx+1, dbErr)
+		}
+		C.duckdb_destroy_data_chunk(&chunk)
+	}
+	return nil
+}
+
 var errCouldNotAppend = errors.New("could not append parameter")

appender_test.go

@@ -7,26 +7,28 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 )
 
 const (
 	testAppenderTableDDL = `
   CREATE TABLE test(
-    id BIGINT,
-    uint8 UTINYINT,
-    int8 TINYINT,
-    uint16 USMALLINT,
-    int16 SMALLINT,
+		id BIGINT,
+		uuid UUID,
+		uint8 UTINYINT,
+		int8 TINYINT,
+		uint16 USMALLINT,
+		int16 SMALLINT,
 		uint32 UINTEGER,
-    int32 INTEGER,
-    uint64 UBIGINT,
-    int64 BIGINT,
+		int32 INTEGER,
+		uint64 UBIGINT,
+		int64 BIGINT,
 		timestamp TIMESTAMP,
-    float REAL,
-    double DOUBLE,
-    string VARCHAR,
-    bool BOOLEAN
+		float REAL,
+		double DOUBLE,
+		string VARCHAR,
+		bool BOOLEAN
   )`
 )
 
@@ -66,6 +68,7 @@ func TestAppender(t *testing.T) {
 
 	type dataRow struct {
 		ID        int
+		UUID      UUID
 		UInt8     uint8
 		Int8      int8
 		UInt16    uint16
@@ -86,9 +89,16 @@ func TestAppender(t *testing.T) {
 		if u64 > 9223372036854775807 {
 			u64 = 9223372036854775807
 		}
+
+		b, err := uuid.New().MarshalBinary()
+		require.NoError(t, err)
+		var uuidBytes [16]byte
+		copy(uuidBytes[:], b)
+
 		return dataRow{
 			ID:        i,
 			UInt8:     uint8(randInt(0, 255)),
+			UUID:      UUID(uuidBytes),
 			Int8:      int8(randInt(-128, 127)),
 			UInt16:    uint16(randInt(0, 65535)),
 			Int16:     int16(randInt(-32768, 32767)),
@@ -114,11 +124,15 @@ func TestAppender(t *testing.T) {
 
 	appender, err := NewAppenderFromConn(conn, "", "test")
 	require.NoError(t, err)
-	defer appender.Close()
 
-	for _, row := range rows {
+	for i, row := range rows {
+		if i%1024 == 0 {
+			err = appender.Flush()
+			require.NoError(t, err)
+		}
 		err := appender.AppendRow(
 			row.ID,
+			row.UUID,
 			row.UInt8,
 			row.Int8,
 			row.UInt16,
@@ -135,12 +149,13 @@ func TestAppender(t *testing.T) {
 		)
 		require.NoError(t, err)
 	}
-	err = appender.Flush()
+	err = appender.Close()
 	require.NoError(t, err)
 
 	res, err := db.QueryContext(
 		context.Background(), `
 			SELECT  id,
+							uuid,
 							uint8,
 							int8,
 							uint16,
@@ -160,10 +175,12 @@ func TestAppender(t *testing.T) {
 	defer res.Close()
 
 	i := 0
+	var scannedUUID []byte
 	for res.Next() {
 		r := dataRow{}
 		err := res.Scan(
 			&r.ID,
+			&scannedUUID,
 			&r.UInt8,
 			&r.Int8,
 			&r.UInt16,
@@ -179,6 +196,7 @@ func TestAppender(t *testing.T) {
 			&r.Bool,
 		)
 		require.NoError(t, err)
+		copy(r.UUID[:], scannedUUID)
 		require.Equal(t, rows[i], r)
 		i++
 	}

types.go

@@ -13,6 +13,8 @@ import (
 	"github.com/mitchellh/mapstructure"
 )
 
+type UUID [16]byte
+
 // duckdb_hugeint is composed of (lower, upper) components.
 // The value is computed as: upper * 2^64 + lower
 
@@ -24,6 +26,16 @@ func hugeIntToUUID(hi C.duckdb_hugeint) []byte {
 	return uuid[:]
 }
 
+func uuidToHugeInt(uuid UUID) C.duckdb_hugeint {
+	var dt C.duckdb_hugeint
+	upper := binary.BigEndian.Uint64(uuid[:8])
+	// flip the sign bit
+	upper = upper ^ (1 << 63)
+	dt.upper = C.int64_t(upper)
+	dt.lower = C.uint64_t(binary.BigEndian.Uint64(uuid[8:]))
+	return dt
+}
+
 func hugeIntToNative(hi C.duckdb_hugeint) *big.Int {
 	i := big.NewInt(int64(hi.upper))
 	i.Lsh(i, 64)