Merge pull request #254 from taniabogatsch/caching

Cache the validity mask, data pointers, and the chunk size

appender.go

@@ -139,7 +139,7 @@ func (a *Appender) AppendRow(args ...driver.Value) error {
 
 func (a *Appender) addDataChunk() error {
 	var chunk DataChunk
-	if err := chunk.initFromTypes(a.ptr, a.types); err != nil {
+	if err := chunk.initFromTypes(a.ptr, a.types, true); err != nil {
 		return err
 	}
 	a.chunks = append(a.chunks, chunk)

data_chunk.go

@@ -18,6 +18,8 @@ type DataChunk struct {
 	columns []vector
 	// columnNames holds the column names, if known.
 	columnNames []string
+	// size caches the size after initialization.
+	size int
 }
 
 // GetDataChunkCapacity returns the capacity of a data chunk.
@@ -27,7 +29,8 @@ func GetDataChunkCapacity() int {
 
 // GetSize returns the internal size of the data chunk.
 func (chunk *DataChunk) GetSize() int {
-	return int(C.duckdb_data_chunk_get_size(chunk.data))
+	chunk.size = int(C.duckdb_data_chunk_get_size(chunk.data))
+	return chunk.size
 }
 
 // SetSize sets the internal size of the data chunk. Cannot exceed GetCapacity().
@@ -71,7 +74,7 @@ func (chunk *DataChunk) SetValue(colIdx int, rowIdx int, val any) error {
 	return nil
 }
 
-func (chunk *DataChunk) initFromTypes(ptr unsafe.Pointer, types []C.duckdb_logical_type) error {
+func (chunk *DataChunk) initFromTypes(ptr unsafe.Pointer, types []C.duckdb_logical_type, writable bool) error {
 	// NOTE: initFromTypes does not initialize the column names.
 	columnCount := len(types)
 
@@ -93,14 +96,13 @@ func (chunk *DataChunk) initFromTypes(ptr unsafe.Pointer, types []C.duckdb_logic
 
 	// Initialize the vectors and their child vectors.
 	for i := 0; i < columnCount; i++ {
-		duckdbVector := C.duckdb_data_chunk_get_vector(chunk.data, C.idx_t(i))
-		chunk.columns[i].duckdbVector = duckdbVector
-		chunk.columns[i].getChildVectors(duckdbVector)
+		v := C.duckdb_data_chunk_get_vector(chunk.data, C.idx_t(i))
+		chunk.columns[i].initVectors(v, writable)
 	}
 	return nil
 }
 
-func (chunk *DataChunk) initFromDuckDataChunk(data C.duckdb_data_chunk) error {
+func (chunk *DataChunk) initFromDuckDataChunk(data C.duckdb_data_chunk, writable bool) error {
 	columnCount := int(C.duckdb_data_chunk_get_column_count(data))
 	chunk.columns = make([]vector, columnCount)
 	chunk.data = data
@@ -117,10 +119,11 @@ func (chunk *DataChunk) initFromDuckDataChunk(data C.duckdb_data_chunk) error {
 			break
 		}
 
-		// Initialize the vectors and their child vectors.
-		chunk.columns[i].duckdbVector = duckdbVector
-		chunk.columns[i].getChildVectors(duckdbVector)
+		// Initialize the vector and its child vectors.
+		chunk.columns[i].initVectors(duckdbVector, writable)
 	}
+
+	chunk.GetSize()
 	return err
 }
 

errors_test.go

@@ -407,6 +407,7 @@ type wrappedDuckDBError struct {
 func (w *wrappedDuckDBError) Error() string {
 	return w.e.Error()
 }
+
 func (w *wrappedDuckDBError) Unwrap() error {
 	return w.e
 }

rows.go

@@ -55,13 +55,13 @@ func (r *rows) Columns() []string {
 }
 
 func (r *rows) Next(dst []driver.Value) error {
-	for r.rowCount == r.chunk.GetSize() {
+	for r.rowCount == r.chunk.size {
 		r.chunk.close()
 		if r.chunkIdx == r.chunkCount {
 			return io.EOF
 		}
 		data := C.duckdb_result_get_chunk(r.res, r.chunkIdx)
-		if err := r.chunk.initFromDuckDataChunk(data); err != nil {
+		if err := r.chunk.initFromDuckDataChunk(data, false); err != nil {
 			return getError(err, nil)
 		}
 

vector.go

@@ -18,6 +18,10 @@ import (
 type vector struct {
 	// The underlying DuckDB vector.
 	duckdbVector C.duckdb_vector
+	// The underlying data ptr.
+	ptr unsafe.Pointer
+	// The vector's validity mask.
+	mask *C.uint64_t
 	// A callback function to get a value from this vector.
 	getFn fnGetVectorValue
 	// A callback function to write to this vector.
@@ -311,19 +315,27 @@ func (vec *vector) init(logicalType C.duckdb_logical_type, colIdx int) error {
 	return nil
 }
 
-func (vec *vector) getChildVectors(vector C.duckdb_vector) {
+func (vec *vector) initVectors(v C.duckdb_vector, writable bool) {
+	vec.duckdbVector = v
+	vec.ptr = C.duckdb_vector_get_data(v)
+	if writable {
+		C.duckdb_vector_ensure_validity_writable(v)
+	}
+	vec.mask = C.duckdb_vector_get_validity(v)
+	vec.getChildVectors(v, writable)
+}
+
+func (vec *vector) getChildVectors(v C.duckdb_vector, writable bool) {
 	switch vec.duckdbType {
 
 	case C.DUCKDB_TYPE_LIST, C.DUCKDB_TYPE_MAP:
-		child := C.duckdb_list_vector_get_child(vector)
-		vec.childVectors[0].duckdbVector = child
-		vec.childVectors[0].getChildVectors(child)
+		child := C.duckdb_list_vector_get_child(v)
+		vec.childVectors[0].initVectors(child, writable)
 
 	case C.DUCKDB_TYPE_STRUCT:
 		for i := 0; i < len(vec.childVectors); i++ {
-			child := C.duckdb_struct_vector_get_child(vector, C.idx_t(i))
-			vec.childVectors[i].duckdbVector = child
-			vec.childVectors[i].getChildVectors(child)
+			child := C.duckdb_struct_vector_get_child(v, C.idx_t(i))
+			vec.childVectors[i].initVectors(child, writable)
 		}
 	}
 }

vector_getters.go

@@ -16,13 +16,20 @@ import (
 type fnGetVectorValue func(vec *vector, rowIdx C.idx_t) any
 
 func (vec *vector) getNull(rowIdx C.idx_t) bool {
-	mask := C.duckdb_vector_get_validity(vec.duckdbVector)
-	return !bool(C.duckdb_validity_row_is_valid(mask, rowIdx))
+	mask := unsafe.Pointer(vec.mask)
+	if mask == nil {
+		return false
+	}
+
+	entryIdx := rowIdx / 64
+	idxInEntry := rowIdx % 64
+	maskPtr := (*[1 << 31]C.uint64_t)(mask)
+	isValid := maskPtr[entryIdx] & (C.uint64_t(1) << idxInEntry)
+	return uint64(isValid) == 0
 }
 
 func getPrimitive[T any](vec *vector, rowIdx C.idx_t) T {
-	ptr := C.duckdb_vector_get_data(vec.duckdbVector)
-	xs := (*[1 << 31]T)(ptr)
+	xs := (*[1 << 31]T)(vec.ptr)
 	return xs[rowIdx]
 }
 

vector_setters.go

@@ -19,9 +19,7 @@ const secondsPerDay = 24 * 60 * 60
 type fnSetVectorValue func(vec *vector, rowIdx C.idx_t, val any)
 
 func (vec *vector) setNull(rowIdx C.idx_t) {
-	C.duckdb_vector_ensure_validity_writable(vec.duckdbVector)
-	mask := C.duckdb_vector_get_validity(vec.duckdbVector)
-	C.duckdb_validity_set_row_invalid(mask, rowIdx)
+	C.duckdb_validity_set_row_invalid(vec.mask, rowIdx)
 
 	if vec.duckdbType == C.DUCKDB_TYPE_STRUCT {
 		for i := 0; i < len(vec.childVectors); i++ {
@@ -31,8 +29,7 @@ func (vec *vector) setNull(rowIdx C.idx_t) {
 }
 
 func setPrimitive[T any](vec *vector, rowIdx C.idx_t, v T) {
-	ptr := C.duckdb_vector_get_data(vec.duckdbVector)
-	xs := (*[1 << 31]T)(ptr)
+	xs := (*[1 << 31]T)(vec.ptr)
 	xs[rowIdx] = v
 }