Default implementation of column.IterableOrderedMap

examples/clickhouse_api/map.go

@@ -20,9 +20,8 @@ package clickhouse_api
 import (
 	"context"
 	"fmt"
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column/orderedmap"
 	"strconv"
-
-	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 )
 
 func MapInsertRead() error {
@@ -108,7 +107,7 @@ func IterableOrderedMapInsertRead() error {
 	}
 	var i int64
 	for i = 0; i < 10; i++ {
-		om := NewOrderedMap()
+		om := &orderedmap.Map[string, string]{}
 		kv1 := strconv.Itoa(int(i))
 		kv2 := strconv.Itoa(int(i + 1))
 		om.Put(kv1, kv1)
@@ -126,7 +125,7 @@ func IterableOrderedMapInsertRead() error {
 		return err
 	}
 	for rows.Next() {
-		var col1 OrderedMap
+		var col1 orderedmap.Map[string, string]
 		if err := rows.Scan(&col1); err != nil {
 			return err
 		}
@@ -135,44 +134,3 @@ func IterableOrderedMapInsertRead() error {
 	rows.Close()
 	return rows.Err()
 }
-
-// OrderedMap is a simple (non thread safe) ordered map
-type OrderedMap struct {
-	Keys   []any
-	Values []any
-}
-
-func NewOrderedMap() column.IterableOrderedMap {
-	return &OrderedMap{}
-}
-
-func (om *OrderedMap) Put(key any, value any) {
-	om.Keys = append(om.Keys, key)
-	om.Values = append(om.Values, value)
-}
-
-func (om *OrderedMap) Iterator() column.MapIterator {
-	return NewOrderedMapIterator(om)
-}
-
-type OrderedMapIter struct {
-	om        *OrderedMap
-	iterIndex int
-}
-
-func NewOrderedMapIterator(om *OrderedMap) column.MapIterator {
-	return &OrderedMapIter{om: om, iterIndex: -1}
-}
-
-func (i *OrderedMapIter) Next() bool {
-	i.iterIndex++
-	return i.iterIndex < len(i.om.Keys)
-}
-
-func (i *OrderedMapIter) Key() any {
-	return i.om.Keys[i.iterIndex]
-}
-
-func (i *OrderedMapIter) Value() any {
-	return i.om.Values[i.iterIndex]
-}

lib/column/orderedmap/orderedmap.go

@@ -0,0 +1,111 @@
+package orderedmap
+
+import (
+	"cmp"
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
+	"slices"
+)
+
+// Map is a simple implementation of [column.IterableOrderedMap] interface.
+// It is intended to be used as a serdes wrapper for map[K]V and not as a general purpose container.
+type Map[K comparable, V any] []entry[K, V]
+
+type entry[K comparable, V any] struct {
+	key   K
+	value V
+}
+
+type iterator[K comparable, V any] struct {
+	om Map[K, V]
+	i  int
+}
+
+func FromMap[M ~map[K]V, K cmp.Ordered, V any](m M) *Map[K, V] {
+	return FromMapFunc(m, cmp.Compare)
+}
+
+func FromMapFunc[M ~map[K]V, K comparable, V any](m M, compare func(K, K) int) *Map[K, V] {
+	om := Map[K, V](make([]entry[K, V], 0, len(m)))
+	for k, v := range m {
+		om.Put(k, v)
+	}
+	slices.SortFunc(om, func(i, j entry[K, V]) int { return compare(i.key, j.key) })
+	return &om
+}
+
+// Collect creates a Map from an iter.Seq2[K,V] iterator.
+func Collect[K cmp.Ordered, V any](seq func(yield func(K, V) bool)) *Map[K, V] {
+	return CollectFunc(seq, cmp.Compare)
+}
+
+// CollectN creates a Map, pre-sized for n entries, from an iter.Seq2[K,V] iterator.
+func CollectN[K cmp.Ordered, V any](seq func(yield func(K, V) bool), n int) *Map[K, V] {
+	return CollectNFunc(seq, n, cmp.Compare)
+}
+
+// CollectFunc creates a Map from an iter.Seq2[K,V] iterator with a custom compare function.
+func CollectFunc[K comparable, V any](seq func(yield func(K, V) bool), compare func(K, K) int) *Map[K, V] {
+	return CollectNFunc(seq, 8, compare)
+}
+
+// CollectNFunc creates a Map, pre-sized for n entries, from an iter.Seq2[K,V] iterator with a custom compare function.
+func CollectNFunc[K comparable, V any](seq func(yield func(K, V) bool), n int, compare func(K, K) int) *Map[K, V] {
+	om := Map[K, V](make([]entry[K, V], 0, n))
+	seq(func(k K, v V) bool {
+		om.Put(k, v)
+		return true
+	})
+	slices.SortFunc(om, func(i, j entry[K, V]) int { return compare(i.key, j.key) })
+	return &om
+}
+
+func (om *Map[K, V]) ToMap() map[K]V {
+	m := make(map[K]V, len(*om))
+	for _, e := range *om {
+		m[e.key] = e.value
+	}
+	return m
+}
+
+// Put is part of [column.IterableOrderedMap] interface, it expects to be called by the driver itself,
+// provides no type safety and expects the keys to be given in order.
+// It is recommended to use [FromMap] and [Collect] to initialize [Map].
+func (om *Map[K, V]) Put(key any, value any) {
+	*om = append(*om, entry[K, V]{key.(K), value.(V)})
+}
+
+// All is an iter.Seq[K,V] iterator that yields all key-value pairs in order.
+func (om *Map[K, V]) All(yield func(k K, v V) bool) {
+	for _, e := range *om {
+		if !yield(e.key, e.value) {
+			return
+		}
+	}
+}
+
+// Keys is an iter.Seq[K] iterator that yields all keys in order.
+func (om *Map[K, V]) Keys(yield func(k K) bool) {
+	for _, e := range *om {
+		if !yield(e.key) {
+			return
+		}
+	}
+}
+
+// Values is an iter.Seq[V] iterator that yields all values in key order.
+func (om *Map[K, V]) Values(yield func(v V) bool) {
+	for _, e := range *om {
+		if !yield(e.value) {
+			return
+		}
+	}
+}
+
+// Iterator is part of [column.IterableOrderedMap] interface, it expects to be called by the driver itself.
+func (om *Map[K, V]) Iterator() column.MapIterator { return &iterator[K, V]{om: *om, i: -1} }
+
+func (i *iterator[K, V]) Next() bool { i.i++; return i.i < len(i.om) }
+
+func (i *iterator[K, V]) Key() any { return i.om[i.i].key }
+
+func (i *iterator[K, V]) Value() any { return i.om[i.i].value }

lib/column/orderedmap/orderedmap_test.go

@@ -0,0 +1,88 @@
+package orderedmap
+
+import (
+	"cmp"
+	"github.com/stretchr/testify/assert"
+	"slices"
+	"testing"
+)
+
+func TestMap(t *testing.T) {
+	m := map[int]int{1: 2, 3: 4, 5: 6, 7: 8, 9: 10, 11: 12, 13: 14, 15: 16, 17: 18, 19: 20}
+	var keys []int
+	var values []int
+	var entries []entry[int, int]
+	for k, v := range m {
+		entries = append(entries, entry[int, int]{k, v})
+	}
+	slices.SortFunc(entries, func(a, b entry[int, int]) int { return cmp.Compare(a.key, b.key) })
+	for _, e := range entries {
+		keys = append(keys, e.key)
+		values = append(values, e.value)
+	}
+
+	// Simple FromMap
+	om1 := FromMap(m)
+
+	// Collect go1.23+ iter.Seq2 iterator
+	om2 := Collect(iterMap(m))
+
+	// Manual fill (e.g. when the map is being read back from ClickHouse)
+	om3 := new(Map[int, int])
+	for _, e := range entries {
+		om3.Put(e.key, e.value)
+	}
+
+	// Custom sort func
+	omR := FromMapFunc(m, func(a, b int) int { return -cmp.Compare(a, b) })
+
+	testMap := func(om *Map[int, int]) {
+		assert.Equal(t, m, om.ToMap())
+		assert.Equal(t, m, collectMap(om.All))
+		assert.Equal(t, keys, collect(om.Keys))
+		assert.Equal(t, values, collect(om.Values))
+		iter, i := om.Iterator(), 0
+		for iter.Next() {
+			assert.Equal(t, keys[i], iter.Key())
+			assert.Equal(t, values[i], iter.Value())
+			i++
+		}
+	}
+
+	testMap(om1)
+	testMap(om2)
+	testMap(om3)
+
+	assert.Equal(t, m, omR.ToMap())
+	keysR := slices.Clone(keys)
+	slices.Reverse(keysR)
+	assert.Equal(t, keysR, collect(omR.Keys))
+}
+
+// go1.23+ helper reimplementations
+func iterMap[K comparable, V any, M ~map[K]V](m M) func(yield func(K, V) bool) {
+	return func(yield func(K, V) bool) {
+		for k, v := range m {
+			if !yield(k, v) {
+				break
+			}
+		}
+	}
+}
+
+func collect[V any](seq func(yield func(V) bool)) (s []V) {
+	seq(func(v V) bool {
+		s = append(s, v)
+		return true
+	})
+	return
+}
+
+func collectMap[K comparable, V any](seq func(yield func(K, V) bool)) (m map[K]V) {
+	m = make(map[K]V)
+	seq(func(k K, v V) bool {
+		m[k] = v
+		return true
+	})
+	return
+}

tests/map_test.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"database/sql/driver"
 	"fmt"
+	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 	"reflect"
 	"testing"
 
@@ -436,16 +437,10 @@ func (om *OrderedMap) KeysUseSlice() []any {
 	return om.keys
 }
 
-func (om *OrderedMap) Iter() MapIter {
+func (om *OrderedMap) Iter() column.MapIterator {
 	return &mapIter{om: om, iterIndex: -1}
 }
 
-type MapIter interface {
-	Next() bool
-	Key() any
-	Value() any
-}
-
 type mapIter struct {
 	om        *OrderedMap
 	iterIndex int