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