1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
|
// +build !clz4
package binary
import (
"encoding/binary"
"io"
"github.com/ClickHouse/clickhouse-go/lib/cityhash102"
"github.com/ClickHouse/clickhouse-go/lib/lz4"
)
type compressWriter struct {
writer io.Writer
// data uncompressed
data []byte
// data position
pos int
// data compressed
zdata []byte
}
// NewCompressWriter wrap the io.Writer
func NewCompressWriter(w io.Writer) *compressWriter {
p := &compressWriter{writer: w}
p.data = make([]byte, BlockMaxSize, BlockMaxSize)
zlen := lz4.CompressBound(BlockMaxSize) + HeaderSize
p.zdata = make([]byte, zlen, zlen)
return p
}
func (cw *compressWriter) Write(buf []byte) (int, error) {
var n int
for len(buf) > 0 {
// Accumulate the data to be compressed.
m := copy(cw.data[cw.pos:], buf)
cw.pos += m
buf = buf[m:]
if cw.pos == len(cw.data) {
err := cw.Flush()
if err != nil {
return n, err
}
}
n += m
}
return n, nil
}
func (cw *compressWriter) Flush() (err error) {
if cw.pos == 0 {
return
}
// write the headers
compressedSize, err := lz4.Encode(cw.zdata[HeaderSize:], cw.data[:cw.pos])
if err != nil {
return err
}
compressedSize += CompressHeaderSize
// fill the header, compressed_size_32 + uncompressed_size_32
cw.zdata[16] = LZ4
binary.LittleEndian.PutUint32(cw.zdata[17:], uint32(compressedSize))
binary.LittleEndian.PutUint32(cw.zdata[21:], uint32(cw.pos))
// fill the checksum
checkSum := cityhash102.CityHash128(cw.zdata[16:], uint32(compressedSize))
binary.LittleEndian.PutUint64(cw.zdata[0:], checkSum.Lower64())
binary.LittleEndian.PutUint64(cw.zdata[8:], checkSum.Higher64())
cw.writer.Write(cw.zdata[:compressedSize+ChecksumSize])
if w, ok := cw.writer.(WriteFlusher); ok {
err = w.Flush()
}
cw.pos = 0
return
}
|
