From efa54292e7a7ab7e8c162863a67add504e189eab Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Fri, 2 Aug 2024 16:04:46 +0300
Subject: [PATCH] feat(ch-gen-col): generate `byteRange` implementation

---
 proto/block.go                      |  6 +++++
 proto/cmd/ch-gen-col/test.go.tmpl   | 34 +++++++++++++++++++++++++++++
 proto/cmd/ch-gen-col/unsafe.go.tmpl | 32 +++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/proto/block.go b/proto/block.go
index 7d698f06..6483b9f8 100644
--- a/proto/block.go
+++ b/proto/block.go
@@ -192,6 +192,12 @@ type BlockWriter struct {
 	Vec net.Buffers
 }
 
+// byteRange represents columns that stored in the memory as-is.
+type byteRange interface {
+	isByteRange() bool
+	appendSlice(net.Buffers) net.Buffers
+}
+
 func (b Block) WriteBlock(w *BlockWriter, version int, input []InputColumn) error {
 	prevLen := len(w.Buf.Buf)
 	cutBuffer := func() {
diff --git a/proto/cmd/ch-gen-col/test.go.tmpl b/proto/cmd/ch-gen-col/test.go.tmpl
index 8747cfeb..05369a83 100644
--- a/proto/cmd/ch-gen-col/test.go.tmpl
+++ b/proto/cmd/ch-gen-col/test.go.tmpl
@@ -77,6 +77,40 @@ func Test{{ .Type }}_DecodeColumn(t *testing.T) {
 	})
 }
 
+
+func Test{{ .Type }}_byteRange(t *testing.T) {
+	t.Parallel()
+
+	const rows = 50
+	var data {{ .Type }}
+	br, ok := any(data).(byteRange)
+	if !ok || !br.isByteRange() {
+		t.Skipf("%T could not be written as-is", data)
+		return
+	}
+
+	for i := 0; i < rows; i++ {
+		{{- if .DateTime }}
+		data.Data = append(data.Data, {{ .New }}(i))
+		{{- else if .Date }}
+		data = append(data, {{ .New }}(i))
+		{{- else }}
+		v := {{ .New }}(i)
+		data.Append(v)
+		require.Equal(t, v, data.Row(i))
+		{{- end }}
+	}
+
+	var expect Buffer
+	data.EncodeColumn(&expect)
+
+	var got Buffer
+	for _, part := range any(data).(byteRange).appendSlice(nil) {
+		got.Buf = append(got.Buf, part...)
+	}
+	require.Equal(t, expect.Buf, got.Buf)
+}
+
 {{- if not .Time }}
 func Test{{ .Type }}Array(t *testing.T) {
 	const rows = 50
diff --git a/proto/cmd/ch-gen-col/unsafe.go.tmpl b/proto/cmd/ch-gen-col/unsafe.go.tmpl
index 3971891e..8a2e5f23 100644
--- a/proto/cmd/ch-gen-col/unsafe.go.tmpl
+++ b/proto/cmd/ch-gen-col/unsafe.go.tmpl
@@ -5,6 +5,7 @@
 package proto
 
 import (
+	"net"
 	"unsafe"
 
 	"github.com/go-faster/errors"
@@ -68,3 +69,34 @@ func (c {{ .Type }}) EncodeColumn(b *Buffer) {
 	dst := b.Buf[offset:]
 	copy(dst, src)
 }
+
+func ({{ .Type }}) isByteRange() bool { return true }
+
+func (c {{ .Type }}) appendSlice(buf net.Buffers) net.Buffers {
+	{{- if .DateTime }}
+	v := c.Data
+	{{- else }}
+	v := c
+	{{- end }}
+	if len(v) == 0 {
+		return buf
+	}
+	{{- if .SingleByte }}
+		src := *(*[]byte)(unsafe.Pointer(&v))
+	{{- else }}
+		{{- if .FixedStr }}
+			const size = {{ .Bytes }}
+		{{- else }}
+			const size = {{ .Bits }} / 8
+		{{- end }}
+
+		s := *(*slice)(unsafe.Pointer(&v))
+		{{- if not .SingleByte }}
+			s.Len *= size
+			s.Cap *= size
+		{{- end }}
+
+		src := *(*[]byte)(unsafe.Pointer(&s))
+	{{- end }}
+	return append(buf, src)
+}