diff --git a/appender.go b/appender.go index c035d5ea..f7e8d184 100644 --- a/appender.go +++ b/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) { diff --git a/appender_test.go b/appender_test.go index c8bbe14a..01ae8367 100644 --- a/appender_test.go +++ b/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) } diff --git a/data_chunk.go b/data_chunk.go index 966ac329..e372c329 100644 --- a/data_chunk.go +++ b/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))