diff --git a/examples/clickhouse_api/map.go b/examples/clickhouse_api/map.go index 432986a378..3b1e287a1c 100644 --- a/examples/clickhouse_api/map.go +++ b/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] -} diff --git a/lib/column/orderedmap/orderedmap.go b/lib/column/orderedmap/orderedmap.go new file mode 100644 index 0000000000..01a99465c4 --- /dev/null +++ b/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 } diff --git a/lib/column/orderedmap/orderedmap_test.go b/lib/column/orderedmap/orderedmap_test.go new file mode 100644 index 0000000000..5dc68d0999 --- /dev/null +++ b/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 +} diff --git a/tests/map_test.go b/tests/map_test.go index 9bd7eac74f..f3be3105dc 100644 --- a/tests/map_test.go +++ b/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