block.go

package data

import (
	"database/sql/driver"
	"fmt"
	"io"
	"reflect"
	"strings"

	"github.com/kshvakov/clickhouse/lib/binary"
	"github.com/kshvakov/clickhouse/lib/column"
	wb "github.com/kshvakov/clickhouse/lib/writebuffer"
)

type offset [][]int

type Block struct {
	Values     [][]interface{}
	Columns    []column.Column
	NumRows    uint64
	NumColumns uint64
	offsets    []offset
	buffers    []*buffer
	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 {
		names = append(names, column.Name())
	}
	return names
}

func (block *Block) Read(serverInfo *ServerInfo, decoder *binary.Decoder) (err error) {
	if err = block.info.read(decoder); err != nil {
		return err
	}

	if block.NumColumns, err = decoder.Uvarint(); err != nil {
		return err
	}
	if block.NumRows, err = decoder.Uvarint(); err != nil {
		return err
	}
	block.Values = make([][]interface{}, block.NumColumns)
	if block.NumRows > 10 {
		for i := 0; i < int(block.NumColumns); i++ {
			block.Values[i] = make([]interface{}, 0, block.NumRows)
		}
	}
	for i := 0; i < int(block.NumColumns); i++ {
		var (
			value      interface{}
			columnName string
			columnType string
		)
		if columnName, err = decoder.String(); err != nil {
			return err
		}
		if columnType, err = decoder.String(); err != nil {
			return err
		}
		c, err := column.Factory(columnName, columnType, serverInfo.Timezone)
		if err != nil {
			return err
		}
		block.Columns = append(block.Columns, c)
		switch column := c.(type) {
		case *column.Array:
			if block.Values[i], err = column.ReadArray(decoder, int(block.NumRows)); err != nil {
				return err
			}
		case *column.Nullable:
			if block.Values[i], err = column.ReadNull(decoder, int(block.NumRows)); err != nil {
				return err
			}
		default:
			for row := 0; row < int(block.NumRows); row++ {
				if value, err = column.Read(decoder); err != nil {
					return err
				}
				block.Values[i] = append(block.Values[i], value)
			}
		}
	}
	return nil
}

func (block *Block) writeArray(column column.Column, value reflect.Value, num, level int) error {
	switch {
	case value.Kind() == reflect.Slice:
		if len(block.offsets[num]) < level {
			block.offsets[num] = append(block.offsets[num], []int{value.Len()})
		} else {
			block.offsets[num][level-1] = append(
				block.offsets[num][level-1],
				block.offsets[num][level-1][len(block.offsets[num][level-1])-1]+value.Len(),
			)
		}
		for i := 0; i < value.Len(); i++ {
			if err := block.writeArray(column, value.Index(i), num, level+1); err != nil {
				return err
			}
		}
	default:
		if err := column.Write(block.buffers[num].Column, value.Interface()); err != nil {
			return err
		}
	}
	return nil
}

func (block *Block) AppendRow(args []driver.Value) error {
	if len(block.Columns) != len(args) {
		return fmt.Errorf("block: expected %d arguments (columns: %s), got %d", len(block.Columns), strings.Join(block.ColumnNames(), ", "), len(args))
	}
	block.Reserve()
	{
		block.NumRows++
	}
	for num, c := range block.Columns {
		switch column := c.(type) {
		case *column.Array:
			value := reflect.ValueOf(args[num])
			if value.Kind() != reflect.Slice {
				return fmt.Errorf("unsupported Array(T) type [%T]", value.Interface())
			}
			if err := block.writeArray(c, value, num, 1); err != nil {
				return err
			}
		case *column.Nullable:
			if err := column.WriteNull(block.buffers[num].Offset, block.buffers[num].Column, args[num]); err != nil {
				return err
			}
		default:
			if err := column.Write(block.buffers[num].Column, args[num]); err != nil {
				return err
			}
		}
	}
	return nil
}

func (block *Block) Reserve() {
	if len(block.buffers) == 0 {
		block.buffers = make([]*buffer, len(block.Columns))
		block.offsets = make([]offset, len(block.Columns))
		for i := 0; i < len(block.Columns); i++ {
			var (
				offsetBuffer = wb.New(wb.InitialSize)
				columnBuffer = wb.New(wb.InitialSize)
			)
			block.buffers[i] = &buffer{
				Offset:       binary.NewEncoder(offsetBuffer),
				Column:       binary.NewEncoder(columnBuffer),
				offsetBuffer: offsetBuffer,
				columnBuffer: columnBuffer,
			}
		}
	}
}

func (block *Block) Reset() {
	block.NumRows = 0
	block.NumColumns = 0
	for _, buffer := range block.buffers {
		buffer.reset()
	}
	{
		block.offsets = nil
		block.buffers = nil
	}
}

func (block *Block) Write(serverInfo *ServerInfo, encoder *binary.Encoder) error {
	if err := block.info.write(encoder); err != nil {
		return err
	}
	encoder.Uvarint(block.NumColumns)
	encoder.Uvarint(block.NumRows)
	defer func() {
		block.NumRows = 0
		for i := range block.offsets {
			block.offsets[i] = offset{}
		}
	}()
	for i, column := range block.Columns {
		encoder.String(column.Name())
		encoder.String(column.CHType())
		if len(block.buffers) == len(block.Columns) {
			for _, offsets := range block.offsets[i] {
				for _, offset := range offsets {
					if err := encoder.UInt64(uint64(offset)); err != nil {
						return err
					}
				}
			}
			if _, err := block.buffers[i].WriteTo(encoder); err != nil {
				return err
			}
		}
	}
	return nil
}

type blockInfo struct {
	num1        uint64
	isOverflows bool
	num2        uint64
	bucketNum   int32
	num3        uint64
}

func (info *blockInfo) read(decoder *binary.Decoder) error {
	var err error
	if info.num1, err = decoder.Uvarint(); err != nil {
		return err
	}
	if info.isOverflows, err = decoder.Bool(); err != nil {
		return err
	}
	if info.num2, err = decoder.Uvarint(); err != nil {
		return err
	}
	if info.bucketNum, err = decoder.Int32(); err != nil {
		return err
	}
	if info.num3, err = decoder.Uvarint(); err != nil {
		return err
	}
	return nil
}

func (info *blockInfo) write(encoder *binary.Encoder) error {
	if err := encoder.Uvarint(1); err != nil {
		return err
	}
	if err := encoder.Bool(info.isOverflows); err != nil {
		return err
	}
	if err := encoder.Uvarint(2); err != nil {
		return err
	}
	if info.bucketNum == 0 {
		info.bucketNum = -1
	}
	if err := encoder.Int32(info.bucketNum); err != nil {
		return err
	}
	if err := encoder.Uvarint(0); err != nil {
		return err
	}
	return nil
}

type buffer struct {
	Offset       *binary.Encoder
	Column       *binary.Encoder
	offsetBuffer *wb.WriteBuffer
	columnBuffer *wb.WriteBuffer
}

func (buf *buffer) WriteTo(w io.Writer) (int64, error) {
	var size int64
	{
		ln, err := buf.offsetBuffer.WriteTo(w)
		if err != nil {
			return size, err
		}
		size += ln
	}
	{
		ln, err := buf.columnBuffer.WriteTo(w)
		if err != nil {
			return size, err
		}
		size += ln
	}
	return size, nil
}

func (buf *buffer) reset() {
	buf.offsetBuffer.Reset()
	buf.columnBuffer.Reset()
}