write blocks concurrently with columnar writer

this changes the interface to enable accessing clickhouse connection and block interfaces directly, so the caller can create and write additional blocks concurrently

Author: vavrusa

bootstrap.go

@@ -37,6 +37,10 @@ func (d *bootstrap) Open(dsn string) (driver.Conn, error) {
 }
 
 func Open(dsn string) (driver.Conn, error) {
+	return open(dsn)
+}
+
+func open(dsn string) (*clickhouse, error) {
 	url, err := url.Parse(dsn)
 	if err != nil {
 		return nil, err

clickhouse.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"regexp"
+	"sync"
 	"time"
 
 	"github.com/kshvakov/clickhouse/lib/binary"
@@ -27,6 +28,7 @@ var (
 type logger func(format string, v ...interface{})
 
 type clickhouse struct {
+	sync.Mutex
 	data.ServerInfo
 	data.ClientInfo
 	logf          logger
@@ -99,7 +101,7 @@ type txOptions struct {
 	ReadOnly  bool
 }
 
-func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (driver.Tx, error) {
+func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
 	ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
 	switch {
 	case ch.inTransaction:

clickhouse_columnar_test.go

@@ -0,0 +1,111 @@
+package clickhouse_test
+
+import (
+	"database/sql/driver"
+	"testing"
+	"time"
+
+	"github.com/kshvakov/clickhouse"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_ColumnarInsert(t *testing.T) {
+	const (
+		ddl = `
+			CREATE TABLE clickhouse_test_columnar_insert (
+				uint8  UInt8,
+				uint16 UInt16,
+				uint32 UInt32,
+				uint64 UInt64,
+				float32 Float32,
+				float64 Float64,
+				string  String,
+				fString FixedString(2),
+				date    Date,
+				datetime DateTime,
+				enum8    Enum8 ('a' = 1, 'b' = 2),
+				enum16   Enum16('c' = 1, 'd' = 2)
+			) Engine=Memory
+		`
+		dml = `
+			INSERT INTO clickhouse_test_columnar_insert (
+				uint8,
+				uint16,
+				uint32,
+				uint64,
+				float32,
+				float64,
+				string,
+				fString,
+				date,
+				datetime,
+				enum8,
+				enum16
+			) VALUES (
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?,
+				?
+			)
+		`
+	)
+	if connect, err := clickhouse.OpenDirect("tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) {
+		{
+			connect.Begin()
+			stmt, _ := connect.Prepare("DROP TABLE clickhouse_test_columnar_insert")
+			stmt.Exec([]driver.Value{})
+			connect.Commit()
+		}
+		{
+			if _, err := connect.Begin(); assert.NoError(t, err) {
+				if stmt, err := connect.Prepare(ddl); assert.NoError(t, err) {
+					if _, err := stmt.Exec([]driver.Value{}); assert.NoError(t, err) {
+						assert.NoError(t, connect.Commit())
+					}
+				}
+			}
+		}
+		{
+			if _, err := connect.Begin(); assert.NoError(t, err) {
+				if _, err := connect.Prepare(dml); assert.NoError(t, err) {
+					block, err := connect.Block()
+					assert.NoError(t, err)
+					block.Reserve()
+					block.NumRows = 100
+
+					for i := 0; i < 100; i++ {
+						block.WriteUInt8(0, uint8(i))
+						block.WriteUInt16(1, uint16(i))
+						block.WriteUInt32(2, uint32(i))
+						block.WriteUInt64(3, uint64(i))
+
+						block.WriteFloat32(4, float32(i))
+						block.WriteFloat64(5, float64(i))
+
+						block.WriteString(6, "string")
+						block.WriteFixedString(7, []byte("CH"))
+						block.WriteDate(8, time.Now())
+						block.WriteDateTime(9, time.Now())
+
+						block.WriteUInt8(10, 1)
+						block.WriteUInt16(11, 2)
+
+						if !assert.NoError(t, err) {
+							return
+						}
+					}
+
+					assert.NoError(t, connect.Commit())
+				}
+			}
+		}
+	}
+}

clickhouse_write_block.go

@@ -6,6 +6,8 @@ import (
 )
 
 func (ch *clickhouse) writeBlock(block *data.Block) error {
+	ch.Lock()
+	defer ch.Unlock()
 	if err := ch.encoder.Uvarint(protocol.ClientData); err != nil {
 		return err
 	}

examples/columnar.go

@@ -3,18 +3,20 @@ package main
 import (
 	"database/sql/driver"
 	"log"
+	"sync"
 	"time"
 
 	"github.com/kshvakov/clickhouse"
+	data "github.com/kshvakov/clickhouse/lib/data"
 )
 
 func main() {
-	connect, err := clickhouse.Open("tcp://127.0.0.1:9000?username=&debug=true")
+	connect, err := clickhouse.OpenDirect("tcp://127.0.0.1:9000?username=&debug=true")
 	if err != nil {
 		log.Fatal(err)
 	}
 	{
-		tx, _ := connect.Begin()
+		connect.Begin()
 		stmt, _ := connect.Prepare(`
 			CREATE TABLE IF NOT EXISTS example (
 				os_id        UInt8,
@@ -27,48 +29,92 @@ func main() {
 		if _, err := stmt.Exec([]driver.Value{}); err != nil {
 			log.Fatal(err)
 		}
-		tx.Commit()
+
+		if err := connect.Commit(); err != nil {
+			log.Fatal(err)
+		}
 	}
 	{
-		tx, _ := connect.Begin()
-		stmt, _ := connect.Prepare("INSERT INTO example (os_id, action_day, tags, categories) VALUES (?, ?, ?, ?)")
-		cstmt, ok := stmt.(clickhouse.ColumnarStatement)
-		if !ok {
-			log.Fatal("Column writer is not supported")
-		}
+		connect.Begin()
+		connect.Prepare("INSERT INTO example (os_id, action_day, tags, categories) VALUES (?, ?, ?, ?)")
 
-		w := cstmt.ColumnWriter()
-		for i := 0; i < 1000; i++ {
-			w.WriteUInt8(0, uint8(10+i))
+		block, err := connect.Block()
+		if err != nil {
+			log.Fatal(err)
 		}
 
-		for i := 0; i < 1000; i++ {
-			w.WriteDate(1, time.Now())
-		}
+		blocks := []*data.Block{block, block.Copy()}
+
+		var wg sync.WaitGroup
+		wg.Add(len(blocks))
 
-		for i := 0; i < 1000; i++ {
-			w.WriteArray(2, clickhouse.Array([]string{"A", "B", "C"}))
+		for i := range blocks {
+			b := blocks[i]
+			go func() {
+				defer wg.Done()
+				writeBatch(b, 1000)
+				if err := connect.WriteBlock(b); err != nil {
+					log.Fatal(err)
+				}
+			}()
 		}
 
-		for i := 0; i < 1000; i++ {
-			w.WriteArray(3, clickhouse.Array([]uint8{1, 2, 3, 4, 5}))
+		wg.Wait()
+
+		if err := connect.Commit(); err != nil {
+			log.Fatal(err)
 		}
+	}
+	{
+		connect.Begin()
+		stmt, _ := connect.Prepare(`SELECT count() FROM example`)
 
-		if err := cstmt.ColumnWriterEnd(1000); err != nil {
+		rows, err := stmt.Query([]driver.Value{})
+		if err != nil {
 			log.Fatal(err)
 		}
 
-		if err := tx.Commit(); err != nil {
+		columns := rows.Columns()
+		row := make([]driver.Value, 1)
+		for rows.Next(row) == nil {
+			for i, c := range columns {
+				log.Print(c, " : ", row[i])
+			}
+		}
+
+		if err := connect.Commit(); err != nil {
 			log.Fatal(err)
 		}
 	}
 	{
-		tx, _ := connect.Begin()
+		connect.Begin()
 		stmt, _ := connect.Prepare(`DROP TABLE example`)
-
 		if _, err := stmt.Exec([]driver.Value{}); err != nil {
 			log.Fatal(err)
 		}
-		tx.Commit()
+		if err := connect.Commit(); err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
+func writeBatch(block *data.Block, n int) {
+	block.Reserve()
+	block.NumRows += uint64(n)
+
+	for i := 0; i < n; i++ {
+		block.WriteUInt8(0, uint8(10+i))
+	}
+
+	for i := 0; i < n; i++ {
+		block.WriteDate(1, time.Now())
+	}
+
+	for i := 0; i < n; i++ {
+		block.WriteArray(2, clickhouse.Array([]string{"A", "B", "C"}))
+	}
+
+	for i := 0; i < n; i++ {
+		block.WriteArray(3, clickhouse.Array([]uint8{1, 2, 3, 4, 5}))
 	}
 }

lib/data/block.go

@@ -22,6 +22,14 @@ type Block struct {
 	info       blockInfo
 }
 
+func (block *Block) Copy() *Block {
+	return &Block{
+		Columns:    block.Columns,
+		NumColumns: block.NumColumns,
+		info:       block.info,
+	}
+}
+
 func (block *Block) ColumnNames() []string {
 	names := make([]string, 0, len(block.Columns))
 	for _, column := range block.Columns {

write_column.go

@@ -1,30 +1,26 @@
 package clickhouse
 
-// @todo: restore block_size functionality
-
 import (
+	"database/sql"
+	"database/sql/driver"
 	"time"
 
+	"github.com/kshvakov/clickhouse/lib/data"
 	"github.com/kshvakov/clickhouse/lib/types"
 )
 
-// Statement supporting columnar writer
-type ColumnarStatement interface {
-	ColumnWriter() ColumnWriter
-	ColumnWriterEnd(rows uint64) error
-}
-
-func (stmt *stmt) ColumnWriter() ColumnWriter {
-	stmt.ch.block.Reserve()
-	return stmt.ch.block
+// Interface for Clickhouse driver
+type Clickhouse interface {
+	Block() (*data.Block, error)
+	Prepare(query string) (driver.Stmt, error)
+	Begin() (driver.Tx, error)
+	Commit() error
+	Rollback() error
+	Close() error
+	WriteBlock(block *data.Block) error
 }
 
-func (stmt *stmt) ColumnWriterEnd(rows uint64) error {
-	stmt.ch.block.NumRows += rows
-	return nil
-}
-
-// Interface for block writer allowing writes to individual columns
+// Interface for Block allowing writes to individual columns
 type ColumnWriter interface {
 	WriteDate(c int, v time.Time) error
 	WriteDateTime(c int, v time.Time) error
@@ -39,3 +35,21 @@ type ColumnWriter interface {
 	WriteString(c int, v string) error
 	WriteFixedString(c int, v []byte) error
 }
+
+func OpenDirect(dsn string) (Clickhouse, error) {
+	return open(dsn)
+}
+
+func (ch *clickhouse) Block() (*data.Block, error) {
+	if ch.block == nil {
+		return nil, sql.ErrTxDone
+	}
+	return ch.block, nil
+}
+
+func (ch *clickhouse) WriteBlock(block *data.Block) error {
+	if block == nil {
+		return sql.ErrTxDone
+	}
+	return ch.writeBlock(block)
+}