Merge branch 'main' of github.com:ClickHouse/clickhouse-go into reduce-memory-usage-compressor

Author: pablomatiasgomez

clickhouse_options.go

@@ -42,6 +42,8 @@ func (c CompressionMethod) String() string {
 		return "zstd"
 	case CompressionLZ4:
 		return "lz4"
+	case CompressionLZ4HC:
+		return "lz4hc"
 	case CompressionGZIP:
 		return "gzip"
 	case CompressionDeflate:
@@ -56,6 +58,7 @@ func (c CompressionMethod) String() string {
 const (
 	CompressionNone    = CompressionMethod(compress.None)
 	CompressionLZ4     = CompressionMethod(compress.LZ4)
+	CompressionLZ4HC   = CompressionMethod(compress.LZ4HC)
 	CompressionZSTD    = CompressionMethod(compress.ZSTD)
 	CompressionGZIP    = CompressionMethod(0x95)
 	CompressionDeflate = CompressionMethod(0x96)
@@ -66,6 +69,7 @@ var compressionMap = map[string]CompressionMethod{
 	"none":    CompressionNone,
 	"zstd":    CompressionZSTD,
 	"lz4":     CompressionLZ4,
+	"lz4hc":   CompressionLZ4HC,
 	"gzip":    CompressionGZIP,
 	"deflate": CompressionDeflate,
 	"br":      CompressionBrotli,
@@ -79,7 +83,7 @@ type Auth struct { // has_control_character
 
 type Compression struct {
 	Method CompressionMethod
-	// this only applies to zlib and brotli compression algorithms
+	// this only applies to lz4, lz4hc, zlib, and brotli compression algorithms
 	Level int
 }
 

conn.go

@@ -73,14 +73,22 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 		}
 	}
 
-	compression := CompressionNone
+	var (
+		compression CompressionMethod
+		compressor  *compress.Writer
+	)
 	if opt.Compression != nil {
 		switch opt.Compression.Method {
-		case CompressionLZ4, CompressionZSTD, CompressionNone:
+		case CompressionLZ4, CompressionLZ4HC, CompressionZSTD, CompressionNone:
 			compression = opt.Compression.Method
 		default:
 			return nil, fmt.Errorf("unsupported compression method for native protocol")
 		}
+
+		compressor = compress.NewWriter(compress.Level(opt.Compression.Level), compress.Method(opt.Compression.Method))
+	} else {
+		compression = CompressionNone
+		compressor = compress.NewWriter(compress.LevelZero, compress.None)
 	}
 
 	var (
@@ -95,7 +103,7 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 			structMap:            &structMap{},
 			compression:          compression,
 			connectedAt:          time.Now(),
-			compressor:           compress.NewWriterWithMethods(0, compress.Method(compression)),
+			compressor:           compressor,
 			readTimeout:          opt.ReadTimeout,
 			blockBufferSize:      opt.BlockBufferSize,
 			maxCompressionBuffer: opt.MaxCompressionBuffer,
@@ -247,7 +255,7 @@ func (c *connect) exception() error {
 func (c *connect) compressBuffer(start int) error {
 	if c.compression != CompressionNone && len(c.buffer.Buf) > 0 {
 		data := c.buffer.Buf[start:]
-		if err := c.compressor.Compress(compress.Method(c.compression), data); err != nil {
+		if err := c.compressor.Compress(data); err != nil {
 			return errors.Wrap(err, "compress")
 		}
 		c.buffer.Buf = append(c.buffer.Buf[:start], c.compressor.Data...)

conn_http.go

@@ -230,7 +230,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		url:             u,
 		buffer:          new(chproto.Buffer),
 		compression:     opt.Compression.Method,
-		blockCompressor: compress.NewWriterWithMethods(0, compress.Method(opt.Compression.Method)),
+		blockCompressor: compress.NewWriter(compress.Level(opt.Compression.Level), compress.Method(opt.Compression.Method)),
 		compressionPool: compressionPool,
 		blockBufferSize: opt.BlockBufferSize,
 		headers:         headers,
@@ -256,7 +256,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		url:             u,
 		buffer:          new(chproto.Buffer),
 		compression:     opt.Compression.Method,
-		blockCompressor: compress.NewWriterWithMethods(0, compress.Method(opt.Compression.Method)),
+		blockCompressor: compress.NewWriter(compress.Level(opt.Compression.Level), compress.Method(opt.Compression.Method)),
 		compressionPool: compressionPool,
 		location:        location,
 		blockBufferSize: opt.BlockBufferSize,
@@ -377,7 +377,7 @@ func (h *httpConnect) writeData(block *proto.Block) error {
 	if h.compression == CompressionLZ4 || h.compression == CompressionZSTD {
 		// Performing compression. Supported and requires
 		data := h.buffer.Buf[start:]
-		if err := h.blockCompressor.Compress(compress.Method(h.compression), data); err != nil {
+		if err := h.blockCompressor.Compress(data); err != nil {
 			return errors.Wrap(err, "compress")
 		}
 		h.buffer.Buf = append(h.buffer.Buf[:start], h.blockCompressor.Data...)

tests/compression_test.go

@@ -19,27 +19,33 @@ package tests
 
 import (
 	"context"
+	"github.com/ClickHouse/ch-go/compress"
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"testing"
 )
 
 func TestZSTDCompression(t *testing.T) {
-	CompressionTest(t, clickhouse.CompressionZSTD)
+	CompressionTest(t, compress.LevelZero, clickhouse.CompressionZSTD)
 }
 
 func TestLZ4Compression(t *testing.T) {
-	CompressionTest(t, clickhouse.CompressionLZ4)
+	CompressionTest(t, compress.Level(3), clickhouse.CompressionLZ4)
+}
+
+func TestLZ4HCCompression(t *testing.T) {
+	CompressionTest(t, compress.LevelLZ4HCDefault, clickhouse.CompressionLZ4HC)
 }
 
 func TestNoCompression(t *testing.T) {
-	CompressionTest(t, clickhouse.CompressionNone)
+	CompressionTest(t, compress.LevelZero, clickhouse.CompressionNone)
 }
 
-func CompressionTest(t *testing.T, method clickhouse.CompressionMethod) {
+func CompressionTest(t *testing.T, level compress.Level, method clickhouse.CompressionMethod) {
 	conn, err := GetNativeConnection(nil, nil, &clickhouse.Compression{
 		Method: method,
+		Level:  int(level),
 	})
 	ctx := context.Background()
 	require.NoError(t, err)
@@ -57,7 +63,7 @@ func CompressionTest(t *testing.T, method clickhouse.CompressionMethod) {
 	var (
 		col1Data = []string{"A", "b", "c"}
 	)
-	for i := 0; i < 10; i++ {
+	for i := 0; i < 100; i++ {
 		require.NoError(t, batch.Append(col1Data))
 	}
 	require.NoError(t, batch.Send())