Improve ordered map performance (#1151)

bind.go

@@ -20,12 +20,12 @@ package clickhouse
 import (
 	std_driver "database/sql/driver"
 	"fmt"
-	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 	"reflect"
 	"regexp"
 	"strings"
 	"time"
 
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 )
 
@@ -320,6 +320,22 @@ func format(tz *time.Location, scale TimeUnit, v any) (string, error) {
 			values = append(values, fmt.Sprintf("%s, %s", name, val))
 		}
 
+		return "map(" + strings.Join(values, ", ") + ")", nil
+	case column.OrderedMapV2:
+		values := make([]string, 0)
+		for key := range v.Keys() {
+			name, err := format(tz, scale, key)
+			if err != nil {
+				return "", err
+			}
+			value, _ := v.Get(key)
+			val, err := format(tz, scale, value)
+			if err != nil {
+				return "", err
+			}
+			values = append(values, fmt.Sprintf("%s, %s", name, val))
+		}
+
 		return "map(" + strings.Join(values, ", ") + ")", nil
 	}
 	switch v := reflect.ValueOf(v); v.Kind() {

lib/column/map.go

@@ -20,10 +20,11 @@ package column
 import (
 	"database/sql/driver"
 	"fmt"
-	"github.com/ClickHouse/ch-go/proto"
 	"reflect"
 	"strings"
 	"time"
+
+	"github.com/ClickHouse/ch-go/proto"
 )
 
 // https://github.com/ClickHouse/ClickHouse/blob/master/src/Columns/ColumnMap.cpp
@@ -42,6 +43,12 @@ type OrderedMap interface {
 	Keys() <-chan any
 }
 
+type OrderedMapV2 interface {
+	Get(key any) (any, bool)
+	Put(key any, value any)
+	Keys() []any
+}
+
 func (col *Map) Reset() {
 	col.keys.Reset()
 	col.values.Reset()
@@ -94,6 +101,13 @@ func (col *Map) ScanRow(dest any, i int) error {
 		value.Set(col.row(i))
 		return nil
 	}
+	if om, ok := dest.(OrderedMapV2); ok {
+		keys, values := col.orderedRow(i)
+		for i := range keys {
+			om.Put(keys[i], values[i])
+		}
+		return nil
+	}
 	if om, ok := dest.(OrderedMap); ok {
 		keys, values := col.orderedRow(i)
 		for i := range keys {
@@ -163,6 +177,29 @@ func (col *Map) AppendRow(v any) error {
 		return nil
 	}
 
+	if orderedMap, ok := v.(OrderedMapV2); ok {
+		var size int64
+		for key := range orderedMap.Keys() {
+			value, ok := orderedMap.Get(key)
+			if !ok {
+				return fmt.Errorf("ordered map has key %v but no corresponding value", key)
+			}
+			size++
+			if err := col.keys.AppendRow(key); err != nil {
+				return err
+			}
+			if err := col.values.AppendRow(value); err != nil {
+				return err
+			}
+		}
+		var prev int64
+		if n := col.offsets.Rows(); n != 0 {
+			prev = col.offsets.col.Row(n - 1)
+		}
+		col.offsets.col.Append(prev + size)
+		return nil
+	}
+
 	if orderedMap, ok := v.(OrderedMap); ok {
 		var size int64
 		for key := range orderedMap.Keys() {

tests/map_test.go

@@ -21,9 +21,10 @@ import (
 	"context"
 	"database/sql/driver"
 	"fmt"
-	"github.com/stretchr/testify/require"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/stretchr/testify/assert"
 )
@@ -354,3 +355,58 @@ func TestMapValuer(t *testing.T) {
 	}
 	require.Equal(t, 1000, i)
 }
+
+func (om *OrderedMap) KeysUseChanNoGo() <-chan any {
+	ch := make(chan any, len(om.keys))
+	for _, key := range om.keys {
+		ch <- key
+	}
+	close(ch)
+	return ch
+}
+
+func (om *OrderedMap) KeysUseSlice() []any {
+	return om.keys
+}
+
+func BenchmarkOrderedMapUseChanGo(b *testing.B) {
+	m := NewOrderedMap()
+	for i := 0; i < 10; i++ {
+		m.Put(i, i)
+	}
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for key := range m.Keys() {
+			_, _ = m.Get(key)
+		}
+	}
+}
+
+func BenchmarkOrderedMapKeysUseChanNoGo(b *testing.B) {
+	m := NewOrderedMap()
+	for i := 0; i < 10; i++ {
+		m.Put(i, i)
+	}
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for key := range m.KeysUseChanNoGo() {
+			_, _ = m.Get(key)
+		}
+	}
+}
+
+func BenchmarkOrderedMapKeysUseSlice(b *testing.B) {
+	m := NewOrderedMap()
+	for i := 0; i < 10; i++ {
+		m.Put(i, i)
+	}
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for key := range m.KeysUseSlice() {
+			_, _ = m.Get(key)
+		}
+	}
+}