fix: sync with original code reservoir items sketch

sampling/reservoir_items_sketch.go

@@ -20,8 +20,11 @@ package sampling
 import (
 	"encoding/binary"
 	"errors"
+	"math"
 	"math/rand"
+	"slices"
 
+	"github.com/apache/datasketches-go/common"
 	"github.com/apache/datasketches-go/internal"
 )
 
@@ -38,7 +41,10 @@ const (
 	ResizeX8 ResizeFactor = 8
 
 	defaultResizeFactor = ResizeX8
-	minK                = 1
+	minK                = 2
+
+	// smallest sampling array allocated: 16
+	minLgArrItems = 4
 )
 
 // ReservoirItemsSketch provides a uniform random sample of items
@@ -52,26 +58,59 @@ const (
 type ReservoirItemsSketch[T any] struct {
 	k    int   // maximum reservoir size
 	n    int64 // total items seen
-	data []T   // reservoir storage
+	rf   ResizeFactor
+	data []T // reservoir storage
+}
+
+type reservoirItemsSketchOptions struct {
+	resizeFactor ResizeFactor
+}
+
+// ReservoirItemsSketchOptionFunc defines a functional option for configuring reservoirItemsSketchOptions.
+type ReservoirItemsSketchOptionFunc func(*reservoirItemsSketchOptions)
+
+// WithReservoirItemsSketchResizeFactor sets the resize factor for the internal array.
+func WithReservoirItemsSketchResizeFactor(rf ResizeFactor) ReservoirItemsSketchOptionFunc {
+	return func(r *reservoirItemsSketchOptions) {
+		r.resizeFactor = rf
+	}
 }
 
 // NewReservoirItemsSketch creates a new reservoir sketch with the given capacity k.
-func NewReservoirItemsSketch[T any](k int) (*ReservoirItemsSketch[T], error) {
+func NewReservoirItemsSketch[T any](
+	k int, opts ...ReservoirItemsSketchOptionFunc,
+) (*ReservoirItemsSketch[T], error) {
 	if k < minK {
-		return nil, errors.New("k must be at least 1")
+		return nil, errors.New("k must be at least 2")
 	}
 
+	options := &reservoirItemsSketchOptions{
+		resizeFactor: defaultResizeFactor,
+	}
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	ceilingLgK, _ := internal.ExactLog2(common.CeilingPowerOf2(k))
+	initialLgSize := startingSubMultiple(
+		ceilingLgK, int(math.Log2(float64(options.resizeFactor))), minLgArrItems,
+	)
 	return &ReservoirItemsSketch[T]{
 		k:    k,
 		n:    0,
-		data: make([]T, 0, min(k, int(defaultResizeFactor))),
+		rf:   options.resizeFactor,
+		data: make([]T, 0, adjustedSamplingAllocationSize(k, 1<<initialLgSize)),
 	}, nil
 }
 
 // Update adds an item to the sketch using reservoir sampling algorithm.
 func (s *ReservoirItemsSketch[T]) Update(item T) {
 	if s.n < int64(s.k) {
 		// Initial phase: store all items until reservoir is full
+		if s.n >= int64(cap(s.data)) {
+			s.growReservoir()
+		}
+
 		s.data = append(s.data, item)
 	} else {
 		// Steady state: replace with probability k/n
@@ -83,6 +122,11 @@ func (s *ReservoirItemsSketch[T]) Update(item T) {
 	s.n++
 }
 
+func (s *ReservoirItemsSketch[T]) growReservoir() {
+	adjustedSize := adjustedSamplingAllocationSize(s.k, cap(s.data)<<int(s.rf))
+	s.data = slices.Grow(s.data, adjustedSize)
+}
+
 // K returns the maximum reservoir capacity.
 func (s *ReservoirItemsSketch[T]) K() int {
 	return s.k
@@ -112,8 +156,13 @@ func (s *ReservoirItemsSketch[T]) IsEmpty() bool {
 
 // Reset clears the sketch while preserving capacity k.
 func (s *ReservoirItemsSketch[T]) Reset() {
+	ceilingLgK, _ := internal.ExactLog2(common.CeilingPowerOf2(s.k))
+	initialLgSize := startingSubMultiple(
+		ceilingLgK, int(math.Log2(float64(s.rf))), minLgArrItems,
+	)
+
 	s.n = 0
-	s.data = s.data[:0]
+	s.data = make([]T, 0, adjustedSamplingAllocationSize(s.k, 1<<initialLgSize))
 }
 
 // ImplicitSampleWeight returns N/K when in sampling mode, or 1.0 in exact mode.

sampling/reservoir_items_sketch_test.go

@@ -18,8 +18,11 @@
 package sampling
 
 import (
+	"math"
 	"testing"
 
+	"github.com/apache/datasketches-go/common"
+	"github.com/apache/datasketches-go/internal"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -69,63 +72,100 @@ func TestReservoirItemsSketchWithStruct(t *testing.T) {
 
 func TestReservoirItemsSketchInvalidK(t *testing.T) {
 	_, err := NewReservoirItemsSketch[int64](0)
-	assert.Error(t, err)
+	assert.ErrorContains(t, err, "k must be at least 2")
 
-	_, err = NewReservoirItemsSketch[int64](-1)
-	assert.Error(t, err)
+	_, err = NewReservoirItemsSketch[int64](1)
+	assert.ErrorContains(t, err, "k must be at least 2")
 }
 
-func TestReservoirItemsSketchUpdateBelowK(t *testing.T) {
-	sketch, _ := NewReservoirItemsSketch[int64](10)
-
-	for i := int64(1); i <= 5; i++ {
-		sketch.Update(i)
-	}
-
-	assert.Equal(t, int64(5), sketch.N())
-	assert.Equal(t, 5, sketch.NumSamples())
-
-	samples := sketch.Samples()
-	for i := int64(1); i <= 5; i++ {
-		assert.Contains(t, samples, i)
-	}
+func TestReservoirItemsSketch_Update(t *testing.T) {
+	t.Run("BelowKStoresAllItems", func(t *testing.T) {
+		sketch, err := NewReservoirItemsSketch[int64](10)
+		assert.NoError(t, err)
+
+		for i := int64(1); i <= 5; i++ {
+			sketch.Update(i)
+		}
+
+		assert.Equal(t, int64(5), sketch.N())
+		assert.Equal(t, 5, sketch.NumSamples())
+		assert.Equal(t, 1.0, sketch.ImplicitSampleWeight())
+
+		samples := sketch.Samples()
+		for i := int64(1); i <= 5; i++ {
+			assert.Contains(t, samples, i)
+		}
+	})
+
+	t.Run("AtKStoresKItems", func(t *testing.T) {
+		sketch, err := NewReservoirItemsSketch[int64](8)
+		assert.NoError(t, err)
+
+		for i := int64(1); i <= 8; i++ {
+			sketch.Update(i)
+		}
+
+		assert.Equal(t, int64(8), sketch.N())
+		assert.Equal(t, 8, sketch.NumSamples())
+		assert.Equal(t, 1.0, sketch.ImplicitSampleWeight())
+
+		samples := sketch.Samples()
+		for i := int64(1); i <= 8; i++ {
+			assert.Contains(t, samples, i)
+		}
+	})
+
+	t.Run("AboveKMaintainsKAndIncrementsN", func(t *testing.T) {
+		k := 10
+		total := 1000
+
+		sketch, err := NewReservoirItemsSketch[int64](k)
+		assert.NoError(t, err)
+
+		for i := 1; i <= total; i++ {
+			sketch.Update(int64(i))
+		}
+
+		assert.Equal(t, int64(total), sketch.N())
+		assert.Equal(t, k, sketch.NumSamples())
+		assert.Equal(t, float64(total)/float64(k), sketch.ImplicitSampleWeight())
+
+		samples := sketch.Samples()
+		seen := make(map[int64]struct{}, len(samples))
+		for _, sample := range samples {
+			assert.True(t, sample >= 1 && sample <= int64(total))
+			_, exists := seen[sample]
+			assert.False(t, exists)
+			seen[sample] = struct{}{}
+		}
+	})
 }
 
-func TestReservoirItemsSketchUpdateAboveK(t *testing.T) {
-	sketch, _ := NewReservoirItemsSketch[int64](10)
-
-	for i := int64(1); i <= 1000; i++ {
-		sketch.Update(i)
-	}
+func TestReservoirItemsSketchReset(t *testing.T) {
+	k := 1024
 
-	assert.Equal(t, int64(1000), sketch.N())
-	assert.Equal(t, 10, sketch.NumSamples())
-}
+	sketch, err := NewReservoirItemsSketch[int64](k)
+	assert.NoError(t, err)
 
-func TestReservoirItemsSketchReset(t *testing.T) {
-	sketch, _ := NewReservoirItemsSketch[int64](10)
+	ceilingLgK, _ := internal.ExactLog2(common.CeilingPowerOf2(k))
+	initialLgSize := startingSubMultiple(
+		ceilingLgK, int(math.Log2(float64(defaultResizeFactor))), minLgArrItems,
+	)
+	expectedInitialCap := adjustedSamplingAllocationSize(k, 1<<initialLgSize)
 
-	for i := int64(1); i <= 5; i++ {
+	for i := int64(1); i <= int64(expectedInitialCap)+1; i++ {
 		sketch.Update(i)
 	}
 
+	assert.Greater(t, cap(sketch.data), expectedInitialCap)
+
 	sketch.Reset()
+
 	assert.True(t, sketch.IsEmpty())
 	assert.Equal(t, int64(0), sketch.N())
-	assert.Equal(t, 10, sketch.K())
-}
-
-func TestReservoirItemsSketchKEqualsOne(t *testing.T) {
-	// Edge case: k=1 should keep exactly one sample
-	sketch, err := NewReservoirItemsSketch[int64](1)
-	assert.NoError(t, err)
-
-	for i := int64(1); i <= 100; i++ {
-		sketch.Update(i)
-	}
-
-	assert.Equal(t, 1, sketch.NumSamples())
-	assert.Equal(t, int64(100), sketch.N())
+	assert.Equal(t, k, sketch.K())
+	assert.Equal(t, 0, len(sketch.data))
+	assert.Equal(t, expectedInitialCap, cap(sketch.data))
 }
 
 func TestReservoirItemsSketchGetSamplesIsCopy(t *testing.T) {

sampling/utils.go

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sampling
+
+func startingSubMultiple(lgTarget, lgRf, lgMin int) int {
+	if lgTarget <= lgMin {
+		return lgMin
+	}
+	if lgRf == 0 {
+		return lgTarget
+	}
+	return (lgTarget-lgMin)%lgRf + lgMin
+}
+
+// adjustedSamplingAllocationSize checks target sampling allocation is more than
+// 50% of max sampling size. If so, return max sampling size, otherwise passes
+// through the target size.
+func adjustedSamplingAllocationSize(
+	maxSize, resizeTarget int,
+) int {
+	if maxSize-(resizeTarget<<1) < 0 {
+		return maxSize
+	}
+	return resizeTarget
+}