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], 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 }

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