go-clickhouse/ch/chschema/column_array.go

package chschema

import (
	"fmt"
	"reflect"

	"github.com/uptrace/go-clickhouse/ch/chproto"
)

type ArrayColumnar interface {
	WriteOffset(wr *chproto.Writer, offset int) int
	WriteData(wr *chproto.Writer) error
}

type ArrayLCStringColumn struct {
	*LCStringColumn
}

func (c ArrayLCStringColumn) Type() reflect.Type {
	return stringSliceType
}

func (c *ArrayLCStringColumn) WriteTo(wr *chproto.Writer) error {
	c.writeData(wr)
	return nil
}

func (c *ArrayLCStringColumn) ReadFrom(rd *chproto.Reader, numRow int) error {
	if numRow == 0 {
		return nil
	}
	return c.readData(rd, numRow)
}

//------------------------------------------------------------------------------

type ArrayColumn struct {
	Column reflect.Value

	typ       reflect.Type
	elem      Columnar
	arrayElem ArrayColumnar
}

var _ Columnar = (*ArrayColumn)(nil)

func NewArrayColumn(typ reflect.Type, chType string, numRow int) Columnar {
	elemType := chArrayElemType(chType)
	if elemType == "" {
		panic(fmt.Errorf("invalid array type: %q (Go type is %s)",
			chType, typ.String()))
	}

	elem := NewColumn(typ.Elem(), elemType, 0)
	var arrayElem ArrayColumnar

	if _, ok := elem.(*LCStringColumn); ok {
		panic("not reached")
	}
	arrayElem, _ = elem.(ArrayColumnar)

	c := &ArrayColumn{
		typ:       reflect.SliceOf(typ),
		elem:      elem,
		arrayElem: arrayElem,
	}

	c.Column = reflect.MakeSlice(c.typ, 0, numRow)

	return c
}

func (c ArrayColumn) Type() reflect.Type {
	return c.typ.Elem()
}

func (c *ArrayColumn) Reset(numRow int) {
	if c.Column.Cap() >= numRow {
		c.Column = c.Column.Slice(0, 0)
	} else {
		c.Column = reflect.MakeSlice(c.typ, 0, numRow)
	}
}

func (c *ArrayColumn) Set(v any) {
	c.Column = reflect.ValueOf(v)
}

func (c *ArrayColumn) Value() any {
	return c.Column.Interface()
}

func (c *ArrayColumn) Nullable(nulls Uint8Column) any {
	panic("not implemented")
}

func (c *ArrayColumn) Len() int {
	return c.Column.Len()
}

func (c *ArrayColumn) Index(idx int) any {
	return c.Column.Index(idx).Interface()
}

func (c ArrayColumn) Slice(s, e int) any {
	return c.Column.Slice(s, e).Interface()
}

func (c *ArrayColumn) ConvertAssign(idx int, v reflect.Value) error {
	v.Set(c.Column.Index(idx))
	return nil
}

func (c *ArrayColumn) AppendValue(v reflect.Value) {
	c.Column = reflect.Append(c.Column, v)
}

func (c *ArrayColumn) ReadFrom(rd *chproto.Reader, numRow int) error {
	if c.Column.Cap() >= numRow {
		c.Column = c.Column.Slice(0, numRow)
	} else {
		c.Column = reflect.MakeSlice(c.typ, numRow, numRow)
	}

	if numRow == 0 {
		return nil
	}

	offsets := make([]int, numRow)
	for i := 0; i < len(offsets); i++ {
		offset, err := rd.Uint64()
		if err != nil {
			return err
		}
		offsets[i] = int(offset)
	}

	if err := c.elem.ReadFrom(rd, offsets[len(offsets)-1]); err != nil {
		return err
	}

	var prev int
	for i, offset := range offsets {
		c.Column.Index(i).Set(reflect.ValueOf(c.elem.Slice(prev, offset)))
		prev = offset
	}

	return nil
}

func (c *ArrayColumn) WriteTo(wr *chproto.Writer) error {
	_ = c.WriteOffset(wr, 0)

	colLen := c.Column.Len()
	for i := 0; i < colLen; i++ {
		// TODO: add SetValue or SetPointer
		c.elem.Set(c.Column.Index(i).Interface())

		var err error
		if c.arrayElem != nil {
			err = c.arrayElem.WriteData(wr)
		} else {
			err = c.elem.WriteTo(wr)
		}
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *ArrayColumn) WriteOffset(wr *chproto.Writer, offset int) int {
	colLen := c.Column.Len()

	for i := 0; i < colLen; i++ {
		el := c.Column.Index(i)
		offset += el.Len()
		wr.Uint64(uint64(offset))
	}

	if c.arrayElem == nil {
		return offset
	}

	offset = 0
	for i := 0; i < colLen; i++ {
		el := c.Column.Index(i)
		c.elem.Set(el.Interface()) // Use SetValue or SetPointer
		offset = c.arrayElem.WriteOffset(wr, offset)
	}

	return offset
}

//------------------------------------------------------------------------------

type StringArrayColumn struct {
	Column     [][]string
	elem       Columnar
	stringElem *StringColumn
	lcElem     *LCStringColumn
}

var _ Columnar = (*StringArrayColumn)(nil)

func NewStringArrayColumn(typ reflect.Type, chType string, numRow int) Columnar {
	if _, funcType := aggFuncNameAndType(chType); funcType != "" {
		chType = funcType
	}
	elemType := chArrayElemType(chType)
	if elemType == "" {
		panic(fmt.Errorf("invalid array type: %q (Go type is %s)",
			chType, typ.String()))
	}

	columnar := NewColumn(typ.Elem(), elemType, 0)
	var stringElem *StringColumn
	var lcElem *LCStringColumn

	switch v := columnar.(type) {
	case *StringColumn:
		stringElem = v
	case *LCStringColumn:
		stringElem = &v.StringColumn
		lcElem = v
		columnar = &ArrayLCStringColumn{v}
	case *EnumColumn:
		stringElem = &v.StringColumn
	default:
		panic(fmt.Errorf("unsupported column: %T", v))
	}

	return &StringArrayColumn{
		Column:     make([][]string, 0, numRow),
		elem:       columnar,
		stringElem: stringElem,
		lcElem:     lcElem,
	}
}

func (c *StringArrayColumn) Reset(numRow int) {
	if cap(c.Column) >= numRow {
		c.Column = c.Column[:0]
	} else {
		c.Column = make([][]string, 0, numRow)
	}
}

func (c *StringArrayColumn) Type() reflect.Type {
	return stringSliceType
}

func (c *StringArrayColumn) Set(v any) {
	c.Column = v.([][]string)
}

func (c *StringArrayColumn) Value() any {
	return c.Column
}

func (c *StringArrayColumn) Nullable(nulls Uint8Column) any {
	panic("not implemented")
}

func (c *StringArrayColumn) Len() int {
	return len(c.Column)
}

func (c *StringArrayColumn) Index(idx int) any {
	return c.Column[idx]
}

func (c StringArrayColumn) Slice(s, e int) any {
	return c.Column[s:e]
}

func (c *StringArrayColumn) ConvertAssign(idx int, v reflect.Value) error {
	v.Set(reflect.ValueOf(c.Column[idx]))
	return nil
}

func (c *StringArrayColumn) AppendValue(v reflect.Value) {
	c.Column = append(c.Column, v.Interface().([]string))
}

func (c *StringArrayColumn) ReadFrom(rd *chproto.Reader, numRow int) error {
	if numRow == 0 {
		return nil
	}

	if cap(c.Column) >= numRow {
		c.Column = c.Column[:numRow]
	} else {
		c.Column = make([][]string, numRow)
	}

	if c.lcElem != nil {
		if err := c.lcElem.readPrefix(rd, numRow); err != nil {
			return err
		}
	}

	offsets := make([]int, numRow)

	for i := 0; i < len(offsets); i++ {
		offset, err := rd.Uint64()
		if err != nil {
			return err
		}
		offsets[i] = int(offset)
	}

	if err := c.elem.ReadFrom(rd, offsets[len(offsets)-1]); err != nil {
		return err
	}

	var prev int
	for i, offset := range offsets {
		c.Column[i] = c.stringElem.Column[prev:offset]
		prev = offset
	}

	return nil
}

func (c *StringArrayColumn) WriteTo(wr *chproto.Writer) error {
	if c.lcElem != nil {
		c.lcElem.writePrefix(wr)
	}

	_ = c.WriteOffset(wr, 0)
	return c.WriteData(wr)
}

var _ ArrayColumnar = (*StringArrayColumn)(nil)

func (c *StringArrayColumn) WriteOffset(wr *chproto.Writer, offset int) int {
	for _, el := range c.Column {
		offset += len(el)
		wr.Uint64(uint64(offset))
	}
	return offset
}

func (c *StringArrayColumn) WriteData(wr *chproto.Writer) error {
	for _, ss := range c.Column {
		c.stringElem.Column = ss
		if err := c.elem.WriteTo(wr); err != nil {
			return err
		}
	}
	return nil
}