Merge pull request #236 from taniabogatsch/appender-fix

Fix appending multiple data chunks

appender.go

@@ -153,10 +153,11 @@ func (a *Appender) appendRowSlice(args []driver.Value) error {
 	}
 
 	// Create a new data chunk if the current chunk is full.
-	if C.idx_t(a.rowCount) == C.duckdb_vector_size() || len(a.chunks) == 0 {
+	if a.rowCount == GetDataChunkCapacity() || len(a.chunks) == 0 {
 		if err := a.addDataChunk(); err != nil {
 			return err
 		}
+		a.rowCount = 0
 	}
 
 	// Set all values.
@@ -177,9 +178,8 @@ func (a *Appender) appendDataChunks() error {
 	var err error
 
 	for i, chunk := range a.chunks {
-
 		// All data chunks except the last are at maximum capacity.
-		size := chunk.GetCapacity()
+		size := GetDataChunkCapacity()
 		if i == len(a.chunks)-1 {
 			size = a.rowCount
 		}
@@ -194,16 +194,13 @@ func (a *Appender) appendDataChunks() error {
 		}
 	}
 
-	a.closeDataChunks()
-	a.rowCount = 0
-	return err
-}
-
-func (a *Appender) closeDataChunks() {
 	for _, chunk := range a.chunks {
 		chunk.close()
 	}
+
 	a.chunks = a.chunks[:0]
+	a.rowCount = 0
+	return err
 }
 
 func mallocTypeSlice(count int) (unsafe.Pointer, []C.duckdb_logical_type) {

appender_test.go

@@ -15,8 +15,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const numAppenderTestRows = 10000
-
 type simpleStruct struct {
 	A int32
 	B string
@@ -169,6 +167,8 @@ func TestAppenderPrimitive(t *testing.T) {
 			bool BOOLEAN
 	  	)`)
 
+	// Test appending a few data chunks.
+	rowCount := GetDataChunkCapacity() * 5
 	type row struct {
 		ID          int64
 		UInt8       uint8
@@ -198,8 +198,8 @@ func TestAppenderPrimitive(t *testing.T) {
 	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++ {
+	rowsToAppend := make([]row, rowCount)
+	for i := 0; i < rowCount; 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).
@@ -248,10 +248,6 @@ func TestAppenderPrimitive(t *testing.T) {
 			rowsToAppend[i].String,
 			rowsToAppend[i].Bool,
 		))
-
-		if i%1000 == 0 {
-			require.NoError(t, a.Flush())
-		}
 	}
 	require.NoError(t, a.Flush())
 
@@ -291,7 +287,7 @@ func TestAppenderPrimitive(t *testing.T) {
 		i++
 	}
 
-	require.Equal(t, numAppenderTestRows, i)
+	require.Equal(t, rowCount, i)
 	require.NoError(t, res.Close())
 	cleanupAppender(t, c, con, a)
 }

data_chunk.go

@@ -18,14 +18,14 @@ type DataChunk struct {
 	columns []vector
 }
 
-// GetCapacity returns the capacity of a data chunk.
-func (chunk *DataChunk) GetCapacity() int {
+// GetDataChunkCapacity returns the capacity of a data chunk.
+func GetDataChunkCapacity() int {
 	return int(C.duckdb_vector_size())
 }
 
 // SetSize sets the internal size of the data chunk. Cannot exceed GetCapacity().
 func (chunk *DataChunk) SetSize(size int) error {
-	if size > chunk.GetCapacity() {
+	if size > GetDataChunkCapacity() {
 		return getError(errAPI, errVectorSize)
 	}
 	C.duckdb_data_chunk_set_size(chunk.data, C.idx_t(size))